(fsdp2 dev.) feat(moe): expert parallelism for MoE drafters on the FSDP2 backend (stacked on #915 + #931)

configs/deepseek-v4-flash-dspark-moe.json

@@ -0,0 +1,60 @@
+{
+  "architectures": ["DSparkDraftModel"],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {"AutoModel": "dspark.DSparkDraftModel"},
+  "block_size": 7,
+  "bos_token_id": 0,
+  "dflash_config": {
+    "attention_mode": "gqa",
+    "confidence_head_alpha": 1.0,
+    "confidence_head_with_markov": true,
+    "enable_confidence_head": true,
+    "markov_head_type": "vanilla",
+    "markov_rank": 256,
+    "mask_token_id": 128799,
+    "moe_bias_update_rate": 0.001,
+    "moe_dispatch": "grouped_mm",
+    "projector_type": "dspark",
+    "target_layer_ids": [1, 11, 21, 31, 41]
+  },
+  "dtype": "bfloat16",
+  "eos_token_id": 1,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 4096,
+  "initializer_range": 0.02,
+  "intermediate_size": 12288,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 1048576,
+  "max_window_layers": 5,
+  "model_type": "qwen3",
+  "moe_intermediate_size": 2048,
+  "moe_preset": "deepseek_v4",
+  "n_routed_experts": 64,
+  "n_shared_experts": 1,
+  "num_attention_heads": 32,
+  "num_experts_per_tok": 6,
+  "num_hidden_layers": 5,
+  "num_key_value_heads": 8,
+  "num_target_layers": 43,
+  "pad_token_id": 1,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "factor": 16.0,
+    "original_max_position_embeddings": 65536,
+    "rope_theta": 10000.0,
+    "rope_type": "yarn"
+  },
+  "sliding_window": null,
+  "tie_word_embeddings": false,
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 129280
+}

docs/recipes/deepseek-v4-flash-dspark-disaggregated.md

@@ -103,6 +103,41 @@ MI355X node over the full two epochs (1,885 optimizer steps): 6.32 s per
 128-sample step on average, 6.1-6.4 s at steady state (about 3.1 s waiting for
 capture and 3.0 s of trainer compute), 3 h 18 min end to end.
 
