Add FSDP-aware hybrid Muon optimizer

docs/basic_usage/training.md

@@ -88,6 +88,33 @@ specforge train \
 Unknown config fields and unknown override paths are errors. This keeps a
 misspelled or retired option from being silently ignored.
 
+## Hybrid Muon optimizer
+
+Set `training.optimizer: muon` to apply PyTorch Muon to hidden
+`nn.Linear.weight` matrices. Embeddings, embedding projections, output and
+confidence heads, normalization parameters, and biases remain on an auxiliary
+AdamW optimizer. AdamW remains the default, so existing recipes are unchanged.
+
+```yaml
+training:
+  strategy: dspark
+  optimizer: muon
+  learning_rate: 1.0e-4       # auxiliary AdamW
+  muon_learning_rate: 2.0e-2  # null reuses learning_rate
+  weight_decay: 0.0
+  muon_weight_decay: 0.1
+```
+
+Both groups use the configured cosine-warmup schedule and global gradient
+clipping. Tracking reports `lr_muon` and `lr_adamw` in addition to the existing
+primary `lr` metric. Under FSDP, FP32 master weights and momentum stay sharded;
+only one BF16 matrix update is temporarily gathered for each Newton--Schulz
+transform. Optimizer CPU offload is therefore unavailable in Muon mode.
+
+Muon and AdamW checkpoints are deliberately type-checked and are not
+interchangeable. The hybrid checkpoint also records the parameter partition so
+a changed model cannot silently load order-dependent optimizer state.
+
 ## Run config
 
 A run config has seven typed sections (`model`, `data`, `training`, `tracking`,

examples/configs/README.md

@@ -155,6 +155,7 @@ should make their training strategy and topology explicit.
 | `model.cache_dir` | `null` | Model/tokenizer download cache. This is distinct from `data.cache_dir`. |
 | `model.mask_token_id` | `null` | DFlash-family/P-EAGLE mask token override. Otherwise it resolves from the draft config and then the tokenizer. |
 | `model.tokenizer_pad_token_id` | `null` | Explicit non-negative tokenizer pad ID. Use it for released tokenizers that omit padding metadata. |
+| `model.use_liger_kernel` | `false` | Enable the optional Liger fused kernel for supported DFlash-family drafts. |
 | `model.sglang_attention_backend` | `flashinfer` | SGLang attention implementation for an in-process or managed capture server. |
 | `model.sglang_mem_fraction_static` | `0.4` | SGLang static-memory fraction in `(0, 1]`; inherited by managed capture servers unless they override it. |
 | `model.sglang_context_length` | `null` | Positive explicit context limit. Managed capture requires at least `data.max_length + 7`; omitting it derives that value. |
@@ -215,10 +216,18 @@ Common fields:
 | `training.batch_size` | `1` | Per-rank microbatch size. P-EAGLE and USP require 1. |
 | `training.accumulation_steps` | `1` | Positive microbatches per optimizer update. |
 | `training.fsdp_sharding` | `SHARD_GRAD_OP` | Trainer FSDP mode: `SHARD_GRAD_OP`, `FULL_SHARD`, or `NO_SHARD`. |
+| `training.optimizer` | `adamw` | `adamw` or the hybrid `muon`/AdamW optimizer. |
 | `training.learning_rate` | `1e-4` | Positive peak learning rate. |
+| `training.weight_decay` | `0.0` | Non-negative AdamW weight decay, including the auxiliary AdamW group in Muon mode. |
 | `training.warmup_ratio` | `0.015` | Fraction in `[0, 1]` used for scheduler warmup. |
 | `training.max_grad_norm` | `0.5` | Positive gradient-clipping norm. |
-| `training.optimizer_cpu_offload` | `false` | Keep the optimizer's FP32 master parameters and Adam state on CPU. |
+| `training.optimizer_cpu_offload` | `false` | Keep FP32 masters and Adam state on CPU. It is not supported with Muon. |
+| `training.muon_learning_rate` | `null` | Positive Muon peak learning rate; `null` reuses `training.learning_rate`. |
+| `training.muon_weight_decay` | `0.1` | Non-negative decoupled weight decay for Muon matrices. |
+| `training.muon_momentum` | `0.95` | Muon momentum in `[0, 1)`. |
+| `training.muon_nesterov` | `true` | Apply Nesterov momentum before Muon's Newton--Schulz transform. |
+| `training.muon_ns_steps` | `5` | Newton--Schulz iteration count in `[1, 99]`. |
+| `training.muon_adjust_lr_fn` | `match_rms_adamw` | Shape-aware Muon scaling: `original` or `match_rms_adamw`. |
 | `training.attention_backend` | `flex_attention` | `eager`, `sdpa`, `flex_attention`, `fa`, or `usp`; the selected strategy must support it. |
 | `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. |

specforge/config/schema.py

@@ -486,12 +486,21 @@ class TrainingConfig(StrictConfigModel):
     batch_size: int = Field(default=1, gt=0)
     accumulation_steps: int = Field(default=1, gt=0)
     fsdp_sharding: Literal["SHARD_GRAD_OP", "FULL_SHARD", "NO_SHARD"] = "SHARD_GRAD_OP"
+    optimizer: Literal["adamw", "muon"] = "adamw"
     learning_rate: float = Field(default=1e-4, gt=0.0)
+    weight_decay: float = Field(default=0.0, ge=0.0)
     warmup_ratio: float = Field(default=0.015, ge=0.0, le=1.0)
     max_grad_norm: float = Field(default=0.5, gt=0.0)
     #: Keep FP32 Adam masters and moments on CPU while the trainable draft
     #: remains on the accelerator.
     optimizer_cpu_offload: bool = False
+    #: Muon updates hidden linear matrices; auxiliary tensors stay on AdamW.
+    muon_learning_rate: Optional[float] = Field(default=None, gt=0.0)
+    muon_weight_decay: float = Field(default=0.1, ge=0.0)
+    muon_momentum: float = Field(default=0.95, ge=0.0, lt=1.0)
+    muon_nesterov: bool = True
+    muon_ns_steps: int = Field(default=5, gt=0, lt=100)
+    muon_adjust_lr_fn: Literal["original", "match_rms_adamw"] = "match_rms_adamw"
     ttt_length: int = Field(default=7, gt=0)
     attention_backend: Literal["eager", "sdpa", "flex_attention", "fa", "usp"] = (
         "flex_attention"
@@ -548,6 +557,10 @@ class TrainingConfig(StrictConfigModel):
 
     @model_validator(mode="after")
     def _validate_training_shape(self):
+        if self.optimizer == "muon" and self.optimizer_cpu_offload:
+            raise ValueError(
+                "training.optimizer_cpu_offload is not supported with Muon"
+            )
         if not 0.0 <= self.dpace_alpha <= 1.0:
             raise ValueError("training.dpace_alpha must be in [0, 1]")
         if not 0.0 < self.down_sample_ratio <= 1.0: