[feat](gpt-oss): Eagle3 speculative decoding support for gpt-oss-120b#1605
Open
ProgMastermind wants to merge 5 commits into
Open
[feat](gpt-oss): Eagle3 speculative decoding support for gpt-oss-120b#1605ProgMastermind wants to merge 5 commits into
ProgMastermind wants to merge 5 commits into
Conversation
- Pass dtype and rope_scaling to RoPE in Eagle3LlamaAttention. - Add norm_before_fc input RMSNorm before the Eagle3 fc projection. - Capture the final target layer auxiliary hidden state in GptOssModel. - Fix fc_norm chunking to use the contiguous concatenated tensor.
…t when missing NVIDIA's LlamaForCausalLMEagle3 checkpoints for gpt-oss-120b (e.g. nvidia/gpt-oss-120b-Eagle3-short-context) do not ship embed_tokens.weight; the draft is expected to share the target model's tokenizer and embedding matrix. EagleProposer.load_model now falls back to the target model's embed_tokens (and lm_head if absent) for eagle3 when the draft checkpoint did not load them.
Contributor
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
Contributor
There was a problem hiding this comment.
Pull request overview
Adds Eagle3 speculative decoding support for openai/gpt-oss-120b, including wiring target aux hidden states into the Eagle3 draft model and improving Eagle3 draft model compatibility with draft checkpoints that omit certain weights.
Changes:
- Capture the final aux hidden state output for GPT-OSS to support Eagle3 aux-hidden-state drafts.
- Update Eagle3 Llama draft internals (RoPE kwargs, optional pre-FC RMSNorm, and fc_norm chunking behavior).
- Allow Eagle3 draft loading to reuse target
embed_tokens/lm_headwhen missing from the draft checkpoint, and document Eagle3 usage in the GPT-OSS recipe.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| recipes/GPT-OSS.md | Documents Eagle3 speculative decoding launch example and notes weight sharing behavior. |
| atom/spec_decode/eagle.py | Loads Eagle3 draft weights and conditionally shares missing embedding / lm_head from the target model. |
| atom/models/gpt_oss.py | Captures an additional aux hidden state at the end of the local layer range for Eagle3 aux-hidden-state support. |
| atom/models/eagle3_llama.py | Passes RoPE dtype/scaling through draft config and adds optional input RMSNorm + fixes fc_norm chunking path. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+300
to
+320
| # Eagle3 checkpoints may omit the embedding matrix (e.g. NVIDIA's | ||
| # LlamaForCausalLMEagle3 speculators share the target tokenizer). | ||
| # Share embed_tokens/lm_head from the target when the draft didn't | ||
| # load them, so the draft can run without needing extra weights. | ||
| target_base = getattr(target_model, "language_model", target_model) | ||
| self._share_if_not_loaded( | ||
| self.model, | ||
| "embed_tokens", | ||
| target_base.model.embed_tokens, | ||
| loaded, | ||
| "embed_tokens.weight", | ||
| "Eagle3 embed_tokens", | ||
| ) | ||
| self._share_if_not_loaded( | ||
| self.model, | ||
| "lm_head", | ||
| target_base.lm_head, | ||
| loaded, | ||
| "lm_head.weight", | ||
| "Eagle3 lm_head", | ||
| ) |
…e parallelism With pipeline parallelism, target_base.model.embed_tokens can be a PPMissingLayer placeholder on non-first ranks. Only share the embedding when the target actually owns the real weight on this rank. The lm_head sharing is left unconditional because the draft and lm_head live on the same PP rank (last rank).
Comment on lines
+304
to
+327
| target_base = getattr(target_model, "language_model", target_model) | ||
| target_embed = getattr( | ||
| getattr(target_base, "model", None), "embed_tokens", None | ||
| ) | ||
| # Only share the target embed_tokens when the target actually owns a | ||
| # real embedding on this rank. With pipeline parallelism, ranks other | ||
| # than the first may hold a PPMissingLayer placeholder instead of the | ||
| # actual embedding, and using that would break the draft. | ||
| if ( | ||
| get_pp_group().world_size == 1 | ||
| and target_embed is not None | ||
| and hasattr(target_embed, "weight") | ||
| and getattr(self.model, "embed_tokens", None) is not None | ||
| and hasattr(self.model.embed_tokens, "weight") | ||
| and self.model.embed_tokens.weight.shape == target_embed.weight.shape | ||
| ): | ||
| self._share_if_not_loaded( | ||
| self.model, | ||
| "embed_tokens", | ||
| target_embed, | ||
| loaded, | ||
| "embed_tokens.weight", | ||
| "Eagle3 embed_tokens", | ||
| ) |
Comment on lines
+62
to
+63
| The draft model shares the target tokenizer/embedding and uses the target's | ||
| `lm_head` when those weights are not present in the draft checkpoint. |
Raise a clear RuntimeError when the Eagle3 draft checkpoint is missing embed_tokens.weight and ATOM cannot share the target embedding (e.g. pipeline parallelism or incompatible tensor-parallel sharding). This prevents silent use of randomly-initialized draft embeddings. Also update recipes/GPT-OSS.md to document the sharing constraints.
valarLip
reviewed
Jul 15, 2026
| ) | ||
| else: | ||
| fc_in = aux_hidden_states | ||
| if self.input_norm is not None: |
Collaborator
There was a problem hiding this comment.
it‘s same thing like self.fc_norm?
Author
There was a problem hiding this comment.
Not the same , input_norm is one RMSNorm over the whole concatenated aux tensor, while fc_norm is a ModuleList that normalizes each aux chunk separately. The gpt-oss-120b checkpoint config has norm_before_fc=True, so the model expects the full [N, target_hidden_size * num_aux] tensor to be normalized before fc. fc_norm alone would normalize per chunk, which is a different operation, so input_norm is needed here.
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 Eagle3 speculative decoding support for
openai/gpt-oss-120b.Changes
atom/models/gpt_oss.py: capture final target layeraux_hidden_stateand feed it to the Eagle3 draft.atom/models/eagle3_llama.py: passdtypeandrope_scalingto the draft RoPE; addnorm_before_fcinput RMSNorm before the Eagle3fcprojection; fixfc_normchunking.atom/spec_decode/eagle.py: share targetembed_tokens/lm_headwith the Eagle3 draft when the draft checkpoint does not include them.recipes/GPT-OSS.md: add Eagle3 server example and accuracy test command.Verification
Server starts with:
python -m atom.benchmarks.benchmark_serving(single stream, random ISL=9592 / OSL=1500, CONC=1): 158.5 tok/s output throughput.lm_evalGSM8K 5-shot: flexible-extract = 0.9030 (CI threshold >= 0.88).lm_evalGSM8K 3-shot: flexible-extract = 0.9007 (CI threshold >= 0.88).