feat: --trim-backbone-rows — run TTT steps >= 1 only on rows that can still emit loss

[julyanghar]

Motivation

Follow-up to #705 (--trim-loss-positions, merged) and the RFC in #706. On prompt-heavy training data (agent/RAG traces, long-context SFT — often ~90% loss-masked), #705 removed the vocab-sized teacher/logits/loss work at masked positions, but the draft backbone still runs full-length attention + MLP for every position at every TTT step, even though masked positions' outputs never reach a live gradient.

One correction to the RFC first: #706 proposed forwarding "only the supervised rows" at steps 2..k, and quoted an equivalence check and a memory table for a prototype built that way. That row selection is wrong — loss rows shift left once per TTT step (the same shifted-mask semantics that #705's review uncovered), and the prototype's self-check shared the same wrong assumption, so it validated itself. Both the old equivalence claim and the old numbers should be considered retracted; this PR is a from-scratch implementation of the corrected math, with tests designed so that the historical wrong reading is explicitly detected (a negative control fails on it).

The corrected row set: at step i, forward R_i = ∪_{j≥i} { s − j : s ∈ sup, 0 ≤ s − j < chunk } — every row that can still emit loss at this or a later step (hidden states chain between steps, so a row needed later must be forwarded now). The sets are nested (R_i ⊆ R_{i−1}), which keeps all row bookkeeping to searchsorted lookups. Step 0 always stays full-length: it produces the K/V context every later row attends to. On contiguous supervised blocks (real SFT data) the union barely inflates: 8% supervised ⇒ steps ≥ 1 forward ~8% of rows.

Modifications

• training.trim_backbone_rows (default off; requires trim_loss_positions; validated against flex_attention, compact_teacher, mrope drafts, non-contiguous position_ids), wired through capability/planning/strategy exactly like trim_loss_positions.
• _build_trim_pack additionally emits per-step row sets and nested-set index maps (dead steps padded so kernels and collectives stay rank-aligned).
• One shared trimmed-attention path for sdpa/fa/usp in llama3_eagle.py: grouped-query layout, explicit causal mask from absolute positions (non-contiguous row sets are exactly what flash kernels cannot express), diagonal-private K/V from later steps row-aligned via the index maps, activation checkpointing per step/row-chunk.
• Long step-0 K/V (≥ 8192) use a single fused memory-efficient SDPA call per step: query groups fold into the row axis (fused kernels require equal head counts), diagonal K/V join as extra columns behind a boolean mask, so no [rows × L] tensor ever reaches HBM.
• Under USP, trimmed steps bypass ring/Ulysses entirely (per-rank row counts differ; all-to-all needs equal lengths): step-0 K/V are all-gathered once via the autograd-aware all_gather (backward = reduce-scatter sum — correct when each rank owns a different loss), and enabling is agreed rank-uniformly (all_reduce MIN) so mixed trim/full ranks cannot deadlock.

Related Issues

RFC #706 (this implements it, with the row-set math corrected as described above). Builds on #705.

Accuracy Test

tests/test_runtime/test_trim_backbone_rows.py:

• CPU golden tables pin the R_i union math (overlap-tail bound, dead-step padding, unreachable-supervision fallback, non-USP back-compat) — these run in CI even without GPUs.
• Per-step loss equivalence, trim ON vs OFF (both with trim_loss_positions on): all-supervised mask (trivial-equality boundary — any row/position error is exposed without tolerance cover), prompt-heavy / block-boundary / isolated-point masks, plus a total-grad-norm check; on sdpa, fa, and 4-rank ring4 USP. The fused SDPA core is additionally forced through every case via a patched threshold.
• A negative control (kept as an experiment script, not in CI) re-creates the RFC's original fixed-sup row selection and confirms the equivalence tests detect it (rel. err 1.7e-2 at step 1, step 0 exact — the fingerprint of the shifted-mask bug).

Benchmark & Profiling

Llama3-8B dims (hidden 4096, target vocab 128256, draft vocab 32000), TTT 7, 16-block 8% supervised mask, RTX 6000 Ada 48 GB, batch 1. Baseline = trim_loss_positions only; trimmed = + trim_backbone_rows. Peak = torch.cuda.max_memory_allocated per rank, mean of 3 steps after warmup:

config	baseline	+ trim_backbone_rows
8k, single GPU, sdpa	OOM	38.2 GB, 1.16 s/step (unlocked)
16k, single GPU, sdpa	OOM	OOM (step-0's materialized sdpa attention is the binding constraint; fa's step 0 does not have it)
32k, ring4	17.9 GB, 1.68 s/step	8.6 GB, 0.62 s/step
64k, ring4	32.9 GB, 4.60 s/step	14.9 GB, 1.81 s/step

