Rebuild the e2e training benchmarks around per-op MaxSim VRAM (ColQwen2 + PyLate)

[tonywu71]

Context

The previous e2e training benches (bench_colpali_training/realdata.py, bench_pylate_training/realdata/lateon.py) ran hand-rolled step loops on synthetic data or small eval subsets and reported average step time plus whole-step peak memory. On encoder-dominated training steps that mostly measures parity (~1×), which is not where the kernel's value is.

colpali PR #412 built a better harness for its own validation: real trainer steps on real train data, the MaxSim op instrumented per call (forward transient, bytes held for backward), every recorded shape replayed on an isolated graph so the op's backward is bracketed exactly, and OOM treated as a recorded sweep outcome instead of a crash. What that surfaces (the op's B²-growing memory term and the batch-size ceiling it sets) survives encoder dilution where the step-time speedup does not.

This PR replaces the five old benches with two ports of that harness, kept as symmetrical as possible:

• bench_colpali_e2e.py — ColQwen2 + LoRA through colpali-engine's own ColModelTraining recipe, on vidore/colpali_train_set shards (released colpali-engine==0.3.16, toggled with patch_colpali_engine()).
• bench_pylate_e2e.py — GTE-ModernColBERT-v1 + Contrastive through SentenceTransformerTrainer, on MS MARCO triplets (pylate==1.5.0, toggled with patch_pylate()), with a --grad-checkpoint regime and --score-mini-batch-size cells covering PyLate's own chunking mitigation.

Benchmarking current PyLate required a product fix that rides along: PyLate 1.5 renamed the scoring module and rerouted its losses through ColBERTScores, which broke patch_pylate(). The old benches pinned 1.3.3 because of this; the patch now detects the layout, so the >=1.3.3,<2 range in the extra is accurate again.

Results (1× H100 80 GB)

Full tables and readings live in docs/benchmarks.md; per-cell JSONs are reproducible via the SkyPilot jobs below.

ColQwen2 reproduces colpali PR #412's published numbers to the MiB. The MaxSim op costs vanilla 7.81 GiB of VRAM at B=128 vs 61 MiB with LIK (~130×, ×4 vs ×2 per batch doubling); step time is at parity since the 2B encoder hides the op. Vanilla OOMs at B=128 (a 1.81 GiB request failing while 25 GiB sit reserved-but-unallocated, fragmentation caused by the multi-GiB score grid) where LIK trains it at 7.98 samples/s:

batch size	vanilla	LIK
64	29.54 GiB	29.54 GiB
128	OOM	54.38 GiB

PyLate has two regimes and two different stories. On the canonical recipe (no grad-checkpointing) the encoder activations fill the card at B=256 before the score grid matters: vanilla, LIK, and PyLate's score_mini_batch_size chunking all OOM identically, so LIK only trims the op transient (1.8 GiB → 56 MiB at B=128) without moving the ceiling. With --grad-checkpoint, the regime matching the ColQwen2 bench, the B² transient binds and the 149 M encoder is too small to hide the op:

batch size	vanilla	vanilla + score_mini_batch_size=64	LIK
512	54.10 GiB · 2.69 s/step	30.60 GiB · 2.70 s/step	29.72 GiB · 2.40 s/step
1024	OOM (one 56.25 GiB grid alloc)	62.18 GiB · 6.02 s/step	57.67 GiB · 4.81 s/step

The harness pins vanilla's B=1024 death inside the score call: a single contiguous allocation of exactly the [1024, 1024, 48, 300] fp32 per-slot grid. LIK is also 1.07–1.12× faster per step than vanilla here, and 1.25× faster than PyLate's own chunking at matched B=1024. Chunking buys the ceiling with serialized recompute; the fused kernel buys it while speeding the step up.

Quick start

sky launch scripts/sky_colpali_e2e.yaml -c lik-colpali-e2e -y
sky launch scripts/sky_pylate_e2e.yaml  -c lik-pylate-e2e  -y

rsync -avP 'lik-colpali-e2e:sky_workdir/benchmarks/results/' ./benchmarks/results/
python benchmarks/summarize_colpali_e2e.py --results-dir benchmarks/results --plot benchmarks/results/colpali_e2e_vram.png
python benchmarks/summarize_pylate_e2e.py  --results-dir benchmarks/results --regime ckpt

Note

• bench_colpali_e2e.py ships a small ContrastiveTrainer shim for two colpali-engine 0.3.16 bugs under transformers 5.x (_get_train_sampler signature; single-dataset loss prefixes). Both are fixed upstream in colpali PR #412 but unreleased; the shim goes away once colpali-engine > 0.3.16 ships.
• bench_colpali_loss.py (loss-head isolation, no encoder) is kept: it backs the 2.5–4.3× explicit-negative numbers already in the docs.
• Historical result tables in docs/benchmarks.md stay; only references to removed scripts were updated.

Changes

Click here to expand

