feat(loss): add WithGradCache wrapper for contrastive losses

[whybe-choi]

Summary

Adds a single WithGradCache wrapper that enables large effective batch sizes under a fixed memory budget for both bi-encoder and late-interaction contrastive losses (GradCache, 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.

Context

This follows up on #399 (where I offered to continue the work) and the original prototype in #198. In #399, you laid out what it would take to merge:

I'll probably only merge in main if: either it gets perf improvements, or the code is very self contained and doesn't add too much bloat. It might be worth it not to just add a "gradcache flag" but really construct a class (WithGradCache) to really control the induced bloating.

Since you also noted that GradCache worked but didn't beat offline hard-negative mining, this PR targets the second criterion — self-contained, minimal bloat — rather than making a performance claim. Concretely:

• A WithGradCache class, not a flag. All GradCache logic lives in one new file, colpali_engine/loss/gradcache.py, and is opt-in by wrapping a loss. No global "gradcache flag" is threaded through the training code.
• No bloat in existing code. The contrastive losses are untouched: the wrapper delegates all scoring/loss math to the wrapped loss and adds no loss-specific logic, so it composes with every current and future bi-encoder / late-interaction loss.
• Standard path left intact. ContrastiveTrainer routes to the gradcache path only when the loss exposes gradcache_enabled (a ~16-line addition); the existing non-cached path is unchanged.
• Backed by parity tests, not just plumbing. The wrapper is proven to reproduce the non-cached path's loss and gradients exactly (see Tests), so "self-contained and equivalent" is verified rather than asserted.

I reimplemented from scratch rather than rebasing #198 to keep the footprint minimal and self-contained, but I'm happy to align with that branch if you'd prefer.

How it works

1. Forward (no grad): each branch (query / positive / negatives) is embedded in mini-batches under torch.no_grad(); embeddings are detached and marked requires_grad. An RNG snapshot (RandContext) is captured per mini-batch.
2. Cache: the wrapped loss runs on the concatenated embeddings and is back-propagated to fill the per-embedding gradient cache. This graph never touches model parameters, so no activation memory for the full batch is held.
3. Replay: a backward hook recomputes each mini-batch with gradients (restoring the captured RNG so dropout matches) and applies the cached gradient via a surrogate dot-product, yielding exact parameter gradients one mini-batch of activations at a time.

All scoring/loss math is delegated to the wrapped loss, so the wrapper is 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. The ambient autocast policy is replayed in the backward hook, and the distributed positive-gather (with per-rank offset) is preserved.

Tests

tests/loss/test_gradcache.py (toy encoder, runs on CPU; one CUDA-gated case):

• Equivalence — gradcache vs. non-cached path match on loss and parameter gradients for BiEncoderLoss, BiNegativeCELoss, ColbertLoss, ColbertNegativeCELoss, across mini_batch_size smaller / equal / larger than the batch.
• Eval / no_grad path — returns the loss with no autograd hook and leaves parameter grads untouched.
• RandContext — dropout is reproduced across the two forward passes; gradcache with dropout is deterministic and finite.
• Distributed — 2-process gloo (no GPU needed) parity over the all-gather + offset path, including the late-interaction front-padding branch for unequal doc lengths.
• Autocast (CUDA, slow) — bf16 / fp16 mixed-precision parity.

pytest:

• pytest tests/loss/test_gradcache.py — CPU + distributed (gloo)
• pytest tests/loss/test_gradcache.py -m slow — include the CUDA autocast parity case

Note

Claude Code helped with parts of the implementation and drafted the test cases; everything's been manually reviewed and verified end-to-end. But there may be something I could have missed, so feel free to add/delete/update features.

Thank you!

[whybe-choi]

Hello, @QuentinJGMace and @ManuelFay !

Whenever you have some time, would you mind taking a look at this PR?
I’d really appreciate any feedback or review when it’s convenient for you.

Thank you!

[whybe-choi]

Experiment: memory usage and retrieval quality with GradCache

I ran a small experiment on ColQwen2 to check GradCache's practical payoff on a fixed GPU: (1) whether it actually lowers memory and lets me fit a batch that plain training OOMs on, and (2) whether the larger in-batch-negative pool improves retrieval.

Setup

• Model: ColQwen2 (LoRA, Qwen/Qwen2-VL-2B-Instruct base)
• Training data: trained on vidore/colpali_train_set (in-batch negatives only)
• Loss: ColbertLoss (temperature = 0.02)
• Hardware: 1× NVIDIA RTX PRO 6000 (96 GiB)
• Held constant across runs: epochs (1), learning rate, seed, temperature, gradient_checkpointing=True. Only the batch size (and the GradCache wrapper) changes.

The memory wall

On this GPU, plain training (no GradCache) OOMs at batch size 192; the largest batch that fits is 128. With GradCache wrapping the same loss, batch size 192 trains fine by embedding mini_batch_size = 8 items at a time — and, as the numbers below show, at a lower peak VRAM than the 128-batch baseline.

Run	GradCache	per_device_train_batch_size	mini_batch_size	Peak VRAM	Mean GPU util	Fits?
baseline	✗	128	–	88.7 GiB	98%	✅
baseline	✗	192	–	–	–	❌ OOM
gradcache	✓	192	8	83.4 GiB	88%	✅

GradCache fits the larger 192-batch at lower peak VRAM than the baseline's 128-batch (83.4 vs 88.7 GiB), while keeping the GPU near-saturated (median util ≥ 96%). The ~10-point lower mean utilization reflects the mini-batch embed loop / double-forward — i.e. the extra wall-clock is real compute, not idle stalls.

Training cost & quality

GradCache trades compute for memory (each branch is forwarded twice: once under no_grad to build the cache, once in the backward hook), so I'm also reporting wall-clock to keep the comparison honest.

Run	per_device_train_batch_size	mini_batch_size	Wall-clock / epoch
baseline	128	–	4h 17m
gradcache	192	8	6h 48m

Retrieval quality (ViDoRe v1 / v2, NDCG@5; bold = better of the two runs). ViDoRe V2 task scores are averaged over their language subsets (fr / es / en / de; ESGHL is English-only):

Run	ViDoRe V1										ViDoRe V2				Average
	ArXiv	DocVQA	InfoVQA	Shift	AI	Ene	Gov	Hlt	TabF	TatD	BMed	ESG	ESGHL	Econ	V1	V2	Avg
baseline (bs=128)	85.8	59.5	89.6	82.7	98.0	95.0	92.3	97.9	83.5	76.9	54.4	47.6	51.2	39.1	86.1	48.1	75.2
gradcache (bs=192, mbz=8)	86.3	60.9	91.4	86.1	98.7	95.2	95.0	97.8	87.4	77.2	55.9	45.1	46.0	45.9	87.6	48.2	76.3

Takeaway

GradCache trains the 192-batch that plain training OOMs on, and does so at lower peak VRAM than the largest batch (128) the baseline could fit. The larger in-batch-negative pool clearly helps in-domain ViDoRe V1 (86.1 → 87.6 NDCG@5, +1.5), while out-of-domain V2 is essentially flat (48.1 → 48.2) — overall 75.2 → 76.3 across the 14 tasks, at ~1.6× wall-clock (4h 17m → 6h 48m). So on a fixed GPU, GradCache turns a hard memory ceiling into a modest, paid-for quality gain, concentrated on V1.

Evaluation is on the ViDoRe v1 and v2 benchmarks (NDCG@5) using mteb; training itself runs with run_eval=False.

[whybe-choi]

Hi, @ManuelFay and @QuentinJGMace !
just a gentle ping on this PR. I’d really appreciate a review whenever you have time. Thanks!