+## MoE-FFN arm (ablation)
+
+`examples/configs/online/disaggregated/external/deepseek-v4-flash-dspark-moe-disaggregated.yaml`
+is the same recipe with `configs/deepseek-v4-flash-dspark-moe.json`: the
+five-layer GQA decoder keeps its attention, and each layer's dense MLP becomes
+the target's MoE (`moe_preset: deepseek_v4`: sqrt-softplus scores,
+aux-loss-free top-k with the sign-controlled selection bias, combine weights
+renormalized and scaled by 1.5, one ungated shared expert, SwiGLU clamped at
+10). Sizes are per run: 64 routed experts, top-6, width 2048, so the activated
+FFN width (6 x 2048 + 2048 shared) matches the dense 12288 at ~10x the FFN
+parameters. Run it against the dense recipe with identical hparams; the two
+YAMLs differ only in the draft JSON and run names. The capture servers are
+shared by both arms unchanged.
+
+Training-only knobs live under the draft JSON's `dflash_config`:
+`moe_bias_update_rate` (0.001, the balancing controller's step) and
+`moe_dispatch` (`grouped_mm` runs the experts as grouped GEMMs with no host
+sync; `sorted_loop` is the portable fallback). The trainer logs `moe/*` load
+metrics (max/min load ratios, unused-expert fraction, balancing-bias
+magnitude) alongside the usual scalars. Checkpoints keep the official
+per-expert naming (`layers.N.mlp.experts.{i}.w{1,2,3}.weight`,
+`layers.N.mlp.gate.bias`, `layers.N.mlp.shared_experts.w{1,2,3}.weight`), so
+exports load into SGLang's DeepSeek-V4 MoE unchanged.
+
+To train the experts instead of freezing warm-started ones, use the
+expert-parallel variant
+`examples/configs/online/disaggregated/external/deepseek-v4-flash-dspark-moe-ep4-disaggregated.yaml`
+(`training.backend: fsdp2`, `training.expert_parallel_size: 4`). With the
+experts sharded as plain FSDP parameters, every micro-batch re-gathers all 64
+experts of every layer; under EP each of the four trainer ranks owns 16 experts
+per layer, the group exchanges its anchor tokens (fixed-size all-gather and
+reduce-scatter per MoE layer) and the expert weights never move. Checkpoints,
+`moe/*` metrics, warm start and exports are unchanged. See the expert
+parallelism notes in `docs/sections/advanced_features/customization.md`.
+
 ## Fresh attempts
 
 Delete the run's `outputs/` directory and, whenever a capture server was

docs/sections/advanced_features/customization.md

@@ -128,6 +128,70 @@ drafts currently implements the GQA/MHA layout only, so plan benchmarks
 accordingly. DFlash2 otherwise follows the same mode selection; its convolution
 and selector do not change the attention projection contract.
 
+## MoE FFN for DFlash-family drafts
+
+Any DFlash-family draft (DFlash, DFlash2, DSpark) swaps its dense MLP for a
+sparse MoE FFN when the draft JSON sets `n_routed_experts > 0`. The MoE is one
+configurable layer (`specforge/modeling/draft/moe/`, see its `DESIGN.md`):
+a `moe_preset` names a target family's routing recipe, and the architecture
+keys use the target checkpoints' native HF names so they can be copied from
+the target's `config.json`:
+
+```json
+{
+  "moe_preset": "deepseek_v4",
+  "n_routed_experts": 64,
+  "num_experts_per_tok": 6,
+  "moe_intermediate_size": 2048,
+  "n_shared_experts": 1,
+  "dflash_config": {"moe_bias_update_rate": 0.001, "moe_dispatch": "grouped_mm"}
+}
+```
+
+Top-level keys override the preset (for ablations: `scoring_func`,
+`norm_topk_prob`, `routed_scaling_factor`, `balance`, `shared_expert_gate`,
+`swiglu_limit`, ...). Training-only knobs live under `dflash_config` with an
+`moe_` prefix and never change the checkpoint. Checkpoints, warm starts and
+exports keep the official per-expert naming (`experts.{i}.w{1,2,3}.weight`),
+so an exported drafter loads into SGLang unchanged. Dense drafts are
+unaffected: with no `n_routed_experts` the kernel provider's MLP is used as-is.
+
+`deepseek_v4` is the checked-in preset (DeepSeek-V4 routing:
+`sqrtsoftplus` scores, aux-loss-free `noaux_tc` balancing, combine weights
+renormalized and scaled by 1.5, one ungated shared expert, SwiGLU clamp 10);
+`configs/deepseek-v4-flash-dspark-moe.json` uses it. A new target family is a
+preset registration plus whichever components it needs (score function,
+balance controller, experts backend, shared expert); each registers by name
+from its own module.
+
+### Expert parallelism
+
+Experts that are trained (not `moe_freeze_experts`) can be sliced across ranks
+on the FSDP2 backend:
+
+```yaml
+training:
+  backend: fsdp2
+  expert_parallel_size: 4
+```
+
+Each MoE layer's routed experts are split across `expert_parallel_size`
+consecutive ranks (the EP group). Inside the layer the group all-gathers its
+tokens, every rank routes them with the replicated router and computes only the
+experts it owns, and a reduce-scatter returns each rank its own tokens summed
+over all owners. The expert slices are `DTensor`s over the `ep` mesh axis and
+FSDP2 shards them again over the ranks that hold the same slice (`efsdp`);
+everything else stays on the full data-parallel mesh. The EP axis is carved out
+of data parallelism, so ranks keep distinct data and the dense part of the draft
+is not computed twice. Checkpoints still gather to the full `[E, ...]` tensors
+and keep the official naming; the balance controller, warm start and exports
+are unchanged. Requirements: a MoE draft JSON, `backend: fsdp2`, a sharded
+`fsdp_sharding`, `tp_size: 1`, no sequence parallelism, and `n_routed_experts`
+divisible by `expert_parallel_size`. Expert parallelism adds one host sync per
+MoE layer (this rank's slot bounds in the sorted routing) and two fixed-size
+collectives; it pays off when the gathered expert weights, not the tokens,
+dominate the step.
+
 ## Draft architectures
 
 Draft classes register through `@register_draft`. The key defaults to the

docs/sections/basic_usage/training.md

@@ -487,6 +487,39 @@ publishes the exact schedule horizon and the consumer trains to EOF.
 
 ## Parallel topologies
 
+`training.backend` selects `fsdp` (the existing FSDP1 implementation, default)
+or `fsdp2` (PyTorch composable `fully_shard`). For example:
+
+```yaml
+training:
+  backend: fsdp2
+  fsdp_sharding: SHARD_GRAD_OP
+```
+
+Both backends retain BF16 compute, FP32 optimizer masters, gradient accumulation,
+and `training.optimizer_cpu_offload`. `SHARD_GRAD_OP` keeps parameters gathered
+through backward and between gradient-accumulation micro-steps, resharding at
+the optimizer boundary. This trades higher live memory between micro-steps for
+fewer all-gathers. `FULL_SHARD` reshards child blocks after forward but keeps the
+root gathered for backward; all units reshard after each backward. Both
+backends use DDP for `NO_SHARD`. Configured sharding takes precedence over the
+legacy `FSDP_SHARDING` environment fallback used by direct Python builders.
+
+FSDP2 uses per-parameter DTensor shards and explicit block-level sharding. Its
+memory management avoids FSDP1's CPU all-gather rate limiter and provides a
+foundation for future DTensor-based parallelism. Throughput and peak memory
+still depend on the model and sharding policy; selecting FSDP2 alone does not
+guarantee a speedup or enable tensor parallelism.
+
+MoE drafts can add expert parallelism on FSDP2 with
+`training.expert_parallel_size`: each MoE layer's routed experts are sliced
+across that many consecutive ranks, the group exchanges its tokens around the
+expert computation, and FSDP2 shards the slice over the ranks that own the same
+experts while dense parameters stay on the full mesh. It requires a sharded
+`fsdp_sharding`, `tp_size: 1`, no sequence parallelism, and a world size and
+`n_routed_experts` divisible by it; see the MoE section of the customization
+guide.
+
 The launcher creates every process group from the typed run config:
 
 - Online target TP/EP belongs to each external SGLang capture server, not the
@@ -658,6 +691,14 @@ The producer itself is not restarted or resumed. Optimizer/FSDP checkpoints
 currently require the same trainer world size; control-plane ref redistribution
 does not imply optimizer-state resharding.
 
+Resume also requires the same training backend and sharding strategy. Old
+checkpoints without backend metadata are treated as FSDP1 checkpoints. FSDP1's
+flat-parameter optimizer shards cannot be loaded into FSDP2's per-parameter
+layout. To switch backends, export the draft and start a new run from its model
+weights; optimizer moments, scheduler position, and FP32 master precision are
+not transferred by that workflow. FSDP2 retains the existing full draft-weight
+checkpoint format, so HF and SGLang export commands remain the same.
+
 Training metrics are printed every `training.log_interval` steps and forwarded
 to the configured tracking backend.
 

examples/configs/README.md

@@ -88,6 +88,16 @@ servers. Its
 [runbook](../../docs/recipes/deepseek-v4-flash-dspark-disaggregated.md)
 covers the v0.5.18 SGLang capture patch and the bundled `deepseek-v4` chat
 template (the checkpoint ships no Jinja template).
+`deepseek-v4-flash-dspark-moe-disaggregated.yaml` is the MoE-FFN arm of the
+drafter-architecture ablation: the same recipe with
+`configs/deepseek-v4-flash-dspark-moe.json`, whose `moe_preset: deepseek_v4`
+swaps the dense MLP for the target's routing (64 routed + 1 shared experts,
+top-6, width 2048); see the runbook's MoE section.
+`deepseek-v4-flash-dspark-moe-ep4-disaggregated.yaml` is that arm with the
+experts trained under expert parallelism on the FSDP2 backend
+(`training.backend: fsdp2`, `training.expert_parallel_size: 4`): each of the
+four trainer ranks owns 16 of the 64 experts instead of re-gathering all of
+them every micro-batch.
 
 `qwen3.8-27b-dflash2-disaggregated.yaml` (external services, two nodes) and
 its managed-local siblings `qwen3.8-27b-dflash2-4server-dp4-disaggregated.yaml`
@@ -273,6 +283,7 @@ Common fields:
 | Field | Default | What to write |
 | --- | --- | --- |
 | `training.strategy` | `eagle3` | `eagle3`, `peagle`, `dflash`, `domino`, or `dspark`. |
+| `training.backend` | `fsdp` | `fsdp` for FSDP1 or `fsdp2` for composable sharding. Both retain the DDP path for `NO_SHARD`; resume requires the original backend. |
 | `training.num_epochs` | `1` | Positive passes over a finite source. |
 | `training.max_steps` | `null` | Positive hard stop in optimizer steps. If it is set while `total_steps` is omitted, it is also the fallback schedule horizon. |
 | `training.total_steps` | `null` | Positive optimizer/loss schedule horizon; it does not itself stop an online stream. A finite online disaggregated run may omit both fields: the producer publishes the exact horizon derived from prepared prompts, epochs, DP size, batch size, and accumulation. |
@@ -288,6 +299,7 @@ Common fields:
 | `training.tp_size` | `1` | Online disaggregated consumers must keep it at 1; configure target TP on capture servers. Offline non-USP ranks consume disjoint data. |
 | `training.sp_ulysses_size` | `1` | Ulysses sequence-parallel factor for offline EAGLE3 USP. |
 | `training.sp_ring_size` | `1` | Ring sequence-parallel factor for offline EAGLE3 USP. |
+| `training.expert_parallel_size` | `1` | Expert parallelism for MoE drafts on the FSDP2 backend: each MoE layer's routed experts are sliced across this many consecutive ranks (the group all-gathers its tokens and reduce-scatters the expert outputs) and FSDP2 shards the slice over the remaining ranks. Requires `backend: fsdp2`, a sharded `fsdp_sharding`, `tp_size: 1`, no sequence parallelism, a world size divisible by it, and `n_routed_experts` divisible by it. |
 | `training.dist_timeout` | `10` | Positive distributed-operation timeout in minutes. |
 | `training.save_interval` | `0` | Save every N optimizer steps; 0 disables periodic saves. A final checkpoint is still written. |
 | `training.eval_interval` | `0` | Evaluate every N optimizer steps; 0 disables evaluation. |

examples/configs/online/disaggregated/external/deepseek-v4-flash-dspark-moe-disaggregated.yaml

@@ -0,0 +1,97 @@
+# MoE-FFN arm of the DSpark drafter-architecture ablation for
+# DeepSeek-V4-Flash. Identical to deepseek-v4-flash-dspark-disaggregated.yaml
+# (the dense arm) except the draft JSON and the run/store names, so the A/B
+# diff is the FFN: 64 routed experts + 1 shared, top-6, moe_intermediate 2048
+# (activated width 6x2048 + 2048 ~= the dense 12288) with DeepSeek-V4 routing.
+# If a from-scratch MoE run destabilizes at the shared hparams, lower the LR
+# and clip on BOTH arms rather than on this file alone.
+model:
+  target_model_path: deepseek-ai/DeepSeek-V4-Flash-0731
+  draft_model_config: configs/deepseek-v4-flash-dspark-moe.json
+  target_backend: sglang
+  trust_remote_code: true
+  # DeepSeek-V4 checkpoints use the native inference weight layout.
+  embedding_key: embed.weight
+  lm_head_key: head.weight
+  mask_token_id: 128799
+  torch_dtype: bfloat16
+  sglang_mem_fraction_static: 0.85
+  # data.max_length plus headroom: capture rejects inputs at exactly context length.
+  sglang_context_length: 8704
+  sglang_max_running_requests: 8
+  # Routed experts are fp4; the default MoE path cannot run them on B200.
+  sglang_moe_runner_backend: flashinfer_mxfp4
+
+data:
+  train_data_path: ./cache/dataset/sharegpt_train.jsonl
+  max_length: 8192
+  chat_template: deepseek-v4
+  cache_dir: cache
+  build_dataset_num_proc: 64
+  # Async prefetch: overlaps per-sample Mooncake feature fetches with compute.
+  dataloader_num_workers: 8
+
+training:
+  strategy: dspark
+  num_epochs: 2
+  # 4 ranks x 32 microbatches -> global batch 128.
+  batch_size: 1
+  accumulation_steps: 32
+  learning_rate: 0.0006
+  lr_scheduler: constant
+  warmup_ratio: 0
+  max_grad_norm: 1
+  attention_backend: flex_attention
+  num_anchors: 512
+  loss_decay_gamma: 4.0
+  objective_chunk_blocks: 128
+  dspark_ce_loss_alpha: 0.1
+  dspark_l1_loss_alpha: 0.9
+  dspark_confidence_head_alpha: 1.0
+  save_interval: 128
+  log_interval: 10
+  dist_timeout: 30
+  seed: 42
+  prompt_seed: 1
+
+tracking:
+  report_to: wandb
+  wandb_project: specforge
+  wandb_name: deepseek-v4-flash-dspark-moe-disaggregated
+  wandb_dir: outputs/deepseek-v4-flash-dspark-moe-disaggregated/wandb
+
+runtime:
+  producer_lease: 8
+  producer_concurrency: 8
+  # Keep two 128-sample optimizer quanta in flight (~0.4 GiB features/sample);
+  # the low watermark is one full quantum so complete windows always dispatch.
+  in_flight_high_watermark: 256
+  in_flight_low_watermark: 128
+  resident_high_watermark_bytes: 137438953472
+  resident_low_watermark_bytes: 103079215104
+  feature_store_max_resident_bytes: 171798691840
+
+run_id: deepseek-v4-flash-dspark-moe-disaggregated
+output_dir: outputs/deepseek-v4-flash-dspark-moe-disaggregated
+
+deployment:
+  mode: disaggregated
+  trainer:
+    nnodes: 1
+    nproc_per_node: 4
+  disaggregated:
+    control_dir: outputs/deepseek-v4-flash-dspark-moe-disaggregated/control
+    consumer_state_dir: outputs/deepseek-v4-flash-dspark-moe-disaggregated/consumer-state
+    backend: mooncake
+    store_id: deepseek-v4-flash-dspark-moe-disaggregated
+    # Two TP2 servers out-produce one TP4 server (TP prefill scaling is sublinear).
+    server_urls:
+      - http://127.0.0.1:30000
+      - http://127.0.0.1:30001
+    mooncake_metadata_server: http://127.0.0.1:35880/metadata
+    mooncake_master_server_addr: 127.0.0.1:35551
+    mooncake_local_hostname: 127.0.0.1
+    mooncake_protocol: tcp
+    client_buffer_size: 1073741824
+    idle_timeout_s: 7200
+    peer_wait_timeout_s: 7200