From 3c1c41a742eeefa4d7d54494856270efc20ae4c7 Mon Sep 17 00:00:00 2001 From: maocheng23 Date: Wed, 22 Jul 2026 16:56:47 -0700 Subject: [PATCH] feat: add FSDP-aware hybrid Muon optimizer --- docs/basic_usage/training.md | 27 + examples/configs/README.md | 11 +- specforge/config/schema.py | 13 + specforge/muon.py | 376 ++++++++++++++ specforge/optimizer.py | 473 +++++++++++++++--- specforge/training/assembly.py | 19 + specforge/training/backend.py | 8 + specforge/training/controller.py | 4 + tests/test_config/test_schema.py | 14 + .../test_bf16_optimizer_clip_grad_norm.py | 28 ++ tests/test_optimizer/test_muon_optimizer.py | 325 ++++++++++++ 11 files changed, 1219 insertions(+), 79 deletions(-) create mode 100644 specforge/muon.py create mode 100644 tests/test_optimizer/test_muon_optimizer.py diff --git a/docs/basic_usage/training.md b/docs/basic_usage/training.md index 9093cc8cf..46b714ff8 100644 --- a/docs/basic_usage/training.md +++ b/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`, diff --git a/examples/configs/README.md b/examples/configs/README.md index 7be1b8f17..a14a03905 100644 --- a/examples/configs/README.md +++ b/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. | diff --git a/specforge/config/schema.py b/specforge/config/schema.py index d3681560e..67e5e6720 100644 --- a/specforge/config/schema.py +++ b/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: diff --git a/specforge/muon.py b/specforge/muon.py new file mode 100644 index 000000000..f0c3bf494 --- /dev/null +++ b/specforge/muon.py @@ -0,0 +1,376 @@ +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Mapping, Sequence + +import torch +import torch.distributed as dist +import torch.nn as nn + +ADAMW_OPTIMIZER = "adamw" +MUON_OPTIMIZER = "muon" +SUPPORTED_OPTIMIZERS = (ADAMW_OPTIMIZER, MUON_OPTIMIZER) + +# Muon is intended for hidden-layer matrices. These path components identify +# embedding projections and output-like heads that stay on AdamW even when their +# weights are two-dimensional. +DEFAULT_MUON_EXCLUDED_MODULES = ( + "classifier", + "confidence_head", + "embed_proj", + "embed_tokens", + "lm_head", + "markov_head", + "output", + "score", +) + + +@dataclass(frozen=True) +class NamedTrainableParameter: + name: str + parameter: nn.Parameter + logical_shape: torch.Size + + +@dataclass(frozen=True) +class MuonParameterMetadata: + """Pre-FSDP parameter classification and logical matrix shapes.""" + + logical_shapes_by_id: Mapping[int, torch.Size] + muon_parameter_ids: frozenset[int] + + +@dataclass(frozen=True) +class MuonParameterPartition: + """Named trainable parameters assigned to Muon or auxiliary AdamW.""" + + muon: tuple[NamedTrainableParameter, ...] + adamw: tuple[NamedTrainableParameter, ...] + + def names(self) -> dict[str, tuple[str, ...]]: + return { + MUON_OPTIMIZER: tuple(item.name for item in self.muon), + ADAMW_OPTIMIZER: tuple(item.name for item in self.adamw), + } + + +def partition_parameters_for_muon( + model: nn.Module, + *, + excluded_module_names: Sequence[str] = DEFAULT_MUON_EXCLUDED_MODULES, + metadata: MuonParameterMetadata | None = None, +) -> MuonParameterPartition: + """Split trainable parameters into hidden matrices and auxiliary tensors. + + Module ownership is checked rather than tensor rank alone because + embeddings and output heads are also commonly two-dimensional. + """ + + modules = dict(model.named_modules()) + excluded_names = frozenset(excluded_module_names) + muon_parameters: list[NamedTrainableParameter] = [] + adamw_parameters: list[NamedTrainableParameter] = [] + + for name, parameter in model.named_parameters(): + if not parameter.requires_grad: + continue + + if metadata is None: + module_name, separator, parameter_name = name.rpartition(".") + if not separator: + module_name = "" + parameter_name = name + owner = modules.get(module_name) + module_path = frozenset(part for part in module_name.split(".") if part) + use_muon = ( + isinstance(owner, nn.Linear) + and parameter_name == "weight" + and parameter.ndim == 2 + and module_path.isdisjoint(excluded_names) + ) + logical_shape = parameter.shape + else: + parameter_id = id(parameter) + try: + logical_shape = metadata.logical_shapes_by_id[parameter_id] + except KeyError as error: + raise ValueError( + f"Parameter {name!r} was not present when Muon metadata " + "was captured before distributed wrapping" + ) from error + use_muon = parameter_id in metadata.muon_parameter_ids + + item = NamedTrainableParameter( + name=name, + parameter=parameter, + logical_shape=torch.Size(logical_shape), + ) + if use_muon: + muon_parameters.append(item) + else: + adamw_parameters.append(item) + + return MuonParameterPartition( + muon=tuple(muon_parameters), + adamw=tuple(adamw_parameters), + ) + + +def capture_muon_parameter_metadata( + model: nn.Module, + *, + excluded_module_names: Sequence[str] = DEFAULT_MUON_EXCLUDED_MODULES, +) -> MuonParameterMetadata: + """Capture Muon eligibility before FSDP exposes flattened local shards.""" + + partition = partition_parameters_for_muon( + model, excluded_module_names=excluded_module_names + ) + trainable_parameters = partition.muon + partition.adamw + return MuonParameterMetadata( + logical_shapes_by_id={ + id(item.parameter): item.logical_shape for item in trainable_parameters + }, + muon_parameter_ids=frozenset(id(item.parameter) for item in partition.muon), + ) + + +def zeropower_via_newton_schulz( + update: torch.Tensor, + *, + ns_steps: int, + eps: float = 1e-7, +) -> torch.Tensor: + """Match the Newton--Schulz transform used by ``torch.optim.Muon``.""" + + if update.ndim != 2: + raise ValueError(f"Muon update must be 2-D, got shape {tuple(update.shape)}") + + orthogonalized = update.to(torch.bfloat16) + transposed = orthogonalized.size(0) > orthogonalized.size(1) + if transposed: + orthogonalized = orthogonalized.T + orthogonalized.div_(orthogonalized.norm().clamp(min=eps)) + + a, b, c = 3.4445, -4.7750, 2.0315 + for _ in range(ns_steps): + gram_matrix = orthogonalized @ orthogonalized.T + gram_update = torch.addmm( + gram_matrix, gram_matrix, gram_matrix, beta=b, alpha=c + ) + orthogonalized = torch.addmm( + orthogonalized, gram_update, orthogonalized, beta=a + ) + return orthogonalized.T if transposed else orthogonalized + + +def adjust_muon_learning_rate( + learning_rate: float, + adjust_lr_fn: str, + logical_shape: torch.Size, +) -> float: + rows, columns = logical_shape + if adjust_lr_fn == "original": + ratio = math.sqrt(max(1, rows / columns)) + elif adjust_lr_fn == "match_rms_adamw": + ratio = 0.2 * math.sqrt(max(rows, columns)) + else: + raise ValueError(f"Unsupported Muon learning-rate adjustment: {adjust_lr_fn}") + return learning_rate * ratio + + +@dataclass(frozen=True) +class _ShardLayout: + sizes: tuple[int, ...] + offset: int + + +class FSDPShardedMuon(torch.optim.Optimizer): + """Muon over FSDP1 local shards with sharded persistent state. + + Newton--Schulz needs a logical 2-D update. This optimizer keeps FP32 + momentum sharded, gathers one BF16 update at a time, applies the same + transform as native Muon, and writes only the local slice back. + """ + + def __init__( + self, + params: Sequence[torch.Tensor], + logical_shapes: Sequence[torch.Size], + *, + lr: float, + weight_decay: float, + momentum: float, + nesterov: bool, + ns_steps: int, + adjust_lr_fn: str, + ) -> None: + parameters = list(params) + if len(parameters) != len(logical_shapes): + raise ValueError("Every sharded Muon parameter needs a logical shape") + if lr < 0: + raise ValueError(f"Muon learning rate must be non-negative, got {lr}") + if weight_decay < 0: + raise ValueError( + f"Muon weight decay must be non-negative, got {weight_decay}" + ) + if not 0 <= momentum < 1: + raise ValueError(f"Muon momentum must be in [0, 1), got {momentum}") + if not 0 < ns_steps < 100: + raise ValueError(f"Muon ns_steps must be in [1, 99], got {ns_steps}") + if adjust_lr_fn not in ("original", "match_rms_adamw"): + raise ValueError( + f"Unsupported Muon learning-rate adjustment: {adjust_lr_fn}" + ) + + defaults = { + "lr": lr, + "weight_decay": weight_decay, + "momentum": momentum, + "nesterov": nesterov, + "ns_steps": ns_steps, + "adjust_lr_fn": adjust_lr_fn, + } + super().__init__(parameters, defaults) + self._logical_shapes = { + id(parameter): torch.Size(shape) + for parameter, shape in zip(parameters, logical_shapes) + } + self._shard_layouts: dict[int, _ShardLayout] = {} + self._process_group = None + self._process_group_configured = False + + def configure_process_group(self, process_group=None) -> None: + """Bind the FSDP group and discover each rank-local shard layout.""" + + if not dist.is_available() or not dist.is_initialized(): + raise RuntimeError("Sharded Muon requires an initialized process group") + if dist.get_world_size(group=process_group) <= 1: + raise RuntimeError("Sharded Muon requires more than one process") + + self._process_group = process_group + self._process_group_configured = True + self._shard_layouts = {} + for group in self.param_groups: + for parameter in group["params"]: + logical_shape = self._logical_shapes[id(parameter)] + self._shard_layouts[id(parameter)] = self._build_shard_layout( + parameter, logical_shape + ) + + def _build_shard_layout( + self, parameter: torch.Tensor, logical_shape: torch.Size + ) -> _ShardLayout: + local_size = torch.tensor( + parameter.numel(), dtype=torch.int64, device=parameter.device + ) + world_size = dist.get_world_size(group=self._process_group) + gathered_sizes = [torch.zeros_like(local_size) for _ in range(world_size)] + dist.all_gather(gathered_sizes, local_size, group=self._process_group) + sizes = tuple(int(size.item()) for size in gathered_sizes) + expected_numel = math.prod(logical_shape) + if sum(sizes) != expected_numel: + raise RuntimeError( + "FSDP Muon shards do not reconstruct the logical matrix: " + f"shape={tuple(logical_shape)}, shard_sizes={sizes}" + ) + group_rank = dist.get_rank(group=self._process_group) + return _ShardLayout(sizes=sizes, offset=sum(sizes[:group_rank])) + + def _all_gather_flat( + self, local_tensor: torch.Tensor, layout: _ShardLayout + ) -> torch.Tensor: + max_size = max(layout.sizes) + padded = local_tensor.new_zeros(max_size) + padded[: local_tensor.numel()].copy_(local_tensor.reshape(-1)) + gathered = [torch.empty_like(padded) for _ in layout.sizes] + dist.all_gather(gathered, padded, group=self._process_group) + return torch.cat( + [tensor[:size] for tensor, size in zip(gathered, layout.sizes)] + ) + + @torch.no_grad() + def step(self, closure=None): + if not self._process_group_configured: + raise RuntimeError( + "Sharded Muon process group was not configured by the training backend" + ) + + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + learning_rate = float(group["lr"]) + for parameter in group["params"]: + has_local_grad = torch.tensor( + parameter.grad is not None, + dtype=torch.uint8, + device=parameter.device, + ) + dist.all_reduce( + has_local_grad, + op=dist.ReduceOp.MAX, + group=self._process_group, + ) + if not bool(has_local_grad.item()): + continue + + local_gradient = ( + torch.zeros_like(parameter) + if parameter.grad is None + else parameter.grad + ) + if local_gradient.shape != parameter.shape: + raise RuntimeError( + "FSDP Muon gradient and local parameter shard differ: " + f"gradient={tuple(local_gradient.shape)}, " + f"parameter={tuple(parameter.shape)}" + ) + + state = self.state[parameter] + if "momentum_buffer" not in state: + state["momentum_buffer"] = torch.zeros_like(local_gradient) + momentum_buffer = state["momentum_buffer"] + momentum_buffer.lerp_(local_gradient, 1 - group["momentum"]) + local_update = ( + local_gradient.lerp(momentum_buffer, group["momentum"]) + if group["nesterov"] + else momentum_buffer + ) + + logical_shape = self._logical_shapes[id(parameter)] + layout = self._shard_layouts[id(parameter)] + full_update = self._all_gather_flat( + local_update.to(torch.bfloat16), layout + ).reshape(logical_shape) + orthogonalized = zeropower_via_newton_schulz( + full_update, ns_steps=group["ns_steps"] + ).reshape(-1) + local_orthogonalized = orthogonalized.narrow( + 0, layout.offset, parameter.numel() + ) + adjusted_lr = adjust_muon_learning_rate( + learning_rate, group["adjust_lr_fn"], logical_shape + ) + + parameter.mul_(1 - learning_rate * group["weight_decay"]) + parameter.add_(local_orthogonalized, alpha=-adjusted_lr) + return loss + + +__all__ = [ + "ADAMW_OPTIMIZER", + "DEFAULT_MUON_EXCLUDED_MODULES", + "FSDPShardedMuon", + "MUON_OPTIMIZER", + "MuonParameterMetadata", + "MuonParameterPartition", + "NamedTrainableParameter", + "SUPPORTED_OPTIMIZERS", + "capture_muon_parameter_metadata", + "partition_parameters_for_muon", +] diff --git a/specforge/optimizer.py b/specforge/optimizer.py index 3d5e1aabe..3c2b3b0dc 100644 --- a/specforge/optimizer.py +++ b/specforge/optimizer.py @@ -1,17 +1,78 @@ +from __future__ import annotations + import logging +import math +from dataclasses import dataclass +from typing import Mapping, Sequence import torch import torch.distributed as dist from specforge.lr_scheduler import CosineAnnealingWarmupLR +from specforge.muon import ( + ADAMW_OPTIMIZER, + DEFAULT_MUON_EXCLUDED_MODULES, + MUON_OPTIMIZER, + SUPPORTED_OPTIMIZERS, + FSDPShardedMuon, + MuonParameterMetadata, + NamedTrainableParameter, + partition_parameters_for_muon, +) from specforge.utils import print_on_rank0 logger = logging.getLogger(__name__) +@dataclass +class _SchedulerCollection: + """Expose multiple schedulers through the existing scheduler interface.""" + + schedulers: Mapping[str, CosineAnnealingWarmupLR] + + _FORMAT_VERSION = 1 + + def step(self) -> None: + for scheduler in self.schedulers.values(): + scheduler.step() + + def state_dict(self) -> dict: + return { + "format_version": self._FORMAT_VERSION, + "schedulers": { + name: scheduler.state_dict() + for name, scheduler in self.schedulers.items() + }, + } + + def load_state_dict(self, state_dict: dict) -> None: + if state_dict.get("format_version") != self._FORMAT_VERSION: + raise ValueError( + "Unsupported hybrid scheduler state format: " + f"{state_dict.get('format_version')!r}" + ) + scheduler_states = state_dict.get("schedulers") + if not isinstance(scheduler_states, dict): + raise ValueError("Hybrid scheduler state is missing 'schedulers'") + if set(scheduler_states) != set(self.schedulers): + raise ValueError( + "Hybrid scheduler groups do not match: " + f"expected={sorted(self.schedulers)}, " + f"received={sorted(scheduler_states)}" + ) + for name, scheduler in self.schedulers.items(): + scheduler.load_state_dict(scheduler_states[name]) + + class BF16Optimizer: - """AdamW over fp32 master copies of the bf16 trainable params, with grad - clipping and cosine warmup scheduling.""" + """FP32-master AdamW optimizer with opt-in hybrid Muon support. + + The default AdamW path preserves the historical state format and training + behavior. Muon mode applies Muon to hidden linear matrices and an auxiliary + AdamW to embeddings, heads, normalization parameters, and biases. + """ + + _HYBRID_STATE_FORMAT_VERSION = 1 def __init__( self, @@ -22,51 +83,223 @@ def __init__( total_steps=800_000, warmup_ratio=0.015, offload_master=False, + *, + optimizer_type: str = ADAMW_OPTIMIZER, + muon_lr: float | None = None, + muon_weight_decay: float = 0.1, + muon_momentum: float = 0.95, + muon_nesterov: bool = True, + muon_ns_steps: int = 5, + muon_adjust_lr_fn: str = "match_rms_adamw", + muon_excluded_module_names: Sequence[str] = DEFAULT_MUON_EXCLUDED_MODULES, + muon_metadata: MuonParameterMetadata | None = None, ): - # defaults copied from EAGLE traineagle3 ds_config.json + optimizer_type = optimizer_type.lower() + if optimizer_type not in SUPPORTED_OPTIMIZERS: + raise ValueError( + f"Unknown optimizer_type={optimizer_type!r}; " + f"expected one of {SUPPORTED_OPTIMIZERS}" + ) + if optimizer_type == MUON_OPTIMIZER and offload_master: + raise ValueError( + "Muon does not support optimizer CPU offload because its " + "Newton-Schulz step must run on the accelerator" + ) + self.model = model - self.model_params = [p for p in model.parameters() if p.requires_grad] + self.optimizer_type = optimizer_type self.max_grad_norm = max_grad_norm self.offload_master = bool(offload_master) - self.fp32_params = [ - ( - p.detach().to(device="cpu", dtype=torch.float32).clone() - if self.offload_master - else p.detach().clone().to(torch.float32) + self.last_grad_norm = None + self._grad_norm_process_group = None + self._reduce_grad_norm_across_ranks = True + + logical_shapes = ( + muon_metadata.logical_shapes_by_id if muon_metadata is not None else {} + ) + named_parameters = tuple( + NamedTrainableParameter( + name=name, + parameter=parameter, + logical_shape=torch.Size( + logical_shapes.get(id(parameter), parameter.shape) + ), ) - for p in self.model_params + for name, parameter in model.named_parameters() + if parameter.requires_grad + ) + if not named_parameters: + raise ValueError( + "Cannot construct an optimizer with no trainable parameters" + ) + + self._model_param_names = tuple(item.name for item in named_parameters) + self._logical_shapes_by_name = { + item.name: item.logical_shape for item in named_parameters + } + self.model_params = [item.parameter for item in named_parameters] + self.fp32_params = [ + self._new_master_parameter(item.parameter) for item in named_parameters ] - for mp in self.fp32_params: - mp.requires_grad = True + master_by_name = { + item.name: master + for item, master in zip(named_parameters, self.fp32_params) + } + + if optimizer_type == ADAMW_OPTIMIZER: + self._init_adamw( + named_parameters=named_parameters, + lr=lr, + weight_decay=weight_decay, + total_steps=total_steps, + warmup_ratio=warmup_ratio, + ) + else: + partition = partition_parameters_for_muon( + model, + excluded_module_names=muon_excluded_module_names, + metadata=muon_metadata, + ) + self._init_muon( + partition=partition, + master_by_name=master_by_name, + adamw_lr=lr, + adamw_weight_decay=weight_decay, + muon_lr=lr if muon_lr is None else muon_lr, + muon_weight_decay=muon_weight_decay, + muon_momentum=muon_momentum, + muon_nesterov=muon_nesterov, + muon_ns_steps=muon_ns_steps, + muon_adjust_lr_fn=muon_adjust_lr_fn, + total_steps=total_steps, + warmup_ratio=warmup_ratio, + ) + + def _new_master_parameter(self, parameter: torch.Tensor) -> torch.Tensor: + master = ( + parameter.detach().to(device="cpu", dtype=torch.float32).clone() + if self.offload_master + else parameter.detach().clone().to(torch.float32) + ) + master.requires_grad = True + return master + + def _init_adamw( + self, + *, + named_parameters: Sequence[NamedTrainableParameter], + lr: float, + weight_decay: float, + total_steps: int, + warmup_ratio: float, + ) -> None: self.optimizer = torch.optim.AdamW( self.fp32_params, lr=lr, weight_decay=weight_decay ) - self.last_grad_norm = None - self._grad_norm_process_group = None - self._reduce_grad_norm_across_ranks = True + self.aux_optimizer = None + self._optimizers = {ADAMW_OPTIMIZER: self.optimizer} + self._parameter_group_names = { + ADAMW_OPTIMIZER: tuple(item.name for item in named_parameters) + } self.scheduler = CosineAnnealingWarmupLR( self.optimizer, total_steps=total_steps, warmup_steps=int(warmup_ratio * total_steps), ) + def _init_muon( + self, + *, + partition, + master_by_name: Mapping[str, torch.Tensor], + adamw_lr: float, + adamw_weight_decay: float, + muon_lr: float, + muon_weight_decay: float, + muon_momentum: float, + muon_nesterov: bool, + muon_ns_steps: int, + muon_adjust_lr_fn: str, + total_steps: int, + warmup_ratio: float, + ) -> None: + if not partition.muon: + raise ValueError( + "Muon mode found no eligible hidden nn.Linear weight matrices" + ) + muon_class = getattr(torch.optim, "Muon", None) + if muon_class is None: + raise RuntimeError( + "optimizer_type='muon' requires torch.optim.Muon (PyTorch >= 2.9)" + ) + + muon_parameters = [master_by_name[item.name] for item in partition.muon] + adamw_parameters = [master_by_name[item.name] for item in partition.adamw] + locally_sharded = any( + tuple(parameter.shape) != tuple(item.logical_shape) + for parameter, item in zip(muon_parameters, partition.muon) + ) + if locally_sharded: + self.optimizer = FSDPShardedMuon( + muon_parameters, + [item.logical_shape for item in partition.muon], + lr=muon_lr, + weight_decay=muon_weight_decay, + momentum=muon_momentum, + nesterov=muon_nesterov, + ns_steps=muon_ns_steps, + adjust_lr_fn=muon_adjust_lr_fn, + ) + else: + self.optimizer = muon_class( + muon_parameters, + lr=muon_lr, + weight_decay=muon_weight_decay, + momentum=muon_momentum, + nesterov=muon_nesterov, + ns_steps=muon_ns_steps, + adjust_lr_fn=muon_adjust_lr_fn, + ) + + self.aux_optimizer = ( + torch.optim.AdamW( + adamw_parameters, + lr=adamw_lr, + weight_decay=adamw_weight_decay, + ) + if adamw_parameters + else None + ) + self._optimizers = {MUON_OPTIMIZER: self.optimizer} + if self.aux_optimizer is not None: + self._optimizers[ADAMW_OPTIMIZER] = self.aux_optimizer + self._parameter_group_names = partition.names() + self.scheduler = _SchedulerCollection( + { + name: CosineAnnealingWarmupLR( + optimizer, + total_steps=total_steps, + warmup_steps=int(warmup_ratio * total_steps), + ) + for name, optimizer in self._optimizers.items() + } + ) + def configure_grad_norm_reduction( self, *, process_group=None, enabled: bool = True ) -> None: - """Configure the group that owns disjoint gradient shards. + """Configure collectives for the group that owns parameter shards.""" - FSDP backends disable the reduction for replicated/NO_SHARD parameters. - """ self._grad_norm_process_group = process_group self._reduce_grad_norm_across_ranks = enabled + if isinstance(self.optimizer, FSDPShardedMuon): + if not enabled: + raise RuntimeError( + "Flattened Muon parameters require sharded gradient reduction" + ) + self.optimizer.configure_process_group(process_group) def _reduce_grad_norm(self, total_norm_sq): - """All-reduce the squared L2 norm across shard ranks and derive the - clip coefficient. - - ``total_norm_sq`` must already live on a device the process group can - reduce (e.g. CUDA for NCCL). Returns ``(total_norm, clip_coef)``. - """ if ( self._reduce_grad_norm_across_ranks and dist.is_available() @@ -82,9 +315,6 @@ def _reduce_grad_norm(self, total_norm_sq): return total_norm, clip_coef def _grad_norm_and_clip_coefficient(self): - """Compute the global grad norm from the model params on their own - device, where NCCL can reduce it safely, without materialising master - gradients first.""" grads = [p.grad.detach() for p in self.model_params if p.grad is not None] if grads: total_norm_sq = torch.stack( @@ -96,12 +326,8 @@ def _grad_norm_and_clip_coefficient(self): return self._reduce_grad_norm(total_norm_sq) def _clip_grad_norm(self): - """Clip already-populated FP32 master gradients in place. + """Clip populated master gradients; retained for custom loops/tests.""" - Convenience entry point for optimizer tests and custom loops. When - masters are CPU-offloaded, only the scalar norm is moved to the model - device so a NCCL process group can still participate in the reduction. - """ grads = [master.grad for master in self.fp32_params if master.grad is not None] if grads: local_norm_sq = torch.stack( @@ -132,34 +358,43 @@ def step(self): float(clip_coefficient.item()) if self.offload_master else None ) with torch.no_grad(): - for p, mp in zip(self.model_params, self.fp32_params): - if p.grad is None: - mp.grad = None + for name, parameter, master in zip( + self._model_param_names, self.model_params, self.fp32_params + ): + if parameter.grad is None: + master.grad = None continue - master_grad = p.grad.detach().to( - device=mp.device, - dtype=torch.float32, + if parameter.grad.shape != master.shape: + raise RuntimeError( + "Optimizer gradient shape changed after distributed " + f"wrapping for {name!r}: gradient={tuple(parameter.grad.shape)}, " + f"master={tuple(master.shape)}" + ) + master_grad = parameter.grad.detach().to( + device=master.device, dtype=torch.float32 ) master_grad.mul_( cpu_clip_coefficient if cpu_clip_coefficient is not None else clip_coefficient ) - mp.grad = master_grad + master.grad = master_grad + self.last_grad_norm = grad_norm.detach() - self.optimizer.step() - self.optimizer.zero_grad() + for optimizer in self._optimizers.values(): + optimizer.step() + optimizer.zero_grad() self.scheduler.step() + with torch.no_grad(): - for p, mp in zip(self.model_params, self.fp32_params): - p.data.copy_(mp.data.to(device=p.device, dtype=p.dtype)) - p.grad = None + for parameter, master in zip(self.model_params, self.fp32_params): + parameter.data.copy_( + master.data.to(device=parameter.device, dtype=parameter.dtype) + ) + parameter.grad = None return self.last_grad_norm - def load_state_dict(self, state_dict): - """Restore optimizer/scheduler state and, when present, the rank-local - fp32 master params; without them the masters are re-cloned from the - bf16 weights and the resume is not numerically faithful.""" + def _validate_max_grad_norm(self, state_dict: dict) -> None: saved_max_grad_norm = state_dict.get("max_grad_norm") if saved_max_grad_norm is not None and float(saved_max_grad_norm) != float( self.max_grad_norm @@ -169,45 +404,127 @@ def load_state_dict(self, state_dict): f"{saved_max_grad_norm} but this run has " f"max_grad_norm={self.max_grad_norm}" ) - # offload_master is a pure device-placement choice: restored fp32 - # masters and Adam moments are relocated to the current master device, - # so toggling it on resume is safe and intentionally not gated here. - self.optimizer.load_state_dict(state_dict["optimizer_state_dict"]) - print_on_rank0("Successfully loaded optimizer state_dict.") - self.scheduler.load_state_dict(state_dict["scheduler_state_dict"]) - print_on_rank0("Successfully loaded scheduler state_dict.") + + def _restore_fp32_params(self, state_dict: dict) -> None: saved_fp32 = state_dict.get("fp32_params") - if saved_fp32 is not None: - if len(saved_fp32) != len(self.fp32_params): - raise ValueError( - f"checkpoint carries {len(saved_fp32)} fp32 master params " - f"but this rank has {len(self.fp32_params)}" - ) - with torch.no_grad(): - for i, (saved, mp) in enumerate(zip(saved_fp32, self.fp32_params)): - if saved.shape != mp.shape: - raise ValueError( - f"fp32 master param {i} shape mismatch: checkpoint " - f"{tuple(saved.shape)} vs current {tuple(mp.shape)}" - ) - mp.data.copy_(saved.to(mp.device, mp.dtype)) - else: + if saved_fp32 is None: logger.warning( "checkpoint has no fp32_params; re-cloning master params from " "bf16 weights — resume will not be numerically faithful" ) - with torch.no_grad(): - for p, mp in zip(self.model_params, self.fp32_params): - mp.data.copy_(p.detach().to(device=mp.device, dtype=mp.dtype)) + saved_fp32 = [parameter.detach() for parameter in self.model_params] + if len(saved_fp32) != len(self.fp32_params): + raise ValueError( + f"checkpoint carries {len(saved_fp32)} fp32 master params " + f"but this rank has {len(self.fp32_params)}" + ) + with torch.no_grad(): + for index, (saved, master) in enumerate(zip(saved_fp32, self.fp32_params)): + if saved.shape != master.shape: + raise ValueError( + f"fp32 master param {index} shape mismatch: checkpoint " + f"{tuple(saved.shape)} vs current {tuple(master.shape)}" + ) + master.data.copy_(saved.to(master.device, master.dtype)) + + def load_state_dict(self, state_dict): + self._validate_max_grad_norm(state_dict) + checkpoint_type = state_dict.get("optimizer_type", ADAMW_OPTIMIZER) + if self.optimizer_type == ADAMW_OPTIMIZER: + if checkpoint_type != ADAMW_OPTIMIZER: + raise ValueError("Cannot load a Muon optimizer state into AdamW") + self.optimizer.load_state_dict(state_dict["optimizer_state_dict"]) + else: + self._load_hybrid_optimizer_state(state_dict) + print_on_rank0("Successfully loaded optimizer state_dict.") + self.scheduler.load_state_dict(state_dict["scheduler_state_dict"]) + print_on_rank0("Successfully loaded scheduler state_dict.") + self._restore_fp32_params(state_dict) + + def _load_hybrid_optimizer_state(self, state_dict: dict) -> None: + if state_dict.get("optimizer_type") != MUON_OPTIMIZER: + raise ValueError("Cannot load a non-Muon optimizer state into Muon") + optimizer_state = state_dict.get("optimizer_state_dict") + if not isinstance(optimizer_state, dict): + raise ValueError("Muon checkpoint is missing 'optimizer_state_dict'") + if optimizer_state.get("format_version") != self._HYBRID_STATE_FORMAT_VERSION: + raise ValueError( + "Unsupported hybrid optimizer state format: " + f"{optimizer_state.get('format_version')!r}" + ) + + saved_names = optimizer_state.get("parameter_group_names") + current_names = { + name: list(names) for name, names in self._parameter_group_names.items() + } + if saved_names != current_names: + raise ValueError( + "Muon parameter partition differs from the checkpoint; refusing " + "an order-dependent optimizer-state load" + ) + + saved_optimizers = optimizer_state.get("optimizers") + if not isinstance(saved_optimizers, dict): + raise ValueError("Muon checkpoint is missing optimizer group states") + if set(saved_optimizers) != set(self._optimizers): + raise ValueError( + "Muon optimizer groups do not match: " + f"expected={sorted(self._optimizers)}, " + f"received={sorted(saved_optimizers)}" + ) + for name, optimizer in self._optimizers.items(): + optimizer.load_state_dict(saved_optimizers[name]) def state_dict(self): - return { - "optimizer_state_dict": self.optimizer.state_dict(), + common_state = { "scheduler_state_dict": self.scheduler.state_dict(), "max_grad_norm": self.max_grad_norm, - # rank-local fp32 masters; without them a resume re-quantizes from bf16 - "fp32_params": [t.detach().cpu() for t in self.fp32_params], + "fp32_params": [tensor.detach().cpu() for tensor in self.fp32_params], + } + if self.optimizer_type == ADAMW_OPTIMIZER: + return { + "optimizer_state_dict": self.optimizer.state_dict(), + **common_state, + } + return { + "optimizer_type": MUON_OPTIMIZER, + "optimizer_state_dict": { + "format_version": self._HYBRID_STATE_FORMAT_VERSION, + "parameter_group_names": { + name: list(names) + for name, names in self._parameter_group_names.items() + }, + "optimizers": { + name: optimizer.state_dict() + for name, optimizer in self._optimizers.items() + }, + }, + **common_state, } def get_learning_rate(self): - return self.optimizer.param_groups[0]["lr"] + primary_name = ( + MUON_OPTIMIZER if self.optimizer_type == MUON_OPTIMIZER else ADAMW_OPTIMIZER + ) + return self._optimizers[primary_name].param_groups[0]["lr"] + + def get_learning_rates(self) -> dict[str, float]: + return { + name: float(optimizer.param_groups[0]["lr"]) + for name, optimizer in self._optimizers.items() + } + + def get_parameter_group_summary(self) -> dict[str, dict[str, object]]: + return { + group_name: { + "parameter_count": len(names), + "numel": sum( + math.prod(self._logical_shapes_by_name[name]) for name in names + ), + "names": names, + } + for group_name, names in self._parameter_group_names.items() + } + + +__all__ = ["BF16Optimizer"] diff --git a/specforge/training/assembly.py b/specforge/training/assembly.py index 5fa2e3091..4901ecb93 100644 --- a/specforge/training/assembly.py +++ b/specforge/training/assembly.py @@ -247,6 +247,16 @@ class _ConfiguredOptimizerFactory: def __init__(self, cfg: Config) -> None: self.cfg = cfg self.total_steps = cfg.training.total_steps or cfg.training.max_steps + self.muon_metadata = None + + def capture_parameter_metadata(self, draft_module) -> None: + """Record logical matrix shapes before FSDP exposes local shards.""" + + if self.cfg.training.optimizer != "muon": + return + from specforge.muon import capture_muon_parameter_metadata + + self.muon_metadata = capture_muon_parameter_metadata(draft_module) def configure_total_steps(self, total_steps: int) -> None: if self.total_steps is None: @@ -266,10 +276,19 @@ def __call__(self, draft_module): return BF16Optimizer( draft_module, lr=t.learning_rate, + weight_decay=t.weight_decay, max_grad_norm=t.max_grad_norm, warmup_ratio=t.warmup_ratio, total_steps=self.total_steps, offload_master=t.optimizer_cpu_offload, + optimizer_type=t.optimizer, + muon_lr=t.muon_learning_rate, + muon_weight_decay=t.muon_weight_decay, + muon_momentum=t.muon_momentum, + muon_nesterov=t.muon_nesterov, + muon_ns_steps=t.muon_ns_steps, + muon_adjust_lr_fn=t.muon_adjust_lr_fn, + muon_metadata=self.muon_metadata, ) diff --git a/specforge/training/backend.py b/specforge/training/backend.py index 853c08730..d98a04cd3 100644 --- a/specforge/training/backend.py +++ b/specforge/training/backend.py @@ -201,6 +201,14 @@ def prepare_model( Replicated ``NO_SHARD`` recipes use DDP; sharded recipes use FSDP. """ + if self._optimizer_factory is not None: + capture_metadata = getattr( + self._optimizer_factory, "capture_parameter_metadata", None + ) + if callable(capture_metadata): + capture_metadata( + optimizer_target if optimizer_target is not None else model + ) if not wrap: self.module = model self._wrapped = False diff --git a/specforge/training/controller.py b/specforge/training/controller.py index de6fcb29e..e4024e935 100644 --- a/specforge/training/controller.py +++ b/specforge/training/controller.py @@ -596,6 +596,10 @@ def _fit(self, data: Iterable[TrainBatch], progress: Optional[Any]) -> int: get_learning_rate = getattr(optimizer, "get_learning_rate", None) if callable(get_learning_rate): log_metrics["lr"] = float(get_learning_rate()) + get_learning_rates = getattr(optimizer, "get_learning_rates", None) + if callable(get_learning_rates): + for name, learning_rate in get_learning_rates().items(): + log_metrics[f"lr_{name}"] = float(learning_rate) self.logger(log_metrics, self.global_step) eval_metrics: Optional[Dict[str, Any]] = None if eval_enabled and self.global_step % self.eval_interval == 0: diff --git a/tests/test_config/test_schema.py b/tests/test_config/test_schema.py index 3e4d21c73..606b88f64 100644 --- a/tests/test_config/test_schema.py +++ b/tests/test_config/test_schema.py @@ -105,6 +105,20 @@ def test_fsdp_sharding_is_typed(self): with self.assertRaises(ValidationError): Config.model_validate(payload) + def test_muon_optimizer_is_typed_and_rejects_cpu_offload(self): + payload = copy.deepcopy(MINIMAL) + payload["training"] = { + "optimizer": "muon", + "muon_learning_rate": 0.02, + } + training = Config.model_validate(payload).training + self.assertEqual(training.optimizer, "muon") + self.assertEqual(training.muon_learning_rate, 0.02) + + payload["training"]["optimizer_cpu_offload"] = True + with self.assertRaisesRegex(ValidationError, "not supported with Muon"): + Config.model_validate(payload) + def test_removed_vlm_pixel_knobs_are_rejected(self): for field in ("min_pixels", "max_pixels"): payload = copy.deepcopy(MINIMAL) diff --git a/tests/test_optimizer/test_bf16_optimizer_clip_grad_norm.py b/tests/test_optimizer/test_bf16_optimizer_clip_grad_norm.py index 377ae7687..bcafe10a7 100644 --- a/tests/test_optimizer/test_bf16_optimizer_clip_grad_norm.py +++ b/tests/test_optimizer/test_bf16_optimizer_clip_grad_norm.py @@ -71,6 +71,34 @@ def configure_grad_norm_reduction(self, **kwargs): backend.set_optimizer(replicated) self.assertFalse(replicated.config["enabled"]) + def test_backend_captures_optimizer_metadata_before_factory_call(self): + events = [] + + class RecordingOptimizer: + def configure_grad_norm_reduction(self, **_kwargs): + pass + + class RecordingFactory: + def capture_parameter_metadata(self, target): + events.append(("capture", target.weight.shape)) + + def __call__(self, target): + events.append(("create", target.weight.shape)) + return RecordingOptimizer() + + model = torch.nn.Linear(8, 8, bias=False) + backend = FSDPTrainingBackend( + ParallelConfig(sharding_strategy="NO_SHARD"), + optimizer_factory=RecordingFactory(), + ) + + backend.prepare_model(model, wrap=False, optimizer_target=model) + + self.assertEqual( + events, + [("capture", torch.Size([8, 8])), ("create", torch.Size([8, 8]))], + ) + def test_cpu_offload_matches_resident_optimizer_update(self): resident_model, resident = _make_optimizer(seed=7, offload_master=False) offload_model, offload = _make_optimizer(seed=7, offload_master=True) diff --git a/tests/test_optimizer/test_muon_optimizer.py b/tests/test_optimizer/test_muon_optimizer.py new file mode 100644 index 000000000..ae28ea632 --- /dev/null +++ b/tests/test_optimizer/test_muon_optimizer.py @@ -0,0 +1,325 @@ +import copy +import os +import tempfile +import unittest +from unittest.mock import patch + +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +import torch.nn as nn +from transformers import Qwen3Config + +from specforge.modeling.draft.dspark import DSparkDraftModel +from specforge.muon import ( + ADAMW_OPTIMIZER, + MUON_OPTIMIZER, + capture_muon_parameter_metadata, + partition_parameters_for_muon, +) +from specforge.optimizer import BF16Optimizer + + +class _TinyMarkovHead(nn.Module): + def __init__(self) -> None: + super().__init__() + self.markov_w1 = nn.Embedding(8, 3) + self.markov_w2 = nn.Linear(3, 8, bias=False) + + def forward(self, token_ids: torch.Tensor) -> torch.Tensor: + return self.markov_w2(self.markov_w1(token_ids)) + + +class _TinyDraft(nn.Module): + def __init__(self) -> None: + super().__init__() + self.fc = nn.Linear(6, 4, bias=False) + self.layers = nn.ModuleList([nn.Linear(4, 4)]) + self.norm = nn.LayerNorm(4) + self.embed_proj = nn.Sequential(nn.Linear(4, 4, bias=False)) + self.markov_head = _TinyMarkovHead() + self.confidence_head = nn.Linear(4, 1) + self.lm_head = nn.Linear(4, 8, bias=False) + self.frozen_projection = nn.Linear(4, 4, bias=False) + self.frozen_projection.requires_grad_(False) + + def forward(self, features: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor: + hidden = self.norm(self.layers[0](self.fc(features))) + return ( + self.lm_head(hidden).square().mean() + + self.confidence_head(hidden).square().mean() + + self.markov_head(token_ids).square().mean() + ) + + +class TestMuonParameterPartition(unittest.TestCase): + def test_dspark_backbone_and_heads_are_partitioned_as_intended(self): + config = Qwen3Config( + vocab_size=32, + hidden_size=16, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=4, + max_position_embeddings=128, + block_size=4, + num_target_layers=4, + dflash_config={ + "attention_mode": "gqa", + "projector_type": "dspark", + "target_layer_ids": [0], + "markov_rank": 4, + "markov_head_type": "vanilla", + "enable_confidence_head": True, + "confidence_head_with_markov": True, + }, + ) + + partition = partition_parameters_for_muon(DSparkDraftModel(config)) + + muon_names = {item.name for item in partition.muon} + adamw_names = {item.name for item in partition.adamw} + self.assertIn("fc.weight", muon_names) + self.assertIn("layers.0.self_attn.q_proj.weight", muon_names) + self.assertIn("layers.0.mlp.down_proj.weight", muon_names) + self.assertIn("markov_head.markov_w2.weight", adamw_names) + self.assertIn("confidence_head.proj.weight", adamw_names) + + def test_only_hidden_linear_weights_use_muon(self): + partition = partition_parameters_for_muon(_TinyDraft()) + + self.assertEqual( + tuple(item.name for item in partition.muon), + ("fc.weight", "layers.0.weight"), + ) + adamw_names = {item.name for item in partition.adamw} + self.assertIn("layers.0.bias", adamw_names) + self.assertIn("norm.weight", adamw_names) + self.assertIn("embed_proj.0.weight", adamw_names) + self.assertIn("markov_head.markov_w1.weight", adamw_names) + self.assertIn("markov_head.markov_w2.weight", adamw_names) + self.assertIn("confidence_head.weight", adamw_names) + self.assertIn("lm_head.weight", adamw_names) + self.assertNotIn("frozen_projection.weight", adamw_names) + + def test_pre_fsdp_metadata_preserves_logical_matrix_shape(self): + model = _TinyDraft() + metadata = capture_muon_parameter_metadata(model) + logical_shape = model.fc.weight.shape + model.fc.weight.data = model.fc.weight.data.reshape(-1) + + partition = partition_parameters_for_muon(model, metadata=metadata) + + fc_weight = next(item for item in partition.muon if item.name == "fc.weight") + self.assertEqual(fc_weight.parameter.ndim, 1) + self.assertEqual(fc_weight.logical_shape, logical_shape) + + def test_muon_rejects_a_model_without_hidden_matrices(self): + model = nn.Sequential(nn.Embedding(8, 4), nn.LayerNorm(4)) + with self.assertRaisesRegex(ValueError, "no eligible hidden"): + BF16Optimizer(model, lr=1e-3, optimizer_type=MUON_OPTIMIZER) + + +class TestBF16MuonOptimizer(unittest.TestCase): + @staticmethod + def _loss(model: nn.Module) -> torch.Tensor: + features = torch.arange(12, dtype=torch.float32).reshape(2, 6) / 10 + token_ids = torch.tensor([1, 3]) + return model(features, token_ids) + + def test_hybrid_step_updates_both_groups_and_round_trips_state(self): + torch.manual_seed(0) + model = _TinyDraft() + optimizer = BF16Optimizer( + model, + lr=1e-3, + weight_decay=0.0, + optimizer_type=MUON_OPTIMIZER, + muon_lr=2e-3, + muon_weight_decay=0.0, + max_grad_norm=1.0, + warmup_ratio=0.2, + total_steps=10, + ) + before = { + name: parameter.detach().clone() + for name, parameter in model.named_parameters() + } + + self._loss(model).backward() + grad_norm = optimizer.step() + + self.assertTrue(torch.isfinite(grad_norm)) + self.assertFalse(torch.equal(before["fc.weight"], model.fc.weight)) + self.assertFalse(torch.equal(before["lm_head.weight"], model.lm_head.weight)) + self.assertTrue(all(parameter.grad is None for parameter in model.parameters())) + self.assertTrue(optimizer.optimizer.state) + self.assertTrue(optimizer.aux_optimizer.state) + + saved_state = copy.deepcopy(optimizer.state_dict()) + restored_model = copy.deepcopy(model) + restored = BF16Optimizer( + restored_model, + lr=1e-3, + weight_decay=0.0, + optimizer_type=MUON_OPTIMIZER, + muon_lr=2e-3, + muon_weight_decay=0.0, + max_grad_norm=1.0, + warmup_ratio=0.2, + total_steps=10, + ) + with patch("specforge.optimizer.print_on_rank0"): + restored.load_state_dict(saved_state) + + self.assertEqual(restored.get_learning_rates(), optimizer.get_learning_rates()) + for restored_master, saved_master in zip( + restored.fp32_params, saved_state["fp32_params"] + ): + torch.testing.assert_close(restored_master, saved_master) + + self._loss(model).backward() + self._loss(restored_model).backward() + optimizer.step() + restored.step() + for expected, actual in zip(model.parameters(), restored_model.parameters()): + torch.testing.assert_close(actual, expected) + + def test_group_summary_is_complete_and_non_overlapping(self): + model = _TinyDraft() + optimizer = BF16Optimizer( + model, + lr=1e-3, + optimizer_type=MUON_OPTIMIZER, + muon_weight_decay=0.0, + warmup_ratio=0.0, + total_steps=10, + ) + + summary = optimizer.get_parameter_group_summary() + muon_names = set(summary[MUON_OPTIMIZER]["names"]) + adamw_names = set(summary[ADAMW_OPTIMIZER]["names"]) + trainable_names = { + name + for name, parameter in model.named_parameters() + if parameter.requires_grad + } + self.assertFalse(muon_names & adamw_names) + self.assertEqual(muon_names | adamw_names, trainable_names) + + def test_muon_rejects_cpu_master_offload(self): + with self.assertRaisesRegex(ValueError, "does not support optimizer CPU"): + BF16Optimizer( + _TinyDraft(), + lr=1e-3, + optimizer_type=MUON_OPTIMIZER, + offload_master=True, + ) + + def test_adamw_checkpoint_schema_remains_backward_compatible(self): + optimizer = BF16Optimizer( + _TinyDraft(), + lr=1e-3, + optimizer_type=ADAMW_OPTIMIZER, + warmup_ratio=0.0, + total_steps=10, + ) + state = optimizer.state_dict() + + self.assertEqual( + set(state), + { + "optimizer_state_dict", + "scheduler_state_dict", + "max_grad_norm", + "fp32_params", + }, + ) + + +def _run_sharded_muon_parity(rank: int, world_size: int, init_file: str) -> None: + dist.init_process_group( + "gloo", + init_method=f"file://{init_file}", + rank=rank, + world_size=world_size, + ) + try: + full_parameter = torch.arange(12, dtype=torch.float32).reshape(4, 3) / 10 + full_gradient = torch.linspace(-0.5, 0.6, 12).reshape(4, 3) + shard_sizes = (7, 5, 0) + offset = sum(shard_sizes[:rank]) + shard_size = shard_sizes[rank] + + model = nn.Linear(3, 4, bias=False) + model.weight.data.copy_(full_parameter) + metadata = capture_muon_parameter_metadata(model) + model.weight.data = full_parameter.reshape(-1)[ + offset : offset + shard_size + ].clone() + optimizer = BF16Optimizer( + model, + lr=2e-3, + weight_decay=0.0, + optimizer_type=MUON_OPTIMIZER, + muon_weight_decay=0.1, + muon_momentum=0.95, + muon_nesterov=True, + muon_ns_steps=5, + muon_adjust_lr_fn="match_rms_adamw", + muon_metadata=metadata, + max_grad_norm=1e9, + warmup_ratio=0.0, + total_steps=10, + ) + optimizer.configure_grad_norm_reduction(enabled=True) + + expected_parameter = nn.Parameter(full_parameter.clone()) + expected_optimizer = torch.optim.Muon( + [expected_parameter], + lr=optimizer.get_learning_rate(), + weight_decay=0.1, + momentum=0.95, + nesterov=True, + ns_steps=5, + adjust_lr_fn="match_rms_adamw", + ) + expected_parameter.grad = full_gradient.clone() + expected_optimizer.step() + + if shard_size: + model.weight.grad = full_gradient.reshape(-1)[ + offset : offset + shard_size + ].clone() + optimizer.step() + + expected_local = expected_parameter.detach().reshape(-1)[ + offset : offset + shard_size + ] + torch.testing.assert_close(model.weight, expected_local) + expected_momentum = expected_optimizer.state[expected_parameter][ + "momentum_buffer" + ].reshape(-1)[offset : offset + shard_size] + torch.testing.assert_close( + optimizer.optimizer.state[optimizer.fp32_params[0]]["momentum_buffer"], + expected_momentum, + ) + finally: + dist.destroy_process_group() + + +class TestFSDPShardedMuon(unittest.TestCase): + def test_sharded_update_matches_native_full_matrix_muon(self): + with tempfile.TemporaryDirectory() as temporary_directory: + init_file = os.path.join(temporary_directory, "process-group") + mp.spawn( + _run_sharded_muon_parity, + args=(3, init_file), + nprocs=3, + join=True, + ) + + +if __name__ == "__main__": + unittest.main(verbosity=2)