From c023c2a1186b8e54de89bf0ed70e68ad1e64ffca Mon Sep 17 00:00:00 2001 From: Yongbin Choi Date: Tue, 16 Jun 2026 13:07:09 +0900 Subject: [PATCH 1/4] feat(loss): add WithGradCache wrapper for contrastive losses Introduce a single GradCache wrapper that enables large effective batch sizes under a fixed memory budget for both bi-encoder and late-interaction contrastive losses (Gao et al., 2021). Embeddings are computed in mini-batches without retaining the full activation graph; per-embedding gradients are cached from the wrapped loss and replayed through a backward hook that recomputes each mini-batch with gradients enabled. All scoring and loss math is delegated to the wrapped loss, so the wrapper stays compatible with every pooled or token-level contrastive loss and adds no loss-specific logic. ContrastiveTrainer delegates to the gradcache path when the loss exposes `gradcache_enabled`, leaving the standard non-cached path unchanged. --- colpali_engine/loss/__init__.py | 1 + colpali_engine/loss/gradcache.py | 212 ++++++++++++ colpali_engine/trainer/contrastive_trainer.py | 16 + tests/loss/test_gradcache.py | 316 ++++++++++++++++++ 4 files changed, 545 insertions(+) create mode 100644 colpali_engine/loss/gradcache.py create mode 100644 tests/loss/test_gradcache.py diff --git a/colpali_engine/loss/__init__.py b/colpali_engine/loss/__init__.py index 0e3ecbc2a..b24690a90 100644 --- a/colpali_engine/loss/__init__.py +++ b/colpali_engine/loss/__init__.py @@ -6,6 +6,7 @@ BiPairwiseNegativeCELoss, BiSigmoidLoss, ) +from .gradcache import WithGradCache from .late_interaction_losses import ( ColbertLoss, ColbertModule, diff --git a/colpali_engine/loss/gradcache.py b/colpali_engine/loss/gradcache.py new file mode 100644 index 000000000..1204db760 --- /dev/null +++ b/colpali_engine/loss/gradcache.py @@ -0,0 +1,212 @@ +import contextlib +from functools import partial +from typing import Dict, List, Tuple + +import torch +import torch.nn as nn +import tqdm +from torch.distributed.nn.functional import all_gather +from torch.utils.checkpoint import get_device_states, set_device_states + + +class RandContext: + """ + Captures the RNG state so a forward pass can be replayed deterministically. + + GradCache runs every forward pass twice (once without gradients to build the cache, + once with gradients in the backward hook). Restoring the RNG state guarantees that + stochastic ops such as dropout produce identical outputs across both passes. + """ + + def __init__(self, *tensors: torch.Tensor) -> None: + self.fwd_cpu_state = torch.get_rng_state() + self.fwd_gpu_devices, self.fwd_gpu_states = get_device_states(*tensors) + + def __enter__(self) -> None: + self._fork = torch.random.fork_rng(devices=self.fwd_gpu_devices, enabled=True) + self._fork.__enter__() + torch.set_rng_state(self.fwd_cpu_state) + set_device_states(self.fwd_gpu_devices, self.fwd_gpu_states) + + def __exit__(self, exc_type, exc_val, exc_tb) -> None: + self._fork.__exit__(exc_type, exc_val, exc_tb) + self._fork = None + + +def _is_distributed() -> bool: + return torch.distributed.is_available() and torch.distributed.is_initialized() + + +def _gather_doc_embeddings(doc_embeddings: torch.Tensor, local_batch_size: int) -> Tuple[torch.Tensor, int]: + """ + Only positive documents are gathered, so each rank can use other ranks' positives as + additional in-batch negatives. Explicit hard negatives are intentionally not gathered + (they are consumed locally), matching the standard non-cached loss path. + + Returns the gathered embeddings and the offset of this rank's positives inside them. + Token-level (late-interaction) embeddings carry a sequence dim that may differ across + ranks, so it is padded to the global max; pooled (bi-encoder) embeddings need no padding. + """ + if not _is_distributed() or local_batch_size <= 0: + return doc_embeddings, 0 + + # Late-interaction embeddings are (batch, seq_len, dim); pad the variable sequence length. + if doc_embeddings.dim() == 3: + max_len = torch.tensor(doc_embeddings.size(1), device=doc_embeddings.device) + torch.distributed.all_reduce(max_len, op=torch.distributed.ReduceOp.MAX) + pad = int(max_len.item()) - doc_embeddings.size(1) + if pad > 0: + shape = list(doc_embeddings.shape) + shape[1] = pad + padding = torch.zeros(*shape, device=doc_embeddings.device, dtype=doc_embeddings.dtype) + doc_embeddings = torch.cat((padding, doc_embeddings), dim=1) + + gathered = torch.cat(all_gather(doc_embeddings), dim=0) + offset = torch.distributed.get_rank() * local_batch_size + return gathered, offset + + +class WithGradCache(nn.Module): + """ + GradCache wrapper for contrastive losses (Gao et al., 2021). + + Wraps any bi-encoder or late-interaction loss (e.g. ``BiEncoderLoss``, ``BiNegativeCELoss``, + ``ColbertLoss``, ``ColbertNegativeCELoss``) to enable large effective batch sizes under a + fixed memory budget. Embeddings are computed in mini-batches without retaining the full + activation graph; per-embedding gradients are cached from the loss and replayed through a + backward hook that recomputes each mini-batch with gradients enabled. + + All scoring and loss math is delegated to ``loss``, so this wrapper stays compatible with + every pooled or token-level contrastive loss and adds no loss-specific logic. + + Args: + loss: The contrastive loss module to wrap. Its ``forward`` must accept + ``query_embeddings``, ``doc_embeddings``, an ``offset`` keyword, and optionally + ``neg_doc_embeddings``. + mini_batch_size: Number of items embedded per mini-batch. + show_progress_bar: Show a progress bar while embedding mini-batches. + """ + + def __init__(self, loss: nn.Module, mini_batch_size: int = 32, show_progress_bar: bool = False): + super().__init__() + self.loss = loss + self.mini_batch_size = mini_batch_size + self.show_progress_bar = show_progress_bar + # Read by ContrastiveTrainer to route through the GradCache code path. + self.gradcache_enabled = True + # Toggled by ContrastiveTrainer; only gathers when running distributed. + self.gather_across_processes = True + + @staticmethod + def _autocast_ctx(device_type: str): + """Replay the ambient autocast policy so both forward passes use the same dtype.""" + try: + enabled = torch.is_autocast_enabled(device_type) + dtype = torch.get_autocast_dtype(device_type) + except TypeError: + # torch < 2.4 has no device-type argument; fall back to the per-device queries. + if device_type == "cpu": + enabled, dtype = torch.is_autocast_cpu_enabled(), torch.get_autocast_cpu_dtype() + else: + enabled, dtype = torch.is_autocast_enabled(), torch.get_autocast_gpu_dtype() + if enabled: + return partial(torch.autocast, device_type=device_type, dtype=dtype) + return contextlib.nullcontext + + def _embed_in_minibatches( + self, model: nn.Module, features: Dict[str, torch.Tensor], autocast + ) -> Tuple[List[torch.Tensor], List[RandContext]]: + bsz = features["input_ids"].size(0) + reps: List[torch.Tensor] = [] + rand_states: List[RandContext] = [] + for start in tqdm.trange( + 0, bsz, self.mini_batch_size, desc="Embedding minibatches", disable=not self.show_progress_bar + ): + mini = {k: v[start : start + self.mini_batch_size] for k, v in features.items()} + rand_states.append(RandContext(*mini.values())) + with torch.no_grad(), autocast(): + embeds = model(**mini) + reps.append(embeds.detach().requires_grad_(True)) + return reps, rand_states + + def _compute_loss(self, reps: List[List[torch.Tensor]], num_neg_docs: int, with_backward: bool) -> torch.Tensor: + query_embeddings = torch.cat(reps[0], dim=0) + doc_embeddings = torch.cat(reps[1], dim=0) + gathered_doc, offset = _gather_doc_embeddings( + doc_embeddings, query_embeddings.size(0) if self.gather_across_processes else 0 + ) + + kwargs = {"query_embeddings": query_embeddings, "doc_embeddings": gathered_doc, "offset": offset} + if num_neg_docs: + neg_embeddings = torch.cat(reps[2], dim=0) + neg_embeddings = neg_embeddings.reshape(-1, num_neg_docs, *neg_embeddings.shape[1:]) + kwargs["neg_doc_embeddings"] = neg_embeddings + + loss = self.loss(**kwargs) + if with_backward: + loss.backward() + return loss + + def _backward_hook(self, grad_output, model, branches, rand_states, cache, autocast): + with torch.enable_grad(): + for features, branch_cache, branch_states in zip(branches, cache, rand_states): + bsz = features["input_ids"].size(0) + for i, start in enumerate(range(0, bsz, self.mini_batch_size)): + mini = {k: v[start : start + self.mini_batch_size] for k, v in features.items()} + with branch_states[i], autocast(): + embeds = model(**mini) + # Replay the cached gradient: d(surrogate)/d(params) == d(loss)/d(params). + surrogate = torch.dot(embeds.flatten(), branch_cache[i].flatten()) * grad_output + surrogate.backward() + + def forward( + self, + model: nn.Module, + inputs: Dict[str, torch.Tensor], + query_prefix: str = "query_", + pos_doc_prefix: str = "doc_", + neg_doc_prefix: str = "neg_doc_", + ) -> torch.Tensor: + query_features = {k[len(query_prefix) :]: v for k, v in inputs.items() if k.startswith(query_prefix)} + doc_features = {k[len(pos_doc_prefix) :]: v for k, v in inputs.items() if k.startswith(pos_doc_prefix)} + neg_features = {k[len(neg_doc_prefix) :]: v for k, v in inputs.items() if k.startswith(neg_doc_prefix)} + + # Flatten negatives from (batch, num_negs, ...) to (batch * num_negs, ...) for embedding. + num_neg_docs = 0 + if neg_features: + num_neg_docs = neg_features["input_ids"].size(1) + neg_features = {k: v.reshape(-1, *v.shape[2:]) for k, v in neg_features.items()} + + branches = [query_features, doc_features] + if num_neg_docs: + branches.append(neg_features) + + # First pass: embed every branch in mini-batches without gradients. + # Capture the ambient autocast policy so the backward-hook re-forward (which runs + # outside autocast during backprop) reproduces the same dtype, keeping grads exact. + autocast = self._autocast_ctx(next(iter(inputs.values())).device.type) + reps: List[List[torch.Tensor]] = [] + rand_states: List[List[RandContext]] = [] + for features in branches: + branch_reps, branch_states = self._embed_in_minibatches(model, features, autocast) + reps.append(branch_reps) + rand_states.append(branch_states) + + if not torch.is_grad_enabled(): + return self._compute_loss(reps, num_neg_docs, with_backward=False) + + # Build the gradient cache, then replay it through a backward hook. + loss = self._compute_loss(reps, num_neg_docs, with_backward=True) + cache = [[mini.grad for mini in branch] for branch in reps] + loss = loss.detach().requires_grad_() + loss.register_hook( + partial( + self._backward_hook, + model=model, + branches=branches, + rand_states=rand_states, + cache=cache, + autocast=autocast, + ) + ) + return loss diff --git a/colpali_engine/trainer/contrastive_trainer.py b/colpali_engine/trainer/contrastive_trainer.py index 3d4c80713..2b6656c5b 100644 --- a/colpali_engine/trainer/contrastive_trainer.py +++ b/colpali_engine/trainer/contrastive_trainer.py @@ -192,6 +192,16 @@ def _reshape_neg_doc_outputs(self, neg_doc_outputs, num_neg_docs): return neg_doc_outputs def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): + if getattr(self.loss_func, "gradcache_enabled", False): + if self.compute_symetric_loss: + raise ValueError("GradCache losses do not support compute_symetric_loss.") + if hasattr(self.loss_func, "gather_across_processes"): + self.loss_func.gather_across_processes = self.accelerator.num_processes > 1 and bool( + self.accelerator.sync_gradients + ) + loss = self.loss_func(model, inputs, self.query_prefix, self.pos_prefix, self.neg_prefix) + return (loss, None) if return_outputs else loss + query_inputs = {k[len(self.query_prefix) :]: v for k, v in inputs.items() if k.startswith(self.query_prefix)} query_outputs = model(**query_inputs) # feed only kwargs with 'doc_' prefix @@ -223,6 +233,12 @@ def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=True) raise ValueError("prediction_step is only called with prediction_loss_only=True") with torch.no_grad(): + if getattr(self.loss_func, "gradcache_enabled", False): + # Eval loss is computed per-process; never gather docs across processes here. + if hasattr(self.loss_func, "gather_across_processes"): + self.loss_func.gather_across_processes = False + loss = self.loss_func(model, inputs, self.query_prefix, self.pos_prefix, self.neg_prefix) + return loss, None, None # feed only kwargs with 'doc_' prefix doc_outputs = model(**{k[4:]: v for k, v in inputs.items() if k.startswith("doc")}) query_outputs = model(input_ids=inputs["query_input_ids"], attention_mask=inputs["query_attention_mask"]) diff --git a/tests/loss/test_gradcache.py b/tests/loss/test_gradcache.py new file mode 100644 index 000000000..38dbcfc11 --- /dev/null +++ b/tests/loss/test_gradcache.py @@ -0,0 +1,316 @@ +# ruff: noqa: N806, N812 +import os +import socket + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +import torch.nn as nn + +from colpali_engine.loss import ( + BiEncoderLoss, + BiNegativeCELoss, + ColbertLoss, + ColbertNegativeCELoss, + WithGradCache, +) +from colpali_engine.loss.gradcache import RandContext, _gather_doc_embeddings + + +class ToyEncoder(nn.Module): + """Tiny deterministic encoder: returns (B, S, D) tokens, or (B, D) if pooled.""" + + def __init__(self, vocab: int = 64, dim: int = 16, pooled: bool = False, dropout: float = 0.0): + super().__init__() + self.emb = nn.Embedding(vocab, dim) + self.lin = nn.Linear(dim, dim) + self.drop = nn.Dropout(dropout) + self.pooled = pooled + + def forward(self, input_ids, attention_mask=None): + x = self.drop(self.lin(torch.tanh(self.emb(input_ids)))) + if attention_mask is not None: + x = x * attention_mask.unsqueeze(-1).to(x.dtype) + if self.pooled: + denom = attention_mask.sum(1, keepdim=True).clamp(min=1).to(x.dtype) + x = torch.nn.functional.normalize(x.sum(1) / denom, dim=-1) + return x + + +def make_inputs(batch: int, num_neg: int = 0, seqlen: int = 12, vocab: int = 64, device="cpu"): + """Build query/doc(/neg) inputs with full attention (lengths handled per-test).""" + out = {} + + def field(prefix, extra=None): + shape = (batch, seqlen) if extra is None else (batch, extra, seqlen) + out[f"{prefix}_input_ids"] = torch.randint(1, vocab, shape, device=device) + out[f"{prefix}_attention_mask"] = torch.ones(shape, dtype=torch.long, device=device) + + field("query") + field("doc") + if num_neg: + field("neg_doc", extra=num_neg) + return out + + +def reference_loss(model, loss_fn, inputs, num_neg, doc_embeddings=None, offset=0): + """Standard non-cached path: full-batch forward, then loss. + + ``doc_embeddings``/``offset`` let a distributed caller substitute gathered docs. + """ + q = model(input_ids=inputs["query_input_ids"], attention_mask=inputs["query_attention_mask"]) + d = model(input_ids=inputs["doc_input_ids"], attention_mask=inputs["doc_attention_mask"]) + kwargs = dict(query_embeddings=q, doc_embeddings=doc_embeddings if doc_embeddings is not None else d, offset=offset) + if num_neg: + ids, am = inputs["neg_doc_input_ids"], inputs["neg_doc_attention_mask"] + b, n, s = ids.shape + neg = model(input_ids=ids.view(b * n, s), attention_mask=am.view(b * n, s)) + kwargs["neg_doc_embeddings"] = neg.view(b, n, *neg.shape[1:]) + return loss_fn(**kwargs), d, q.size(0) + + +def grads_of(model): + return {n: p.grad.detach().clone() for n, p in model.named_parameters() if p.grad is not None} + + +def max_rel_grad_diff(g1, g2): + assert g1.keys() == g2.keys() + worst = 0.0 + for n in g1: + a, b = g1[n], g2[n] + denom = a.abs().max().item() + 1e-12 + worst = max(worst, (a - b).abs().max().item() / denom) + return worst + + +# Factories (not instances) so the loss can be rebuilt inside spawned workers. +def _bi_encoder(): + return BiEncoderLoss(temperature=1.0) + + +def _bi_negative(): + return BiNegativeCELoss(temperature=1.0, in_batch_term_weight=0.5) + + +def _colbert(): + return ColbertLoss(temperature=1.0, normalize_scores=False, use_smooth_max=False) + + +def _colbert_negative(): + return ColbertNegativeCELoss( + temperature=1.0, normalize_scores=False, use_smooth_max=False, in_batch_term_weight=0.5 + ) + + +_CASES = [ # (factory, pooled, num_neg, id) + (_bi_encoder, True, 0, "BiEncoderLoss"), + (_bi_negative, True, 3, "BiNegativeCELoss"), + (_colbert, False, 0, "ColbertLoss"), + (_colbert_negative, False, 3, "ColbertNegativeCELoss"), +] + + +class TestGradCacheEquivalence: + """Single-process: GradCache must match the non-cached path bit-for-bit (up to fp noise).""" + + @pytest.mark.parametrize("mini_batch_size", [3, 8, 16], ids=["mbsbatch"]) + @pytest.mark.parametrize("factory,pooled,num_neg,name", _CASES, ids=[c[3] for c in _CASES]) + def test_loss_and_grad_match(self, factory, pooled, num_neg, name, mini_batch_size): + batch = 8 + torch.manual_seed(0) + model = ToyEncoder(pooled=pooled) + torch.manual_seed(1) + inputs = make_inputs(batch, num_neg=num_neg) + loss_fn = factory() + + model.zero_grad(set_to_none=True) + ref_loss, _, _ = reference_loss(model, loss_fn, inputs, num_neg) + ref_loss.backward() + ref_grads = grads_of(model) + + model.zero_grad(set_to_none=True) + gc = WithGradCache(loss_fn, mini_batch_size=mini_batch_size) + gc.gather_across_processes = False + gc_loss = gc(model, inputs) + gc_loss.backward() + gc_grads = grads_of(model) + + assert torch.allclose(ref_loss, gc_loss, atol=1e-6, rtol=0), f"{name}: loss mismatch" + assert max_rel_grad_diff(ref_grads, gc_grads) < 1e-4, f"{name}: grad mismatch" + + def test_eval_no_grad_path_returns_loss_only(self): + """Under ``no_grad`` the wrapper returns the loss with no autograd hook.""" + batch = 6 + torch.manual_seed(0) + model = ToyEncoder(pooled=False) + torch.manual_seed(1) + inputs = make_inputs(batch) + loss_fn = _colbert() + + with torch.no_grad(): + ref_loss, _, _ = reference_loss(model, loss_fn, inputs, num_neg=0) + gc = WithGradCache(loss_fn, mini_batch_size=4) + gc.gather_across_processes = False + gc_loss = gc(model, inputs) + + assert not gc_loss.requires_grad + assert torch.allclose(ref_loss, gc_loss, atol=1e-6, rtol=0) + + def test_no_grad_does_not_touch_param_grads(self): + torch.manual_seed(0) + model = ToyEncoder(pooled=True) + inputs = make_inputs(4) + gc = WithGradCache(_bi_encoder(), mini_batch_size=2) + gc.gather_across_processes = False + with torch.no_grad(): + gc(model, inputs) + assert all(p.grad is None for p in model.parameters()) + + +class TestRandContext: + """RandContext is what makes GradCache's two forward passes agree under dropout.""" + + def test_dropout_is_reproduced(self): + torch.manual_seed(0) + model = ToyEncoder(pooled=False, dropout=0.5).train() + ids = torch.randint(1, 64, (4, 10)) + mask = torch.ones(4, 10, dtype=torch.long) + + ctx = RandContext(ids, mask) + with ctx: + first = model(input_ids=ids, attention_mask=mask) + # Without restoring RNG, a second dropout draw would differ; with the context it must not. + with ctx: + second = model(input_ids=ids, attention_mask=mask) + assert torch.equal(first, second) + + def test_gradcache_with_dropout_is_deterministic_and_finite(self): + """With dropout on, repeated GradCache calls under the same seed must agree exactly.""" + inputs = make_inputs(8, num_neg=0) + + def run(): + torch.manual_seed(0) + model = ToyEncoder(pooled=False, dropout=0.3).train() + torch.manual_seed(123) # fix the stochastic forward + gc = WithGradCache(_colbert(), mini_batch_size=3) + gc.gather_across_processes = False + model.zero_grad(set_to_none=True) + loss = gc(model, inputs) + loss.backward() + return loss.detach(), grads_of(model) + + loss1, g1 = run() + loss2, g2 = run() + assert torch.isfinite(loss1) and all(torch.isfinite(g).all() for g in g1.values()) + assert torch.equal(loss1, loss2) + assert max_rel_grad_diff(g1, g2) == 0.0 + + +# Distributed (gloo / CPU) tests for the all-gather + offset path. +def _find_free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def _dist_worker(rank, world_size, port, case, results): + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(port) + dist.init_process_group("gloo", rank=rank, world_size=world_size) + try: + factory, pooled, num_neg, mbs, seqlen = case + torch.manual_seed(0) # identical params on every rank + model = ToyEncoder(pooled=pooled) + torch.manual_seed(100 + rank) # rank-specific data (and per-rank seqlen) + inputs = make_inputs(6, num_neg=num_neg, seqlen=seqlen[rank]) + loss_fn = factory() + + # Reference: gather docs across ranks exactly like the wrapper does, then full-batch loss. + model.zero_grad(set_to_none=True) + q = model(input_ids=inputs["query_input_ids"], attention_mask=inputs["query_attention_mask"]) + d = model(input_ids=inputs["doc_input_ids"], attention_mask=inputs["doc_attention_mask"]) + d_gathered, offset = _gather_doc_embeddings(d, q.size(0)) + ref_loss, _, _ = reference_loss(model, loss_fn, inputs, num_neg, doc_embeddings=d_gathered, offset=offset) + ref_loss.backward() + ref_grads = grads_of(model) + + model.zero_grad(set_to_none=True) + gc = WithGradCache(loss_fn, mini_batch_size=mbs) + gc.gather_across_processes = True + gc_loss = gc(model, inputs) + gc_loss.backward() + gc_grads = grads_of(model) + + results[rank] = { + "loss_diff": (ref_loss - gc_loss).abs().item(), + "grad_diff": max_rel_grad_diff(ref_grads, gc_grads), + "error": None, + } + except Exception: + import traceback + + results[rank] = {"loss_diff": None, "grad_diff": None, "error": traceback.format_exc()} + finally: + dist.destroy_process_group() + + +def _run_distributed(case, world_size=2): + manager = mp.Manager() + results = manager.dict() + mp.spawn(_dist_worker, args=(world_size, _find_free_port(), case, results), nprocs=world_size, join=True) + return dict(results) + + +@pytest.mark.parametrize( + "factory,pooled,num_neg,mbs,seqlen,name", + [ + (_bi_encoder, True, 0, 4, (12, 12), "BiEncoder-equal-len"), + (_bi_negative, True, 3, 4, (12, 12), "BiNegative-equal-len"), + (_colbert, False, 0, 4, (12, 12), "Colbert-equal-len"), + (_colbert, False, 0, 4, (10, 14), "Colbert-unequal-len-frontpad"), + ], + ids=lambda v: v if isinstance(v, str) else None, +) +def test_gradcache_distributed_matches_reference(factory, pooled, num_neg, mbs, seqlen, name): + """2-process gloo: GradCache == distributed non-cached path (gather + offset, incl. front-pad).""" + results = _run_distributed((factory, pooled, num_neg, mbs, seqlen)) + assert len(results) == 2 + for rank, r in results.items(): + assert r["error"] is None, f"rank {rank} failed:\n{r['error']}" + assert r["loss_diff"] < 1e-5, f"rank {rank} loss mismatch: {r['loss_diff']}" + assert r["grad_diff"] < 1e-4, f"rank {rank} grad mismatch: {r['grad_diff']}" + + +# CUDA-only: autocast (mixed-precision) parity. +@pytest.mark.slow +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16], ids=["bf16", "fp16"]) +def test_gradcache_autocast_parity_cuda(dtype): + """The backward-hook re-forward runs outside autocast; the wrapper replays the policy itself.""" + batch = 8 + device = "cuda" + torch.manual_seed(0) + model = ToyEncoder(pooled=False).to(device) + torch.manual_seed(1) + inputs = make_inputs(batch, num_neg=0, device=device) + loss_fn = _colbert() + + with torch.autocast(device_type="cuda", dtype=dtype): + model.zero_grad(set_to_none=True) + ref_loss, _, _ = reference_loss(model, loss_fn, inputs, num_neg=0) + ref_loss.backward() + ref_grads = grads_of(model) + + with torch.autocast(device_type="cuda", dtype=dtype): + model.zero_grad(set_to_none=True) + gc = WithGradCache(loss_fn, mini_batch_size=3) + gc.gather_across_processes = False + gc_loss = gc(model, inputs) + gc_loss.backward() + gc_grads = grads_of(model) + + # Mixed precision: compare loss within a loose tolerance and require finite, close grads. + assert torch.allclose(ref_loss, gc_loss, atol=1e-2, rtol=1e-2) + assert all(torch.isfinite(g).all() for g in gc_grads.values()) + assert max_rel_grad_diff(ref_grads, gc_grads) < 5e-2 From 757f8bd3491982e14359ab0ff479b34f2d9671f0 Mon Sep 17 00:00:00 2001 From: Yongbin Choi Date: Tue, 21 Jul 2026 03:06:20 +0900 Subject: [PATCH 2/4] docs: document GradCache training usage --- README.md | 62 +++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/README.md b/README.md index fd6a4e133..b91fc1193 100644 --- a/README.md +++ b/README.md @@ -337,6 +337,68 @@ sbatch --nodes=1 --time=5:00:00 -A cad15443 --gres=gpu:8 --constraint=MI250 -- +GradCache can reduce activation memory and enable larger in-batch negative pools by embedding each training batch in +smaller mini-batches. This trades additional computation time for lower peak memory. + +
+🔽 Example 3: Training with GradCache + +Wrap an existing contrastive loss with `WithGradCache` and pass it as the `loss_func` in +`ColModelTrainingConfig`: + +```python +from colpali_engine.loss import ColbertLoss, WithGradCache +from colpali_engine.trainer import ColModelTrainingConfig + +training_config = ColModelTrainingConfig( + model=model, + processor=processor, + train_dataset=train_dataset, + tr_args=training_args, + loss_func=WithGradCache( + loss=ColbertLoss(temperature=0.02), + mini_batch_size=8, + ), +) +``` + +To disable GradCache, use the original loss directly in the same training configuration: + +```python +training_config.loss_func = ColbertLoss(temperature=0.02) +``` + +`WithGradCache` delegates scoring and loss computation to the wrapped loss, so the wrapped loss's hyperparameters still +apply. + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `loss` | Required | A compatible bi-encoder or late-interaction contrastive loss to wrap. | +| `mini_batch_size` | `32` | Positive integer number of items embedded in each forward mini-batch. Smaller values generally use less memory but take longer. | +| `show_progress_bar` | `False` | Whether to show progress while embedding mini-batches. | + +`per_device_train_batch_size` still controls the contrastive batch and its in-batch negative pool; +`mini_batch_size` only controls how many items are embedded at once. `ContrastiveTrainer` detects the wrapper +automatically, while an unwrapped loss follows the standard training path. GradCache is not compatible with +`compute_symetric_loss=True`. + +In a ColQwen2 LoRA experiment on one NVIDIA RTX PRO 6000 (96 GiB), GradCache made a batch size of 192 fit where the +standard path ran out of memory. Both runs used one epoch on `vidore/colpali_train_set`, gradient checkpointing, the +same learning rate, seed, and a `ColbertLoss` temperature of `0.02`: + +| Mode | Batch size | Mini-batch size | Peak VRAM | Time per epoch | Average NDCG@5 | +|------|-----------:|----------------:|----------:|---------------:|---------------:| +| Standard | 128 | - | 88.7 GiB | 4h 17m | 75.2 | +| Standard | 192 | - | OOM | - | - | +| GradCache | 192 | 8 | 83.4 GiB | 6h 48m | 76.3 | + +Average NDCG@5 is calculated across the 14 ViDoRe v1 and v2 tasks. These measurements illustrate the memory-compute +trade-off for this setup and are not general performance guarantees. GradCache primarily reduces model activation +memory; full-batch embeddings and loss tensors still scale with the batch size, so larger batches can still run out of +memory. + +
+ ## Contributing We welcome contributions to ColPali! 🤗 From 1cc94076bfe67a30b08a4df043283fe86050d01b Mon Sep 17 00:00:00 2001 From: Yongbin Choi Date: Tue, 21 Jul 2026 22:44:37 +0900 Subject: [PATCH 3/4] fix: validate mini_batch_size is positive --- colpali_engine/loss/gradcache.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/colpali_engine/loss/gradcache.py b/colpali_engine/loss/gradcache.py index 1204db760..6043134e8 100644 --- a/colpali_engine/loss/gradcache.py +++ b/colpali_engine/loss/gradcache.py @@ -89,6 +89,8 @@ class WithGradCache(nn.Module): def __init__(self, loss: nn.Module, mini_batch_size: int = 32, show_progress_bar: bool = False): super().__init__() + if mini_batch_size <= 0: + raise ValueError("mini_batch_size must be a positive integer.") self.loss = loss self.mini_batch_size = mini_batch_size self.show_progress_bar = show_progress_bar From ac4fc927d7818b9c8faf6d8e4ac6fc68608ed4f3 Mon Sep 17 00:00:00 2001 From: Yongbin Choi Date: Wed, 22 Jul 2026 04:48:41 +0000 Subject: [PATCH 4/4] test: cover GradCache DDP gradient accumulation --- tests/loss/test_gradcache.py | 130 +++++++++++++++++++++++++++++++++++ 1 file changed, 130 insertions(+) diff --git a/tests/loss/test_gradcache.py b/tests/loss/test_gradcache.py index 38dbcfc11..07b574575 100644 --- a/tests/loss/test_gradcache.py +++ b/tests/loss/test_gradcache.py @@ -1,6 +1,9 @@ # ruff: noqa: N806, N812 import os +import shutil import socket +import subprocess +import sys import pytest import torch @@ -314,3 +317,130 @@ def test_gradcache_autocast_parity_cuda(dtype): assert torch.allclose(ref_loss, gc_loss, atol=1e-2, rtol=1e-2) assert all(torch.isfinite(g).all() for g in gc_grads.values()) assert max_rel_grad_diff(ref_grads, gc_grads) < 5e-2 + + +class _ToyPairDataset(torch.utils.data.Dataset): + """Deterministic pairs used by the launched Accelerate integration test.""" + + def __len__(self): + return 64 + + def __getitem__(self, index): + generator = torch.Generator().manual_seed(index) + query_ids = torch.randint(1, 64, (8,), generator=generator) + doc_ids = torch.randint(1, 64, (8,), generator=generator) + return { + "query_input_ids": query_ids, + "query_attention_mask": torch.ones_like(query_ids), + "doc_input_ids": doc_ids, + "doc_attention_mask": torch.ones_like(doc_ids), + } + + +class _ToyPairCollator: + query_prefix = "query_" + pos_doc_prefix = "doc_" + neg_doc_prefix = "neg_doc_" + + def __call__(self, features): + return {key: torch.stack([feature[key] for feature in features]) for key in features[0]} + + +class _RecordingGradCache(WithGradCache): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.gather_history = [] + + def forward(self, *args, **kwargs): + self.gather_history.append(self.gather_across_processes) + return super().forward(*args, **kwargs) + + +def _accelerate_ddp_gradacc_worker(output_dir): + """Run by ``accelerate launch``; assertions execute independently on both ranks.""" + from accelerate.utils import DistributedType + from transformers import TrainingArguments + + from colpali_engine.trainer import ContrastiveTrainer + + loss_fn = _RecordingGradCache(_bi_encoder(), mini_batch_size=2) + model = ToyEncoder(pooled=True, dropout=0.1) + initial_params = {name: param.detach().clone() for name, param in model.named_parameters()} + ddp_observations = [] + + class RecordingTrainer(ContrastiveTrainer): + def compute_loss(self, wrapped_model, inputs, *args, **kwargs): + ddp_observations.append(isinstance(wrapped_model, torch.nn.parallel.DistributedDataParallel)) + return super().compute_loss(wrapped_model, inputs, *args, **kwargs) + + args = TrainingArguments( + output_dir=output_dir, + per_device_train_batch_size=8, + gradient_accumulation_steps=2, + max_steps=2, + learning_rate=5e-2, + save_strategy="no", + logging_strategy="no", + report_to="none", + disable_tqdm=True, + dataloader_num_workers=0, + seed=17, + ) + trainer = RecordingTrainer( + model=model, + args=args, + train_dataset=_ToyPairDataset(), + eval_dataset=None, + data_collator=_ToyPairCollator(), + loss_func=loss_fn, + is_vision_model=False, + ) + result = trainer.train() + + assert trainer.accelerator.distributed_type == DistributedType.MULTI_GPU + assert dist.is_initialized() and dist.get_backend() == "nccl" + assert result.global_step == 2 + assert ddp_observations == [True, True, True, True] + assert loss_fn.gather_history == [False, True, False, True] + + params_changed = False + for name, param in model.named_parameters(): + assert torch.isfinite(param).all(), f"non-finite parameter on rank {dist.get_rank()}: {name}" + params_changed |= not torch.equal(param.detach().cpu(), initial_params[name]) + gathered = [torch.empty_like(param) for _ in range(dist.get_world_size())] + dist.all_gather(gathered, param.detach()) + for other_rank, other_param in enumerate(gathered): + assert torch.equal(param, other_param), f"parameter {name} differs on rank {other_rank}" + assert params_changed, f"optimizer did not update parameters on rank {dist.get_rank()}" + + +@pytest.mark.slow +@pytest.mark.skipif(not torch.cuda.is_available() or torch.cuda.device_count() < 2, reason="requires 2 CUDA GPUs") +def test_gradcache_accelerate_ddp_nccl_with_gradient_accumulation(tmp_path): + """End-to-end HF Trainer run through Accelerate, 2-process DDP/NCCL, and grad accumulation.""" + accelerate = shutil.which("accelerate") or os.path.join(os.path.dirname(sys.executable), "accelerate") + if not os.path.isfile(accelerate): + pytest.skip("requires the Accelerate CLI") + + command = [ + accelerate, + "launch", + "--multi_gpu", + "--num_processes=2", + f"--main_process_port={_find_free_port()}", + "--mixed_precision=no", + os.path.abspath(__file__), + "--accelerate-ddp-gradacc-worker", + str(tmp_path), + ] + completed = subprocess.run(command, capture_output=True, text=True, timeout=180, check=False) + assert completed.returncode == 0, ( + f"Accelerate DDP worker failed with exit code {completed.returncode}\n" + f"stdout:\n{completed.stdout}\n" + f"stderr:\n{completed.stderr}" + ) + + +if __name__ == "__main__" and "--accelerate-ddp-gradacc-worker" in sys.argv: + marker_index = sys.argv.index("--accelerate-ddp-gradacc-worker") + _accelerate_ddp_gradacc_worker(sys.argv[marker_index + 1])