Eagle3: compact supervised positions before vocabulary projection

[julyanghar]

Summary

This PR ports SpecForge's A-level Eagle3 loss-position compaction into DeepSpec's cached-teacher training path.

• keep the teacher LM-head projection full-length;
• keep every TTT backbone step full-length, including Q/K/V, attention, cache updates, O projection, and MLP;
• compact only teacher probability conversion, draft norm / LM-head, metrics, and fused soft cross-entropy;
• preserve the original local_mean = batch_size * sequence_length denominator;
• keep the feature opt-in with model.trim_loss_positions=false by default.

The implementation builds exact per-step draft/teacher row maps and computes one shared compact teacher-probability table. Indexed Triton kernels for metrics and soft cross-entropy read target rows directly from that table, so no per-step [N_t, V] target tensor is materialized in forward or backward. Empty local steps execute one zero-masked dummy row so FSDP ranks retain aligned module execution.

Correctness and scope

• A direct-index oracle compares indexed metrics, loss, and dlogits with an explicitly materialized reference. A test-only dispatch counter observes teacher-table index_select in the reference and zero such materializations in the candidate; shifted indices are detected as a negative control.
• Tiny Qwen3 full-vs-compact forward/backward testing compares total loss, every trainable gradient, and all Eagle3 training metrics, with a deliberately shifted row map as a second negative control.
• Scope hooks verify three q_proj calls remain length 16, the teacher LM-head remains length 16, only draft LM-head rows become [5, 4, 4], and DynamicCache grows [16, 32, 48].
• A two-rank FSDP test includes a rank with zero valid rows at later TTT steps; optimizer updates match the full path (max_parameter_diff=0) and parameters remain rank-consistent.
• Qwen3 and Gemma4 share the same optional post-backbone logit-index interface.

Performance

Synthetic values with real Qwen3-8B draft shapes, BF16, one Eagle3 layer, sequence length 4096, TTT length 7, contiguous-tail supervision, RTX 6000 Ada. Timing covers one Eagle3 forward + backward, excluding model construction, data generation, optimizer, I/O, and profiler. Each result uses 3 alternating warm-up pairs and 7 alternating measured pairs.

supervised density	paired-median speedup	full / compact median	full / compact peak allocated
30%	1.735x	1350.0 / 775.5 ms	27.696 / 13.055 GiB
60%	1.294x	1524.3 / 1171.4 ms	27.696 / 16.400 GiB
90%	1.062x	1671.0 / 1571.9 ms	27.696 / 19.748 GiB

The direct-index path removes the previous dense-supervision crossover as well: at 90% supervision it now reduces peak allocated memory by 7.948 GiB and improves forward+backward time by 6.2%.

Usage

python train.py --config config/eagle3/eagle3_qwen3_8b.py \
  --opts model.trim_loss_positions=true \
  --opts data.target_cache_path=/path/to/target_cache

Tests

CUDA_VISIBLE_DEVICES=6 pytest -q tests/test_eagle3_trim_loss_positions.py
CUDA_VISIBLE_DEVICES=6,7 torchrun --standalone --nproc_per_node=2 \
  tests/distributed/run_eagle3_trim_loss_positions_fsdp.py

[julyanghar]

Hi @Hannibal046, could you please take a look when you have time?

This PR only compacts EAGLE3's vocabulary-side supervised rows. The teacher LM-head projection and every TTT backbone/QKV/attention step remain full-length.

It includes:

• loss, gradient, and metric equivalence with negative controls;
• two-rank FSDP optimizer-update equivalence;
• direct-index kernels with zero teacher-table row materialization;
• Qwen3-8B, 4K, k=7 measurements across 30/60/90% supervision density.

I'm happy to adjust the API or split the benchmark/tests if a smaller PR is preferred.