Persistent MXFP8 for experts layers in MoE

.claude/docs/backends/megatron.md

@@ -24,6 +24,18 @@ Key strategies:
 
 Note: Sequence parallelism is auto-enabled when `tensor_model_parallel_size > 1` — there is no separate config field for it.
 
+## Expert MXFP8
+
+Set `trainer.policy.model.expert_mxfp8.enabled=true` on SM100/SM103 to run
+routed expert GEMMs with Transformer Engine MXFP8 and configure vLLM for
+expert-only online MXFP8. Checkpoints and weight sync remain high precision.
+Set `trainer.policy.model.expert_mxfp8.persistent=true` with
+`trainer.policy.megatron_config.ddp_config.fp8_param_gather=true` to keep
+routed-expert primary parameters in MXFP8 between optimizer steps. Set
+`generator.inference_engine.fp8_weight_sync_mode=serialized_mxfp8` to send
+routed experts to vLLM as MXFP8 data and E8M0 scales. Unmerged LoRA is not
+supported.
+
 ## Test Requirements
 
 Megatron GPU tests need: `NVTE_FLASH_ATTN=0`

docs/content/docs/examples/quantized_rollouts.mdx

@@ -31,6 +31,41 @@ generator.inference_engine.engine_init_kwargs.quantization=fp8
 
 This uses vLLM's [online dynamic FP8 quantization](https://docs.vllm.ai/en/latest/features/quantization/fp8.html), so no calibration data or pre-quantized checkpoint is required.
 
+## Expert-only MXFP8
+
+On Blackwell GPUs, the Megatron backend can run routed MoE expert GEMMs in
+MXFP8 while vLLM serves the same experts with online MXFP8 quantization:
+
+```bash
+trainer.strategy=megatron \
+trainer.policy.model.expert_mxfp8.enabled=true \
+trainer.policy.model.expert_mxfp8.persistent=true \
+trainer.policy.megatron_config.ddp_config.fp8_param_gather=true \
+generator.inference_engine.fp8_weight_sync_mode=serialized_mxfp8
+```
+
+Persistent expert weights are quantized after each optimizer step and reused by
+Transformer Engine until the next update. Weight sync sends routed experts as
+MXFP8 data with E8M0 scales, while checkpoints remain high precision. Attention,
+dense layers, shared experts, the router, and the LM head remain high precision.
+
+This mode requires SM100 or SM103 and does not support unmerged LoRA. Set
+`training=false` or `rollout=false` under `expert_mxfp8` to enable only one
+side for benchmarking. Set `persistent=false` to quantize expert weights during
+each forward pass. Set `fp8_weight_sync_mode=null` to let vLLM quantize experts
+after each weight sync.
+
+Run the Modal benchmark with identical BF16 and MXFP8 workloads:
+
+```bash
+uv run --isolated --with modal modal run --detach \
+  examples/train/megatron/modal_expert_mxfp8_benchmark.py \
+  --mode both --steps 10
+```
+
+Use `--mode bf16` or `--mode mxfp8` for one side. Set `MODAL_GPU` to override
+the default `B200:8` allocation.
+
 ## Enabling off-policy correction (TIS)
 
 To apply TIS, we need the inference engine to return the rollout logprobs for the generated tokens, and we configure the correction on the policy loss:

examples/train/megatron/modal_expert_mxfp8_benchmark.py

@@ -0,0 +1,254 @@
+"""Benchmark BF16 versus expert-only MXFP8 on DAPO with one Modal node.
+
+Run:
+    uv run --isolated --with modal modal run --detach \
+        examples/train/megatron/modal_expert_mxfp8_benchmark.py \
+        --mode both --steps 10
+"""
+
+from __future__ import annotations
+
+import os
+import pathlib
+
+import modal
+
+APP_NAME = "skyrl-expert-mxfp8-benchmark"
+MODEL = "Qwen/Qwen3-30B-A3B"
+REMOTE_REPO = "/root/SkyRL"
+HF_HOME = "/root/hf-cache"
+DATA_DIR = "/root/data/dapo"
+TRAIN_FILE = f"{DATA_DIR}/dapo-math-17k-cleaned.parquet"
+VAL_FILE = f"{DATA_DIR}/aime-2024-cleaned.parquet"
+PROFILE_ROOT = "/root/profiles"  # profiler
+PROFILE_COMMIT_INTERVAL_S = 120  # profiler
+GPU = os.environ.get("MODAL_GPU", "B200:8")
+
+
+def _repo_root() -> pathlib.Path:
+    for start in (pathlib.Path(__file__).resolve(), pathlib.Path.cwd().resolve()):
+        candidate = start if start.is_dir() else start.parent
+        for path in (candidate, *candidate.parents):
+            if (path / "pyproject.toml").exists() and (path / "skyrl").exists():
+                return path
+    raise RuntimeError("Run the benchmark from a SkyRL checkout")
+
+
+repo_root = _repo_root()
+hf_volume = modal.Volume.from_name("skyrl-hf-cache", create_if_missing=True)
+data_volume = modal.Volume.from_name("skyrl-expert-mxfp8-data", create_if_missing=True)
+profile_volume = modal.Volume.from_name("skyrl-expert-mxfp8-profiles", create_if_missing=True)  # profiler
+
+image = (
+    modal.Image.from_registry("nvidia/cuda:12.9.1-devel-ubuntu22.04", add_python="3.12")
+    .apt_install("git", "curl", "build-essential", "ca-certificates", "libnuma1", "numactl")
+    .pip_install("huggingface-hub")
+    .run_commands("curl -LsSf https://astral.sh/uv/install.sh | sh")
+    .env(
+        {
+            "PATH": "/root/.local/bin:/usr/local/cuda/bin:${PATH}",
+            "HF_HOME": HF_HOME,
+            "HF_XET_HIGH_PERFORMANCE": "1",
+            "NVTE_FLASH_ATTN": "0",
+            "SKYRL_WAIT_UNTIL_INFERENCE_SERVER_HEALTHY_TIMEOUT_S": "1800",
+            "SKYRL_DUMP_INFRA_LOG_TO_STDOUT": "1",
+            "UV_LINK_MODE": "copy",
+            "UV_PROJECT_ENVIRONMENT": f"{REMOTE_REPO}/.venv",
+            "VLLM_USE_FLASHINFER_MOE_FP16": "0",
+            "VLLM_USE_FLASHINFER_SAMPLER": "0",
+        }
+    )
+    .add_local_dir(
+        str(repo_root),
+        REMOTE_REPO,
+        copy=True,
+        ignore=[".venv", ".git", "**/__pycache__"],
+    )
+    .workdir(REMOTE_REPO)
+    .run_commands("uv sync --extra megatron", gpu="any")
+    .run_commands(f"rm -rf {HF_HOME}")
+)
+
+app = modal.App(APP_NAME)
+
+
+def _run(mode: str, steps: int, max_prompt_length: int, max_generate_length: int) -> None:
+    import subprocess
+    import time
+
+    enabled = str(mode == "mxfp8").lower()
+    run_name = f"qwen3-30b-a3b-dapo-{mode}"
+    command = [
+        "uv",
+        "run",
+        "--frozen",
+        "--no-sync",
+        "--extra",
+        "megatron",
+        "-m",
+        "examples.train.algorithms.dapo.main_dapo",
+        f"data.train_data=['{TRAIN_FILE}']",
+        f"data.val_data=['{VAL_FILE}']",
+        "trainer.strategy=megatron",
+        f"trainer.policy.model.path={MODEL}",
+        f"trainer.policy.model.expert_mxfp8.enabled={enabled}",
+        "trainer.placement.colocate_all=true",
+        "trainer.placement.policy_num_nodes=1",
+        "trainer.placement.policy_num_gpus_per_node=8",
+        "generator.inference_engine.num_engines=1",
+        "generator.inference_engine.tensor_parallel_size=1",
+        "generator.inference_engine.data_parallel_size=8",
+        "generator.inference_engine.expert_parallel_size=8",
+        "generator.inference_engine.distributed_executor_backend=mp",
+        "generator.inference_engine.enforce_eager=false",
+        "generator.inference_engine.weight_sync_backend=nccl",
+        "generator.inference_engine.gpu_memory_utilization=0.6",
+        "trainer.policy.megatron_config.tensor_model_parallel_size=1",
+        "trainer.policy.megatron_config.pipeline_model_parallel_size=1",
+        "trainer.policy.megatron_config.context_parallel_size=1",
+        "trainer.policy.megatron_config.expert_model_parallel_size=8",
+        "trainer.policy.megatron_config.expert_tensor_parallel_size=1",
+        "trainer.algorithm.advantage_estimator=grpo",
+        "trainer.algorithm.policy_loss_type=dual_clip",
+        "trainer.algorithm.loss_reduction=token_mean_legacy",
+        "trainer.algorithm.overlong_buffer_len=4096",
+        "trainer.algorithm.overlong_buffer_penalty_factor=1.0",
+        "trainer.algorithm.eps_clip_low=0.2",
+        "trainer.algorithm.eps_clip_high=0.28",
+        "trainer.algorithm.use_kl_loss=false",
+        "trainer.train_batch_size=64",
+        "trainer.policy_mini_batch_size=64",
+        "trainer.micro_forward_batch_size_per_gpu=1",
+        "trainer.micro_train_batch_size_per_gpu=1",
+        "trainer.eval_before_train=false",
+        "trainer.eval_interval=0",
+        "trainer.ckpt_interval=0",
+        "trainer.hf_save_interval=0",
+        "trainer.resume_mode=none",
+        f"trainer.max_training_steps={steps}",
+        f"trainer.max_prompt_length={max_prompt_length}",
+        f"generator.sampling_params.max_generate_length={max_generate_length}",
+        "generator.apply_overlong_filtering=true",
+        "generator.n_samples_per_prompt=8",
+        "generator.batched=true",
+        "environment.env_class=aime",
+        # profiler
+        "trainer.policy.torch_profiler_config.enable=false",
+        "trainer.policy.torch_profiler_config.ranks=[0]",
+        f"trainer.policy.torch_profiler_config.save_path={PROFILE_ROOT}/trainer/{mode}",
+        "trainer.policy.torch_profiler_config.skip_first=3",
+        "trainer.policy.torch_profiler_config.wait=0",
+        "trainer.policy.torch_profiler_config.warmup=1",
+        "trainer.policy.torch_profiler_config.active=2",
+        "trainer.policy.torch_profiler_config.repeat=1",
+        "trainer.policy.torch_profiler_config.record_shapes=false",
+        "trainer.policy.torch_profiler_config.with_stack=false",
+        "trainer.logger=wandb",
+        "trainer.project_name=skyrl-expert-mxfp8-dapo",
+        f"trainer.run_name={run_name}",
+    ]
+    if mode == "mxfp8":
+        command.extend(
+            [
+                "trainer.policy.model.expert_mxfp8.persistent=true",
+                "trainer.policy.megatron_config.ddp_config.fp8_param_gather=true",
+                "trainer.policy.megatron_config.ddp_config.reuse_grad_buf_for_mxfp8_param_ag=true",
+                "trainer.policy.megatron_config.optimizer_config_kwargs.fp8_recipe=mxfp8",
+                "trainer.policy.megatron_config.optimizer_config_kwargs.reuse_grad_buf_for_mxfp8_param_ag=true",
+                "generator.inference_engine.fp8_weight_sync_mode=serialized_mxfp8",
+            ]
+        )
+    started = time.perf_counter()
+    process = subprocess.Popen(command, cwd=REMOTE_REPO)
+    while True:
+        try:
+            returncode = process.wait(timeout=PROFILE_COMMIT_INTERVAL_S)
+            break
+        except subprocess.TimeoutExpired:
+            profile_volume.commit()  # profiler
+    if returncode != 0:
+        raise subprocess.CalledProcessError(returncode, command)
+    print(f"{mode} benchmark completed in {time.perf_counter() - started:.1f}s")
+
+
+@app.function(
+    image=image,
+    cpu=4,
+    memory=16384,
+    volumes={HF_HOME: hf_volume, "/root/data": data_volume},
+    timeout=2 * 60 * 60,
+)
+def prepare_assets() -> None:
+    import shutil
+    import subprocess
+
+    from huggingface_hub import hf_hub_download, snapshot_download
+
+    snapshot_download(MODEL)
+    hf_volume.commit()
+    os.makedirs(DATA_DIR, exist_ok=True)
+    if not os.path.exists(TRAIN_FILE) or not os.path.exists(VAL_FILE):
+        train_download = hf_hub_download(
+            repo_id="BytedTsinghua-SIA/DAPO-Math-17k",
+            filename="data/dapo-math-17k.parquet",
+            repo_type="dataset",
+        )
+        val_download = hf_hub_download(
+            repo_id="BytedTsinghua-SIA/AIME-2024",
+            filename="data/aime-2024.parquet",
+            repo_type="dataset",
+        )
+        shutil.copyfile(train_download, f"{DATA_DIR}/dapo-math-17k.parquet")
+        shutil.copyfile(val_download, f"{DATA_DIR}/aime-2024.parquet")
+        subprocess.run(
+            [
+                "uv",
+                "run",
+                "--frozen",
+                "--no-sync",
+                "--extra",
+                "megatron",
+                "-m",
+                "examples.train.algorithms.dapo.data_preprocess_dapo_aime",
+                "--data-dir",
+                DATA_DIR,
+            ],
+            cwd=REMOTE_REPO,
+            check=True,
+        )
+        data_volume.commit()
+
+
+@app.function(
+    image=image,
+    gpu=GPU,
+    secrets=[modal.Secret.from_name("wandb-secret")],
+    volumes={
+        HF_HOME: hf_volume,
+        "/root/data": data_volume,
+        PROFILE_ROOT: profile_volume,
+    },  # profiler
+    timeout=24 * 60 * 60,
+)
+def benchmark(mode: str, steps: int, max_prompt_length: int, max_generate_length: int) -> None:
+    try:
+        _run(mode, steps, max_prompt_length, max_generate_length)
+    finally:
+        profile_volume.commit()  # profiler
+
+
+@app.local_entrypoint()
+def main(
+    mode: str = "both",
+    steps: int = 10,
+    max_prompt_length: int = 2048,
+    max_generate_length: int = 4096,
+) -> None:
+    if mode not in {"bf16", "mxfp8", "both"}:
+        raise ValueError("mode must be bf16, mxfp8, or both")
+    if max_prompt_length <= 0 or max_generate_length <= 0:
+        raise ValueError("max_prompt_length and max_generate_length must be positive")
+    prepare_assets.remote()
+    modes = ("bf16", "mxfp8") if mode == "both" else (mode,)
+    for selected in modes:
+        benchmark.remote(selected, steps, max_prompt_length, max_generate_length)