[sherlockwu]

Stacked on #798 (kan/pr2-dsv4-dspark); retarget as the stack merges. First of two: this PR lays out the MoE FFN package (modules, contracts, config, and every trainer/checkpoint seam) with no routing math; the follow-up adds the DeepSeek-V4 implementation and the DSV4-Flash ablation arm. Together they supersede #799.

Design

One configurable MoE layer with swappable components, not one block per target family (the Megatron-Core split; see specforge/modeling/draft/moe/DESIGN.md). Family differences are a handful of small functions (score function, balancing policy, renorm/scale, shared-expert gate, activation clamp); the expensive parts (dispatch, checkpoint naming, FSDP layout, deferred balance-update timing) are shared.

specforge/modeling/draft/moe/
  config.py      MoEConfig + moe_preset registry; resolve_moe_config(draft_config)
  router.py      Router contract (x -> RoutingResult) + score-function registry
  balance.py     BalanceController contract (selection shift / observe / deferred apply / aux loss / metrics)
  experts.py     RoutedExperts contract (weights + dispatch)
  shared.py      SharedExpert contract
  layer.py       MoELayer = gate + experts + shared_experts; build_ffn() is the dense/MoE switch
  hooks.py       apply_pending_balance_updates / collect_moe_aux_loss / collect_moe_metrics
  state_dict.py  to/from_checkpoint_state_dict: module layout <-> official file naming
  init.py        WarmStartPlan / select_target_experts

Config contract. Architecture keys use the target checkpoints' native HF names at the draft-JSON top level (n_routed_experts, num_experts_per_tok, moe_intermediate_size, scoring_func, routed_scaling_factor, ...) so a draft JSON can copy them from the target's config.json; moe_preset supplies a family's defaults and any explicit key overrides it for ablations. Training-only knobs live under dflash_config as moe_* and never touch the checkpoint.

Seams wired here (so the implementation PR is components only):

