feat(eagle3): add per-layer FC normalization (--fc-norm)

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` / `--no-norm-before-fc`** (flag, default: `True` for eagle3, `False` otherwise) Apply RMSNorm before the FC layer in the draft path.
+- **`--norm-before-fc` / `--no-norm-before-fc`** (flag, default: `True` for eagle3, `False` otherwise) 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).
 
 - **`--norm-output` / `--no-norm-output`** (flag, default: `True` for eagle3, `False` otherwise) Feed post-norm hidden states back across TTT steps to stabilize magnitude drift across speculation depths.
 

scripts/train.py

@@ -916,10 +916,19 @@ def parse_args():
         "--norm-before-fc",
         action=argparse.BooleanOptionalAction,
         default=None,
-        help="Apply RMSNorm before the FC layer in the draft path "
+        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. "
         "(default: True for eagle3, False otherwise). "
         "Disable with --no-norm-before-fc.",
     )
+    parser.add_argument(
+        "--fc-norm",
+        action="store_true",
+        default=False,
+        help="Apply per-layer RMSNorm to each auxiliary hidden state before "
+        "concatenation and FC projection (Eagle 3.1 paper approach).",
+    )
     parser.add_argument(
         "--norm-output",
         action=argparse.BooleanOptionalAction,

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,

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
 
@@ -58,8 +58,18 @@ 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)."
+        ),
+    )
+
+    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."
         ),
     )
 
@@ -76,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."""

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,12 +102,24 @@ 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:
             self.input_norm = None
 
+        self.fc_norm: torch.nn.ModuleList | None = None
+        if config.fc_norm:
+            self.fc_norm = torch.nn.ModuleList(
+                [
+                    self._model_definitions.norm_class(
+                        self.hidden_size,
+                        eps=config.transformer_layer_config.rms_norm_eps,
+                    )
+                    for _ in range(num_aux)
+                ]
+            )
+
         self.post_init()
 
     @property
@@ -139,7 +158,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]
@@ -178,6 +197,12 @@ def forward(  # noqa: C901
 
         if self.input_norm is not None:
             hidden_states = self.input_norm(hidden_states)
+        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_norm, chunks, strict=True)],
+                dim=-1,
+            )
         hidden_states = self.fc(hidden_states)
         # shape: [1, total_seq_len, hidden_size]
 
@@ -329,6 +354,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,

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(
@@ -117,6 +124,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_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)],
+                dim=-1,
+            )
         sampled_hidden = self.fc(sampled_hidden)  # [1, total_sampled, hidden_size]
 
         layer_input = torch.cat(
@@ -216,6 +229,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),

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,

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_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)
+        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_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)
+
+        assert loss.isfinite()
+        loss.backward()
+
     def test_peagle_norm_before_fc(self):
         model = make_peagle_model()
         assert model.input_norm is None

tests/unit/test_config.py

@@ -538,3 +538,31 @@ 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_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(
+        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