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! 🤗
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..6043134e8
--- /dev/null
+++ b/colpali_engine/loss/gradcache.py
@@ -0,0 +1,214 @@
+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__()
+ 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
+ # 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..07b574575
--- /dev/null
+++ b/tests/loss/test_gradcache.py
@@ -0,0 +1,446 @@
+# ruff: noqa: N806, N812
+import os
+import shutil
+import socket
+import subprocess
+import sys
+
+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
+
+
+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])