diff --git a/3.test_cases/pytorch/slime/README.md b/3.test_cases/pytorch/slime/README.md index 72abf202c..60ed1348f 100644 --- a/3.test_cases/pytorch/slime/README.md +++ b/3.test_cases/pytorch/slime/README.md @@ -522,6 +522,8 @@ slime/ # 3.test_cases/pytorch/slime ├── env_vars.colocated.example # Base config: colocated train+rollout, built-in reward ├── env_vars.disaggregated.example # Overlay: reward model on a CPU pool + heavier GRPO ├── slime.Dockerfile # SLIME + SGLang + Megatron + EFA image +├── patches/ +│ └── apply_slime_patches.py # Self-neutralizing in-place fixes for the pinned upstream SLIME checkout (no-op once upstream merges them) ├── requirements.txt # Pinned Python RL dependencies ├── reward_service.Dockerfile # CPU-only image for the remote reward service ├── reward_service/ @@ -532,8 +534,10 @@ slime/ # 3.test_cases/pytorch/slime │ ├── reward-service.yaml # CPU reward service Deployment + Service │ └── data-prep-pod.yaml # Utility pod for data preparation ├── recipe/ -│ ├── run_grpo_qwen3_4b.sh # GRPO training launcher (Qwen3-4B, colocated) -│ └── run_grpo_qwen3_30b_a3b.sh # GRPO training launcher (Qwen3-30B-A3B MoE, disaggregated) +│ ├── run_grpo_qwen3_4b.sh # GRPO submit script (Qwen3-4B, colocated) +│ ├── run_grpo_qwen3_30b_a3b.sh # GRPO submit script (Qwen3-30B-A3B MoE, disaggregated) +│ └── launcher/ +│ └── grpo_launch.sh # Ray job entrypoint: sources the model script, expands MODEL_ARGS, execs train.py └── scripts/ ├── convert_checkpoint.sh # HF <-> Megatron conversion helper └── evaluate.sh # Evaluation launcher diff --git a/3.test_cases/pytorch/slime/patches/apply_slime_patches.py b/3.test_cases/pytorch/slime/patches/apply_slime_patches.py new file mode 100644 index 000000000..760d8c36f --- /dev/null +++ b/3.test_cases/pytorch/slime/patches/apply_slime_patches.py @@ -0,0 +1,406 @@ +#!/usr/bin/env python3 +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +"""Self-neutralizing patches for the pinned upstream SLIME checkout. + +This test case installs SLIME straight from upstream (``THUDM/slime`` at the +pinned ``SLIME_VERSION`` -- no fork, no forked URL). A handful of upstream bugs +still block a clean end-to-end run on B300 / H200 (CUDA 13). Rather than fork +SLIME or hard-code a fork URL, this script applies each fix *in place* against +the upstream checkout, and only when the unfixed pattern is actually present. + +Design goals (why this shape): + * Default is upstream. The image installs upstream SLIME as-is; this runs + afterwards as a thin, auditable layer. + * Self-neutralizing. Every patch first checks whether the upstream code still + exhibits the bug. Once upstream merges the corresponding fix, the pattern is + gone and the patch becomes a no-op automatically -- nothing to remember, no + version pin to bump, no fork to track. Upstream simply wins. + * Idempotent. Re-running (or running against an already-patched tree) is safe; + an applied patch is detected and skipped. + * Minimal blast radius. Each patch is scoped to the smallest possible edit and + is a no-op unless its exact precondition matches, so an unexpected upstream + refactor makes the patch skip (and say so) rather than corrupt the file. + +Each patch links to the upstream issue/PR that will make it unnecessary. When +all patches report "already fixed upstream", this file can be deleted. + +Usage (from the Dockerfile, right after the upstream SLIME install): + python3 patches/apply_slime_patches.py --slime-root /opt/slime + +Exit code is 0 whenever every patch ends in a known-good state (applied, +already-applied, or already-fixed-upstream). It is non-zero only if a patch's +target file is missing or a patch is genuinely unable to reach a good state, so +a broken image fails the build loudly instead of silently shipping the bug. +""" +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + + +# --- helper injected into slime/ray/actor_group.py -------------------------- +# Identical in behavior to the upstream fix proposed in the accompanying PR: +# delegate .so selection to torch_memory_saver's own CUDA-aware resolver, with a +# loadability-based fallback for older torch_memory_saver builds that predate it. +_TMS_HELPER = ''' + +def _resolve_tms_preload_lib(torch_memory_saver): + """Path to the torch_memory_saver preload .so that matches the CUDA runtime. + + Prefer the library's own CUDA-aware resolver. Fall back -- for older + torch_memory_saver builds that lack it -- to a candidate list that includes + the cu variant for the *detected* CUDA major and picks the first that + actually loads (existence is not loadability: a cu12 .so exists on a CUDA 13 + box but cannot be dlopen'd, which is what makes this fail on CUDA 13). + + Injected by awsome-distributed-training patches/apply_slime_patches.py as a + stopgap until the equivalent fix lands upstream in THUDM/slime. + """ + import os as _os + + stem = "torch_memory_saver_hook_mode_preload" + + try: + from torch_memory_saver.utils import get_binary_path_from_package + + return str(get_binary_path_from_package(stem)) + except Exception: + pass + + import ctypes + + base = _os.path.dirname(_os.path.dirname(torch_memory_saver.__file__)) + try: + import torch + + cuda = getattr(torch.version, "cuda", None) + major = cuda.split(".", 1)[0] if cuda else None + except Exception: + major = None + + candidates = [] + if major: + candidates.append(f"{stem}_cu{major}.abi3.so") + candidates += [f"{stem}.abi3.so", f"{stem}_cu12.abi3.so", f"{stem}_cu13.abi3.so"] + + tried = [] + for name in candidates: + path = _os.path.join(base, name) + if not _os.path.exists(path): + continue + try: + ctypes.CDLL(path) + except OSError as exc: + tried.append(f"{name}: {exc}") + continue + return path + + raise FileNotFoundError( + "Could not find a loadable torch_memory_saver preload library for the " + f"current CUDA runtime under {base}. Tried: {tried or candidates}" + ) +''' + +_HELPER_MARKER = "def _resolve_tms_preload_lib(" + +# The unfixed upstream pattern hard-codes the preload .so filename(s) and selects +# by existence. We match on the filename token that only appears in that unfixed +# path; once upstream selects by CUDA runtime, this token is gone and the patch +# self-neutralizes. +_BROKEN_TOKEN = '"torch_memory_saver_hook_mode_preload.abi3.so"' +_ENV_ASSIGN = 'env_vars["LD_PRELOAD"] = dynlib_path' +_ENV_ASSIGN_REPLACEMENT = ( + "dynlib_path = _resolve_tms_preload_lib(torch_memory_saver)\n\n" + ' env_vars["LD_PRELOAD"] = dynlib_path' +) + + +class PatchResult: + """Outcome of one patch: (status, message). status in known-good set => ok.""" + + GOOD = {"applied", "already-applied", "already-fixed-upstream"} + + def __init__(self, name: str, status: str, message: str): + self.name = name + self.status = status + self.message = message + + @property + def ok(self) -> bool: + return self.status in self.GOOD + + def __str__(self) -> str: + flag = "OK " if self.ok else "ERR" + return f"[{flag}] {self.name}: {self.status} -- {self.message}" + + +def patch_tms_preload_selection(slime_root: Path) -> PatchResult: + """Fix: pick the torch_memory_saver LD_PRELOAD .so by CUDA runtime, not by + filename existence (upstream issue/PR: torch_memory_saver cu13 selection). + + On CUDA 13, upstream selects a cu12-linked .so and every child dies with + 'libcudart.so.12: cannot open shared object file'. See the accompanying + SLIME PR for the full analysis. + """ + name = "tms-preload-cuda-aware" + target = slime_root / "slime" / "ray" / "actor_group.py" + if not target.is_file(): + return PatchResult(name, "error", f"target not found: {target}") + + src = target.read_text() + + # (a) already patched by us? + if _HELPER_MARKER in src: + return PatchResult(name, "already-applied", f"{target} already has the helper") + + # (b) upstream already fixed it? The unfixed path is identified by the + # hard-coded preload filename token; if it is gone, there is nothing to fix. + if _BROKEN_TOKEN not in src: + return PatchResult( + name, + "already-fixed-upstream", + "unfixed pattern absent; leaving upstream code untouched", + ) + + # (c) the assignment site we hang the helper call on must be present and + # unique, or we refuse to edit (an unexpected refactor -> skip, do not guess). + if src.count(_ENV_ASSIGN) != 1: + return PatchResult( + name, + "error", + f'expected exactly one occurrence of `{_ENV_ASSIGN}`, ' + f"found {src.count(_ENV_ASSIGN)}; refusing to edit", + ) + + # Insert the helper after the import block (after the last top-level import), + # and route the existing assignment through it. + lines = src.splitlines(keepends=True) + insert_at = 0 + for i, line in enumerate(lines): + if line.startswith(("import ", "from ")): + insert_at = i + 1 + patched = "".join(lines[:insert_at]) + _TMS_HELPER + "".join(lines[insert_at:]) + patched = patched.replace(_ENV_ASSIGN, _ENV_ASSIGN_REPLACEMENT, 1) + + # Byte-compile to guarantee we did not produce invalid Python. + try: + compile(patched, str(target), "exec") + except SyntaxError as exc: + return PatchResult(name, "error", f"patched file fails to compile: {exc}") + + target.write_text(patched) + return PatchResult(name, "applied", f"routed LD_PRELOAD through {_HELPER_MARKER[:-1]}()") + + +# --- wall: Megatron validate_args probes the GPU on a GPU-less Ray driver ----- +# Megatron's validate_args eagerly probes the CUDA device during pure argument +# validation (which SLIME runs on the Ray driver -- intentionally GPU-less in +# this test case: the head is num-gpus:0, GCS/dashboard only). Two probes are +# only reached by MoE models with tensor/context parallelism (the 30B recipe: +# --moe-grouped-gemm, TP=2, CP=2), so the 4B path never hits them: +# * torch.cuda.get_device_capability() (moe_grouped_gemm compute-cap assert) +# * megatron.training.utils.get_device_arch_version() (imported by name into +# megatron.training.arguments; used for the CUDA_DEVICE_MAX_CONNECTIONS note) +# On the GPU-less driver both raise "Found no NVIDIA driver". The real GPU actors +# (torch.cuda.is_available() == True) probe the real device unchanged. +# +# This injects a guard at the top of SLIME's own validate_args wrapper +# (slime/backends/megatron_utils/arguments.py) that, ONLY when no CUDA device is +# present (i.e. the driver), makes those two probes return safe values so +# argument validation can complete: +# * get_device_capability -> (8, 0): satisfies the moe_grouped_gemm assert +# (dc[0] >= 8) without a device; does not affect any arch-dependent branch. +# * get_device_arch_version -> a sentinel (9999) that is NOT any real GPU +# generation (Ampere=8, Hopper=9, Blackwell=10, ...), so it never mislabels +# the hardware -- on a GPU-less driver the arch is genuinely unknown. Being +# >= 10, it makes the driver SKIP the arch<10 CUDA_DEVICE_MAX_CONNECTIONS +# branch, deferring that decision to the real GPU actors. Returning a real +# generation (9 or 10) would falsely claim a specific arch; the sentinel does +# not. The actual requirement is enforced on the actors from their real arch +# (H200 = sm_90 requires it, satisfied via the recipe env; B300 = sm_100 does not). +# Guarded so substitution happens only while is_available() is False, and +# idempotent via a _slime_cpu_guard marker. +_VALIDATE_ANCHOR = '''def validate_args(args): + """Run megatron\'s own validate_args plus slime-specific megatron validations.""" + _megatron_validate_args(args)''' + +_VALIDATE_INJECT = '''def validate_args(args): + """Run megatron\'s own validate_args plus slime-specific megatron validations.""" + # Megatron validate_args eagerly probes the CUDA device (get_device_capability + # for moe_grouped_gemm; get_device_arch_version for the TP/CP + # CUDA_DEVICE_MAX_CONNECTIONS note). SLIME runs validate_args on the Ray + # driver, which is intentionally GPU-less in this test case, so those probes + # raise "Found no NVIDIA driver". Guard them for the GPU-less driver only; the + # real GPU actors probe the real device unchanged. Injected as a stopgap until + # Megatron guards these probes with torch.cuda.is_available() upstream. + import torch as _torch + + if not _torch.cuda.is_available(): + import megatron.training.arguments as _ma + + if not getattr(_torch.cuda.get_device_capability, "_slime_cpu_guard", False): + _orig_cap = _torch.cuda.get_device_capability + + def _cap(*a, **k): + return (8, 0) if not _torch.cuda.is_available() else _orig_cap(*a, **k) + + _cap._slime_cpu_guard = True + _torch.cuda.get_device_capability = _cap + + if not getattr(_ma.get_device_arch_version, "_slime_cpu_guard", False): + _orig_arch = _ma.get_device_arch_version + # A GPU-less driver has no GPU arch to report. Return a value that is + # deliberately NOT any real GPU generation (Ampere=8, Hopper=9, + # Blackwell=10, ...) so we never mislabel the hardware; it just has to + # be >= 10 so the arch-gated CUDA_DEVICE_MAX_CONNECTIONS branch is + # skipped on the driver. The real requirement is decided on the GPU + # actors from their real arch (they never re-run validate_args). + _ARCH_UNKNOWN_ON_GPULESS_DRIVER = 9999 + + def _arch(*a, **k): + return ( + _ARCH_UNKNOWN_ON_GPULESS_DRIVER + if not _torch.cuda.is_available() + else _orig_arch(*a, **k) + ) + + _arch._slime_cpu_guard = True + _ma.get_device_arch_version = _arch + + _megatron_validate_args(args)''' + + +def _megatron_probe_is_unguarded() -> bool: + """True if the installed Megatron's validate_args still probes the CUDA + device eagerly without a torch.cuda.is_available() guard. + + This is the real defect this patch works around. When Megatron guards the + probe upstream (the permanent fix), this returns False and the SLIME-side + patch self-neutralizes. Locating Megatron via import keeps this robust to + the install path. If Megatron cannot be located or its arguments module has + been restructured beyond recognition, we conservatively assume the probe is + still unguarded (better a harmless is_available()-gated guard than a crash). + """ + try: + import importlib.util + + spec = importlib.util.find_spec("megatron.training.arguments") + if spec is None or not spec.origin: + return True + src = Path(spec.origin).read_text() + except Exception: + return True + + if "get_device_capability()" not in src and "get_device_arch_version()" not in src: + # Probe removed/renamed entirely -> upstream changed it; nothing to guard. + return False + + # A probe is considered guarded if a torch.cuda.is_available() check appears + # on the probe's own line OR within the few lines immediately preceding it + # (upstream's likely fix is `if ...is_available(): dc = get_device_capability()` + # -- the guard sits on the enclosing `if`, not the probe line itself). If any + # eager get_device_capability() call has no is_available() nearby, the probe + # is still unguarded and the SLIME-side workaround is needed. + lines = src.splitlines() + WINDOW = 3 + for i, line in enumerate(lines): + if "get_device_capability()" not in line: + continue + context = lines[max(0, i - WINDOW): i + 1] + if not any("is_available" in c for c in context): + return True + return False + + +def patch_gpuless_driver_validate(slime_root: Path) -> PatchResult: + """Fix: let Megatron validate_args run on the GPU-less Ray driver. + + Guards the two eager CUDA device probes (get_device_capability, + get_device_arch_version) inside SLIME's validate_args wrapper so that on a + GPU-less driver they return safe values instead of crashing with + "Found no NVIDIA driver" (upstream Megatron issue/PR filed separately). + """ + name = "gpuless-driver-validate" + target = slime_root / "slime" / "backends" / "megatron_utils" / "arguments.py" + if not target.is_file(): + return PatchResult(name, "error", f"target not found: {target}") + + src = target.read_text() + + if "_slime_cpu_guard" in src: + return PatchResult(name, "already-applied", f"{target} already guards the driver probes") + + # True self-neutralization: the real defect is Megatron's *unguarded* eager + # device probe in validate_args. If Megatron has guarded it upstream (the + # permanent fix), this patch is unnecessary and must not touch SLIME. Detect + # the unguarded probe in the installed Megatron; if it is gone, no-op. + if not _megatron_probe_is_unguarded(): + return PatchResult( + name, + "already-fixed-upstream", + "Megatron validate_args no longer has an unguarded CUDA device probe; leaving SLIME untouched", + ) + + # The probe is still unguarded upstream, so the SLIME-side guard is needed. + # It must be injectable into SLIME's validate_args wrapper in its known shape. + if _VALIDATE_ANCHOR not in src: + return PatchResult( + name, + "error", + "Megatron probe is unguarded but SLIME validate_args is not in the expected shape to inject a guard; " + "refusing to edit (needs a refreshed anchor)", + ) + if src.count(_VALIDATE_ANCHOR) != 1: + return PatchResult( + name, + "error", + f"expected exactly one validate_args wrapper, found {src.count(_VALIDATE_ANCHOR)}", + ) + + patched = src.replace(_VALIDATE_ANCHOR, _VALIDATE_INJECT, 1) + try: + compile(patched, str(target), "exec") + except SyntaxError as exc: + return PatchResult(name, "error", f"patched file fails to compile: {exc}") + + target.write_text(patched) + return PatchResult(name, "applied", "guarded get_device_capability / get_device_arch_version on the GPU-less driver") + + +PATCHES = [ + patch_tms_preload_selection, + patch_gpuless_driver_validate, +] + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument( + "--slime-root", + default="/opt/slime", + type=Path, + help="Path to the upstream SLIME checkout (default: /opt/slime).", + ) + args = ap.parse_args() + + if not args.slime_root.is_dir(): + print(f"[patch] SLIME root not found: {args.slime_root}", file=sys.stderr) + return 2 + + print(f"[patch] applying self-neutralizing SLIME patches under {args.slime_root}") + results = [patch(args.slime_root) for patch in PATCHES] + for r in results: + print(f"[patch] {r}") + + failed = [r for r in results if not r.ok] + if failed: + print(f"[patch] {len(failed)} patch(es) failed", file=sys.stderr) + return 1 + print("[patch] all patches in a known-good state") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/3.test_cases/pytorch/slime/recipe/launcher/grpo_launch.sh b/3.test_cases/pytorch/slime/recipe/launcher/grpo_launch.sh new file mode 100755 index 000000000..1400b0fc6 --- /dev/null +++ b/3.test_cases/pytorch/slime/recipe/launcher/grpo_launch.sh @@ -0,0 +1,67 @@ +#!/usr/bin/env bash +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +# ============================================================ +# SLIME GRPO training entrypoint (runs on the Ray worker). +# +# This is the single shell that the Ray job entrypoint executes. It is +# uploaded to the Ray cluster via `ray job submit --working-dir ` +# and invoked as: ... -- bash grpo_launch.sh +# +# WHY A DEDICATED LAUNCHER (design note): +# SLIME's per-model scripts (scripts/models/*.sh) define MODEL_ARGS as a bash +# ARRAY. `ray job submit` re-joins everything after `--` with +# subprocess.list2cmdline and runs it through an outer `/bin/sh -c` +# (Popen(shell=True)). If the array were referenced in the submitted string +# (e.g. `-- bash -c "... ${MODEL_ARGS[@]} ..."`), that outer shell would +# expand ${MODEL_ARGS[@]} BEFORE this script sources the definition, so it +# would expand to zero elements and train.py would receive no model config. +# Keeping the `source` and the array expansion inside this one launcher — +# which the outer shell never expands, because the entrypoint tokens are just +# `bash grpo_launch.sh ` — makes the whole class of shell +# escaping bug impossible. This matches how SLIME's own upstream launch +# scripts (scripts/run-*.sh) expand ${MODEL_ARGS[@]} in the same shell that +# sourced it. +# +# Inputs: +# $@ : all train.py flags assembled by the recipe (scalar +# argv tokens, safely quoted by Ray). +# env MODEL_SCRIPT : the SLIME model script to source (e.g. qwen3-4B.sh), +# passed through the Ray runtime env by the recipe. +# env SLIME_DIR : SLIME install dir in the image (default /opt/slime). +# ============================================================ + +set -euo pipefail + +SLIME_DIR="${SLIME_DIR:-/opt/slime}" + +if [[ -z "${MODEL_SCRIPT:-}" ]]; then + echo "[grpo_launch] ERROR: MODEL_SCRIPT env var is not set." >&2 + exit 1 +fi + +cd "${SLIME_DIR}" + +MODEL_SCRIPT_PATH="scripts/models/${MODEL_SCRIPT}" +if [[ ! -f "${MODEL_SCRIPT_PATH}" ]]; then + echo "[grpo_launch] ERROR: model script ${SLIME_DIR}/${MODEL_SCRIPT_PATH} not found." >&2 + exit 1 +fi + +# Sourcing defines the MODEL_ARGS bash array in THIS shell. +# shellcheck disable=SC1090 +source "${MODEL_SCRIPT_PATH}" + +# Fail fast if the model script did not populate MODEL_ARGS. This is the exact +# condition that previously slipped through to train.py as "hidden_size None". +if [[ "${#MODEL_ARGS[@]}" -eq 0 ]]; then + echo "[grpo_launch] ERROR: MODEL_ARGS is empty after sourcing ${MODEL_SCRIPT_PATH}." >&2 + exit 1 +fi + +echo "[grpo_launch] MODEL_SCRIPT=${MODEL_SCRIPT} MODEL_ARGS count=${#MODEL_ARGS[@]}" +echo "[grpo_launch] launching: python3 train.py <${#MODEL_ARGS[@]} model args> $# recipe args" + +# MODEL_ARGS is expanded here (same shell that sourced it); the recipe-provided +# flags arrive as "$@". Quote both so values containing spaces stay single tokens. +exec python3 train.py "${MODEL_ARGS[@]}" "$@" diff --git a/3.test_cases/pytorch/slime/recipe/run_grpo_qwen3_30b_a3b.sh b/3.test_cases/pytorch/slime/recipe/run_grpo_qwen3_30b_a3b.sh index 574c5ca0f..f8ca6cdba 100644 --- a/3.test_cases/pytorch/slime/recipe/run_grpo_qwen3_30b_a3b.sh +++ b/3.test_cases/pytorch/slime/recipe/run_grpo_qwen3_30b_a3b.sh @@ -53,102 +53,133 @@ echo " Rollout BS: ${ROLLOUT_BATCH_SIZE} x ${N_SAMPLES_PER_PROMPT}" echo " Global BS: ${GLOBAL_BATCH_SIZE}" echo "============================================================" +# Build the train.py flags as a bash ARRAY (not a single string). Each element +# is one argv token, so values are never re-split by a shell. The array is +# expanded into the `ray job submit -- ...` argv below; MODEL_ARGS itself is +# expanded inside recipe/launcher/grpo_launch.sh, in the same shell that sources +# the SLIME model script. See that launcher for why this avoids the shell +# escaping trap that a `-- bash -c "...${MODEL_ARGS[@]}..."` string would hit. +# # When RM_TYPE=remote_rm, point SLIME at the CPU-hosted reward Service via # --rm-url (see kubernetes/reward-service.yaml). Otherwise scoring is in-process. -RM_ARGS="--rm-type ${RM_TYPE}" +RM_ARGS=(--rm-type "${RM_TYPE}") if [ "${RM_TYPE}" = "remote_rm" ]; then if [ -z "${RM_URL}" ]; then echo "[ERROR] RM_TYPE=remote_rm but RM_URL is not set. Configure it in env_vars." exit 1 fi - RM_ARGS="${RM_ARGS} --rm-url ${RM_URL}" + RM_ARGS+=(--rm-url "${RM_URL}") echo " Reward: remote_rm @ ${RM_URL}" fi -TRAIN_CMD="cd /opt/slime && source scripts/models/${MODEL_SCRIPT} && python3 train.py \ - \${MODEL_ARGS[@]} \ - --hf-checkpoint ${MODEL_LOCAL} \ - --ref-load ${MODEL_DIST} \ - --load ${CHECKPOINT_DIR}/qwen3-30b-a3b-grpo/ \ - --save ${CHECKPOINT_DIR}/qwen3-30b-a3b-grpo/ \ - --save-interval ${SAVE_INTERVAL} \ - \ - --prompt-data ${PROMPT_DATA} \ - --input-key prompt \ - --label-key label \ - --apply-chat-template \ - --rollout-shuffle \ - \ - ${RM_ARGS} \ - \ - --num-rollout ${NUM_ROLLOUT} \ - --rollout-batch-size ${ROLLOUT_BATCH_SIZE} \ - --n-samples-per-prompt ${N_SAMPLES_PER_PROMPT} \ - --num-steps-per-rollout ${NUM_STEPS_PER_ROLLOUT:-1} \ - --global-batch-size ${GLOBAL_BATCH_SIZE} \ - \ - --rollout-max-response-len ${ROLLOUT_MAX_RESPONSE_LEN} \ - --rollout-temperature ${ROLLOUT_TEMPERATURE} \ - --balance-data \ - \ - --eval-interval 10 \ - --eval-prompt-data aime ${EVAL_DATA} \ - --n-samples-per-eval-prompt 4 \ - --eval-max-response-len 16384 \ - --eval-top-p 1 \ - \ - --tensor-model-parallel-size ${TP_SIZE} \ - --pipeline-model-parallel-size ${PP_SIZE} \ - --context-parallel-size ${CP_SIZE} \ - --expert-model-parallel-size ${EP_SIZE} \ - --expert-tensor-parallel-size 1 \ - --sequence-parallel \ - \ - --recompute-granularity full \ - --recompute-method uniform \ - --recompute-num-layers 1 \ - \ - --use-dynamic-batch-size \ - --max-tokens-per-gpu ${MAX_TOKENS_PER_GPU} \ - \ - --advantage-estimator grpo \ - --use-kl-loss \ - --kl-loss-coef 0.00 \ - --kl-loss-type low_var_kl \ - --entropy-coef 0.00 \ - --eps-clip 0.2 \ - --eps-clip-high 0.28 \ - \ - --optimizer adam \ - --lr ${LEARNING_RATE} \ - --lr-decay-style constant \ - --weight-decay 0.1 \ - --adam-beta1 0.9 \ - --adam-beta2 0.98 \ - \ - --actor-num-nodes ${ACTOR_NUM_NODES} \ - --actor-num-gpus-per-node ${ACTOR_GPUS_PER_NODE} \ - --rollout-num-gpus ${ROLLOUT_NUM_GPUS} \ - --rollout-num-gpus-per-engine ${ROLLOUT_GPUS_PER_ENGINE} \ - \ - --sglang-mem-fraction-static 0.85 \ - --sglang-log-level WARN \ - --sglang-enable-ep-moe" +TRAIN_ARGS=( + --hf-checkpoint "${MODEL_LOCAL}" + --ref-load "${MODEL_DIST}" + --load "${CHECKPOINT_DIR}/qwen3-30b-a3b-grpo/" + --save "${CHECKPOINT_DIR}/qwen3-30b-a3b-grpo/" + --save-interval "${SAVE_INTERVAL}" + + --prompt-data "${PROMPT_DATA}" + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + + "${RM_ARGS[@]}" + + --num-rollout "${NUM_ROLLOUT}" + --rollout-batch-size "${ROLLOUT_BATCH_SIZE}" + --n-samples-per-prompt "${N_SAMPLES_PER_PROMPT}" + --num-steps-per-rollout "${NUM_STEPS_PER_ROLLOUT:-1}" + --global-batch-size "${GLOBAL_BATCH_SIZE}" + + --rollout-max-response-len "${ROLLOUT_MAX_RESPONSE_LEN}" + --rollout-temperature "${ROLLOUT_TEMPERATURE}" + --balance-data + + --eval-interval 10 + --eval-prompt-data aime "${EVAL_DATA}" + --n-samples-per-eval-prompt 4 + --eval-max-response-len 16384 + --eval-top-p 1 + + --tensor-model-parallel-size "${TP_SIZE}" + --pipeline-model-parallel-size "${PP_SIZE}" + --context-parallel-size "${CP_SIZE}" + --expert-model-parallel-size "${EP_SIZE}" + --expert-tensor-parallel-size 1 + --sequence-parallel + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + --use-dynamic-batch-size + --max-tokens-per-gpu "${MAX_TOKENS_PER_GPU}" + + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 + + --optimizer adam + --lr "${LEARNING_RATE}" + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + + --actor-num-nodes "${ACTOR_NUM_NODES}" + --actor-num-gpus-per-node "${ACTOR_GPUS_PER_NODE}" + --rollout-num-gpus "${ROLLOUT_NUM_GPUS}" + --rollout-num-gpus-per-engine "${ROLLOUT_GPUS_PER_ENGINE}" + + --sglang-mem-fraction-static 0.85 + # SGLang forwards this to uvicorn's log_level, whose LOG_LEVELS dict is keyed + # by lowercase names only (critical/error/warning/info/debug/trace) with no + # "warn" key. Uppercase "WARN" (the original value) raises a KeyError and + # uvicorn dies before the rollout HTTP server binds, so the rollout health + # check never passes and training hangs before it starts. Use lowercase + # "warning" to preserve the original intended verbosity. + --sglang-log-level warning + # NOTE: the original recipe passed `--sglang-enable-ep-moe`, which SGLang + # 0.5.12 removed. SLIME v0.2.4 registers --sglang-* flags from SGLang's live + # ServerArgs (parse_known_args / ignore_unknown_args), so the dead flag is + # silently ignored rather than erroring -- but it configures nothing, so it + # is dropped here. No replacement flag is needed for this recipe: the rollout + # engine serves the Qwen3-30B-A3B MoE correctly with the SGLang defaults + # (moe_runner_backend=auto resolves to the triton runner for bf16 on H200; + # ep_size defaults to 1, which is a valid serving mode). Both were verified + # unnecessary by a full end-to-end run with neither flag set. +) + +# Submit via Ray job API. +# +# The entrypoint after `--` is `bash grpo_launch.sh ` (plain argv tokens, +# no shell array crosses the ray boundary). --working-dir uploads the launcher +# to the Ray workers; the SLIME code itself is already in the image at +# /opt/slime. MODEL_SCRIPT is forwarded so the launcher can source the right +# model definition. echo "[INFO] Submitting Ray job for MoE GRPO training..." ray job submit \ --address="http://127.0.0.1:8265" \ + --working-dir "${SCRIPT_DIR}/launcher" \ --runtime-env-json="{ \"env_vars\": { \"PYTHONPATH\": \"/opt/Megatron-LM\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", \"HF_TOKEN\": \"${HF_TOKEN}\", + \"MODEL_SCRIPT\": \"${MODEL_SCRIPT}\", \"TOKENIZERS_PARALLELISM\": \"false\", \"NCCL_DEBUG\": \"WARN\", \"FI_PROVIDER\": \"efa\", \"FI_EFA_USE_DEVICE_RDMA\": \"1\" } }" \ - -- bash -c "${TRAIN_CMD}" + -- bash grpo_launch.sh "${TRAIN_ARGS[@]}" echo "[INFO] Job submitted. Monitor at http://localhost:8265" diff --git a/3.test_cases/pytorch/slime/recipe/run_grpo_qwen3_4b.sh b/3.test_cases/pytorch/slime/recipe/run_grpo_qwen3_4b.sh index c2cdb635a..f8202e23f 100644 --- a/3.test_cases/pytorch/slime/recipe/run_grpo_qwen3_4b.sh +++ b/3.test_cases/pytorch/slime/recipe/run_grpo_qwen3_4b.sh @@ -53,102 +53,122 @@ echo " Global BS: ${GLOBAL_BATCH_SIZE}" echo " Num rollouts: ${NUM_ROLLOUT}" echo "============================================================" -# Build the training command +# Build the train.py flags as a bash ARRAY (not a single string). Each element +# is one argv token, so values are never re-split by a shell. The array is +# expanded into the `ray job submit -- ...` argv below; MODEL_ARGS itself is +# expanded inside recipe/launcher/grpo_launch.sh, in the same shell that sources +# the SLIME model script. See that launcher for why this avoids the shell +# escaping trap that a `-- bash -c "...${MODEL_ARGS[@]}..."` string would hit. +# # When RM_TYPE=remote_rm, point SLIME at the CPU-hosted reward Service via # --rm-url (see kubernetes/reward-service.yaml). Otherwise scoring is in-process. -RM_ARGS="--rm-type ${RM_TYPE}" +RM_ARGS=(--rm-type "${RM_TYPE}") if [ "${RM_TYPE}" = "remote_rm" ]; then if [ -z "${RM_URL}" ]; then echo "[ERROR] RM_TYPE=remote_rm but RM_URL is not set. Configure it in env_vars." exit 1 fi - RM_ARGS="${RM_ARGS} --rm-url ${RM_URL}" + RM_ARGS+=(--rm-url "${RM_URL}") echo " Reward: remote_rm @ ${RM_URL}" fi -TRAIN_CMD="cd /opt/slime && source scripts/models/${MODEL_SCRIPT} && python3 train.py \ - \${MODEL_ARGS[@]} \ - --hf-checkpoint ${MODEL_LOCAL} \ - --ref-load ${MODEL_DIST} \ - --load ${CHECKPOINT_DIR}/qwen3-4b-grpo/ \ - --save ${CHECKPOINT_DIR}/qwen3-4b-grpo/ \ - --save-interval ${SAVE_INTERVAL} \ - \ - --prompt-data ${PROMPT_DATA} \ - --input-key prompt \ - --label-key label \ - --apply-chat-template \ - --rollout-shuffle \ - \ - ${RM_ARGS} \ - \ - --num-rollout ${NUM_ROLLOUT} \ - --rollout-batch-size ${ROLLOUT_BATCH_SIZE} \ - --n-samples-per-prompt ${N_SAMPLES_PER_PROMPT} \ - --num-steps-per-rollout ${NUM_STEPS_PER_ROLLOUT:-1} \ - --global-batch-size ${GLOBAL_BATCH_SIZE} \ - \ - --rollout-max-response-len ${ROLLOUT_MAX_RESPONSE_LEN} \ - --rollout-temperature ${ROLLOUT_TEMPERATURE} \ - --balance-data \ - \ - --eval-interval 10 \ - --eval-prompt-data aime ${EVAL_DATA} \ - --n-samples-per-eval-prompt 8 \ - --eval-max-response-len 16384 \ - --eval-top-p 1 \ - \ - --tensor-model-parallel-size ${TP_SIZE} \ - --pipeline-model-parallel-size ${PP_SIZE} \ - --context-parallel-size ${CP_SIZE} \ - --expert-model-parallel-size ${EP_SIZE} \ - --sequence-parallel \ - \ - --recompute-granularity full \ - --recompute-method uniform \ - --recompute-num-layers 1 \ - \ - --use-dynamic-batch-size \ - --max-tokens-per-gpu ${MAX_TOKENS_PER_GPU} \ - \ - --advantage-estimator grpo \ - --use-kl-loss \ - --kl-loss-coef 0.00 \ - --kl-loss-type low_var_kl \ - --entropy-coef 0.00 \ - --eps-clip 0.2 \ - --eps-clip-high 0.28 \ - \ - --optimizer adam \ - --lr ${LEARNING_RATE} \ - --lr-decay-style constant \ - --weight-decay 0.1 \ - --adam-beta1 0.9 \ - --adam-beta2 0.98 \ - \ - --actor-num-nodes ${ACTOR_NUM_NODES} \ - --actor-num-gpus-per-node ${ACTOR_GPUS_PER_NODE} \ - --colocate \ - --rollout-num-gpus-per-engine ${ROLLOUT_GPUS_PER_ENGINE} \ - \ - --sglang-mem-fraction-static 0.8 \ - --sglang-log-level WARN" - -# Submit via Ray job API +TRAIN_ARGS=( + --hf-checkpoint "${MODEL_LOCAL}" + --ref-load "${MODEL_DIST}" + --load "${CHECKPOINT_DIR}/qwen3-4b-grpo/" + --save "${CHECKPOINT_DIR}/qwen3-4b-grpo/" + --save-interval "${SAVE_INTERVAL}" + + --prompt-data "${PROMPT_DATA}" + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + + "${RM_ARGS[@]}" + + --num-rollout "${NUM_ROLLOUT}" + --rollout-batch-size "${ROLLOUT_BATCH_SIZE}" + --n-samples-per-prompt "${N_SAMPLES_PER_PROMPT}" + --num-steps-per-rollout "${NUM_STEPS_PER_ROLLOUT:-1}" + --global-batch-size "${GLOBAL_BATCH_SIZE}" + + --rollout-max-response-len "${ROLLOUT_MAX_RESPONSE_LEN}" + --rollout-temperature "${ROLLOUT_TEMPERATURE}" + --balance-data + + --eval-interval 10 + --eval-prompt-data aime "${EVAL_DATA}" + --n-samples-per-eval-prompt 8 + --eval-max-response-len 16384 + --eval-top-p 1 + + --tensor-model-parallel-size "${TP_SIZE}" + --pipeline-model-parallel-size "${PP_SIZE}" + --context-parallel-size "${CP_SIZE}" + --expert-model-parallel-size "${EP_SIZE}" + --sequence-parallel + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + --use-dynamic-batch-size + --max-tokens-per-gpu "${MAX_TOKENS_PER_GPU}" + + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 + + --optimizer adam + --lr "${LEARNING_RATE}" + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + + --actor-num-nodes "${ACTOR_NUM_NODES}" + --actor-num-gpus-per-node "${ACTOR_GPUS_PER_NODE}" + --colocate + --rollout-num-gpus-per-engine "${ROLLOUT_GPUS_PER_ENGINE}" + + --sglang-mem-fraction-static 0.8 + # SGLang forwards this to uvicorn's log_level, whose LOG_LEVELS dict is keyed + # by lowercase names only (critical/error/warning/info/debug/trace) with no + # "warn" key. Uppercase "WARN" (the original value) raises a KeyError and + # uvicorn dies before the rollout HTTP server binds, so the rollout health + # check never passes and training hangs before it starts. Use lowercase + # "warning" to preserve the original intended verbosity. + --sglang-log-level warning +) + +# Submit via Ray job API. +# +# The entrypoint after `--` is `bash grpo_launch.sh ` (plain argv tokens, +# no shell array crosses the ray boundary). --working-dir uploads the launcher +# to the Ray workers; the SLIME code itself is already in the image at +# /opt/slime. MODEL_SCRIPT is forwarded so the launcher can source the right +# model definition. echo "[INFO] Submitting Ray job..." ray job submit \ --address="http://127.0.0.1:8265" \ + --working-dir "${SCRIPT_DIR}/launcher" \ --runtime-env-json="{ \"env_vars\": { \"PYTHONPATH\": \"/opt/Megatron-LM\", \"HF_TOKEN\": \"${HF_TOKEN}\", + \"MODEL_SCRIPT\": \"${MODEL_SCRIPT}\", \"TOKENIZERS_PARALLELISM\": \"false\", \"NCCL_DEBUG\": \"WARN\", \"FI_PROVIDER\": \"efa\", \"FI_EFA_USE_DEVICE_RDMA\": \"1\" } }" \ - -- bash -c "${TRAIN_CMD}" + -- bash grpo_launch.sh "${TRAIN_ARGS[@]}" echo "[INFO] Job submitted. Monitor at http://localhost:8265" diff --git a/3.test_cases/pytorch/slime/requirements.txt b/3.test_cases/pytorch/slime/requirements.txt index b602ef3cf..cc0ea374b 100644 --- a/3.test_cases/pytorch/slime/requirements.txt +++ b/3.test_cases/pytorch/slime/requirements.txt @@ -2,6 +2,26 @@ # SPDX-License-Identifier: MIT-0 # Reinforcement-learning Python dependencies for the SLIME image. + +# Megatron-LM asserts numpy 1.x at init (slime .../megatron_utils/initialize.py: +# `assert np.__version__.startswith("1.")`, per NVIDIA/Megatron-LM#1563), but +# sglang[all] (installed just before this file in slime.Dockerfile) pulls in +# numpy 2.x transitively. Pin numpy < 2 so this requirements install downgrades +# it back to 1.x; without it the train worker aborts during Megatron init with +# "Megatron does not support numpy 2.x". Upstream SLIME's own docker/Dockerfile +# does the same (`pip install "numpy<2"`), so this mirrors the sanctioned pin. +# +# Why 1.x survives to runtime: the only pip steps after this one in the +# Dockerfile are the slime and sgl-router installs (--no-deps) and +# ring_flash_attn==0.1.8 (which declares no dependencies), so none of them +# reintroduce numpy 2.x. If a later step is added that pulls numpy WITHOUT +# --no-deps, re-pin numpy after it. +# +# TODO(numpy<2): remove this pin once MEGATRON_LM_VERSION is bumped to a commit +# that no longer asserts numpy 1.x (drops the numpy 2.x assert); until then the +# pin is required. +numpy<2 + nltk==3.9.4 awscli==1.45.26 pynvml==13.0.1 diff --git a/3.test_cases/pytorch/slime/slime.Dockerfile b/3.test_cases/pytorch/slime/slime.Dockerfile index cd28cebc7..48ef014be 100644 --- a/3.test_cases/pytorch/slime/slime.Dockerfile +++ b/3.test_cases/pytorch/slime/slime.Dockerfile @@ -180,12 +180,37 @@ RUN pip install --no-cache-dir --force-reinstall --no-deps "${SGL_ROUTER_WHEEL}" ##################### # SLIME (RL post-training framework) +# +# Installed straight from upstream THUDM/slime at the pinned SLIME_VERSION -- +# no fork, no forked URL. A few upstream bugs still block a clean end-to-end run +# on CUDA 13; they are healed in the next step against this upstream checkout. ##################### RUN cd /opt && \ git clone --depth 1 --branch ${SLIME_VERSION} https://github.com/THUDM/slime.git && \ cd slime && \ pip install --no-cache-dir -e . --no-deps +# mbridge: required by SLIME's HF->torch_dist converter +# (tools/convert_hf_to_torch_dist.py imports `slime_plugins.mbridge` and +# `from mbridge import AutoBridge`). It is NOT pulled by the `--no-deps` slime +# install above, so the 30B MoE checkpoint conversion fails with +# `ModuleNotFoundError: No module named 'mbridge'` without it. Installed with +# --no-deps so it cannot drag numpy 2.x (or any other pinned dep) back in. +RUN pip install --no-cache-dir --no-deps mbridge + +##################### +# Self-neutralizing upstream patches. +# +# Applies small, in-place fixes to the upstream SLIME checkout above -- only +# where the unfixed pattern is actually present. Each patch checks upstream +# first, so once the corresponding fix lands in THUDM/slime the patch becomes a +# no-op automatically (no fork to track, no URL to bump; upstream simply wins). +# The build fails loudly if a patch cannot reach a known-good state. See +# patches/apply_slime_patches.py for the per-patch rationale and upstream links. +##################### +COPY patches/apply_slime_patches.py /opt/slime-patches/apply_slime_patches.py +RUN python3 /opt/slime-patches/apply_slime_patches.py --slime-root /opt/slime + ## Set Open MPI variables to exclude network interface and conduit. ENV OMPI_MCA_pml=^ucx \ OMPI_MCA_btl=tcp,self \