From fd2230b2f89315f6fa4c6579026e511ffd343b52 Mon Sep 17 00:00:00 2001 From: Orestis Zambounis Date: Tue, 23 Jun 2026 15:09:17 +0000 Subject: [PATCH 1/9] Add per-layer FC normalization (--fc-norm) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add fc_norm: per-layer RMSNorm on each auxiliary hidden state before concatenation and FC projection — concat(Norm(h_a), Norm(h_b), Norm(h_c)) — matching the Eagle 3.1 paper specification. This differs from norm_before_fc which applies a single norm to the full concatenation. Defaults to True for llama draft arch (along with norm_output). norm_before_fc no longer defaults to True for llama arch since fc_norm is the more principled alternative. Co-Authored-By: Claude Opus 4.6 Signed-off-by: Orestis Zambounis --- docs/cli/train.md | 4 ++- scripts/train.py | 35 +++++++++++++++---- .../convert/eagle/eagle3_converter.py | 4 +++ src/speculators/models/eagle3/config.py | 10 ++++++ src/speculators/models/eagle3/core.py | 23 ++++++++++++ src/speculators/models/peagle/core.py | 10 ++++++ tests/integration/conftest.py | 4 +++ .../integration/models/test_model_forward.py | 24 +++++++++++++ tests/unit/test_config.py | 17 +++++++++ 9 files changed, 123 insertions(+), 8 deletions(-) diff --git a/docs/cli/train.md b/docs/cli/train.md index e127e6af7..8afbea91a 100644 --- a/docs/cli/train.md +++ b/docs/cli/train.md @@ -136,7 +136,9 @@ torchrun --standalone --nproc_per_node=4 scripts/train.py \ - **`--embed-requires-grad` / `--no-embed-requires-grad`** (flag, default: `False`) Whether to train embedding layer weights. -- **`--norm-before-fc`** (flag, default: `False`) Use RMSNorm before FC layer in draft path (e.g., for Eagle 3.1 / gpt-oss models). +- **`--norm-before-fc` / `--no-norm-before-fc`** (flag, default: `False`) Apply a single RMSNorm to concatenated target hidden states before the FC projection layer. See `--fc-norm` for the per-layer alternative. + +- **`--fc-norm` / `--no-fc-norm`** (flag, default: `True` for llama arch, `False` otherwise) Apply per-layer RMSNorm to each auxiliary hidden state before concatenation and FC projection — i.e. `concat(Norm(h_a), Norm(h_b), Norm(h_c))`. - **`--norm-output`** (flag, default: `False`) Feed post-norm hidden states back across TTT steps to stabilize magnitude drift across speculation depths (Eagle 3.1). diff --git a/scripts/train.py b/scripts/train.py index 4a1582c47..fc549f9d2 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -484,6 +484,19 @@ def main(args: argparse.Namespace): # noqa: C901 ) hidden_states_dtype = getattr(torch, args.hidden_states_dtype) + # Default to Eagle 3.1 (fc_norm + norm_output) for llama arch + is_llama_eagle = ( + args.speculator_type in ("eagle3", "peagle") + and getattr(args, "draft_arch", None) == "llama" + ) + if args.norm_before_fc is None: + args.norm_before_fc = False + if args.fc_norm is None: + args.fc_norm = is_llama_eagle + if args.norm_output is None: + args.norm_output = is_llama_eagle + + if args.speculator_type == "mtp": if args.draft_attn_impl != "simple_flex_attention": raise ValueError( @@ -1016,17 +1029,25 @@ def parse_args(): ) parser.add_argument( "--norm-before-fc", - action="store_true", - default=False, - help="Use RMSNorm before FC layer in draft path " - "(e.g., for Eagle 3.1 / gpt-oss models).", + action=argparse.BooleanOptionalAction, + default=None, + help="Apply a single RMSNorm to concatenated target hidden states before " + "the FC projection layer. See --fc-norm for per-layer alternative.", + ) + parser.add_argument( + "--fc-norm", + action=argparse.BooleanOptionalAction, + default=None, + help="Apply per-layer RMSNorm to each auxiliary hidden state before " + "concatenation and FC projection. Defaults to True for llama draft arch.", ) parser.add_argument( "--norm-output", - action="store_true", - default=False, + action=argparse.BooleanOptionalAction, + default=None, help="Feed post-norm hidden states back across TTT steps to stabilize " - "magnitude drift across speculation depths (Eagle 3.1).", + "magnitude drift across speculation depths (Eagle 3.1). " + "Defaults to True for llama draft arch.", ) # D-Flash specific parameters parser.add_argument( diff --git a/src/speculators/convert/eagle/eagle3_converter.py b/src/speculators/convert/eagle/eagle3_converter.py index 98f417658..f25d560fb 100644 --- a/src/speculators/convert/eagle/eagle3_converter.py +++ b/src/speculators/convert/eagle/eagle3_converter.py @@ -40,6 +40,7 @@ def convert( validate: bool = True, norm_before_residual: bool = False, norm_before_fc: bool = False, + fc_norm: bool = False, norm_output: bool = False, eagle_aux_hidden_state_layer_ids: list[int] | None = None, cache_dir: str | Path | None = None, @@ -80,6 +81,7 @@ def convert( base_model, norm_before_residual, norm_before_fc, + fc_norm, norm_output, eagle_aux_hidden_state_layer_ids, ) @@ -111,6 +113,7 @@ def _build_eagle3_speculator_config( base_model: str, norm_before_residual: bool = False, norm_before_fc: bool = False, + fc_norm: bool = False, norm_output: bool = False, eagle_aux_hidden_state_layer_ids: list[int] | None = None, ) -> Eagle3SpeculatorConfig: @@ -137,6 +140,7 @@ def _build_eagle3_speculator_config( draft_vocab_size=eagle_config.get("draft_vocab_size", 32000), norm_before_residual=norm_before_residual, norm_before_fc=norm_before_fc or eagle_config.get("norm_before_fc", False), + fc_norm=fc_norm or eagle_config.get("fc_norm", False), norm_output=norm_output or eagle_config.get("norm_output", False), target_hidden_size=eagle_config.get("target_hidden_size"), eagle_aux_hidden_state_layer_ids=eagle_aux_hidden_state_layer_ids, diff --git a/src/speculators/models/eagle3/config.py b/src/speculators/models/eagle3/config.py index c394c2fd2..17666f409 100644 --- a/src/speculators/models/eagle3/config.py +++ b/src/speculators/models/eagle3/config.py @@ -63,6 +63,16 @@ class Eagle3SpeculatorConfig(SpeculatorModelConfig): ), ) + fc_norm: bool = Field( + default=False, + description=( + "Apply per-layer RMSNorm to each auxiliary hidden state before " + "concatenation and FC projection — i.e. " + "concat(Norm(h_a), Norm(h_b), Norm(h_c)) instead of the single " + "Norm(concat(h_a, h_b, h_c)) used by norm_before_fc." + ), + ) + norm_output: bool = Field( default=False, description=( diff --git a/src/speculators/models/eagle3/core.py b/src/speculators/models/eagle3/core.py index fe51e7ff0..974b9216e 100644 --- a/src/speculators/models/eagle3/core.py +++ b/src/speculators/models/eagle3/core.py @@ -101,6 +101,19 @@ def __init__(self, config: Eagle3SpeculatorConfig): else: self.input_norm = None + if config.fc_norm: + self.fc_norms = torch.nn.ModuleList( + [ + self._model_definitions.norm_class( + self.hidden_size, + eps=config.transformer_layer_config.rms_norm_eps, + ) + for _ in range(3) + ] + ) + else: + self.fc_norms = None + self.post_init() @property @@ -179,6 +192,15 @@ def forward( # noqa: C901 if self.input_norm is not None: hidden_states = self.input_norm(hidden_states) + if self.fc_norms is not None: + chunks = hidden_states.chunk(len(self.fc_norms), dim=-1) + hidden_states = torch.cat( + [ + norm(chunk) + for norm, chunk in zip(self.fc_norms, chunks, strict=True) + ], + dim=-1, + ) hidden_states = self.fc(hidden_states) # shape: [1, total_seq_len, hidden_size] @@ -330,6 +352,7 @@ def from_training_args( draft_vocab_size=kwargs["draft_vocab_size"], norm_before_residual=kwargs["norm_before_residual"], norm_before_fc=kwargs.get("norm_before_fc", False), + fc_norm=kwargs.get("fc_norm", False), norm_output=kwargs.get("norm_output", False), embed_requires_grad=kwargs.get("embed_requires_grad", False), eagle_aux_hidden_state_layer_ids=target_layer_ids, diff --git a/src/speculators/models/peagle/core.py b/src/speculators/models/peagle/core.py index c83f59ba9..cdb2b4678 100644 --- a/src/speculators/models/peagle/core.py +++ b/src/speculators/models/peagle/core.py @@ -117,6 +117,15 @@ def forward( # Project concatenated hidden states (3*hidden_size) -> hidden_size if self.input_norm is not None: sampled_hidden = self.input_norm(sampled_hidden) + if self.fc_norms is not None: + chunks = sampled_hidden.chunk(len(self.fc_norms), dim=-1) + sampled_hidden = torch.cat( + [ + norm(chunk) + for norm, chunk in zip(self.fc_norms, chunks, strict=True) + ], + dim=-1, + ) sampled_hidden = self.fc(sampled_hidden) # [1, total_sampled, hidden_size] layer_input = torch.cat( @@ -216,6 +225,7 @@ def from_training_args( draft_vocab_size=kwargs["draft_vocab_size"], norm_before_residual=kwargs.get("norm_before_residual", False), norm_before_fc=kwargs.get("norm_before_fc", False), + fc_norm=kwargs.get("fc_norm", False), norm_output=kwargs.get("norm_output", False), eagle_aux_hidden_state_layer_ids=target_layer_ids, num_depths=kwargs.get("num_depths", 8), diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 934d92e99..55df572eb 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -98,6 +98,7 @@ def make_eagle3_model( draft_vocab_size: int = 64, norm_before_residual: bool = False, norm_before_fc: bool = False, + fc_norm: bool = False, norm_output: bool = False, draft_attn_impl: str | None = None, device: str = "cuda:0", @@ -111,6 +112,7 @@ def make_eagle3_model( draft_vocab_size=draft_vocab_size, norm_before_residual=norm_before_residual, norm_before_fc=norm_before_fc, + fc_norm=fc_norm, norm_output=norm_output, embed_requires_grad=False, speculators_config=SpeculatorsConfig( @@ -171,6 +173,7 @@ def make_peagle_model( num_depths: int = 4, down_sample_ratio: float = 0.7, norm_before_fc: bool = False, + fc_norm: bool = False, norm_output: bool = False, draft_attn_impl: str | None = None, device: str = "cuda:0", @@ -184,6 +187,7 @@ def make_peagle_model( draft_vocab_size=draft_vocab_size, norm_before_residual=False, norm_before_fc=norm_before_fc, + fc_norm=fc_norm, norm_output=norm_output, embed_requires_grad=True, num_depths=num_depths, diff --git a/tests/integration/models/test_model_forward.py b/tests/integration/models/test_model_forward.py index 7647b53ce..1f6ad7d73 100644 --- a/tests/integration/models/test_model_forward.py +++ b/tests/integration/models/test_model_forward.py @@ -382,6 +382,30 @@ def test_norm_output_without_norm_before_fc(self): assert loss.isfinite() loss.backward() + def test_fc_norm(self): + model = make_eagle3_model(fc_norm=True, norm_output=True) + assert model.fc_norms is not None + assert len(model.fc_norms) == 3 + assert model.input_norm is None + samples = _make_samples([128]) + batch = make_batch(max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE) + draft_tokens, loss, _metrics = model(**batch, ttt_steps=3) + + assert len(draft_tokens) == 3 + assert loss.isfinite() + loss.backward() + + def test_peagle_fc_norm(self): + model = make_peagle_model(fc_norm=True) + assert model.fc_norms is not None + assert len(model.fc_norms) == 3 + samples = _make_samples([128]) + batch = make_batch(max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE) + _draft_tokens, loss, _metrics = model(**batch) + + assert loss.isfinite() + loss.backward() + def test_peagle_norm_before_fc(self): model = make_peagle_model() assert model.input_norm is None diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index 3506da76b..8019dbea0 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -538,3 +538,20 @@ def test_eagle3_config_norm_output_defaults(): speculators_config=_make_eagle3_speculators_config(), ) assert config.norm_output is False + assert config.fc_norm is False + + +@pytest.mark.sanity +def test_eagle3_config_fc_norm_roundtrip(): + original = Eagle3SpeculatorConfig( + transformer_layer_config=copy.deepcopy(TINY_LLAMA_CONFIG), + draft_vocab_size=32000, + fc_norm=True, + speculators_config=_make_eagle3_speculators_config(), + ) + + config_dict = original.to_dict() + assert config_dict["fc_norm"] is True + + reloaded = SpeculatorModelConfig.from_dict(config_dict) + assert reloaded.fc_norm is True From 30422f48fa91d0d72b70c3f1b7023e5e063ca110 Mon Sep 17 00:00:00 2001 From: Orestis Zambounis Date: Fri, 26 Jun 2026 12:18:15 +0000 Subject: [PATCH 2/9] refactor: rename fc_norms to fc_norm for vLLM weight compatibility The vLLM model uses `self.fc_norm` (singular) for the ModuleList, producing weight names like `fc_norm.0.weight`. Align the speculators attribute name so checkpoints are directly loadable without remapping. Co-Authored-By: Claude Opus 4.6 Signed-off-by: Orestis Zambounis --- scripts/train.py | 1 - src/speculators/models/eagle3/core.py | 13 +++++-------- src/speculators/models/peagle/core.py | 6 +++--- tests/integration/models/test_model_forward.py | 8 ++++---- 4 files changed, 12 insertions(+), 16 deletions(-) diff --git a/scripts/train.py b/scripts/train.py index fc549f9d2..fa27257c0 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -496,7 +496,6 @@ def main(args: argparse.Namespace): # noqa: C901 if args.norm_output is None: args.norm_output = is_llama_eagle - if args.speculator_type == "mtp": if args.draft_attn_impl != "simple_flex_attention": raise ValueError( diff --git a/src/speculators/models/eagle3/core.py b/src/speculators/models/eagle3/core.py index 974b9216e..36fcab639 100644 --- a/src/speculators/models/eagle3/core.py +++ b/src/speculators/models/eagle3/core.py @@ -102,7 +102,7 @@ def __init__(self, config: Eagle3SpeculatorConfig): self.input_norm = None if config.fc_norm: - self.fc_norms = torch.nn.ModuleList( + self.fc_norm = torch.nn.ModuleList( [ self._model_definitions.norm_class( self.hidden_size, @@ -112,7 +112,7 @@ def __init__(self, config: Eagle3SpeculatorConfig): ] ) else: - self.fc_norms = None + self.fc_norm = None self.post_init() @@ -192,13 +192,10 @@ def forward( # noqa: C901 if self.input_norm is not None: hidden_states = self.input_norm(hidden_states) - if self.fc_norms is not None: - chunks = hidden_states.chunk(len(self.fc_norms), dim=-1) + if self.fc_norm is not None: + chunks = hidden_states.chunk(len(self.fc_norm), dim=-1) hidden_states = torch.cat( - [ - norm(chunk) - for norm, chunk in zip(self.fc_norms, chunks, strict=True) - ], + [norm(chunk) for norm, chunk in zip(self.fc_norm, chunks, strict=True)], dim=-1, ) hidden_states = self.fc(hidden_states) diff --git a/src/speculators/models/peagle/core.py b/src/speculators/models/peagle/core.py index cdb2b4678..188b6746c 100644 --- a/src/speculators/models/peagle/core.py +++ b/src/speculators/models/peagle/core.py @@ -117,12 +117,12 @@ def forward( # Project concatenated hidden states (3*hidden_size) -> hidden_size if self.input_norm is not None: sampled_hidden = self.input_norm(sampled_hidden) - if self.fc_norms is not None: - chunks = sampled_hidden.chunk(len(self.fc_norms), dim=-1) + if self.fc_norm is not None: + chunks = sampled_hidden.chunk(len(self.fc_norm), dim=-1) sampled_hidden = torch.cat( [ norm(chunk) - for norm, chunk in zip(self.fc_norms, chunks, strict=True) + for norm, chunk in zip(self.fc_norm, chunks, strict=True) ], dim=-1, ) diff --git a/tests/integration/models/test_model_forward.py b/tests/integration/models/test_model_forward.py index 1f6ad7d73..876221617 100644 --- a/tests/integration/models/test_model_forward.py +++ b/tests/integration/models/test_model_forward.py @@ -384,8 +384,8 @@ def test_norm_output_without_norm_before_fc(self): def test_fc_norm(self): model = make_eagle3_model(fc_norm=True, norm_output=True) - assert model.fc_norms is not None - assert len(model.fc_norms) == 3 + assert model.fc_norm is not None + assert len(model.fc_norm) == 3 assert model.input_norm is None samples = _make_samples([128]) batch = make_batch(max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE) @@ -397,8 +397,8 @@ def test_fc_norm(self): def test_peagle_fc_norm(self): model = make_peagle_model(fc_norm=True) - assert model.fc_norms is not None - assert len(model.fc_norms) == 3 + assert model.fc_norm is not None + assert len(model.fc_norm) == 3 samples = _make_samples([128]) batch = make_batch(max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE) _draft_tokens, loss, _metrics = model(**batch) From e5652d6eba9d27a7478df97aa459b6b7639fe8ec Mon Sep 17 00:00:00 2001 From: Orestis Zambounis Date: Fri, 26 Jun 2026 14:09:02 +0000 Subject: [PATCH 3/9] refactor: simplify norm flags to store_true, remove defaulting logic Co-Authored-By: Claude Opus 4.6 Signed-off-by: Orestis Zambounis --- docs/cli/train.md | 4 ++-- scripts/train.py | 33 ++++++++++----------------------- 2 files changed, 12 insertions(+), 25 deletions(-) diff --git a/docs/cli/train.md b/docs/cli/train.md index 8afbea91a..ba3764ac2 100644 --- a/docs/cli/train.md +++ b/docs/cli/train.md @@ -136,9 +136,9 @@ torchrun --standalone --nproc_per_node=4 scripts/train.py \ - **`--embed-requires-grad` / `--no-embed-requires-grad`** (flag, default: `False`) Whether to train embedding layer weights. -- **`--norm-before-fc` / `--no-norm-before-fc`** (flag, default: `False`) Apply a single RMSNorm to concatenated target hidden states before the FC projection layer. See `--fc-norm` for the per-layer alternative. +- **`--norm-before-fc`** (flag, default: `False`) Use RMSNorm before FC layer in draft path (e.g., for Eagle 3.1 / gpt-oss models). -- **`--fc-norm` / `--no-fc-norm`** (flag, default: `True` for llama arch, `False` otherwise) Apply per-layer RMSNorm to each auxiliary hidden state before concatenation and FC projection — i.e. `concat(Norm(h_a), Norm(h_b), Norm(h_c))`. +- **`--fc-norm`** (flag, default: `False`) Apply per-layer RMSNorm to each auxiliary hidden state before concatenation and FC projection (Eagle 3.1 paper approach). - **`--norm-output`** (flag, default: `False`) Feed post-norm hidden states back across TTT steps to stabilize magnitude drift across speculation depths (Eagle 3.1). diff --git a/scripts/train.py b/scripts/train.py index fa27257c0..d44deca5e 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -484,18 +484,6 @@ def main(args: argparse.Namespace): # noqa: C901 ) hidden_states_dtype = getattr(torch, args.hidden_states_dtype) - # Default to Eagle 3.1 (fc_norm + norm_output) for llama arch - is_llama_eagle = ( - args.speculator_type in ("eagle3", "peagle") - and getattr(args, "draft_arch", None) == "llama" - ) - if args.norm_before_fc is None: - args.norm_before_fc = False - if args.fc_norm is None: - args.fc_norm = is_llama_eagle - if args.norm_output is None: - args.norm_output = is_llama_eagle - if args.speculator_type == "mtp": if args.draft_attn_impl != "simple_flex_attention": raise ValueError( @@ -1028,25 +1016,24 @@ def parse_args(): ) parser.add_argument( "--norm-before-fc", - action=argparse.BooleanOptionalAction, - default=None, - help="Apply a single RMSNorm to concatenated target hidden states before " - "the FC projection layer. See --fc-norm for per-layer alternative.", + action="store_true", + default=False, + help="Use RMSNorm before FC layer in draft path " + "(e.g., for Eagle 3.1 / gpt-oss models).", ) parser.add_argument( "--fc-norm", - action=argparse.BooleanOptionalAction, - default=None, + action="store_true", + default=False, help="Apply per-layer RMSNorm to each auxiliary hidden state before " - "concatenation and FC projection. Defaults to True for llama draft arch.", + "concatenation and FC projection (Eagle 3.1 paper approach).", ) parser.add_argument( "--norm-output", - action=argparse.BooleanOptionalAction, - default=None, + action="store_true", + default=False, help="Feed post-norm hidden states back across TTT steps to stabilize " - "magnitude drift across speculation depths (Eagle 3.1). " - "Defaults to True for llama draft arch.", + "magnitude drift across speculation depths (Eagle 3.1).", ) # D-Flash specific parameters parser.add_argument( From fbedd79de9508451b891c2d32de15709216b5509 Mon Sep 17 00:00:00 2001 From: Orestis Zambounis Date: Mon, 29 Jun 2026 10:28:47 +0000 Subject: [PATCH 4/9] style: fix ruff formatting in peagle/core.py Co-Authored-By: Claude Opus 4.6 Signed-off-by: Orestis Zambounis --- src/speculators/models/peagle/core.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/src/speculators/models/peagle/core.py b/src/speculators/models/peagle/core.py index 188b6746c..149ffe6f8 100644 --- a/src/speculators/models/peagle/core.py +++ b/src/speculators/models/peagle/core.py @@ -120,10 +120,7 @@ def forward( if self.fc_norm is not None: chunks = sampled_hidden.chunk(len(self.fc_norm), dim=-1) sampled_hidden = torch.cat( - [ - norm(chunk) - for norm, chunk in zip(self.fc_norm, chunks, strict=True) - ], + [norm(chunk) for norm, chunk in zip(self.fc_norm, chunks, strict=True)], dim=-1, ) sampled_hidden = self.fc(sampled_hidden) # [1, total_sampled, hidden_size] From c6d960c0efcae65f5c75710b78bfbeaca44d00a1 Mon Sep 17 00:00:00 2001 From: Orestis Zambounis Date: Mon, 29 Jun 2026 10:34:00 +0000 Subject: [PATCH 5/9] fix: reject enabling both norm_before_fc and fc_norm MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The two flags apply conflicting normalization strategies to the FC input — using both causes double-norming. Add a Pydantic model validator to raise early with a clear error message. Signed-off-by: Orestis Zambounis Co-Authored-By: Claude Opus 4.6 Signed-off-by: Orestis Zambounis --- src/speculators/models/eagle3/config.py | 11 ++++++++++- tests/unit/test_config.py | 11 +++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/src/speculators/models/eagle3/config.py b/src/speculators/models/eagle3/config.py index 17666f409..4cb8fa263 100644 --- a/src/speculators/models/eagle3/config.py +++ b/src/speculators/models/eagle3/config.py @@ -1,6 +1,6 @@ from typing import Any, Literal -from pydantic import Field, field_serializer, field_validator +from pydantic import Field, field_serializer, field_validator, model_validator from transformers import AutoConfig, PretrainedConfig from transformers.models.qwen3.configuration_qwen3 import Qwen3Config @@ -86,6 +86,15 @@ class Eagle3SpeculatorConfig(SpeculatorModelConfig): description="Whether embedding layer weights require gradients during training", ) + @model_validator(mode="after") + def _check_norm_flags(self) -> "Eagle3SpeculatorConfig": + if self.norm_before_fc and self.fc_norm: + raise ValueError( + "norm_before_fc and fc_norm are mutually exclusive — " + "enable one or the other, not both." + ) + return self + @field_serializer("transformer_layer_config") def serialize_transformer_config(self, value: PretrainedConfig) -> dict: """Serialize transformer config to dict.""" diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index 8019dbea0..7fe250780 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -541,6 +541,17 @@ def test_eagle3_config_norm_output_defaults(): assert config.fc_norm is False +@pytest.mark.sanity +def test_eagle3_config_rejects_both_norm_before_fc_and_fc_norm(): + with pytest.raises(ValidationError, match="mutually exclusive"): + Eagle3SpeculatorConfig( + transformer_layer_config=copy.deepcopy(TINY_LLAMA_CONFIG), + norm_before_fc=True, + fc_norm=True, + speculators_config=_make_eagle3_speculators_config(), + ) + + @pytest.mark.sanity def test_eagle3_config_fc_norm_roundtrip(): original = Eagle3SpeculatorConfig( From 8989c059ccefe704c71ef8126d95e4627fae394b Mon Sep 17 00:00:00 2001 From: Orestis Zambounis Date: Mon, 29 Jun 2026 10:44:50 +0000 Subject: [PATCH 6/9] fix: add type annotation for optional fc_norm ModuleList Fixes mypy error: assigning None to a variable inferred as ModuleList. Signed-off-by: Orestis Zambounis Co-Authored-By: Claude Opus 4.6 Signed-off-by: Orestis Zambounis --- src/speculators/models/eagle3/core.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/speculators/models/eagle3/core.py b/src/speculators/models/eagle3/core.py index 36fcab639..e21a962fb 100644 --- a/src/speculators/models/eagle3/core.py +++ b/src/speculators/models/eagle3/core.py @@ -102,7 +102,7 @@ def __init__(self, config: Eagle3SpeculatorConfig): self.input_norm = None if config.fc_norm: - self.fc_norm = torch.nn.ModuleList( + self.fc_norm: torch.nn.ModuleList | None = torch.nn.ModuleList( [ self._model_definitions.norm_class( self.hidden_size, From 2ec6f4333ec73b0b95344818ca768e963fca6d7a Mon Sep 17 00:00:00 2001 From: Orestis Zambounis Date: Mon, 29 Jun 2026 10:49:29 +0000 Subject: [PATCH 7/9] =?UTF-8?q?docs:=20fix=20norm-before-fc=20help=20text?= =?UTF-8?q?=20=E2=80=94=20gpt-oss=20style,=20not=20Eagle=203.1=20paper?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Orestis Zambounis Co-Authored-By: Claude Opus 4.6 Signed-off-by: Orestis Zambounis --- docs/cli/train.md | 2 +- scripts/train.py | 5 +++-- src/speculators/models/eagle3/config.py | 4 ++-- 3 files changed, 6 insertions(+), 5 deletions(-) diff --git a/docs/cli/train.md b/docs/cli/train.md index ba3764ac2..569f59a8a 100644 --- a/docs/cli/train.md +++ b/docs/cli/train.md @@ -136,7 +136,7 @@ torchrun --standalone --nproc_per_node=4 scripts/train.py \ - **`--embed-requires-grad` / `--no-embed-requires-grad`** (flag, default: `False`) Whether to train embedding layer weights. -- **`--norm-before-fc`** (flag, default: `False`) Use RMSNorm before FC layer in draft path (e.g., for Eagle 3.1 / gpt-oss models). +- **`--norm-before-fc`** (flag, default: `False`) Apply a single RMSNorm to the concatenated auxiliary hidden states before the FC projection (gpt-oss style). See `--fc-norm` for the per-layer alternative from the Eagle 3.1 paper. - **`--fc-norm`** (flag, default: `False`) Apply per-layer RMSNorm to each auxiliary hidden state before concatenation and FC projection (Eagle 3.1 paper approach). diff --git a/scripts/train.py b/scripts/train.py index d44deca5e..b3a4a6a40 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -1018,8 +1018,9 @@ def parse_args(): "--norm-before-fc", action="store_true", default=False, - help="Use RMSNorm before FC layer in draft path " - "(e.g., for Eagle 3.1 / gpt-oss models).", + help="Apply a single RMSNorm to the concatenated auxiliary hidden states " + "before the FC projection (gpt-oss style). See --fc-norm for the " + "per-layer alternative from the Eagle 3.1 paper.", ) parser.add_argument( "--fc-norm", diff --git a/src/speculators/models/eagle3/config.py b/src/speculators/models/eagle3/config.py index 4cb8fa263..5c657b3f2 100644 --- a/src/speculators/models/eagle3/config.py +++ b/src/speculators/models/eagle3/config.py @@ -58,8 +58,8 @@ class Eagle3SpeculatorConfig(SpeculatorModelConfig): norm_before_fc: bool = Field( default=False, description=( - "Use RMSNorm before FC layer in draft path " - "(e.g., for Eagle 3.1 / gpt-oss models)." + "Apply a single RMSNorm to the concatenated auxiliary hidden states " + "before the FC projection (gpt-oss style)." ), ) From 1b3d55680a9b62c139bb0b74ddcd9fbe170c610f Mon Sep 17 00:00:00 2001 From: Orestis Zambounis Date: Thu, 2 Jul 2026 11:35:23 +0000 Subject: [PATCH 8/9] refactor: derive aux layer count from config instead of hardcoding 3 Replace hardcoded `3` with `len(eagle_aux_hidden_state_layer_ids)` for the FC layer input dim, input_norm dim, fc_norm ModuleList count, and peagle mask_hidden param so users can experiment with different numbers of auxiliary hidden state layers. Co-Authored-By: Claude Opus 4.6 Signed-off-by: Orestis Zambounis --- src/speculators/models/eagle3/core.py | 15 +++++++++++---- src/speculators/models/peagle/core.py | 9 ++++++++- 2 files changed, 19 insertions(+), 5 deletions(-) diff --git a/src/speculators/models/eagle3/core.py b/src/speculators/models/eagle3/core.py index 378208f2a..ac3f84e93 100644 --- a/src/speculators/models/eagle3/core.py +++ b/src/speculators/models/eagle3/core.py @@ -61,7 +61,14 @@ def __init__(self, config: Eagle3SpeculatorConfig): self.embed_tokens.weight.requires_grad = self.config.embed_requires_grad # FC LAYER - self.fc = torch.nn.Linear(3 * self.hidden_size, self.hidden_size, bias=False) + num_aux = ( + len(config.eagle_aux_hidden_state_layer_ids) + if config.eagle_aux_hidden_state_layer_ids + else 3 + ) + self.fc = torch.nn.Linear( + num_aux * self.hidden_size, self.hidden_size, bias=False + ) # DECODER LAYERS num_layers = tl_config.num_hidden_layers @@ -95,7 +102,7 @@ def __init__(self, config: Eagle3SpeculatorConfig): if config.norm_before_fc: self.input_norm = self._model_definitions.norm_class( - 3 * self.hidden_size, + num_aux * self.hidden_size, eps=config.transformer_layer_config.rms_norm_eps, ) else: @@ -108,7 +115,7 @@ def __init__(self, config: Eagle3SpeculatorConfig): self.hidden_size, eps=config.transformer_layer_config.rms_norm_eps, ) - for _ in range(3) + for _ in range(num_aux) ] ) else: @@ -152,7 +159,7 @@ def load_verifier_weights(self): @conditional_torch_compile def forward( # noqa: C901 self, - hidden_states: torch.Tensor, # shape: [1, total_seq_len, 3 * hidden_size] + hidden_states: torch.Tensor, # shape: [1, total_seq_len, num_aux * hidden_size] input_ids: torch.Tensor, # shape: [1, total_seq_len] document_ids: torch.Tensor, # shape: [1, total_seq_len] loss_mask: torch.Tensor | None = None, # shape: [1, total_seq_len] diff --git a/src/speculators/models/peagle/core.py b/src/speculators/models/peagle/core.py index 2e3c5b857..42a144c9b 100644 --- a/src/speculators/models/peagle/core.py +++ b/src/speculators/models/peagle/core.py @@ -44,7 +44,14 @@ def __init__( self.mask_token_id = config.mask_token_id # Learnable mask_hidden parameter for padding unsampled positions - self.mask_hidden = torch.nn.Parameter(torch.randn(1, 1, 3 * self.hidden_size)) + num_aux = ( + len(self.config.eagle_aux_hidden_state_layer_ids) + if self.config.eagle_aux_hidden_state_layer_ids + else 3 + ) + self.mask_hidden = torch.nn.Parameter( + torch.randn(1, 1, num_aux * self.hidden_size) + ) @conditional_torch_compile def forward( From 0d69b5f5ac2210fd4e704b9c7e80c583e86aba99 Mon Sep 17 00:00:00 2001 From: Orestis Zambounis Date: Thu, 2 Jul 2026 11:37:52 +0000 Subject: [PATCH 9/9] style: declare fc_norm annotation before conditional Co-Authored-By: Claude Opus 4.6 Signed-off-by: Orestis Zambounis --- src/speculators/models/eagle3/core.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/speculators/models/eagle3/core.py b/src/speculators/models/eagle3/core.py index ac3f84e93..5ab870965 100644 --- a/src/speculators/models/eagle3/core.py +++ b/src/speculators/models/eagle3/core.py @@ -108,8 +108,9 @@ def __init__(self, config: Eagle3SpeculatorConfig): else: self.input_norm = None + self.fc_norm: torch.nn.ModuleList | None = None if config.fc_norm: - self.fc_norm: torch.nn.ModuleList | None = torch.nn.ModuleList( + self.fc_norm = torch.nn.ModuleList( [ self._model_definitions.norm_class( self.hidden_size, @@ -118,8 +119,6 @@ def __init__(self, config: Eagle3SpeculatorConfig): for _ in range(num_aux) ] ) - else: - self.fc_norm = None self.post_init()