feat(draft): support sequence packing with online eagle3 draft training - #3463
Open
yuekaizhang wants to merge 1 commit into
Open
feat(draft): support sequence packing with online eagle3 draft training#3463yuekaizhang wants to merge 1 commit into
yuekaizhang wants to merge 1 commit into
Conversation
Contributor
Author
|
/ok to test 114598c |
github-actions
Bot
added
the
Documentation
Online draft co-training previously forced sequence_packing off, which roughly halves training/logprob throughput on MoE policies (forfeits permute fusion) and made the eagle3 run slower end-to-end than a packed baseline despite the spec-decode generation speedup. - loss/utils: add roll_packed_seq_dim (per-segment left shift via one global roll + zeroed segment boundaries) and pack_rolled_draft_token_mask; factor the d2t teacher-vocab mapping out of prepare_loss_input into map_teacher_logits_to_draft_vocab. - DraftLossWrapper: packed mode computes the draft soft-CE once over the packed layout (teacher shifted within segments, mask packed with the same shift); numerically identical to the unpacked path. - train.py: in packed mode, build the draft's shifted input embeddings by re-embedding per-sequence-shifted token ids (_pack_input_ids roll_shift=-1) instead of rolling captured embeddings across packing boundaries; attach DraftLossWrapper on top of the packing loss wrappers (fused and non-fused) and pass packed_seq_params to the draft forward. - EagleModel: forward packed_seq_params into the modelopt EagleModule and temporarily override the eagle layers' AttnMaskType.arbitrary with padding_causal for THD (TE's THD kernels reject arbitrary; mcore only auto-converts causal/no_mask). - lm_policy: replace the draft-vs-packing guard with an explicit draft-vs-context-parallelism guard (CP was previously blocked only transitively through packing). Validated by unit tests asserting packed == unpacked draft loss (with and without d2t) and a 2-step 8B GPU run per packing loss wrapper; step-1 losses match bit-for-bit between the fused and non-fused wrappers. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Yuekai Zhang <zhangyuekai@foxmail.com>
Contributor
Author
|
/ok to test 01699bc |
yuekaizhang
added
CI:L0
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Adds sequence packing support to online EAGLE-3 draft training (previously the
two features were mutually exclusive). With packing enabled, policy training
time is cut roughly in half (~2100 vs ~1040 tok/s/GPU on a Qwen3-30B-A3B 8K
GRPO run), so draft co-training no longer pays the unpacked-training penalty.
Key changes:
(boundary-safe rolls; no leakage across packing boundaries)
DraftLossWrapperwith a packed token maskpadding_causalmask-type override for TE kernels)
Unit tests prove packed/unpacked numerical equivalence (full-vocab and
d2t-mapped drafts); verified end-to-end on 1n8g (8B) and 4n8g (30B-A3B) GRPO
runs with no acceptance regression.
Blue curve: without the PR
Orange curve: with the PR