feat: add EAGLE3 draft training for Kimi K2.7 Code - #1
Open
maocheng23 wants to merge 1 commit into
Open
Conversation
Adds an EAGLE3-online recipe for the Kimi K2.7 Code target (DeepseekV3-arch MoE, vocab 163840), stacked on the DFlash support in sgl-project#593, plus the sglang / transformers compatibility fixes needed to run eagle3-online against a newer sglang and a multimodal (KimiK25) target. - configs/kimi-k2.7-code-eagle3.json: EAGLE3 draft config (1-layer Llama, hidden 7168, eagle_aux_hidden_state_layer_ids [1,29,57], draft_vocab 32000). - examples/run_kimi_k2.7_code_eagle3.sh: multi-node launch; pure-TP capture + max_length 2048 so the 1T target's per-rank capture activations fit on 184GB. - sglang_backend/utils.py: wrap_eagle3_logits_processors_in_module snapshots named_modules() before mutating and resolves dotted (nested) module names, so a multimodal target's nested language_model.logits_processor is actually replaced; the replaced forward reads multi_item_delimiter_indices off the ForwardBatch (newer sglang dropped LogitsProcessor.multi_item_delimiter). - eagle3_target_model.py: set logprob_start_len=-1 + set_extend_input_len in extend()/extend_vlm() so prepare_mlp_sync_batch_raw's per-request logprob assert holds under DP attention; defensive import of prepare_mlp_sync_batch_raw. - auto.py: guard the transformers-4.57-only custom_backend import so AutoEagle3DraftModel loads on transformers>=5. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
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
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.
Summary
Adds an EAGLE3 draft-training recipe for the Kimi K2.7 Code target, stacked on the DFlash support in sgl-project#593. Includes the sglang / transformers compatibility fixes required to run
train_eagle3.py(online capture) against a current sglang and a multimodal (KimiK25ForConditionalGeneration) target.Companion to the DFlash draft from sgl-project#593 — same target, data, and tokenize cache — so EAGLE3 vs DFlash can be compared head-to-head (accept-length).
What's included
configs/kimi-k2.7-code-eagle3.json— EAGLE3 draft config: 1-layer Llama draft, hidden 7168,eagle_aux_hidden_state_layer_ids [1,29,57](valid for the 61-layer K2.7 target),vocab_size 163840,draft_vocab_size 32000, K2.7 token ids.examples/run_kimi_k2.7_code_eagle3.sh— multi-node launch. Uses pure TP for the capture andmax_length 2048so the 1T target's per-rank capture activations fit on 184 GB GPUs (rationale in comments).sglang_backend/utils.py:wrap_eagle3_logits_processors_in_modulenow snapshotsnamed_modules()before mutating and resolves dotted (nested) names, so a multimodal target's nestedlanguage_model.logits_processoris actually replaced — previously it raiseddictionary changed size during iteration, and even past that would have silently failed to wrap the nested processor. The replaced forward readsmulti_item_delimiter_indicesoff theForwardBatch(newer sglang droppedLogitsProcessor.multi_item_delimiter).eagle3_target_model.py:logprob_start_len = -1+set_extend_input_len(...)inextend()/extend_vlm()soprepare_mlp_sync_batch_raw's per-request logprob assertion holds under DP attention; plus a defensive import ofprepare_mlp_sync_batch_raw.auto.py: guard the transformers-4.57-onlycustom_backendimport soAutoEagle3DraftModelloads on transformers ≥ 5.Base
Stacked on sgl-project#593 (
reproduction) so the diff is only the EAGLE3 additions. Once sgl-project#593 lands, this can be retargeted tomain.Status
Training in progress on 4×GB200 (WORLD=16); draft loss is converging.