Enable the SLIME GRPO test case to run end-to-end on B300 / H200 - #1164
Conversation
littlemex
changed the title
littlemex
changed the title
First fix in this epic:
|
| Launch path | how ${MODEL_ARGS[@]} expands |
args reaching the process |
|---|---|---|
single bash -c "..." |
one shell | 25 (correct) |
ray job submit -- bash -c "..." (current recipe) |
outer /bin/sh expands first |
0 (hidden_size None) |
ray job submit -- bash grpo_launch.sh ... (this fix) |
inside the launcher | 25 (correct) |
End-to-end with the real train.py, Megatron receives the config and passes hf_validate_args, so training starts:
[grpo_launch] MODEL_SCRIPT=qwen3-4B.sh MODEL_ARGS count=25
[grpo_launch] launching: python3 train.py <25 model args> 100 recipe args
ffn_hidden_size ................................. 9728
hidden_size ..................................... 2560 # was: None
… starts The GRPO recipes passed --sglang-log-level WARN (uppercase). SLIME forwards the value verbatim into SGLang's ServerArgs, and SGLang hands it to uvicorn, whose LOG_LEVELS dict is keyed by lowercase names only (no 'warn' key). uvicorn raises KeyError while building its Config, before the socket binds, so the rollout HTTP server never comes up: the engine core idles, the port never listens, and the RolloutManager waits on an unreachable /health_generate forever (GPU idle). Use lowercase 'info' in both run_grpo_qwen3_4b.sh and run_grpo_qwen3_30b_a3b.sh. The proper normalization belongs upstream in SGLang (it accepts WARN for its own stdlib logger but not for uvicorn); this recipe change is the immediate unblock.
… of bash -c
The GRPO recipes built the whole train.py invocation into a single string and
submitted it as `ray job submit ... -- bash -c "${TRAIN_CMD}"`, with the
model config deferred as `\${MODEL_ARGS[@]}`. ray job submit re-joins the
tokens after `--` with subprocess.list2cmdline and runs them through an outer
`/bin/sh -c` (Popen(shell=True)), so that outer shell expands ${MODEL_ARGS[@]}
to zero elements before the inner bash sources the model script. train.py then
receives no model args and aborts at hf_validate_args ("hidden_size ... None").
Move the launch into recipe/launcher/grpo_launch.sh, which sources the model
script and expands "${MODEL_ARGS[@]}" in the same shell, and submit it as
`-- bash grpo_launch.sh <flags>` with --working-dir. No bash array ever crosses
the ray-submit boundary, matching how SLIME's own scripts/run-*.sh launch. The
recipes now assemble train.py flags as a bash array (argv tokens); the README
file tree lists the new launcher.
Applies to both run_grpo_qwen3_4b.sh and run_grpo_qwen3_30b_a3b.sh.
fb58e3c to
5a3777f
Compare
Verified that uvicorn accepts lowercase 'warning' (its LOG_LEVELS has a 'warning' key but no 'warn' key). The original recipe intent was WARN-level verbosity, so use lowercase 'warning' rather than 'info' to keep that intent while avoiding the uppercase KeyError that hangs the rollout HTTP server.
fix(slime): use lowercase sglang log level so the rollout HTTP server starts
…in-place On CUDA 13 the upstream SLIME actor_group.py picks the torch_memory_saver preload .so by filename existence (cu12, then unsuffixed), not by CUDA runtime. Both are cu12-linked, so LD_PRELOAD makes the train worker die with 'libcudart.so.12: cannot open shared object file'; the correct cu13 build ships in the wheel but is never enumerated. Rather than fork SLIME or hard-code a fork URL, install upstream SLIME as-is and run a self-neutralizing patch step that rewrites the .so selection to delegate to torch_memory_saver's own CUDA-aware resolver (get_binary_path_from_package), with a loadability-based fallback for older torch_memory_saver. The patch only touches the file when the unfixed pattern is present, so once the equivalent fix lands upstream in THUDM/slime it becomes a no-op automatically. Verified on H200 (CUDA 13): LD_PRELOAD resolves to the cu13 build and the train worker starts (no libcudart.so.12 error), reaching Megatron init.
fix(slime): heal the CUDA-13 torch_memory_saver LD_PRELOAD selection without forking SLIME
…orker Megatron-LM asserts numpy 1.x at init (per NVIDIA/Megatron-LM#1563), but sglang[all] -- installed just before requirements.txt in slime.Dockerfile -- pulls numpy 2.x transitively, so MegatronTrainRayActor dies during init with 'AssertionError: Megatron does not support numpy 2.x'. Pin numpy<2 in requirements.txt. Because the requirements install runs after sglang[all] and the later slime/sgl-router installs use --no-deps, this downgrades numpy back to 1.26.4 and nothing reintroduces 2.x. Verified on H200 (CUDA 13): numpy 1.26.4, torch/megatron import cleanly, Megatron init passes and the train worker proceeds into the rollout/train cycle.
Fixes an inaccurate comment (the later ring_flash_attn install is not --no-deps) and records why numpy 1.x survives to runtime: the only post-pin pip steps are slime/sgl-router (--no-deps) and ring_flash_attn 0.1.8 (declares no deps), so none reintroduce numpy 2.x. Notes that upstream SLIME's own Dockerfile pins numpy<2 the same way, and adds a TODO tying the pin's removal to a future MEGATRON_LM_VERSION that drops the numpy 1.x assert. No functional change.
fix(slime): pin numpy < 2 so Megatron init does not abort the train worker
Bring up recipe/run_grpo_qwen3_30b_a3b.sh (Qwen3-30B-A3B, disaggregated) to a full GRPO training step. The 4B colocated path already works; the 30B MoE recipe hit several MoE/parallelism-specific walls that 4B does not: - Checkpoint conversion needs mbridge (convert_hf_to_torch_dist.py imports it), not pulled by the --no-deps slime install. Install it in the Dockerfile with --no-deps so it cannot drag numpy 2.x back in. - SGLang 0.5.12 removed --enable-ep-moe; use --sglang-moe-runner-backend triton + --sglang-expert-parallel-size instead so the rollout engine starts. - Megatron validate_args eagerly probes the CUDA device (get_device_capability for moe_grouped_gemm, get_device_arch_version for the TP/CP note) on the Ray driver, which is intentionally GPU-less (head num-gpus:0), crashing with 'Found no NVIDIA driver'. Guard both probes in SLIME's validate_args wrapper only when torch.cuda.is_available() is False (the driver); GPU actors probe the real device unchanged. get_device_arch_version returns 10 on the driver so it defers the CUDA_DEVICE_MAX_CONNECTIONS decision to the real actors rather than faking a pre-Blackwell arch (which would wrongly force it on B300). - TP/CP>1 requires CUDA_DEVICE_MAX_CONNECTIONS=1; add it to the recipe runtime-env. - Cap --sglang-cuda-graph-max-bs 8 so large-HBM capture is not pathologically slow. The driver-probe guard is added to the existing self-neutralizing patch mechanism (patches/apply_slime_patches.py): it edits the upstream checkout only when the un-guarded probe is present, is idempotent, byte-compiles, and no-ops once Megatron guards the probes upstream (issue/PR to be filed there). Verified on H200 (2x p5en.48xlarge, CUDA 13): the 30B MoE recipe reaches a full GRPO loop (rollout -> ref/actor log-probs -> Timer train end -> weight sync), with no wall regressions and no MoE weight-sync 400.
…lizing The wall-9 guard was keyed only on SLIME's validate_args wrapper shape, so it would keep applying (harmlessly, since it is is_available()-gated) even after Megatron guards the probe upstream. Make it detect the actual defect: inspect the installed Megatron's validate_args and apply the SLIME-side guard only while an eager get_device_capability() probe remains unguarded (no torch.cuda.is_available() on or just before its line). Once Megatron guards or removes the probe, the patch reports already-fixed-upstream and leaves SLIME untouched -- matching the self-neutralizing behavior of the tms-preload patch. Verified all states on H200: applied / already-applied / already-fixed-upstream (guarded same-line, guarded enclosing-if, and probe-removed).
An audit + end-to-end retest showed two of the SGLang flags added while bringing up the 30B MoE recipe are not needed, so remove them and keep the recipe minimal: - --sglang-moe-runner-backend triton: redundant. On SGLang 0.5.12 the default moe_runner_backend=auto resolves to the same triton runner for an unquantized bf16 MoE on H200 (sm_90); forcing triton changes nothing. - --sglang-expert-parallel-size (sglang ep_size): not required. The rollout engine's expert parallelism is independent of Megatron's training EP, and ep_size=1 (the SGLang default) is a valid serving mode for Qwen3-30B-A3B. Also correct the comment about --sglang-enable-ep-moe: SGLang 0.5.12 removed the flag, but SLIME v0.2.4 parses --sglang-* leniently (parse_known_args / ignore_unknown_args), so the dead flag is silently ignored, not rejected. It is dropped because it configures nothing, not because it errors. Verified on H200: the 30B MoE disaggregated recipe reaches a full GRPO loop (Timer train end) with none of these flags set.
… on the GPU-less driver Return _ARCH_UNKNOWN_ON_GPULESS_DRIVER = 9999 instead of 10 from the GPU-less-driver guard. 9999 is deliberately not any real GPU generation (Ampere=8, Hopper=9, Blackwell=10), so the driver never mislabels the hardware; being >= 10 it still skips the arch<10 CUDA_DEVICE_MAX_CONNECTIONS branch and defers that decision to the real GPU actors. Matches the image already verified end-to-end on H200 (get_device_arch_version = 9999).
fix(slime): enable the 30B MoE disaggregated recipe to run end-to-end
littlemex
changed the title
allela-roy
left a comment
allela-roy
left a comment
There was a problem hiding this comment.
@littlemex , thanks for the PR. These are clean, well-engineered fixes and verified the end-end execution on my cluster. LGTM. Approved.
|
@allela-roy Thank you so much for taking the time to thoroughly verify the changes end-to-end on your cluster — I really appreciate it! |
2 participants
Purpose
Relates to #1163.
This is the integration PR for a small epic: making the SLIME GRPO test case (
3.test_cases/pytorch/slime, added in #1129) run through to training on NVIDIA B300 (p6) and H200 (p5en). Bringing the test case up surfaced a series of issues, each of which is a small, self-contained fix. Rather than one large "B300 support" change, the epic lands them as focused commits on this branch so each is easy to review, and this PR tracks the overall progress.Thanks to the authors and reviewers of #1129 for the foundation this builds on: the container image, the reward-service split, the RayCluster manifests and the docs. The reward path works as designed, and every change here is scoped to getting training to start and run.
How this PR is organized
Each fix is one focused commit. The mechanism, reproduction commands, and evidence for a fix live in a dedicated comment on this PR, not in this description. This description only tracks the overall goal, the scope, and the progress checklist.
Scope and progress
The findings from #1163 are grouped by root cause. Most are general (they affect the documented p5/HyperPod path too — CUDA 13 and dependency issues, not GPU-generation-specific); a few are scoped to the 30B MoE recipe and are tracked separately below. Each item links to its detail comment as it lands.
General fixes (affect the documented p5/HyperPod path too; CUDA 13 / dependency issues, not GPU-generation-specific):
MODEL_ARGS reaches
train.pyas zero elements, so training never starts. Fixed by launching through a small launcher script instead ofray job submit -- bash -c "...".--sglang-log-level WARNbreaks uvicorn startup (uvicorn accepts lowercase log levels only), so the rollout HTTP server never binds and training hangs. Fixed by using lowercasewarning(preserving the original intended verbosity).The train worker dies on CUDA 13 because SLIME
LD_PRELOADs a cu12-linkedtorch_memory_saver.so. This is a CUDA-13 issue, not GPU-generation-specific: it reproduces on any CUDA-13 image (the test case's NGC 26.02 base), on H100/H200/B300 alike. Proper fix filed upstream in SLIME (delegate.soselection totorch_memory_saver's own CUDA-aware resolver); test-case side applies the same fix in place via a self-neutralizing build-time patch (no fork, no forked URL, no-op once upstream merges it).slime.Dockerfile; resolves the_cu13.so, loads cleanly, train worker reaches training)Megatron aborts on numpy 2.x pulled in transitively by
sglang[all]. Fixed by pinningnumpy<2inrequirements.txt(the last dependency step that touches numpy), matching upstream SLIME's own Dockerfile. numpy is a pure-CPU package, so this is GPU-generation- and CUDA-independent (identical on H100/H200/B300).slime.Dockerfile; numpy 1.26.4, Megatron init passes)30B MoE recipe only (not exercised by the Qwen3-4B path; needed only for
run_grpo_qwen3_30b_a3b.sh). All three are bundled into one experiment-scoped feature PR because the recipe does not reach training until all are in place:30B MoE
torch_distconversion needsmbridge, which the--no-depsslime install does not pull. Fix:pip install --no-deps mbridgeinslime.Dockerfile.Megatron
validate_argseagerly probes the CUDA device (get_device_capabilityfor--moe-grouped-gemm;get_device_arch_versionfor the TP/CP note), but SLIME runsvalidate_argson the GPU-less Ray driver, so it crashes withFound no NVIDIA driver. Only the MoE + TP/CP>1 recipe reaches these probes. Fix: a self-neutralizing build-time patch guards both probes in SLIME'svalidate_argswrapper only whentorch.cuda.is_available()is False (real GPU actors unchanged); the permanent fix belongs upstream in Megatron.--sglang-enable-ep-moeis a dead flag on SGLang 0.5.12 (the arg was removed; SLIME parses--sglang-*leniently, so it is silently ignored, not rejected) and the recipe omitsCUDA_DEVICE_MAX_CONNECTIONS=1, which Megatron asserts for TP/CP>1. Fix: drop the dead flag (SGLang defaults serve the MoE correctly, verified end-to-end — no replacement flag needed) and addCUDA_DEVICE_MAX_CONNECTIONS=1to the recipe runtime-env.Validation environment (common to the fixes here)
Everything above the cluster layer (image, recipe, model, versions, flags) is the test case as-is; only the cluster layer differs from the documented HyperPod/p5 target. Per-fix results are in each fix's comment.
p5en.48xlarge(NVIDIA H200, 8 GPU/node)slime.Dockerfile(with the fixes in this PR)End-to-end result so far
With the general fixes above, the Qwen3-4B colocated GRPO recipe (
recipe/run_grpo_qwen3_4b.sh, built-in reward) now runs end-to-end. On 2xp5en.48xlarge(16x H200), an image built from this branch'sslime.Dockerfilecompleted 3 full GRPO loops (rollout generation → ref/actor log-probs → Megatron training step → weight sync back to SGLang), with no regressions (libcudart.so.12, numpy-2.x assert, uvicornKeyError,hidden_size Noneall absent). Inter-node NCCL used EFA (NET/OFI Selected provider is efa, fabric is efa-direct (found 16 nics), repeated across all 16 ranks). This exercises the documented quick-start path (4B dense, colocated, in-process reward).remote_rmHTTP reward service (reward_service/,kubernetes/reward-service.yaml): also validated. I built the CPU-only reward image fromreward_service.Dockerfile, ran the service (math_verifybackend) on a non-GPU node, and re-ran the 4B recipe withRM_TYPE=remote_rm/RM_URL=http://slime-reward.<ns>.svc.cluster.local:8000/score. SLIME picked uprm_type=remote_rmand the rollout reached the reward service over HTTP: 250+POST /scorerequests, all200 OK, zero non-200, and zeroremote_rmretry/failure log lines on the SLIME side. The/scorecontract behaved as specified (correct answer -> 1.0, wrong -> 0.0). Environment delta:reward-service.yamltargets a HyperPod CPU instance group vianodeSelector: sagemaker.amazonaws.com/instance-group-name, which does not exist on plain EKS; I kept the manifest's GPU-exclusionnodeAffinityand dropped the instance-groupnodeSelectorso the pod lands on a non-GPU node, and added animagePullSecretsfor the cross-account ECR. These are cluster-layer adaptations; the reward service code and the recipe wiring are unchanged.30B MoE disaggregated recipe (
recipe/run_grpo_qwen3_30b_a3b.sh): also validated, on the same 2xp5en.48xlarge(16x H200), in a disaggregated topology (12 train GPUs at TP=2/EP=2/CP=2, 4 rollout GPUs). After the three 30B-only fixes above (mbridge, GPU-less driver validate guard, drop dead--sglang-enable-ep-moe+ addCUDA_DEVICE_MAX_CONNECTIONS=1), the recipe reached a full GRPO loop:Final collected 96 samples from rollout to train→Timer train start→Timer ref_log_probs end (27.4s)→Timer actor_train end (55.4s)→Timer train end (93.2s)→Timer update_weights end (12.1s), with zero occurrences ofFound no NVIDIA driver,not divisible by micro batch size,libcudart.so.12, or400 Bad Request(the MoE online weight sync succeeded on SGLang's defaults). This is the disaggregated + MoE path, complementary to the 4B colocated run above.Open questions for the maintainers
[Don't merge]epic PR the shape you prefer, or would you rather each fix be its own separate PR?torch_memory_saver.so, is a thin test-case-side fallback acceptable until the CUDA-major selection lands upstream in SLIME?Checklist
mainbranch.latest).