feat(engine): add distributed Muon optimizer offload

.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

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={

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
             )

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