diff --git a/docs/cli/train.md b/docs/cli/train.md index b4d4f718c..fc9b93a18 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` / `--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. diff --git a/scripts/train.py b/scripts/train.py index ba83a3332..a9558c95a 100644 --- a/scripts/train.py +++ b/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, 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..5c657b3f2 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 @@ -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.""" diff --git a/src/speculators/models/eagle3/core.py b/src/speculators/models/eagle3/core.py index 14db3ed2d..5ab870965 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,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, diff --git a/src/speculators/models/peagle/core.py b/src/speculators/models/peagle/core.py index a9ae0d2af..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( @@ -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), 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..876221617 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_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 diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index 3506da76b..7fe250780 100644 --- a/tests/unit/test_config.py +++ b/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