• benchmarks/bench_colpali_e2e.py — ColQwen2 harness: per-MaxSim-call VRAM recording at the loss head, isolated-op replay (vanilla replays the release's exact einsum → amax → sum, lik replays lik.autograd.maxsim), step timer, OOM-as-outcome
• benchmarks/bench_pylate_e2e.py — PyLate sibling: wraps the loss's score_metric, replays PyLate 1.5's exact per-slot ColBERTScores math vs the fused per-slot path; --grad-checkpoint and --score-mini-batch-size dimensions
• summarize_colpali_e2e.py / summarize_pylate_e2e.py — op-VRAM and batch-ceiling markdown tables plus log-log plots from the per-cell JSONs (--regime plain|ckpt for pylate)
• scripts/sky_colpali_e2e.yaml / sky_pylate_e2e.yaml — fresh process per (batch, variant) cell so OOMs are isolated; crash markers keep the sweep going; the colpali job is also the new home for bench_colpali_loss.py --sweep
• patch_pylate() 1.5 support — resolve the defining module per layout (scores.scores for <=1.4, scores.colbert for >=1.5), rewrite only Distillation's import-time capture on 1.5; layout-aware CUDA compat tests
• Removed — bench_colpali_training/realdata.py, bench_pylate_training/realdata/lateon.py, _bench_common.py, sky_colpali_benchmark.yaml, sky_pylate_benchmark.yaml; pylate blocks dropped from sky_run_all_benchmarks.yaml
• Docs — new ColQwen2 and PyLate e2e sections in docs/benchmarks.md with the measured tables, README benchmark/memory updates, CHANGELOG entries, benchmarks/README operator guide

Next steps

• Once colpali PR #412 merges and a colpali-engine release ships it, drop the trainer shim here and migrate the colpali job to COLPALI_SCORES_BACKEND (per that PR's plan, the LIK-side patches eventually retire in favor of the native integrations).
• Follow-up PR: reorganize benchmarks/ per stack — kernels/ (synthetic micro benches), plaid/, colpali/, pylate/ — with each summarize_* next to its bench, and rename bench_lateon.py → kernels/bench_longdoc.py (it sweeps long-document shapes, and its e2e namesake is gone). Kept out of this PR so the benchmark work isn't buried in rename noise.

[h-aurelien-lac]

Went through this carefully and cross-checked the harnesses against colpali-engine 0.3.16 and pylate 1.5.0 source — really clean, the hook points and the fp32-grid replay all line up with what the real trainers do. Just a few small things, nothing blocking:

• In bench_pylate_e2e._replay_op_isolated, the replay masks are built with the embeddings dtype, but the real skiplist masks are bool — so the "allocation-identical to the real skiplist masks" comment is slightly off. It's sub-MiB so it doesn't move any number, but maybe cast to bool or soften the comment.
• The slope annotations in summarize_colpali_e2e._write_plot ("×4 per doubling", "×2 per doubling") are hardcoded regardless of the data — fine today, but if a future run drifts the figure would assert a scaling the points don't show. Could derive them from the actual ratios.
• Tiny one: we record step_peak_reserved_mib but neither summarizer surfaces it, even though the fragmentation story ("25 GiB reserved-but-unallocated") is really a reserved-vs-alloc thing. Either show it or drop the field.

The ~150 lines shared between the two benches (recorder/replay/timer) I assume is intentional given the benchmarks/ reorg in your next-steps — just confirming that's the plan and not drift.

[tonywu71]

Went through this carefully and cross-checked the harnesses against colpali-engine 0.3.16 and pylate 1.5.0 source — really clean, the hook points and the fp32-grid replay all line up with what the real trainers do. Just a few small things, nothing blocking:

• In bench_pylate_e2e._replay_op_isolated, the replay masks are built with the embeddings dtype, but the real skiplist masks are bool — so the "allocation-identical to the real skiplist masks" comment is slightly off. It's sub-MiB so it doesn't move any number, but maybe cast to bool or soften the comment.
• The slope annotations in summarize_colpali_e2e._write_plot ("×4 per doubling", "×2 per doubling") are hardcoded regardless of the data — fine today, but if a future run drifts the figure would assert a scaling the points don't show. Could derive them from the actual ratios.
• Tiny one: we record step_peak_reserved_mib but neither summarizer surfaces it, even though the fragmentation story ("25 GiB reserved-but-unallocated") is really a reserved-vs-alloc thing. Either show it or drop the field.

The ~150 lines shared between the two benches (recorder/replay/timer) I assume is intentional given the benchmarks/ reorg in your next-steps — just confirming that's the plan and not drift.

Thanks for the careful pass — all four are right. Addressed in 841335f:

1. Replay masks: confirmed extract_skiplist_mask returns bool (torch.logical_and), so the replay now builds bool masks instead of softening the comment. This also fixed a second subtlety: the lik replay's != 0 conversion was allocating a bool copy that the real patched path never makes (_mask_as_bool is a pass-through on bool). New _as_bool helper mirrors the patch exactly.
2. Slope annotations: now derived from the last two points; a comment keeps the expected ×4/×2 as documentation of intent. On the current results the figure reads ×4.0 / ×2.1, so the claim holds with measured numbers.
3. step_peak_reserved_mib: went with "show it" — both ceiling tables now print alloc / reserved per cell (colpali B=128: vanilla OOM (53.95 / 77.73 GiB pre-OOM) vs LIK 54.38 / 78.37 GiB, which makes the fragmentation mechanism visible in the table itself).
4. Duplication: intentional — the harnesses are deliberately single-file so each reads standalone, same as the PR 412 original. Extracting the shared recorder/replay/timer into a common module is queued for the per-stack benchmarks/ reorg listed in Next steps.

Caveat: 1–2 touch GPU-only paths and are verified by lint + regenerated plots from the existing JSONs; the bool-mask replay gets its first real exercise on the next sweep. Sub-MiB delta, so I didn't relaunch an H100 for it. Will do in follow-up PR.

[h-aurelien-lac]

LGTM !!