diff --git a/src/speculators/models/dflash2/config.py b/src/speculators/models/dflash2/config.py index 8a93c5c3b..842212860 100644 --- a/src/speculators/models/dflash2/config.py +++ b/src/speculators/models/dflash2/config.py @@ -1,6 +1,6 @@ from typing import Literal -from pydantic import Field +from pydantic import Field, model_validator from speculators import SpeculatorModelConfig from speculators.models.dflash.config import DFlashSpeculatorConfig @@ -43,3 +43,51 @@ class DFlash2SpeculatorConfig(DFlashSpeculatorConfig): ge=1, description="Number of unary candidates reranked during inference.", ) + draft_ffn_type: Literal["dense", "moe"] = Field( + default="dense", + description="Feed-forward implementation used by every draft layer.", + ) + num_experts: int = Field( + default=256, + ge=1, + description="Number of routed experts when draft_ffn_type='moe'.", + ) + num_experts_per_tok: int = Field( + default=8, + ge=1, + description="Number of routed experts selected per token.", + ) + moe_intermediate_size: int = Field( + default=512, + ge=1, + description="Intermediate width of each routed expert.", + ) + shared_expert_intermediate_size: int = Field( + default=512, + ge=1, + description="Intermediate width of the always-on shared expert.", + ) + moe_experts_implementation: Literal[ + "grouped_mm", "batched_mm", "deepgemm", "sonicmoe", "reference" + ] = Field( + default="grouped_mm", + description=( + "Kernel used for the routed-expert GEMMs. This is a closed set on " + "purpose: Transformers silently falls back to the per-expert Python " + "loop when it does not recognize the name, and that loop costs about " + "68x a grouped GEMM at 256 experts / hidden 2048 / intermediate 512 " + "(measured on H200, ~505ms vs ~7.4ms per layer forward+backward). " + "'reference' names that loop explicitly for debugging; 'batched_mm' " + "materializes a dense per-expert activation and will exhaust device " + "memory at production expert counts." + ), + ) + + @model_validator(mode="after") + def validate_moe_routing(self) -> "DFlash2SpeculatorConfig": + if self.draft_ffn_type == "moe" and self.num_experts_per_tok > self.num_experts: + raise ValueError( + "num_experts_per_tok cannot exceed num_experts: " + f"{self.num_experts_per_tok} > {self.num_experts}." + ) + return self diff --git a/src/speculators/models/dflash2/core.py b/src/speculators/models/dflash2/core.py index 0f6360f4e..cf83ce272 100644 --- a/src/speculators/models/dflash2/core.py +++ b/src/speculators/models/dflash2/core.py @@ -31,7 +31,7 @@ class DFlash2DraftModel(DFlashDraftModel): """DFlash with local convolution and bilinear candidate reranking.""" config_class: ClassVar[type[DFlash2SpeculatorConfig]] = DFlash2SpeculatorConfig # type: ignore[misc,assignment] - _no_split_modules = ["Qwen3DFlash2DecoderLayer"] + _no_split_modules = ["Qwen3DFlash2DecoderLayer", "Qwen35DFlashMoeBlock"] def __init__(self, config: DFlash2SpeculatorConfig) -> None: target_vocab_size = config.transformer_layer_config.vocab_size @@ -48,13 +48,14 @@ def __init__(self, config: DFlash2SpeculatorConfig) -> None: ) super().__init__(config=config) + initializer_range = getattr( + config.transformer_layer_config, "initializer_range", 0.02 + ) for layer_ in self.layers: assert isinstance(layer_, Qwen3DFlash2DecoderLayer) # noqa: S101 layer_.reset_convolutions() + layer_.reset_moe_parameters(initializer_range) - initializer_range = getattr( - config.transformer_layer_config, "initializer_range", 0.02 - ) self.candidate_selector = CandidateSelector( vocab_size=target_vocab_size, hidden_size=config.transformer_layer_config.hidden_size, @@ -73,6 +74,12 @@ def _make_decoder_layer( block_size=config.block_size, conv_kernel_size=config.conv_kernel_size, conv_group_size=config.conv_group_size, + draft_ffn_type=config.draft_ffn_type, + num_experts=config.num_experts, + num_experts_per_tok=config.num_experts_per_tok, + moe_intermediate_size=config.moe_intermediate_size, + shared_expert_intermediate_size=config.shared_expert_intermediate_size, + moe_experts_implementation=config.moe_experts_implementation, ) @classmethod @@ -90,6 +97,16 @@ def from_training_args( conv_group_size=kwargs.get("conv_group_size", 16), selector_rank=kwargs.get("selector_rank", 256), selector_top_k=kwargs.get("selector_top_k", 16), + draft_ffn_type=kwargs.get("draft_ffn_type", "dense"), + num_experts=kwargs.get("num_experts", 256), + num_experts_per_tok=kwargs.get("num_experts_per_tok", 8), + moe_intermediate_size=kwargs.get("moe_intermediate_size", 512), + shared_expert_intermediate_size=kwargs.get( + "shared_expert_intermediate_size", 512 + ), + moe_experts_implementation=kwargs.get( + "moe_experts_implementation", "grouped_mm" + ), ) model = cls(config=config) model.load_vocab_mappings(t2d, d2t) diff --git a/src/speculators/models/dflash2/model_definitions.py b/src/speculators/models/dflash2/model_definitions.py index 62dab4b79..94d8d4b5f 100644 --- a/src/speculators/models/dflash2/model_definitions.py +++ b/src/speculators/models/dflash2/model_definitions.py @@ -25,6 +25,8 @@ https://github.com/z-lab/dflash/blob/07ebd93db9f472af339b644bb70221ad8428328a/dflash/model.py """ +from copy import copy + import torch from torch import nn from transformers.cache_utils import Cache @@ -32,6 +34,11 @@ FlashAttentionKwargs, Qwen3Config, ) +from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import ( + Qwen3_5MoeExperts, + Qwen3_5MoeMLP, + Qwen3_5MoeTopKRouter, +) from typing_extensions import Unpack from speculators.models.dflash.model_definitions import Qwen3DFlashDecoderLayer @@ -40,10 +47,84 @@ "CandidateSelector", "GroupedDynamicCausalConv", "Qwen3DFlash2DecoderLayer", + "Qwen35DFlashMoeBlock", "grouped_dynamic_conv", ] +class Qwen35DFlashMoeBlock(nn.Module): + """Qwen3.5-compatible routed and shared experts for a DFlash2 layer. + + Only the feed-forward is replaced; the attention, the local convolutions + and the draft-block geometry are the dense DFlash2 ones. Weight names match + the Qwen3.5 MoE contract so a verifier layer warm-starts tensor for tensor. + """ + + def __init__( + self, + config: Qwen3Config, + *, + num_experts: int, + num_experts_per_tok: int, + moe_intermediate_size: int, + shared_expert_intermediate_size: int, + experts_implementation: str = "grouped_mm", + ) -> None: + super().__init__() + # Copy rather than annotate the shared attention config: it is + # serialized as transformer_layer_config, and MoE architecture fields + # have their one source of truth in DFlash2SpeculatorConfig. + moe_config = copy(config) + moe_config.num_experts = num_experts + moe_config.num_experts_per_tok = num_experts_per_tok + moe_config.moe_intermediate_size = moe_intermediate_size + moe_config._experts_implementation = ( # noqa: SLF001 + None if experts_implementation == "reference" else experts_implementation + ) + self.gate = Qwen3_5MoeTopKRouter(moe_config) + self.experts = Qwen3_5MoeExperts(moe_config) + self.shared_expert = Qwen3_5MoeMLP( + moe_config, intermediate_size=shared_expert_intermediate_size + ) + self.shared_expert_gate = nn.Linear(moe_config.hidden_size, 1, bias=False) + + def reset_parameters(self, initializer_range: float) -> None: + """Initialize the raw expert tensors and never leave the router at zero. + + ``Qwen3_5MoeTopKRouter`` allocates its weight with ``torch.zeros`` and + relies on ``Qwen3_5MoePreTrainedModel._init_weights`` to replace it. + A DFlash2 draft is not on that class's MRO, so without this every token + would see identical router logits at step 0 and collapse onto one fixed + expert subset. + """ + nn.init.normal_(self.gate.weight, mean=0.0, std=initializer_range) + nn.init.normal_(self.experts.gate_up_proj, mean=0.0, std=initializer_range) + nn.init.normal_(self.experts.down_proj, mean=0.0, std=initializer_range) + for projection in ( + self.shared_expert.gate_proj, + self.shared_expert.up_proj, + self.shared_expert.down_proj, + self.shared_expert_gate, + ): + nn.init.normal_(projection.weight, mean=0.0, std=initializer_range) + + def _route( + self, flat_hidden: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + router_logits, routing_weights, selected_experts = self.gate(flat_hidden) + routed = self.experts(flat_hidden, selected_experts, routing_weights) + shared_gate = torch.sigmoid(self.shared_expert_gate(flat_hidden)) + shared = shared_gate * self.shared_expert(flat_hidden) + return routed + shared, router_logits, selected_experts, shared_gate + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + original_shape = hidden_states.shape + output, _logits, _selected, _gate = self._route( + hidden_states.reshape(-1, original_shape[-1]) + ) + return output.reshape(original_shape) + + def grouped_dynamic_conv( hidden_states: torch.Tensor, delta_kernel: torch.Tensor, @@ -183,8 +264,24 @@ def __init__( block_size: int, conv_kernel_size: int, conv_group_size: int, + draft_ffn_type: str = "dense", + num_experts: int = 256, + num_experts_per_tok: int = 8, + moe_intermediate_size: int = 512, + shared_expert_intermediate_size: int = 512, + moe_experts_implementation: str = "grouped_mm", ) -> None: super().__init__(config=config, layer_idx=layer_idx) + self.block_size = block_size + if draft_ffn_type == "moe": + self.mlp = Qwen35DFlashMoeBlock( + config, + num_experts=num_experts, + num_experts_per_tok=num_experts_per_tok, + moe_intermediate_size=moe_intermediate_size, + shared_expert_intermediate_size=shared_expert_intermediate_size, + experts_implementation=moe_experts_implementation, + ) conv_kwargs = { "block_size": block_size, "kernel_size": conv_kernel_size, @@ -200,6 +297,11 @@ def reset_convolutions(self) -> None: self.attention_conv.reset_parameters() self.mlp_conv.reset_parameters() + def reset_moe_parameters(self, initializer_range: float) -> None: + """Initialize MoE tensors; a no-op for a dense draft layer.""" + if isinstance(self.mlp, Qwen35DFlashMoeBlock): + self.mlp.reset_parameters(initializer_range) + def forward( self, target_hidden: torch.Tensor | None = None, diff --git a/src/speculators/train/config/schema.py b/src/speculators/train/config/schema.py index e0d51c2ae..8d626665d 100644 --- a/src/speculators/train/config/schema.py +++ b/src/speculators/train/config/schema.py @@ -484,6 +484,34 @@ class DFlash2Args(_Group): ge=0.0, description="DFlash2: weight of the candidate-selector K-way CE term.", ) + draft_ffn_type: Literal["dense", "moe"] = Field( + default="dense", + description="DFlash2: dense MLP or routed MoE plus a shared expert. " + "Only the feed-forward changes; attention, convolutions and the draft " + "block layout stay the same.", + ) + num_experts: int = Field( + default=256, ge=1, description="DFlash2 MoE: number of routed experts." + ) + num_experts_per_tok: int = Field( + default=8, ge=1, description="DFlash2 MoE: routed experts selected per token." + ) + moe_intermediate_size: int = Field( + default=512, ge=1, description="DFlash2 MoE: per-routed-expert width." + ) + shared_expert_intermediate_size: int = Field( + default=512, ge=1, description="DFlash2 MoE: shared-expert width." + ) + moe_experts_implementation: Literal[ + "grouped_mm", "batched_mm", "deepgemm", "sonicmoe", "reference" + ] = Field( + default="grouped_mm", + description="DFlash2 MoE: routed-expert GEMM kernel. Restricted on " + "purpose -- an unrecognized name makes Transformers fall back to a " + "per-expert Python loop that costs ~68x a grouped GEMM at 256 experts. " + "'reference' selects that loop deliberately; 'batched_mm' exhausts " + "device memory at production expert counts.", + ) class DSparkArgs(_Group): @@ -734,6 +762,20 @@ def _validate_dpace(self) -> "TrainConfig": ) return self + @model_validator(mode="after") + def _validate_moe(self) -> "TrainConfig": + """Reject MoE geometry that would otherwise fail deep in a run.""" + if self.dflash2.draft_ffn_type != "moe": + return self + if self.speculator_type != "dflash2": + raise ValueError("--draft-ffn-type=moe is currently supported by dflash2") + if self.dflash2.num_experts_per_tok > self.dflash2.num_experts: + raise ValueError( + "--num-experts-per-tok cannot exceed --num-experts: " + f"{self.dflash2.num_experts_per_tok} > {self.dflash2.num_experts}" + ) + return self + def flatten(self) -> dict[str, Any]: """The flat ``vars(args)``-shaped dict the ``SpeculatorModel`` classes consume via ``**kwargs``. diff --git a/tests/unit/models/test_dflash2_model_definitions.py b/tests/unit/models/test_dflash2_model_definitions.py index 791b26426..2fbfb5d3f 100644 --- a/tests/unit/models/test_dflash2_model_definitions.py +++ b/tests/unit/models/test_dflash2_model_definitions.py @@ -19,6 +19,7 @@ from speculators.models.dflash2.model_definitions import ( CandidateSelector, GroupedDynamicCausalConv, + Qwen35DFlashMoeBlock, grouped_dynamic_conv, ) from speculators.models.dspark.metrics import compute_metrics as compute_unary_metrics @@ -304,6 +305,103 @@ def test_config_round_trip_preserves_dflash2_contract(tmp_path): assert loaded.selector_top_k == config.selector_top_k +def test_config_round_trip_preserves_moe_contract(tmp_path): + config = _tiny_config( + speculators_config=SpeculatorsConfig( + algorithm="dflash2", + proposal_methods=[GreedyTokenProposalConfig(speculative_tokens=3)], + default_proposal_method="greedy", + verifier=VerifierConfig( + name_or_path="Qwen/Qwen3-4B", + architectures=["Qwen3ForCausalLM"], + ), + ), + draft_ffn_type="moe", + num_experts=16, + num_experts_per_tok=2, + moe_intermediate_size=12, + shared_expert_intermediate_size=10, + moe_experts_implementation="grouped_mm", + ) + + config.save_pretrained(tmp_path) + loaded = SpeculatorModelConfig.from_pretrained(tmp_path) + + assert isinstance(loaded, DFlash2SpeculatorConfig) + assert loaded.draft_ffn_type == "moe" + assert loaded.num_experts == 16 + assert loaded.num_experts_per_tok == 2 + assert loaded.moe_intermediate_size == 12 + assert loaded.shared_expert_intermediate_size == 10 + assert loaded.moe_experts_implementation == "grouped_mm" + + +def test_moe_ffn_uses_fused_checkpoint_layout_without_mutating_attention_config(): + config = _tiny_config( + draft_ffn_type="moe", + num_experts=16, + num_experts_per_tok=2, + moe_intermediate_size=12, + shared_expert_intermediate_size=10, + ) + transformer_config = config.transformer_layer_config + + model = DFlash2DraftModel(config) + mlp = model.layers[0].mlp + + assert isinstance(mlp, Qwen35DFlashMoeBlock) + assert mlp.experts.config._experts_implementation == "grouped_mm" + assert not hasattr(transformer_config, "num_experts") + assert not hasattr(transformer_config, "num_experts_per_tok") + assert not hasattr(transformer_config, "moe_intermediate_size") + assert torch.count_nonzero(mlp.gate.weight) + + state_keys = set(model.state_dict()) + assert "layers.0.mlp.gate.weight" in state_keys + assert "layers.0.mlp.experts.gate_up_proj" in state_keys + assert "layers.0.mlp.experts.down_proj" in state_keys + + +def test_dense_ffn_remains_the_default(): + model = DFlash2DraftModel(_tiny_config()) + + assert not isinstance(model.layers[0].mlp, Qwen35DFlashMoeBlock) + assert not any(".mlp.experts." in name for name, _ in model.named_parameters()) + + +def test_moe_ffn_reference_forward_backward(): + torch.manual_seed(0) + config = _tiny_config().transformer_layer_config + mlp = Qwen35DFlashMoeBlock( + config, + num_experts=4, + num_experts_per_tok=2, + moe_intermediate_size=8, + shared_expert_intermediate_size=8, + experts_implementation="reference", + ) + mlp.reset_parameters(config.initializer_range) + hidden_states = torch.randn(2, 3, config.hidden_size, requires_grad=True) + + output = mlp(hidden_states) + + assert output.shape == hidden_states.shape + assert torch.isfinite(output).all() + output.square().mean().backward() + for name, parameter in mlp.named_parameters(): + assert parameter.grad is not None, f"missing gradient for {name}" + assert torch.isfinite(parameter.grad).all(), f"non-finite gradient for {name}" + + +def test_moe_config_rejects_more_selected_than_routed_experts(): + with pytest.raises(ValueError, match="cannot exceed"): + _tiny_config( + draft_ffn_type="moe", + num_experts=4, + num_experts_per_tok=8, + ) + + def test_model_rejects_pruned_draft_vocabulary(): with pytest.raises(ValueError, match="full verifier vocabulary"): DFlash2DraftModel(_tiny_config(draft_vocab_size=32)) diff --git a/tests/unit/train/config/test_schema.py b/tests/unit/train/config/test_schema.py index fcccc8f7a..daa99d3c5 100644 --- a/tests/unit/train/config/test_schema.py +++ b/tests/unit/train/config/test_schema.py @@ -72,6 +72,8 @@ def test_flatten_resolves_dflash2_derived_defaults(): assert flat["selector_rank"] == 256 assert flat["selector_top_k"] == 16 assert flat["selector_loss_alpha"] == pytest.approx(1.0) + assert flat["draft_ffn_type"] == "dense" + assert flat["moe_experts_implementation"] == "grouped_mm" assert flat["sliding_window_non_causal"] is True @@ -126,6 +128,45 @@ def test_dflash2_explicit_values_override_defaults(): assert flat["selector_loss_alpha"] == pytest.approx(0.25) +def test_dflash2_moe_geometry_flattens_for_model_construction(): + cfg = TrainConfig( + speculator_type="dflash2", + dflash2=DFlash2Args( + draft_ffn_type="moe", + num_experts=256, + num_experts_per_tok=8, + moe_intermediate_size=512, + shared_expert_intermediate_size=512, + moe_experts_implementation="grouped_mm", + ), + ) + + flat = cfg.flatten() + assert flat["draft_ffn_type"] == "moe" + assert flat["num_experts"] == 256 + assert flat["num_experts_per_tok"] == 8 + assert flat["moe_intermediate_size"] == 512 + assert flat["shared_expert_intermediate_size"] == 512 + + +def test_dflash2_moe_geometry_fails_closed(): + with pytest.raises(ValueError, match="cannot exceed"): + TrainConfig( + speculator_type="dflash2", + dflash2=DFlash2Args( + draft_ffn_type="moe", + num_experts=4, + num_experts_per_tok=8, + ), + ) + + with pytest.raises(ValueError, match="supported by dflash2"): + TrainConfig( + speculator_type="dspark", + dflash2=DFlash2Args(draft_ffn_type="moe"), + ) + + def test_from_flat_inverts_flatten(): cfg = TrainConfig( speculator_type="dspark",