• Qwen3DFlashDecoderLayer builds its FFN through build_ffn(config, dense=kernels.make_mlp): dense drafts get the kernel provider's MLP verbatim (DFlash, DFlash2 and DSpark all inherit).
• DFlashDraftModel._init_weights reaches MoE bare Parameters; the model forward applies pending balance updates in training, outside activation-checkpoint regions.
• DFlash/DSpark strategies add moe/* load metrics to StepOutput.metrics (DP-averaged and logged like any scalar).
• Checkpoint-naming boundary applied in FSDPTrainingBackend save/load, warm start, materialize_draft, and both exporters. FSDP's full-state-dict hooks index by the module's own FQNs, so the rename cannot live inside state_dict().

Tests

tests/test_modeling/test_moe.py pins the contracts with stub components: config resolution/validation/presets, registries, MoELayer composition (shape, shared-expert add, train-only observe, hook delegation, selection bias affects choice not weights), hooks over module trees, the state-dict boundary, warm-start plans, and DFlash wiring (dense unchanged; MoE layers initialized; balance updates applied per training forward). Existing modeling/training/export/loading suites pass unchanged.

Not in this PR

Any preset, score function, controller, experts backend or shared expert. Aux-loss consumption in an objective (collected, not yet wired; the first preset with an aux-loss policy adds it).

[yushengsu-thu]

(cc catch - human verify)
Resume path bypasses the checkpoint-naming boundary: Trainer.__init__ loads state["draft_state_dict"] (file naming) straight into model.draft_model (L347-L349) and raises on unexpected keys. Once an experts backend registers a converter, every resume_from of an MoE drafter fails. Fix: convert like _load_module_state_dict does.

diff --git a/specforge/training/trainer.py b/specforge/training/trainer.py
index d465b947..594ca5b2 100644
--- a/specforge/training/trainer.py
+++ b/specforge/training/trainer.py
@@ -344,7 +344,12 @@ class Trainer:
                         f"{key}={comparison_current}; "
                         f"resume with the original configuration"
                     )
-            saved_weights = state["draft_state_dict"]
+            # Checkpoint FILES use the official parameter naming; the module may
+            # use a native MoE layout (stacked experts). Convert at this boundary
+            # exactly like FSDPTrainingBackend._load_module_state_dict does.
+            from specforge.modeling.draft.moe import from_checkpoint_state_dict
+
+            saved_weights = from_checkpoint_state_dict(state["draft_state_dict"])
             load_result = model.draft_model.load_state_dict(saved_weights, strict=False)
             loaded = len(saved_weights) - len(load_result.unexpected_keys)
             missing_keys = set(load_result.missing_keys)

[yushengsu-thu]

(cc catch - human verify)
(I know we seldom use DDP training but it's better to modify)
_moe_metrics silently returns {} on the NO_SHARD (DDP) recipe: strategies get the wrapped module (trainer.py L430-L437), and DDP, unlike FSDP, does not forward draft_model (L55-L62). Same root cause already disables the DFlash2 selector loss under DDP (_selector_loss_alpha). Fix: unwrap .module like Eagle3TrainStrategy does.

diff --git a/specforge/training/strategies/base.py b/specforge/training/strategies/base.py
index cbc253da..dcff526c 100644
--- a/specforge/training/strategies/base.py
+++ b/specforge/training/strategies/base.py
@@ -52,11 +52,24 @@ class StepContext:
     total_steps: Optional[int] = None
 
 
+def _unwrap_parallel(model: nn.Module) -> nn.Module:
+    """Reach the composite model behind a parallel wrapper.
+
+    Strategies receive ``backend.prepare_model(...)``'s return value. FSDP
+    forwards attribute lookups to the wrapped module; DDP does not, so
+    ``getattr(ddp, "draft_model")`` / ``getattr(ddp, "selector_loss_alpha")``
+    would silently miss. Same unwrap rule as ``Eagle3TrainStrategy``.
+    """
+    while not hasattr(model, "draft_model") and hasattr(model, "module"):
+        model = model.module
+    return model
+
+
 def _moe_metrics(model_wrapper: nn.Module) -> Dict[str, Any]:
     """``moe/...`` load diagnostics of the wrapped draft; ``{}`` for dense drafts."""
     from specforge.modeling.draft.moe import collect_moe_metrics
 
-    draft_model = getattr(model_wrapper, "draft_model", None)
+    draft_model = getattr(_unwrap_parallel(model_wrapper), "draft_model", None)
     if draft_model is None:
         return {}
     return collect_moe_metrics(draft_model)
@@ -478,20 +491,20 @@ class DFlashTrainStrategy(DraftTrainStrategy):
         return next(self.dflash_model.parameters()).device
 
     def _selector_loss_alpha(self, ctx: Optional[StepContext]) -> float:
-        target = float(getattr(self.dflash_model, "selector_loss_alpha", 0.0))
+        model = _unwrap_parallel(self.dflash_model)
+        target = float(getattr(model, "selector_loss_alpha", 0.0))
         if target <= 0 or ctx is None or not ctx.total_steps:
             return target
 
         total_steps = int(ctx.total_steps)
         warmup_steps = int(
-            total_steps
-            * float(getattr(self.dflash_model, "selector_warmup_ratio", 0.0))
+            total_steps * float(getattr(model, "selector_warmup_ratio", 0.0))
         )
         if ctx.global_step < warmup_steps:
             return 0.0
 
         ramp_steps = int(
-            total_steps * float(getattr(self.dflash_model, "selector_ramp_ratio", 0.0))
+            total_steps * float(getattr(model, "selector_ramp_ratio", 0.0))
         )
         if ramp_steps <= 0:
             return target

[yushengsu-thu]

This doc really help me quickly understand the design but I not sure if it's proper to leave it or put it here.

[yushengsu-thu]

Values like -1 and 0.5 silently select dense layers. Could we reject invalid expert counts instead of defaulting to dense?

resolve_moe_config({"n_routed_experts": -1})   # None
resolve_moe_config({"n_routed_experts": 0.5})  # None