feat(engine): add distributed Muon optimizer offload

.github/workflows/test-areal.yml

@@ -307,6 +307,22 @@ jobs:
     steps:
       - uses: actions/checkout@v6
 
+      - name: Sync AWEX dependency
+        run: |
+          AWEX_REQUIREMENT=$("$VIRTUAL_ENV/bin/python" - <<'PY'
+          import tomllib
+          with open("pyproject.toml", "rb") as f:
+              dependencies = tomllib.load(f)["project"]["dependencies"]
+          print(next(dep for dep in dependencies if dep.startswith("awex==")))
+          PY
+          )
+          uv pip install --python "$VIRTUAL_ENV/bin/python" --no-deps "$AWEX_REQUIREMENT"
+
+      - 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
         run: |
           export PATH="/opt/.venv/bin:$PATH"
@@ -316,7 +332,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={
@@ -914,6 +968,36 @@ class FP8EngineConfig:
     )
 
 
+@dataclass
+class CPUStagedOffloadConfig:
+    """Bounded GPU staging for CPU-resident Megatron optimizer state."""
+
+    enabled: bool = field(
+        default=False,
+        metadata={
+            "help": (
+                "Keep optimizer state in pinned CPU memory and stream bounded "
+                "update units through GPU staging buffers. The optimizer algorithm "
+                "is selected independently by optimizer.type. Megatron only."
+            )
+        },
+    )
+    buffer_count: int = field(
+        default=2,
+        metadata={"help": "Number of reusable GPU optimizer staging buffers."},
+    )
+    bucket_size_mb: float = field(
+        default=128.0,
+        metadata={"help": "Maximum size in MiB of one GPU optimizer staging unit."},
+    )
+
+    def __post_init__(self) -> None:
+        if self.buffer_count < 1:
+            raise ValueError("cpu_staged_offload.buffer_count must be at least 1")
+        if self.bucket_size_mb <= 0:
+            raise ValueError("cpu_staged_offload.bucket_size_mb must be positive")
+
+
 @dataclass
 class MegatronEngineConfig:
     """Configuration for Megatron-LM training framework.
@@ -962,6 +1046,15 @@ class MegatronEngineConfig:
     main_params_dtype: str = "float32"
     exp_avg_dtype: str = "float32"
     exp_avg_sq_dtype: str = "float32"
+    cpu_staged_offload: CPUStagedOffloadConfig = field(
+        default_factory=CPUStagedOffloadConfig,
+        metadata={
+            "help": (
+                "Keep the selected optimizer's state CPU-resident and use bounded "
+                "GPU staging buffers for optimizer steps."
+            )
+        },
+    )
 
     # Checkpointing Configuration
     async_save: bool = field(
@@ -1174,6 +1267,11 @@ class MegatronEngineConfig:
     )
 
     def __post_init__(self) -> None:
+        if self.cpu_staged_offload.enabled and self.async_save:
+            raise ValueError(
+                "megatron.async_save is not supported with "
+                "megatron.cpu_staged_offload.enabled=true"
+            )
         if self.enable_mtp_training and not self.enable_mtp:
             raise ValueError("enable_mtp_training requires enable_mtp=True")
         if self.mtp_only:

areal/dataset/__init__.py

@@ -30,6 +30,7 @@
     "virl39k",
     "hh-rlhf",
     "torl_data",
+    "dapo-math-17k",
     "swe_sft",
 ]
 
@@ -174,6 +175,16 @@ def _get_custom_dataset(
             max_length=max_length,
             **kwargs,
         )
+    elif "dapo-math-17k" in path.lower() and type == "rl":
+        from .dapo_math_17k import get_dapo_math_17k_rl_dataset
+
+        return get_dapo_math_17k_rl_dataset(
+            path=path,
+            split=split,
+            tokenizer=tokenizer,
+            max_length=max_length,
+            **kwargs,
+        )
     elif (
         _SWE_PATH_PATTERN.search(path.lower()) or _is_swe_sft_artifact(path)
     ) and type == "sft":

areal/dataset/dapo_math_17k.py

@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from pathlib import Path
+
+from datasets import Dataset, load_dataset
+
+from areal.utils.hf_utils import apply_chat_template
+
+
+def get_dapo_math_17k_rl_dataset(
+    path: str,
+    split: str,
+    tokenizer,
+    max_length: int | None = None,
+) -> Dataset:
+    """Load local DAPO-Math-17k parquet data for RL training."""
+    del split
+    dataset_root = Path(path)
+    parquet_files = sorted((dataset_root / "data").glob("*.parquet"))
+    json_files = sorted(dataset_root.glob("*.jsonl")) + sorted(
+        dataset_root.glob("*.json")
+    )
+    if parquet_files:
+        dataset = load_dataset(
+            "parquet",
+            data_files=[str(data_file) for data_file in parquet_files],
+            split="train",
+        )
+    elif json_files:
+        dataset = load_dataset(
+            "json",
+            data_files=[str(data_file) for data_file in json_files],
+            split="train",
+        )
+    else:
+        raise FileNotFoundError(
+            f"No parquet or JSON data files found under {dataset_root}"
+        )
+
+    def process(sample):
+        answer = sample.get("label")
+        if answer is None:
+            answer = sample["reward_model"]["ground_truth"]
+        return {"messages": sample["prompt"], "answer": str(answer)}
+
+    dataset = dataset.map(process, remove_columns=dataset.column_names)
+
+    if max_length is not None:
+        if tokenizer is None:
+            raise ValueError("tokenizer is required when max_length is set")
+
+        def filter_length(sample):
+            prompt_ids = apply_chat_template(
+                tokenizer, sample["messages"], add_generation_prompt=True, tokenize=True
+            )
+            return len(prompt_ids) <= max_length
+
+        dataset = dataset.filter(filter_length)
+
+    return dataset