diff --git a/.github/workflows/test-areal.yml b/.github/workflows/test-areal.yml index 7523d3deb4..206b63c64b 100644 --- a/.github/workflows/test-areal.yml +++ b/.github/workflows/test-areal.yml @@ -307,7 +307,14 @@ jobs: steps: - uses: actions/checkout@v6 + - name: Install Muon dependency in the prebuilt environment + run: | + export PATH="/opt/.venv/bin:$PATH" + uv pip install --python /opt/.venv/bin/python --target "$RUNNER_TEMP/muon-deps" --no-deps "emerging-optimizers==0.3.0" + - name: Validate Docker installation + env: + PYTHONPATH: ${{ github.workspace }}:${{ runner.temp }}/muon-deps run: | export PATH="/opt/.venv/bin:$PATH" python areal/tools/validate_docker_installation.py @@ -316,7 +323,7 @@ jobs: env: CI: true HF_TOKEN: ${{ secrets.HF_TOKEN }} - PYTHONPATH: ${{ github.workspace }} + PYTHONPATH: ${{ github.workspace }}:${{ runner.temp }}/muon-deps TOKENIZERS_PARALLELISM: false AREAL_IS_IN_CI: 1 VIRTUAL_ENV: /opt/.venv diff --git a/areal/api/cli_args.py b/areal/api/cli_args.py index c0d9fae750..f96196a38a 100644 --- a/areal/api/cli_args.py +++ b/areal/api/cli_args.py @@ -388,6 +388,55 @@ def to_openai_args_dict( # Train Engine Configs +@dataclass +class MuonOptimizerConfig: + """Algorithm settings for Megatron's distributed Muon optimizer.""" + + momentum: float = 0.95 + use_nesterov: bool = False + fp32_matmul_prec: str = field( + default="medium", + metadata={"choices": ["low", "medium", "high"]}, + ) + coefficient_type: str = "quintic" + num_ns_steps: int = 5 + scale_mode: str = field( + default="spectral", + metadata={"choices": ["spectral", "unit_rms_norm", "shape_scaling"]}, + ) + split_qkv: bool = True + tp_mode: str = field( + default="duplicated", + metadata={"choices": ["blockwise", "duplicated", "distributed"]}, + ) + extra_scale_factor: float = 1.0 + + def __post_init__(self) -> None: + if not 0.0 <= self.momentum < 1.0: + raise ValueError("optimizer.muon.momentum must be in [0, 1)") + if self.fp32_matmul_prec not in {"low", "medium", "high"}: + raise ValueError( + "optimizer.muon.fp32_matmul_prec must be low, medium, or high" + ) + if self.num_ns_steps < 1: + raise ValueError("optimizer.muon.num_ns_steps must be at least 1") + if self.scale_mode not in { + "spectral", + "unit_rms_norm", + "shape_scaling", + }: + raise ValueError( + "optimizer.muon.scale_mode must be spectral, unit_rms_norm, " + "or shape_scaling" + ) + if self.tp_mode not in {"blockwise", "duplicated", "distributed"}: + raise ValueError( + "optimizer.muon.tp_mode must be blockwise, duplicated, or distributed" + ) + if self.extra_scale_factor <= 0: + raise ValueError("optimizer.muon.extra_scale_factor must be positive") + + @dataclass class OptimizerConfig: """Configuration for model optimization during training.""" @@ -399,8 +448,9 @@ class OptimizerConfig: metadata={ "help": "Optimizer type. For FSDP Engine, adam_bf16 enables memory-efficient BF16 optimizer states. " "For Megatron Engine, adam_bf16 requires dtype=bfloat16 and is automatically converted to adam " - "with precision-aware optimizer enabled.", - "choices": ["adam", "sgd", "adam_bf16"], + "with precision-aware optimizer enabled. dist_muon selects Megatron's layer-wise distributed " + "Muon optimizer and uses AdamW for scalar/non-matrix parameters.", + "choices": ["adam", "sgd", "adam_bf16", "dist_muon"], }, ) lr: float = field(default=1e-3, metadata={"help": "Learning rate"}) @@ -423,6 +473,10 @@ class OptimizerConfig: "help": "Adam epsilon parameter. Only effective when optimizer_type is adam/adam_bf16" }, ) + muon: MuonOptimizerConfig = field( + default_factory=MuonOptimizerConfig, + metadata={"help": "Muon algorithm settings used when type=dist_muon."}, + ) min_lr_ratio: float = field( default=0.0, metadata={ diff --git a/areal/engine/megatron_engine.py b/areal/engine/megatron_engine.py index 33476f8ce7..6eab6d4e5c 100644 --- a/areal/engine/megatron_engine.py +++ b/areal/engine/megatron_engine.py @@ -85,6 +85,11 @@ from areal.engine.megatron_utils.checkpointer import MegatronCheckpointManager from areal.engine.megatron_utils.deterministic import set_deterministic_algorithms from areal.engine.megatron_utils.fp8 import FP8BlockwiseTensorHelper +from areal.engine.megatron_utils.gpu_staged_muon import ( + GPUStagedMuonConfig, + get_megatron_optimizer_with_dist_muon, + get_megatron_optimizer_with_gpu_staged_muon, +) from areal.engine.megatron_utils.gpu_staged_optimizer import ( GPUStagedAdamWConfig, get_megatron_optimizer_with_gpu_staged_adamw, @@ -2931,26 +2936,59 @@ def _create_optimizer(self, ft_spec: FinetuneSpec) -> None: staged_settings = self.mcore_config.cpu_staged_offload staged = staged_settings.enabled optimizer_type = self.optimizer_config.type - if staged and optimizer_type != "adam": + use_dist_muon = optimizer_type == "dist_muon" + + if staged and optimizer_type not in {"adam", "dist_muon"}: raise ValueError( - "CPU-staged offload supports optimizer.type='adam', " - f"got {optimizer_type!r}" + "CPU-staged offload supports optimizer.type='adam' or " + f"'dist_muon', got {optimizer_type!r}" ) - if staged and not self.mcore_config.ddp.use_distributed_optimizer: + if use_dist_muon and staged and self.config.use_lora: + raise ValueError("CPU-staged Muon does not support LoRA") + if use_dist_muon: + if self.mcore_config.ddp.use_distributed_optimizer: + raise ValueError( + "optimizer.type='dist_muon' requires " + "megatron.ddp.use_distributed_optimizer=false" + ) + if self.mcore_config.use_precision_aware_optimizer: + raise ValueError( + "optimizer.type='dist_muon' does not support " + "megatron.use_precision_aware_optimizer=true" + ) + if self.mcore_config.overlap_param_gather_with_optimizer_step: + raise ValueError( + "optimizer.type='dist_muon' does not support " + "overlap_param_gather_with_optimizer_step" + ) + if staged and self.mcore_config.ddp.overlap_param_gather: + raise ValueError( + "CPU-staged dist_muon requires " + "megatron.ddp.overlap_param_gather=false" + ) + if staged and self.optimizer_config.muon.tp_mode != "duplicated": + raise ValueError( + "CPU-staged dist_muon currently supports only " + "optimizer.muon.tp_mode='duplicated'" + ) + if self.dtype is not torch.bfloat16: + raise ValueError("optimizer.type='dist_muon' requires dtype=bfloat16") + elif staged and not self.mcore_config.ddp.use_distributed_optimizer: raise ValueError( "CPU-staged AdamW requires megatron.ddp.use_distributed_optimizer=true" ) use_distributed_optimizer = ( False - if self.config.use_lora + if self.config.use_lora or use_dist_muon else self.mcore_config.ddp.use_distributed_optimizer ) assert self.optimizer_config.type in [ "adam", "sgd", - ], "Only AdamW/SGD optimizer is supported in this engine." + "dist_muon", + ], "Only AdamW/SGD/distributed Muon optimizer is supported in this engine." if self.optimizer_config.type == "sgd": self.logger.warning( "Using the 'sgd' optimizer with Megatron may be less stable. Consider using the 'adam' (AdamW) optimizer for improved stability." @@ -2983,9 +3021,17 @@ def _create_optimizer(self, ft_spec: FinetuneSpec) -> None: if staged and optimizer_type == "adam" else None ) + staged_muon_config = ( + GPUStagedMuonConfig( + buffer_count=staged_settings.buffer_count, + slot_size_mb=staged_settings.bucket_size_mb, + ) + if staged and use_dist_muon + else None + ) # Make megatron optimizer config mcore_opt_config = MCoreOptimizerConfig( - optimizer=self.optimizer_config.type, + optimizer="muon" if use_dist_muon else self.optimizer_config.type, lr=self.optimizer_config.lr, min_lr=self.optimizer_config.min_lr_ratio * self.optimizer_config.lr, weight_decay=self.optimizer_config.weight_decay, @@ -2995,17 +3041,25 @@ def _create_optimizer(self, ft_spec: FinetuneSpec) -> None: adam_beta2=self.optimizer_config.beta2, adam_eps=self.optimizer_config.eps, use_distributed_optimizer=use_distributed_optimizer, + use_layer_wise_distributed_optimizer=use_dist_muon, params_dtype=self.dtype, clip_grad=self.optimizer_config.gradient_clipping, fp8_recipe=(self.fp8_config.recipe if self.enable_fp8 else None), overlap_param_gather_with_optimizer_step=( self.mcore_config.overlap_param_gather_with_optimizer_step ), + overlap_param_gather=self.mcore_config.ddp.overlap_param_gather, use_precision_aware_optimizer=( - staged_adamw_config is not None + False + if use_dist_muon + else staged_adamw_config is not None or self.mcore_config.use_precision_aware_optimizer ), - main_grads_dtype=getattr(torch, self.mcore_config.main_grads_dtype), + main_grads_dtype=( + torch.float32 + if use_dist_muon + else getattr(torch, self.mcore_config.main_grads_dtype) + ), main_params_dtype=( torch.float32 if staged_adamw_config is not None @@ -3021,6 +3075,16 @@ def _create_optimizer(self, ft_spec: FinetuneSpec) -> None: if staged_adamw_config is not None else getattr(torch, self.mcore_config.exp_avg_sq_dtype) ), + muon_momentum=self.optimizer_config.muon.momentum, + muon_nesterov=self.optimizer_config.muon.use_nesterov, + muon_fp32_matmul_prec=self.optimizer_config.muon.fp32_matmul_prec, + muon_coefficient_type=self.optimizer_config.muon.coefficient_type, + muon_num_ns_steps=self.optimizer_config.muon.num_ns_steps, + muon_scale_mode=self.optimizer_config.muon.scale_mode, + muon_split_qkv=self.optimizer_config.muon.split_qkv, + muon_tp_mode=self.optimizer_config.muon.tp_mode, + muon_extra_scale_factor=self.optimizer_config.muon.extra_scale_factor, + muon_scalar_optimizer="adam", ) optimizer_kwargs = {} @@ -3035,7 +3099,16 @@ def _create_optimizer(self, ft_spec: FinetuneSpec) -> None: optimizer_kwargs["config_overrides"] = qwen4_exp_optimizer_overrides( mcore_opt_config ) - if staged_adamw_config is None: + if staged_muon_config is not None: + self.optimizer = get_megatron_optimizer_with_gpu_staged_muon( + mcore_opt_config, self.model, staged_muon_config, **optimizer_kwargs + ) + self.optimizer.bind_managed_checkpoint_process_group(self.cpu_group) + elif use_dist_muon: + self.optimizer = get_megatron_optimizer_with_dist_muon( + mcore_opt_config, self.model, **optimizer_kwargs + ) + elif staged_adamw_config is None: self.optimizer = get_megatron_optimizer( mcore_opt_config, self.model, **optimizer_kwargs ) diff --git a/areal/engine/megatron_utils/checkpointer.py b/areal/engine/megatron_utils/checkpointer.py index 71fa3a96ce..b7a2ea3219 100644 --- a/areal/engine/megatron_utils/checkpointer.py +++ b/areal/engine/megatron_utils/checkpointer.py @@ -297,9 +297,14 @@ def __init__( self.lr_scheduler = lr_scheduler self.use_distributed_optimizer = use_distributed_optimizer - assert self.use_distributed_optimizer, ( - "MegatronCheckpointManager now only support distributed optimizer" - ) + if not self.use_distributed_optimizer and not getattr( + optimizer, "supports_non_distributed_checkpoint", False + ): + raise NotImplementedError( + "MegatronCheckpointManager requires either Megatron's distributed " + "optimizer or an optimizer with explicit non-distributed checkpoint " + "support" + ) self.use_checkpoint_opt_param_scheduler = use_checkpoint_opt_param_scheduler self.rank = torch.distributed.get_rank() self.use_dist_checkpointing = use_dist_checkpointing diff --git a/areal/engine/megatron_utils/gpu_staged_muon.py b/areal/engine/megatron_utils/gpu_staged_muon.py new file mode 100644 index 0000000000..91705e9d46 --- /dev/null +++ b/areal/engine/megatron_utils/gpu_staged_muon.py @@ -0,0 +1,3728 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""MCore 0.19 layer-wise Muon with CPU-authoritative optimizer state. + +The official Muon builder remains responsible for parameter classification and +``LayerWiseDistributedOptimizer.shard_params`` ownership. Only after that +ownership decision do we replace each rank-local leaf with a staged backend. +Muon matrices are indivisible units; scalar parameters retain the staged AdamW +backend. Synchronous checkpointing permits DP/expert-DP ownership resharding +while keeping every model-parallel dimension fixed. Async checkpoint and +prefetch are deliberately outside this MVP. +""" + +from __future__ import annotations + +import copy +import hashlib +import importlib.metadata +import itertools +import json +import math +from collections.abc import Callable, Iterable, Mapping, Sequence +from dataclasses import dataclass +from typing import Any + +import torch +from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors + +from areal.engine.megatron_utils.gpu_staged_optimizer import ( + GPUStagedAdamW, + GPUStagedAdamWConfig, +) +from areal.engine.megatron_utils.staged_optimizer_runtime import ( + CUDAStagingSlot, + SlotStateMachine, + StagedOptimizerRuntime, +) + +_SUPPORTED_MEGATRON_CORE_VERSION = "0.19.0" +_SUPPORTED_EMERGING_OPTIMIZERS_VERSION = "0.3.0" +_EMPTY_MUON_LEAF = object() +_MUON_CHECKPOINT_SCHEMA_VERSION = 2 +_MUON_CHECKPOINT_PREFIX = "optimizer.gpu_staged_muon.v2" +_MUON_CHECKPOINT_STATE_KINDS = { + "muon": ("master_param", "momentum_buffer"), + "scalar_adamw": ("master_param", "exp_avg", "exp_avg_sq"), +} + + +def _require_muon_checkpoint_versions() -> tuple[str, str]: + """Fail closed at every Muon checkpoint entry on the pinned private APIs.""" + try: + mcore_version = importlib.metadata.version("megatron-core") + eo_version = importlib.metadata.version("emerging-optimizers") + except importlib.metadata.PackageNotFoundError as error: + raise RuntimeError( + "Muon checkpoint requires megatron-core 0.19.0 and " + "emerging-optimizers 0.3.0" + ) from error + if mcore_version != _SUPPORTED_MEGATRON_CORE_VERSION: + raise RuntimeError( + f"Muon checkpoint requires megatron-core 0.19.0, found {mcore_version}" + ) + if eo_version != _SUPPORTED_EMERGING_OPTIMIZERS_VERSION: + raise RuntimeError( + f"Muon checkpoint requires emerging-optimizers 0.3.0, found {eo_version}" + ) + return mcore_version, eo_version + + +@dataclass(frozen=True) +class GPUStagedMuonConfig: + """GPU staging capacity for the CPU-resident Muon backend.""" + + buffer_count: int = 2 + slot_size_mb: float = 128.0 + + def __post_init__(self) -> None: + if self.buffer_count < 1: + raise ValueError("Muon staging buffer_count must be at least 1") + if self.slot_size_mb <= 0: + raise ValueError("Muon staging slot_size_mb must be positive") + + @property + def slot_numel(self) -> int: + return max(1, int(self.slot_size_mb * 1024 * 1024) // 4) + + +@dataclass +class MuonCPUSlabs: + """Pinned FP32 master and momentum slabs authoritative between steps.""" + + master: torch.Tensor + momentum: torch.Tensor + + @classmethod + def allocate(cls, numel: int) -> MuonCPUSlabs: + if numel < 0: + raise ValueError("Muon slab numel must be non-negative") + kwargs = {"dtype": torch.float32, "device": "cpu", "pin_memory": True} + return cls( + master=torch.empty(numel, **kwargs), + momentum=torch.zeros(numel, **kwargs), + ) + + +@dataclass(frozen=True) +class MuonOwnedUnit: + """One complete owner-held matrix; Newton--Schulz never sees fragments.""" + + param: torch.nn.Parameter + group_index: int + slab_offset: int + numel: int + + +@dataclass(frozen=True) +class _FrozenOwnerParameter: + """Immutable local authority for one layer-wise all-gather parameter.""" + + param: torch.Tensor + domain: str + owner_rank: int + ordinal: int + shape: tuple[int, ...] + numel: int + stride: tuple[int, ...] + layout: torch.layout + dtype: torch.dtype + device_type: str + device_index: int | None + storage: Any + storage_cdata: int + storage_nbytes: int + storage_offset: int + + +def _freeze_owner_schema( + domain: str, + owner_lists: Sequence[Sequence[torch.Tensor]] | None, +) -> tuple[tuple[_FrozenOwnerParameter, ...], ...] | None: + """Capture owner ordering and tensor/storage metadata at bind time.""" + if owner_lists is None: + return None + frozen: list[tuple[_FrozenOwnerParameter, ...]] = [] + ordinal = 0 + for owner_rank, params in enumerate(owner_lists): + owner: list[_FrozenOwnerParameter] = [] + for param in params: + storage = param.untyped_storage() + owner.append( + _FrozenOwnerParameter( + param=param, + domain=domain, + owner_rank=owner_rank, + ordinal=ordinal, + shape=tuple(param.shape), + numel=param.numel(), + stride=tuple(param.stride()), + layout=param.layout, + dtype=param.dtype, + device_type=param.device.type, + device_index=param.device.index, + storage=storage, + storage_cdata=int(storage._cdata), + storage_nbytes=int(storage.nbytes()), + storage_offset=param.storage_offset(), + ) + ) + ordinal += 1 + frozen.append(tuple(owner)) + return tuple(frozen) + + +def _owner_schema_digest( + domain: str, + schema: tuple[tuple[_FrozenOwnerParameter, ...], ...] | None, +) -> tuple[int, int, int, int]: + """Return a rank-stable structural digest; local storage identity is separate.""" + structural = ( + domain, + None + if schema is None + else tuple( + tuple( + ( + entry.owner_rank, + entry.ordinal, + entry.shape, + entry.numel, + entry.stride, + str(entry.layout), + str(entry.dtype), + entry.device_type, + ) + for entry in owner + ) + for owner in schema + ), + ) + digest = hashlib.sha256(repr(structural).encode()).digest() + words = tuple( + int.from_bytes(digest[offset : offset + 8], "little") & ((1 << 63) - 1) + for offset in range(0, 32, 8) + ) + if len(words) != 4: + raise AssertionError("SHA-256 owner schema digest must contain four words") + return words[0], words[1], words[2], words[3] + + +def _preflight_tp_gradient_participation( + units: Sequence[MuonOwnedUnit], + tp_groups: Mapping[str, Any] | None, +) -> tuple[bool, str | None]: + """Run fixed-size raw-gradient votes in both official TP domains.""" + if tp_groups is None: + return ( + any( + getattr(unit.param, "main_grad", None) is not None + or unit.param.grad is not None + for unit in units + ), + None, + ) + if not torch.distributed.is_initialized(): + raise RuntimeError( + "staged Muon TP groups are bound but torch.distributed is not initialized" + ) + + any_active = False + errors: list[str] = [] + for domain in ("dense", "expert"): + group = tp_groups[domain] + group_name = "expt_tp" if domain == "expert" else "tp" + group_size, group_rank = _group_size_and_rank(group, name=group_name) + params = tuple( + unit.param + for unit in units + if ("expert" if getattr(unit.param, "expert_tp", False) else "dense") + == domain + ) + local_presence = tuple( + getattr(param, "main_grad", None) is not None or param.grad is not None + for param in params + ) + local_active = any(local_presence) + any_active |= local_active + if group_size == 1: + continue + + digest = hashlib.sha256(repr(local_presence).encode()).digest() + digest_words = [ + int.from_bytes(digest[offset : offset + 8], "little") & ((1 << 63) - 1) + for offset in range(0, 32, 8) + ] + vote = torch.tensor( + [len(params), int(local_active), *digest_words], + dtype=torch.int64, + device=params[0].device + if params + else torch.device("cuda", torch.cuda.current_device()), + ) + minimum = vote.clone() + maximum = vote.clone() + torch.distributed.all_reduce( + minimum, op=torch.distributed.ReduceOp.MIN, group=group + ) + torch.distributed.all_reduce( + maximum, op=torch.distributed.ReduceOp.MAX, group=group + ) + if not torch.equal(minimum, maximum): + errors.append( + "Muon TP peers have inconsistent gradient participation: " + f"domain={domain}, group={group_name}, group_rank={group_rank}, " + f"local={local_presence}" + ) + return any_active, "; ".join(errors) if errors else None + + +class GPUStagedEmptyOptimizer: + """State-free leaf used when one official partition owns no parameters.""" + + manages_cpu_residency = True + manages_master_weight = True + + def __init__(self, optimizer_kind: str) -> None: + if optimizer_kind not in {"muon", "scalar_adamw"}: + raise ValueError(f"unsupported empty optimizer kind: {optimizer_kind!r}") + self.optimizer_kind = optimizer_kind + self.param_groups: list[dict[str, Any]] = [] + self.state: dict[Any, Any] = {} + self._checkpoint_load_error: BaseException | None = None + + @property + def residency(self) -> str: + return "CPU_RESIDENT" + + @property + def cuda_state_numel(self) -> int: + return 0 + + @property + def gpu_staging_state_numel(self) -> int: + return 0 + + @property + def gpu_staging_numel(self) -> int: + return 0 + + def step(self, closure: Callable[[], float] | None = None) -> float | None: + if self._checkpoint_load_error is not None: + raise RuntimeError( + "cannot step after a failed checkpoint load; restart and recover" + ) from self._checkpoint_load_error + return closure() if closure is not None else None + + def bind_parallel_groups(self, *, tp: Any, expt_tp: Any) -> None: + if tp is None or expt_tp is None: + raise RuntimeError("staged Muon requires explicit tp and expt_tp groups") + + def preflight_step_activity(self) -> tuple[bool, str | None]: + return False, None + + def drain(self) -> None: + return None + + def offload_to_cpu(self) -> None: + return None + + def restore_from_cpu(self) -> None: + return None + + def prepare_checkpoint_save(self) -> None: + if self._checkpoint_load_error is not None: + raise RuntimeError( + "cannot save after a failed checkpoint load; restart and recover" + ) from self._checkpoint_load_error + + def begin_checkpoint_load(self) -> None: + if self._checkpoint_load_error is not None: + raise RuntimeError( + "checkpoint load already failed; restart the process to recover" + ) from self._checkpoint_load_error + + def complete_checkpoint_load(self) -> None: + return None + + def mark_checkpoint_load_failed(self, error: BaseException) -> None: + self._checkpoint_load_error = error + + def reset_from_model_params(self) -> None: + return None + + def state_dict(self) -> dict[str, Any]: + self.prepare_checkpoint_save() + return {"state": {}, "param_groups": []} + + def load_state_dict(self, state_dict: Mapping[str, Any]) -> None: + if not isinstance(state_dict, Mapping) or state_dict != { + "state": {}, + "param_groups": [], + }: + raise ValueError("empty staged optimizer checkpoint must contain no state") + + +class GPUStagedMuon(torch.optim.Optimizer): + """Muon shell whose master and momentum live in pinned CPU slabs.""" + + manages_cpu_residency = True + manages_master_weight = True + optimizer_kind = "muon" + + def __init__( + self, + params: Iterable[torch.Tensor] | Iterable[dict[str, Any]], + *, + staged_config: GPUStagedMuonConfig, + weight_decay_method: str, + native_optimizer: torch.optim.Optimizer, + ) -> None: + if weight_decay_method != "decoupled": + raise ValueError("staged Muon MVP requires decoupled weight decay") + super().__init__(params, defaults={}) + self.staged_config = staged_config + self._native_template = native_optimizer + self._native_unit_optimizers: dict[int, torch.optim.Optimizer] = {} + self.cpu_slabs: MuonCPUSlabs | None = None + self._units: tuple[MuonOwnedUnit, ...] = () + self._runtime = StagedOptimizerRuntime( + self.staged_config.buffer_count, + ("master", "momentum", "grad"), + ) + self._tp_groups: dict[str, Any] | None = None + self._bound = False + + @property + def residency(self) -> str: + return self._runtime.residency + + @property + def _slots(self) -> list[CUDAStagingSlot]: + return self._runtime.slots + + @property + def _slot_machine(self) -> SlotStateMachine | None: + return self._runtime.slot_machine + + @property + def units(self) -> tuple[MuonOwnedUnit, ...]: + return self._units + + @property + def cuda_state_numel(self) -> int: + return sum( + value.numel() + for state in self.state.values() + for value in state.values() + if isinstance(value, torch.Tensor) and value.is_cuda + ) + + @property + def gpu_staging_state_numel(self) -> int: + return sum(slot.master.numel() + slot.momentum.numel() for slot in self._slots) + + @property + def gpu_staging_numel(self) -> int: + """Total bounded slot storage, including gradients.""" + return sum( + slot.master.numel() + slot.momentum.numel() + slot.grad.numel() + for slot in self._slots + ) + + def bind_owned_params( + self, + param_groups: list[dict[str, Any]], + *, + empty_device: torch.device | None = None, + ) -> None: + """Bind only after official layer-wise ownership has been established.""" + if self._bound: + raise RuntimeError("GPU-staged Muon is already bound") + if len(param_groups) != len(self.param_groups): + raise ValueError("Muon bound groups do not match official groups") + self.param_groups = param_groups + + units: list[MuonOwnedUnit] = [] + devices: set[torch.device] = set() + offset = 0 + seen: set[int] = set() + for group_index, group in enumerate(param_groups): + for param in group["params"]: + if id(param) in seen: + raise ValueError("duplicate parameter in Muon owner groups") + seen.add(id(param)) + if param.ndim != 2: + raise ValueError( + "official Muon owner parameter must be a full 2D matrix" + ) + if not param.is_cuda: + raise ValueError("staged Muon parameters must be CUDA tensors") + if param.numel() > self.staged_config.slot_numel: + raise ValueError( + "Muon slot is too small for an indivisible owner matrix: " + f"matrix_numel={param.numel()}, " + f"slot_numel={self.staged_config.slot_numel}" + ) + devices.add(param.device) + units.append(MuonOwnedUnit(param, group_index, offset, param.numel())) + offset += param.numel() + if len(devices) > 1: + raise ValueError("one staged Muon leaf cannot own multiple CUDA devices") + if not devices and empty_device is None: + raise ValueError("empty Muon owner requires an explicit CUDA device") + device = next(iter(devices), empty_device) + assert device is not None + if device.type != "cuda": + raise ValueError("Muon staging device must be CUDA") + + self._units = tuple(units) + self.state.clear() + if not units: + self.cpu_slabs = None + self._runtime.bind(capacity=None, device=device) + self._bound = True + return + + self.cpu_slabs = MuonCPUSlabs.allocate(offset) + for unit in units: + state = self.state[unit.param] + state["master_param"] = self.cpu_slabs.master.narrow( + 0, unit.slab_offset, unit.numel + ).view_as(unit.param) + state["momentum_buffer"] = self.cpu_slabs.momentum.narrow( + 0, unit.slab_offset, unit.numel + ).view_as(unit.param) + + self._runtime.bind(capacity=self.staged_config.slot_numel, device=device) + self._bind_native_unit_optimizers() + self._bound = True + self._runtime.schedule_units( + self._units, + self._schedule_master_initialization, + wait_for_compute=False, + ) + self.drain() + + def bind_parallel_groups(self, *, tp: Any, expt_tp: Any) -> None: + """Bind the official TP groups used by MCore's Muon implementation.""" + if self._bound: + raise RuntimeError("Muon parallel groups must be bound before CPU state") + if tp is None or expt_tp is None: + raise RuntimeError("staged Muon requires explicit tp and expt_tp groups") + self._tp_groups = {"dense": tp, "expert": expt_tp} + + def preflight_step_activity(self) -> tuple[bool, str | None]: + """Validate raw TP gradient participation before MCore prepares gradients.""" + return _preflight_tp_gradient_participation(self._units, self._tp_groups) + + def _bind_native_unit_optimizers(self) -> None: + """Bind native Muon steps to persistent views of reusable slot buffers.""" + for unit_index, unit in enumerate(self._units): + slot = self._slots[unit_index % len(self._slots)] + proxy = slot.master[: unit.numel].view_as(unit.param).detach() + for name in ( + "expert_tp", + "partition_dim", + "partition_stride", + "tensor_model_parallel", + "is_qkv", + ): + if hasattr(unit.param, name): + setattr(proxy, name, getattr(unit.param, name)) + momentum = slot.momentum[: unit.numel].view_as(unit.param) + native = copy.copy(self._native_template) + # Optimizer.__getstate__ omits Muon's algorithm and topology attributes. + vars(native).update(vars(self._native_template)) + group = dict(self.param_groups[unit.group_index]) + group["params"] = [proxy] + native.param_groups = [group] + native.state = {proxy: {"momentum_buffer": momentum}} + self._native_unit_optimizers[id(unit.param)] = native + + def _schedule_master_initialization( + self, + unit: MuonOwnedUnit, + slot_index: int, + params_ready: torch.cuda.Event, + ) -> None: + assert self.cpu_slabs is not None + slot = self._runtime.acquire_slot(slot_index) + with torch.cuda.stream(slot.h2d_stream): + slot.h2d_stream.wait_event(params_ready) + slot.master[: unit.numel].copy_(unit.param.detach().view(-1)) + slot.h2d_done.record(slot.h2d_stream) + with torch.cuda.stream(slot.d2h_stream): + slot.d2h_stream.wait_event(slot.h2d_done) + self.cpu_slabs.master.narrow(0, unit.slab_offset, unit.numel).copy_( + slot.master[: unit.numel], non_blocking=True + ) + slot.d2h_done.record(slot.d2h_stream) + self._runtime.mark_d2h_pending(slot_index) + + def _schedule_update( + self, + unit: MuonOwnedUnit, + slot_index: int, + grads_ready: torch.cuda.Event, + ) -> None: + assert self.cpu_slabs is not None + slot = self._runtime.acquire_slot(slot_index) + slab_slice = slice(unit.slab_offset, unit.slab_offset + unit.numel) + with torch.cuda.stream(slot.h2d_stream): + slot.h2d_stream.wait_event(grads_ready) + slot.master[: unit.numel].copy_( + self.cpu_slabs.master[slab_slice], non_blocking=True + ) + slot.momentum[: unit.numel].copy_( + self.cpu_slabs.momentum[slab_slice], non_blocking=True + ) + slot.h2d_done.record(slot.h2d_stream) + + group = self.param_groups[unit.group_index] + with torch.cuda.stream(slot.compute_stream): + slot.compute_stream.wait_event(grads_ready) + slot.compute_stream.wait_event(slot.h2d_done) + grad = getattr(unit.param, "decoupled_grad", None) + if grad is None: + grad = unit.param.grad + if grad is not None: + slot.grad[: unit.numel].copy_(grad.detach().view(-1)) + grad_matrix = slot.grad[: unit.numel].view_as(unit.param) + native = self._native_unit_optimizers[id(unit.param)] + native_group = native.param_groups[0] + # Scheduler and checkpoint updates target the outer groups. + # Keep the slot's proxy parameter while refreshing settings. + native_group.update( + {key: value for key, value in group.items() if key != "params"} + ) + proxy = native_group["params"][0] + proxy.grad = grad_matrix + native.step() + unit.param.copy_(proxy) + slot.compute_done.record(slot.compute_stream) + + with torch.cuda.stream(slot.d2h_stream): + slot.d2h_stream.wait_event(slot.compute_done) + self.cpu_slabs.master[slab_slice].copy_( + slot.master[: unit.numel], non_blocking=True + ) + self.cpu_slabs.momentum[slab_slice].copy_( + slot.momentum[: unit.numel], non_blocking=True + ) + slot.d2h_done.record(slot.d2h_stream) + self._runtime.mark_d2h_pending(slot_index) + + @torch.no_grad() + def step(self, closure: Callable[[], float] | None = None) -> float | None: + if not self._bound: + raise RuntimeError("bind_owned_params() must be called before Muon step") + if self._runtime.checkpoint_load_error is not None: + raise RuntimeError( + "cannot step after a failed checkpoint load; restart and recover" + ) from self._runtime.checkpoint_load_error + if closure is not None: + raise ValueError("staged Muon does not support closures") + if not self._units: + self._runtime.residency = "CPU_RESIDENT" + return None + self._runtime.schedule_units( + self._units, self._schedule_update, wait_for_compute=True + ) + return None + + def drain(self) -> None: + self._runtime.drain() + + def offload_to_cpu(self) -> None: + self.drain() + self._native_unit_optimizers.clear() + self._runtime.release_slots() + + def restore_from_cpu(self) -> None: + if not self._bound or not self._units or self._slots: + return + device = self._units[0].param.device + self._runtime.restore_slots( + capacity=self.staged_config.slot_numel, device=device + ) + self._bind_native_unit_optimizers() + + def _validate_bound_state_views(self) -> None: + if not self._bound: + raise RuntimeError("GPU-staged Muon must be bound before checkpointing") + if not self._units: + if self.cpu_slabs is not None or self.state: + raise RuntimeError("empty Muon ownership unexpectedly has state") + return + if self.cpu_slabs is None: + raise RuntimeError("bound Muon optimizer has no CPU slabs") + expected = {"master_param", "momentum_buffer"} + for name, slab in ( + ("master", self.cpu_slabs.master), + ("momentum", self.cpu_slabs.momentum), + ): + if slab.device.type != "cpu" or slab.dtype is not torch.float32: + raise RuntimeError(f"Muon {name} slab must remain CPU FP32") + if not slab.is_pinned(): + raise RuntimeError(f"Muon {name} slab lost pinned residency") + for unit in self._units: + state = self.state.get(unit.param) + if not isinstance(state, Mapping) or set(state) != expected: + raise RuntimeError("Muon checkpoint state schema changed after bind") + for key, slab_name in ( + ("master_param", "master"), + ("momentum_buffer", "momentum"), + ): + value = state[key] + slab = getattr(self.cpu_slabs, slab_name) + if ( + value.device.type != "cpu" + or value.dtype is not torch.float32 + or value.numel() != unit.numel + or value.untyped_storage().data_ptr() + != slab.untyped_storage().data_ptr() + ): + raise RuntimeError( + f"Muon checkpoint state {key} lost its CPU slab alias" + ) + + def prepare_checkpoint_save(self) -> None: + """Drain staging before synchronous DCP reads authoritative CPU slabs.""" + if self._runtime.checkpoint_load_error is not None: + raise RuntimeError( + "cannot save after a failed checkpoint load; restart and recover" + ) from self._runtime.checkpoint_load_error + self.drain() + self._validate_bound_state_views() + if self.cuda_state_numel != 0: + raise RuntimeError("staged Muon checkpoint source contains CUDA state") + + def begin_checkpoint_load(self) -> None: + """Prepare for one in-place synchronous load without rollback state.""" + if self._runtime.checkpoint_load_error is not None: + raise RuntimeError( + "checkpoint load already failed; restart the process to recover" + ) from self._runtime.checkpoint_load_error + self.drain() + self._validate_bound_state_views() + + def complete_checkpoint_load(self) -> None: + """Validate CPU slab aliases after a successful load.""" + self.drain() + self._validate_bound_state_views() + + def mark_checkpoint_load_failed(self, error: BaseException) -> None: + """Fail-stop after an in-place load error; no rollback or retry.""" + self._runtime.mark_checkpoint_load_failed(error) + + def reset_from_model_params(self) -> None: + """Rebuild masters and clear momentum after a model-only load.""" + if not self._bound: + raise RuntimeError("GPU-staged Muon must be bound before state reset") + if not self._units: + return + assert self.cpu_slabs is not None + slots_were_released = not self._slots + if slots_were_released: + self.restore_from_cpu() + try: + self.drain() + self._runtime.schedule_units( + self._units, + self._schedule_master_initialization, + wait_for_compute=False, + ) + self.drain() + self.cpu_slabs.momentum.zero_() + finally: + if slots_were_released: + self.offload_to_cpu() + + def get_unscaled_state(self, param: torch.Tensor, key: str) -> torch.Tensor: + self.prepare_checkpoint_save() + return self.state[param][key] + + def set_scaled_state( + self, param: torch.Tensor, key: str, value: torch.Tensor + ) -> None: + if self._runtime.checkpoint_load_error is not None: + raise RuntimeError( + "cannot mutate after a failed checkpoint load; restart and recover" + ) from self._runtime.checkpoint_load_error + if key not in {"master_param", "momentum_buffer"}: + raise KeyError(f"unsupported staged Muon state: {key}") + if not isinstance(value, torch.Tensor): + raise TypeError(f"loaded {key} must be a tensor") + if value.device.type != "cpu" or value.dtype is not torch.float32: + raise TypeError(f"loaded {key} must be a CPU FP32 tensor") + destination = self.state[param][key] + if destination.shape != value.shape: + raise ValueError( + f"checkpoint state shape mismatch for {key}: " + f"expected {tuple(destination.shape)}, got {tuple(value.shape)}" + ) + destination.copy_(value) + + def state_dict(self) -> dict[str, Any]: + self.prepare_checkpoint_save() + return super().state_dict() + + def load_state_dict(self, state_dict: Mapping[str, Any]) -> None: + """Load CPU state in place without moving it to parameter devices.""" + if self._runtime.checkpoint_load_error is not None: + raise RuntimeError( + "cannot reload after a failed checkpoint load; restart and recover" + ) from self._runtime.checkpoint_load_error + if not isinstance(state_dict, Mapping) or set(state_dict) != { + "state", + "param_groups", + }: + raise ValueError("staged Muon checkpoint fields mismatch") + loaded_groups = state_dict["param_groups"] + if not isinstance(loaded_groups, (list, tuple)) or len(loaded_groups) != len( + self.param_groups + ): + raise ValueError("staged Muon parameter groups mismatch") + id_to_param: dict[int, torch.Tensor] = {} + for current_group, loaded_group in zip( + self.param_groups, loaded_groups, strict=True + ): + if not isinstance(loaded_group, Mapping) or set(loaded_group) != set( + current_group + ): + raise ValueError("staged Muon parameter-group metadata mismatch") + loaded_ids = loaded_group["params"] + if not isinstance(loaded_ids, (list, tuple)) or len(loaded_ids) != len( + current_group["params"] + ): + raise ValueError("staged Muon parameter-group ownership mismatch") + for loaded_id, param in zip( + loaded_ids, current_group["params"], strict=True + ): + if not isinstance(loaded_id, int) or isinstance(loaded_id, bool): + raise TypeError("staged Muon parameter ids must be integers") + if loaded_id in id_to_param: + raise ValueError("staged Muon parameter id is duplicated") + id_to_param[loaded_id] = param + loaded_state = state_dict["state"] + if not isinstance(loaded_state, Mapping) or set(loaded_state) != set( + id_to_param + ): + raise ValueError("staged Muon state ownership mismatch") + tensor_plan: list[tuple[torch.Tensor, str, torch.Tensor]] = [] + for loaded_id, values in loaded_state.items(): + if not isinstance(values, Mapping) or set(values) != { + "master_param", + "momentum_buffer", + }: + raise ValueError("staged Muon state schema mismatch") + param = id_to_param[loaded_id] + for key, value in values.items(): + destination = self.state[param][key] + if ( + not isinstance(value, torch.Tensor) + or value.device.type != "cpu" + or value.dtype is not torch.float32 + or value.shape != destination.shape + ): + raise ValueError(f"loaded staged Muon {key} is incompatible") + tensor_plan.append((param, key, value)) + self._validate_bound_state_views() + for current_group, loaded_group in zip( + self.param_groups, loaded_groups, strict=True + ): + params = current_group["params"] + current_group.clear() + current_group.update( + {key: value for key, value in loaded_group.items() if key != "params"} + ) + current_group["params"] = params + for param, key, value in tensor_plan: + self.state[param][key].copy_(value) + + +def _make_layerwise_leaf_class(): + from megatron.core.optimizer.optimizer import MegatronOptimizer + + class GPUStagedLayerWiseLeaf(MegatronOptimizer): + """MCore leaf adapter around one staged Muon or scalar optimizer.""" + + def __init__(self, optimizer: Any, config: Any, device: torch.device) -> None: + super().__init__(optimizer, config, None) + self.is_stub_optimizer = False + self._scale = torch.tensor([1.0], dtype=torch.float32, device=device) + self._step_activity: bool | None = None + + def _physical_parameters(self) -> list[torch.nn.Parameter]: + return [ + param + for group in self.optimizer.param_groups + for param in group["params"] + ] + + def preflight_step_activity(self) -> tuple[bool, str | None]: + if getattr(self.optimizer, "optimizer_kind", None) == "muon": + return self.optimizer.preflight_step_activity() + return ( + any( + getattr(param, "main_grad", None) is not None + or param.grad is not None + for param in self._physical_parameters() + ), + None, + ) + + def set_step_activity(self, active: bool) -> None: + if self._step_activity is not None: + raise RuntimeError("staged Muon leaf already has an activity plan") + self._step_activity = active + + def clear_step_activity(self) -> None: + self._step_activity = None + + def get_parameters(self) -> list[torch.nn.Parameter]: + if self._step_activity is False: + return [] + return self._physical_parameters() + + def zero_grad(self, set_to_none: bool = True) -> None: + for param in self.get_parameters(): + if set_to_none: + param.grad = None + param.decoupled_grad = None + else: + if param.grad is not None: + param.grad.zero_() + if getattr(param, "decoupled_grad", None) is not None: + param.decoupled_grad.zero_() + + def get_loss_scale(self) -> torch.Tensor: + return self._scale + + def prepare_grads(self) -> bool: + for param in self._physical_parameters(): + grad = getattr(param, "main_grad", None) + if grad is None: + grad = param.grad + param.decoupled_grad = ( + grad if self._step_activity is not False else None + ) + return False + + def step_with_ready_grads(self) -> bool: + if self._step_activity is False: + return True + self.optimizer.step() + return True + + def step(self): + return True, None, None + + def reload_model_params(self, state_dict=None) -> None: + del state_dict + raise RuntimeError( + "Muon staged model-only optimizer reset is not supported in this MVP" + ) + + def state_dict(self): + raise RuntimeError("Muon staged checkpoint is not supported in this MVP") + + def load_state_dict(self, state_dict): + del state_dict + raise RuntimeError("Muon staged checkpoint is not supported in this MVP") + + def sharded_state_dict(self, *args, **kwargs): + del args, kwargs + raise RuntimeError("Muon staged checkpoint is not supported in this MVP") + + def offload_to_cpu(self) -> None: + self.optimizer.offload_to_cpu() + + def restore_from_cpu(self) -> None: + self.optimizer.restore_from_cpu() + + return GPUStagedLayerWiseLeaf + + +@torch.no_grad() +def _allgather_layerwise_params_empty_safe(self: Any) -> None: + """All-gather layer-wise parameters when the first owner shard is empty.""" + + def _allgather(params_list: list[list[torch.Tensor]], group: Any) -> None: + first_param = next( + (param for owner_params in params_list for param in owner_params), None + ) + if first_param is None: + return + rank = torch.distributed.get_rank(group) + world_size = torch.distributed.get_world_size(group) + if len(params_list) != world_size: + raise RuntimeError( + "Muon owner list does not match its process group: " + f"owners={len(params_list)}, world_size={world_size}" + ) + local_params = params_list[rank] + source = ( + _flatten_dense_tensors(local_params) + if local_params + else torch.empty( + 0, + device=first_param.device, + dtype=first_param.dtype, + ) + ) + flat_sizes = [sum(param.numel() for param in params) for params in params_list] + gathered = [ + source + if owner_rank == rank + else torch.empty( + size, + device=first_param.device, + dtype=first_param.dtype, + ) + for owner_rank, size in enumerate(flat_sizes) + ] + torch.distributed.all_gather(gathered, source, group=group) + for owner_rank, owner_params in enumerate(params_list): + if owner_rank == rank or not owner_params: + continue + updated = _unflatten_dense_tensors(gathered[owner_rank], owner_params) + for source_param, model_param in zip(updated, owner_params, strict=True): + model_param.data.copy_(source_param) + + if self.pg_collection is None: + return + if self.dp_cp_params_list: + _allgather(self.dp_cp_params_list, self.pg_collection.dp_cp) + if self.expt_dp_params_list: + _allgather(self.expt_dp_params_list, self.pg_collection.expt_dp) + + +def get_megatron_optimizer_with_dist_muon( + mcore_config: Any, + model: list[Any], + *, + pg_collection: Any | None = None, + config_overrides: dict[Any, Any] | None = None, +) -> Any: + """Build MCore's native layer-wise Muon without mutating caller config.""" + version = importlib.metadata.version("megatron-core") + if version != _SUPPORTED_MEGATRON_CORE_VERSION: + raise RuntimeError( + f"distributed Muon supports megatron-core 0.19.0 exactly, found {version}" + ) + if ( + mcore_config.optimizer != "muon" + or not mcore_config.use_layer_wise_distributed_optimizer + ): + raise ValueError( + "distributed Muon requires the MCore layer-wise Muon optimizer" + ) + if mcore_config.use_distributed_optimizer: + raise ValueError("distributed Muon does not use Megatron dist-opt") + if mcore_config.muon_scalar_optimizer != "adam": + raise ValueError("AReaL distributed Muon supports only scalar AdamW") + if not mcore_config.bf16 or mcore_config.fp16: + raise ValueError("distributed Muon requires BF16 without FP16 loss scaling") + if mcore_config.use_precision_aware_optimizer: + raise ValueError("distributed Muon does not support precision-aware optimizer") + + from megatron.core.optimizer import get_megatron_optimizer + from megatron.core.optimizer.layer_wise_optimizer import ( + LayerWiseDistributedOptimizer, + ) + + build_config = copy.copy(mcore_config) + optimizer = get_megatron_optimizer( + build_config, + model, + use_gloo_process_groups=True, + pg_collection=pg_collection, + config_overrides=config_overrides, + ) + if type(optimizer) is not LayerWiseDistributedOptimizer: + raise TypeError( + "official Muon builder did not return LayerWiseDistributedOptimizer" + ) + # MCore 0.19 implements sharded checkpointing for this optimizer at fixed + # DP topology, despite not using Megatron's distributed Adam optimizer. + optimizer.supports_non_distributed_checkpoint = True + optimizer.allgather_params = _allgather_layerwise_params_empty_safe.__get__( + optimizer, type(optimizer) + ) + return optimizer + + +def _checkpoint_group_identity( + group: Any, *, name: str, singleton_global_rank: int +) -> dict[str, Any]: + size, rank = _group_size_and_rank(group, name=name) + if group is None: + members = [singleton_global_rank] + else: + ranks_fn = getattr(torch.distributed, "get_process_group_ranks", None) + if not callable(ranks_fn): + raise RuntimeError( + f"Muon checkpoint cannot enumerate the {name} process group" + ) + members = list(ranks_fn(group)) + if len(members) != size or len(set(members)) != size: + raise RuntimeError( + f"Muon checkpoint {name} group membership is inconsistent: " + f"size={size}, members={members}" + ) + return {"size": size, "rank": rank, "members": members} + + +def _validate_muon_checkpoint_group_metadata( + checkpoint: Any, runtime: Mapping[str, Any], *, location: str +) -> None: + if not isinstance(checkpoint, Mapping): + raise TypeError(f"{location} must be a mapping") + expected_keys = set(runtime) - {"params"} + if set(checkpoint) != expected_keys: + raise ValueError( + f"{location} fields mismatch: missing={sorted(expected_keys - set(checkpoint))}, " + f"unexpected={sorted(set(checkpoint) - expected_keys)}" + ) + for name, value in checkpoint.items(): + if name in { + "is_expert_parallel", + "is_decoupled_lr", + "maximize", + "capturable", + "differentiable", + "fused", + "nesterov", + }: + if value is not None and type(value) is not bool: + raise TypeError(f"{location} field {name} must be bool or None") + elif name == "betas": + if not isinstance(value, (tuple, list)) or len(value) != 2: + raise TypeError(f"{location} field betas must be a pair") + if any( + isinstance(beta, bool) + or not isinstance(beta, (int, float)) + or not math.isfinite(float(beta)) + or not 0.0 <= float(beta) < 1.0 + for beta in value + ): + raise ValueError(f"{location} field betas must be finite in [0, 1)") + elif name == "step": + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise TypeError(f"{location} field step must be a non-negative int") + elif isinstance(value, (int, float)) and not isinstance(value, bool): + if not math.isfinite(float(value)): + raise ValueError(f"{location} field {name} must be finite") + if ( + name + in { + "lr", + "initial_lr", + "max_lr", + "min_lr", + "weight_decay", + "momentum", + "eps", + } + and float(value) < 0.0 + ): + raise ValueError(f"{location} field {name} must be non-negative") + if name == "momentum" and not float(value) < 1.0: + raise ValueError(f"{location} field momentum must be less than 1") + if name == "eps" and float(value) <= 0.0: + raise ValueError(f"{location} field eps must be positive") + + +_MUON_CHECKPOINT_INVARIANT_GROUPS = ("tp", "expt_tp", "pp", "cp", "ep") +_MUON_CHECKPOINT_OWNERSHIP_GROUPS = ("dp", "dp_cp", "expt_dp") +_MUON_CHECKPOINT_GROUPS = ( + "dp", + "dp_cp", + "tp", + "cp", + "ep", + "expt_tp", + "expt_dp", + "pp", +) + + +def _muon_checkpoint_integer( + value: Any, + *, + location: str, + minimum: int = 0, + maximum: int | None = None, +) -> int: + if type(value) is not int: + raise ValueError(f"{location} must be an integer, got {value!r}") + if value < minimum or (maximum is not None and value >= maximum): + interval = f"[{minimum}, {maximum})" if maximum is not None else f">= {minimum}" + raise ValueError(f"{location} must be in {interval}, got {value!r}") + return value + + +def _validate_muon_checkpoint_cartesian_topology( + participants: Sequence[Mapping[str, Any]], + *, + group_sizes: Mapping[str, int], +) -> None: + """Prove that metadata has the two required WORLD rank spaces.""" + from megatron.core.parallel_state import RankGenerator + + world_size = len(participants) + order = "tp-cp-ep-dp-pp" + dense_dimensions = ("tp", "pp", "cp", "dp") + expert_dimensions = ("expt_tp", "pp", "ep", "expt_dp") + + def validate_coordinates(dimensions: tuple[str, ...], *, label: str) -> None: + product = math.prod(group_sizes[name] for name in dimensions) + if product != world_size: + raise ValueError( + f"Muon checkpoint {label} topology is not a WORLD Cartesian " + f"product: dimensions={dimensions!r}, product={product}, " + f"world_size={world_size}" + ) + actual = { + tuple(participant["groups"][name]["rank"] for name in dimensions) + for participant in participants + } + expected = set( + itertools.product(*(range(group_sizes[name]) for name in dimensions)) + ) + if len(actual) != world_size or actual != expected: + raise ValueError( + f"Muon checkpoint {label} coordinates do not cover the WORLD " + "Cartesian product exactly: " + f"duplicates={world_size - len(actual)}, " + f"missing={sorted(expected - actual)!r}, " + f"unexpected={sorted(actual - expected)!r}" + ) + + validate_coordinates(dense_dimensions, label="dense") + validate_coordinates(expert_dimensions, label="expert") + if group_sizes["dp_cp"] != group_sizes["dp"] * group_sizes["cp"]: + raise ValueError( + "Muon checkpoint dp_cp size does not match MCore's derived DP x CP " + f"layout: actual={group_sizes['dp_cp']}, " + f"expected={group_sizes['dp'] * group_sizes['cp']}" + ) + + dense_generator = RankGenerator( + tp=group_sizes["tp"], + ep=1, + dp=group_sizes["dp"], + pp=group_sizes["pp"], + cp=group_sizes["cp"], + order=order, + ) + expert_generator = RankGenerator( + tp=group_sizes["expt_tp"], + ep=group_sizes["ep"], + dp=group_sizes["expt_dp"], + pp=group_sizes["pp"], + cp=1, + order=order, + ) + if dense_generator.world_size != world_size: + raise ValueError("Muon checkpoint dense MCore topology does not cover WORLD") + if expert_generator.world_size != world_size: + raise ValueError("Muon checkpoint expert MCore topology does not cover WORLD") + if dense_generator.get_ranks("pp") != expert_generator.get_ranks("pp"): + raise ValueError("Muon checkpoint dense and expert pipeline topology disagree") + + actual_groups = { + name: sorted( + { + tuple(participant["groups"][name]["members"]) + for participant in participants + } + ) + for name in _MUON_CHECKPOINT_GROUPS + } + expected_groups = { + "tp": dense_generator.get_ranks("tp"), + "cp": dense_generator.get_ranks("cp"), + "dp": dense_generator.get_ranks("dp"), + "dp_cp": dense_generator.get_ranks("dp-cp"), + "pp": dense_generator.get_ranks("pp"), + "expt_tp": expert_generator.get_ranks("tp"), + "ep": expert_generator.get_ranks("ep"), + "expt_dp": expert_generator.get_ranks("dp"), + } + for name, expected in expected_groups.items(): + normalized_expected = sorted(tuple(members) for members in expected) + if actual_groups[name] != normalized_expected: + raise ValueError( + f"Muon checkpoint {name} topology does not match the expected " + f"rank layout: expected={normalized_expected!r}, " + f"actual={actual_groups[name]!r}" + ) + + +def _validate_muon_checkpoint_participant_topologies( + topologies: Sequence[Any], + *, + trusted_global_ranks: Sequence[int], +) -> tuple[list[dict[str, Any]], dict[str, int]]: + """Validate rank authority and every process-group WORLD partition. + + ``trusted_global_ranks`` comes from the receive slots of the explicit + WORLD-sized checkpoint control group. Metadata's self-reported rank is + never used to establish contributor authority. + """ + participant_count = len(topologies) + if participant_count == 0: + raise ValueError("Muon checkpoint metadata participants are empty") + if len(trusted_global_ranks) != participant_count: + raise ValueError( + "Muon checkpoint trusted participant count mismatch: " + f"metadata={participant_count}, trusted={len(trusted_global_ranks)}" + ) + trusted = [ + _muon_checkpoint_integer( + rank, + location=f"Muon checkpoint trusted participant slot {slot} rank", + maximum=participant_count, + ) + for slot, rank in enumerate(trusted_global_ranks) + ] + if set(trusted) != set(range(participant_count)): + raise ValueError( + "Muon checkpoint trusted participants must cover every WORLD rank " + f"exactly once: actual={trusted!r}" + ) + + canonical: list[dict[str, Any] | None] = [None] * participant_count + for slot, (trusted_rank, topology) in enumerate( + zip(trusted, topologies, strict=True) + ): + location = ( + f"Muon checkpoint participant slot {slot} (rank {trusted_rank}) topology" + ) + if not isinstance(topology, Mapping) or set(topology) != { + "world_size", + "global_rank", + "groups", + }: + raise ValueError(f"{location} fields mismatch") + world_size = _muon_checkpoint_integer( + topology["world_size"], + location=f"{location}.world_size", + minimum=1, + ) + if world_size != participant_count: + raise ValueError( + f"{location}.world_size mismatch: declared={world_size}, " + f"participants={participant_count}" + ) + declared_rank = _muon_checkpoint_integer( + topology["global_rank"], + location=f"{location}.global_rank", + maximum=participant_count, + ) + if declared_rank != trusted_rank: + raise ValueError( + f"{location}.global_rank does not match its trusted contributor: " + f"declared={declared_rank}, trusted={trusted_rank}" + ) + groups = topology["groups"] + if not isinstance(groups, Mapping) or set(groups) != set( + _MUON_CHECKPOINT_GROUPS + ): + raise ValueError(f"{location}.groups fields mismatch") + normalized_groups: dict[str, Any] = {} + for group_name in _MUON_CHECKPOINT_GROUPS: + identity = groups[group_name] + group_location = f"{location}.groups.{group_name}" + if not isinstance(identity, Mapping) or set(identity) != { + "size", + "rank", + "members", + }: + raise ValueError(f"{group_location} fields mismatch") + size = _muon_checkpoint_integer( + identity["size"], + location=f"{group_location}.size", + minimum=1, + maximum=participant_count + 1, + ) + group_rank = _muon_checkpoint_integer( + identity["rank"], + location=f"{group_location}.rank", + maximum=size, + ) + members_value = identity["members"] + if not isinstance(members_value, list): + raise ValueError(f"{group_location}.members must be a list") + members = [ + _muon_checkpoint_integer( + member, + location=f"{group_location}.members[{index}]", + maximum=participant_count, + ) + for index, member in enumerate(members_value) + ] + if len(members) != size or len(set(members)) != size: + raise ValueError( + f"{group_location} membership mismatch: size={size}, " + f"members={members!r}" + ) + if members != sorted(members): + raise ValueError( + f"{group_location}.members do not use MCore's global-rank " + f"ordering: members={members!r}" + ) + if trusted_rank not in members: + raise ValueError( + f"{group_location} does not contain participant rank {trusted_rank}" + ) + if members[group_rank] != trusted_rank: + raise ValueError( + f"{group_location}.rank has the wrong member index: " + f"rank={group_rank}, members={members!r}, " + f"participant={trusted_rank}" + ) + normalized_groups[group_name] = { + "size": size, + "rank": group_rank, + "members": members, + } + canonical[trusted_rank] = { + "world_size": participant_count, + "global_rank": trusted_rank, + "groups": normalized_groups, + } + + if any(topology is None for topology in canonical): + raise ValueError( + "Muon checkpoint participant WORLD rank coverage is incomplete" + ) + typed_canonical = [topology for topology in canonical if topology is not None] + group_sizes: dict[str, int] = {} + world = set(range(participant_count)) + for group_name in _MUON_CHECKPOINT_GROUPS: + declarations: dict[tuple[int, ...], set[int]] = {} + sizes: set[int] = set() + for participant_rank, topology in enumerate(typed_canonical): + identity = topology["groups"][group_name] + members = tuple(identity["members"]) + declarations.setdefault(members, set()).add(participant_rank) + sizes.add(identity["size"]) + if len(sizes) != 1: + raise ValueError( + f"Muon checkpoint {group_name} group sizes conflict across WORLD: " + f"sizes={sorted(sizes)!r}" + ) + covered: set[int] = set() + for members, declaring_ranks in declarations.items(): + member_set = set(members) + if declaring_ranks != member_set: + raise ValueError( + f"Muon checkpoint {group_name} membership declarations are " + "inconsistent: " + f"members={members!r}, declared_by={sorted(declaring_ranks)!r}" + ) + overlap = covered & member_set + if overlap: + raise ValueError( + f"Muon checkpoint {group_name} groups overlap on WORLD ranks " + f"{sorted(overlap)!r}" + ) + covered.update(member_set) + if covered != world: + raise ValueError( + f"Muon checkpoint {group_name} groups do not partition WORLD: " + f"missing={sorted(world - covered)!r}, " + f"extra={sorted(covered - world)!r}" + ) + group_sizes[group_name] = sizes.pop() + _validate_muon_checkpoint_cartesian_topology( + typed_canonical, group_sizes=group_sizes + ) + return typed_canonical, group_sizes + + +def _muon_checkpoint_coordinate( + topology: Mapping[str, Any], *, domain: str +) -> dict[str, int]: + groups = topology["groups"] + coordinate = { + "pp": groups["pp"]["rank"], + } + if domain == "dense": + coordinate["tp"] = groups["tp"]["rank"] + elif domain == "expert": + coordinate["ep"] = groups["ep"]["rank"] + coordinate["expt_tp"] = groups["expt_tp"]["rank"] + else: + raise ValueError(f"unsupported Muon checkpoint domain: {domain!r}") + return coordinate + + +def _normalize_muon_checkpoint_parameter_coordinate( + coordinate: Any, + *, + domain: str, + topology: Mapping[str, Any] | None = None, +) -> dict[str, int]: + if domain == "dense": + fields = {"pp": "pp", "tp": "tp"} + elif domain == "expert": + fields = {"pp": "pp", "ep": "ep", "expt_tp": "expt_tp"} + else: + raise ValueError(f"Muon checkpoint parameter domain is invalid: {domain!r}") + if not isinstance(coordinate, Mapping) or set(coordinate) != set(fields): + raise ValueError( + f"Muon checkpoint {domain} parameter coordinate fields mismatch" + ) + normalized = {} + for field, group_name in fields.items(): + maximum = ( + topology["groups"][group_name]["size"] if topology is not None else None + ) + normalized[field] = _muon_checkpoint_integer( + coordinate[field], + location=f"Muon checkpoint {domain} parameter coordinate {field}", + maximum=maximum, + ) + return normalized + + +def _muon_checkpoint_parameter_identity( + leaf_index: int, parameter: Mapping[str, Any] +) -> tuple[Any, ...]: + domain = parameter.get("domain") + coordinate = _normalize_muon_checkpoint_parameter_coordinate( + parameter.get("coordinate"), domain=domain + ) + return ( + leaf_index, + parameter.get("name"), + domain, + tuple(sorted(coordinate.items())), + ) + + +def _muon_checkpoint_parameter_invariant( + parameter: Mapping[str, Any], +) -> dict[str, Any]: + expected = { + "name", + "domain", + "coordinate", + "shape", + "dtype", + "state_kinds", + "source_owner", + } + actual = set(parameter) if isinstance(parameter, Mapping) else set() + if not isinstance(parameter, Mapping) or actual != expected: + raise ValueError( + "Muon checkpoint parameter fields mismatch: " + f"missing={sorted(expected - actual)}, " + f"unexpected={sorted(actual - expected)}" + ) + if not isinstance(parameter["name"], str) or not parameter["name"]: + raise ValueError("Muon checkpoint parameter name must be a non-empty string") + domain = parameter["domain"] + coordinate = _normalize_muon_checkpoint_parameter_coordinate( + parameter["coordinate"], domain=domain + ) + shape = parameter["shape"] + if not isinstance(shape, list): + raise ValueError("Muon checkpoint parameter shape must be a list") + normalized_shape = [ + _muon_checkpoint_integer( + dimension, + location=f"Muon checkpoint parameter shape[{index}]", + ) + for index, dimension in enumerate(shape) + ] + if not isinstance(parameter["dtype"], str): + raise ValueError("Muon checkpoint parameter dtype must be a string") + state_kinds = parameter["state_kinds"] + if not isinstance(state_kinds, list) or any( + not isinstance(state_kind, str) for state_kind in state_kinds + ): + raise ValueError("Muon checkpoint parameter state_kinds must be a string list") + return { + "name": parameter["name"], + "domain": domain, + "coordinate": coordinate, + "shape": normalized_shape, + "dtype": parameter["dtype"], + "state_kinds": list(state_kinds), + } + + +def _validate_muon_checkpoint_leaf_metadata( + leaf: Any, *, leaf_index: int, participant_rank: int +) -> None: + expected = {"tree_path", "kind", "parameters", "param_groups"} + if not isinstance(leaf, Mapping) or set(leaf) != expected: + raise ValueError( + f"Muon checkpoint rank {participant_rank} leaf {leaf_index} fields mismatch" + ) + tree_path = leaf["tree_path"] + if not isinstance(tree_path, list): + raise ValueError( + f"Muon checkpoint rank {participant_rank} leaf {leaf_index} tree_path " + "must be a list" + ) + normalized_path = [ + _muon_checkpoint_integer( + value, + location=( + f"Muon checkpoint rank {participant_rank} leaf {leaf_index} " + f"tree_path[{path_index}]" + ), + ) + for path_index, value in enumerate(tree_path) + ] + if normalized_path != [leaf_index]: + raise ValueError( + f"Muon checkpoint rank {participant_rank} leaf {leaf_index} tree_path " + f"mismatch: actual={normalized_path!r}" + ) + if leaf["kind"] not in _MUON_CHECKPOINT_STATE_KINDS: + raise ValueError( + f"Muon checkpoint rank {participant_rank} leaf {leaf_index} kind is invalid" + ) + if not isinstance(leaf["parameters"], list) or not isinstance( + leaf["param_groups"], list + ): + raise ValueError( + f"Muon checkpoint rank {participant_rank} leaf {leaf_index} containers " + "must be lists" + ) + + +def _validate_muon_checkpoint_source_owner( + source_owner: Any, + *, + domain: str, + ownership_sizes: Mapping[str, int], + world_size: int, + contributor_rank: int | None = None, + contributor_topology: Mapping[str, Any] | None = None, +) -> None: + expected = { + "global_rank", + "owner_rank", + "owner_ordinal", + "group_index", + "parameter_index", + "unit_order", + } + if not isinstance(source_owner, Mapping) or set(source_owner) != expected: + raise ValueError("Muon checkpoint source_owner fields mismatch") + owner_group = "expt_dp" if domain == "expert" else "dp_cp" + limits = { + "global_rank": world_size, + "owner_rank": ownership_sizes[owner_group], + } + for field in ( + "global_rank", + "owner_rank", + "owner_ordinal", + "group_index", + "parameter_index", + ): + value = source_owner[field] + _muon_checkpoint_integer( + value, + location=f"Muon checkpoint source owner {field}", + maximum=limits.get(field), + ) + unit_order = source_owner["unit_order"] + if unit_order is not None: + _muon_checkpoint_integer( + unit_order, + location="Muon checkpoint source owner unit_order", + ) + if contributor_rank is not None and source_owner["global_rank"] != contributor_rank: + raise ValueError( + "Muon checkpoint source owner global_rank does not match its trusted " + f"payload contributor: owner={source_owner['global_rank']}, " + f"contributor={contributor_rank}" + ) + if contributor_topology is not None: + expected_owner_rank = contributor_topology["groups"][owner_group]["rank"] + if source_owner["owner_rank"] != expected_owner_rank: + raise ValueError( + f"Muon checkpoint {domain} source owner coordinate mismatch: " + f"owner_rank={source_owner['owner_rank']}, " + f"participant_{owner_group}_rank={expected_owner_rank}" + ) + + +def merge_muon_checkpoint_metadata( + local_metadata: Sequence[Mapping[str, Any]], + *, + trusted_global_ranks: Sequence[int], +) -> dict[str, Any]: + """Create one DP-independent schema from rank-local official ownership.""" + if not local_metadata: + raise ValueError("Muon checkpoint metadata participants are empty") + participant_topologies, group_sizes = ( + _validate_muon_checkpoint_participant_topologies( + [ + metadata.get("topology") if isinstance(metadata, Mapping) else None + for metadata in local_metadata + ], + trusted_global_ranks=trusted_global_ranks, + ) + ) + trusted_by_slot = list(trusted_global_ranks) + first = local_metadata[0] + required = { + "schema_version", + "megatron_core_version", + "emerging_optimizers_version", + "topology", + "algorithm", + "leaf_tree", + } + if set(first) != required: + raise ValueError("Muon checkpoint local metadata fields mismatch") + schema_version = first["schema_version"] + if schema_version != _MUON_CHECKPOINT_SCHEMA_VERSION: + raise ValueError( + f"unsupported Muon checkpoint schema version: {schema_version!r}" + ) + invariant_sizes = { + name: group_sizes[name] for name in _MUON_CHECKPOINT_INVARIANT_GROUPS + } + ownership_sizes = { + name: group_sizes[name] for name in _MUON_CHECKPOINT_OWNERSHIP_GROUPS + } + if not isinstance(first["leaf_tree"], list): + raise ValueError("Muon checkpoint leaf_tree must be a list") + leaf_count = len(first["leaf_tree"]) + for leaf_index, leaf in enumerate(first["leaf_tree"]): + _validate_muon_checkpoint_leaf_metadata( + leaf, leaf_index=leaf_index, participant_rank=trusted_by_slot[0] + ) + merged_leaves = [ + { + "tree_path": list(first["leaf_tree"][index]["tree_path"]), + "kind": first["leaf_tree"][index]["kind"], + "parameters": [], + "param_groups": [], + } + for index in range(leaf_count) + ] + parameters: list[dict[tuple[Any, ...], dict[str, Any]]] = [ + {} for _ in range(leaf_count) + ] + groups: list[dict[int, dict[str, Any]]] = [{} for _ in range(leaf_count)] + owned_groups: list[set[int]] = [set() for _ in range(leaf_count)] + empty_groups: list[dict[int, list[dict[str, Any]]]] = [ + {} for _ in range(leaf_count) + ] + group_counts: dict[int, int] = {} + + for participant_slot, metadata in enumerate(local_metadata): + participant_rank = trusted_by_slot[participant_slot] + participant_topology = participant_topologies[participant_rank] + if set(metadata) != required: + raise ValueError( + f"Muon checkpoint rank {participant_rank} metadata fields mismatch" + ) + for field in ( + "schema_version", + "megatron_core_version", + "emerging_optimizers_version", + "algorithm", + ): + if metadata[field] != first[field]: + raise ValueError( + f"Muon checkpoint rank {participant_rank} {field} conflict" + ) + leaves = metadata["leaf_tree"] + if not isinstance(leaves, list) or len(leaves) != leaf_count: + raise ValueError("Muon checkpoint leaf tree length conflict") + for leaf_index, leaf in enumerate(leaves): + _validate_muon_checkpoint_leaf_metadata( + leaf, + leaf_index=leaf_index, + participant_rank=participant_rank, + ) + merged_leaf = merged_leaves[leaf_index] + if ( + list(leaf["tree_path"]) != merged_leaf["tree_path"] + or leaf["kind"] != merged_leaf["kind"] + ): + raise ValueError( + f"Muon checkpoint rank {participant_rank} leaf {leaf_index} conflict" + ) + local_owned_groups: set[int] = set() + group_count = len(leaf["param_groups"]) + if ( + group_count + and group_counts.setdefault(leaf_index, group_count) != group_count + ): + raise ValueError( + "Muon checkpoint param-group count conflicts across owners" + ) + for parameter in leaf["parameters"]: + invariant = _muon_checkpoint_parameter_invariant(parameter) + domain = invariant["domain"] + if domain not in {"dense", "expert"}: + raise ValueError( + f"Muon checkpoint parameter domain is invalid: {domain!r}" + ) + expected_coordinate = ( + {"pp", "tp"} if domain == "dense" else {"pp", "ep", "expt_tp"} + ) + if set(invariant["coordinate"]) != expected_coordinate: + raise ValueError( + "Muon checkpoint parameter coordinate fields mismatch" + ) + contributor_coordinate = _muon_checkpoint_coordinate( + participant_topology, domain=domain + ) + if invariant["coordinate"] != contributor_coordinate: + raise ValueError( + f"Muon checkpoint {domain} logical coordinate does not match " + f"trusted participant rank {participant_rank}: " + f"declared={invariant['coordinate']!r}, " + f"expected={contributor_coordinate!r}" + ) + expected_states = list( + _MUON_CHECKPOINT_STATE_KINDS[merged_leaf["kind"]] + ) + if invariant["state_kinds"] != expected_states: + raise ValueError( + "Muon checkpoint parameter state kinds mismatch: " + f"expected={expected_states!r}, " + f"actual={invariant['state_kinds']!r}" + ) + identity = _muon_checkpoint_parameter_identity(leaf_index, parameter) + previous = parameters[leaf_index].get(identity) + if previous is not None: + raise ValueError( + "Muon checkpoint logical parameter has duplicate owners: " + f"identity={identity!r}, first={previous['source_owner']!r}, " + f"duplicate={parameter['source_owner']!r}" + ) + source_owner = parameter["source_owner"] + _validate_muon_checkpoint_source_owner( + source_owner, + domain=domain, + ownership_sizes=ownership_sizes, + world_size=len(local_metadata), + contributor_rank=participant_rank, + contributor_topology=participant_topology, + ) + group_index = source_owner["group_index"] + if group_index >= len(leaf["param_groups"]): + raise ValueError( + "Muon checkpoint source owner group_index is out of range" + ) + local_owned_groups.add(group_index) + parameters[leaf_index][identity] = copy.deepcopy(dict(parameter)) + + for group_index, group in enumerate(leaf["param_groups"]): + normalized_group = copy.deepcopy(dict(group)) + previous_group = groups[leaf_index].setdefault( + group_index, normalized_group + ) + has_params = group_index in local_owned_groups + had_params = group_index in owned_groups[leaf_index] + # Empty AdamW groups do not advance step. Compare their shared + # settings, but take the step counter from a real owner. + compared_previous = previous_group + compared_current = normalized_group + if not (has_params and had_params): + compared_previous = { + key: value + for key, value in previous_group.items() + if key != "step" + } + compared_current = { + key: value + for key, value in normalized_group.items() + if key != "step" + } + if compared_previous != compared_current: + raise ValueError( + f"Muon checkpoint leaf {leaf_index} param-group " + f"{group_index} conflicts across owners" + ) + if has_params: + groups[leaf_index][group_index] = normalized_group + owned_groups[leaf_index].add(group_index) + else: + empty_groups[leaf_index].setdefault(group_index, []).append( + normalized_group + ) + + for leaf_index, merged_leaf in enumerate(merged_leaves): + for group_index, candidates in empty_groups[leaf_index].items(): + if group_index not in owned_groups[leaf_index] and any( + candidate != candidates[0] for candidate in candidates[1:] + ): + raise ValueError( + f"Muon checkpoint leaf {leaf_index} param-group " + f"{group_index} conflicts across owners" + ) + merged_leaf["parameters"] = [ + parameters[leaf_index][identity] + for identity in sorted(parameters[leaf_index], key=repr) + ] + merged_leaf["param_groups"] = [ + groups[leaf_index][index] for index in sorted(groups[leaf_index]) + ] + + return { + "schema_version": schema_version, + "megatron_core_version": first["megatron_core_version"], + "emerging_optimizers_version": first["emerging_optimizers_version"], + "topology": { + "invariant_group_sizes": invariant_sizes, + "source_ownership_group_sizes": ownership_sizes, + "source_world_size": len(local_metadata), + "source_participant_groups": [ + copy.deepcopy(topology["groups"]) for topology in participant_topologies + ], + }, + "algorithm": copy.deepcopy(first["algorithm"]), + "leaf_tree": merged_leaves, + } + + +def _make_staged_layerwise_class(): + from megatron.core.optimizer.layer_wise_optimizer import ( + LayerWiseDistributedOptimizer, + ) + from megatron.core.optimizer.optimizer import ChainedOptimizer + + class GPUStagedLayerWiseDistributedOptimizer(LayerWiseDistributedOptimizer): + """Layer-wise wrapper preserving official ownership and collectives.""" + + manages_cpu_residency = True + supports_non_distributed_checkpoint = True + + def __init__(self, official: Any, leaves: list[Any]) -> None: + self.pg_collection = official.pg_collection + self.dp_cp = official.dp_cp + self.expt_dp = official.expt_dp + self.dp_cp_params_list = official.dp_cp_params_list + self.expt_dp_params_list = official.expt_dp_params_list + self._staged_owner_schema = { + "dense": _freeze_owner_schema("dense", self.dp_cp_params_list), + "expert": _freeze_owner_schema("expert", self.expt_dp_params_list), + } + self._staged_leaves = tuple(leaves) + self.overlap_param_gather = official.overlap_param_gather + self.use_buffer_param_sync = official.use_buffer_param_sync + self._checkpoint_parameter_names: dict[torch.Tensor, str] | None = None + self._checkpoint_algorithm: dict[str, Any] | None = None + ChainedOptimizer.__init__(self, leaves) + self.model_chunks = official.model_chunks + self._checkpoint_process_group: Any | None = None + + def bind_managed_checkpoint_process_group(self, group: Any) -> None: + """Bind the explicit WORLD-sized Gloo group used by checkpoint phases.""" + if group is None: + raise RuntimeError( + "Muon checkpoint requires an explicit checkpoint process group" + ) + self._checkpoint_process_group = group + + def configure_managed_checkpoint_schema( + self, + parameter_names: Mapping[torch.Tensor, str], + *, + algorithm: Mapping[str, Any] | None = None, + ) -> None: + names: dict[torch.Tensor, str] = {} + params_by_name: dict[str, torch.Tensor] = {} + for leaf in self._staged_leaves: + for group in leaf.optimizer.param_groups: + for param in group["params"]: + try: + stable_name = parameter_names[param] + except KeyError as error: + raise RuntimeError( + "Muon checkpoint parameter lacks a stable model name" + ) from error + previous = names.setdefault(param, stable_name) + if previous != stable_name: + raise RuntimeError( + "Muon checkpoint parameter has ambiguous model names" + ) + other = params_by_name.setdefault(stable_name, param) + if other is not param: + raise RuntimeError( + "Muon checkpoint stable model name is not unique: " + f"{stable_name!r}" + ) + self._checkpoint_parameter_names = names + if algorithm is not None: + self._checkpoint_algorithm = dict(algorithm) + + def _checkpoint_topology(self) -> dict[str, Any]: + if self._checkpoint_parameter_names is None: + raise RuntimeError("Muon checkpoint schema was not configured") + if torch.distributed.is_available() and torch.distributed.is_initialized(): + checkpoint_group = self._checkpoint_process_group + if checkpoint_group is None: + raise RuntimeError( + "Muon checkpoint process group was not bound before topology " + "construction" + ) + ranks_fn = getattr(torch.distributed, "get_process_group_ranks", None) + if not callable(ranks_fn): + raise RuntimeError( + "Muon checkpoint cannot enumerate its explicit checkpoint " + "process group" + ) + checkpoint_members = list(ranks_fn(checkpoint_group)) + checkpoint_rank = torch.distributed.get_rank(checkpoint_group) + if len(checkpoint_members) != torch.distributed.get_world_size( + checkpoint_group + ) or not 0 <= checkpoint_rank < len(checkpoint_members): + raise RuntimeError( + "Muon checkpoint process-group membership is inconsistent" + ) + global_rank = checkpoint_members[checkpoint_rank] + world_size = len(checkpoint_members) + else: + global_rank = 0 + world_size = 1 + groups = { + name: _checkpoint_group_identity( + getattr(self.pg_collection, name), + name=name, + singleton_global_rank=global_rank, + ) + for name in ( + "tp", + "expt_tp", + "dp", + "dp_cp", + "expt_dp", + "pp", + "cp", + "ep", + ) + } + return { + "world_size": world_size, + "global_rank": global_rank, + "groups": groups, + } + + def _checkpoint_owner_lookup( + self, + ) -> dict[torch.Tensor, tuple[str, int, int]]: + result: dict[torch.Tensor, tuple[str, int, int]] = {} + for domain in ("dense", "expert"): + schema = self._staged_owner_schema[domain] + if schema is None: + continue + for owner in schema: + for entry in owner: + value = (domain, entry.owner_rank, entry.ordinal) + previous = result.setdefault(entry.param, value) + if previous != value: + raise RuntimeError( + "Muon checkpoint parameter occurs in multiple owner domains" + ) + # MCore represents size-one DP/CP ownership with ``None`` rather + # than a one-entry owner list. In that exact case every local + # parameter is the official owner; this is not a fallback for a + # missing distributed ownership schema. + next_ordinal = {"dense": 0, "expert": 0} + for leaf in self._staged_leaves: + for group in leaf.optimizer.param_groups: + domain = ( + "expert" if group.get("is_expert_parallel", False) else "dense" + ) + owner_schema = self._staged_owner_schema[domain] + owner_group = ( + self.pg_collection.expt_dp + if domain == "expert" + else self.pg_collection.dp_cp + ) + group_size, _ = _group_size_and_rank( + owner_group, + name="expt_dp" if domain == "expert" else "dp_cp", + ) + for param in group["params"]: + if param in result: + continue + if owner_schema is not None or group_size != 1: + raise RuntimeError( + f"Muon checkpoint {domain} ownership is incomplete" + ) + ordinal = next_ordinal[domain] + result[param] = (domain, 0, ordinal) + next_ordinal[domain] = ordinal + 1 + return result + + def _checkpoint_local_metadata(self) -> dict[str, Any]: + mcore_version, eo_version = _require_muon_checkpoint_versions() + if self._checkpoint_parameter_names is None: + raise RuntimeError("Muon checkpoint schema was not configured") + if self._checkpoint_algorithm is None: + raise RuntimeError( + "Muon checkpoint algorithm schema was not configured" + ) + topology = self._checkpoint_topology() + owner_lookup = self._checkpoint_owner_lookup() + leaves: list[dict[str, Any]] = [] + for leaf_index, leaf in enumerate(self._staged_leaves): + base = leaf.optimizer + kind = base.optimizer_kind + state_kinds = list(_MUON_CHECKPOINT_STATE_KINDS[kind]) + parameters = [] + if kind == "muon": + unit_order = { + unit.param: unit_index + for unit_index, unit in enumerate(base.units) + } + else: + unit_order = { + param: parameter_index + for parameter_index, param in enumerate( + param + for group in base.param_groups + for param in group["params"] + ) + } + for group_index, group in enumerate(base.param_groups): + for parameter_index, param in enumerate(group["params"]): + try: + domain, owner_rank, owner_ordinal = owner_lookup[param] + name = self._checkpoint_parameter_names[param] + except KeyError as error: + raise RuntimeError( + "Muon checkpoint owner metadata is incomplete" + ) from error + parameters.append( + { + "name": name, + "domain": domain, + "coordinate": _muon_checkpoint_coordinate( + topology, domain=domain + ), + "shape": list(param.shape), + "dtype": str(param.dtype), + "state_kinds": state_kinds, + "source_owner": { + "global_rank": topology["global_rank"], + "owner_rank": owner_rank, + "owner_ordinal": owner_ordinal, + "group_index": group_index, + "parameter_index": parameter_index, + "unit_order": unit_order.get(param), + }, + } + ) + leaves.append( + { + "tree_path": [leaf_index], + "kind": kind, + "parameters": parameters, + "param_groups": [ + { + key: value + for key, value in group.items() + if key != "params" + } + for group in base.param_groups + ], + } + ) + return { + "schema_version": _MUON_CHECKPOINT_SCHEMA_VERSION, + "megatron_core_version": mcore_version, + "emerging_optimizers_version": eo_version, + "topology": topology, + "algorithm": self._checkpoint_algorithm, + "leaf_tree": leaves, + } + + def _checkpoint_metadata(self) -> dict[str, Any]: + if not ( + torch.distributed.is_available() and torch.distributed.is_initialized() + ): + return merge_muon_checkpoint_metadata( + [self._checkpoint_local_metadata()], trusted_global_ranks=[0] + ) + group = self._checkpoint_process_group + if group is None: + raise RuntimeError( + "Muon checkpoint process group was not bound before metadata gather" + ) + try: + local_payload = { + "error": None, + "metadata": self._checkpoint_local_metadata(), + } + except BaseException as error: + local_payload = { + "error": f"{type(error).__name__}: {error}", + "metadata": None, + } + participants: list[Any] = [ + None for _ in range(torch.distributed.get_world_size(group)) + ] + torch.distributed.all_gather_object( + participants, local_payload, group=group + ) + failures = [ + (rank, payload["error"]) + for rank, payload in enumerate(participants) + if payload["error"] is not None + ] + if failures: + raise RuntimeError( + f"Muon checkpoint local metadata validation failed: {failures!r}" + ) + ranks_fn = getattr(torch.distributed, "get_process_group_ranks", None) + if not callable(ranks_fn): + raise RuntimeError( + "Muon checkpoint cannot bind metadata participants to explicit " + "checkpoint process-group ranks" + ) + trusted_global_ranks = list(ranks_fn(group)) + return merge_muon_checkpoint_metadata( + [payload["metadata"] for payload in participants], + trusted_global_ranks=trusted_global_ranks, + ) + + def managed_checkpoint_identities( + self, model_parameter_names: Mapping[torch.Tensor, str] + ) -> dict[tuple[int, ...], dict[str, Any]]: + self.configure_managed_checkpoint_schema(model_parameter_names) + metadata = self._checkpoint_metadata() + result = {} + for leaf_index, leaf in enumerate(metadata["leaf_tree"]): + result[(leaf_index,)] = { + "version": _MUON_CHECKPOINT_SCHEMA_VERSION, + "tree_path": [leaf_index], + "kind": leaf["kind"], + "invariant_group_sizes": metadata["topology"][ + "invariant_group_sizes" + ], + "parameters": [ + _muon_checkpoint_parameter_invariant(parameter) + for parameter in leaf["parameters"] + ], + } + return result + + def _checkpoint_rank_prefix(self) -> str: + return _MUON_CHECKPOINT_PREFIX + + def managed_checkpoint_outer_template(self) -> dict[str, Any]: + from megatron.core.dist_checkpointing.mapping import ShardedObject + + prefix = self._checkpoint_rank_prefix() + return { + "metadata": ShardedObject( + f"{prefix}.metadata", + None, + (1,), + (0,), + replica_id=self._checkpoint_topology()["global_rank"], + ) + } + + def validate_managed_checkpoint_outer_state(self, state: Any) -> None: + if not isinstance(state, Mapping) or set(state) != {"metadata"}: + raise ValueError("Muon checkpoint outer fields must be {'metadata'}") + checkpoint = state["metadata"] + if not isinstance(checkpoint, Mapping): + raise TypeError("Muon checkpoint metadata must be a mapping") + runtime = self._checkpoint_metadata() + if set(checkpoint) != set(runtime): + raise ValueError("Muon checkpoint metadata fields mismatch") + for field in ( + "schema_version", + "megatron_core_version", + "emerging_optimizers_version", + "algorithm", + ): + if checkpoint[field] != runtime[field]: + raise ValueError( + f"Muon checkpoint {field} mismatch: " + f"expected={runtime[field]!r}, actual={checkpoint[field]!r}" + ) + checkpoint_topology = checkpoint["topology"] + if not isinstance(checkpoint_topology, Mapping) or set( + checkpoint_topology + ) != { + "invariant_group_sizes", + "source_ownership_group_sizes", + "source_world_size", + "source_participant_groups", + }: + raise ValueError("Muon checkpoint topology fields mismatch") + if ( + checkpoint_topology["invariant_group_sizes"] + != runtime["topology"]["invariant_group_sizes"] + ): + raise ValueError( + "Muon checkpoint invariant topology mismatch: " + f"expected={runtime['topology']['invariant_group_sizes']!r}, " + f"actual={checkpoint_topology['invariant_group_sizes']!r}" + ) + invariant_sizes = checkpoint_topology["invariant_group_sizes"] + if not isinstance(invariant_sizes, Mapping) or set(invariant_sizes) != set( + _MUON_CHECKPOINT_INVARIANT_GROUPS + ): + raise ValueError("Muon checkpoint invariant topology fields mismatch") + normalized_invariant_sizes = { + name: _muon_checkpoint_integer( + invariant_sizes[name], + location=f"Muon checkpoint invariant topology {name} size", + minimum=1, + ) + for name in _MUON_CHECKPOINT_INVARIANT_GROUPS + } + if ( + normalized_invariant_sizes + != runtime["topology"]["invariant_group_sizes"] + ): + raise ValueError( + "Muon checkpoint invariant topology mismatch: " + f"expected={runtime['topology']['invariant_group_sizes']!r}, " + f"actual={normalized_invariant_sizes!r}" + ) + source_sizes = checkpoint_topology["source_ownership_group_sizes"] + if not isinstance(source_sizes, Mapping) or set(source_sizes) != set( + _MUON_CHECKPOINT_OWNERSHIP_GROUPS + ): + raise ValueError( + "Muon checkpoint source ownership topology fields mismatch" + ) + source_sizes = { + name: _muon_checkpoint_integer( + source_sizes[name], + location=f"Muon checkpoint source ownership {name} size", + minimum=1, + ) + for name in _MUON_CHECKPOINT_OWNERSHIP_GROUPS + } + source_world_size = checkpoint_topology["source_world_size"] + source_world_size = _muon_checkpoint_integer( + source_world_size, + location="Muon checkpoint source world size", + minimum=1, + ) + source_participant_groups = checkpoint_topology["source_participant_groups"] + if not isinstance(source_participant_groups, list): + raise ValueError( + "Muon checkpoint source participant groups must be a list" + ) + source_topologies, verified_source_sizes = ( + _validate_muon_checkpoint_participant_topologies( + [ + { + "world_size": source_world_size, + "global_rank": global_rank, + "groups": participant_groups, + } + for global_rank, participant_groups in enumerate( + source_participant_groups + ) + ], + trusted_global_ranks=list(range(source_world_size)), + ) + ) + if { + name: verified_source_sizes[name] + for name in _MUON_CHECKPOINT_OWNERSHIP_GROUPS + } != source_sizes: + raise ValueError( + "Muon checkpoint source ownership group size metadata conflicts " + "with participant topology" + ) + if { + name: verified_source_sizes[name] + for name in _MUON_CHECKPOINT_INVARIANT_GROUPS + } != checkpoint_topology["invariant_group_sizes"]: + raise ValueError( + "Muon checkpoint invariant group size metadata conflicts with " + "participant topology" + ) + checkpoint_leaves = checkpoint["leaf_tree"] + expected_leaves = runtime["leaf_tree"] + if not isinstance(checkpoint_leaves, list) or len(checkpoint_leaves) != len( + expected_leaves + ): + raise ValueError("Muon checkpoint leaf tree mismatch") + for leaf_index, (loaded, runtime) in enumerate( + zip(checkpoint_leaves, expected_leaves, strict=True) + ): + if set(loaded) != set(runtime): + raise ValueError( + f"Muon checkpoint leaf {leaf_index} fields mismatch" + ) + for field in ("tree_path", "kind"): + if loaded[field] != runtime[field]: + raise ValueError( + f"Muon checkpoint leaf {leaf_index} {field} mismatch" + ) + loaded_parameters = { + _muon_checkpoint_parameter_identity( + leaf_index, parameter + ): _muon_checkpoint_parameter_invariant(parameter) + for parameter in loaded["parameters"] + } + if len(loaded_parameters) != len(loaded["parameters"]): + raise ValueError( + f"Muon checkpoint leaf {leaf_index} has duplicate parameters" + ) + for parameter in loaded["parameters"]: + source_owner = parameter["source_owner"] + source_rank = _muon_checkpoint_integer( + source_owner.get("global_rank"), + location="Muon checkpoint source owner global_rank", + maximum=source_world_size, + ) + _validate_muon_checkpoint_source_owner( + source_owner, + domain=parameter["domain"], + ownership_sizes=source_sizes, + world_size=source_world_size, + contributor_rank=source_rank, + contributor_topology=source_topologies[source_rank], + ) + expected_coordinate = _muon_checkpoint_coordinate( + source_topologies[source_rank], domain=parameter["domain"] + ) + if parameter["coordinate"] != expected_coordinate: + raise ValueError( + "Muon checkpoint logical parameter coordinate conflicts " + "with its source participant topology" + ) + runtime_parameters = { + _muon_checkpoint_parameter_identity( + leaf_index, parameter + ): _muon_checkpoint_parameter_invariant(parameter) + for parameter in runtime["parameters"] + } + if loaded_parameters != runtime_parameters: + raise ValueError( + f"Muon checkpoint leaf {leaf_index} parameter identity mismatch" + ) + groups = loaded["param_groups"] + if not isinstance(groups, list) or len(groups) != len( + runtime["param_groups"] + ): + raise ValueError( + f"Muon checkpoint leaf {leaf_index} param-group count mismatch" + ) + base = self._staged_leaves[leaf_index].optimizer + if not base.param_groups: + if base.state: + raise RuntimeError( + f"Muon checkpoint empty leaf {leaf_index} retained state" + ) + continue + for group_index, (group, runtime_group) in enumerate( + zip(groups, base.param_groups, strict=True) + ): + _validate_muon_checkpoint_group_metadata( + group, + runtime_group, + location=( + f"Muon checkpoint leaf {leaf_index} group {group_index}" + ), + ) + + def _checkpoint_tensor_key( + self, + *, + leaf_index: int, + parameter_name: str, + domain: str, + coordinate: Mapping[str, int], + state_kind: str, + ) -> str: + name_digest = hashlib.sha256(parameter_name.encode()).hexdigest() + coordinate_digest = hashlib.sha256( + json.dumps( + dict(coordinate), sort_keys=True, separators=(",", ":") + ).encode() + ).hexdigest() + return ( + f"{_MUON_CHECKPOINT_PREFIX}.leaf_{leaf_index}.{domain}." + f"coord_{coordinate_digest}." + f"param_{name_digest}.{state_kind}" + ) + + def sharded_state_dict( + self, + model_sharded_state_dict: Any, + is_loading: bool = False, + **kwargs: Any, + ) -> dict[str, Any]: + del model_sharded_state_dict, is_loading, kwargs + from megatron.core.dist_checkpointing.mapping import ( + ShardedObject, + ShardedTensor, + ) + + _require_muon_checkpoint_versions() + self.drain() + metadata = self._checkpoint_metadata() + local_metadata = self._checkpoint_local_metadata() + prefix = self._checkpoint_rank_prefix() + tensor_state: dict[str, Any] = {} + for leaf_index, leaf in enumerate(self._staged_leaves): + base = leaf.optimizer + state_kinds = _MUON_CHECKPOINT_STATE_KINDS[base.optimizer_kind] + descriptors = { + parameter["name"]: parameter + for parameter in local_metadata["leaf_tree"][leaf_index][ + "parameters" + ] + } + for group in base.param_groups: + for param in group["params"]: + name = self._checkpoint_parameter_names[param] + descriptor = descriptors[name] + for state_kind in state_kinds: + value = base.state[param][state_kind] + flat_value = value.view(-1) + tensor_key = self._checkpoint_tensor_key( + leaf_index=leaf_index, + parameter_name=name, + domain=descriptor["domain"], + coordinate=descriptor["coordinate"], + state_kind=state_kind, + ) + if tensor_key in tensor_state: + raise RuntimeError( + "Muon checkpoint tensor key collision for " + f"parameter {name!r} state {state_kind!r}" + ) + tensor_state[tensor_key] = ShardedTensor( + tensor_key, + flat_value, + flat_value.dtype, + tuple(flat_value.shape), + tuple(flat_value.shape), + (0,), + (1,), + replica_id=0, + ) + return { + "metadata": ShardedObject( + f"{prefix}.metadata", + metadata, + (1,), + (0,), + replica_id=self._checkpoint_topology()["global_rank"], + ), + "state": tensor_state, + } + + def state_dict(self) -> dict[str, Any]: + self.drain() + metadata = self._checkpoint_metadata() + local_metadata = self._checkpoint_local_metadata() + tensor_state: dict[str, torch.Tensor] = {} + for leaf_index, leaf in enumerate(self._staged_leaves): + base = leaf.optimizer + state_kinds = _MUON_CHECKPOINT_STATE_KINDS[base.optimizer_kind] + descriptors = { + parameter["name"]: parameter + for parameter in local_metadata["leaf_tree"][leaf_index][ + "parameters" + ] + } + for group in base.param_groups: + for param in group["params"]: + name = self._checkpoint_parameter_names[param] + descriptor = descriptors[name] + for state_kind in state_kinds: + tensor_key = self._checkpoint_tensor_key( + leaf_index=leaf_index, + parameter_name=name, + domain=descriptor["domain"], + coordinate=descriptor["coordinate"], + state_kind=state_kind, + ) + if tensor_key in tensor_state: + raise RuntimeError( + "Muon checkpoint tensor key collision for " + f"parameter {name!r} state {state_kind!r}" + ) + tensor_state[tensor_key] = base.state[param][ + state_kind + ].view(-1) + return { + "metadata": metadata, + "state": tensor_state, + } + + def load_state_dict(self, state_dict: Any) -> None: + if not isinstance(state_dict, Mapping): + raise TypeError("Muon staged checkpoint must be a mapping") + if "metadata" not in state_dict or set(state_dict) - { + "metadata", + "state", + }: + raise ValueError("Muon staged checkpoint fields mismatch") + self.validate_managed_checkpoint_outer_state( + {"metadata": state_dict["metadata"]} + ) + # torch_dist omits an empty nested mapping from a rank-local load + # result. Normalize only that container; exact local tensor keys + # remain mandatory below for every owned parameter. + tensor_state = state_dict.get("state", {}) + if not isinstance(tensor_state, Mapping): + raise TypeError("Muon checkpoint tensor state must be a mapping") + metadata_leaves = state_dict["metadata"]["leaf_tree"] + local_metadata = self._checkpoint_local_metadata() + apply_plan: list[ + tuple[ + Any, + list[Mapping[str, Any]], + list[tuple[torch.Tensor, str, torch.Tensor]], + ] + ] = [] + expected_tensor_keys: set[str] = set() + for leaf_index, (leaf, leaf_metadata) in enumerate( + zip(self._staged_leaves, metadata_leaves, strict=True) + ): + base = leaf.optimizer + groups = leaf_metadata["param_groups"] + name_to_param = { + self._checkpoint_parameter_names[param]: param + for group in base.param_groups + for param in group["params"] + } + if len(name_to_param) != sum( + len(group["params"]) for group in base.param_groups + ): + raise RuntimeError( + f"Muon checkpoint leaf {leaf_index} stable names are not unique" + ) + state_kinds = _MUON_CHECKPOINT_STATE_KINDS[base.optimizer_kind] + descriptors = { + parameter["name"]: parameter + for parameter in local_metadata["leaf_tree"][leaf_index][ + "parameters" + ] + } + tensor_plan: list[tuple[torch.Tensor, str, torch.Tensor]] = [] + for name, param in name_to_param.items(): + descriptor = descriptors[name] + for state_kind in state_kinds: + tensor_key = self._checkpoint_tensor_key( + leaf_index=leaf_index, + parameter_name=name, + domain=descriptor["domain"], + coordinate=descriptor["coordinate"], + state_kind=state_kind, + ) + expected_tensor_keys.add(tensor_key) + value = tensor_state.get(tensor_key) + if not isinstance(value, torch.Tensor): + raise TypeError( + f"Muon checkpoint tensor {tensor_key!r} must be a tensor" + ) + if ( + value.device.type != "cpu" + or value.dtype is not torch.float32 + ): + raise TypeError( + f"Muon checkpoint parameter {name!r} {state_kind} " + "must be CPU FP32" + ) + expected_shape = (param.numel(),) + if tuple(value.shape) != expected_shape: + raise ValueError( + f"Muon checkpoint parameter {name!r} {state_kind} " + f"shape mismatch: expected={expected_shape}, " + f"actual={tuple(value.shape)}" + ) + tensor_plan.append((param, state_kind, value)) + apply_plan.append( + (base, groups if base.param_groups else [], tensor_plan) + ) + + if set(tensor_state) != expected_tensor_keys: + raise ValueError( + "Muon checkpoint tensor key mismatch: " + f"missing={sorted(expected_tensor_keys - set(tensor_state))}, " + f"unexpected={sorted(set(tensor_state) - expected_tensor_keys)}" + ) + + # Every leaf, group, and tensor has passed validation before mutation. + for base, groups, tensor_plan in apply_plan: + for group, loaded_group in zip(base.param_groups, groups, strict=True): + params = group["params"] + group.clear() + group.update(loaded_group) + group["params"] = params + for param, state_kind, value in tensor_plan: + base.set_scaled_state(param, state_kind, value.view_as(param)) + + def _validate_owner_domain( + self, + *, + domain: str, + params_list: Sequence[Sequence[torch.Tensor]] | None, + group: Any, + ) -> tuple[list[int], torch.Tensor | None, int, str | None]: + """Validate one frozen domain and complete its fixed subgroup vote.""" + group_size, group_rank = _group_size_and_rank( + group, + name="expt_dp" if domain == "expert" else "dp_cp", + ) + expected = self._staged_owner_schema[domain] + local_error: str | None = None + flat_sizes: list[int] = [] + first_param: torch.Tensor | None = None + try: + if (params_list is None) != (expected is None): + raise RuntimeError( + f"staged Muon {domain} owner schema changed after bind" + ) + if params_list is None: + if expected is not None: + raise AssertionError("owner schema presence mismatch") + elif not isinstance(params_list, Sequence) or isinstance( + params_list, (str, bytes) + ): + raise RuntimeError( + f"staged Muon {domain} owner schema is not rank-indexed" + ) + elif expected is None or len(params_list) != group_size: + raise RuntimeError( + f"staged Muon {domain} owner schema length " + f"{len(params_list)} does not match group size {group_size}" + ) + elif len(expected) != len(params_list): + raise RuntimeError( + f"staged Muon {domain} owner schema changed after bind" + ) + else: + seen: set[int] = set() + expected_device: torch.device | None = None + expected_dtype: torch.dtype | None = None + for owner_rank, (params, expected_params) in enumerate( + zip(params_list, expected, strict=True) + ): + if not isinstance(params, Sequence) or isinstance( + params, (str, bytes) + ): + raise RuntimeError( + f"staged Muon {domain} owner rank {owner_rank} " + "entry is not a sequence" + ) + if len(params) != len(expected_params): + raise RuntimeError( + f"staged Muon {domain} owner rank {owner_rank} " + "parameter order changed after bind" + ) + size = 0 + for param, frozen in zip(params, expected_params, strict=True): + if not isinstance(param, torch.Tensor): + raise RuntimeError( + f"staged Muon {domain} owner rank " + f"{owner_rank} contains a non-tensor" + ) + if param is not frozen.param: + raise RuntimeError( + f"staged Muon {domain} owner rank " + f"{owner_rank} parameter order changed after bind" + ) + if id(param) in seen: + raise RuntimeError( + f"staged Muon {domain} owner schema contains a " + f"duplicate parameter on owner rank {owner_rank}" + ) + seen.add(id(param)) + storage = param.untyped_storage() + metadata = ( + ("shape", frozen.shape, tuple(param.shape)), + ("numel", frozen.numel, param.numel()), + ("stride", frozen.stride, tuple(param.stride())), + ("layout", frozen.layout, param.layout), + ("dtype", frozen.dtype, param.dtype), + ("device_type", frozen.device_type, param.device.type), + ( + "device_index", + frozen.device_index, + param.device.index, + ), + ( + # AWEX may reallocate this storage on resume; + # its identity and tensor aliases must survive. + "storage_object_identity", + True, + storage is frozen.storage, + ), + ( + "storage_cdata", + frozen.storage_cdata, + int(storage._cdata), + ), + ( + "storage_nbytes", + frozen.storage_nbytes, + int(storage.nbytes()), + ), + ( + "storage_offset", + frozen.storage_offset, + param.storage_offset(), + ), + ) + drift = [ + f"{field}: expected={expected_value!r}, " + f"actual={actual_value!r}" + for field, expected_value, actual_value in metadata + if actual_value != expected_value + ] + if drift: + raise RuntimeError( + f"staged Muon {domain} owner rank {owner_rank} " + "parameter metadata changed after bind: " + f"ordinal={frozen.ordinal}, fields=[" + f"{'; '.join(drift)}]" + ) + if expected_device is None: + expected_device = param.device + expected_dtype = param.dtype + elif ( + param.device != expected_device + or param.dtype != expected_dtype + ): + raise RuntimeError( + f"staged Muon {domain} owner rank {owner_rank} " + "has inconsistent parameter device or dtype" + ) + first_param = param if first_param is None else first_param + size += param.numel() + flat_sizes.append(size) + except BaseException as error: + local_error = f"{type(error).__name__}: {error}" + + digest = _owner_schema_digest(domain, expected) + if group_size > 1: + vote_device = ( + first_param.device + if first_param is not None and first_param.is_cuda + else torch.device("cuda", torch.cuda.current_device()) + ) + vote = torch.tensor( + [int(local_error is None), *digest], + dtype=torch.int64, + device=vote_device, + ) + minimum = vote.clone() + maximum = vote.clone() + torch.distributed.all_reduce( + minimum, op=torch.distributed.ReduceOp.MIN, group=group + ) + torch.distributed.all_reduce( + maximum, op=torch.distributed.ReduceOp.MAX, group=group + ) + if not torch.equal(minimum, maximum) or minimum[0] == 0: + local_error = local_error or ( + f"staged Muon {domain} owner metadata drifted on another rank" + ) + return flat_sizes, first_param, group_rank, local_error + + def _validate_owner_metadata( + self, + ) -> dict[str, tuple[list[int], torch.Tensor | None, int]]: + """Validate both domains, then publish failures to the full chain.""" + if self.pg_collection is None: + raise RuntimeError("staged Muon all-gather has no process groups") + validated: dict[str, tuple[list[int], torch.Tensor | None, int]] = {} + errors: list[str] = [] + for domain, params_list, group in ( + ("dense", self.dp_cp_params_list, self.pg_collection.dp_cp), + ("expert", self.expt_dp_params_list, self.pg_collection.expt_dp), + ): + flat_sizes, first_param, group_rank, error = ( + self._validate_owner_domain( + domain=domain, params_list=params_list, group=group + ) + ) + validated[domain] = flat_sizes, first_param, group_rank + if error is not None: + errors.append(error) + + if torch.distributed.is_initialized(): + status = torch.tensor( + [int(not errors)], + dtype=torch.int64, + device=torch.cuda.current_device(), + ) + torch.distributed.all_reduce( + status, + op=torch.distributed.ReduceOp.MIN, + group=torch.distributed.group.WORLD, + ) + if status[0] == 0: + raise RuntimeError( + "; ".join(errors) + if errors + else "staged Muon owner metadata drifted on another rank" + ) + elif errors: + raise RuntimeError("; ".join(errors)) + return validated + + @torch.no_grad() + def allgather_params(self) -> None: + """All-gather with empty owner shards handled safely.""" + validated = self._validate_owner_metadata() + for domain, params_list, group in ( + ("dense", self.dp_cp_params_list, self.pg_collection.dp_cp), + ("expert", self.expt_dp_params_list, self.pg_collection.expt_dp), + ): + flat_sizes, first_param, group_rank = validated[domain] + if params_list is None or first_param is None or max(flat_sizes) == 0: + continue + device = first_param.device + dtype = first_param.dtype + if device.type != "cuda": + raise RuntimeError( + f"staged Muon {domain} all-gather requires CUDA parameters" + ) + local_params = params_list[group_rank] + src = ( + _flatten_dense_tensors(local_params) + if local_params + else torch.empty(0, device=device, dtype=dtype) + ) + if src.numel() != flat_sizes[group_rank]: + raise RuntimeError( + f"staged Muon {domain} local flat size does not match " + f"owner schema on rank {group_rank}" + ) + gather_list = [ + src + if owner_rank == group_rank + else torch.empty(size, device=device, dtype=dtype) + for owner_rank, size in enumerate(flat_sizes) + ] + torch.distributed.all_gather(gather_list, src, group=group) + for owner_rank, params in enumerate(params_list): + if owner_rank == group_rank or not params: + continue + updated = _unflatten_dense_tensors(gather_list[owner_rank], params) + for source, target in zip(updated, params, strict=True): + target.data.copy_(source) + + @torch.no_grad() + def step(self): + # This validation must precede ChainedOptimizer.prepare_grads(), the + # WORLD norm/clip collectives, and every model/state mutation. + self._validate_owner_metadata() + activities: list[bool] = [] + errors: list[str] = [] + if tuple(self.chained_optimizers) != self._staged_leaves: + errors.append("staged Muon optimizer leaf order changed after bind") + for leaf_index, leaf in enumerate(self._staged_leaves): + try: + active, error = leaf.preflight_step_activity() + activities.append(active) + if error is not None: + errors.append(f"leaf {leaf_index}: {error}") + except BaseException as error: + activities.append(False) + errors.append(f"leaf {leaf_index}: {type(error).__name__}: {error}") + + chain_kinds = tuple( + getattr(leaf.optimizer, "optimizer_kind", type(leaf).__name__) + for leaf in self._staged_leaves + ) + chain_digest = hashlib.sha256(repr(chain_kinds).encode()).digest() + digest_words = [ + int.from_bytes(chain_digest[offset : offset + 8], "little") + & ((1 << 63) - 1) + for offset in range(0, 32, 8) + ] + world = torch.distributed.group.WORLD + vote = torch.tensor( + [int(not errors), int(any(activities)), *digest_words], + dtype=torch.int64, + device=torch.cuda.current_device(), + ) + minimum = vote.clone() + maximum = vote.clone() + torch.distributed.all_reduce( + minimum, op=torch.distributed.ReduceOp.MIN, group=world + ) + torch.distributed.all_reduce( + maximum, op=torch.distributed.ReduceOp.MAX, group=world + ) + if not torch.equal(minimum[2:], maximum[2:]) or minimum[0] == 0: + raise RuntimeError( + "; ".join(errors) + if errors + else "staged Muon step preflight failed on another rank" + ) + + planned: list[Any] = [] + try: + for leaf, active in zip(self._staged_leaves, activities, strict=True): + leaf.set_step_activity(active) + planned.append(leaf) + if maximum[1] == 0: + return ( + True, + 0.0, + 0 if self.config.log_num_zeros_in_grad else None, + ) + return super().step() + finally: + self.drain() + for leaf in reversed(planned): + leaf.clear_step_activity() + + def drain(self) -> None: + for leaf in self.chained_optimizers: + inner = getattr(leaf, "optimizer", None) + if inner is not None: + inner.drain() + + @property + def residency(self) -> str: + states = { + leaf.optimizer.residency + for leaf in self.chained_optimizers + if getattr(leaf, "optimizer", None) is not None + } + return "CPU_RESIDENT" if states <= {"CPU_RESIDENT"} else "STEP_ACTIVE" + + @property + def cuda_state_numel(self) -> int: + return sum( + getattr(leaf.optimizer, "cuda_state_numel", 0) + for leaf in self.chained_optimizers + if getattr(leaf, "optimizer", None) is not None + ) + + return GPUStagedLayerWiseDistributedOptimizer + + +def _group_size_and_rank(group: Any, *, name: str) -> tuple[int, int]: + if group is None: + return 1, 0 + size_fn = getattr(group, "size", None) + rank_fn = getattr(group, "rank", None) + if not callable(size_fn) or not callable(rank_fn): + raise RuntimeError(f"official {name} process group has no size/rank capability") + size = size_fn() + rank = rank_fn() + if isinstance(size, bool) or not isinstance(size, int) or size < 1: + raise RuntimeError(f"official {name} process group has invalid size {size!r}") + if isinstance(rank, bool) or not isinstance(rank, int) or not 0 <= rank < size: + raise RuntimeError( + f"official {name} process group rank {rank!r} is outside [0, {size})" + ) + return size, rank + + +def _validate_owner_lists( + owner_lists: Any, + *, + name: str, + group_size: int, +) -> tuple[tuple[torch.Tensor, ...], ...]: + if not isinstance(owner_lists, Sequence) or isinstance(owner_lists, (str, bytes)): + raise RuntimeError( + f"official {name} owner list must be a rank-indexed sequence" + ) + if len(owner_lists) != group_size: + raise RuntimeError( + f"official {name} owner list length {len(owner_lists)} does not match " + f"group size {group_size}" + ) + normalized = [] + for owner_rank, params in enumerate(owner_lists): + if not isinstance(params, Sequence) or isinstance(params, (str, bytes)): + raise RuntimeError( + f"official {name} owner rank {owner_rank} entry must be a sequence" + ) + if any(not isinstance(param, torch.Tensor) for param in params): + raise RuntimeError( + f"official {name} owner rank {owner_rank} contains a non-tensor" + ) + normalized.append(tuple(params)) + return tuple(normalized) + + +def _parameter_identities(params: Iterable[torch.Tensor]) -> tuple[str, ...]: + return tuple( + f"id={id(param)},shape={tuple(param.shape)},dtype={param.dtype}" + for param in params + ) + + +def _require_same_domain_params( + *, + domain: str, + owner_rank: int, + expected: Iterable[torch.Tensor], + actual: Iterable[torch.Tensor], + scope: str, +) -> None: + expected_tuple = tuple(expected) + actual_tuple = tuple(actual) + expected_ids = {id(param) for param in expected_tuple} + actual_ids = {id(param) for param in actual_tuple} + if expected_ids != actual_ids: + missing = tuple( + param for param in expected_tuple if id(param) not in actual_ids + ) + extra = tuple(param for param in actual_tuple if id(param) not in expected_ids) + raise RuntimeError( + f"official layer-wise {scope} ownership does not match owner lists: " + f"domain={domain}, scope={scope}, owner_rank={owner_rank}, " + f"expected={_parameter_identities(expected_tuple)}, " + f"local={_parameter_identities(actual_tuple)}, " + f"missing={_parameter_identities(missing)}, " + f"extra={_parameter_identities(extra)}" + ) + + +def _validate_official_ownership( + official: Any, + expected_params: Iterable[torch.Tensor] | None = None, +) -> None: + local_dense: list[torch.Tensor] = [] + local_expert: list[torch.Tensor] = [] + for leaf in official.chained_optimizers: + for group in leaf.param_groups: + is_expert_parallel = group.get("is_expert_parallel") + if not isinstance(is_expert_parallel, bool): + raise RuntimeError( + "official layer-wise param group has invalid ownership flag: " + f"is_expert_parallel={is_expert_parallel!r}" + ) + target = local_expert if is_expert_parallel else local_dense + target.extend(group["params"]) + local_params = local_dense + local_expert + if len({id(param) for param in local_dense}) != len(local_dense): + raise RuntimeError( + "official layer-wise ownership contains duplicate local params: " + "domain=dp_cp" + ) + if len({id(param) for param in local_expert}) != len(local_expert): + raise RuntimeError( + "official layer-wise ownership contains duplicate local params: " + "domain=expt_dp" + ) + local_cross_domain = {id(param) for param in local_dense} & { + id(param) for param in local_expert + } + if local_cross_domain: + raise RuntimeError( + "official layer-wise local parameter appears in both ownership domains: " + f"parameter_ids={sorted(local_cross_domain)}" + ) + pg = official.pg_collection + if pg is None or not hasattr(pg, "dp_cp") or not hasattr(pg, "expt_dp"): + raise RuntimeError("official layer-wise ownership has no process groups") + dp_size, dp_rank = _group_size_and_rank(pg.dp_cp, name="dp_cp") + expt_size, expt_rank = _group_size_and_rank(pg.expt_dp, name="expt_dp") + + expected_all = tuple(local_params if expected_params is None else expected_params) + if len({id(param) for param in expected_all}) != len(expected_all): + raise RuntimeError("expected model parameter ownership contains duplicates") + expected_dense = tuple( + param for param in expected_all if getattr(param, "allreduce", True) + ) + expected_expert = tuple( + param for param in expected_all if not getattr(param, "allreduce", True) + ) + + dp_owner_lists: tuple[tuple[torch.Tensor, ...], ...] | None = None + if official.dp_cp_params_list is None: + if dp_size != 1: + raise RuntimeError( + "official dp_cp owner list is missing for ownership group size " + f"{dp_size}" + ) + else: + if dp_size == 1: + raise RuntimeError( + "official dp_cp owner list must be None for ownership group size 1" + ) + dp_owner_lists = _validate_owner_lists( + official.dp_cp_params_list, + name="dp_cp", + group_size=dp_size, + ) + + expt_owner_lists: tuple[tuple[torch.Tensor, ...], ...] | None = None + expert_list_required = dp_size > 1 and expt_size > 1 and bool(expected_expert) + if official.expt_dp_params_list is None: + if expert_list_required: + raise RuntimeError( + "official expt_dp owner list is missing for non-empty expert ownership: " + f"domain=expt_dp, owner_rank={expt_rank}, " + f"expected={_parameter_identities(expected_expert)}" + ) + else: + if not expert_list_required: + raise RuntimeError( + "official expt_dp owner list must be None here: " + f"dp_cp_size={dp_size}, expt_dp_size={expt_size}, " + f"expert_param_count={len(expected_expert)}" + ) + expt_owner_lists = _validate_owner_lists( + official.expt_dp_params_list, + name="expt_dp", + group_size=expt_size, + ) + + dense_owner_params = [param for params in dp_owner_lists or () for param in params] + expert_owner_params = [ + param for params in expt_owner_lists or () for param in params + ] + if len({id(param) for param in dense_owner_params}) != len(dense_owner_params): + raise RuntimeError( + "official dp_cp owner lists contain duplicate params: domain=dp_cp" + ) + if len({id(param) for param in expert_owner_params}) != len(expert_owner_params): + raise RuntimeError( + "official expt_dp owner lists contain duplicate params: domain=expt_dp" + ) + cross_domain = {id(param) for param in dense_owner_params} & { + id(param) for param in expert_owner_params + } + if cross_domain: + raise RuntimeError( + "official layer-wise parameter appears in both ownership domains: " + f"parameter_ids={sorted(cross_domain)}" + ) + + dp_schema = tuple(dense_owner_params) + if dp_owner_lists is None: + dp_schema = tuple(local_dense) + _require_same_domain_params( + domain="dp_cp", + owner_rank=dp_rank, + expected=expected_dense, + actual=dp_schema, + scope="global", + ) + + expt_schema = tuple(expert_owner_params) + if expt_owner_lists is None: + expt_schema = tuple(local_expert) + _require_same_domain_params( + domain="expt_dp", + owner_rank=expt_rank, + expected=expected_expert, + actual=expt_schema, + scope="global", + ) + + expected_local_dense = tuple( + dp_owner_lists[dp_rank] if dp_owner_lists is not None else expected_dense + ) + expected_local_expert = tuple( + expt_owner_lists[expt_rank] if expt_owner_lists is not None else expected_expert + ) + _require_same_domain_params( + domain="dp_cp", + owner_rank=dp_rank, + expected=expected_local_dense, + actual=local_dense, + scope="local", + ) + _require_same_domain_params( + domain="expt_dp", + owner_rank=expt_rank, + expected=expected_local_expert, + actual=local_expert, + scope="local", + ) + + +def _has_materialized_parameter_state(optimizer: Any) -> bool: + """Accept empty defaultdict entries but reject any real parameter state.""" + for state in optimizer.state.values(): + if not isinstance(state, Mapping) or bool(state): + return True + return False + + +def _validate_muon_parallel_topology( + official: Any, + muon_optimizers: Sequence[Any], + expected_params: Sequence[torch.Tensor], + *, + tp_mode: str, +) -> None: + """Validate MCore's TP metadata and DP-owner collective coherence. + + MCore/EO own the actual Newton--Schulz implementation. This preflight only + proves that every TP peer will enter that implementation for the same + owner-held TP-local parameter units and that expert parameters select the + expert TP group. It runs after official layer-wise sharding and before any + staged CPU or CUDA storage is allocated. + """ + pg = official.pg_collection + required_groups = ("tp", "expt_tp", "dp_cp", "expt_dp") + if pg is None: + raise RuntimeError("official Muon topology has no process-group collection") + groups: dict[str, Any] = {} + group_metadata: dict[str, tuple[int, int]] = {} + for name in required_groups: + group = getattr(pg, name, None) + if group is None: + raise RuntimeError( + f"official Muon topology is missing process group {name}" + ) + groups[name] = group + group_metadata[name] = _group_size_and_rank(group, name=name) + + expected_muon = tuple( + param for param in expected_params if _is_muon_parameter(param) + ) + expected_ids = {id(param) for param in expected_muon} + parameter_ordinals = { + id(param): index for index, param in enumerate(expected_params) + } + local_by_dp_domain: dict[str, list[torch.Tensor]] = {"dense": [], "expert": []} + local_by_tp_domain: dict[str, list[torch.Tensor]] = {"dense": [], "expert": []} + for optimizer in muon_optimizers: + if optimizer.pg_collection is not pg: + raise RuntimeError( + "official Muon optimizer and LayerWise wrapper use different " + "process-group collections" + ) + if optimizer.tp_mode != tp_mode: + raise RuntimeError( + "official Muon TP mode does not match staged configuration: " + f"official={optimizer.tp_mode!r}, staged={tp_mode!r}" + ) + for group_index, param_group in enumerate(optimizer.param_groups): + is_expert = param_group.get("is_expert_parallel") + if not isinstance(is_expert, bool): + raise RuntimeError( + "official Muon param group has invalid expert ownership flag: " + f"group_index={group_index}, is_expert_parallel={is_expert!r}" + ) + dp_domain = "expert" if is_expert else "dense" + local_by_dp_domain[dp_domain].extend(param_group["params"]) + for param in param_group["params"]: + tp_domain = "expert" if getattr(param, "expert_tp", False) else "dense" + local_by_tp_domain[tp_domain].append(param) + + local_muon = local_by_dp_domain["dense"] + local_by_dp_domain["expert"] + if any(id(param) not in expected_ids for param in local_muon): + raise RuntimeError( + "official Muon TP topology contains an unknown local parameter" + ) + + for param in expected_muon: + expert_tp = getattr(param, "expert_tp", False) + if not isinstance(expert_tp, bool): + raise RuntimeError( + "official Muon parameter has invalid expert_tp flag: " + f"parameter={_parameter_identities((param,))}, " + f"expert_tp={expert_tp!r}" + ) + tp_group_name = "expt_tp" if expert_tp else "tp" + tp_size, tp_rank = group_metadata[tp_group_name] + tensor_parallel = getattr(param, "tensor_model_parallel", False) + if not isinstance(tensor_parallel, bool): + raise RuntimeError( + "official Muon parameter has invalid tensor_model_parallel flag: " + f"parameter={_parameter_identities((param,))}, " + f"value={tensor_parallel!r}" + ) + partition_dim = getattr(param, "partition_dim", -1) + if isinstance(partition_dim, bool) or partition_dim not in {-1, 0, 1}: + raise RuntimeError( + "official Muon parameter has invalid TP partition_dim: " + f"parameter={_parameter_identities((param,))}, " + f"partition_dim={partition_dim!r}" + ) + has_partition_axis = partition_dim in {0, 1} + if tensor_parallel and not has_partition_axis: + raise RuntimeError( + "official Muon tensor-parallel metadata is inconsistent: " + f"parameter={_parameter_identities((param,))}, " + f"tensor_model_parallel={tensor_parallel}, " + f"partition_dim={partition_dim}, group={tp_group_name}, " + f"group_size={tp_size}, group_rank={tp_rank}" + ) + if not tensor_parallel and has_partition_axis and not expert_tp and tp_size > 1: + raise RuntimeError( + "official dense Muon TP metadata has a partition axis without " + "tensor_model_parallel: " + f"parameter={_parameter_identities((param,))}, " + f"partition_dim={partition_dim}, group_size={tp_size}, " + f"group_rank={tp_rank}" + ) + if has_partition_axis: + partition_stride = getattr(param, "partition_stride", None) + if ( + isinstance(partition_stride, bool) + or not isinstance(partition_stride, int) + or partition_stride < 1 + ): + raise RuntimeError( + "official Muon parameter has invalid TP partition_stride: " + f"parameter={_parameter_identities((param,))}, " + f"partition_stride={partition_stride!r}" + ) + if param.shape[partition_dim] < 1: + raise RuntimeError( + "official Muon TP-local parameter shard is empty: " + f"parameter={_parameter_identities((param,))}, " + f"partition_dim={partition_dim}" + ) + + if not torch.distributed.is_initialized(): + return + + def _require_tp_owner_consensus(domain: str, group_name: str) -> None: + group = groups[group_name] + group_size, group_rank = group_metadata[group_name] + signature = tuple( + ( + parameter_ordinals[id(param)], + tuple(param.shape), + getattr(param, "partition_dim", -1), + ) + for param in local_by_tp_domain[domain] + ) + if group_size == 1: + return + gathered: list[Any] = [None] * group_size + torch.distributed.all_gather_object(gathered, signature, group=group) + if any(peer != gathered[0] for peer in gathered[1:]): + raise RuntimeError( + "official Muon TP peers selected different DP-owned units: " + f"domain={domain}, group={group_name}, group_rank={group_rank}, " + f"local={signature}, peers={gathered}" + ) + + _require_tp_owner_consensus("dense", "tp") + _require_tp_owner_consensus("expert", "expt_tp") + + def _require_dp_schema_consensus( + domain: str, + group_name: str, + owner_lists: Sequence[Sequence[torch.Tensor]] | None, + ) -> None: + group = groups[group_name] + group_size, group_rank = group_metadata[group_name] + if group_size == 1: + return + signature = ( + None + if owner_lists is None + else tuple( + tuple( + ( + parameter_ordinals[id(param)], + tuple(param.shape), + str(param.dtype), + ) + for param in owner_params + ) + for owner_params in owner_lists + ) + ) + gathered: list[Any] = [None] * group_size + torch.distributed.all_gather_object(gathered, signature, group=group) + if any(peer != gathered[0] for peer in gathered[1:]): + raise RuntimeError( + "official Muon DP owner schemas differ across replicas: " + f"domain={domain}, group={group_name}, group_rank={group_rank}, " + f"local={signature}, peers={gathered}" + ) + + _require_dp_schema_consensus("dense", "dp_cp", official.dp_cp_params_list) + _require_dp_schema_consensus("expert", "expt_dp", official.expt_dp_params_list) + + +def _is_muon_parameter(param: torch.Tensor) -> bool: + """Match MCore's Muon predicate for the empty-Muon edge.""" + return ( + not getattr(param, "is_embedding_or_output_parameter", False) + and param.ndim == 2 + ) + + +def get_megatron_optimizer_with_gpu_staged_muon( + mcore_config: Any, + model: list[Any], + staged_config: GPUStagedMuonConfig, + *, + pg_collection: Any | None = None, + config_overrides: dict[Any, Any] | None = None, +) -> Any: + """Run official classification/ownership, then bind staged rank-local state.""" + version = importlib.metadata.version("megatron-core") + if version != _SUPPORTED_MEGATRON_CORE_VERSION: + raise RuntimeError( + f"GPU-staged Muon supports megatron-core 0.19.0 exactly, found {version}" + ) + if mcore_config.use_distributed_optimizer: + raise ValueError("staged Muon uses official layer-wise ownership, not dist-opt") + if ( + mcore_config.optimizer != "muon" + or not mcore_config.use_layer_wise_distributed_optimizer + ): + raise ValueError("staged Muon requires the MCore layer-wise Muon optimizer") + if mcore_config.muon_scalar_optimizer != "adam": + raise ValueError("staged Muon requires AdamW for scalar parameters") + if not mcore_config.decoupled_weight_decay: + raise ValueError("staged Muon MVP requires decoupled weight decay") + if not mcore_config.bf16 or mcore_config.fp16: + raise ValueError("staged Muon MVP requires BF16 without FP16 loss scaling") + if mcore_config.optimizer_cuda_graph: + raise ValueError("staged Muon does not support optimizer CUDA graphs") + if mcore_config.overlap_param_gather: + raise ValueError("staged Muon MVP does not support overlap_param_gather") + if getattr(mcore_config, "overlap_param_gather_with_optimizer_step", False): + raise ValueError( + "staged Muon MVP does not support overlap_param_gather_with_optimizer_step" + ) + if mcore_config.use_precision_aware_optimizer: + raise ValueError("staged Muon owns precision explicitly; PAO must be disabled") + if mcore_config.muon_tp_mode != "duplicated": + raise ValueError( + "staged Muon currently supports only the verified duplicated TP mode; " + f"found {mcore_config.muon_tp_mode!r}" + ) + + # This is deliberately before importing/calling either optimizer builder and + # before touching ``model``. Multiple staging streams have not yet been + # validated against either dense-TP or expert-TP collective ordering. An + # explicit collection is authoritative; otherwise MCore's already-created + # topology is queried without inspecting model chunks. + from megatron.core.process_groups_config import ProcessGroupCollection + + topology_pg_collection = ( + pg_collection + if pg_collection is not None + else ProcessGroupCollection.use_mpu_process_groups( + required_pgs=["tp", "expt_tp"] + ) + ) + tp_size, _ = _group_size_and_rank(topology_pg_collection.tp, name="tp") + expt_tp_size, _ = _group_size_and_rank( + topology_pg_collection.expt_tp, name="expt_tp" + ) + if (tp_size > 1 or expt_tp_size > 1) and staged_config.buffer_count != 1: + raise ValueError( + "staged Muon TP/expert-TP collectives require buffer_count=1; " + f"tp_size={tp_size}, expt_tp_size={expt_tp_size}, " + f"buffer_count={staged_config.buffer_count}" + ) + + from megatron.core.optimizer import ( + HAVE_EMERGING_OPTIMIZERS, + get_megatron_optimizer, + ) + from megatron.core.optimizer.emerging_optimizers import TensorParallelMuon + from megatron.core.optimizer.layer_wise_optimizer import ( + LayerWiseDistributedOptimizer, + ) + + if not HAVE_EMERGING_OPTIMIZERS: + raise ImportError( + "MCore 0.19 Muon requires its optional emerging-optimizers backend; " + "the staged factory will not substitute a different algorithm" + ) + eo_version = importlib.metadata.version("emerging-optimizers") + if eo_version != _SUPPORTED_EMERGING_OPTIMIZERS_VERSION: + raise RuntimeError( + "staged Muon supports emerging-optimizers 0.3.0 exactly, " + f"found {eo_version}" + ) + trainable = [] + checkpoint_parameter_names: dict[torch.Tensor, str] = {} + for model_index, chunk in enumerate(model): + module = chunk.module if hasattr(chunk, "module") else chunk + prefix = "model" if len(model) == 1 else f"model{model_index}" + named_parameters = getattr(module, "named_parameters", None) + if callable(named_parameters): + parameters_with_names = named_parameters() + else: + parameters = getattr(module, "parameters", None) + if not callable(parameters): + raise TypeError("Muon model chunk has no parameter iterator") + parameters_with_names = ( + (f"parameter_{index}", param) + for index, param in enumerate(parameters()) + ) + for name, param in parameters_with_names: + stable_name = f"{prefix}.{name}" + previous = checkpoint_parameter_names.setdefault(param, stable_name) + if previous != stable_name: + raise RuntimeError( + "Muon checkpoint parameter has multiple stable model names: " + f"{previous!r} and {stable_name!r}" + ) + if param.requires_grad: + trainable.append(param) + if not trainable: + raise ValueError("staged Muon requires at least one trainable parameter") + if any(not param.is_cuda for param in trainable): + raise ValueError("staged Muon requires CUDA model parameters") + if len({param.device for param in trainable}) != 1: + raise ValueError("staged Muon requires one CUDA device per process") + device = trainable[0].device + + build_config = copy.copy(mcore_config) + build_config.bf16 = False + build_config.fp16 = False + build_config.use_distributed_optimizer = False + build_config.use_precision_aware_optimizer = False + build_config.use_precision_aware_optimizer_no_fp8_or_ds_fp8 = False + build_config.main_grads_dtype = torch.float32 + build_config.main_params_dtype = torch.float32 + build_config.exp_avg_dtype = torch.float32 + build_config.exp_avg_sq_dtype = torch.float32 + build_config.muon_scalar_optimizer = "adam" + build_config.overlap_param_gather = False + build_config.overlap_param_gather_with_optimizer_step = False + + has_muon_parameters = any(_is_muon_parameter(p) for p in trainable) + official = get_megatron_optimizer( + build_config, + model, + use_gloo_process_groups=False, + pg_collection=pg_collection, + config_overrides=config_overrides, + ) + + if type(official) is not LayerWiseDistributedOptimizer: + raise TypeError( + "official Muon builder did not return LayerWiseDistributedOptimizer" + ) + if official.overlap_param_gather or official.use_buffer_param_sync: + raise RuntimeError("staged Muon MVP requires synchronous official all-gather") + _validate_official_ownership(official, trainable) + + # MCore 0.19 shards the raw Muon/Adam optimizers before wrapping BF16 + # leaves. The copied build config disables that wrapper so these leaves + # expose the official rank-local parameter groups without GPU master state. + official_bases = list(official.chained_optimizers) + for base in official_bases: + if _has_materialized_parameter_state(base): + raise RuntimeError("official Muon leaf allocated state before staged bind") + muon_leaf_count = sum( + isinstance(base, TensorParallelMuon) for base in official_bases + ) + expected_muon_leaf_count = 1 if has_muon_parameters else 0 + if muon_leaf_count != expected_muon_leaf_count: + raise RuntimeError( + "official Muon classification produced " + f"{muon_leaf_count} Muon leaves, expected {expected_muon_leaf_count}" + ) + _validate_muon_parallel_topology( + official, + tuple(base for base in official_bases if isinstance(base, TensorParallelMuon)), + tuple(trainable), + tp_mode=mcore_config.muon_tp_mode, + ) + if not has_muon_parameters: + official_bases.insert(0, _EMPTY_MUON_LEAF) + + # MCore's clip/norm helpers use this derived capability to read the FP32 + # ``decoupled_grad`` attached by our leaf wrapper. Keep this on an + # instance-local config copy so a rejected factory never mutates the + # caller's configuration. + staged_wrapper_config = copy.copy(mcore_config) + staged_wrapper_config.use_precision_aware_optimizer_no_fp8_or_ds_fp8 = True + staged_wrapper_config.overlap_param_gather = False + staged_wrapper_config.overlap_param_gather_with_optimizer_step = False + leaf_cls = _make_layerwise_leaf_class() + staged_leaves = [] + for base in official_bases: + if base is _EMPTY_MUON_LEAF: + staged_inner = GPUStagedEmptyOptimizer("muon") + staged_inner.bind_parallel_groups( + tp=official.pg_collection.tp, + expt_tp=official.pg_collection.expt_tp, + ) + elif isinstance(base, TensorParallelMuon): + staged_inner = GPUStagedMuon( + base.param_groups, + staged_config=staged_config, + weight_decay_method=base.weight_decay_method, + native_optimizer=base, + ) + staged_inner.bind_parallel_groups( + tp=official.pg_collection.tp, + expt_tp=official.pg_collection.expt_tp, + ) + staged_inner.bind_owned_params(base.param_groups, empty_device=device) + elif base is None or not any( + group.get("params") for group in base.param_groups + ): + staged_inner = GPUStagedEmptyOptimizer("scalar_adamw") + else: + staged_inner = GPUStagedAdamW( + base.param_groups, + lr=mcore_config.lr, + betas=(mcore_config.adam_beta1, mcore_config.adam_beta2), + eps=mcore_config.adam_eps, + weight_decay=mcore_config.weight_decay, + staged_config=GPUStagedAdamWConfig( + buffer_count=staged_config.buffer_count, + bucket_size_mb=staged_config.slot_size_mb, + ), + adam_w_mode=mcore_config.decoupled_weight_decay, + master_weights=True, + use_decoupled_grad=True, + ) + staged_inner.bind_owned_params( + base.param_groups, + empty_device=device, + ) + staged_inner.optimizer_kind = "scalar_adamw" + staged_leaves.append(leaf_cls(staged_inner, staged_wrapper_config, device)) + + staged_kinds = [leaf.optimizer.optimizer_kind for leaf in staged_leaves] + if ( + not staged_kinds + or staged_kinds[0] != "muon" + or any(kind != "scalar_adamw" for kind in staged_kinds[1:]) + ): + raise RuntimeError( + "staged Muon chain topology must be [muon, scalar_adamw...], " + f"found {staged_kinds}" + ) + + staged_layerwise_cls = _make_staged_layerwise_class() + result = staged_layerwise_cls(official, staged_leaves) + result.configure_managed_checkpoint_schema( + checkpoint_parameter_names, + algorithm={ + "momentum": mcore_config.muon_momentum, + "use_nesterov": mcore_config.muon_nesterov, + "fp32_matmul_prec": mcore_config.muon_fp32_matmul_prec, + "coefficient_type": mcore_config.muon_coefficient_type, + "num_ns_steps": mcore_config.muon_num_ns_steps, + "scale_mode": mcore_config.muon_scale_mode, + "split_qkv": mcore_config.muon_split_qkv, + "tp_mode": mcore_config.muon_tp_mode, + "extra_scale_factor": mcore_config.muon_extra_scale_factor, + }, + ) + if result.cuda_state_numel != 0 or result.residency != "CPU_RESIDENT": + raise RuntimeError("staged Muon initialization did not reach CPU residency") + return result diff --git a/areal/tools/validate_docker_installation.py b/areal/tools/validate_docker_installation.py index 0b98010baf..ec4d9aff6f 100755 --- a/areal/tools/validate_docker_installation.py +++ b/areal/tools/validate_docker_installation.py @@ -37,6 +37,10 @@ class DockerInstallationValidator(BaseInstallationValidator): "megatron.core.parallel_state", "megatron.core.tensor_parallel", ], + "emerging-optimizers": [ + "emerging_optimizers.orthogonalized_optimizers", + "emerging_optimizers.orthogonalized_optimizers.muon_utils", + ], # DeepSeek-V3 related packages "flash_mla": ["flash_mla"], "deep_gemm": ["deep_gemm"], @@ -55,6 +59,7 @@ class DockerInstallationValidator(BaseInstallationValidator): "transformer_engine", "flash_attn_3", "megatron-core", + "emerging-optimizers", "mbridge", "megatron-bridge", "causal_conv1d", @@ -138,6 +143,11 @@ def parse_pyproject(self): ) # Megatron packages (from cuda-train > megatron extra) + self.add_additional_package( + "emerging-optimizers", + opt_versions.get("emerging-optimizers", ""), + required=True, + ) self.add_additional_package( "megatron-core", opt_versions.get("megatron-core", ""), diff --git a/docs/en/cli_reference.md b/docs/en/cli_reference.md index dce6a3866b..57f6befa74 100644 --- a/docs/en/cli_reference.md +++ b/docs/en/cli_reference.md @@ -87,6 +87,7 @@ For detailed examples, see the experiment configurations in the `examples/` dire - [MOPDTeacher Specification](section-mopd-teacher) - [MegatronEngine Configuration](section-megatron-engine) - [MemoryProfiler Configuration](section-memory-profiler) +- [MuonOptimizer Configuration](section-muon-optimizer) - [PRMAdvantageShaping Configuration](section-prm-advantage-shaping) - [PRM Configuration](section-prm) - [PRMScorer Configuration](section-prm-scorer) @@ -336,23 +337,24 @@ Configuration for reward/advantage normalization. Configuration for model optimization during training. -| Parameter | Type | Default | Description | -| ------------------------- | --------------- | ------------ | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `type` | string | `"adam"` | Optimizer type. For FSDP Engine, adam_bf16 enables memory-efficient BF16 optimizer states. For Megatron Engine, adam_bf16 requires dtype=bfloat16 and is automatically converted to adam with precision-aware optimizer enabled. **Choices:** `adam`, `sgd`, `adam_bf16` | -| `lr` | float | `0.001` | Learning rate | -| `weight_decay` | float | `0.01` | Weight decay | -| `beta1` | float | `0.9` | Adam beta1 parameter. Only effective when optimizer_type is adam/adam_bf16 | -| `beta2` | float | `0.999` | Adam beta2 parameter. Only effective when optimizer_type is adam/adam_bf16 | -| `eps` | float | `1e-08` | Adam epsilon parameter. Only effective when optimizer_type is adam/adam_bf16 | -| `min_lr_ratio` | float | `0.0` | Minimum learning rate ratio after annealing | -| `lr_scheduler_type` | string | `"constant"` | Learning rate scheduler type **Choices:** `linear`, `cosine`, `constant` | -| `warmup_steps_proportion` | float | `0.001` | Non-negative proportion of training steps for warmup. Ignored when warmup_steps is set. For Megatron, the resolved warmup steps must be less than the total training steps. | -| `warmup_steps` | integer \| None | `None` | Fixed number of learning-rate scheduler steps for warmup. Must be non-negative. For Megatron, it must also be less than the total training steps. When both options are explicitly configured, warmup_steps takes precedence over warmup_steps_proportion and a warning is emitted. | -| `initial_loss_scale` | float | `4294967296` | Initial loss scaling factor | -| `min_loss_scale` | float | `1.0` | Minimum loss scaling factor | -| `loss_scale_window` | float | `5` | Window size for loss scaling adjustment | -| `hysteresis` | integer | `2` | Hysteresis (scaling factor) for loss scaling | -| `gradient_clipping` | float | `1.0` | Gradient clipping threshold | +| Parameter | Type | Default | Description | +| ------------------------- | ----------------------------------------------- | --------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `type` | string | `"adam"` | Optimizer type. For FSDP Engine, adam_bf16 enables memory-efficient BF16 optimizer states. For Megatron Engine, adam_bf16 requires dtype=bfloat16 and is automatically converted to adam with precision-aware optimizer enabled. dist_muon selects Megatron's layer-wise distributed Muon optimizer and uses AdamW for scalar/non-matrix parameters. **Choices:** `adam`, `sgd`, `adam_bf16`, `dist_muon` | +| `lr` | float | `0.001` | Learning rate | +| `weight_decay` | float | `0.01` | Weight decay | +| `beta1` | float | `0.9` | Adam beta1 parameter. Only effective when optimizer_type is adam/adam_bf16 | +| `beta2` | float | `0.999` | Adam beta2 parameter. Only effective when optimizer_type is adam/adam_bf16 | +| `eps` | float | `1e-08` | Adam epsilon parameter. Only effective when optimizer_type is adam/adam_bf16 | +| `muon` | [`MuonOptimizerConfig`](section-muon-optimizer) | *MuonOptimizerConfig* | Muon algorithm settings used when type=dist_muon. | +| `min_lr_ratio` | float | `0.0` | Minimum learning rate ratio after annealing | +| `lr_scheduler_type` | string | `"constant"` | Learning rate scheduler type **Choices:** `linear`, `cosine`, `constant` | +| `warmup_steps_proportion` | float | `0.001` | Non-negative proportion of training steps for warmup. Ignored when warmup_steps is set. For Megatron, the resolved warmup steps must be less than the total training steps. | +| `warmup_steps` | integer \| None | `None` | Fixed number of learning-rate scheduler steps for warmup. Must be non-negative. For Megatron, it must also be less than the total training steps. When both options are explicitly configured, warmup_steps takes precedence over warmup_steps_proportion and a warning is emitted. | +| `initial_loss_scale` | float | `4294967296` | Initial loss scaling factor | +| `min_loss_scale` | float | `1.0` | Minimum loss scaling factor | +| `loss_scale_window` | float | `5` | Window size for loss scaling adjustment | +| `hysteresis` | integer | `2` | Hysteresis (scaling factor) for loss scaling | +| `gradient_clipping` | float | `1.0` | Gradient clipping threshold | (section-ppo-actor)= @@ -1329,6 +1331,24 @@ entries for torch.cuda.memory.\_record_memory_history. | `profile_steps` | list of integer | *list* | List of global steps to capture memory snapshots. | | `max_entries` | integer | `100000` | Max entries for memory history ring buffer. | +(section-muon-optimizer)= + +## MuonOptimizer Configuration + +Algorithm settings for Megatron's distributed Muon optimizer. + +| Parameter | Type | Default | Description | +| -------------------- | ------- | -------------- | ----------------------------------------------------------- | +| `momentum` | float | `0.95` | - | +| `use_nesterov` | boolean | `False` | - | +| `fp32_matmul_prec` | string | `"medium"` | - **Choices:** `low`, `medium`, `high` | +| `coefficient_type` | string | `"quintic"` | - | +| `num_ns_steps` | integer | `5` | - | +| `scale_mode` | string | `"spectral"` | - **Choices:** `spectral`, `unit_rms_norm`, `shape_scaling` | +| `split_qkv` | boolean | `True` | - | +| `tp_mode` | string | `"duplicated"` | - **Choices:** `blockwise`, `duplicated`, `distributed` | +| `extra_scale_factor` | float | `1.0` | - | + (section-prm-advantage-shaping)= ## PRMAdvantageShaping Configuration diff --git a/docs/zh/cli_reference.md b/docs/zh/cli_reference.md index 49a7b67098..06ccf923c4 100644 --- a/docs/zh/cli_reference.md +++ b/docs/zh/cli_reference.md @@ -85,6 +85,7 @@ python3 train.py --config path/to/config.yaml actor.lr=1e-4 seed=42 - [MOPDTeacher Specification](section-mopd-teacher) - [MegatronEngine Configuration](section-megatron-engine) - [MemoryProfiler Configuration](section-memory-profiler) +- [MuonOptimizer Configuration](section-muon-optimizer) - [PRMAdvantageShaping Configuration](section-prm-advantage-shaping) - [PRM Configuration](section-prm) - [PRMScorer Configuration](section-prm-scorer) @@ -334,23 +335,24 @@ Configuration for reward/advantage normalization. Configuration for model optimization during training. -| Parameter | Type | Default | Description | -| ------------------------- | --------------- | ------------ | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `type` | string | `"adam"` | Optimizer type. For FSDP Engine, adam_bf16 enables memory-efficient BF16 optimizer states. For Megatron Engine, adam_bf16 requires dtype=bfloat16 and is automatically converted to adam with precision-aware optimizer enabled. **Choices:** `adam`, `sgd`, `adam_bf16` | -| `lr` | float | `0.001` | Learning rate | -| `weight_decay` | float | `0.01` | Weight decay | -| `beta1` | float | `0.9` | Adam beta1 parameter. Only effective when optimizer_type is adam/adam_bf16 | -| `beta2` | float | `0.999` | Adam beta2 parameter. Only effective when optimizer_type is adam/adam_bf16 | -| `eps` | float | `1e-08` | Adam epsilon parameter. Only effective when optimizer_type is adam/adam_bf16 | -| `min_lr_ratio` | float | `0.0` | Minimum learning rate ratio after annealing | -| `lr_scheduler_type` | string | `"constant"` | Learning rate scheduler type **Choices:** `linear`, `cosine`, `constant` | -| `warmup_steps_proportion` | float | `0.001` | Non-negative proportion of training steps for warmup. Ignored when warmup_steps is set. For Megatron, the resolved warmup steps must be less than the total training steps. | -| `warmup_steps` | integer \| None | `None` | Fixed number of learning-rate scheduler steps for warmup. Must be non-negative. For Megatron, it must also be less than the total training steps. When both options are explicitly configured, warmup_steps takes precedence over warmup_steps_proportion and a warning is emitted. | -| `initial_loss_scale` | float | `4294967296` | Initial loss scaling factor | -| `min_loss_scale` | float | `1.0` | Minimum loss scaling factor | -| `loss_scale_window` | float | `5` | Window size for loss scaling adjustment | -| `hysteresis` | integer | `2` | Hysteresis (scaling factor) for loss scaling | -| `gradient_clipping` | float | `1.0` | Gradient clipping threshold | +| Parameter | Type | Default | Description | +| ------------------------- | ----------------------------------------------- | --------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `type` | string | `"adam"` | Optimizer type. For FSDP Engine, adam_bf16 enables memory-efficient BF16 optimizer states. For Megatron Engine, adam_bf16 requires dtype=bfloat16 and is automatically converted to adam with precision-aware optimizer enabled. dist_muon selects Megatron's layer-wise distributed Muon optimizer and uses AdamW for scalar/non-matrix parameters. **Choices:** `adam`, `sgd`, `adam_bf16`, `dist_muon` | +| `lr` | float | `0.001` | Learning rate | +| `weight_decay` | float | `0.01` | Weight decay | +| `beta1` | float | `0.9` | Adam beta1 parameter. Only effective when optimizer_type is adam/adam_bf16 | +| `beta2` | float | `0.999` | Adam beta2 parameter. Only effective when optimizer_type is adam/adam_bf16 | +| `eps` | float | `1e-08` | Adam epsilon parameter. Only effective when optimizer_type is adam/adam_bf16 | +| `muon` | [`MuonOptimizerConfig`](section-muon-optimizer) | *MuonOptimizerConfig* | Muon algorithm settings used when type=dist_muon. | +| `min_lr_ratio` | float | `0.0` | Minimum learning rate ratio after annealing | +| `lr_scheduler_type` | string | `"constant"` | Learning rate scheduler type **Choices:** `linear`, `cosine`, `constant` | +| `warmup_steps_proportion` | float | `0.001` | Non-negative proportion of training steps for warmup. Ignored when warmup_steps is set. For Megatron, the resolved warmup steps must be less than the total training steps. | +| `warmup_steps` | integer \| None | `None` | Fixed number of learning-rate scheduler steps for warmup. Must be non-negative. For Megatron, it must also be less than the total training steps. When both options are explicitly configured, warmup_steps takes precedence over warmup_steps_proportion and a warning is emitted. | +| `initial_loss_scale` | float | `4294967296` | Initial loss scaling factor | +| `min_loss_scale` | float | `1.0` | Minimum loss scaling factor | +| `loss_scale_window` | float | `5` | Window size for loss scaling adjustment | +| `hysteresis` | integer | `2` | Hysteresis (scaling factor) for loss scaling | +| `gradient_clipping` | float | `1.0` | Gradient clipping threshold | (section-ppo-actor)= @@ -1327,6 +1329,24 @@ entries for torch.cuda.memory.\_record_memory_history. | `profile_steps` | list of integer | *list* | List of global steps to capture memory snapshots. | | `max_entries` | integer | `100000` | Max entries for memory history ring buffer. | +(section-muon-optimizer)= + +## MuonOptimizer Configuration + +Algorithm settings for Megatron's distributed Muon optimizer. + +| Parameter | Type | Default | Description | +| -------------------- | ------- | -------------- | ----------------------------------------------------------- | +| `momentum` | float | `0.95` | - | +| `use_nesterov` | boolean | `False` | - | +| `fp32_matmul_prec` | string | `"medium"` | - **Choices:** `low`, `medium`, `high` | +| `coefficient_type` | string | `"quintic"` | - | +| `num_ns_steps` | integer | `5` | - | +| `scale_mode` | string | `"spectral"` | - **Choices:** `spectral`, `unit_rms_norm`, `shape_scaling` | +| `split_qkv` | boolean | `True` | - | +| `tp_mode` | string | `"duplicated"` | - **Choices:** `blockwise`, `duplicated`, `distributed` | +| `extra_scale_factor` | float | `1.0` | - | + (section-prm-advantage-shaping)= ## PRMAdvantageShaping Configuration diff --git a/examples/cpu_staged_offload/README.md b/examples/cpu_staged_offload/README.md index 0b95d0256e..40c43ec42d 100644 --- a/examples/cpu_staged_offload/README.md +++ b/examples/cpu_staged_offload/README.md @@ -24,6 +24,31 @@ optimizer, precision-aware optimizer semantics, and FP32 master/moment state. En it selects precision-aware mode automatically. Staged optimizer checkpoint saves are synchronous; `megatron.async_save=true` is rejected. +The Muon optimizer can use the same staging infrastructure: + +```yaml +actor: + optimizer: + type: dist_muon + muon: + momentum: 0.95 + num_ns_steps: 5 + tp_mode: duplicated + megatron: + ddp: + use_distributed_optimizer: false + cpu_staged_offload: + enabled: true + buffer_count: 1 + bucket_size_mb: 128 +``` + +Set `cpu_staged_offload.enabled: false` for native layer-wise Muon. Both modes use +Megatron-Core 0.19.0 and emerging-optimizers 0.3.0 and require BF16 and synchronous +parameter gather. CPU staging supports `tp_mode: duplicated`; TP or expert-TP greater +than one requires `buffer_count: 1`. Staged checkpoints support DP ownership resharding +with fixed model-parallel topology; native checkpoints require fixed DP topology. + Checkpoint loading is fail-stop. DCP writes optimizer state into the authoritative CPU slabs in place. If loading fails, the process must terminate and AReaL recovery starts a new process from the last complete checkpoint; no in-process disk snapshot, rollback, or diff --git a/pyproject.toml b/pyproject.toml index 364ce596c6..463304d6d9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -159,6 +159,7 @@ kernels = [ ] megatron = [ "megatron-core==0.19.0; python_version >= '3.12' and sys_platform == 'linux' and platform_machine == 'x86_64'", + "emerging-optimizers==0.3.0; python_version >= '3.12' and sys_platform == 'linux' and platform_machine == 'x86_64'", "mbridge @ git+https://github.com/ISEEKYAN/mbridge.git@310e8fb; sys_platform == 'linux' and platform_machine == 'x86_64'", "megatron-bridge==0.6.0; python_version >= '3.12' and sys_platform == 'linux' and platform_machine == 'x86_64'", ] diff --git a/pyproject.vllm.toml b/pyproject.vllm.toml index 7ac2442c3e..ceb58a061d 100644 --- a/pyproject.vllm.toml +++ b/pyproject.vllm.toml @@ -171,6 +171,7 @@ kernels = [ ] megatron = [ "megatron-core==0.19.0; python_version >= '3.12' and sys_platform == 'linux' and platform_machine == 'x86_64'", + "emerging-optimizers==0.3.0; python_version >= '3.12' and sys_platform == 'linux' and platform_machine == 'x86_64'", "mbridge @ git+https://github.com/ISEEKYAN/mbridge.git@310e8fb; sys_platform == 'linux' and platform_machine == 'x86_64'", "megatron-bridge==0.6.0; python_version >= '3.12' and sys_platform == 'linux' and platform_machine == 'x86_64'", ] diff --git a/tests/megatron/test_gpu_staged_muon_dp.py b/tests/megatron/test_gpu_staged_muon_dp.py new file mode 100644 index 0000000000..9c23b29441 --- /dev/null +++ b/tests/megatron/test_gpu_staged_muon_dp.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Layer-wise DP-owner parity for GPU-staged Muon.""" + +from __future__ import annotations + +import os +import subprocess +import sys + +import pytest +import torch + +from areal.infra.platforms import current_platform +from areal.utils.network import find_free_ports + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="CUDA not available" +) + +_PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +_WORKER = "tests/megatron/torchrun/run_gpu_staged_muon_dp.py" + + +@pytest.mark.multi_gpu +@pytest.mark.slow +@pytest.mark.parametrize( + ("world_size", "tp_size"), + [(2, 1), (4, 2)], + ids=["dp2-tp1", "dp2-tp2"], +) +def test_gpu_staged_muon_matches_layerwise_dp2_owners( + tmp_path_factory, world_size: int, tp_size: int +) -> None: + """Dense, expert, scalar, and empty-owner DP paths match native MCore.""" + if current_platform.device_count() < world_size: + pytest.skip(f"staged Muon layer-wise parity requires {world_size} GPUs") + output = ( + tmp_path_factory.mktemp(f"gpu_staged_muon_dp_{world_size}_{tp_size}") + / "result.txt" + ) + port = find_free_ports(1)[0] + env = os.environ.copy() + env["PYTHONPATH"] = _PROJECT_ROOT + os.pathsep + env.get("PYTHONPATH", "") + env["NCCL_DEBUG"] = "WARN" + subprocess.run( + [ + sys.executable, + "-m", + "torch.distributed.run", + f"--nproc_per_node={world_size}", + "--nnodes=1", + "--master-addr=localhost", + f"--master-port={port}", + _WORKER, + f"--output={output}", + f"--tp-size={tp_size}", + ], + check=True, + cwd=_PROJECT_ROOT, + env=env, + text=True, + stdout=sys.stdout, + stderr=sys.stdout, + ) + assert output.read_text().strip() == "Passed" diff --git a/tests/megatron/test_gpu_staged_muon_tp.py b/tests/megatron/test_gpu_staged_muon_tp.py new file mode 100644 index 0000000000..ba7b83621c --- /dev/null +++ b/tests/megatron/test_gpu_staged_muon_tp.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Numerical parity between staged and official Muon under real TP collectives.""" + +from __future__ import annotations + +import os +import subprocess +import sys + +import pytest +import torch + +from areal.infra.platforms import current_platform +from areal.utils.network import find_free_ports + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="CUDA not available" +) + +_PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +_WORKER = "tests/megatron/torchrun/run_gpu_staged_muon_tp.py" + + +@pytest.mark.multi_gpu +@pytest.mark.slow +@pytest.mark.parametrize( + ("world_size", "tp_size"), + [(2, 2), (4, 2)], + ids=["dp1-tp2", "dp2-tp2"], +) +def test_gpu_staged_muon_matches_official_parallel_grid( + tmp_path_factory, world_size: int, tp_size: int +) -> None: + """Dense, expert, QKV, and DP-replicated updates match official Muon.""" + if current_platform.device_count() < world_size: + pytest.skip(f"staged Muon parity requires {world_size} GPUs") + output = ( + tmp_path_factory.mktemp(f"gpu_staged_muon_{world_size}_{tp_size}") + / "result.txt" + ) + port = find_free_ports(1)[0] + env = os.environ.copy() + env["PYTHONPATH"] = _PROJECT_ROOT + os.pathsep + env.get("PYTHONPATH", "") + env["NCCL_DEBUG"] = "WARN" + subprocess.run( + [ + sys.executable, + "-m", + "torch.distributed.run", + f"--nproc_per_node={world_size}", + "--nnodes=1", + "--master-addr=localhost", + f"--master-port={port}", + _WORKER, + f"--output={output}", + f"--tp-size={tp_size}", + ], + check=True, + cwd=_PROJECT_ROOT, + env=env, + text=True, + stdout=sys.stdout, + stderr=sys.stdout, + ) + assert output.read_text().strip() == "Passed" diff --git a/tests/megatron/torchrun/run_gpu_staged_muon_dp.py b/tests/megatron/torchrun/run_gpu_staged_muon_dp.py new file mode 100644 index 0000000000..57a98ef602 --- /dev/null +++ b/tests/megatron/torchrun/run_gpu_staged_muon_dp.py @@ -0,0 +1,381 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 + +"""torchrun worker for staged Muon layer-wise DP numerical parity.""" + +from __future__ import annotations + +import argparse +import copy +import os +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +from megatron.core import parallel_state +from megatron.core.dist_checkpointing.mapping import ShardedTensor +from megatron.core.optimizer import get_standard_config_overrides +from megatron.core.optimizer.optimizer_config import OptimizerConfig, ParamKey +from megatron.core.optimizer_param_scheduler import ( + OptimizerParamScheduler, + ParamGroupOverride, +) +from megatron.core.process_groups_config import ProcessGroupCollection + +from areal.engine.megatron_utils.checkpointer import MegatronCheckpointManager +from areal.engine.megatron_utils.gpu_staged_muon import ( + GPUStagedMuonConfig, + get_megatron_optimizer_with_dist_muon, + get_megatron_optimizer_with_gpu_staged_muon, +) + + +class _TinyMuonModel(torch.nn.Module): + def __init__(self, initial: dict[str, torch.Tensor], tp_size: int) -> None: + super().__init__() + self.dense_weight = torch.nn.Parameter(initial["dense_weight"].clone()) + self.experts_weight = torch.nn.Parameter(initial["experts_weight"].clone()) + self.experts_weight.allreduce = False + self.bias = torch.nn.Parameter(initial["bias"].clone()) + self.config = SimpleNamespace( + num_attention_heads=4, + num_query_groups=2, + kv_channels=2, + context_parallel_size=1, + ) + self.ddp_config = SimpleNamespace( + use_megatron_fsdp=False, + use_distributed_optimizer=False, + num_distributed_optimizer_instances=1, + ) + if tp_size > 1: + self.dense_weight.tensor_model_parallel = True + self.dense_weight.partition_dim = 0 + self.dense_weight.partition_stride = 1 + self.experts_weight.tensor_model_parallel = False + self.experts_weight.partition_dim = 0 + self.experts_weight.partition_stride = 1 + + def sharded_state_dict(self) -> dict: + # This checkpoint regression uses TP1 and fixed DP2. + return { + name: ShardedTensor.from_rank_offsets( + name, param, replica_id=(0, 0, parallel_state.get_data_parallel_rank()) + ) + for name, param in self.named_parameters() + } + + +class _SingleParamModel(torch.nn.Module): + def __init__(self, initial: torch.Tensor, *, expert: bool, tp_size: int) -> None: + super().__init__() + name = "experts_weight" if expert else "bias" + param = torch.nn.Parameter(initial.clone()) + param.no_weight_decay = True + setattr(self, name, param) + if expert: + param.allreduce = False + if tp_size > 1: + param.partition_dim = 0 + param.partition_stride = 1 + self.config = SimpleNamespace( + num_attention_heads=4, + num_query_groups=2, + kv_channels=2, + context_parallel_size=1, + ) + self.ddp_config = SimpleNamespace( + use_megatron_fsdp=False, + use_distributed_optimizer=False, + num_distributed_optimizer_instances=1, + ) + + +def _make_pg_collection() -> ProcessGroupCollection: + return ProcessGroupCollection.use_mpu_process_groups( + required_pgs=[ + "tp", + "expt_tp", + "pp", + "cp", + "ep", + "mp", + "tp_ep_pp", + "dp", + "dp_cp", + "expt_dp", + ] + ) + + +def _make_config() -> OptimizerConfig: + return OptimizerConfig( + optimizer="muon", + use_layer_wise_distributed_optimizer=True, + lr=0.02, + min_lr=0.0, + weight_decay=0.015, + adam_beta1=0.8, + adam_beta2=0.95, + adam_eps=1e-6, + bf16=True, + fp16=False, + use_distributed_optimizer=False, + use_precision_aware_optimizer=False, + decoupled_weight_decay=True, + muon_scalar_optimizer="adam", + muon_momentum=0.82, + muon_nesterov=True, + muon_split_qkv=False, + muon_fp32_matmul_prec="highest", + muon_num_ns_steps=3, + muon_tp_mode="duplicated", + overlap_param_gather=False, + overlap_param_gather_with_optimizer_step=False, + clip_grad=0.0, + ) + + +def _run_dp2_outer(tp_size: int, checkpoint_path: Path) -> None: + if dist.get_world_size() != 2 * tp_size: + raise ValueError("staged Muon outer worker requires DP size 2") + device = torch.device("cuda", int(os.environ["LOCAL_RANK"])) + tp_rank = parallel_state.get_tensor_model_parallel_rank() + generator = torch.Generator(device=device).manual_seed(20260902 + tp_rank) + initial = { + "dense_weight": torch.randn( + 5, 7, generator=generator, device=device, dtype=torch.bfloat16 + ), + "experts_weight": torch.randn( + 6, 4, generator=generator, device=device, dtype=torch.bfloat16 + ), + "bias": torch.randn( + 5, generator=generator, device=device, dtype=torch.bfloat16 + ), + } + staged_model = _TinyMuonModel(initial, tp_size).to(device) + baseline_model = _TinyMuonModel(initial, tp_size).to(device) + pg_collection = _make_pg_collection() + + staged_config = _make_config() + staged = get_megatron_optimizer_with_gpu_staged_muon( + staged_config, + [staged_model], + GPUStagedMuonConfig(buffer_count=1, slot_size_mb=1), + pg_collection=pg_collection, + ) + baseline_config = _make_config() + baseline = get_megatron_optimizer_with_dist_muon( + baseline_config, + [baseline_model], + pg_collection=pg_collection, + ) + + if staged.dp_cp_params_list is None or staged.expt_dp_params_list is None: + raise AssertionError("DP2 must materialize both dense and expert owner lists") + if not any(len(owner) == 0 for owner in staged.expt_dp_params_list): + raise AssertionError("expert DP case must contain an empty-owner rank") + staged_kinds = [leaf.optimizer.optimizer_kind for leaf in staged.chained_optimizers] + if staged_kinds != ["muon", "scalar_adamw"]: + raise AssertionError(f"unexpected staged optimizer chain: {staged_kinds}") + + staged_by_name = dict(staged_model.named_parameters()) + baseline_by_name = dict(baseline_model.named_parameters()) + for step in range(3): + if step == 1: + _reallocate_model_storage(staged_model.dense_weight) + for param_index, name in enumerate(staged_by_name): + grad = torch.randn( + staged_by_name[name].shape, + generator=generator, + device=device, + dtype=torch.bfloat16, + ).mul_(step + param_index + 1) + staged_by_name[name].main_grad = grad + baseline_by_name[name].main_grad = grad.clone() + + baseline_success, _, _ = baseline.step() + staged_success, _, _ = staged.step() + if not baseline_success or not staged_success: + raise AssertionError("layer-wise optimizer step failed") + for name, staged_param in staged_by_name.items(): + torch.testing.assert_close( + staged_param, + baseline_by_name[name], + rtol=4e-3, + atol=4e-3, + ) + replicas = [ + torch.empty_like(staged_param) + for _ in range(dist.get_world_size(pg_collection.dp)) + ] + dist.all_gather(replicas, staged_param.detach(), group=pg_collection.dp) + for replica in replicas[1:]: + torch.testing.assert_close(replica, replicas[0], rtol=0.0, atol=0.0) + + if tp_size == 1: + _check_checkpoint(baseline_model, baseline, checkpoint_path) + staged.bind_managed_checkpoint_process_group(dist.new_group(backend="gloo")) + _check_checkpoint( + staged_model, staged, checkpoint_path.parent / "staged_checkpoint" + ) + _check_owner_rebinding_rejected(staged) + _run_single_param_case(pg_collection, tp_size, expert=False) + _run_single_param_case(pg_collection, tp_size, expert=True) + + +def _reallocate_model_storage(param: torch.nn.Parameter) -> None: + """Reproduce AWEX release/resume while retaining Tensor and Storage identity.""" + value = param.detach().cpu().clone() + storage = param.untyped_storage() + nbytes, old_ptr = storage.nbytes(), storage.data_ptr() + torch.cuda.synchronize() + storage.resize_(0) + # Hold allocations so resume cannot reuse the released address. + blockers = [ + torch.empty(nbytes, dtype=torch.uint8, device=param.device) for _ in range(32) + ] + storage.resize_(nbytes) + assert param.untyped_storage() is storage + assert storage.data_ptr() != old_ptr + with torch.no_grad(): + param.copy_(value) + del blockers + + +def _check_owner_rebinding_rejected(optimizer) -> None: + owners = optimizer.dp_cp_params_list + owner = next(params for params in owners if params) + original = owner[0] + owner[0] = torch.nn.Parameter(original.detach().clone()) + try: + with pytest.raises(RuntimeError, match="changed after bind"): + optimizer.allgather_params() + finally: + owner[0] = original + storage = original.untyped_storage() + original.data = original.detach().clone() + try: + with pytest.raises(RuntimeError, match="storage"): + optimizer.allgather_params() + finally: + original.data = torch.empty( + 0, dtype=original.dtype, device=original.device + ).set_(storage, original.storage_offset(), original.shape, original.stride()) + + +def _check_checkpoint(model, optimizer, checkpoint_path: Path) -> None: + """The real manager saves and resumes either Muon path at fixed DP.""" + manager = MegatronCheckpointManager( + torch.nn.ModuleList([model]), optimizer, None, use_distributed_optimizer=False + ) + checkpoint_path.mkdir(parents=True, exist_ok=True) + manager.save_checkpoint(str(checkpoint_path), with_rng=False) + saved_params = copy.deepcopy(model.state_dict()) + assert optimizer.step()[0] + expected = copy.deepcopy(model.state_dict()) + manager.load_checkpoint(str(checkpoint_path), with_rng=False) + for name, param in model.named_parameters(): + torch.testing.assert_close(param, saved_params[name], rtol=0.0, atol=0.0) + assert optimizer.step()[0] + for name, param in model.named_parameters(): + torch.testing.assert_close(param, expected[name], rtol=0.0, atol=0.0) + manager.close() + + +def _run_single_param_case( + pg_collection: ProcessGroupCollection, tp_size: int, *, expert: bool +) -> None: + """Cover scalar-only staging and the native all-expert gather edge.""" + tp_rank = parallel_state.get_tensor_model_parallel_rank() + device = torch.device("cuda", int(os.environ["LOCAL_RANK"])) + generator = torch.Generator(device=device).manual_seed(20260915 + tp_rank) + shape = (4, 4) if expert else (4,) + initial = torch.randn( + shape, generator=generator, device=device, dtype=torch.bfloat16 + ) + staged_model = _SingleParamModel(initial, expert=expert, tp_size=tp_size) + baseline_model = _SingleParamModel(initial, expert=expert, tp_size=tp_size) + overrides = get_standard_config_overrides(_make_config()) + overrides[ParamKey(attr="no_weight_decay")] = ParamGroupOverride(wd_mult=0.0) + staged = get_megatron_optimizer_with_gpu_staged_muon( + _make_config(), + [staged_model], + GPUStagedMuonConfig(buffer_count=1, slot_size_mb=1), + pg_collection=pg_collection, + config_overrides=dict(overrides), + ) + baseline = get_megatron_optimizer_with_dist_muon( + _make_config(), + [baseline_model], + pg_collection=pg_collection, + config_overrides=dict(overrides), + ) + config = _make_config() + for optimizer in (baseline, staged): + assert all(group["wd_mult"] == 0.0 for group in optimizer.param_groups) + OptimizerParamScheduler( + optimizer, + init_lr=0.0, + max_lr=config.lr, + min_lr=config.min_lr, + lr_warmup_steps=0, + lr_decay_steps=2, + lr_decay_style="constant", + start_wd=config.weight_decay, + end_wd=config.weight_decay, + wd_incr_steps=2, + wd_incr_style="constant", + ) + assert all(group["weight_decay"] == 0.0 for group in optimizer.param_groups) + expected_kinds = ["muon"] if expert else ["muon", "scalar_adamw"] + assert [leaf.optimizer.optimizer_kind for leaf in staged.chained_optimizers] == ( + expected_kinds + ) + if expert: + assert staged.dp_cp_params_list is not None + assert all(not owner for owner in staged.dp_cp_params_list) + staged_param = next(staged_model.parameters()) + baseline_param = next(baseline_model.parameters()) + for step in range(2): + grad = torch.randn( + shape, generator=generator, device=device, dtype=torch.bfloat16 + ) + staged_param.main_grad = grad.mul(step + 1) + baseline_param.main_grad = staged_param.main_grad.clone() + assert baseline.step()[0] + assert staged.step()[0] + torch.testing.assert_close(staged_param, baseline_param, rtol=4e-3, atol=4e-3) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--output", required=True) + parser.add_argument("--tp-size", type=int, required=True) + args = parser.parse_args() + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group("nccl", device_id=torch.device("cuda", local_rank)) + parallel_state.initialize_model_parallel( + tensor_model_parallel_size=args.tp_size, + pipeline_model_parallel_size=1, + context_parallel_size=1, + expert_model_parallel_size=1, + expert_tensor_parallel_size=args.tp_size, + create_gloo_process_groups=False, + ) + try: + _run_dp2_outer(args.tp_size, Path(args.output).parent / "checkpoint") + dist.barrier() + if dist.get_rank() == 0: + with open(args.output, "w") as file: + file.write("Passed") + finally: + parallel_state.destroy_model_parallel() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/megatron/torchrun/run_gpu_staged_muon_tp.py b/tests/megatron/torchrun/run_gpu_staged_muon_tp.py new file mode 100644 index 0000000000..d2dde0487c --- /dev/null +++ b/tests/megatron/torchrun/run_gpu_staged_muon_tp.py @@ -0,0 +1,220 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 + +"""torchrun worker for staged Muon parallel numerical parity.""" + +from __future__ import annotations + +import argparse +import os +from types import SimpleNamespace + +import torch +import torch.distributed as dist +from megatron.core.optimizer.emerging_optimizers import TensorParallelMuon + +from areal.engine.megatron_utils.gpu_staged_muon import ( + GPUStagedMuon, + GPUStagedMuonConfig, +) + + +def _make_parallel_groups(tp_size: int) -> tuple[dist.ProcessGroup, dist.ProcessGroup]: + world_size = dist.get_world_size() + if world_size % tp_size != 0: + raise ValueError( + f"world size {world_size} is not divisible by TP size {tp_size}" + ) + + rank = dist.get_rank() + tp_group = None + for dp_rank in range(world_size // tp_size): + ranks = list(range(dp_rank * tp_size, (dp_rank + 1) * tp_size)) + group = dist.new_group(ranks) + if rank in ranks: + tp_group = group + + dp_group = None + for tp_rank in range(tp_size): + ranks = list(range(tp_rank, world_size, tp_size)) + group = dist.new_group(ranks) + if rank in ranks: + dp_group = group + + assert tp_group is not None + assert dp_group is not None + return tp_group, dp_group + + +def _make_singleton_groups() -> dist.ProcessGroup: + rank = dist.get_rank() + local_group = None + for group_rank in range(dist.get_world_size()): + group = dist.new_group([group_rank]) + if rank == group_rank: + local_group = group + assert local_group is not None + return local_group + + +def _run_case( + *, + name: str, + tp_group: dist.ProcessGroup, + expt_tp_group: dist.ProcessGroup, + dp_group: dist.ProcessGroup, + shape: tuple[int, int], + partition_dim: int, + expert_tp: bool = False, + split_qkv: bool = False, +) -> None: + tp_rank = dist.get_rank(expt_tp_group if expert_tp else tp_group) + device = torch.device("cuda", int(os.environ["LOCAL_RANK"])) + generator = torch.Generator(device=device).manual_seed( + 20260831 + sum(ord(char) for char in name) + tp_rank + ) + initial = torch.randn( + shape, generator=generator, device=device, dtype=torch.bfloat16 + ) + staged_param = torch.nn.Parameter(initial.clone()) + baseline_param = torch.nn.Parameter(initial.float()) + for param in (staged_param, baseline_param): + param.tensor_model_parallel = not expert_tp + param.partition_dim = partition_dim + param.partition_stride = 1 + param.expert_tp = expert_tp + param.is_qkv = split_qkv + + pg_collection = SimpleNamespace( + tp=tp_group, + expt_tp=expt_tp_group, + gtp_remat=None, + expt_gtp_remat=None, + ) + kwargs = { + "lr": 0.025, + "momentum": 0.82, + "nesterov": True, + "weight_decay": 0.015, + "use_decoupled_weight_decay": True, + "fp32_matmul_prec": "highest", + "num_ns_steps": 3, + "pg_collection": pg_collection, + "tp_mode": "duplicated", + "split_qkv": split_qkv, + "is_qkv_fn": lambda param: getattr(param, "is_qkv", False), + "qkv_split_shapes": (4, 2, 2) if split_qkv else None, + } + param_group = { + "params": [baseline_param], + "is_expert_parallel": expert_tp, + } + baseline = TensorParallelMuon([param_group], **kwargs) + staged = GPUStagedMuon( + [ + { + "params": [staged_param], + "lr": kwargs["lr"], + "momentum": kwargs["momentum"], + "weight_decay": kwargs["weight_decay"], + "is_expert_parallel": expert_tp, + } + ], + staged_config=GPUStagedMuonConfig( + buffer_count=1, + slot_size_mb=staged_param.numel() * 4 / (1024 * 1024), + ), + weight_decay_method=baseline.weight_decay_method, + native_optimizer=baseline, + ) + staged.bind_parallel_groups(tp=tp_group, expt_tp=expt_tp_group) + staged.bind_owned_params(staged.param_groups) + + for step in range(3): + grad = torch.randn( + shape, generator=generator, device=device, dtype=torch.bfloat16 + ).mul_(step + 1) + staged_param.decoupled_grad = grad + baseline_param.grad = grad.float() + baseline.step() + staged.step() + staged.drain() + + staged_state = staged.state[staged_param] + baseline_state = baseline.state[baseline_param] + torch.testing.assert_close( + staged_state["master_param"], + baseline_param.detach().cpu(), + rtol=3e-6, + atol=3e-6, + ) + torch.testing.assert_close( + staged_state["momentum_buffer"], + baseline_state["momentum_buffer"].cpu(), + rtol=3e-6, + atol=3e-6, + ) + torch.testing.assert_close( + staged_param, + baseline_param.detach().bfloat16(), + rtol=0.0, + atol=0.0, + ) + if dist.get_world_size(dp_group) > 1: + replicas = [ + torch.empty_like(staged_param) + for _ in range(dist.get_world_size(dp_group)) + ] + dist.all_gather(replicas, staged_param.detach(), group=dp_group) + for replica in replicas[1:]: + torch.testing.assert_close(replica, replicas[0], rtol=0.0, atol=0.0) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--output", required=True) + parser.add_argument("--tp-size", type=int, required=True) + args = parser.parse_args() + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group("nccl", device_id=torch.device("cuda", local_rank)) + try: + tp_group, dp_group = _make_parallel_groups(args.tp_size) + singleton_group = _make_singleton_groups() + for partition_dim, shape in ((0, (5, 7)), (1, (7, 5))): + _run_case( + name=f"dense-axis-{partition_dim}", + tp_group=tp_group, + expt_tp_group=tp_group, + dp_group=dp_group, + shape=shape, + partition_dim=partition_dim, + ) + _run_case( + name="qkv-split", + tp_group=tp_group, + expt_tp_group=tp_group, + dp_group=dp_group, + shape=(16, 6), + partition_dim=0, + split_qkv=True, + ) + _run_case( + name="expert-tp", + tp_group=singleton_group, + expt_tp_group=tp_group, + dp_group=dp_group, + shape=(6, 7), + partition_dim=0, + expert_tp=True, + ) + dist.barrier() + if dist.get_rank() == 0: + with open(args.output, "w") as file: + file.write("Passed") + finally: + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/test_gpu_staged_muon.py b/tests/test_gpu_staged_muon.py new file mode 100644 index 0000000000..1b96675336 --- /dev/null +++ b/tests/test_gpu_staged_muon.py @@ -0,0 +1,472 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import gc +from types import SimpleNamespace + +import pytest +import torch +from megatron.core.optimizer.emerging_optimizers import TensorParallelMuon + +from areal.engine.megatron_utils.gpu_staged_muon import ( + GPUStagedMuon, + GPUStagedMuonConfig, + _validate_muon_parallel_topology, + _validate_official_ownership, +) +from areal.engine.megatron_utils.gpu_staged_optimizer import ( + GPUStagedAdamW, + GPUStagedAdamWConfig, +) + + +def _config(*, slot_numel: int = 64, buffer_count: int = 2) -> GPUStagedMuonConfig: + return GPUStagedMuonConfig( + buffer_count=buffer_count, + slot_size_mb=slot_numel * 4 / (1024 * 1024), + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_muon_owned_matrices_are_indivisible_pinned_slab_units() -> None: + """Each official-owned 2D matrix must map to one complete staged unit.""" + params = [ + torch.nn.Parameter(torch.randn(4, 5, device="cuda", dtype=torch.bfloat16)), + torch.nn.Parameter(torch.randn(3, 7, device="cuda", dtype=torch.bfloat16)), + ] + groups = [ + { + "params": params, + "lr": 0.02, + "momentum": 0.9, + "weight_decay": 0.01, + } + ] + optimizer = GPUStagedMuon( + groups, + native_optimizer=TensorParallelMuon(groups, use_decoupled_weight_decay=True), + staged_config=_config(slot_numel=32), + weight_decay_method="decoupled", + ) + optimizer.bind_owned_params(optimizer.param_groups) + + assert [unit.param for unit in optimizer.units] == params + assert [unit.numel for unit in optimizer.units] == [20, 21] + assert optimizer.residency == "CPU_RESIDENT" + assert optimizer.cuda_state_numel == 0 + assert optimizer.cpu_slabs is not None + assert optimizer.cpu_slabs.master.dtype is torch.float32 + assert optimizer.cpu_slabs.momentum.dtype is torch.float32 + assert optimizer.cpu_slabs.master.is_pinned() + assert optimizer.cpu_slabs.momentum.is_pinned() + for param in params: + assert set(optimizer.state[param]) == {"master_param", "momentum_buffer"} + for value in optimizer.state[param].values(): + assert value.device.type == "cpu" + assert value.dtype is torch.float32 + first_state = optimizer.state[params[0]]["master_param"] + assert ( + first_state.untyped_storage().data_ptr() + == optimizer.cpu_slabs.master.untyped_storage().data_ptr() + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_muon_and_scalar_adamw_state_schemas_are_isolated() -> None: + """Official classification leaves receive disjoint Muon and AdamW schemas.""" + matrix = torch.nn.Parameter(torch.ones(4, 4, device="cuda", dtype=torch.bfloat16)) + scalar = torch.nn.Parameter(torch.ones(4, device="cuda", dtype=torch.bfloat16)) + muon = GPUStagedMuon( + [{"params": [matrix], "lr": 0.1, "momentum": 0.9, "weight_decay": 0.0}], + native_optimizer=TensorParallelMuon( + [{"params": [matrix], "lr": 0.1, "momentum": 0.9, "weight_decay": 0.0}], + use_decoupled_weight_decay=True, + ), + staged_config=_config(), + weight_decay_method="decoupled", + ) + muon.bind_owned_params(muon.param_groups) + adam = GPUStagedAdamW( + [scalar], + lr=0.1, + staged_config=GPUStagedAdamWConfig(buffer_count=1, bucket_size_mb=1), + ) + adam.bind_owned_params(adam.param_groups) + + assert set(muon.state[matrix]) == {"master_param", "momentum_buffer"} + assert set(adam.state[scalar]) == {"master_param", "exp_avg", "exp_avg_sq"} + assert "exp_avg" not in muon.state[matrix] + assert "momentum_buffer" not in adam.state[scalar] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_muon_gpu_residency_is_bounded_by_max_unit_not_total_state() -> None: + """More owner matrices grow only CPU slabs, not resident CUDA slots.""" + observed_staging = [] + observed_cpu_state = [] + observed_optimizer_peak_bytes = [] + for matrix_count in (1, 8): + gc.collect() + torch.cuda.empty_cache() + params = [ + torch.nn.Parameter(torch.ones(8, 8, device="cuda", dtype=torch.bfloat16)) + for _ in range(matrix_count) + ] + before_optimizer = torch.cuda.memory_allocated() + torch.cuda.reset_peak_memory_stats() + optimizer = GPUStagedMuon( + [ + { + "params": params, + "lr": 0.1, + "momentum": 0.9, + "weight_decay": 0.0, + } + ], + native_optimizer=TensorParallelMuon( + [ + { + "params": params, + "lr": 0.1, + "momentum": 0.9, + "weight_decay": 0.0, + } + ], + use_decoupled_weight_decay=True, + ), + staged_config=_config(slot_numel=64, buffer_count=2), + weight_decay_method="decoupled", + ) + optimizer.bind_owned_params(optimizer.param_groups) + observed_optimizer_peak_bytes.append( + torch.cuda.max_memory_allocated() - before_optimizer + ) + observed_staging.append(optimizer.gpu_staging_numel) + assert optimizer.cpu_slabs is not None + observed_cpu_state.append( + optimizer.cpu_slabs.master.numel() + optimizer.cpu_slabs.momentum.numel() + ) + assert optimizer.cuda_state_numel == 0 + del optimizer, params + + assert observed_staging == [2 * 64 * 3, 2 * 64 * 3] + assert observed_cpu_state == [2 * 64, 2 * 8 * 64] + assert observed_optimizer_peak_bytes[1] == observed_optimizer_peak_bytes[0] + + +@pytest.mark.parametrize("use_nesterov", [False, True]) +@pytest.mark.parametrize("schedule_hyperparameters", [False, True]) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_muon_steps_match_official_tensor_parallel_muon( + use_nesterov: bool, + schedule_hyperparameters: bool, +) -> None: + """Staged updates match native Muon with fixed or scheduled settings.""" + torch.manual_seed(20260829) + initial = [ + torch.randn(shape, device="cuda", dtype=torch.bfloat16) + for shape in ((5, 7), (8, 3)) + ] + staged_params = [torch.nn.Parameter(value.clone()) for value in initial] + baseline_params = [torch.nn.Parameter(value.float()) for value in initial] + kwargs = { + "lr": 0.03, + "momentum": 0.8, + "nesterov": use_nesterov, + "weight_decay": 0.02, + "use_decoupled_weight_decay": True, + "fp32_matmul_prec": "highest", + "coefficient_type": "quintic", + "num_ns_steps": 3, + "scale_mode": "spectral", + "extra_scale_factor": 1.0, + "tp_mode": "duplicated", + } + baseline = TensorParallelMuon(baseline_params, **kwargs) + staged = GPUStagedMuon( + [ + { + "params": staged_params, + "lr": kwargs["lr"], + "momentum": kwargs["momentum"], + "weight_decay": kwargs["weight_decay"], + } + ], + staged_config=_config(slot_numel=64, buffer_count=1), + weight_decay_method=baseline.weight_decay_method, + native_optimizer=baseline, + ) + staged.bind_owned_params(staged.param_groups) + + for step in range(5): + if schedule_hyperparameters: + settings = { + "lr": (0.0, 0.01, 0.03, 0.015, 0.0)[step], + "weight_decay": (0.02, 0.04, 0.0, 0.01, 0.03)[step], + "momentum": (0.8, 0.85, 0.9, 0.85, 0.8)[step], + } + for group in (*staged.param_groups, *baseline.param_groups): + group.update(settings) + for param_index, (staged_param, baseline_param) in enumerate( + zip(staged_params, baseline_params, strict=True) + ): + grad = torch.randn_like(staged_param).mul_(0.1 + step + param_index) + staged_param.decoupled_grad = grad + baseline_param.grad = grad.float() + staged.step() + baseline.step() + staged.drain() + + for staged_param, baseline_param in zip( + staged_params, baseline_params, strict=True + ): + staged_state = staged.state[staged_param] + baseline_state = baseline.state[baseline_param] + torch.testing.assert_close( + staged_state["master_param"], + baseline_param.detach().cpu(), + rtol=3e-6, + atol=3e-6, + ) + torch.testing.assert_close( + staged_state["momentum_buffer"], + baseline_state["momentum_buffer"].cpu(), + rtol=3e-6, + atol=3e-6, + ) + torch.testing.assert_close( + staged_param, + baseline_param.detach().bfloat16(), + rtol=0.0, + atol=0.0, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_muon_checkpoint_resume_matches_official_tensor_parallel_muon() -> None: + """Muon CPU state resumes on the official optimization trajectory.""" + torch.manual_seed(20260830) + initial = torch.randn(6, 5, device="cuda", dtype=torch.bfloat16) + kwargs = { + "lr": 0.02, + "momentum": 0.85, + "nesterov": True, + "weight_decay": 0.03, + "use_decoupled_weight_decay": True, + "fp32_matmul_prec": "highest", + "num_ns_steps": 3, + "tp_mode": "duplicated", + } + baseline_param = torch.nn.Parameter(initial.float()) + baseline = TensorParallelMuon([baseline_param], **kwargs) + source_param = torch.nn.Parameter(initial.clone()) + + def make_staged(param: torch.nn.Parameter) -> GPUStagedMuon: + optimizer = GPUStagedMuon( + [ + { + "params": [param], + "lr": kwargs["lr"], + "momentum": kwargs["momentum"], + "weight_decay": kwargs["weight_decay"], + } + ], + staged_config=_config(slot_numel=32, buffer_count=1), + weight_decay_method=baseline.weight_decay_method, + native_optimizer=baseline, + ) + optimizer.bind_owned_params(optimizer.param_groups) + return optimizer + + source = make_staged(source_param) + for _ in range(3): + grad = torch.randn_like(source_param) + source_param.decoupled_grad = grad + baseline_param.grad = grad.float() + source.step() + baseline.step() + source.drain() + + live_checkpoint = source.state_dict() + checkpoint = { + "state": { + state_id: {key: value.clone() for key, value in state.items()} + for state_id, state in live_checkpoint["state"].items() + }, + "param_groups": [dict(group) for group in live_checkpoint["param_groups"]], + } + resumed_param = torch.nn.Parameter(source_param.detach().clone()) + resumed = make_staged(resumed_param) + resumed.begin_checkpoint_load() + resumed.load_state_dict(checkpoint) + resumed.complete_checkpoint_load() + resumed.offload_to_cpu() + resumed.restore_from_cpu() + assert resumed.cuda_state_numel == 0 + assert resumed.cpu_slabs.master.is_pinned() + assert resumed.cpu_slabs.momentum.is_pinned() + + for _ in range(3): + grad = torch.randn_like(resumed_param) + resumed_param.decoupled_grad = grad + baseline_param.grad = grad.float() + resumed.step() + baseline.step() + resumed.drain() + + resumed_state = resumed.state[resumed_param] + baseline_state = baseline.state[baseline_param] + torch.testing.assert_close( + resumed_state["master_param"], + baseline_param.detach().cpu(), + rtol=3e-6, + atol=3e-6, + ) + torch.testing.assert_close( + resumed_state["momentum_buffer"], + baseline_state["momentum_buffer"].cpu(), + rtol=3e-6, + atol=3e-6, + ) + torch.testing.assert_close( + resumed_param, + baseline_param.detach().bfloat16(), + rtol=0.0, + atol=0.0, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_muon_rejects_duplicate_nonmatrix_and_undersized_slot() -> None: + """Malformed ownership and split-prone capacity fail before state is usable.""" + matrix = torch.nn.Parameter(torch.ones(8, 8, device="cuda", dtype=torch.bfloat16)) + + def make(groups, slot_numel=128): + return GPUStagedMuon( + groups, + native_optimizer=TensorParallelMuon( + groups, use_decoupled_weight_decay=True + ), + staged_config=_config(slot_numel=slot_numel), + weight_decay_method="decoupled", + ) + + duplicate = make( + [{"params": [matrix, matrix], "lr": 0.1, "momentum": 0.9, "weight_decay": 0.0}] + ) + with pytest.raises(ValueError, match="duplicate"): + duplicate.bind_owned_params(duplicate.param_groups) + + vector = torch.nn.Parameter(torch.ones(8, device="cuda", dtype=torch.bfloat16)) + nonmatrix = make( + [{"params": [vector], "lr": 0.1, "momentum": 0.9, "weight_decay": 0.0}] + ) + with pytest.raises(ValueError, match="2D matrix"): + nonmatrix.bind_owned_params(nonmatrix.param_groups) + + undersized = make( + [{"params": [matrix], "lr": 0.1, "momentum": 0.9, "weight_decay": 0.0}], + slot_numel=63, + ) + with pytest.raises(ValueError, match="slot is too small"): + undersized.bind_owned_params(undersized.param_groups) + + +def test_official_ownership_validation_rejects_duplicate_and_missing_owner() -> None: + """Official owner lists must map each local leaf parameter exactly once.""" + param = torch.nn.Parameter(torch.ones(2, 2)) + leaf = SimpleNamespace( + param_groups=[{"params": [param], "is_expert_parallel": False}] + ) + process_group = SimpleNamespace(rank=lambda: 0, size=lambda: 2) + pg_collection = SimpleNamespace(dp_cp=process_group, expt_dp=process_group) + + duplicate = SimpleNamespace( + chained_optimizers=[leaf], + pg_collection=pg_collection, + dp_cp_params_list=[[param], [param]], + expt_dp_params_list=None, + ) + with pytest.raises(RuntimeError, match="owner lists contain duplicate"): + _validate_official_ownership(duplicate) + + missing = SimpleNamespace( + chained_optimizers=[leaf], + pg_collection=pg_collection, + dp_cp_params_list=[[], []], + expt_dp_params_list=None, + ) + with pytest.raises(RuntimeError, match="does not match owner lists"): + _validate_official_ownership(missing) + + distributed_group = SimpleNamespace(rank=lambda: 0, size=lambda: 2) + distributed_without_owner_lists = SimpleNamespace( + chained_optimizers=[leaf], + pg_collection=SimpleNamespace( + dp_cp=distributed_group, + expt_dp=distributed_group, + ), + dp_cp_params_list=None, + expt_dp_params_list=None, + ) + param_before = param.detach().clone() + group_params_before = tuple(leaf.param_groups[0]["params"]) + with pytest.raises(RuntimeError, match="owner list"): + _validate_official_ownership(distributed_without_owner_lists) + torch.testing.assert_close(param, param_before, rtol=0.0, atol=0.0) + assert all( + actual is expected + for actual, expected in zip( + leaf.param_groups[0]["params"], group_params_before, strict=True + ) + ) + + +def test_official_ownership_validation_accepts_missing_lists_only_for_dp1() -> None: + """MCore's list elision is legal only for an actual size-one DP group.""" + param = torch.nn.Parameter(torch.ones(2, 2)) + leaf = SimpleNamespace( + param_groups=[{"params": [param], "is_expert_parallel": False}] + ) + process_group = SimpleNamespace(rank=lambda: 0, size=lambda: 1) + official = SimpleNamespace( + chained_optimizers=[leaf], + pg_collection=SimpleNamespace( + dp_cp=process_group, + expt_dp=process_group, + ), + dp_cp_params_list=None, + expt_dp_params_list=None, + ) + + _validate_official_ownership(official, [param]) + + +def test_muon_topology_accepts_explicit_expert_partition_metadata() -> None: + """MCore expert comm keeps the partition axis without the dense TP flag.""" + param = torch.nn.Parameter(torch.ones(2, 2)) + param.expert_tp = True + param.tensor_model_parallel = False + param.partition_dim = 0 + param.partition_stride = 1 + singleton = SimpleNamespace(rank=lambda: 0, size=lambda: 1) + pg_collection = SimpleNamespace( + tp=singleton, + expt_tp=singleton, + dp_cp=singleton, + expt_dp=singleton, + ) + optimizer = SimpleNamespace( + pg_collection=pg_collection, + tp_mode="duplicated", + param_groups=[{"params": [param], "is_expert_parallel": True}], + ) + official = SimpleNamespace(pg_collection=pg_collection) + + _validate_muon_parallel_topology( + official, + [optimizer], + [param], + tp_mode="duplicated", + ) diff --git a/tests/test_gpu_staged_muon_checkpoint.py b/tests/test_gpu_staged_muon_checkpoint.py new file mode 100644 index 0000000000..b6313c44b4 --- /dev/null +++ b/tests/test_gpu_staged_muon_checkpoint.py @@ -0,0 +1,112 @@ +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from areal.engine.megatron_utils.gpu_staged_muon import merge_muon_checkpoint_metadata + + +def _metadata(rank: int) -> dict: + groups = {} + for name in ("dp", "dp_cp", "tp", "cp", "ep", "expt_tp", "expt_dp", "pp"): + members = [0, 1] if name in {"dp", "dp_cp", "expt_dp"} else [rank] + groups[name] = { + "size": len(members), + "rank": members.index(rank), + "members": members, + } + parameters = [] + for index, name in enumerate( + (f"norm_{rank}", "embedding") if rank == 0 else (f"norm_{rank}",) + ): + parameters.append( + { + "name": name, + "domain": "dense", + "coordinate": {"pp": 0, "tp": 0}, + "shape": [2], + "dtype": "torch.bfloat16", + "state_kinds": ["master_param", "exp_avg", "exp_avg_sq"], + "source_owner": { + "global_rank": rank, + "owner_rank": rank, + "owner_ordinal": index, + "group_index": index, + "parameter_index": 0, + "unit_order": index, + }, + } + ) + return { + "schema_version": 2, + "megatron_core_version": "0.19.0", + "emerging_optimizers_version": "0.3.0", + "topology": {"world_size": 2, "global_rank": rank, "groups": groups}, + "algorithm": {}, + "leaf_tree": [ + { + "tree_path": [0], + "kind": "scalar_adamw", + "parameters": parameters, + "param_groups": [ + {"lr": 0.02, "weight_decay": 0.0, "step": 1}, + {"lr": 0.02, "weight_decay": 0.1, "step": 1 - rank}, + ], + } + ], + } + + +@pytest.mark.parametrize("ranks", [[0, 1], [1, 0]]) +def test_checkpoint_partially_empty_groups_select_nonempty_owner( + ranks: list[int], +) -> None: + """Empty-owner step counters cannot override authoritative optimizer state.""" + metadata = [_metadata(rank) for rank in ranks] + merged = merge_muon_checkpoint_metadata(metadata, trusted_global_ranks=ranks) + groups = merged["leaf_tree"][0]["param_groups"] + assert [group["step"] for group in groups] == [1, 1] + assert len(merged["leaf_tree"][0]["parameters"]) == 3 + assert metadata[ranks.index(1)]["leaf_tree"][0]["param_groups"][1]["step"] == 0 + + +@pytest.mark.parametrize("field,value", [("step", 2), ("lr", 0.03)]) +def test_checkpoint_nonempty_owners_conflict_raises(field: str, value: float) -> None: + """Real owners must still agree on counters and optimizer settings.""" + metadata = [_metadata(0), _metadata(1)] + metadata[1]["leaf_tree"][0]["param_groups"][0][field] = value + with pytest.raises(ValueError, match="conflicts across owners"): + merge_muon_checkpoint_metadata(metadata, trusted_global_ranks=[0, 1]) + + +@pytest.mark.parametrize("conflicting", [False, True]) +def test_checkpoint_globally_empty_group_preserves_validation( + conflicting: bool, +) -> None: + """Groups with no owner retain their metadata and reject inconsistencies.""" + metadata = [_metadata(0), _metadata(1)] + for rank, item in enumerate(metadata): + item["leaf_tree"][0]["param_groups"].append( + {"lr": 0.02, "step": int(conflicting and rank == 1)} + ) + if conflicting: + with pytest.raises(ValueError, match="conflicts across owners"): + merge_muon_checkpoint_metadata(metadata, trusted_global_ranks=[0, 1]) + else: + merged = merge_muon_checkpoint_metadata(metadata, trusted_global_ranks=[0, 1]) + assert merged["leaf_tree"][0]["param_groups"][2] == {"lr": 0.02, "step": 0} + + +def test_checkpoint_invalid_owner_group_index_raises() -> None: + """A parameter cannot claim ownership of a nonexistent group.""" + metadata = [_metadata(0), _metadata(1)] + metadata[0]["leaf_tree"][0]["parameters"][0]["source_owner"]["group_index"] = 2 + with pytest.raises(ValueError, match="group_index is out of range"): + merge_muon_checkpoint_metadata(metadata, trusted_global_ranks=[0, 1]) + + +def test_checkpoint_nonempty_group_count_mismatch_raises() -> None: + """Partially populated leaves must agree on the group structure.""" + metadata = [_metadata(0), _metadata(1)] + metadata[1]["leaf_tree"][0]["param_groups"].pop() + with pytest.raises(ValueError, match="param-group count conflicts"): + merge_muon_checkpoint_metadata(metadata, trusted_global_ranks=[0, 1]) diff --git a/tests/test_megatron_async_save.py b/tests/test_megatron_async_save.py index 5d88b7a190..dd767dc0bc 100644 --- a/tests/test_megatron_async_save.py +++ b/tests/test_megatron_async_save.py @@ -141,6 +141,34 @@ def test_async_disabled_creates_no_queue(): m.close() +def test_non_distributed_optimizer_requires_checkpoint_capability(): + mod = _import_checkpointer() + with ( + patch("torch.distributed.get_rank", return_value=0), + pytest.raises(NotImplementedError, match="explicit non-distributed"), + ): + mod.MegatronCheckpointManager( + model=MagicMock(), + optimizer=MagicMock(spec=[]), + lr_scheduler=None, + use_distributed_optimizer=False, + ) + + +def test_checkpoint_capable_non_distributed_optimizer_is_supported(): + mod = _import_checkpointer() + optimizer = MagicMock() + optimizer.supports_non_distributed_checkpoint = True + with patch("torch.distributed.get_rank", return_value=0): + manager = mod.MegatronCheckpointManager( + model=MagicMock(), + optimizer=optimizer, + lr_scheduler=None, + use_distributed_optimizer=False, + ) + assert manager.optimizer is optimizer + + def test_save_schedules_async_request(patched_checkpointer, tmp_path): mod, manager, queue = patched_checkpointer fake_request = object() diff --git a/tests/test_megatron_optimizer_config.py b/tests/test_megatron_optimizer_config.py index 1bff29b124..d3b7af19de 100644 --- a/tests/test_megatron_optimizer_config.py +++ b/tests/test_megatron_optimizer_config.py @@ -9,7 +9,9 @@ from areal.api import FinetuneSpec from areal.api.cli_args import ( CPUStagedOffloadConfig, + DistributedDataParallelConfig, MegatronEngineConfig, + MuonOptimizerConfig, OptimizerConfig, ) from areal.engine import megatron_engine as megatron_engine_module @@ -398,6 +400,105 @@ def capture_optimizer(config, model, config_arg, **kwargs): assert config.exp_avg_sq_dtype is torch.float32 +def test_cpu_staged_dist_muon_selects_layerwise_factory(monkeypatch) -> None: + """The core Megatron config directly selects the staged Muon factory.""" + captured = {} + engine = _make_test_engine( + OptimizerConfig( + type="dist_muon", + muon=MuonOptimizerConfig(momentum=0.91, use_nesterov=True, num_ns_steps=4), + ) + ) + engine.mcore_config = MegatronEngineConfig( + ddp=DistributedDataParallelConfig(use_distributed_optimizer=False), + cpu_staged_offload=CPUStagedOffloadConfig( + enabled=True, + buffer_count=1, + bucket_size_mb=8, + ), + ) + + class FakeMuonOptimizer: + def bind_managed_checkpoint_process_group(self, group) -> None: + captured["checkpoint_group"] = group + + def capture_optimizer(config, model, config_arg): + captured["config"] = config + captured["model"] = model + captured["staged_config"] = config_arg + return FakeMuonOptimizer() + + monkeypatch.setattr( + megatron_engine_module, + "get_megatron_optimizer_with_gpu_staged_muon", + capture_optimizer, + ) + monkeypatch.setattr( + megatron_engine_module, + "OptimizerParamScheduler", + lambda *args, **kwargs: object(), + ) + monkeypatch.setattr( + megatron_engine_module, + "MegatronCheckpointManager", + lambda **kwargs: object(), + ) + + engine._create_optimizer( + FinetuneSpec(total_train_epochs=1, dataset_size=2, train_batch_size=1) + ) + + config = captured["config"] + assert captured["model"] is engine.model + assert captured["staged_config"].buffer_count == 1 + assert captured["staged_config"].slot_size_mb == 8 + assert config.use_distributed_optimizer is False + assert config.use_precision_aware_optimizer is False + assert config.main_grads_dtype is torch.float32 + assert config.optimizer == "muon" + assert config.use_layer_wise_distributed_optimizer is True + assert config.muon_momentum == 0.91 + assert config.muon_nesterov is True + assert config.muon_num_ns_steps == 4 + assert captured["checkpoint_group"] is engine.cpu_group + + +def test_native_dist_muon_selects_official_layerwise_factory(monkeypatch) -> None: + captured = {} + engine = _make_test_engine(OptimizerConfig(type="dist_muon")) + engine.mcore_config.ddp.use_distributed_optimizer = False + + def capture_optimizer(config, model): + captured["config"] = config + captured["model"] = model + return object() + + monkeypatch.setattr( + megatron_engine_module, + "get_megatron_optimizer_with_dist_muon", + capture_optimizer, + ) + monkeypatch.setattr( + megatron_engine_module, + "OptimizerParamScheduler", + lambda *args, **kwargs: object(), + ) + monkeypatch.setattr( + megatron_engine_module, + "MegatronCheckpointManager", + lambda **kwargs: object(), + ) + + engine._create_optimizer( + FinetuneSpec(total_train_epochs=1, dataset_size=2, train_batch_size=1) + ) + + assert captured["model"] is engine.model + assert captured["config"].optimizer == "muon" + assert captured["config"].use_layer_wise_distributed_optimizer is True + assert captured["config"].use_distributed_optimizer is False + + def test_scheduler_uses_fixed_warmup_and_resume_safe_initial_lr( monkeypatch, ) -> None: diff --git a/uv.lock b/uv.lock index ea6ff854e1..d6f52dfa7b 100644 --- a/uv.lock +++ b/uv.lock @@ -471,6 +471,7 @@ cli = [ { name = "colorlog", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] cuda = [ + { name = "emerging-optimizers", marker = "python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "flash-linear-attention", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "kernels", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "mbridge", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, @@ -486,6 +487,7 @@ cuda = [ { name = "torchvision", version = "0.24.1+cu129", source = { registry = "https://download.pytorch.org/whl/cu129" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, ] cuda-train = [ + { name = "emerging-optimizers", marker = "python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "flash-linear-attention", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "kernels", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "mbridge", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, @@ -499,6 +501,7 @@ kernels = [ { name = "kernels", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] megatron = [ + { name = "emerging-optimizers", marker = "python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "mbridge", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "megatron-bridge", marker = "python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "megatron-core", marker = "python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, @@ -558,6 +561,7 @@ requires-dist = [ { name = "daytona", marker = "extra == 'sandbox'", specifier = ">=0.167.0" }, { name = "distro-info", specifier = ">=1.0" }, { name = "einops" }, + { name = "emerging-optimizers", marker = "python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'megatron'", specifier = "==0.3.0" }, { name = "fastapi", specifier = ">=0.115.12" }, { name = "flash-linear-attention", marker = "platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'cuda-train'", specifier = "==0.4.2" }, { name = "flask" }, @@ -1753,6 +1757,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/de/15/545e2b6cf2e3be84bc1ed85613edd75b8aea69807a71c26f4ca6a9258e82/email_validator-2.3.0-py3-none-any.whl", hash = "sha256:80f13f623413e6b197ae73bb10bf4eb0908faf509ad8362c5edeb0be7fd450b4", size = 35604, upload-time = "2025-08-26T13:09:05.858Z" }, ] +[[package]] +name = "emerging-optimizers" +version = "0.3.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py", marker = "python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "torch", version = "2.9.1+cu129", source = { registry = "https://download.pytorch.org/whl/cu129" }, marker = "python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c5/bb/3fb2bfacc364b82356f886ce1babbc9b9041bda32bc46ae4a12f77a24d9c/emerging_optimizers-0.3.0.tar.gz", hash = "sha256:4eb96f38280dab17d1fa5bf46bdd7a110b1ab399e98891d0f90c467fb134e0ce", size = 101699, upload-time = "2026-07-08T19:45:16.371Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3a/fd/b8aa6cad45beae4d6b585ad55dad32ed95ac74097dd7bfb8952a72421fee/emerging_optimizers-0.3.0-py3-none-any.whl", hash = "sha256:e993682a8b93b6953b3056b0d4e2714e3a3182b0c85f456bcc78d76cd7b47429", size = 102898, upload-time = "2026-07-08T19:45:15.36Z" }, +] + [[package]] name = "eval-type-backport" version = "0.3.1" diff --git a/uv.vllm.lock b/uv.vllm.lock index 9a5213c3d8..a4c15172d1 100644 --- a/uv.vllm.lock +++ b/uv.vllm.lock @@ -503,6 +503,7 @@ cli = [ { name = "colorlog", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] cuda = [ + { name = "emerging-optimizers", marker = "python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "flash-linear-attention", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "kernels", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "mbridge", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, @@ -517,6 +518,7 @@ cuda = [ { name = "vllm", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, ] cuda-train = [ + { name = "emerging-optimizers", marker = "python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "flash-linear-attention", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "kernels", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "mbridge", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, @@ -530,6 +532,7 @@ kernels = [ { name = "kernels", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] megatron = [ + { name = "emerging-optimizers", marker = "python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "mbridge", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "megatron-bridge", marker = "python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "megatron-core", marker = "python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, @@ -588,6 +591,7 @@ requires-dist = [ { name = "daytona", marker = "extra == 'sandbox'", specifier = ">=0.167.0" }, { name = "distro-info", specifier = ">=1.0" }, { name = "einops" }, + { name = "emerging-optimizers", marker = "python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'megatron'", specifier = "==0.3.0" }, { name = "fastapi", specifier = ">=0.115.12" }, { name = "flash-linear-attention", marker = "platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'cuda-train'", specifier = "==0.4.2" }, { name = "flask" }, @@ -1873,6 +1877,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/de/15/545e2b6cf2e3be84bc1ed85613edd75b8aea69807a71c26f4ca6a9258e82/email_validator-2.3.0-py3-none-any.whl", hash = "sha256:80f13f623413e6b197ae73bb10bf4eb0908faf509ad8362c5edeb0be7fd450b4", size = 35604, upload-time = "2025-08-26T13:09:05.858Z" }, ] +[[package]] +name = "emerging-optimizers" +version = "0.3.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py", marker = "python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "torch", version = "2.10.0+cu129", source = { registry = "https://download.pytorch.org/whl/cu129" }, marker = "python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c5/bb/3fb2bfacc364b82356f886ce1babbc9b9041bda32bc46ae4a12f77a24d9c/emerging_optimizers-0.3.0.tar.gz", hash = "sha256:4eb96f38280dab17d1fa5bf46bdd7a110b1ab399e98891d0f90c467fb134e0ce", size = 101699, upload-time = "2026-07-08T19:45:16.371Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3a/fd/b8aa6cad45beae4d6b585ad55dad32ed95ac74097dd7bfb8952a72421fee/emerging_optimizers-0.3.0-py3-none-any.whl", hash = "sha256:e993682a8b93b6953b3056b0d4e2714e3a3182b0c85f456bcc78d76cd7b47429", size = 102898, upload-time = "2026-07-08T19:45:15.36Z" }, +] + [[package]] name = "eval-type-backport" version = "0.3.1"