Skip to content

Skip unsupervised teacher projections in pruned EAGLE3 targets - #15

Open
julyanghar wants to merge 1 commit into
Tencent:mainfrom
julyanghar:perf/defer-teacher-projection
Open

julyanghar wants to merge 1 commit into
Tencent:mainfrom
julyanghar:perf/defer-teacher-projection

Conversation

@julyanghar

Copy link
Copy Markdown

With vocabulary pruning and sparse supervision, compute_target_p_padded currently projects every teacher row into the draft vocabulary and softmaxes the full sequence, even though TTT loss only reads supervised rows. This change projects supervised hidden rows in the existing chunks and scatters their probabilities into the original dense lookup table. Unsupervised rows are zero; shapes, padding, and position-mask semantics are preserved.

When all rows are supervised, retain the dense path to avoid gather/scatter overhead. The lazy target path and draft backbone are unchanged.

Validation:

  • 14 new target-projection tests passed. Expanded CPU check: 54 passed, 1 skipped, 11 deselected with TORCH_COMPILE_DISABLE=1 python -m pytest tests/test_target_projection.py tests/test_mtp.py -k "not TestMTPServeAlignment and not TestMTPUlyssesAllToAll" -q.
  • New tests cover sparse/full/empty/last-token masks, multiple batches, chunk boundaries, projection row counts, shifted target/mask lookup, and a real four-round EAGLE3 model comparison of losses and parameter gradients.
  • Ruff check/format and git diff --check passed for changed files.
  • Broader MTP testing encountered four TestMTPServeAlignment import failures due to missing system libibverbs.so.1; all four also reproduce on untouched base d3412be. Distributed MTP coverage was not completed and is excluded above. No full training/distributed validation is claimed.

Function-only CUDA benchmark (not end-to-end training): RTX 6000 Ada, PyTorch 2.11.0+cu130, BF16, B=1, H=1024, full vocabulary=65536, draft vocabulary=32768, depth=7, chunk=4096. Identical inputs per arm, three warmup calls, four alternating AB/BA pairs with five calls per arm, synchronized host wall time and no profiler. Times below are averages in milliseconds; reduction denominator is baseline time. BF16 valid-row probabilities and position masks were also compared.

Sequence rows Supervised rows Before ms After ms Time reduction
2048 204 3.295 1.958 40.58%
2048 2048 6.220 6.219 0.01%
8192 819 12.205 7.163 41.31%
8192 8192 24.017 24.020 -0.02%

The full probability lookup allocation remains; this PR does not introduce compact target-table storage. Full training throughput and training quality were not measured.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant