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!

[ManuelFay]

Hello ! PR looks clean, quite similar to the one I had opened at the time so looks familiar ! Thanks !
Sorry for not responding before, I am away from work for a few months so just looking at this now.

It looks to be well tested, the only thing that I would like is if it's a bit more documented:

• how should one activate / deactivate gradcache from the training loop ? basicqlly if there is a little part in the readme or a dedicated md file that explains gradcahce: a little code snippet, the rationale/the stats you give in this PR, the default behaviors, hparams, it would help usage a lot !

The other thing is that SentenceTransformers is bound to include these visual retrivers at terms and supports gradcache out of the box, so I was hesistating to merge this now (but probably worth it to do so regardless).

Thanks again for everything and sorry for the delays

[whybe-choi]

Thanks for the review!

I added a collapsible “Training with GradCache” example to the README. It now covers how to enable and disable WithGradCache, its default parameters and behavior, the memory-compute trade-off, and the experiment results shared in this PR.I kept the high-level explanation visible and placed the detailed usage and statistics under a toggle to keep the README concise.

I’d appreciate another look when you have time. Thanks!

[ManuelFay]

Thanks, great !

Two minor technical follow-ups, but not preventing the merge:

• validate that mini_batch_size > 0;

• is this tested with an end-to-end Accelerate/DDP run with NCCL and gradient accumulation whcih would be the breaking situation I could envision ? worst case, we can follow this up later, but i know gradacc is non trivial with HF trainers

[whybe-choi]

Thanks for the follow-up!

I’ve added validation to ensure that mini_batch_size is greater than 0. For the end-to-end Accelerate/DDP run with NCCL and gradient accumulation, I’ll run additional tests and share the results here in a follow-up.

[whybe-choi]

Following up on my previous comment, I ran the additional end-to-end Accelerate/DDP test with NCCL and gradient accumulation and added it to the test suite.

I added an end-to-end GPU integration test using the actual ContrastiveTrainer.train() path.

Test configuration

Setting	Value
Launcher	accelerate launch --multi_gpu
Distributed backend	NCCL
DDP processes / GPUs	2
GradCache mini-batch size	2
Per-device training batch size	8
Global batch size per microstep	16
Gradient accumulation steps	2
Effective global batch size per optimizer step	32
Optimizer steps	2
Training microsteps per rank	4
Total training samples	64

Each per-device batch of 8 samples is processed as four GradCache chunks of 2 samples. The test covers all 64 samples without repeating the dataset.

Assertions

The test verifies that:

• the actual Hugging Face Trainer.train() and Accelerate paths are used;
• the distributed backend is NCCL;
• the model passed to GradCache is DDP-wrapped;
• cross-process document gathering follows the expected accumulation pattern:
  [False, True, False, True];
• two optimizer updates are performed;
• the optimizer actually changes the model parameters;
• all final parameters remain finite; and
• the final parameters are bitwise identical across both DDP ranks.

The run completed without deadlocks, NCCL collective errors, nested-backward errors, or rank divergence.

Environment

• 2× NVIDIA RTX PRO 6000
• PyTorch 2.11.0+cu130
• Accelerate 1.14.0
• Transformers 5.14.1

Results

The new end-to-end test passed:

1 passed in 12.25s

The complete GradCache test suite also passed:

23 passed in 33.28s

This full suite includes:

• single-process loss and gradient parity;
• 2-process distributed parity with gloo;
• CUDA fp16/bf16 autocast parity; and
• the new Accelerate/DDP/NCCL gradient-accumulation run.

One implementation detail worth noting is that cross-rank document gathering occurs only on synchronization microsteps, matching the existing non-GradCache trainer path. Non-synchronization microsteps use the local per-rank batch. Therefore, the effective training batch size per optimizer step is 32 samples, but it should not be interpreted as a single contrastive negative pool of size 32.