From 94397d0d164903539a0122628df7e5dc70a6a161 Mon Sep 17 00:00:00 2001 From: kailash Date: Tue, 23 Jun 2026 21:39:20 +0000 Subject: [PATCH 1/6] nemo rl multinode grpo example --- nemo-rl/README.md | 136 ++++++++++++ nemo-rl/configs/__init__.py | 17 ++ nemo-rl/configs/base.py | 161 ++++++++++++++ nemo-rl/configs/llama3_1_8b_math_2node.py | 36 ++++ .../nemotron_nano_v3_30ba3b_math_2node.py | 39 ++++ nemo-rl/configs/qwen2_5_1_5b_math.py | 33 +++ nemo-rl/configs/qwen2_5_1_5b_math_2node.py | 35 +++ nemo-rl/configs/qwen3_8b_math.py | 33 +++ nemo-rl/modal_helpers/__init__.py | 0 nemo-rl/modal_helpers/run_grpo_multinode.py | 113 ++++++++++ nemo-rl/modal_helpers/utils.py | 191 ++++++++++++++++ nemo-rl/modal_train.py | 204 ++++++++++++++++++ 12 files changed, 998 insertions(+) create mode 100644 nemo-rl/README.md create mode 100644 nemo-rl/configs/__init__.py create mode 100644 nemo-rl/configs/base.py create mode 100644 nemo-rl/configs/llama3_1_8b_math_2node.py create mode 100644 nemo-rl/configs/nemotron_nano_v3_30ba3b_math_2node.py create mode 100644 nemo-rl/configs/qwen2_5_1_5b_math.py create mode 100644 nemo-rl/configs/qwen2_5_1_5b_math_2node.py create mode 100644 nemo-rl/configs/qwen3_8b_math.py create mode 100644 nemo-rl/modal_helpers/__init__.py create mode 100644 nemo-rl/modal_helpers/run_grpo_multinode.py create mode 100644 nemo-rl/modal_helpers/utils.py create mode 100644 nemo-rl/modal_train.py diff --git a/nemo-rl/README.md b/nemo-rl/README.md new file mode 100644 index 0000000..fcafe21 --- /dev/null +++ b/nemo-rl/README.md @@ -0,0 +1,136 @@ +# nemo-rl — Modal launcher for NeMo-RL training + +Thin Modal launcher that runs [NeMo-RL](https://github.com/NVIDIA-NeMo/RL) RL +training (GRPO, SFT, DPO, …) on multi-node Modal GPU clusters. + +It mirrors the [`slime/`](../slime) launcher in this repo: each experiment is a +Python config, a clustered Modal function brings up a Ray cluster across the +allocation, and the driver runs on the head node. The difference is that NeMo-RL +is driven by a YAML config plus Hydra `key=value` overrides rather than raw CLI +flags. + +## Prerequisites + +- Modal CLI installed and authenticated +- Modal environment selected: `export MODAL_ENVIRONMENT=` +- Modal secrets: + - `huggingface-secret` — `HF_TOKEN` for model/data download (gated models) + - `wandb-secret` — `WANDB_API_KEY` for configs with `logger.wandb_enabled=true` + +Run all commands from this directory (`nemo-rl/`). + +## Common Workflow + +Set the config name once: + +```bash +export EXPERIMENT_CONFIG=qwen2_5_1_5b_math +``` + +List available configs: + +```bash +modal run modal_train.py::list_configs +``` + +Download the model into the HF cache volume: + +```bash +modal run modal_train.py::download_model +``` + +Download dataset: + +```bash +modal run modal_train.py::download_data +``` + +Launch training: + +```bash +modal run -d modal_train.py::train +``` + +Use `-d` (detached) to keep training running after you close your terminal. The +Ray dashboard URL is printed at the start of the run. + +## Launcher Model + +Each experiment lives in `configs/.py` and exposes: + +- `modal`: `ModalConfig` for image, GPU type, and Modal resources +- `nemo_rl`: `NemoRLConfig` for the recipe (run script, base YAML, cluster + shape, and Hydra overrides) + +The launcher runs this on the Ray head node: + +```bash +cd /opt/nemo-rl && uv run python --config +``` + +### How multi-node works + +1. `train` is wrapped in `modal.experimental.clustered(num_nodes, rdma=True)`, + so Modal provisions `num_nodes` RDMA-connected GPU nodes. +2. Each node reads its rank from `modal.experimental.get_cluster_info()`. +3. Rank 0 starts the Ray head and waits until all nodes and GPUs have joined; + ranks 1…N start Ray workers pointed at the head and idle. +4. Rank 0 runs the NeMo-RL driver, which attaches to the existing Ray cluster + (`RAY_ADDRESS`) and schedules its actors across the whole allocation. + +This is the equivalent of NemoRL's Slurm script [`ray.sub`](https://github.com/NVIDIA-NeMo/RL/blob/main/ray.sub). + +## Volumes + +| Volume | Mount path | Purpose | +| --- | --- | --- | +| `huggingface-cache` | `/root/.cache/huggingface` | HF model + dataset cache | +| `nemo-rl-checkpoints` | `/checkpoints` | Training checkpoints | + +## Add A Config + +Create `configs/.py`: + +```python +from configs.base import ModalConfig, NemoRLConfig + +modal = ModalConfig(gpu="H100") + + +class _Recipe(NemoRLConfig): + entrypoint = "examples/run_grpo.py" + base_config = "examples/configs/grpo_math_8B.yaml" + + num_nodes = 2 + gpus_per_node = 8 + + hf_model = "meta-llama/Llama-3.1-8B-Instruct" + hf_datasets = ["nvidia/OpenMathInstruct-2"] + + overrides = { + "policy.model_name": "meta-llama/Llama-3.1-8B-Instruct", + "policy.dtensor_cfg.tensor_parallel_size": 8, + "logger.wandb_enabled": True, + "logger.wandb.name": "my-run", + } + + +nemo_rl = _Recipe() +``` + +The base_config points to a path in the NemoRL repo under examples/configs +## Dev overlay + +To run local NeMo-RL changes without rebuilding the image, point `local_nemo_rl` +at your checkout. It is copied over `/opt/nemo-rl` in the image: + +```python +modal = ModalConfig( + gpu="H100", + local_nemo_rl="/path/to/your/RL", +) +``` + +The container still uses its baked `/opt/nemo_rl_venv`, so this only overlays +source code, not dependencies. If you change dependencies, rebuild/extend the +image (e.g. via `image_run_commands`). diff --git a/nemo-rl/configs/__init__.py b/nemo-rl/configs/__init__.py new file mode 100644 index 0000000..edcaf85 --- /dev/null +++ b/nemo-rl/configs/__init__.py @@ -0,0 +1,17 @@ +import importlib +from pathlib import Path + +_CONFIGS_DIR = Path(__file__).parent +_SKIP = {"base", "__init__"} + + +def get_module(name: str): + try: + return importlib.import_module(f"configs.{name}") + except ModuleNotFoundError as exc: + if exc.name != f"configs.{name}": + raise + available = sorted( + f.stem for f in _CONFIGS_DIR.glob("*.py") if f.stem not in _SKIP + ) + raise ValueError(f"Unknown config {name!r}. Available: {available}") from exc diff --git a/nemo-rl/configs/base.py b/nemo-rl/configs/base.py new file mode 100644 index 0000000..3544cc3 --- /dev/null +++ b/nemo-rl/configs/base.py @@ -0,0 +1,161 @@ +"""Base configuration classes and volume mount paths for NeMo-RL. + +Two separate concerns: + + ModalConfig — Modal infrastructure (gpu model, image, dev overlay) + NemoRLConfig — NeMo-RL recipe: which run script, which base YAML config, + cluster shape, and Hydra overrides + +NeMo-RL is driven by a YAML config plus Hydra-style ``key=value`` overrides +(see ``examples/run_grpo.py``). Unlike slime — where every config attribute +becomes a CLI flag — NeMo-RL configs carry a single ``overrides`` dict of +dotted Hydra keys, because the keys contain dots that are not valid Python +attribute names. + +Each experiment defines one ``ModalConfig`` and one ``NemoRLConfig`` instance. +""" + +from pathlib import Path +from typing import Any, Literal + +# ── Volume mount paths ──────────────────────────────────────────────────────── + +# The NeMo-RL container caches Hugging Face artifacts under the default HF home. +HF_CACHE_PATH = Path("/root/.cache/huggingface") +CHECKPOINTS_PATH = Path("/checkpoints") + +# Where the NeMo-RL repo lives inside the official image (see docker/Dockerfile). +NEMO_RL_ROOT = "/opt/nemo-rl" + +# ── Types ───────────────────────────────────────────────────────────────────── + +GPUType = Literal["H100", "H200", "B200", "B300", "A100"] + +class ModalConfig: + """Modal infrastructure configuration — GPU provisioning and image setup only.""" + + # Official NeMo-RL release image (https://registry.ngc.nvidia.com/orgs/nvidia/containers/nemo-rl). + docker_image: str = "nvcr.io/nvidia/nemo-rl:v0.5.0" + gpu: GPUType = "H100" + memory: tuple[int, int] | None = ( + None # per-container memory in MiB; see https://modal.com/docs/guide/resources#memory-limits + ) + cloud: str | None = None # e.g. "aws", "gcp" + region: str | None = None # e.g. "us-east-2" + local_nemo_rl: str | None = None # path to local NeMo-RL repo for dev overlay + image_run_commands: list[str] = [] # extra commands to run during image build + image_env: dict[str, str] = {} # env vars baked into the image (Modal .env()) + + def __init__(self, **kwargs: Any) -> None: + for k, v in kwargs.items(): + setattr(self, k, v) + +def _hydra_value(val: Any) -> str: + """Serialize a Python value into a Hydra override RHS string.""" + if val is None: + return "null" + if isinstance(val, bool): + return "true" if val else "false" + if isinstance(val, (list, tuple)): + return "[" + ",".join(_hydra_value(v) for v in val) + "]" + return str(val) + + +class NemoRLConfig: + """Base NeMo-RL recipe configuration. + + Subclass and set class attributes to configure an experiment. The launcher + runs:: + + uv run python --config + + on the Ray head node, after the Modal cluster's Ray cluster is up. + + Launcher fields: + entrypoint — run script relative to /opt/nemo-rl (e.g. examples/run_grpo.py) + base_config — YAML config path passed to --config + num_nodes — total Modal/Ray nodes (also set as cluster.num_nodes) + gpus_per_node — GPUs per node (also set as cluster.gpus_per_node) + hf_model — model id to prefetch into the HF cache; defaults to the + policy.model_name override, else the base config's value + hf_datasets — optional dataset repo ids to prefetch in download_data() + environment — extra environment variables for the driver process + + Recipe knobs: + overrides — dict of dotted Hydra keys → values, applied on top of the + base YAML config (e.g. {"policy.model_name": "Qwen/Qwen2.5-1.5B"}) + + Example:: + + class _Recipe(NemoRLConfig): + entrypoint = "examples/run_grpo.py" + base_config = "examples/configs/grpo_math_1B.yaml" + num_nodes = 1 + gpus_per_node = 8 + overrides = { + "policy.model_name": "Qwen/Qwen2.5-1.5B", + "logger.wandb_enabled": True, + } + + nemo_rl = _Recipe() + """ + + # ── Launcher instructions ─────────────────────────────────────────────────── + entrypoint: str = "examples/run_grpo.py" + base_config: str = "examples/configs/grpo_math_1B.yaml" + num_nodes: int = 1 + gpus_per_node: int = 8 + hf_model: str | None = None + hf_datasets: list[str] = [] + environment: dict[str, str] = {} + + # ── Recipe knobs ──────────────────────────────────────────────────────────── + overrides: dict[str, Any] = {} + + def __init__(self, **kwargs: Any) -> None: + # Fresh mutable copies per instance — never mutate the class-level defaults. + self.overrides = dict(type(self).overrides) + self.environment = dict(type(self).environment) + self.hf_datasets = list(type(self).hf_datasets) + for k, v in kwargs.items(): + setattr(self, k, v) + + # ── Public API ────────────────────────────────────────────────────────────── + + def total_nodes(self) -> int: + """Total Modal cluster nodes required by this recipe.""" + return self.num_nodes + + def resolved_overrides(self, experiment: str | None = None) -> dict[str, Any]: + """Hydra overrides with launcher-managed defaults merged in. + + The cluster shape and checkpoint directory are forced to match the Modal + cluster and mounted checkpoints volume. Explicit ``overrides`` win over + the default checkpoint dir, never over cluster shape. ``experiment`` (the + config file stem) keys the default checkpoint dir so recipes don't + collide on the shared volume. + """ + run_name = experiment or type(self).__name__ + merged: dict[str, Any] = { + "checkpointing.checkpoint_dir": f"{CHECKPOINTS_PATH}/{run_name}", + } + merged.update(self.overrides) + # Cluster shape always tracks the actual Modal allocation. + merged["cluster.num_nodes"] = self.num_nodes + merged["cluster.gpus_per_node"] = self.gpus_per_node + return merged + + def cli_args(self, experiment: str | None = None) -> list[str]: + """Argument list for the NeMo-RL run script. + + Produces ``["--config", , "k=v", ...]`` with Hydra-encoded + values (True→true, None→null, lists→[a,b]). + """ + out = ["--config", self.base_config] + for key, val in self.resolved_overrides(experiment).items(): + out.append(f"{key}={_hydra_value(val)}") + return out + + def model_id(self) -> str | None: + """Hugging Face model id this recipe trains, for prefetching.""" + return self.hf_model or self.overrides.get("policy.model_name") diff --git a/nemo-rl/configs/llama3_1_8b_math_2node.py b/nemo-rl/configs/llama3_1_8b_math_2node.py new file mode 100644 index 0000000..d161b5d --- /dev/null +++ b/nemo-rl/configs/llama3_1_8b_math_2node.py @@ -0,0 +1,36 @@ +"""Llama-3.1-8B-Instruct GRPO on OpenMathInstruct-2 — 2 nodes x 8 GPUs. + +Mirrors NeMo-RL's examples/configs/grpo_math_8B.yaml scaled to two nodes. This +is the multi-node example: Modal provisions a 2-node RDMA cluster, the launcher +brings up Ray across both nodes, and the driver runs on the head node with +cluster.num_nodes=2. +""" + +from configs.base import ModalConfig, NemoRLConfig + +modal = ModalConfig(gpu="H100") + + +class _Recipe(NemoRLConfig): + entrypoint = "examples/run_grpo.py" + base_config = "examples/configs/grpo_math_8B.yaml" + + num_nodes = 2 + gpus_per_node = 8 + + hf_model = "meta-llama/Llama-3.1-8B-Instruct" + hf_datasets = ["nvidia/OpenMathInstruct-2"] + + overrides = { + "policy.model_name": "meta-llama/Llama-3.1-8B-Instruct", + "policy.dtensor_cfg.enabled": True, + "policy.dtensor_cfg.tensor_parallel_size": 8, + "policy.dtensor_cfg.sequence_parallel": True, + "policy.dtensor_cfg.activation_checkpointing": True, + "logger.wandb_enabled": True, + "logger.wandb.project": "nemo-rl-grpo", + "logger.wandb.name": "llama3.1-8b-math-2node", + } + + +nemo_rl = _Recipe() diff --git a/nemo-rl/configs/nemotron_nano_v3_30ba3b_math_2node.py b/nemo-rl/configs/nemotron_nano_v3_30ba3b_math_2node.py new file mode 100644 index 0000000..f39d04a --- /dev/null +++ b/nemo-rl/configs/nemotron_nano_v3_30ba3b_math_2node.py @@ -0,0 +1,39 @@ +"""Nemotron-3-Nano-30B-A3B (MoE) GRPO on OpenMathInstruct-2 — 2 nodes x 8 GPUs. + +Uses NeMo-RL's convergence-tested 2-node recipe +``examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-fsdp2.yaml`` verbatim as +the base config: FSDP2 with expert_parallel_size=8 for the 30B/3B-active MoE, +nemo-automodel TransformerEngine backend + DeepEP, and vLLM tensor_parallel_size=4 +for generation. Modal provisions a 2-node RDMA cluster (16xH100) and the driver +runs on the head node. + +This recipe is only shipped in NeMo-RL >= v0.6.0, so this config pins the +matching v0.6.0 image rather than the repo default (v0.5.0). +""" + +from configs.base import ModalConfig, NemoRLConfig + +modal = ModalConfig(gpu="H100", docker_image="nvcr.io/nvidia/nemo-rl:v0.6.0") + + +class _Recipe(NemoRLConfig): + entrypoint = "examples/run_grpo.py" + base_config = "examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-fsdp2.yaml" + + num_nodes = 2 + gpus_per_node = 8 + + # Model and tokenizer live in separate HF repos; only the model is prefetched + # here, the tokenizer is fetched on first use. + hf_model = "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16" + hf_datasets = ["nvidia/OpenMathInstruct-2"] + + # The recipe already sets every MoE/parallelism knob; only override logging. + overrides = { + "logger.wandb_enabled": True, + "logger.wandb.project": "nemo-rl-grpo", + "logger.wandb.name": "nemotron-nano-v3-30ba3b-math-2node", + } + + +nemo_rl = _Recipe() diff --git a/nemo-rl/configs/qwen2_5_1_5b_math.py b/nemo-rl/configs/qwen2_5_1_5b_math.py new file mode 100644 index 0000000..da56cd8 --- /dev/null +++ b/nemo-rl/configs/qwen2_5_1_5b_math.py @@ -0,0 +1,33 @@ +"""Qwen2.5-1.5B GRPO on OpenMathInstruct-2 — single node, 8 GPUs. + +Mirrors NeMo-RL's examples/configs/grpo_math_1B.yaml, bumped to a full 8-GPU +node. The dataset is downloaded automatically by NeMo-RL at runtime into the +mounted HF cache, so download_data is only used to warm that cache. +""" + +from configs.base import ModalConfig, NemoRLConfig + +modal = ModalConfig(gpu="H100") + + +class _Recipe(NemoRLConfig): + entrypoint = "examples/run_grpo.py" + base_config = "examples/configs/grpo_math_1B.yaml" + + num_nodes = 1 + gpus_per_node = 8 + + hf_model = "Qwen/Qwen2.5-1.5B" + hf_datasets = ["nvidia/OpenMathInstruct-2"] + + overrides = { + "policy.model_name": "Qwen/Qwen2.5-1.5B", + # Spread generation + training across all 8 GPUs. + "policy.generation.vllm_cfg.tensor_parallel_size": 1, + "logger.wandb_enabled": True, + "logger.wandb.project": "nemo-rl-grpo", + "logger.wandb.name": "qwen2.5-1.5b-math", + } + + +nemo_rl = _Recipe() diff --git a/nemo-rl/configs/qwen2_5_1_5b_math_2node.py b/nemo-rl/configs/qwen2_5_1_5b_math_2node.py new file mode 100644 index 0000000..f47ca19 --- /dev/null +++ b/nemo-rl/configs/qwen2_5_1_5b_math_2node.py @@ -0,0 +1,35 @@ +"""Qwen2.5-1.5B GRPO on OpenMathInstruct-2 — 2 nodes x 8 GPUs. + +Same recipe as qwen2_5_1_5b_math (grpo_math_1B.yaml), scaled to two nodes. +Modal provisions a 2-node RDMA cluster and NeMo-RL schedules across all 16 GPUs. +""" + +from configs.base import ModalConfig, NemoRLConfig + +modal = ModalConfig(gpu="H100") + + +class _Recipe(NemoRLConfig): + entrypoint = "examples/run_grpo.py" + base_config = "examples/configs/grpo_math_1B.yaml" + + num_nodes = 2 + gpus_per_node = 8 + + hf_model = "Qwen/Qwen2.5-1.5B" + hf_datasets = ["nvidia/OpenMathInstruct-2"] + + overrides = { + "policy.model_name": "Qwen/Qwen2.5-1.5B", + # Base config ties max_input_seq_length, max_model_len, and max_new_tokens + # all to max_total_sequence_length (512). A 512-token prompt then leaves no + # room for generation and vLLM raises. Bump it so prompt + output fit. + "policy.max_total_sequence_length": 1024, + "policy.generation.vllm_cfg.tensor_parallel_size": 1, + "logger.wandb_enabled": True, + "logger.wandb.project": "nemo-rl-grpo", + "logger.wandb.name": "qwen2.5-1.5b-math-2node", + } + + +nemo_rl = _Recipe() diff --git a/nemo-rl/configs/qwen3_8b_math.py b/nemo-rl/configs/qwen3_8b_math.py new file mode 100644 index 0000000..e91652e --- /dev/null +++ b/nemo-rl/configs/qwen3_8b_math.py @@ -0,0 +1,33 @@ +"""Qwen2.5-1.5B GRPO on OpenMathInstruct-2 — single node, 8 GPUs. + +Mirrors NeMo-RL's examples/configs/grpo_math_1B.yaml, bumped to a full 8-GPU +node. The dataset is downloaded automatically by NeMo-RL at runtime into the +mounted HF cache, so download_data is only used to warm that cache. +""" + +from configs.base import ModalConfig, NemoRLConfig + +modal = ModalConfig(gpu="H100") + + +class _Recipe(NemoRLConfig): + entrypoint = "examples/run_grpo.py" + base_config = "examples/configs/grpo_math_1B.yaml" + + num_nodes = 1 + gpus_per_node = 8 + + hf_model = "Qwen/Qwen3-8B" + hf_datasets = ["nvidia/OpenMathInstruct-2"] + + overrides = { + "policy.model_name": "Qwen/Qwen3-8B", + # Spread generation + training across all 8 GPUs. + "policy.generation.vllm_cfg.tensor_parallel_size": 1, + "logger.wandb_enabled": True, + "logger.wandb.project": "nemo-rl-grpo", + "logger.wandb.name": "qwen3-8b-math", + } + + +nemo_rl = _Recipe() diff --git a/nemo-rl/modal_helpers/__init__.py b/nemo-rl/modal_helpers/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/nemo-rl/modal_helpers/run_grpo_multinode.py b/nemo-rl/modal_helpers/run_grpo_multinode.py new file mode 100644 index 0000000..d2d54aa --- /dev/null +++ b/nemo-rl/modal_helpers/run_grpo_multinode.py @@ -0,0 +1,113 @@ +""" +Helpers to run NemoRL multinode jobs on Modal. + +There are two problems that break colocated training across Modal nodes in NemoRL: + +- Placement: Colocated VLLM has fractional Ray GPUs (0.5 per actor) so all 16 workers + could fit in one node's placement group, which leaves the other node empty (in NemoRL, SLURM scheduler ensures each pg goes to a separate node). We force a + single placement group of 16 whole-GPU bundles and spread them across nodes with Ray's + SPREAD, whereas the default NemoRL recipe will create one placement group per node. This way gpus_per_node bundles are put into one node, which ensures one rank per GPU. + +- Once placement is right, Modal gives every container the same NCCL_HOSTID and same hostname, and NemoRL will copy the driver's env onto every worker. +Because of this, when NCCL keys on (hostHash, busId) it will find two gpus on the same rank on different nodes (ie. gpu 2 on nodes 0 and 1) are the same device and give a duplicate GPU error. + +To fix this we manually set each worker's NCCL_HOSTID from the physical Ray node its bundle landed on, so ranks on different nodes look distinct. +""" + +from __future__ import annotations + +import os +import runpy + +# Force bundles to spread across nodes. NemoRL defaults to one placement group per node +# (relying on SLURM to keep them apart); on Modal we use a single unified placement group +# so Ray's SPREAD strategy balances whole-GPU bundles across nodes instead of packing them +# onto one. +def _patch_virtual_cluster() -> None: + from nemo_rl.distributed.virtual_cluster import RayVirtualCluster + + if getattr(RayVirtualCluster, "_modal_multinode_patched", False): + return + + _orig_init_pg = RayVirtualCluster._init_placement_groups + + def _init_placement_groups(self, strategy=None, use_unified_pg=False): + if len(self._bundle_ct_per_node_list) > 1 and self.use_gpus: + use_unified_pg = True + return _orig_init_pg(self, strategy=strategy, use_unified_pg=use_unified_pg) + + RayVirtualCluster._init_placement_groups = _init_placement_groups + RayVirtualCluster._modal_multinode_patched = True + + +def _bundle_node_id(placement_group, bundle_index: int) -> str | None: + """Physical Ray node id hosting a given bundle of a placement group. + + The node id is unique per machine, so it's the right seed for a per-node NCCL + host hash. Best-effort: returns None if the placement table isn't populated. + """ + from ray.util.placement_group import placement_group_table + + try: + table = placement_group_table(placement_group) + return table.get("bundles_to_node_id", {}).get(bundle_index) + except Exception: + return None + + +# Modal gives every container the same NCCL_HOSTID, so NCCL keys on (hostHash, busId) +# and treats same-busId GPUs on different nodes as one device ("Duplicate GPU detected"). +# We override NCCL_HOSTID per physical node. +def _patch_worker_group_hostid() -> None: + """Give each worker an NCCL_HOSTID keyed to its physical node.""" + from nemo_rl.distributed import worker_groups as wg + + if getattr(wg.RayWorkerBuilder, "_modal_hostid_patched", False): + return + + _orig_create_worker_async = wg.RayWorkerBuilder.create_worker_async + + def create_worker_async( + self, + placement_group, + placement_group_bundle_index, + num_gpus, + bundle_indices=None, + **extra_options, + ): + node_id = _bundle_node_id(placement_group, placement_group_bundle_index) + if node_id is not None: + env_vars = extra_options.setdefault("runtime_env", {}).setdefault( + "env_vars", {} + ) + env_vars["NCCL_HOSTID"] = f"modal-node-{node_id}" + print( + f"[modal] bundle={placement_group_bundle_index} " + f"NCCL_HOSTID=modal-node-{node_id}", + flush=True, + ) + return _orig_create_worker_async( + self, + placement_group, + placement_group_bundle_index, + num_gpus, + bundle_indices, + **extra_options, + ) + + wg.RayWorkerBuilder.create_worker_async = create_worker_async + wg.RayWorkerBuilder._modal_hostid_patched = True + + +def main() -> None: + # Surface NCCL's per-rank hostHash/busId so duplicate-GPU errors are diagnosable. + # This is set in the driver env, which NeMo-RL copies onto every worker. + os.environ.setdefault("NCCL_DEBUG", "INFO") + os.environ.setdefault("NCCL_DEBUG_SUBSYS", "INIT,ENV") + _patch_virtual_cluster() + _patch_worker_group_hostid() + runpy.run_path("examples/run_grpo.py", run_name="__main__") + + +if __name__ == "__main__": + main() diff --git a/nemo-rl/modal_helpers/utils.py b/nemo-rl/modal_helpers/utils.py new file mode 100644 index 0000000..305c7cc --- /dev/null +++ b/nemo-rl/modal_helpers/utils.py @@ -0,0 +1,191 @@ +""" +General helpers for Ray multinode +""" +import json +import shlex +import subprocess +import time + +# Glob covering both lib/ and lib64/ python site-packages in the NeMo-RL venv. +_NSIGHT_GLOB = ( + "/opt/nemo_rl_venv/lib*/python*/site-packages/ray/_private/runtime_env/nsight.py" +) + +# Keep Ray worker ports below the OS ephemeral range (see NeMo-RL ray.sub). +_MIN_WORKER_PORT = 10002 +_MAX_WORKER_PORT = 11000 + + +def _ray_node_resources(gpus_per_node: int, node_rank: int | None = None) -> str: + """Custom Ray resources registered by NeMo-RL's ray.sub on every node.""" + resources: dict[str, int] = { + "worker_units": gpus_per_node, + "slurm_managed_ray_cluster": 1, + } + if node_rank is not None: + resources[f"modal_node_{node_rank}"] = 1 + return json.dumps(resources, separators=(",", ":")) + + +def cluster_driver_env(head_ip: str, cluster_ips: list[str] | None = None) -> dict[str, str]: + """Environment for the NeMo-RL driver and Ray actors on a multi-node cluster.""" + no_proxy_hosts = ",".join(dict.fromkeys(["127.0.0.1", head_ip, *(cluster_ips or [])])) + return { + # NeMo-RL uses `uv run` itself; disable Ray's per-task uv runtime env (ray.sub). + "RAY_ENABLE_UV_RUN_RUNTIME_ENV": "0", + "TRAIN_ENABLE_SHARE_CUDA_VISIBLE_DEVICES": "0", + "MASTER_ADDR": head_ip, + "no_proxy": no_proxy_hosts, + "NCCL_NVLS_ENABLE": "0", + "CUDA_DEVICE_MAX_CONNECTIONS": "1", + } + + +def get_modal_cluster_context(n_nodes: int) -> tuple[int, str, str, int, list[str]]: + """Return (rank, head_ip, my_ip, n_nodes, cluster_ips) for the current Modal cluster.""" + if n_nodes == 1: + return 0, "127.0.0.1", "127.0.0.1", 1, ["127.0.0.1"] + + import modal.experimental + + info = modal.experimental.get_cluster_info() + actual_nodes = len(info.container_ipv4_ips) + if actual_nodes != n_nodes: + raise RuntimeError( + f"cluster size mismatch: expected {n_nodes} node(s), got {actual_nodes}" + ) + return ( + info.rank, + info.container_ipv4_ips[0], + info.container_ipv4_ips[info.rank], + actual_nodes, + list(info.container_ipv4_ips), + ) + +# TODO: probably unnecessary and can be removed +def _patch_nsight() -> None: + """Mirror ray.sub's nsight patch so Ray honors NeMo-RL's py_executable. + + NeMo-RL launches workers via ``uv run``; Ray's nsight runtime-env plugin + otherwise hardcodes ``python`` and breaks profiling. Best-effort. + """ + import glob + + sed = ( + r's/context\.py_executable = " "\.join(self\.nsight_cmd) + " python"/' + r'context.py_executable = " ".join(self.nsight_cmd) + f" {context.py_executable}"/g' + ) + for path in glob.glob(_NSIGHT_GLOB): + subprocess.run(["sed", "-i", sed, path], check=False) + + +def _wait_for_ray_cluster( + ray, n_nodes: int, gpus_per_node: int, timeout_s: int = 240 +) -> None: + """Block until each Ray node has registered its GPUs (not just cluster-wide total).""" + expected_gpus = n_nodes * gpus_per_node + for _ in range(timeout_s // 2): + alive = [n for n in ray.nodes() if n["Alive"]] + total_gpus = int(ray.cluster_resources().get("GPU", 0)) + per_node = sorted( + ( + n.get("NodeManagerAddress"), + int(n.get("Resources", {}).get("GPU", 0)), + ) + for n in alive + ) + ready_nodes = sum(1 for _, g in per_node if g >= gpus_per_node) + print( + f"Waiting for cluster: {len(alive)}/{n_nodes} nodes, " + f"{total_gpus}/{expected_gpus} GPUs, " + f"{ready_nodes}/{n_nodes} nodes with >={gpus_per_node} GPU(s)" + ) + print(f" per-node: {per_node}") + if ( + len(alive) >= n_nodes + and total_gpus >= expected_gpus + and ready_nodes >= n_nodes + ): + return + time.sleep(2) + raise RuntimeError( + f"Timed out waiting for {n_nodes} nodes each with {gpus_per_node} GPUs " + f"(cluster total {expected_gpus})" + ) + +# Nemo-RL launches training through a SLURM bash script (https://github.com/NVIDIA-NeMo/RL/blob/main/ray.sub) getting nodes from a list SLURM_JOBS_NODELIST +# instead we break this up into start_ray_head and start_ray_worker +def start_ray_head( + head_ip: str, port: int, n_nodes: int, gpus_per_node: int, node_rank: int = 0 +) -> None: + """Start the Ray head and block until every node and GPU has joined.""" + import ray + + _patch_nsight() + subprocess.Popen( + [ + "ray", + "start", + "--head", + "--disable-usage-stats", + f"--num-gpus={gpus_per_node}", + f"--resources={_ray_node_resources(gpus_per_node, node_rank)}", + f"--node-ip-address={head_ip}", + f"--port={port}", + "--dashboard-host=0.0.0.0", + "--include-dashboard=True", + "--block", + ] + ) + + for _ in range(60): + try: + ray.init(address="auto") + break + except Exception: + time.sleep(2) + else: + raise RuntimeError("Ray head node failed to start") + + _wait_for_ray_cluster(ray, n_nodes, gpus_per_node) + # Detach the driver-side ray handle; the NeMo-RL driver reconnects itself. + ray.shutdown() + + +def start_ray_worker( + head_ip: str, port: int, my_ip: str, gpus_per_node: int, node_rank: int +) -> None: + """Start a Ray worker that joins the head and blocks forever.""" + _patch_nsight() + subprocess.Popen( + [ + "ray", + "start", + f"--node-ip-address={my_ip}", + "--address", + f"{head_ip}:{port}", + "--disable-usage-stats", + f"--num-gpus={gpus_per_node}", + f"--resources={_ray_node_resources(gpus_per_node, node_rank)}", + f"--min-worker-port={_MIN_WORKER_PORT}", + f"--max-worker-port={_MAX_WORKER_PORT}", + "--block", + ] + ) + + + +def build_train_cmd(nemo_rl_cfg, nemo_rl_root: str, experiment: str | None = None) -> str: + """Build the driver command run on the Ray head node.""" + import importlib.util + + args = shlex.join(nemo_rl_cfg.cli_args(experiment)) + if nemo_rl_cfg.num_nodes > 1: + + spec = importlib.util.find_spec("modal_helpers.run_grpo_multinode") + if spec is None or spec.origin is None: + raise RuntimeError("modal_helpers.run_grpo_multinode not found in image") + entrypoint = spec.origin + else: + entrypoint = nemo_rl_cfg.entrypoint + return f"cd {shlex.quote(nemo_rl_root)} && uv run python {shlex.quote(entrypoint)} {args}" diff --git a/nemo-rl/modal_train.py b/nemo-rl/modal_train.py new file mode 100644 index 0000000..3fd759f --- /dev/null +++ b/nemo-rl/modal_train.py @@ -0,0 +1,204 @@ +import asyncio +import os +import subprocess + +import modal +import modal.experimental + +from configs import get_module, _CONFIGS_DIR +from configs.base import ( + CHECKPOINTS_PATH, + HF_CACHE_PATH, + NEMO_RL_ROOT, + ModalConfig, +) + +# ── Experiment (client-side only — feeds decorator params) ──────────────────── + +experiment = os.environ.get("EXPERIMENT_CONFIG", "") +exp_mod = get_module(experiment) if experiment else None +modal_cfg = exp_mod.modal if exp_mod else None +nemo_rl_cfg = exp_mod.nemo_rl if exp_mod else None + +# ── Image ───────────────────────────────────────────────────────────────────── + +image = ( + modal.Image.from_registry( + modal_cfg.docker_image if modal_cfg else ModalConfig.docker_image + ) + .entrypoint([]) + .add_local_python_source("configs", copy=True) + .add_local_python_source("modal_helpers", copy=True) +) +if modal_cfg: + if modal_cfg.local_nemo_rl: + image = image.add_local_dir( + modal_cfg.local_nemo_rl, + remote_path=NEMO_RL_ROOT, + copy=True, + ignore=[ + "**/__pycache__", + "**/*.pyc", + "**/.git", + "**/.venv", + "**/results", + "**/logs", + ], + ) + if modal_cfg.image_run_commands: + image = image.run_commands(*modal_cfg.image_run_commands) + if modal_cfg.image_env: + image = image.env(modal_cfg.image_env) + +with image.imports(): + from modal_helpers.utils import ( + build_train_cmd, + cluster_driver_env, + get_modal_cluster_context, + start_ray_head, + start_ray_worker, + ) + +# ── Volumes ─────────────────────────────────────────────────────────────────── + +hf_cache_volume = modal.Volume.from_name("huggingface-cache", create_if_missing=True) +checkpoints_volume = modal.Volume.from_name( + "nemo-rl-checkpoints", create_if_missing=True +) + +modal_volumes = { + str(HF_CACHE_PATH): hf_cache_volume, + str(CHECKPOINTS_PATH): checkpoints_volume, +} + +# ── App ────────────────────────────────────────────────────────────────────── + +app = modal.App(experiment) + +RAY_PORT = 6379 +RAY_DASHBOARD_PORT = 8265 + + +@app.local_entrypoint() +def list_configs(): + """Print all available experiments.""" + _skip = {"base", "__init__"} + names = sorted(f.stem for f in _CONFIGS_DIR.glob("*.py") if f.stem not in _skip) + print("Available experiments:") + for name in names: + mod = get_module(name) + cfg = mod.nemo_rl + gpu = f"{mod.modal.gpu}:{cfg.gpus_per_node}" + print( + f" {name:<40} {cfg.num_nodes} node(s) × {gpu} " + f"({os.path.basename(cfg.entrypoint)})" + ) + + +@app.function( + image=image, + volumes={str(HF_CACHE_PATH): hf_cache_volume}, + timeout=4 * 60 * 60, + secrets=[modal.Secret.from_name("huggingface-secret")], +) +def download_model(experiment: str = os.environ.get("EXPERIMENT_CONFIG", "")): + """Prefetch the recipe's model into the mounted HF cache.""" + from huggingface_hub import snapshot_download + + cfg = get_module(experiment).nemo_rl + model_id = cfg.model_id() + if not model_id: + raise ValueError( + f"{experiment!r}: set hf_model or overrides['policy.model_name'] to download a model" + ) + hf_cache_volume.reload() + print(f"Downloading model {model_id!r} into the HF cache...") + snapshot_download(model_id) + hf_cache_volume.commit() + + +@app.function( + image=image, + volumes={str(HF_CACHE_PATH): hf_cache_volume}, + timeout=4 * 60 * 60, + secrets=[modal.Secret.from_name("huggingface-secret")], +) +def download_data(experiment: str = os.environ.get("EXPERIMENT_CONFIG", "")): + # fetch dataset from huggingface (optional, nemorl will do this automatically) + cfg = get_module(experiment).nemo_rl + if not cfg.hf_datasets: + print(f"{experiment!r} declares no hf_datasets; nothing to prefetch.") + return + + from datasets import load_dataset + + hf_cache_volume.reload() + for repo_id in cfg.hf_datasets: + print(f"Prefetching dataset {repo_id!r}...") + load_dataset(repo_id) + hf_cache_volume.commit() + + + +@app.function( + image=image, + gpu=f"{modal_cfg.gpu}:{nemo_rl_cfg.gpus_per_node}" if modal_cfg else None, + memory=modal_cfg.memory if modal_cfg and modal_cfg.memory else None, + cloud=modal_cfg.cloud if modal_cfg and modal_cfg.cloud else None, + region=modal_cfg.region if modal_cfg and modal_cfg.region else None, + volumes=modal_volumes, + secrets=[ + modal.Secret.from_name("huggingface-secret"), + modal.Secret.from_name("wandb-secret"), + ], + timeout=24 * 60 * 60, + experimental_options={"efa_enabled": True}, +) +@( + modal.experimental.clustered(nemo_rl_cfg.total_nodes(), rdma=True) + if nemo_rl_cfg + else lambda fn: fn +) +async def train(experiment: str = os.environ.get("EXPERIMENT_CONFIG", "")): + await asyncio.gather( + hf_cache_volume.reload.aio(), + checkpoints_volume.reload.aio(), + ) + exp_mod = get_module(experiment) + cfg = exp_mod.nemo_rl + modal_cfg = exp_mod.modal + + rank, head_ip, my_ip, n_nodes, cluster_ips = get_modal_cluster_context( + cfg.total_nodes() + ) + + # on all ranks that are not the head node, do ray --start and spin + if rank != 0: + + start_ray_worker(head_ip, RAY_PORT, my_ip, cfg.gpus_per_node, rank) + while True: + await asyncio.sleep(10) + + # on the head node, start ray and wait for all nodes to join + start_ray_head(head_ip, RAY_PORT, n_nodes, cfg.gpus_per_node, node_rank=rank) + + + cmd = build_train_cmd(cfg, NEMO_RL_ROOT, experiment) + env = { + **os.environ, + "RAY_ADDRESS": f"{head_ip}:{RAY_PORT}", + **cluster_driver_env(head_ip, cluster_ips), + **cfg.environment, + } + + gpu = f"{modal_cfg.gpu}:{cfg.gpus_per_node}" + print(f"Training {experiment:<40} {n_nodes} node(s) × {gpu}") + print(f"Command: {cmd}") + + async with modal.forward(RAY_DASHBOARD_PORT) as tunnel: + print(f"Ray dashboard: {tunnel.url}") + result = subprocess.run(["bash", "-c", cmd], env=env) + + await checkpoints_volume.commit.aio() + if result.returncode != 0: + raise RuntimeError(f"NeMo-RL driver exited with code {result.returncode}") From 1dec9713107156674cd5c7b9829fa7864ade076a Mon Sep 17 00:00:00 2001 From: kailash Date: Fri, 10 Jul 2026 15:37:24 +0000 Subject: [PATCH 2/6] added working configs for glm5.2 lora on gsm8k --- miles/configs/glm5_2_744b_a40b_lora.py | 190 ++++++++++++++++++ miles/configs/glm5_2_744b_a40b_lora_5layer.py | 152 ++++++++++++++ miles/configs/glm5_2_744b_a40b_lora_dapo.py | 155 ++++++++++++++ miles/modal_train.py | 41 +++- miles/modal_train_glm_test.py | 134 ++++++++++++ 5 files changed, 663 insertions(+), 9 deletions(-) create mode 100644 miles/configs/glm5_2_744b_a40b_lora.py create mode 100644 miles/configs/glm5_2_744b_a40b_lora_5layer.py create mode 100644 miles/configs/glm5_2_744b_a40b_lora_dapo.py create mode 100644 miles/modal_train_glm_test.py diff --git a/miles/configs/glm5_2_744b_a40b_lora.py b/miles/configs/glm5_2_744b_a40b_lora.py new file mode 100644 index 0000000..db5be4c --- /dev/null +++ b/miles/configs/glm5_2_744b_a40b_lora.py @@ -0,0 +1,190 @@ +"""GLM-5.2 (full 744B-A40B) LoRA GRPO — 8 nodes x 8 H200, colocated. + +this configuration is adapted from: +https://github.com/radixark/miles/blob/c4d9d49cbf8a39185f4c80c0f6084836fc759819/launch_glm_rl_att_unfused_moe.sh + +The launcher was later deleted. This config uses its Megatron/BSHD branch with +a smaller padding quantum to control activation memory. The topology is 8 nodes +x 8 H200, EP 32, DP 8, TP 8, PP 1, CP 1. + +on GSm8k dataset with max response length 256. + +to run: + + EXPERIMENT_CONFIG=glm5_2_744b_a40b_lora uv run modal run miles/modal_train.py::download_model + EXPERIMENT_CONFIG=glm5_2_744b_a40b_lora uv run modal run miles/modal_train.py::download_data + EXPERIMENT_CONFIG=glm5_2_744b_a40b_lora uv run modal run miles/modal_train.py::train +""" + +from configs.base import ModalConfig, MilesConfig, DATA_PATH, CHECKPOINTS_PATH, HF_CACHE_PATH + +modal = ModalConfig( + docker_image="radixark/miles:dev-202607090055", # validated versions sglang 0.5.15, Megatron-Bridge 0.5.0, PR #1559 + #1593 for latest lora support + gpu="H200", + memory=(1024, int(2 * 1024 * 1024)), + image_run_commands=[ + f"rm -rf {HF_CACHE_PATH} 2>/dev/null || true", + + "rm -rf /usr/local/lib/python3.12/dist-packages/nvidia/cudnn/ 2>/dev/null || true", + + "pip install --no-cache-dir hf_xet", + ], + image_env={ + "LD_LIBRARY_PATH": "/usr/lib/x86_64-linux-gnu:$LD_LIBRARY_PATH", + + "HF_XET_HIGH_PERFORMANCE": "1", # for downloading + }, +) + + +class _Miles(MilesConfig): + # Architecture only (MODEL_ARGS); --spec inside it is inert under bridge LoRA. + miles_model_script = "scripts/models/glm5.2-744B-A40B_lora.sh" + + environment = { + "PYTHONPATH": "/root/Megatron-LM/", + "CUDA_DEVICE_MAX_CONNECTIONS": "1", + "NCCL_NVLS_ENABLE": "1", + # extra env vars from run_glm5_2_744b_a40b_lora.py + "MILES_EXPERIMENTAL_ROLLOUT_REFACTOR": "1", + "INDEXER_ROPE_NEOX_STYLE": "0", + "SGLANG_NSA_FORCE_MLA": "1", + } + + + hf_checkpoint = "zai-org/GLM-5.2" + megatron_to_hf_mode = "bridge" + + # tilelang + thd backward pass produces nan gradients under nonzero loss (TODO: figure out why?) + # keep megatron for now -- upstream config had data + dsa_attention_backend = "megatron" + qkv_format = "bshd" + data_pad_size_multiplier = 32 + micro_batch_size = 1 + save = f"{CHECKPOINTS_PATH}/GLM-5.2-lora-ckpt" + save_interval = 20 + + + actor_num_nodes = 8 + actor_num_gpus_per_node = 8 + num_gpus_per_node = 8 + colocate = True + use_miles_router = True + calculate_per_token_loss = True + tensor_model_parallel_size = 8 + sequence_parallel = True + pipeline_model_parallel_size = 1 + context_parallel_size = 1 + expert_model_parallel_size = 32 + expert_tensor_parallel_size = 1 + moe_token_dispatcher_type = "alltoall" + + # attention/MLA on every layer + per expert linear_fc1 only on last 10 moe layers + # exclude expert down proj + lora_rank = 8 + lora_alpha = 16 + lora_dropout = 0.0 + target_modules = ( + "q_a_proj,kv_a_proj_with_mqa,q_b_proj,kv_b_proj,o_proj," + "*.layers.68.*.linear_fc1,*.layers.69.*.linear_fc1," + "*.layers.70.*.linear_fc1,*.layers.71.*.linear_fc1," + "*.layers.72.*.linear_fc1,*.layers.73.*.linear_fc1," + "*.layers.74.*.linear_fc1,*.layers.75.*.linear_fc1," + "*.layers.76.*.linear_fc1,*.layers.77.*.linear_fc1" + ) + experts_shared_outer_loras = False + lora_base_cpu_backup = True + no_gradient_accumulation_fusion = True + + + prompt_data = f"{DATA_PATH}/gsm8k/train.parquet" + input_key = "messages" + label_key = "label" + apply_chat_template = True + rollout_shuffle = True + rm_type = "math" + + + num_rollout = 50 + rollout_batch_size = 8 + n_samples_per_prompt = 16 + rollout_max_response_len = 256 + rollout_temperature = 1.0 + global_batch_size = 64 + use_rollout_routing_replay = True + + + advantage_estimator = "grpo" + kl_loss_coef = 0.0 + kl_loss_type = "low_var_kl" + kl_coef = 0.0 + entropy_coef = 0.0 + eps_clip = 0.2 + eps_clip_high = 0.28 + + + optimizer = "adam" + lr = 1e-5 + lr_decay_style = "constant" + weight_decay = 0.1 + adam_beta1 = 0.9 + adam_beta2 = 0.98 + optimizer_cpu_offload = True + overlap_cpu_optimizer_d2h_h2d = True + use_precision_aware_optimizer = True + + + attention_dropout = 0.0 + hidden_dropout = 0.0 + accumulate_allreduce_grads_in_fp32 = True + attention_softmax_in_fp32 = True + attention_backend = "flash" + + # do bf16 sglang rollout -- todo try fp8 rollout and compare logprob diff + rollout_num_gpus_per_engine = 32 + sglang_mem_fraction_static = 0.7 + sglang_enable_dp_attention = True + sglang_ep_size = 32 + sglang_dp_size = 32 + sglang_moe_dense_tp_size = 1 + sglang_enable_dp_lm_head = True + sglang_attention_backend = "nsa" + sglang_nsa_decode_backend = "flashmla_sparse" + sglang_nsa_prefill_backend = "flashmla_sparse" + sglang_page_size = 64 + sglang_cuda_graph_max_bs = 64 + sglang_max_running_requests = 512 + sglang_chunked_prefill_size = 65536 + sglang_watchdog_timeout = 3600 + sglang_moe_runner_backend = "triton" + sglang_disable_shared_experts_fusion = True + sglang_max_lora_rank = 16 + sglang_lora_backend = "triton" + sglang_lora_use_virtual_experts = True + + + use_wandb = True + wandb_project = "miles-run_glm5_2_744b_a40b_lora" + wandb_group = "glm5.2-744B-8node-no-down-proj-megatron-pad32-modal" + disable_wandb_random_suffix = True + + def download_model(self) -> None: + + from huggingface_hub import snapshot_download + + snapshot_download(self.hf_checkpoint, max_workers=32) + + def download_data(self) -> None: + import os + + from huggingface_hub import snapshot_download + + os.makedirs(f"{DATA_PATH}/gsm8k", exist_ok=True) + snapshot_download( + repo_id="zhuzilin/gsm8k", + repo_type="dataset", + local_dir=f"{DATA_PATH}/gsm8k", + ) + + +miles = _Miles() diff --git a/miles/configs/glm5_2_744b_a40b_lora_5layer.py b/miles/configs/glm5_2_744b_a40b_lora_5layer.py new file mode 100644 index 0000000..ebfb0a3 --- /dev/null +++ b/miles/configs/glm5_2_744b_a40b_lora_5layer.py @@ -0,0 +1,152 @@ +"""GLM-5.2 (744B-A40B arch, 5-layer prune) LoRA GRPO — single node, colocated. + +Smoke test for the GLM-5.2 bridge-mode DSA LoRA path. ``Pinaster/GLM-5.2_5layer`` +is a 5-layer prune (3 dense + 2 MoE) of GLM-5.2 that keeps one computing + one +skip layer, so it exercises the same DSA cross-layer index-sharing, MoE, bridge +LoRA, and sglang MoE-LoRA serving path as the full 744B model at toy cost. + +Ports ``scripts/run_glm5_2_744b_a40b_lora.py`` (which the guide launcher does NOT +run directly) into config attributes: the model ``.sh`` supplies architecture +only, every LoRA/DSA/sglang flag is set here and forwarded by ``cli_args()``. + +Requires a miles image built after PR #1559 (GLM-5/5.1/5.2 LoRA) and PR #1593 +(bridge-LoRA recompute fix); the repo default dev-202605291323 predates both. + +Launched by the dedicated smoke harness (pinned to this config): + uv run modal run miles/modal_train_glm_test.py::download_model + uv run modal run miles/modal_train_glm_test.py::download_data + uv run modal run miles/modal_train_glm_test.py::train +""" + +from configs.base import ModalConfig, MilesConfig, DATA_PATH, CHECKPOINTS_PATH, HF_CACHE_PATH + +modal = ModalConfig( + docker_image="radixark/miles:dev-202607090055", + gpu="H200", + memory=(1024, int(2 * 1024 * 1024)), + image_run_commands=[ + f"rm -rf {HF_CACHE_PATH} 2>/dev/null || true", + "rm -rf /usr/local/lib/python3.12/dist-packages/nvidia/cudnn/ 2>/dev/null || true", + ], + image_env={"LD_LIBRARY_PATH": "/usr/lib/x86_64-linux-gnu:$LD_LIBRARY_PATH"}, +) + + +class _Miles(MilesConfig): + miles_model_script = "scripts/models/glm5.2-744B-A40B_5layer_lora.sh" + + environment = { + "PYTHONPATH": "/root/Megatron-LM/", + "CUDA_DEVICE_MAX_CONNECTIONS": "1", + "NCCL_NVLS_ENABLE": "1", + "MILES_EXPERIMENTAL_ROLLOUT_REFACTOR": "1", + "INDEXER_ROPE_NEOX_STYLE": "0", + "SGLANG_NSA_FORCE_MLA": "1", + } + + hf_checkpoint = "Pinaster/GLM-5.2_5layer" + megatron_to_hf_mode = "bridge" + dsa_attention_backend = "tilelang" + qkv_format = "thd" + micro_batch_size = 1 + save = f"{CHECKPOINTS_PATH}/GLM-5.2_5layer-lora-ckpt" + save_interval = 1 + + actor_num_nodes = 1 + actor_num_gpus_per_node = 4 + colocate = True + use_miles_router = True + calculate_per_token_loss = True + tensor_model_parallel_size = 4 + sequence_parallel = True + pipeline_model_parallel_size = 1 + context_parallel_size = 1 + expert_model_parallel_size = 4 + expert_tensor_parallel_size = 1 + + lora_rank = 16 + lora_alpha = 32 + lora_dropout = 0.0 + target_modules = "q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj,q_a_proj,kv_a_proj_with_mqa,q_b_proj,kv_b_proj" + experts_shared_outer_loras = True + lora_base_cpu_backup = True + no_gradient_accumulation_fusion = True + + prompt_data = f"{DATA_PATH}/gsm8k/train.parquet" + input_key = "messages" + label_key = "label" + apply_chat_template = True + rollout_shuffle = True + rm_type = "math" + + num_rollout = 1 + rollout_batch_size = 4 + n_samples_per_prompt = 4 + rollout_max_response_len = 512 + rollout_temperature = 1.0 + global_batch_size = 16 + use_rollout_routing_replay = True + + advantage_estimator = "grpo" + kl_loss_coef = 0.0 + kl_loss_type = "low_var_kl" + kl_coef = 0.0 + entropy_coef = 0.0 + eps_clip = 0.2 + eps_clip_high = 0.28 + + optimizer = "adam" + lr = 1e-5 + lr_decay_style = "constant" + weight_decay = 0.1 + adam_beta1 = 0.9 + adam_beta2 = 0.98 + optimizer_cpu_offload = True + overlap_cpu_optimizer_d2h_h2d = True + use_precision_aware_optimizer = True + + attention_dropout = 0.0 + hidden_dropout = 0.0 + accumulate_allreduce_grads_in_fp32 = True + attention_softmax_in_fp32 = True + attention_backend = "flash" + + rollout_num_gpus_per_engine = 2 + sglang_mem_fraction_static = 0.5 + sglang_enable_dp_attention = True + sglang_ep_size = 2 + sglang_dp_size = 2 + sglang_moe_dense_tp_size = 1 + sglang_enable_dp_lm_head = True + sglang_attention_backend = "nsa" + sglang_nsa_decode_backend = "flashmla_sparse" + sglang_nsa_prefill_backend = "flashmla_sparse" + sglang_page_size = 64 + sglang_cuda_graph_max_bs = 64 + sglang_max_running_requests = 512 + sglang_chunked_prefill_size = 4096 + sglang_watchdog_timeout = 3600 + sglang_moe_runner_backend = "triton" + sglang_disable_shared_experts_fusion = True + sglang_max_lora_rank = 16 + sglang_lora_backend = "triton" + + use_wandb = True + wandb_project = "miles-run_glm5_2_744b_a40b_lora" + wandb_group = "glm5.2-5layer-lora" + disable_wandb_random_suffix = True + + def download_data(self) -> None: + import os + + from huggingface_hub import snapshot_download + + os.makedirs(f"{DATA_PATH}/gsm8k", exist_ok=True) + snapshot_download( + repo_id="zhuzilin/gsm8k", + repo_type="dataset", + local_dir=f"{DATA_PATH}/gsm8k", + ) + + +miles = _Miles() diff --git a/miles/configs/glm5_2_744b_a40b_lora_dapo.py b/miles/configs/glm5_2_744b_a40b_lora_dapo.py new file mode 100644 index 0000000..40bd7bf --- /dev/null +++ b/miles/configs/glm5_2_744b_a40b_lora_dapo.py @@ -0,0 +1,155 @@ +"""GLM-5.2 (full 744B-A40B) LoRA GRPO on dapo-math — long-context DSA variant. + +untested -- same recipe as glm5_2_744b_a40b_lora.py but on dapo-math task instead of gsm8k and increase +rollout respone length to 4096 tokens and context window to 8192 tokens + + EXPERIMENT_CONFIG=glm5_2_744b_a40b_lora_dapo uv run modal run miles/modal_train.py::download_model + EXPERIMENT_CONFIG=glm5_2_744b_a40b_lora_dapo uv run modal run miles/modal_train.py::download_data + EXPERIMENT_CONFIG=glm5_2_744b_a40b_lora_dapo uv run modal run miles/modal_train.py::train + +""" + +from configs.base import ModalConfig, MilesConfig, DATA_PATH, CHECKPOINTS_PATH, HF_CACHE_PATH + +modal = ModalConfig( + docker_image="radixark/miles:dev-202607090055", + gpu="H200", + memory=(1024, int(2 * 1024 * 1024)), + image_run_commands=[ + f"rm -rf {HF_CACHE_PATH} 2>/dev/null || true", + "rm -rf /usr/local/lib/python3.12/dist-packages/nvidia/cudnn/ 2>/dev/null || true", + "pip install --no-cache-dir hf_xet", + ], + image_env={ + "LD_LIBRARY_PATH": "/usr/lib/x86_64-linux-gnu:$LD_LIBRARY_PATH", + "HF_XET_HIGH_PERFORMANCE": "1", + }, +) + + +class _Miles(MilesConfig): + miles_model_script = "scripts/models/glm5.2-744B-A40B_lora.sh" + + environment = { + "PYTHONPATH": "/root/Megatron-LM/", + "CUDA_DEVICE_MAX_CONNECTIONS": "1", + "NCCL_NVLS_ENABLE": "1", + "MILES_EXPERIMENTAL_ROLLOUT_REFACTOR": "1", + "INDEXER_ROPE_NEOX_STYLE": "0", + "SGLANG_NSA_FORCE_MLA": "1", + } + + hf_checkpoint = "zai-org/GLM-5.2" + megatron_to_hf_mode = "bridge" + dsa_attention_backend = "tilelang" + qkv_format = "thd" + micro_batch_size = 1 + save = f"{CHECKPOINTS_PATH}/GLM-5.2-lora-dapo-ckpt" + save_interval = 1 + + actor_num_nodes = 1 + actor_num_gpus_per_node = 8 + colocate = True + use_miles_router = True + calculate_per_token_loss = True + tensor_model_parallel_size = 8 + sequence_parallel = True + pipeline_model_parallel_size = 1 + context_parallel_size = 1 + expert_model_parallel_size = 8 + expert_tensor_parallel_size = 1 + + lora_rank = 16 + lora_alpha = 32 + lora_dropout = 0.0 + target_modules = "q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj,q_a_proj,kv_a_proj_with_mqa,q_b_proj,kv_b_proj" + experts_shared_outer_loras = True + lora_base_cpu_backup = True + no_gradient_accumulation_fusion = True + + prompt_data = f"{DATA_PATH}/dapo-math-17k/dapo-math-17k.jsonl" + input_key = "prompt" + label_key = "label" + apply_chat_template = True + rollout_shuffle = True + rm_type = "math" + + num_rollout = 1 + rollout_batch_size = 4 + n_samples_per_prompt = 4 + rollout_max_response_len = 4096 + seq_length = 8192 + rollout_max_context_len = 8192 + rollout_temperature = 1.0 + global_batch_size = 16 + use_rollout_routing_replay = True + + advantage_estimator = "grpo" + kl_loss_coef = 0.0 + kl_loss_type = "low_var_kl" + kl_coef = 0.0 + entropy_coef = 0.0 + eps_clip = 0.2 + eps_clip_high = 0.28 + + optimizer = "adam" + lr = 1e-5 + lr_decay_style = "constant" + weight_decay = 0.1 + adam_beta1 = 0.9 + adam_beta2 = 0.98 + optimizer_cpu_offload = True + overlap_cpu_optimizer_d2h_h2d = True + use_precision_aware_optimizer = True + + attention_dropout = 0.0 + hidden_dropout = 0.0 + accumulate_allreduce_grads_in_fp32 = True + attention_softmax_in_fp32 = True + attention_backend = "flash" + + rollout_num_gpus_per_engine = 8 + sglang_mem_fraction_static = 0.5 + sglang_enable_dp_attention = True + sglang_ep_size = 8 + sglang_dp_size = 8 + sglang_moe_dense_tp_size = 1 + sglang_enable_dp_lm_head = True + sglang_attention_backend = "nsa" + sglang_nsa_decode_backend = "flashmla_kv" + sglang_nsa_prefill_backend = "flashmla_sparse" + sglang_kv_cache_dtype = "fp8_e4m3" + sglang_page_size = 64 + sglang_cuda_graph_max_bs = 256 + sglang_max_running_requests = 512 + sglang_chunked_prefill_size = 16384 + sglang_watchdog_timeout = 3600 + sglang_moe_runner_backend = "triton" + sglang_disable_shared_experts_fusion = True + sglang_max_lora_rank = 16 + sglang_lora_backend = "triton" + + use_wandb = True + wandb_project = "miles-run_glm5_2_744b_a40b_lora" + wandb_group = "glm5.2-744B-lora-dapo" + disable_wandb_random_suffix = True + + def download_model(self) -> None: + from huggingface_hub import snapshot_download + + snapshot_download(self.hf_checkpoint, max_workers=32) + + def download_data(self) -> None: + import os + + from huggingface_hub import snapshot_download + + os.makedirs(f"{DATA_PATH}/dapo-math-17k", exist_ok=True) + snapshot_download( + repo_id="zhuzilin/dapo-math-17k", + repo_type="dataset", + local_dir=f"{DATA_PATH}/dapo-math-17k", + ) + + +miles = _Miles() diff --git a/miles/modal_train.py b/miles/modal_train.py index 4ab3296..59e72d1 100644 --- a/miles/modal_train.py +++ b/miles/modal_train.py @@ -78,13 +78,39 @@ def run_config_hook(experiment: str, hook_name: str, mounted_volumes) -> None: - """Reload mounted volumes, run a MilesConfig hook, then commit them.""" + """Reload mounted volumes, run a MilesConfig hook, then commit them. + + Commits the mounted volumes every 2 min *while* the hook runs, not just at the + end, so a long HF pull persists completed shards as it goes. If the container dies + mid-download (network stall, disconnect), a re-run reloads the last commit and + snapshot_download skips the already-present files — resuming near where it stopped + instead of restarting from zero. + """ + import threading + miles_cfg = get_module(experiment).miles for volume in mounted_volumes: volume.reload() - getattr(miles_cfg, hook_name)() - for volume in mounted_volumes: - volume.commit() + + stop = threading.Event() + + def _periodic_commit() -> None: + while not stop.wait(120): + for volume in mounted_volumes: + try: + volume.commit() + except Exception as e: # best-effort; never kill the download + print(f"[modal] periodic volume commit failed: {e}", flush=True) + + committer = threading.Thread(target=_periodic_commit, daemon=True) + committer.start() + try: + getattr(miles_cfg, hook_name)() + finally: + stop.set() + committer.join(timeout=10) + for volume in mounted_volumes: + volume.commit() @app.local_entrypoint() @@ -290,11 +316,8 @@ async def train(experiment: str = os.environ.get("EXPERIMENT_CONFIG", "")): start_ray_head(my_ip, n_nodes) prepare_miles_config(miles_cfg, tempfile.mkdtemp()) - if (wandb_key := os.environ.get("WANDB_API_KEY", "")) and getattr( - miles_cfg, "use_wandb", False - ): - miles_cfg.wandb_key = wandb_key - + # W&B reads WANDB_API_KEY from the inherited Modal secret. Do not copy it + # into miles_cfg: CLI args are retained verbatim in Ray and Modal logs. cmd = build_train_cmd(miles_cfg, MILES_ROOT) runtime_env = { "env_vars": { diff --git a/miles/modal_train_glm_test.py b/miles/modal_train_glm_test.py new file mode 100644 index 0000000..eb74b2e --- /dev/null +++ b/miles/modal_train_glm_test.py @@ -0,0 +1,134 @@ +"""Dedicated 5-layer GLM-5.2 LoRA smoke-test launcher. + +glm5_2_744b_a40b_lora_5layer run on 1x8h200 + +Run: + uv run modal run miles/modal_train_glm_test.py::download_model + uv run modal run miles/modal_train_glm_test.py::download_data + uv run modal run -d miles/modal_train_glm_test.py::train +""" + +import asyncio +import os +import tempfile + +import modal + +from configs import get_module +from configs.base import HF_CACHE_PATH, DATA_PATH, CHECKPOINTS_PATH + +EXPERIMENT = "glm5_2_744b_a40b_lora_5layer" + +exp_mod = get_module(EXPERIMENT) +modal_cfg = exp_mod.modal +miles_cfg = exp_mod.miles + +MILES_ROOT = "/root/miles" + +image = ( + modal.Image.from_registry(modal_cfg.docker_image) + .entrypoint([]) + .add_local_python_source("configs", copy=True) + .add_local_python_source("modal_helpers", copy=True) +) +if modal_cfg.image_run_commands: + image = image.run_commands(*modal_cfg.image_run_commands) +if modal_cfg.image_env: + image = image.env(modal_cfg.image_env) + +with image.imports(): + from ray.job_submission import JobSubmissionClient + from modal_helpers.utils import ( + build_train_cmd, + prepare_miles_config, + start_ray_head, + ) + +hf_cache_volume = modal.Volume.from_name("huggingface-cache", create_if_missing=True) +data_volume = modal.Volume.from_name("miles-data", create_if_missing=True) +checkpoints_volume = modal.Volume.from_name("miles-checkpoints", create_if_missing=True) + +modal_volumes = { + str(HF_CACHE_PATH): hf_cache_volume, + str(DATA_PATH): data_volume, + str(CHECKPOINTS_PATH): checkpoints_volume, +} + +app = modal.App(f"{EXPERIMENT}-test") + +RAY_DASHBOARD_PORT = 8265 + + +def run_config_hook(hook_name: str, mounted_volumes) -> None: + cfg = get_module(EXPERIMENT).miles + for volume in mounted_volumes: + volume.reload() + getattr(cfg, hook_name)() + for volume in mounted_volumes: + volume.commit() + + +@app.function( + image=image, + volumes={str(HF_CACHE_PATH): hf_cache_volume}, + timeout=4 * 60 * 60, + secrets=[modal.Secret.from_name("huggingface-secret")], +) +def download_model(): + run_config_hook("download_model", (hf_cache_volume,)) + + +@app.function( + image=image, + volumes={str(DATA_PATH): data_volume}, + timeout=4 * 60 * 60, + secrets=[modal.Secret.from_name("huggingface-secret")], +) +def download_data(): + run_config_hook("download_data", (data_volume,)) + + +@app.function( + image=image, + gpu=f"{modal_cfg.gpu}:{miles_cfg.actor_num_gpus_per_node}", + memory=modal_cfg.memory if modal_cfg.memory else None, + cloud=modal_cfg.cloud if modal_cfg.cloud else None, + region=modal_cfg.region if modal_cfg.region else None, + volumes=modal_volumes, + secrets=[modal.Secret.from_name("wandb-secret")], + timeout=24 * 60 * 60, +) +async def train(): + await asyncio.gather( + hf_cache_volume.reload.aio(), + data_volume.reload.aio(), + checkpoints_volume.reload.aio(), + ) + cfg = get_module(EXPERIMENT).miles + my_ip = "127.0.0.1" + os.environ["MILES_HOST_IP"] = my_ip + os.environ["SGLANG_HOST_IP"] = my_ip + os.environ["HOST_IP"] = my_ip + + start_ray_head(my_ip, 1) + prepare_miles_config(cfg, tempfile.mkdtemp()) + + cmd = build_train_cmd(cfg, MILES_ROOT) + runtime_env = { + "env_vars": { + "no_proxy": f"127.0.0.1,{my_ip}", + "MASTER_ADDR": my_ip, + **cfg.environment, + } + } + + client = JobSubmissionClient("http://127.0.0.1:8265") + job_id = client.submit_job(entrypoint=cmd, runtime_env=runtime_env) + print(f"Job submitted: {job_id}") + print(f"Training {EXPERIMENT} on 1 node x {modal_cfg.gpu}:{cfg.actor_num_gpus_per_node}") + print(f"Command: {cmd}") + + async with modal.forward(RAY_DASHBOARD_PORT) as tunnel: + print(f"Ray dashboard: {tunnel.url}") + async for line in client.tail_job_logs(job_id): + print(line, end="", flush=True) From cab0aa5bb8bba21cf1eaa9ddf7717b19900221e4 Mon Sep 17 00:00:00 2001 From: kailash Date: Fri, 10 Jul 2026 15:50:32 +0000 Subject: [PATCH 3/6] revert modal_train --- miles/modal_train.py | 41 +++++++++-------------------------------- 1 file changed, 9 insertions(+), 32 deletions(-) diff --git a/miles/modal_train.py b/miles/modal_train.py index 59e72d1..4ab3296 100644 --- a/miles/modal_train.py +++ b/miles/modal_train.py @@ -78,39 +78,13 @@ def run_config_hook(experiment: str, hook_name: str, mounted_volumes) -> None: - """Reload mounted volumes, run a MilesConfig hook, then commit them. - - Commits the mounted volumes every 2 min *while* the hook runs, not just at the - end, so a long HF pull persists completed shards as it goes. If the container dies - mid-download (network stall, disconnect), a re-run reloads the last commit and - snapshot_download skips the already-present files — resuming near where it stopped - instead of restarting from zero. - """ - import threading - + """Reload mounted volumes, run a MilesConfig hook, then commit them.""" miles_cfg = get_module(experiment).miles for volume in mounted_volumes: volume.reload() - - stop = threading.Event() - - def _periodic_commit() -> None: - while not stop.wait(120): - for volume in mounted_volumes: - try: - volume.commit() - except Exception as e: # best-effort; never kill the download - print(f"[modal] periodic volume commit failed: {e}", flush=True) - - committer = threading.Thread(target=_periodic_commit, daemon=True) - committer.start() - try: - getattr(miles_cfg, hook_name)() - finally: - stop.set() - committer.join(timeout=10) - for volume in mounted_volumes: - volume.commit() + getattr(miles_cfg, hook_name)() + for volume in mounted_volumes: + volume.commit() @app.local_entrypoint() @@ -316,8 +290,11 @@ async def train(experiment: str = os.environ.get("EXPERIMENT_CONFIG", "")): start_ray_head(my_ip, n_nodes) prepare_miles_config(miles_cfg, tempfile.mkdtemp()) - # W&B reads WANDB_API_KEY from the inherited Modal secret. Do not copy it - # into miles_cfg: CLI args are retained verbatim in Ray and Modal logs. + if (wandb_key := os.environ.get("WANDB_API_KEY", "")) and getattr( + miles_cfg, "use_wandb", False + ): + miles_cfg.wandb_key = wandb_key + cmd = build_train_cmd(miles_cfg, MILES_ROOT) runtime_env = { "env_vars": { From 916b17c2cedbef6abf1e73fc4b71c8dbc9a2b166 Mon Sep 17 00:00:00 2001 From: kailash Date: Mon, 13 Jul 2026 20:37:31 +0000 Subject: [PATCH 4/6] fp 8 rollout config --- ...44b_a40b_lora_dapo_tilelang_32k_fp8roll.py | 48 +++++++++++++++++++ 1 file changed, 48 insertions(+) create mode 100644 miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k_fp8roll.py diff --git a/miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k_fp8roll.py b/miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k_fp8roll.py new file mode 100644 index 0000000..1d78743 --- /dev/null +++ b/miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k_fp8roll.py @@ -0,0 +1,48 @@ +"""GLM-5.2 DAPO 32k padded — FP8 rollout A/B test, 8 nodes x 8 H200. + +Identical to glm5_2_744b_a40b_lora_dapo_tilelang_32k_padstress except SGLang +serves the frozen base FP8-quantized (--quantization fp8, dynamic per-token +activation quant). Training stays bf16. Goals: + + 1. rollout/prefill+decode speedup vs the bf16 baseline (~9 min/rollout) + 2. train_rollout_logprob_abs_diff — bf16 baseline is ~0.01; this measures + the train/rollout policy mismatch introduced by FP8 serving + +Short run (3 rollouts = 6 train steps), no checkpoints. + + EXPERIMENT_CONFIG=glm5_2_744b_a40b_lora_dapo_tilelang_32k_fp8roll \ + uv run modal run --detach miles/modal_train.py::train +""" + +from configs import glm5_2_744b_a40b_lora_dapo_tilelang_32k_padstress as _base + +from configs.base import ModalConfig + +modal = ModalConfig( + docker_image=_base.modal.docker_image, + gpu=_base.modal.gpu, + memory=_base.modal.memory, + cloud=_base.modal.cloud, + region=_base.modal.region, + patch_files=[*_base.modal.patch_files, "miles/sglang_fp8_lora_fix.py"], + image_run_commands=[ + *_base.modal.image_run_commands, + # LoRA-B buffer sizing fix for quantized column-parallel layers + # (validated on the 5-layer FP8 diag; see sglang_fp8_lora_fix.py). + "python /tmp/sglang_fp8_lora_fix.py", + ], + image_env=dict(_base.modal.image_env), +) + + +class _Miles(_base._Miles): + sglang_quantization = "fp8" + + num_rollout = 3 + save = None + save_interval = None + + wandb_group = "glm5.2-744B-8node-tilelang-dapo-32k-fp8roll" + + +miles = _Miles() From 56a6b3dd65e6bfd078b3f74c6652761220bf9696 Mon Sep 17 00:00:00 2001 From: kailash Date: Wed, 15 Jul 2026 16:24:38 +0000 Subject: [PATCH 5/6] fp8 rollouts working on glm5.2 lora, misc bugfixing on sglang --- ...glm5_2_744b_a40b_lora_dapo_tilelang_32k.py | 92 ++++++++++ ..._2_744b_a40b_lora_dapo_tilelang_32k_fp8.py | 101 +++++++++++ ...44b_a40b_lora_dapo_tilelang_32k_fp8roll.py | 48 ----- miles/glm5_tilelang_safe_indices.patch | 157 +++++++++++++++++ miles/modal_sglang_serve_test.py | 164 ++++++++++++++++++ miles/modal_train_glm_test.py | 30 ++-- miles/sglang_fp8_lora_fix.py | 92 ++++++++++ miles/sglang_tp1_shared_expert_fix.py | 72 ++++++++ 8 files changed, 698 insertions(+), 58 deletions(-) create mode 100644 miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k.py create mode 100644 miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k_fp8.py delete mode 100644 miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k_fp8roll.py create mode 100644 miles/glm5_tilelang_safe_indices.patch create mode 100644 miles/modal_sglang_serve_test.py create mode 100644 miles/sglang_fp8_lora_fix.py create mode 100644 miles/sglang_tp1_shared_expert_fix.py diff --git a/miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k.py b/miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k.py new file mode 100644 index 0000000..b039559 --- /dev/null +++ b/miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k.py @@ -0,0 +1,92 @@ +"""GLM-5.2 full-model TileLang/THD on dapo-math — 32k context + 1k rollout, 8 nodes x 8 H200. + +Uses the validated 64-GPU TileLang setup (safe-indices kernel patch, full +activation recompute, extended SGLang load timeout) with the dapo-math task. +No context parallelism (CP=1); 32k context fits without it (validated 2026-07-12). + + EXPERIMENT_CONFIG=glm5_2_744b_a40b_lora_dapo_tilelang_32k \ + uv run modal run miles/modal_train.py::download_data + EXPERIMENT_CONFIG=glm5_2_744b_a40b_lora_dapo_tilelang_32k \ + uv run modal run --detach miles/modal_train.py::train +""" + +from configs import glm5_2_744b_a40b_lora as _base +from configs.base import CHECKPOINTS_PATH, DATA_PATH, ModalConfig + +_PATCH = "miles/glm5_tilelang_safe_indices.patch" + +modal = ModalConfig( + docker_image=_base.modal.docker_image, + gpu=_base.modal.gpu, + memory=_base.modal.memory, + cloud=_base.modal.cloud, + region=_base.modal.region, + patch_files=[_PATCH], + image_run_commands=[ + *_base.modal.image_run_commands, + ( + "cd /usr/local/lib/python3.12/dist-packages && " + "git apply --check /tmp/glm5_tilelang_safe_indices.patch && " + "git apply /tmp/glm5_tilelang_safe_indices.patch" + ), + ( + "python -c \"from pathlib import Path; " + "p = Path('/sgl-workspace/sglang/python/sglang/srt/model_executor/model_runner.py'); " + "s = p.read_text(); " + "old = 'UNBALANCED_MODEL_LOADING_TIMEOUT_S = 480'; " + "assert s.count(old) == 1; " + "p.write_text(s.replace(old, 'UNBALANCED_MODEL_LOADING_TIMEOUT_S = 1800'))\"" + ), + # LoRA checkpoint fix: dp_rank_0's mkdir on the shared volume is not + # visible to containers on other nodes, so every rank must mkdir itself + # before writing its training_state_rank{N}.pt. + ( + "python -c \"from pathlib import Path; " + "p = Path('/root/miles/miles/backends/megatron_utils/lora_utils.py'); " + "s = p.read_text(); " + "old = 'if is_dp_rank_0:\\n save_path.mkdir(parents=True, exist_ok=True)'; " + "new = 'save_path.mkdir(parents=True, exist_ok=True)'; " + "assert s.count(old) == 1; " + "p.write_text(s.replace(old, new))\"" + ), + ], + image_env=dict(_base.modal.image_env), +) + + +class _Miles(_base._Miles): + dsa_attention_backend = "tilelang" + qkv_format = "thd" + data_pad_size_multiplier = None + recompute_granularity = "full" + recompute_method = "uniform" + recompute_num_layers = 1 + + prompt_data = f"{DATA_PATH}/dapo-math-17k/dapo-math-17k.jsonl" + input_key = "prompt" + label_key = "label" + + seq_length = 32768 + rollout_max_context_len = 32768 + rollout_max_response_len = 1024 + + num_rollout = 50 + save = f"{CHECKPOINTS_PATH}/GLM-5.2-lora-dapo-32k-ckpt" + save_interval = 10 + + wandb_group = "glm5.2-744B-8node-tilelang-dapo-32k" + + def download_data(self) -> None: + import os + + from huggingface_hub import snapshot_download + + os.makedirs(f"{DATA_PATH}/dapo-math-17k", exist_ok=True) + snapshot_download( + repo_id="zhuzilin/dapo-math-17k", + repo_type="dataset", + local_dir=f"{DATA_PATH}/dapo-math-17k", + ) + + +miles = _Miles() diff --git a/miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k_fp8.py b/miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k_fp8.py new file mode 100644 index 0000000..5ec0f86 --- /dev/null +++ b/miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k_fp8.py @@ -0,0 +1,101 @@ +"""GLM-5.2 LoRA DAPO training with FP8 rollout — 8 nodes x 8 H200, TileLang/THD, 32k context. + +The finalized FP8-rollout setup: the rollout engines serve the official +block-quantized ``zai-org/GLM-5.2-FP8`` checkpoint while the trainer keeps +bf16 ``zai-org/GLM-5.2``. LoRA adapters (bf16) sync to the engines every +rollout. Measured vs the bf16/bf16 baseline: generation ~25-40% faster +(351-450s vs ~540s per rollout), train_rollout_logprob_abs_diff ~0.040 +(bf16 ~0.010, online-quant fp8 ~0.057), no sampling NaNs, rewards learn. + +Fixes this configuration depends on (see each patch's docstring): + - glm5_tilelang_safe_indices.patch: TileLang sparse-MLA backward NaNs + (unsafe padded-index access + aggressive shared-memory merge miscompile). + - sglang_fp8_lora_fix.py: LoRA-B buffer mis-sizing on quantized + column-parallel layers crashed engine init under --quantization fp8 / + quantized checkpoints. + - sglang_tp1_shared_expert_fix.py: SGLANG_SHARED_EXPERT_TP1-replicated + shared expert was double-added (once per TP rank) whenever the post-MoE + all-reduce is deferred/replaced (FlashInfer AllReduce Fusion, + dp-attention reduce-scatterv); folded into the pre-reduction add. + - SGLANG_SHARED_EXPERT_TP1=1: the 128x128 block-quantized checkpoint cannot + TP32-shard the shared expert (2048/32 = 64-row shards < one scale block), + so replicate it instead (~25 MB/rank). + + EXPERIMENT_CONFIG=glm5_2_744b_a40b_lora_dapo_tilelang_32k_fp8 \ + uv run modal run miles/modal_train.py::download_model + EXPERIMENT_CONFIG=glm5_2_744b_a40b_lora_dapo_tilelang_32k_fp8 \ + uv run modal run miles/modal_train.py::download_data + EXPERIMENT_CONFIG=glm5_2_744b_a40b_lora_dapo_tilelang_32k_fp8 \ + uv run modal run --detach miles/modal_train.py::train +""" + +from configs import glm5_2_744b_a40b_lora_dapo_tilelang_32k as _base +from configs.base import CHECKPOINTS_PATH, ModalConfig + +_FP8_CHECKPOINT = "zai-org/GLM-5.2-FP8" + +modal = ModalConfig( + docker_image=_base.modal.docker_image, + gpu=_base.modal.gpu, + memory=_base.modal.memory, + cloud=_base.modal.cloud, + region=_base.modal.region, + patch_files=[ + *_base.modal.patch_files, + "miles/sglang_fp8_lora_fix.py", + "miles/sglang_tp1_shared_expert_fix.py", + ], + image_run_commands=[ + *_base.modal.image_run_commands, + "python /tmp/sglang_fp8_lora_fix.py", + "python /tmp/sglang_tp1_shared_expert_fix.py", + ], + image_env=dict(_base.modal.image_env), +) + + +class _Miles(_base._Miles): + environment = { + **_base._Miles.environment, + "SGLANG_SHARED_EXPERT_TP1": "1", + } + + # The pre-quantized checkpoint carries its own quantization_config; + # do NOT also force online quantization. + sglang_quantization = None + + # Rollout base differs from hf_checkpoint, so adapter sync must be + # opted into explicitly. + sglang_config = { + "sglang": [ + { + "name": "actor", + "model_path": _FP8_CHECKPOINT, + "update_weights": True, + "num_gpus_per_engine": 32, + "server_groups": [ + {"worker_type": "regular", "num_gpus": 64}, + ], + } + ] + } + + # 4096-token response budget so completions can finish and reward can + # improve (1024 truncated heavily on dapo-math). + rollout_max_response_len = 4096 + + # 2 gradient steps per rollout (128 samples / global_batch_size 64). + num_rollout = 50 + save = f"{CHECKPOINTS_PATH}/GLM-5.2-lora-dapo-32k-fp8-ckpt" + save_interval = 10 + + wandb_group = "glm5.2-744B-8node-tilelang-dapo-32k-fp8" + + def download_model(self) -> None: + from huggingface_hub import snapshot_download + + snapshot_download(self.hf_checkpoint, max_workers=32) + snapshot_download(_FP8_CHECKPOINT, max_workers=32) + + +miles = _Miles() diff --git a/miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k_fp8roll.py b/miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k_fp8roll.py deleted file mode 100644 index 1d78743..0000000 --- a/miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k_fp8roll.py +++ /dev/null @@ -1,48 +0,0 @@ -"""GLM-5.2 DAPO 32k padded — FP8 rollout A/B test, 8 nodes x 8 H200. - -Identical to glm5_2_744b_a40b_lora_dapo_tilelang_32k_padstress except SGLang -serves the frozen base FP8-quantized (--quantization fp8, dynamic per-token -activation quant). Training stays bf16. Goals: - - 1. rollout/prefill+decode speedup vs the bf16 baseline (~9 min/rollout) - 2. train_rollout_logprob_abs_diff — bf16 baseline is ~0.01; this measures - the train/rollout policy mismatch introduced by FP8 serving - -Short run (3 rollouts = 6 train steps), no checkpoints. - - EXPERIMENT_CONFIG=glm5_2_744b_a40b_lora_dapo_tilelang_32k_fp8roll \ - uv run modal run --detach miles/modal_train.py::train -""" - -from configs import glm5_2_744b_a40b_lora_dapo_tilelang_32k_padstress as _base - -from configs.base import ModalConfig - -modal = ModalConfig( - docker_image=_base.modal.docker_image, - gpu=_base.modal.gpu, - memory=_base.modal.memory, - cloud=_base.modal.cloud, - region=_base.modal.region, - patch_files=[*_base.modal.patch_files, "miles/sglang_fp8_lora_fix.py"], - image_run_commands=[ - *_base.modal.image_run_commands, - # LoRA-B buffer sizing fix for quantized column-parallel layers - # (validated on the 5-layer FP8 diag; see sglang_fp8_lora_fix.py). - "python /tmp/sglang_fp8_lora_fix.py", - ], - image_env=dict(_base.modal.image_env), -) - - -class _Miles(_base._Miles): - sglang_quantization = "fp8" - - num_rollout = 3 - save = None - save_interval = None - - wandb_group = "glm5.2-744B-8node-tilelang-dapo-32k-fp8roll" - - -miles = _Miles() diff --git a/miles/glm5_tilelang_safe_indices.patch b/miles/glm5_tilelang_safe_indices.patch new file mode 100644 index 0000000..4337905 --- /dev/null +++ b/miles/glm5_tilelang_safe_indices.patch @@ -0,0 +1,157 @@ +diff --git a/megatron/bridge/models/glm5/tilelang/sparse_mla.py b/megatron/bridge/models/glm5/tilelang/sparse_mla.py +--- a/megatron/bridge/models/glm5/tilelang/sparse_mla.py ++++ b/megatron/bridge/models/glm5/tilelang/sparse_mla.py +@@ -59,6 +59,14 @@ class SparseMLA(torch.autograd.Function): + scaling = ctx.scaling + + tl_dq, tl_dkv = sparse_mla_bwd(q, kv, tl_out, grad_output.contiguous(), indices, tl_lse, sm_scale=scaling) ++ for name, grad in (("dQ", tl_dq), ("dKV", tl_dkv)): ++ finite = torch.isfinite(grad) ++ if not finite.all(): ++ nonfinite = (~finite).sum().item() ++ raise FloatingPointError( ++ f"TileLang SparseMLA {name} contains {nonfinite} non-finite values " ++ f"(shape={tuple(grad.shape)}, dtype={grad.dtype})" ++ ) + + # Return gradients for each input (None for indices as it's not differentiable) + return tl_dq, tl_dkv, None, None +diff --git a/megatron/bridge/models/glm5/tilelang/tilelang_sparse_mla_fwd.py b/megatron/bridge/models/glm5/tilelang/tilelang_sparse_mla_fwd.py +--- a/megatron/bridge/models/glm5/tilelang/tilelang_sparse_mla_fwd.py ++++ b/megatron/bridge/models/glm5/tilelang/tilelang_sparse_mla_fwd.py +@@ -109,6 +109,7 @@ def sparse_mla_fwd( + O_shared = T.alloc_shared([H_per_block, D], dtype) + Lse_shared = T.alloc_shared([H_per_block], accum_dtype) + mask = T.alloc_fragment([BI], "bool") ++ safe_indices = T.alloc_shared([BI], indices_dtype) + + acc_o = T.alloc_fragment([H_per_block, D], accum_dtype) + acc_s = T.alloc_fragment([H_per_block, BI], accum_dtype) +@@ -136,13 +137,25 @@ def sparse_mla_fwd( + + for i_i in T.Pipelined(NI, num_stages=num_stages): + for bi_i in T.Parallel(BI): +- # Changed here for thd +- mask[bi_i] = Indices[b_i, s_i, g_i, i_i * BI + bi_i] != -1 ++ mask[bi_i] = ( ++ Indices[b_i, s_i, g_i, i_i * BI + bi_i] >= 0 ++ and Indices[b_i, s_i, g_i, i_i * BI + bi_i] < seq_len_kv ++ ) ++ safe_indices[bi_i] = T.if_then_else( ++ mask[bi_i], ++ Indices[b_i, s_i, g_i, i_i * BI + bi_i], ++ 0, ++ ) ++ T.sync_threads() + + for bi_i, d_i in T.Parallel(BI, D): +- KV_shared[bi_i, d_i] = KV[b_i, Indices[b_i, s_i, g_i, i_i * BI + bi_i], g_i, d_i] ++ KV_shared[bi_i, d_i] = T.if_then_else( ++ mask[bi_i], KV[b_i, safe_indices[bi_i], g_i, d_i], 0 ++ ) + for bi_i, d_i in T.Parallel(BI, D_tail): +- K_tail_shared[bi_i, d_i] = KV[b_i, Indices[b_i, s_i, g_i, i_i * BI + bi_i], g_i, D + d_i] ++ K_tail_shared[bi_i, d_i] = T.if_then_else( ++ mask[bi_i], KV[b_i, safe_indices[bi_i], g_i, D + d_i], 0 ++ ) + + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.if_then_else(mask[bi_i], 0, -T.infinity(acc_s.dtype)) +diff --git a/megatron/bridge/models/glm5/tilelang/tilelang_sparse_mla_bwd.py b/megatron/bridge/models/glm5/tilelang/tilelang_sparse_mla_bwd.py +--- a/megatron/bridge/models/glm5/tilelang/tilelang_sparse_mla_bwd.py ++++ b/megatron/bridge/models/glm5/tilelang/tilelang_sparse_mla_bwd.py +@@ -98,7 +98,7 @@ def bwd( + pass_configs={ + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, +- tilelang.PassConfigKey.TL_ENABLE_AGGRESSIVE_SHARED_MEMORY_MERGE: True, ++ tilelang.PassConfigKey.TL_ENABLE_AGGRESSIVE_SHARED_MEMORY_MERGE: False, + }, + ) + def bwd( +@@ -168,6 +168,7 @@ def bwd( + KV_tail_shared = T.alloc_shared([BS, D_tail], dtype) + dO_shared = T.alloc_shared([block_H, D], dtype) + mask = T.alloc_fragment([BS], "bool") ++ safe_indices = T.alloc_shared([BS], indices_dtype) + + P_shared_cast = T.alloc_shared([block_H, BS], dtype) + dP_shared_cast = T.alloc_shared([block_H, BS], dtype) +@@ -195,21 +196,33 @@ def bwd( + # Process each block of indices + for i_i in T.Pipelined(NS, num_stages=num_stages): + # Check which indices are valid + for bi_i in T.Parallel(BS): +- # Changed here for thd +- mask[bi_i] = Indices[by, s_i, bz // NH, i_i * BS + bi_i] != -1 ++ mask[bi_i] = ( ++ Indices[by, s_i, bz // NH, i_i * BS + bi_i] >= 0 ++ and Indices[by, s_i, bz // NH, i_i * BS + bi_i] < S_kv ++ ) ++ safe_indices[bi_i] = T.if_then_else( ++ mask[bi_i], ++ Indices[by, s_i, bz // NH, i_i * BS + bi_i], ++ 0, ++ ) ++ T.sync_threads() + + # Compute attention scores + for h_i, bi_i in T.Parallel(block_H, BS): + acc_p[h_i, bi_i] = T.if_then_else(mask[bi_i], 0, -T.infinity(acc_p.dtype)) + + # Load KV, V for this block of indices + for bi_i, d_i in T.Parallel(BS, D): +- KV_shared[bi_i, d_i] = KV[by, Indices[by, s_i, bz // NH, i_i * BS + bi_i], bz // NH, d_i] ++ KV_shared[bi_i, d_i] = T.if_then_else( ++ mask[bi_i], KV[by, safe_indices[bi_i], bz // NH, d_i], 0 ++ ) + + T.gemm(Q_shared, KV_shared, acc_p, transpose_B=True, policy=T.GemmWarpPolicy.FullCol) + + for bi_i, d_i in T.Parallel(BS, D_tail): +- KV_tail_shared[bi_i, d_i] = KV[by, Indices[by, s_i, bz // NH, i_i * BS + bi_i], bz // NH, D + d_i] ++ KV_tail_shared[bi_i, d_i] = T.if_then_else( ++ mask[bi_i], KV[by, safe_indices[bi_i], bz // NH, D + d_i], 0 ++ ) + T.gemm(Q_tail_shared, KV_tail_shared, acc_p, transpose_B=True, policy=T.GemmWarpPolicy.FullCol) + +@@ -256,27 +269,21 @@ def bwd( + acc_dkv_tail_shared[bi_i, d_i] = acc_dkv_tail[bi_i + s * (BS // split_store), d_i] + + for bi_i, d_i in T.Parallel(BS // split_store, D // 4): +- T.atomic_addx4( +- dKV[ +- by, +- Indices[by, s_i, bz // NH, i_i * BS + bi_i + s * (BS // split_store)], +- bz // NH, +- d_i * 4, +- ], +- acc_dkv_shared[bi_i, d_i * 4], +- ) ++ index_offset = bi_i + s * (BS // split_store) ++ if mask[index_offset]: ++ T.atomic_addx4( ++ dKV[by, safe_indices[index_offset], bz // NH, d_i * 4], ++ acc_dkv_shared[bi_i, d_i * 4], ++ ) + + # Atomically update dKV, dKV_tail tensors + for bi_i, d_i in T.Parallel(BS // split_store, D_tail // 4): +- T.atomic_addx4( +- dKV[ +- by, +- Indices[by, s_i, bz // NH, i_i * BS + bi_i + s * (BS // split_store)], +- bz // NH, +- D + d_i * 4, +- ], +- acc_dkv_tail_shared[bi_i, d_i * 4], +- ) ++ index_offset = bi_i + s * (BS // split_store) ++ if mask[index_offset]: ++ T.atomic_addx4( ++ dKV[by, safe_indices[index_offset], bz // NH, D + d_i * 4], ++ acc_dkv_tail_shared[bi_i, d_i * 4], ++ ) + + # Store the accumulated dQ + T.copy(acc_dq, dQ_shared) diff --git a/miles/modal_sglang_serve_test.py b/miles/modal_sglang_serve_test.py new file mode 100644 index 0000000..4a8a723 --- /dev/null +++ b/miles/modal_sglang_serve_test.py @@ -0,0 +1,164 @@ +"""Standalone 1-node SGLang serving test for GLM-5.2 FP8 debugging. + +Serves the model on 8xH200 with the same SGLang flags as the miles rollout +engines (minus dp-attention, which needs 32 GPUs), then runs greedy +completions with logprobs to judge output quality directly. Bisect levers +are exposed via --server-args / --env-json so variants (shared-expert TP1, +moe runner, nsa backends, online vs offline quant) don't need code changes. + + uv run modal run --detach miles/modal_sglang_serve_test.py::serve_test + uv run modal run --detach miles/modal_sglang_serve_test.py::serve_test \ + --env-json '{"SGLANG_SHARED_EXPERT_TP1": "1"}' +""" + +import json +import os +import subprocess +import time +import urllib.request + +import modal + +from configs import get_module +from configs.base import HF_CACHE_PATH + +# Reuse the FP8 experiment's image (same sglang build + patches as the +# 64-GPU rollout engines, incl. sglang_fp8_lora_fix and the load-timeout bump). +EXPERIMENT = "glm5_2_744b_a40b_lora_dapo_tilelang_32k_fp8" +modal_cfg = get_module(EXPERIMENT).modal + +image = ( + modal.Image.from_registry(modal_cfg.docker_image) + .entrypoint([]) + .add_local_python_source("configs", copy=True) + .add_local_python_source("modal_helpers", copy=True) +) +for patch in modal_cfg.patch_files: + image = image.add_local_file(patch, f"/tmp/{os.path.basename(patch)}", copy=True) +if modal_cfg.image_run_commands: + image = image.run_commands(*modal_cfg.image_run_commands) +if modal_cfg.image_env: + image = image.env(modal_cfg.image_env) + +hf_cache_volume = modal.Volume.from_name("huggingface-cache", create_if_missing=True) + +app = modal.App("glm52-fp8-serve-test") + +PORT = 30000 + +DEFAULT_SERVER_ARGS = ( + "--model-path zai-org/GLM-5.2-FP8 " + "--tp-size 8 " + "--trust-remote-code " + "--attention-backend nsa " + "--nsa-decode-backend flashmla_sparse " + "--nsa-prefill-backend flashmla_sparse " + "--moe-runner-backend triton " + "--disable-shared-experts-fusion " + "--mem-fraction-static 0.80 " + "--context-length 8192 " + # Quality test only: skip graph capture (15 min) and its NVLS multicast + # setup, which Fabric Manager on these hosts cannot provide. + "--disable-cuda-graph " + f"--port {PORT} --host 127.0.0.1" +) + +# Same rollout-engine env as the miles config. +DEFAULT_ENV = { + "SGLANG_NSA_FORCE_MLA": "1", + "INDEXER_ROPE_NEOX_STYLE": "0", + "NCCL_NVLS_ENABLE": "0", +} + +PROMPTS = [ + "The capital of France is", + ( + "Question: Natalia sold clips to 48 of her friends in April, and then " + "she sold half as many clips in May. How many clips did Natalia sell " + "altogether in April and May?\nAnswer:" + ), + "def fibonacci(n):\n", + "1 + 1 = 2, 2 + 2 = 4, 4 + 4 =", +] + + +def _generate(prompt: str, max_new_tokens: int = 96) -> dict: + payload = { + "text": prompt, + "sampling_params": {"temperature": 0, "max_new_tokens": max_new_tokens}, + "return_logprob": True, + } + req = urllib.request.Request( + f"http://127.0.0.1:{PORT}/generate", + data=json.dumps(payload).encode(), + headers={"Content-Type": "application/json"}, + ) + return json.loads(urllib.request.urlopen(req, timeout=900).read()) + + +def _post(path: str, payload: dict | None = None) -> str: + req = urllib.request.Request( + f"http://127.0.0.1:{PORT}/{path}", + data=json.dumps(payload or {}).encode(), + headers={"Content-Type": "application/json"}, + ) + return urllib.request.urlopen(req, timeout=900).read().decode() + + +def _run_prompts(phase: str) -> None: + for prompt in PROMPTS: + out = _generate(prompt) + lps = [t[0] for t in out["meta_info"]["output_token_logprobs"]] + mean_lp = sum(lps) / max(len(lps), 1) + print("=" * 60) + print(f"[{phase}] PROMPT: {prompt[:120]!r}") + print(f"[{phase}] OUTPUT: {out['text'][:400]!r}") + print(f"[{phase}] greedy mean output logprob: {mean_lp:.3f} over {len(lps)} tokens") + + +@app.function( + image=image, + gpu="H200:8", + volumes={str(HF_CACHE_PATH): hf_cache_volume}, + timeout=3 * 60 * 60, + secrets=[modal.Secret.from_name("huggingface-secret")], +) +def serve_test(server_args: str = "", env_json: str = "{}", cycle_memory: bool = False): + hf_cache_volume.reload() + args = server_args or DEFAULT_SERVER_ARGS + if cycle_memory: + # Mirror the miles rollout engines' memory-saver setup so we can + # exercise the release/resume weight backup path the training loop + # performs before rollout 0. + args += " --enable-memory-saver --enable-weights-cpu-backup" + env = {**os.environ, **DEFAULT_ENV, **json.loads(env_json)} + print(f"server args: {args}") + print(f"env overrides: {json.loads(env_json)}") + + proc = subprocess.Popen(f"python -m sglang.launch_server {args}", shell=True, env=env) + try: + deadline = time.time() + 45 * 60 + while True: + if proc.poll() is not None: + raise RuntimeError(f"server exited during startup: {proc.returncode}") + try: + urllib.request.urlopen(f"http://127.0.0.1:{PORT}/health_generate", timeout=5) + break + except Exception: + if time.time() > deadline: + raise TimeoutError("server did not become healthy in 45 min") + time.sleep(10) + print("server healthy, running prompts") + + _run_prompts("fresh") + + if cycle_memory: + print("cycling memory occupation (release -> resume)...") + print(_post("release_memory_occupation")[:200]) + time.sleep(10) + print(_post("resume_memory_occupation")[:200]) + _run_prompts("after-cycle") + + print("DONE_SERVE_TEST") + finally: + proc.terminate() diff --git a/miles/modal_train_glm_test.py b/miles/modal_train_glm_test.py index eb74b2e..2025047 100644 --- a/miles/modal_train_glm_test.py +++ b/miles/modal_train_glm_test.py @@ -17,7 +17,12 @@ from configs import get_module from configs.base import HF_CACHE_PATH, DATA_PATH, CHECKPOINTS_PATH -EXPERIMENT = "glm5_2_744b_a40b_lora_5layer" +_ALLOWED_EXPERIMENTS = { + "glm5_2_744b_a40b_lora_5layer", +} +EXPERIMENT = os.environ.get("EXPERIMENT_CONFIG", "glm5_2_744b_a40b_lora_5layer") +if EXPERIMENT not in _ALLOWED_EXPERIMENTS: + raise ValueError(f"This launcher only supports 5-layer GLM diagnostics; got {EXPERIMENT!r}") exp_mod = get_module(EXPERIMENT) modal_cfg = exp_mod.modal @@ -31,6 +36,10 @@ .add_local_python_source("configs", copy=True) .add_local_python_source("modal_helpers", copy=True) ) +for patch in modal_cfg.patch_files: + image = image.add_local_file( + patch, f"/tmp/{os.path.basename(patch)}", copy=True + ) if modal_cfg.image_run_commands: image = image.run_commands(*modal_cfg.image_run_commands) if modal_cfg.image_env: @@ -59,8 +68,8 @@ RAY_DASHBOARD_PORT = 8265 -def run_config_hook(hook_name: str, mounted_volumes) -> None: - cfg = get_module(EXPERIMENT).miles +def run_config_hook(experiment: str, hook_name: str, mounted_volumes) -> None: + cfg = get_module(experiment).miles for volume in mounted_volumes: volume.reload() getattr(cfg, hook_name)() @@ -74,8 +83,8 @@ def run_config_hook(hook_name: str, mounted_volumes) -> None: timeout=4 * 60 * 60, secrets=[modal.Secret.from_name("huggingface-secret")], ) -def download_model(): - run_config_hook("download_model", (hf_cache_volume,)) +def download_model(experiment: str = EXPERIMENT): + run_config_hook(experiment, "download_model", (hf_cache_volume,)) @app.function( @@ -84,8 +93,8 @@ def download_model(): timeout=4 * 60 * 60, secrets=[modal.Secret.from_name("huggingface-secret")], ) -def download_data(): - run_config_hook("download_data", (data_volume,)) +def download_data(experiment: str = EXPERIMENT): + run_config_hook(experiment, "download_data", (data_volume,)) @app.function( @@ -98,13 +107,14 @@ def download_data(): secrets=[modal.Secret.from_name("wandb-secret")], timeout=24 * 60 * 60, ) -async def train(): +async def train(experiment: str = EXPERIMENT): await asyncio.gather( hf_cache_volume.reload.aio(), data_volume.reload.aio(), checkpoints_volume.reload.aio(), ) - cfg = get_module(EXPERIMENT).miles + exp_mod = get_module(experiment) + cfg = exp_mod.miles my_ip = "127.0.0.1" os.environ["MILES_HOST_IP"] = my_ip os.environ["SGLANG_HOST_IP"] = my_ip @@ -125,7 +135,7 @@ async def train(): client = JobSubmissionClient("http://127.0.0.1:8265") job_id = client.submit_job(entrypoint=cmd, runtime_env=runtime_env) print(f"Job submitted: {job_id}") - print(f"Training {EXPERIMENT} on 1 node x {modal_cfg.gpu}:{cfg.actor_num_gpus_per_node}") + print(f"Training {experiment} on 1 node x {exp_mod.modal.gpu}:{cfg.actor_num_gpus_per_node}") print(f"Command: {cmd}") async with modal.forward(RAY_DASHBOARD_PORT) as tunnel: diff --git a/miles/sglang_fp8_lora_fix.py b/miles/sglang_fp8_lora_fix.py new file mode 100644 index 0000000..3e1c89f --- /dev/null +++ b/miles/sglang_fp8_lora_fix.py @@ -0,0 +1,92 @@ +"""Image-build patcher: fix LoRA-B buffer sizing on quantized column-parallel layers. + +Root cause of "LoRA B output dim ... does not match base partition prefix dim" +under --quantization fp8: + + * mem_pool.get_lora_b_shape derives the effective TP for non-MoE column + modules from _row_parallel_shard_tp, an INPUT-sharding probe + (input_size // input_size_per_partition). + * On bf16, UnquantizedLinearMethod never sets input_size_per_partition, so + the probe falls back to the global tp_size and the code path below then + corrects the output dim via the output-side probe. Works. + * Fp8LinearMethod.create_weights DOES set layer.input_size_per_partition + (== input_size for column-parallel layers, whose input is unsharded), so + the probe returns 1, the whole sharding branch is skipped, and LoRA-B is + sized at the FULL output dim while the base layer stays TP-sharded. + set_lora_info then fails (e.g. shared_experts.gate_up_proj: B=4096 vs + per-rank partitions [1024, 1024] -> 2048). + +Fix: for non-MoE column-parallel modules, use the base module's +output_size_per_partition (probed by _column_parallel_out_partition) as the +authoritative per-rank LoRA-B output dim. It exists on both bf16 and +quantized layers and is exactly what set_lora_info validates against. +Applied at image build via `python /tmp/sglang_fp8_lora_fix.py`. +""" + +from pathlib import Path + +P = Path("/sgl-workspace/sglang/python/sglang/srt/lora/mem_pool.py") + +OLD = ''' if ( + effective_tp_size > 1 + and module_name not in ROW_PARALLELISM_LINEAR_LORA_NAMES + and module_name not in REPLICATED_LINEAR_LORA_NAMES + ): + # If the base column-parallel module is fully REPLICATED (its actual + # output_size_per_partition still equals the full output_dim -- e.g. the + # dense MLP gate_up under --moe-dense-tp-size 1), its output is NOT + # sharded, so keep LoRA-B at the full output dim. Dividing by the global + # tp_size here undersizes B and crashes set_lora_info ("LoRA B output dim + # != base partition prefix dim"). Non-MoE only; MoE shards by moe_tp_size. + probed_out = ( + None + if self.is_moe_module(module_name) + else self._column_parallel_out_partition( + module_name, base_model, layer_idx + ) + ) + if probed_out is not None and probed_out == output_dim: + pass # replicated base: keep full B output dim + else: + output_dim = self._column_parallel_lora_b_per_rank_dim( + module_name, output_dim, effective_tp_size + ) +''' + +NEW = ''' if ( + module_name not in ROW_PARALLELISM_LINEAR_LORA_NAMES + and module_name not in REPLICATED_LINEAR_LORA_NAMES + and not self.is_moe_module(module_name) + ): + # The base module's output_size_per_partition is the ground truth + # for LoRA-B's per-rank output dim (replicated OR TP-sharded), and + # it is quant-independent. Do NOT gate this on the input-sharding + # probe above: Fp8LinearMethod sets input_size_per_partition == + # input_size on column-parallel layers, which makes that probe + # return 1 and previously skipped sharding entirely, sizing LoRA-B + # at the full output dim against a TP-sharded base and crashing + # set_lora_info ("LoRA B output dim != base partition prefix dim"). + probed_out = self._column_parallel_out_partition( + module_name, base_model, layer_idx + ) + if probed_out is not None: + output_dim = probed_out + elif effective_tp_size > 1: + output_dim = self._column_parallel_lora_b_per_rank_dim( + module_name, output_dim, effective_tp_size + ) + elif ( + effective_tp_size > 1 + and module_name not in ROW_PARALLELISM_LINEAR_LORA_NAMES + and module_name not in REPLICATED_LINEAR_LORA_NAMES + ): + # MoE modules keep the moe_tp_size sharding path. + output_dim = self._column_parallel_lora_b_per_rank_dim( + module_name, output_dim, effective_tp_size + ) +''' + +src = P.read_text() +assert src.count(OLD) == 1, f"expected 1 match, got {src.count(OLD)}" +P.write_text(src.replace(OLD, NEW)) +print("patched", P) diff --git a/miles/sglang_tp1_shared_expert_fix.py b/miles/sglang_tp1_shared_expert_fix.py new file mode 100644 index 0000000..5e60e9b --- /dev/null +++ b/miles/sglang_tp1_shared_expert_fix.py @@ -0,0 +1,72 @@ +"""Fix SGLANG_SHARED_EXPERT_TP1 double-add in deepseek_v2.py. + +With a TP1-replicated shared expert, upstream adds its output to the MoE +result "after the all-reduce" so each TP rank contributes it only once. +But `should_skip_post_experts_all_reduce` can skip that explicit all-reduce +in favor of a *deferred/replaced* reduction (FlashInfer AllReduce Fusion, +dp-attention's reduce-scatterv), in which case the "post-all-reduce" add +actually lands *before* the real reduction and the replicated shared output +is summed once per TP rank (8-32x here) — corrupting every MoE layer. + +Fix: fold the shared output into the pre-reduction add scaled by +1/tp_size. Any linear reduction (all-reduce, fused all-reduce, +reduce-scatterv) then reconstitutes exactly one copy. tp_size is a power +of two, so the bf16 scaling is exact and there is no precision cost. + +Verified empirically on GLM-5.2-FP8: TP1 + AllReduce Fusion and +TP1 + dp-attention both produced garbage output before this patch. +""" + +from pathlib import Path + +TARGET = Path("/sgl-workspace/sglang/python/sglang/srt/models/deepseek_v2.py") + +SENTINEL = "MILES_TP1_SHARED_EXPERT_FIX" + +# Pre-reduction arg: pass the scaled shared output instead of None when TP1. +OLD_FUSE_ARG = "None if self._shared_expert_tp1 else shared_output," +NEW_FUSE_ARG = ( + "(shared_output * (1.0 / self.tp_size) if shared_output is not None else None) " + "if self._shared_expert_tp1 else shared_output, # " + SENTINEL +) + +# Post-reduction adds: now double-counting (the scaled copy is already in), +# so neuter them. +OLD_POST_ADD_DUAL = ( + " if self._shared_expert_tp1:\n" + " final_hidden_states += shared_output\n" +) +OLD_POST_ADD_NORMAL = ( + " if shared_output is not None and self._shared_expert_tp1:\n" + " final_hidden_states += shared_output\n" +) +NEW_POST_ADD = ( + " if False: # " + SENTINEL + ": folded into pre-reduction add\n" + " final_hidden_states += shared_output\n" +) + + +def main() -> None: + src = TARGET.read_text() + if SENTINEL in src: + print("already patched, skipping") + return + + n_fuse = src.count(OLD_FUSE_ARG) + assert n_fuse == 2, f"expected 2 maybe_fuse TP1 args, found {n_fuse}" + src = src.replace(OLD_FUSE_ARG, NEW_FUSE_ARG) + + n_dual = src.count(OLD_POST_ADD_DUAL) + assert n_dual == 1, f"expected 1 dual-stream post-add, found {n_dual}" + src = src.replace(OLD_POST_ADD_DUAL, NEW_POST_ADD) + + n_normal = src.count(OLD_POST_ADD_NORMAL) + assert n_normal == 1, f"expected 1 forward_normal post-add, found {n_normal}" + src = src.replace(OLD_POST_ADD_NORMAL, NEW_POST_ADD) + + TARGET.write_text(src) + print(f"patched {TARGET}: 2 pre-reduction args, 2 post-adds neutered") + + +if __name__ == "__main__": + main() From f92ef7dadbd1243d9b9d67eb8cc585ee9717bdda Mon Sep 17 00:00:00 2001 From: kailash Date: Fri, 17 Jul 2026 17:04:10 +0000 Subject: [PATCH 6/6] 16 node wip config --- ...744b_a40b_lora_dapo_tilelang_32k_16node.py | 113 ++++++++++++++++++ ...b_a40b_lora_dapo_tilelang_32k_8node_cp4.py | 29 +++++ miles/megatron_dsa_cp_assert_fix.py | 45 +++++++ 3 files changed, 187 insertions(+) create mode 100644 miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k_16node.py create mode 100644 miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k_8node_cp4.py create mode 100644 miles/megatron_dsa_cp_assert_fix.py diff --git a/miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k_16node.py b/miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k_16node.py new file mode 100644 index 0000000..9e5933c --- /dev/null +++ b/miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k_16node.py @@ -0,0 +1,113 @@ +"""GLM-5.2 LoRA DAPO 32k padded — 16 nodes, NO activation recompute, bf16 rollouts. +tp 8 x cp 4 -> 32 gpu replicas, dp 4, ep 32 unchanged + +""" + +from configs import glm5_2_744b_a40b_lora_dapo_tilelang_32k as _base +from configs.base import CHECKPOINTS_PATH, DATA_PATH, ModalConfig + +modal = ModalConfig( + docker_image=_base.modal.docker_image, + gpu=_base.modal.gpu, + memory=_base.modal.memory, + cloud=_base.modal.cloud, + region=_base.modal.region, + patch_files=[*_base.modal.patch_files, "miles/megatron_dsa_cp_assert_fix.py"], + image_run_commands=[ + *_base.modal.image_run_commands, + # megatron-core blanket-refuses CP for DSA; the TileLang bridge path + # has its own CP collectives (see the patch docstring), so gate the + # assert on the backend instead. + "python /tmp/megatron_dsa_cp_assert_fix.py", + ], + image_env=dict(_base.modal.image_env), +) + + +# filler tokens to mimic 32k context length +# 26k filler + header + question (~0.3-1.5k) ≈ 26.5-27.5k prompt tokens, +# + 4096 response ≤ ~31.7k, under the 32767 limit. +_FILLER_TOKENS = 26000 + + +class _Miles(_base._Miles): + actor_num_nodes = 16 + + # TP8 x CP4 -> 32-GPU replicas, DP 4. EP32 unchanged (128 % 32 == 0). + context_parallel_size = 4 + allgather_cp = True + + # The whole point of this config: no activation recompute. + recompute_granularity = None + recompute_method = None + recompute_num_layers = None + + prompt_data = f"{DATA_PATH}/dapo-math-17k/dapo-math-17k-pad26k.jsonl" + rollout_max_response_len = 4096 + + # 2 gradient steps per rollout (128 samples / global_batch_size 64). + num_rollout = 40 + save = f"{CHECKPOINTS_PATH}/GLM-5.2-lora-dapo-32k-16node-ckpt" + save_interval = 10 + + wandb_group = "glm5.2-744B-16node-tilelang-dapo-32k-norecompute" + + def download_data(self) -> None: + """Generate the padded dapo set (same as the retired padstress config).""" + import json + import os + import random + + from huggingface_hub import snapshot_download + from transformers import AutoTokenizer + + os.makedirs(f"{DATA_PATH}/dapo-math-17k", exist_ok=True) + snapshot_download( + repo_id="zhuzilin/dapo-math-17k", + repo_type="dataset", + local_dir=f"{DATA_PATH}/dapo-math-17k", + ) + + dst = f"{DATA_PATH}/dapo-math-17k/dapo-math-17k-pad26k.jsonl" + if os.path.exists(dst): + print(f"{dst} already exists, skipping generation") + return + + tokenizer = AutoTokenizer.from_pretrained( + "zai-org/GLM-5.2", trust_remote_code=True + ) + words = ( + "system model tensor kernel matrix vector gradient layer token " + "attention memory cache buffer stream block thread warp shard " + "sequence batch epoch metric loss reward policy value state action" + ).split() + + + rng = random.Random(0) + base = " ".join(rng.choice(words) for _ in range(2 * _FILLER_TOKENS)) + filler = tokenizer.decode( + tokenizer.encode(base, add_special_tokens=False)[:_FILLER_TOKENS] + ) + + src = f"{DATA_PATH}/dapo-math-17k/dapo-math-17k.jsonl" + n_written = 0 + with open(src) as fin, open(dst, "w") as fout: + for i, line in enumerate(fin): + sample = json.loads(line) + question = sample["prompt"][0]["content"] + sample["prompt"][0]["content"] = ( + f"Reference log #{i}-{random.Random(i).getrandbits(64):x} " + "follows. It is not relevant to the question; ignore it and " + "solve the question at the end.\n\n" + f"{filler}\n\n{question}" + ) + fout.write(json.dumps(sample) + "\n") + n_written += 1 + + total = len( + tokenizer.encode(json.loads(open(dst).readline())["prompt"][0]["content"]) + ) + print(f"Wrote {n_written} padded samples to {dst}; sample 0 = {total} tokens") + + +miles = _Miles() diff --git a/miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k_8node_cp4.py b/miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k_8node_cp4.py new file mode 100644 index 0000000..9d12c24 --- /dev/null +++ b/miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k_8node_cp4.py @@ -0,0 +1,29 @@ +"""GLM-5.2 LoRA DAPO 32k padded — 8 nodes, CP4, NO activation recompute. + +Same TP8 x CP4 geometry (and hence identical per-rank memory picture) as the +16-node config, on the readily-schedulable 8-node block: replica = 32 GPUs, +DP = 2. See glm5_2_744b_a40b_lora_dapo_tilelang_32k_16node for the CP design +notes and the megatron_dsa_cp_assert_fix rationale. + +Correctness oracle: train_rollout_logprob_abs_diff (~0.01 healthy, >1 means +broken CP attention math — stop the run if seen). + + EXPERIMENT_CONFIG=glm5_2_744b_a40b_lora_dapo_tilelang_32k_8node_cp4 \ + uv run modal run --detach miles/modal_train.py::train +""" + +from configs import glm5_2_744b_a40b_lora_dapo_tilelang_32k_16node as _base +from configs.base import CHECKPOINTS_PATH + +modal = _base.modal + + +class _Miles(_base._Miles): + actor_num_nodes = 8 + + save = f"{CHECKPOINTS_PATH}/GLM-5.2-lora-dapo-32k-8node-cp4-ckpt" + + wandb_group = "glm5.2-744B-8node-cp4-tilelang-dapo-32k-norecompute" + + +miles = _Miles() diff --git a/miles/megatron_dsa_cp_assert_fix.py b/miles/megatron_dsa_cp_assert_fix.py new file mode 100644 index 0000000..2798808 --- /dev/null +++ b/miles/megatron_dsa_cp_assert_fix.py @@ -0,0 +1,45 @@ +"""Relax megatron-core's blanket "no CP for DSA" assert for the TileLang backend. + +megatron-core's MCoreMLATransformerConfig.__post_init__ refuses +context_parallel_size > 1 whenever experimental_attention_variant == "dsa". +That is correct for the generic megatron-core DSAttention module (it has no +CP collectives), but Megatron-Bridge's GLM-5 TileLang path replaces that +module with CP-capable code (see bridge models/glm5/tilelang/tilelang_mla.py: +CP-gathered K, CP-local q RoPE replicate/slice, indexer varlen bounds +scattered over the CP group — the slime/baseten allgather-CP scheme). + +This patch gates the assert on the DSA kernel backend: CP stays forbidden +for the megatron backend, and is allowed for tilelang. finalize() calls +__post_init__ on the provider instance, which carries dsa_attention_backend, +so the getattr sees it; the default keeps the assert enforced otherwise. + +Correctness oracle for runs using this: train_rollout_logprob_abs_diff. +The trainer rescoring of SGLang-generated tokens catches wrong CP attention +math immediately (~0.01 healthy vs >1 broken). +""" + +from pathlib import Path + +TARGET = Path("/root/Megatron-LM/megatron/core/transformer/transformer_config.py") + +OLD = ''' elif self.experimental_attention_variant == "dsa": + assert ( + self.context_parallel_size == 1 + ), "Currently context parallelism is not supported by DSAttention!"''' + +NEW = ''' elif self.experimental_attention_variant == "dsa": + assert ( + self.context_parallel_size == 1 + or getattr(self, "dsa_attention_backend", "megatron") == "tilelang" + ), ( + "Context parallelism with DSA requires the TileLang backend " + "(megatron-core DSAttention has no CP collectives)." + )''' + +src = TARGET.read_text() +if NEW in src: + print("already patched, skipping") +else: + assert src.count(OLD) == 1, f"expected 1 match, got {src.count(OLD)}" + TARGET.write_text(src.replace(OLD, NEW)) + print(f"patched {TARGET}: DSA CP assert gated on tilelang backend")