wehub-resource-sync
456 lines
16 KiB
Python
456 lines
16 KiB
Python
#!/usr/bin/env python3
|
|
"""Verify GradCache produces identical gradients to full-memory chunked reference.
|
|
|
|
The CORRECT comparison for GradCache is NOT "vanilla full-batch" vs "GradCache chunked",
|
|
because chunked forward can produce slightly different embeddings (bf16 accumulation order).
|
|
|
|
Instead, we compare:
|
|
1. REFERENCE: forward each chunk WITH grad (same chunks!) → concat embeddings →
|
|
compute loss → backward. This keeps all activations in memory.
|
|
2. GRADCACHE: the actual 3-step process from train_contrastors.py.
|
|
|
|
Both process the SAME chunks in the SAME order with the SAME precision, so the
|
|
embeddings are identical. The only difference is HOW gradients are computed:
|
|
reference does one backward through the full graph, GradCache uses surrogate loss.
|
|
If GradCache's chain rule decomposition is correct, gradients must match exactly.
|
|
|
|
Usage:
|
|
CUDA_VISIBLE_DEVICES=0 python training/tests/test_grad_equivalence.py
|
|
"""
|
|
|
|
import copy
|
|
import sys
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from pathlib import Path
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
|
|
|
from train_contrastors import (
|
|
LogitScale,
|
|
clip_loss,
|
|
grad_cache_loss,
|
|
chunk_inputs,
|
|
forward_query,
|
|
forward_doc,
|
|
_clear_rope_deltas,
|
|
)
|
|
|
|
|
|
def chunked_reference_forward_backward(model, q_chunks, d_chunks, logit_scale):
|
|
"""Full-memory reference: forward ALL chunks with grad, then single backward.
|
|
|
|
This computes the EXACT SAME embeddings as GradCache step 1 (same chunks,
|
|
same order, same autocast), but keeps all activations for a single backward.
|
|
"""
|
|
# Forward each chunk with grad, same as GradCache would
|
|
q_embs = []
|
|
with torch.autocast("cuda", dtype=torch.bfloat16):
|
|
for chunk in q_chunks:
|
|
_clear_rope_deltas(model)
|
|
q_embs.append(model(**chunk))
|
|
q_emb = torch.cat(q_embs, dim=0)
|
|
|
|
d_embs = []
|
|
with torch.autocast("cuda", dtype=torch.bfloat16):
|
|
for chunk in d_chunks:
|
|
_clear_rope_deltas(model)
|
|
d_embs.append(model(**chunk))
|
|
d_emb = torch.cat(d_embs, dim=0)
|
|
|
|
# Same loss computation as GradCache step 2
|
|
# gather_enabled=True with no dist → effectively gather_enabled=False
|
|
with torch.autocast("cuda", dtype=torch.bfloat16):
|
|
loss, acc = clip_loss(q_emb, d_emb, logit_scale, gather_enabled=True)
|
|
|
|
loss.backward()
|
|
return loss.detach(), acc.detach()
|
|
|
|
|
|
def make_fake_inputs(processor, batch_size, device):
|
|
from PIL import Image
|
|
import numpy as np
|
|
|
|
queries = [f"What is topic number {i}?" for i in range(batch_size)]
|
|
images = []
|
|
for i in range(batch_size):
|
|
arr = np.random.randint(0, 255, (200, 300, 3), dtype=np.uint8)
|
|
images.append(Image.fromarray(arr))
|
|
|
|
from train_contrastors import process_queries, process_doc_images
|
|
|
|
query_inputs = process_queries(processor, queries)
|
|
doc_inputs = process_doc_images(processor, images)
|
|
|
|
query_inputs = {k: v.to(device) for k, v in query_inputs.items()}
|
|
doc_inputs = {k: v.to(device) for k, v in doc_inputs.items()}
|
|
return query_inputs, doc_inputs
|
|
|
|
|
|
def collect_grads(model, logit_scale):
|
|
"""Collect gradients, cast to fp32 for comparison."""
|
|
names, grads = [], []
|
|
for n, p in model.named_parameters():
|
|
if p.requires_grad:
|
|
names.append(n)
|
|
grads.append(p.grad.clone().float() if p.grad is not None else None)
|
|
for n, p in logit_scale.named_parameters():
|
|
names.append(f"logit_scale.{n}")
|
|
grads.append(p.grad.clone().float() if p.grad is not None else None)
|
|
return names, grads
|
|
|
|
|
|
def compare_gradients(grads_ref, grads_gc, names, verbose=True):
|
|
"""Compare gradients. Returns (cosine, rel_l2, max_rel_diff)."""
|
|
flat_ref, flat_gc = [], []
|
|
max_rel = 0.0
|
|
n_compared = 0
|
|
|
|
if verbose:
|
|
print(f"\n{'Parameter':<60} {'MaxDiff':>12} {'MeanDiff':>12} {'RelDiff':>12}")
|
|
print("-" * 100)
|
|
|
|
for name, gr, ggc in zip(names, grads_ref, grads_gc):
|
|
if gr is None and ggc is None:
|
|
continue
|
|
if gr is None or ggc is None:
|
|
print(f" MISMATCH: {name} — one grad is None")
|
|
return 0.0, float("inf"), float("inf")
|
|
if gr.abs().max().item() == 0 and ggc.abs().max().item() == 0:
|
|
continue
|
|
|
|
flat_ref.append(gr.flatten())
|
|
flat_gc.append(ggc.flatten())
|
|
n_compared += 1
|
|
|
|
abs_diff = (gr - ggc).abs()
|
|
max_diff = abs_diff.max().item()
|
|
mean_diff = abs_diff.mean().item()
|
|
scale = gr.abs().mean().item()
|
|
rel_diff = mean_diff / max(scale, 1e-12)
|
|
max_rel = max(max_rel, rel_diff)
|
|
|
|
if verbose:
|
|
print(f"{name:<60} {max_diff:>12.6e} {mean_diff:>12.6e} {rel_diff:>12.6e}")
|
|
|
|
if not flat_ref:
|
|
print("No non-zero gradients to compare.")
|
|
return 1.0, 0.0, 0.0
|
|
|
|
ref_cat = torch.cat(flat_ref)
|
|
gc_cat = torch.cat(flat_gc)
|
|
cosine = F.cosine_similarity(ref_cat.unsqueeze(0), gc_cat.unsqueeze(0)).item()
|
|
l2_diff = (ref_cat - gc_cat).norm().item()
|
|
rel_l2 = l2_diff / max(ref_cat.norm().item(), 1e-12)
|
|
|
|
if verbose:
|
|
print(f"\nCompared {n_compared} parameter groups")
|
|
print(f" Cosine similarity: {cosine:.10f}")
|
|
print(f" Relative L2 diff: {rel_l2:.6e}")
|
|
print(f" Max per-param rel: {max_rel:.6e}")
|
|
|
|
return cosine, rel_l2, max_rel
|
|
|
|
|
|
def run_test(
|
|
label,
|
|
model,
|
|
model_state,
|
|
ls_state,
|
|
query_inputs,
|
|
doc_inputs,
|
|
chunk_size,
|
|
device,
|
|
verbose=True,
|
|
):
|
|
"""Compare chunked-reference vs GradCache for a given chunk_size."""
|
|
|
|
# --- Chunked reference (full-memory backward) ---
|
|
model.load_state_dict(model_state)
|
|
ls_ref = LogitScale(init_value=1 / 0.07).to(device)
|
|
ls_ref.load_state_dict(ls_state)
|
|
model.zero_grad()
|
|
ls_ref.zero_grad()
|
|
|
|
# Create fresh chunks for each path (can't share autograd graphs)
|
|
q_chunks_ref = chunk_inputs(
|
|
{k: v.clone() for k, v in query_inputs.items()}, chunk_size
|
|
)
|
|
d_chunks_ref = chunk_inputs(
|
|
{k: v.clone() for k, v in doc_inputs.items()}, chunk_size
|
|
)
|
|
|
|
# Seed before reference forward so dropout masks are deterministic
|
|
torch.manual_seed(42)
|
|
torch.cuda.manual_seed_all(42)
|
|
loss_ref, _ = chunked_reference_forward_backward(
|
|
model, q_chunks_ref, d_chunks_ref, ls_ref
|
|
)
|
|
names, grads_ref = collect_grads(model, ls_ref)
|
|
|
|
# --- GradCache (actual function from train_contrastors.py) ---
|
|
model.load_state_dict(model_state)
|
|
ls_gc = LogitScale(init_value=1 / 0.07).to(device)
|
|
ls_gc.load_state_dict(ls_state)
|
|
model.zero_grad()
|
|
ls_gc.zero_grad()
|
|
|
|
q_chunks_gc = chunk_inputs(
|
|
{k: v.clone() for k, v in query_inputs.items()}, chunk_size
|
|
)
|
|
d_chunks_gc = chunk_inputs(
|
|
{k: v.clone() for k, v in doc_inputs.items()}, chunk_size
|
|
)
|
|
|
|
# Same seed → GradCache step 1 (no-grad forward) uses the same dropout masks
|
|
# as the reference. Step 3 (replay) uses RandContext to reproduce them.
|
|
torch.manual_seed(42)
|
|
torch.cuda.manual_seed_all(42)
|
|
loss_gc, acc_gc = grad_cache_loss(
|
|
model=model,
|
|
query_chunks=q_chunks_gc,
|
|
doc_chunks=d_chunks_gc,
|
|
logit_scale=ls_gc,
|
|
query_process_fn=forward_query,
|
|
doc_process_fn=forward_doc,
|
|
)
|
|
_, grads_gc = collect_grads(model, ls_gc)
|
|
|
|
if verbose:
|
|
print(f"\n Loss reference: {loss_ref.item():.8f}")
|
|
print(f" Loss GradCache: {loss_gc.item():.8f}")
|
|
print(f" Loss diff: {abs(loss_ref.item() - loss_gc.item()):.2e}")
|
|
|
|
cosine, rel_l2, max_rel = compare_gradients(
|
|
grads_ref, grads_gc, names, verbose=verbose
|
|
)
|
|
return cosine, rel_l2, loss_ref.item(), loss_gc.item()
|
|
|
|
|
|
def main():
|
|
device = torch.device("cuda:0")
|
|
batch_size = 4
|
|
|
|
print("Loading model + processor...")
|
|
from models.biqwen3 import BiQwen3
|
|
from transformers import AutoProcessor
|
|
from peft import LoraConfig, get_peft_model
|
|
|
|
model_name = "Qwen/Qwen3-VL-Embedding-2B"
|
|
base_model = BiQwen3.from_pretrained(model_name, dtype=torch.bfloat16)
|
|
processor = AutoProcessor.from_pretrained(model_name)
|
|
patch_size = processor.image_processor.patch_size
|
|
merge_size = processor.image_processor.merge_size
|
|
tile = patch_size * merge_size
|
|
processor.image_processor.max_pixels = 256 * tile * tile
|
|
processor.image_processor.size["longest_edge"] = (
|
|
processor.image_processor.max_pixels
|
|
)
|
|
processor.tokenizer.padding_side = "left"
|
|
|
|
print("Creating fake data...")
|
|
query_inputs, doc_inputs = make_fake_inputs(processor, batch_size, device)
|
|
|
|
results = {}
|
|
|
|
# =========================================================================
|
|
# PART A: dropout=0 (isolates GradCache chain-rule math)
|
|
# =========================================================================
|
|
print("\n" + "=" * 80)
|
|
print("PART A: lora_dropout=0.0 — tests GradCache chain-rule decomposition")
|
|
print("=" * 80)
|
|
|
|
lora_config = LoraConfig(
|
|
r=8,
|
|
lora_alpha=8,
|
|
target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
|
|
lora_dropout=0.0,
|
|
task_type="FEATURE_EXTRACTION",
|
|
)
|
|
model = get_peft_model(copy.deepcopy(base_model), lora_config).to(device)
|
|
model.train()
|
|
model_state = copy.deepcopy(model.state_dict())
|
|
ls_state = LogitScale(init_value=1 / 0.07).state_dict()
|
|
|
|
for cs in [batch_size, 2, 1]:
|
|
label = f"A_chunk{cs}"
|
|
print(f"\n--- chunk_size={cs} {'(degenerate)' if cs == batch_size else ''} ---")
|
|
cos, rl2, _, _ = run_test(
|
|
label, model, model_state, ls_state, query_inputs, doc_inputs, cs, device
|
|
)
|
|
results[label] = cos
|
|
|
|
# =========================================================================
|
|
# PART B: dropout=0.05 — tests RandContext RNG replay
|
|
# =========================================================================
|
|
print("\n" + "=" * 80)
|
|
print("PART B: lora_dropout=0.05 — tests RandContext RNG state replay")
|
|
print("=" * 80)
|
|
|
|
lora_config_drop = LoraConfig(
|
|
r=8,
|
|
lora_alpha=8,
|
|
target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
|
|
lora_dropout=0.05,
|
|
task_type="FEATURE_EXTRACTION",
|
|
)
|
|
model_drop = get_peft_model(copy.deepcopy(base_model), lora_config_drop).to(device)
|
|
model_drop.train()
|
|
model_drop_state = copy.deepcopy(model_drop.state_dict())
|
|
|
|
for cs in [batch_size, 2, 1]:
|
|
label = f"B_drop_chunk{cs}"
|
|
print(
|
|
f"\n--- chunk_size={cs} with dropout {'(degenerate)' if cs == batch_size else ''} ---"
|
|
)
|
|
cos, rl2, _, _ = run_test(
|
|
label,
|
|
model_drop,
|
|
model_drop_state,
|
|
ls_state,
|
|
query_inputs,
|
|
doc_inputs,
|
|
cs,
|
|
device,
|
|
)
|
|
results[label] = cos
|
|
|
|
# =========================================================================
|
|
# PART C: clip_loss label correctness (unit tests, no model needed)
|
|
# =========================================================================
|
|
print("\n" + "=" * 80)
|
|
print("PART C: clip_loss label arithmetic")
|
|
print("=" * 80)
|
|
|
|
# C1: Basic case — no hard negatives, no gather
|
|
print("\n--- Test C1: basic labels (no hard negs, no gather) ---")
|
|
N, D = 4, 8
|
|
q = F.normalize(torch.randn(N, D, device=device), dim=-1)
|
|
d = F.normalize(torch.randn(N, D, device=device), dim=-1)
|
|
ls = LogitScale(init_value=1 / 0.07).to(device)
|
|
loss, acc = clip_loss(q, d, ls, gather_enabled=False)
|
|
# Labels should be [0, 1, 2, 3] — each query matches its own doc
|
|
sim = ls(q @ d.T)
|
|
expected_labels = torch.arange(N, device=device)
|
|
expected_loss = F.cross_entropy(sim, expected_labels)
|
|
c1_ok = abs(loss.item() - expected_loss.item()) < 1e-5
|
|
print(
|
|
f" loss={loss.item():.6f} expected={expected_loss.item():.6f} "
|
|
f"diff={abs(loss.item() - expected_loss.item()):.2e} [{'PASS' if c1_ok else 'FAIL'}]"
|
|
)
|
|
results["C1_basic_labels"] = 1.0 if c1_ok else 0.0
|
|
|
|
# C2: Hard negatives — 2 hard negs per query, docs interleaved
|
|
print("\n--- Test C2: hard negative labels (2 negs/query) ---")
|
|
num_hard_neg = 2
|
|
docs_per_q = 1 + num_hard_neg # 3
|
|
# doc layout: [pos0, neg0a, neg0b, pos1, neg1a, neg1b, pos2, ..., pos3, ...]
|
|
d_hn = F.normalize(torch.randn(N * docs_per_q, D, device=device), dim=-1)
|
|
loss_hn, _ = clip_loss(q, d_hn, ls, gather_enabled=False)
|
|
# Labels should point to positive positions: [0, 3, 6, 9]
|
|
sim_hn = ls(q @ d_hn.T)
|
|
expected_hn_labels = torch.arange(N, device=device) * docs_per_q
|
|
expected_hn_loss = F.cross_entropy(sim_hn, expected_hn_labels)
|
|
c2_ok = abs(loss_hn.item() - expected_hn_loss.item()) < 1e-5
|
|
print(
|
|
f" labels={expected_hn_labels.tolist()} loss={loss_hn.item():.6f} "
|
|
f"expected={expected_hn_loss.item():.6f} "
|
|
f"diff={abs(loss_hn.item() - expected_hn_loss.item()):.2e} [{'PASS' if c2_ok else 'FAIL'}]"
|
|
)
|
|
results["C2_hard_neg_labels"] = 1.0 if c2_ok else 0.0
|
|
|
|
# C3: Verify assertion fires for bad divisibility
|
|
print("\n--- Test C3: assertion for bad doc/query ratio ---")
|
|
d_bad = F.normalize(torch.randn(N + 1, D, device=device), dim=-1)
|
|
c3_ok = False
|
|
try:
|
|
clip_loss(q, d_bad, ls, gather_enabled=False)
|
|
print(" ERROR: no assertion raised!")
|
|
except AssertionError as e:
|
|
c3_ok = True
|
|
print(f" AssertionError raised as expected: {e}")
|
|
print(" [PASS]")
|
|
results["C3_divisibility_assert"] = 1.0 if c3_ok else 0.0
|
|
|
|
# =========================================================================
|
|
# PART D: _clear_rope_deltas prevents stale state
|
|
# =========================================================================
|
|
print("\n" + "=" * 80)
|
|
print("PART D: _clear_rope_deltas prevents image→text state leakage")
|
|
print("=" * 80)
|
|
|
|
# Use the no-dropout model for this test
|
|
model.load_state_dict(model_state)
|
|
model.eval()
|
|
|
|
# Forward an image batch to populate rope_deltas
|
|
with torch.no_grad():
|
|
with torch.autocast("cuda", dtype=torch.bfloat16):
|
|
_clear_rope_deltas(model)
|
|
_ = model(**{k: v[:2] for k, v in doc_inputs.items()})
|
|
|
|
# Check that rope_deltas is set on the inner model
|
|
inner = model
|
|
while hasattr(inner, "model"):
|
|
inner = inner.model
|
|
has_rope = hasattr(inner, "rope_deltas") and inner.rope_deltas is not None
|
|
print(f" rope_deltas set after image forward: {has_rope}")
|
|
|
|
# Now forward text — WITHOUT clearing, this should fail or give wrong results
|
|
# if rope_deltas shape mismatches the text batch
|
|
# First, demonstrate that clearing prevents the issue:
|
|
_clear_rope_deltas(model)
|
|
cleared = not hasattr(inner, "rope_deltas") or inner.rope_deltas is None
|
|
print(f" rope_deltas cleared after _clear_rope_deltas: {cleared}")
|
|
|
|
d_ok = False
|
|
try:
|
|
with torch.no_grad():
|
|
with torch.autocast("cuda", dtype=torch.bfloat16):
|
|
_clear_rope_deltas(model)
|
|
_ = model(**{k: v[:3] for k, v in query_inputs.items()})
|
|
d_ok = True
|
|
print(" Text forward after clear succeeded: [PASS]")
|
|
except Exception as e:
|
|
print(f" Text forward after clear failed: {e} [FAIL]")
|
|
results["D_rope_deltas"] = 1.0 if d_ok else 0.0
|
|
|
|
model.train()
|
|
|
|
# =========================================================================
|
|
# Summary
|
|
# =========================================================================
|
|
print("\n" + "=" * 80)
|
|
print("SUMMARY")
|
|
print("=" * 80)
|
|
|
|
# Threshold: the reference and GradCache should produce nearly identical gradients
|
|
# since they compute the same function. Only bf16 non-determinism can cause diffs.
|
|
cos_threshold = 0.999
|
|
|
|
all_pass = True
|
|
for name, cos in results.items():
|
|
ok = cos > cos_threshold
|
|
if not ok:
|
|
all_pass = False
|
|
print(f" {name:<25} cosine={cos:.10f} [{'PASS' if ok else 'FAIL'}]")
|
|
|
|
print(f"\n Threshold: cosine > {cos_threshold}")
|
|
if all_pass:
|
|
print("\n ALL TESTS PASSED.")
|
|
print(" GradCache's surrogate backward matches full-memory backward.")
|
|
print(" RandContext correctly replays dropout RNG states.")
|
|
else:
|
|
print("\n SOME TESTS FAILED — GradCache produces incorrect gradients.")
|
|
for name, cos in results.items():
|
|
if cos <= cos_threshold:
|
|
print(f" {name}: cosine={cos:.10f}")
|
|
print("=" * 80)
|
|
|
|
return 0 if all_pass else 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|