feat: --trim-loss-positions (compute teacher/logits/loss only at supe…

[julyanghar]

Motivation

On prompt-heavy training data (agent traces, RAG, long-context SFT), most positions are masked out of the loss — only the assistant/supervised positions contribute. Yet online EAGLE3 training still materializes the vocab-wide teacher target_p, the draft logits, and the per-position loss over the full sequence at every TTT step, so the majority of that compute and memory is spent on positions that are multiplied by zero.

--trim-loss-positions computes the teacher target_p, the draft logits, and the loss only at supervised (loss-masked) positions. It is mathematically equivalent to the full-length path — the mean loss denominator is rescaled from n_sup back to the full length — and it saves memory and compute proportional to the masked fraction.

This is a port of the loss-end position-trimming idea (already present in TorchSpec) to SpecForge's TTT loop + precomputed teacher target_p pipeline. The genuinely SpecForge-specific parts are the sliding-window teacher-table compaction (_build_trim_pack) and the mean-denominator rescale. It only touches the loss end; the draft backbone still runs full-length here (backbone-internal trimming is future work, best discussed as an RFC first).

Modifications

• specforge/core/eagle3.py:
  • OnlineEagle3Model gains a trim_loss_positions flag. When active (and eligible), it builds a compact teacher table over supervised rows + their k-step sliding window (_build_trim_pack), runs the backbone full-length, then gathers only supervised rows into compute_logits and _acc_and_loss.
  • _acc_and_loss gains loss_scale / full_positions so the trimmed mean is rescaled to the full-length semantics (loss_scale = n_sup / full_len).
  • New helpers _build_trim_pack, _compute_target_p_eager, _compute_position_mask_at.
  • Eligibility gate: online path with batch==1, lk_loss_type is None, at least one supervised position, and non-VLM; otherwise it falls back to the existing full-length path (byte-identical when off).
• scripts/train_eagle3.py: --trim-loss-positions CLI flag, wired into the OnlineEagle3Model construction.
• tests/test_runtime/test_equiv_trim_loss_positions.py: GPU equivalence test — the same online forward with trimming off vs. on must produce matching per-step losses (bf16 tolerance). GPU-only, matching the other online EAGLE3 equivalence tests in that directory.

No change to llama3_eagle.py / flex_attention.py (the backbone runs full-length under trimming).

Related Issues

None known.

Accuracy Test

This is an equivalence-preserving change, so it does not affect trained-model accuracy. Correctness is verified by test_equiv_trim_loss_positions: with a prompt-heavy loss mask, the per-step training losses with --trim-loss-positions on vs. off match within bf16 tolerance (relative error ≤ 5e-3). The off-path is byte-identical to main (the flag defaults to off).

Benchmark & Profiling

The teacher target_p and draft logits/loss tensors are the largest per-step activations, all ∝ sequence_length × vocab. Trimming makes them ∝ n_sup × vocab, so the saving scales with the unsupervised fraction.

Measured on Qwen3-32B (TP4, max_length=8192, flex_attention, TTT=3) on prompt-heavy data where ~92% of positions are masked out of the loss (supervised fraction ≈ 3–18% per sample):

metric (torch max_memory_allocated, per card)	off	--trim-loss-positions	saved
peak training activation	39.1–41.2 GiB	36.5–38.0 GiB	≈ 2.3–3.4 GiB/card

The saving is larger on more prompt-heavy data and at longer max_length (it is ∝ masked_positions), and smaller on evenly-supervised data. Because it is a strict activation/compute reduction with equivalent output, it does not change inference latency/throughput. (This is a specific-config data point; the effect scales with the masked fraction, so a standard ShareGPT-shaped workload would show a proportionally smaller saving.)

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[julyanghar]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces the --trim-loss-positions feature to optimize memory and compute by calculating the teacher target_p, draft logits, and loss only at supervised positions for batch size 1. It also includes an equivalence test to verify that this optimization does not alter the training loss. The review feedback suggests several performance and idiomatic improvements in PyTorch, such as using F.softmax instead of instantiating nn.Softmax in a loop, vectorizing tensor operations via broadcasting and unbind(1), utilizing logits.new_tensor for cleaner device allocation, and ensuring temporary test directories are properly cleaned up.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Instantiating nn.Softmax inside a loop is inefficient and non-idiomatic in PyTorch. Use F.softmax instead.

Suggested change

	tps.append(nn.Softmax(dim=2)(dth).detach())
	tps.append(F.softmax(dth, dim=2).detach())

[gemini-code-assist]

This list comprehension and concatenation can be fully vectorized using broadcasting, which is more efficient and avoids creating multiple temporary tensors on the GPU.

        shifted = (sup.unsqueeze(1) + torch.arange(length + 1, device=sup.device)).view(-1)

[gemini-code-assist]

Instead of performing length separate indexing operations in a loop (which launches multiple GPU kernels), you can perform a single 2D advanced indexing operation and then split the result using unbind(1). This is significantly more efficient.

Suggested change

	full_map = torch.full(
	(L + length + 1,), n_real, dtype=torch.long, device=sup.device
	)
	full_map[real] = torch.arange(n_real, device=sup.device)
	idx_steps = [full_map[sup + j] for j in range(length)]
	full_map = torch.full(
	(L + length + 1,), n_real, dtype=torch.long, device=sup.device
	)
	full_map[real] = torch.arange(n_real, device=sup.device)
	idx_steps = list(full_map[sup.unsqueeze(1) + torch.arange(length, device=sup.device)].unbind(1))

[gemini-code-assist]

Using logits.new_tensor is more idiomatic and efficient than torch.tensor(..., device=logits.device) as it automatically inherits the device and other properties of the existing tensor.

        loss_denom = logits.new_tensor(
            logits.shape[0]
            * (full_positions if full_positions is not None else logits.shape[1]),
            dtype=torch.float32,
        )

[gemini-code-assist]

The temporary directory created by tempfile.mkdtemp is never cleaned up, which can leak disk space. Use self.addCleanup with shutil.rmtree to ensure it is automatically deleted after the test runs.

Suggested change

	workdir = tempfile.mkdtemp(prefix="equiv_trim_")
	import shutil
	workdir = tempfile.mkdtemp(prefix="equiv_trim_")
	self.addCleanup(shutil.rmtree, workdir)

[julyanghar]

Addressed 4 of the 5 in f31dce1:

• nn.Softmax inside the loop → switched to F.softmax(dth, dim=2). Agreed —
  instantiating the module once per row-chunk was wasteful.
• Shifted-position concatenation → replaced with a single broadcast
  (sup.unsqueeze(1) + torch.arange(length + 1, device=sup.device)).view(-1).
  This is safe here because the result feeds straight into torch.unique, which
  sorts and de-duplicates, so the ordering change is immaterial.
• Per-step index gather → replaced the k separate index ops with one 2D
  advanced index + unbind(1).
• Test temp dir → now cleaned up via
  self.addCleanup(shutil.rmtree, workdir, ignore_errors=True).

I did not apply the logits.new_tensor() suggestion for loss_denom, for two
reasons:

1. That torch.tensor(...) construct is pre-existing on main — this PR only
   changes the multiplicand — so switching it would widen the diff beyond the
   feature itself.
2. More importantly, new_tensor() inherits the source tensor's dtype, and
   logits is bf16 here. The straightforward replacement would silently change
   loss_denom from fp32 to bf16. Keeping the explicit dtype=torch.float32
   seemed safer. Happy to switch to logits.new_tensor(..., dtype=torch.float32)
   if you'd prefer that form.

Re-ran the GPU equivalence test after the index-math vectorization — per-step loss
with trimming on vs off still matches.

[julyanghar]

Rebased onto current main (force-push) — and while doing so I found and fixed a real bug in my own patch. Details below.

The rebase

The unified-runtime refactor moved both files this PR touched: specforge/core/eagle3.py → specforge/algorithms/eagle3/model.py, and scripts/train_eagle3.py was removed, so the CLI flag is now a TrainingConfig field (training.trim_loss_positions) wired through model_providers.build_eagle3_model.

The bug I found while porting

The original patch assumed the supervised-row set is shared across all TTT steps, with the teacher sliding (row p fixed, teacher at p + j). That is backwards.

The loop applies padding(..., left=False) to position_mask / loss_mask once per step, so at step j the mask seen at row p is mask[p + j], while step_view slices the padded teacher so row p is supervised by the teacher at p + j. A row therefore contributes at step j iff p + j is supervised — i.e.

step j:  rows = {s - j : s in sup, s >= j},  teacher/mask pinned at those s

The rows shift down by one per step; the teacher positions stay put. My original code had it inverted (rows fixed, teacher sliding), which is correct at step 0 and wrong from step 1 onward.

Measured on a small fixture, old code vs the full-length path:

mask	bf16 max rel. err	fp32 max rel. err
all-supervised	25%	25%
half-masked	9.0e-3	9.2e-3

fp32 and bf16 give the same error, so it was a semantic bug, not numerical noise. (Step 0 matched exactly in every case, which is the signature of the per-step shift.)

Why my earlier checks missed it: my equivalence self-check compared the trimmed path against a reference that used the unshifted full-length mask — the reference shared the same wrong assumption, so both sides agreed. A self-check whose reference shares the bug can't detect it. The end-to-end test I had did compare against the real full path, but on data where the shift only moved a contiguous block's boundary by one position, so the difference stayed under the tolerance.

After the fix

Rows now follow sup - j and the teacher is evaluated once at sup (no sliding window needed, so the pack is simpler than before). Verified against the full-length path:

mask	bf16	fp32
all-supervised	0.0	0.0
half-masked (prompt-heavy)	1.1e-7	1.1e-7

Plus 27 existing eagle3 runtime tests and 50 config tests pass; the off path (flag defaults to off) is unchanged.

On the memory numbers I quoted earlier: those were measured on the pre-refactor code path, with the old row selection. The fix changes which rows are selected, not how many — the old version used n_sup rows at every step, the new one uses |{s in sup : s >= j}|, which is at most n_sup and differs by at most k entries — so the memory characteristics are unchanged and the reported saving still holds (if anything it is now marginally conservative). Happy to re-measure on the unified runtime if you'd like a number from the current code path.

Sorry for shipping that first version. Flagging it now rather than after you'd spent review time on it.

[jiapingW]

It's a good idea to reduce VRAM usage in Long context training of eagle3. But now the long context eagle3 training needs sequence parallel so It is important to verify the compatibility of the current method and sequence parallel. You can examine the compatibility with sequence parallelism and the corresponding gains. For instance, the performance boost observed with a ring size of 4 at a sequence length of 64k.

[julyanghar]

Great call asking about sequence parallelism — I dug in, and the honest answer is: as previously written, trim was NOT compatible with USP (it crashed, and in configurations where it didn't crash it would have been silently wrong). I've now made it SP-native, verified equivalence on 4-rank USP including pure ring-4, and measured the gains. Details below.

What was wrong

The trim branch bypassed the USP adapter's step_view, so with attention_backend=usp it fed the backbone this rank's whole local buffer (own chunk + the ttt_length overlap tail) instead of the chunk:

1. the backbone input length disagreed with the collator's position_ids (and with other ranks' ring/Ulysses block sizes) — on current code this fails fast with a shape error in torch.cat (verified on a 4-rank fixture);
2. supervised positions in the overlap tail would have emitted loss rows on this rank while also being the next rank's own rows — cross-rank double counting;
3. the loss was normalized by local_len instead of usp_chunk_size, breaking the per-rank-denominator convention that DDP + accumulation_steps ×= sp_size relies on.

And nothing guarded the combination — usp + trim_loss_positions silently took that path.

The fix (SP-native trim)

The backbone now runs exactly the full path's inputs under USP (chunk-sliced hidden/input_ids/attention-mask, adapter-owned position_ids); trimming applies only at the teacher/logits/loss stage:

• overlap-tail positions may act as teachers for own-chunk rows at deeper TTT steps (that's what the tail is for), but never emit loss rows themselves: rows_j = {s − j : s ∈ sup, j ≤ s, s − j < usp_chunk_size};
• the loss denominator is usp_chunk_size, matching the full path's mean-over-chunk semantics per rank;
• a step whose row set comes out empty (e.g. all local supervision sits in the tail at depth 0) is padded with one zero-masked dummy row, so it contributes exactly zero loss while kernel launches and collectives stay aligned across ranks — including the mixed case where some ranks trim and others (no local supervision) fall back to the full path.

Each (teacher-position, TTT-depth) loss term is computed exactly once globally.

Equivalence evidence (4 ranks, real offline pipeline)

Per-step losses, trim ON vs OFF on identical inputs, via the actual offline strategy path (TargetHead.preprocess → collator with use_usp_preprocess → forward), on both ulysses2 × ring2 and pure ring4, across four mask patterns chosen to be adversarial:

mask pattern	why it's there	result
all positions supervised	trivial-equality boundary: trim must select every chunk row; any row/denominator error is naked, no tolerance can hide it	max diff ≤ 2.4e-7 (1–2 ulp fp32; the smallest possible discrete error is ~5 orders larger)
~30% random supervised	realistic	≤ 1e-3 rel (bf16 GEMM-shape noise)
supervision straddling every rank boundary	overlap-tail-as-teacher on every rank	pass
supervision only in one rank's overlap tail	dead steps + two ranks with zero supervision (mixed trim/full ranks — the collective-alignment hazard)	pass, no hangs

Backward is covered too: total grad-norm trim vs full agrees within 5e-3 (calibrated against flash-attn's own backward nondeterminism floor). As a fixture sanity check, the mean over ranks of full-path per-step losses reproduces a single-GPU full-sequence reference on the un-sharded sample.

A negative control run — the same experiment against the previous trim code — fails immediately (the shape error above), so the test discriminates.

Gains at ring 4 (4× RTX 6000 Ada 48 GB, Llama3-8B-shaped: hidden 4096, target vocab 128256, draft vocab 32000, TTT 7, ~30% supervised)

global len	per-rank chunk	full peak GB/rank	trim peak GB/rank	saving	full step (s)	trim step (s)
8k	2k	8.22	7.54	0.68	0.431	0.385 (−10.7%)
16k	4k	13.60	12.17	1.43	0.939	0.813 (−13.4%)
32k	8k	24.37	21.45	2.92	2.193	1.960 (−10.6%)
64k	16k	OOM (all 4 ranks)	40.32	—	OOM	5.025

And the same sweep at ring 8 (8 GPUs):

global len	per-rank chunk	full peak GB/rank	trim peak GB/rank	saving	full step (s)	trim step (s)
64k	8k	24.37	21.61	2.76	2.903	2.726 (−6.1%)
128k	16k	OOM (all 8 ranks)	40.32	—	OOM	8.200

Takeaways:

• Trim moves the trainable-context frontier on this hardware: ring4/64k and ring8/128k both OOM on the full path and both run under trim (40.3 GB/rank in each case — identical per-rank chunk, which also cross-checks the sharding math: full at ring8/64k reproduces full at ring4/32k to the hundredth of a GB).
• The saving is additive with SP as expected (both attack the same [seq × vocab] activations along orthogonal axes) and scales linearly with length; trim is also consistently faster per step (~6–13%), since the teacher projection and draft lm_head run on ~30% of positions instead of all of them.

Equivalence was verified at world size 8 as well (pure ring-8 and ulysses2×ring4, same four masks — all pass), and beyond the differential tests the branch carries an oracle suite where expected outputs are derived independently of the implementation: hand-written literal row/keep tables for the per-step selection (including the overlap-tail bound and dead steps), an analytic closed form for the loss (zero logits + normalized teachers ⇒ each masked row contributes exactly ln(draft_vocab)), and a NumPy re-derivation of per-step row counts / teacher token ids / step losses straight from the raw feature files. The committed test file (tests/test_runtime/test_equiv_trim_usp.py) keeps the CPU-only golden-table layer running in CI even on hosts without four GPUs.

Caveats, stated plainly: these are 48 GB cards (not 80 GB), the benchmark is the raw module without FSDP (the delta we measure — vocab-sized loss-stage activations — is not sharded by FSDP anyway), and features are synthetic (memory/latency don't depend on tensor contents). The saving scales with 1 − supervised_fraction, so prompt-heavy long-context data benefits more. One more disclosure: the full-path baseline runs with compact_teacher off (the default), and the two features are currently mutually exclusive (the compact branch takes precedence). compact_teacher chunks the teacher's full-vocab fp32 transient over the vocab axis, but the full-length draft logits (+ their grad copies) and the [chunk × draft_vocab] teacher tables stay whole — those are what trim removes — so trim's saving should remain largely additive; we have not measured the full + compact_teacher combination at 64k.

The branch also picks up two robustness fixes that fell out of the verification: a dead-step NaN edge on the single-GPU path (a mask whose supervision is unreachable at deep TTT steps used to produce a mean over zero rows), and training.trim_loss_positions is now documented in examples/configs/README.md (the recipe-README guard).

Happy to adjust scope if you'd rather land the single-GPU part first and the USP part as a follow-up.

[julyanghar]

The CI failure looks unrelated to this PR: both errors come from tests/test_offline_capture/test_sglang_backend.py (its dense and MoE cases — the only 2 errors in the run), which dies in SGLang's model-runner init with

ValueError: The memory capacity is unbalanced. Some GPUs may be occupied by other processes. min_per_gpu_memory=1.71..., local_gpu_memory=17.13...

— i.e. the self-hosted runner appears to have leftover processes holding GPU memory. Recent runs on other branches show the same failure pattern since Aug 1 (last green was Jul 31). This PR's own tests (test_equiv_trim_loss_positions, test_equiv_trim_usp) pass both in this run and locally under the same unittest discover command. Could you re-trigger the workflow once the runner is freed up?