Skip to content

feat: add EAGLE3 draft training for Kimi K2.7 Code - #1

Open
maocheng23 wants to merge 1 commit into
Boreas618:reproductionfrom
maocheng23:eagle3-kimi-k2.7-code
Open

feat: add EAGLE3 draft training for Kimi K2.7 Code#1
maocheng23 wants to merge 1 commit into
Boreas618:reproductionfrom
maocheng23:eagle3-kimi-k2.7-code

Conversation

@maocheng23

Copy link
Copy Markdown

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 and max_length 2048 so the 1T target's per-rank capture activations fit on 184 GB GPUs (rationale in comments).
  • Compatibility fixes (needed for eagle3-online on a newer sglang + a multimodal target):
    • sglang_backend/utils.py: wrap_eagle3_logits_processors_in_module now snapshots named_modules() before mutating and resolves dotted (nested) names, so a multimodal target's nested language_model.logits_processor is actually replaced — previously it raised dictionary changed size during iteration, and even past that would have silently failed to wrap the nested processor. The replaced forward reads multi_item_delimiter_indices off the ForwardBatch (newer sglang dropped LogitsProcessor.multi_item_delimiter).
    • eagle3_target_model.py: logprob_start_len = -1 + set_extend_input_len(...) in extend()/extend_vlm() so prepare_mlp_sync_batch_raw's per-request logprob assertion holds under DP attention; plus a 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.

Base

Stacked on sgl-project#593 (reproduction) so the diff is only the EAGLE3 additions. Once sgl-project#593 lands, this can be retargeted to main.

Status

Training in progress on 4×GB200 (WORLD=16); draft loss is converging.

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>
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