The savings compose with #705's (both are on in the baseline column already).

Applicability boundary and the cost guard

Row trimming is not free: a trimmed step cannot use flash attention (the kept rows are not
contiguous), so it runs the explicit masked-matmul kernel, which costs more per row. Fewer rows
dropped and a higher per-row cost stack up, so the win shrinks as supervision gets denser and
eventually inverts.

Measured on an 8192-token sequence, TTT=7, single RTX 6000 Ada, fa backend, both arms with
--trim-loss-positions:

supervision density	rows kept over steps 1..6	step time vs baseline
2%	1,164	0.10x
8%	4,044	0.30x
25%	12,492	0.64x
40%	19,788	~1.00x (break-even)
100%	49,125	2.45x

Supervision density is a property of the data, not a tunable: it is the fraction of tokens written
by the assistant (_apply_loss_mask_from_chat_template marks assistant spans only). Over the
120,675 ShareGPT samples shipped with this repo, samples under 4096 tokens (99.7% of them) have a
median density of 0.92, while samples at or above 4096 tokens have a median of 0.056. Short
chats are dense; long samples are long because of pasted material, not longer answers.

training.trim_backbone_rows_max_density (default 0.35) therefore gates the feature per sample:
if sum_i |R_i| > tau * (k-1) * chunk, the sample runs the full-length path and an INFO log
explains why. Under USP the decision is agreed across ranks with all_reduce(MIN), since ranks can
otherwise land on opposite sides of the budget and mismatch collectives (verified: removing the
agreement hangs a 4-rank run past a 180 s timeout).

With the guard, low-density data keeps the full benefit (9.03 GB / 0.325 s at 8%, versus 9.03 GB /
0.320 s without the guard) and dense data degrades to the baseline instead of running 2.45x slower
(32.68 GB / 1.586 s versus the baseline's 32.68 GB / 1.548 s).

Practical consequence worth stating plainly: on the default ShareGPT recipe the guard will keep
--trim-backbone-rows switched off. This feature targets long-context and agent-trajectory data,
where supervision is sparse — not the short-chat default.

The decision is per sample rather than per step: |R_i| barely varies across steps for contiguous
supervision, and a mixed trimmed/full unroll would need per-backend K/V layout conversion plus an
extra head-dimension gather under USP — new collectives for a safety guard is a poor trade.

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in Running Unit Tests.
• Update documentation as needed (config README row + schema docstrings).
• Benchmark results above.

[julyanghar]

/gemini review

[julyanghar]

I just pushed two more commits to this branch, and I owe you an explanation for why they landed
here rather than in a follow-up PR.

After the benchmarks in the description, I went looking for the point where row trimming stops
paying, and found one that matters for the default recipe. Merging the previous head alone would
have shipped a flag that makes the default dataset slower without telling anyone why, so I
would rather you review the feature together with its applicability boundary.

What I found. A trimmed step cannot use flash attention: the kept rows are not contiguous in
the sequence, so it falls back to the explicit masked-matmul kernel, which costs more per row.
Two effects then stack — fewer rows can be dropped and each remaining row is more expensive:

supervision density (8192-token sequence)	step time vs fa + --trim-loss-positions
8%	0.30x (the number in the PR description)
40%	~1.00x — break-even
100%	2.45x slower

Why it matters for the default recipe. Supervision density is not a knob, it is the fraction
of tokens written by the assistant. Measured over all 120,675 ShareGPT samples shipped with
SpecForge (Qwen3.6 tokenizer):

sequence length	samples	median supervision density
< 4096	120,303 (99.7%)	0.92
>= 4096	372 (0.31%)	0.056

Short chats are dense because the user asks one line and the model answers at length; long samples
are sparse because they are long on pasted material. And 63 of the 66 example configs use
max_length <= 4096. So the default configuration sits squarely on the losing side.

The fix. A per-sample cost guard, training.trim_backbone_rows_max_density (default 0.35):
compare the trimmed work sum_i |R_i| against tau * (k-1) * chunk; if it exceeds the budget the
whole sample runs the full-length path. Effect:

	baseline fa + A	before the guard	after the guard
8% supervision	18.39 GB / 1.054 s	9.03 GB / 0.320 s	9.03 GB / 0.325 s
100% supervision	32.68 GB / 1.548 s	35.10 GB / 3.795 s	32.68 GB / 1.586 s

Savings at low density are unchanged (the 1.6% delta is the guard's own arithmetic); the loss at
high density is gone.

The two commits are +218 lines over 8 files (56 lines of guard in model.py, 152 lines of tests,
the rest config plumbing) plus a one-line docstring fix where a stale comment still described an
earlier per-step design. Full suite and pre-commit are green locally. Happy to split the guard back
out into its own PR if you would rather review the feature and the guard separately.