[Don't Merge] Add miles grpo test case - #1225
Draft
littlemex wants to merge 15 commits into
Draft
Conversation
Adds a PyTorch test case running GRPO post-training with miles (a direct fork of SLIME targeting CUDA 13 / NVIDIA Blackwell sm_103) on Amazon EKS / SageMaker HyperPod, mirroring the sibling 3.test_cases/pytorch/slime test case so users can run the same GRPO workloads on the miles stack (SGLang rollout + Megatron-LM training via Ray). Verified on p5en.48xlarge (H200): Qwen3-4B dense colocated GRPO on 1 and 2 nodes (16 GPU over EFA, 3 cycles), and Qwen3-30B-A3B MoE colocated on 2 nodes (16 GPU) with --use-distributed-optimizer + triton MoE runner. 2-node EFA NCCL busbw 190-257 GB/s. Components that mirror slime but were not run on this hardware (head-on-CPU manifest, 30B disaggregated which needs B300 HBM, the CPU reward service) are marked UNVERIFIED. Three miles-base pitfalls (CUDA compat shadowing / libcuda in the SGLang subprocess / Ray driver on the head) are fixed in miles.Dockerfile + the manifests and documented in docs/PORT_NOTES.md. Image tags pinned; no secrets or account-specific values committed.
- env_vars: SAVE_INTERVAL default 1000 (> NUM_ROLLOUT) so the shipped default does not hit the known save_model() pickle-truncation crash; comment explains why not to lower it. - 4B recipe: add TP/PP/CP/EP to the required-variable validation (parity with 30B recipe); make --sglang-mem-fraction-static a variable (SGLANG_MEM_FRACTION:-0.8) like the 30B recipe; align the EXTRA_TRAIN_ARGS comment with the shipped --use-tensorboard default. - miles.Dockerfile: correct the weight-sync note (miles uses Ray actor methods, not an SGLang-fork HTTP path); mark NCCL_VERSION/AWS_OFI_NCCL_VERSION as reference-only ARGs (nothing installs from them; they come from the base image / EFA installer). - raycluster.yaml: head /dev/shm sizeLimit 64Gi -> 8Gi to stay within the 16Gi memory limit (medium:Memory counts against it); comment added.
…reshness CI The 30B MoE configuration was written up as "Verified" on the strength of the job completing. It does complete: the actor spans all 16 GPU without OOM and every rollout cycle finishes. But its generation is degenerate -- across four runs the model loops and never reaches an answer (repetition_frac 0.48-0.70 vs 0.0 on dense 4B, reward 0.0 vs 0.42-0.55). Four candidate causes were tested and refuted: response-length cap, trainer/rollout weight divergence, sampling temperature, and checkpoint corruption. Root cause unresolved. Split the conflated claim into "completes without OOM/crash" (Verified) and "produces a usable trained model" (Known Issue), and rescope the same overclaim everywhere it appears: README table, quick-start note, launch-command comment, parallelism section, recipe STATUS banner and runtime warning, env_vars ALTERNATE block, and the RayCluster manifest comment. Also: re-pin MILES_BASE_TAG (the old tag is now 404 on Docker Hub), parameterise the buildkit Job namespace, drop a docs reference to a path that does not exist in this repo, annotate env vars that are exported but never read, and add a GPU-free workflow that fails when a pinned external artifact stops resolving. That workflow deliberately does not build or train -- no runner here has 8-16 H200s -- and its header says so.
… parameterising
evaluate.sh parsed its CLI flags into shell variables but never exported them, and the
Python it runs is inside a quoted heredoc, so it read the environment and fell back to
defaults. Passing --num-samples or --temperature therefore changed nothing while the run
still wrote a results file that looked authoritative. Export the parsed values, and give
the server port a single definition instead of three copies.
The buildkit Job wrote its push target as a literal account/region/tag string built from
variables the pod has no way to expand, so applying it without envsubst pushed to a
malformed registry path; it also pinned a tag that nothing else tracks. Use ${FULL_IMAGE},
which is the same value the RayCluster already consumes. Give that Job and the data-prep
pod the ${NAMESPACE} and ${FSX_CLAIM} treatment the other manifests already had.
Both recipes dereference EVAL_DATA but left it out of the required-variable check, which is
exactly the opaque mid-construction failure that check exists to prevent.
Also: document that the checkpoint conversion has to run inside a miles-image pod with the
shared filesystem mounted, since the script needs both and the quick start had it before
any such pod exists; and drop a stale NUM_NODES line the PR body already claimed was gone.
The conversion step needs the miles image, the shared filesystem, and real GPUs when
--num-gpus is passed. The quick start put it before any such pod exists, and the earlier fix
for that pointed at a path inside the image that does not exist. Copy the script into a Ray
GPU worker and run it there. Not the head: the head runs with num-gpus 0 and
NVIDIA_VISIBLE_DEVICES=void, which is the environment docs/PORT_NOTES.md already documents
as one where importing CUDA fails.
evaluate.sh passed --log-level WARN to uvicorn, which only accepts lowercase names and
raises KeyError on that value. The recipes already used "warning"; this was the same bug
left behind in one more place.
The data-prep pod carried HF_TOKEN as an empty literal in the manifest, inviting a reader to
paste a token into a file they might commit. Read it from the same hf-token Secret the
RayCluster already uses, and apply the manifest through envsubst now that it references
${NAMESPACE} and ${FSX_CLAIM}.
…in CI radixark publishes dev-* as mutable snapshots and has already deleted a tag this file pinned, so a dated tag gives neither reproducibility nor availability. Resolve the tag to its digest and build FROM that, keeping the tag alongside so a reader can still tell which snapshot it is. The Dockerfile carries the curl invocation for resolving a new one. The freshness workflow was checking the tag, which the build no longer uses. It now checks the digest, and separately warns when the documented tag has drifted off it -- expected for a mutable tag, harmless for the build, but it means the comment has gone stale. Verified against the live registry: the digest resolves 200 and currently matches the tag.
Records what was run, on what, and what came back: the submission command, the Ray job id, the flags the job actually received (read out of the job log, not the recipe source), and the metric values read back from the trainer's TensorBoard event files. The run documented here is the reader's path: the shipped env example with only the cluster-specific paths filled in, launched with the shipped recipe unmodified. It came back SUCCEEDED with reward 0.516, repetition 0.0, truncation 0.484, and a finite grad norm. The log also states why reward and repetition are the metrics to check rather than the exit code, with the 30B MoE numbers next to the dense ones: both exit SUCCEEDED, and only one of them trained anything. That distinction is the reason the verification table now has separate rows for completing and for producing a usable model, and an earlier version of the table got it wrong for want of exactly this log.
…d stale prose
- kubernetes/raycluster.yaml: parameterize EFA count (${EFA_PER_NODE}) and worker
replicas (${ACTOR_NUM_NODES}). Both were hardcoded (8 of 15 allocatable EFA devices;
replicas: 1) with a comment telling the reader to edit them for 2-node runs. A
partial EFA request breaks point-to-point NCCL (PP/EP) while passing an all-reduce
smoke test; a missed replica edit leaves the job waiting on placement it cannot get.
- env_vars.colocated.example: add EFA_PER_NODE=15, matching the node's actual
allocatable count (README said 16; node reports 15).
- env_vars.disaggregated.example: define REWARD_IMAGE/REWARD_TAG, referenced by
kubernetes/reward-service.yaml via envsubst but never defined -- rendered
`image: ""` and the Deployment was rejected. Also correct the Spot-tolerance
comment, which contradicted the README's own "miles lacks slime's retry" note.
- scripts/evaluate.sh: fix three scoring bugs found by independent review and
confirmed against the real aime-2024.jsonl and a synthetic pass@k case --
pass@k keyed on a JSONL "idx" field the data does not have (all 30 prompts
collapsed onto key 0; reproduced: reports 1.0000 where the true value is 0.0333);
\boxed{} extraction stopped at the first brace, misreading nested-brace answers
like \frac{1}{2}; unbounded concurrency under a 300s timeout counted most
requests as wrong answers once MAX_TOKENS made generation slower than the queue.
Added bounded concurrency, a generation-sized timeout, and a non-zero exit above
a 5% error rate.
- README.md: fix the head-vs-worker contradiction in step 4 (prose said "exec into
its head", code block correctly targeted a worker); make step 8 explicit that the
shipped SAVE_INTERVAL/NUM_ROLLOUT produce no checkpoint at all; replace fragile
`ray job list | head -3` parsing and an undefined ${HEAD_POD} in the monitoring
block; make the hf-token Secret's required-ness match the manifest instead of the
"skip if public" prose; align the EFA hardware row with allocatable (15, not 16);
document ACTOR_NUM_NODES driving both topology and replica count, drop the two
variables (NUM_NODES/GPUS_PER_NODE) the doc showed that env_vars never defines.
- docs/EFA_2NODE.md: correct prose that called a partial EFA request "a throughput
improvement, not a correctness issue" -- it is a correctness issue on P2P paths.
- docs/VERIFICATION_LOG.md: record this round -- what was found by reading rather
than re-running, and what was re-verified (all four manifests rendered and
validated with kubectl apply --dry-run=server against a live EKS cluster at both
ACTOR_NUM_NODES=1 and =2; the boxed{} extractor and pass@k indexing checked
against the real aime-2024.jsonl; every script parses with bash -n).
…al server on abort Four defects from an adversarial review pass, all cases of a stated setting or claim not matching what the code does: - scripts/evaluate.sh leaked the SGLang server. The high-error-rate abort added in the previous commit exits non-zero from inside the heredoc, and `set -e` then skips the cleanup `kill` at the end of the script -- leaving a multi-GPU server holding the node. Cleanup now runs from an EXIT trap. Verified with a standalone reproduction: the trap fires on a non-zero exit, the child is reaped, and the exit code is preserved. - Both recipes validate COLOCATE, print it, and branch on it, then pass --colocate unconditionally. COLOCATE=false printed "Colocated: false" (4B) or a warning (30B) and ran the colocated layout anyway. Neither recipe builds a disaggregated actor layout, so they now refuse rather than silently disagree with the env. - env_vars.disaggregated.example twice pointed the reader at a README "Prerequisites" section for an `aws sagemaker update-cluster` command. No such section exists, and the command is HyperPod-specific while this test case targets EKS. Replaced with what is actually true: the pool must exist, and how you create it depends on your cluster. - The 2-node "3 rollout cycles / Verified" row cited only docs/EFA_2NODE.md, which measures fabric bandwidth and says nothing about a training loop. The GRPO metrics behind that row (3 cycles SUCCEEDED in 676s, raw_reward 0.477/0.523/0.492) are now in the verification log, with an explicit note that three optimizer steps evidence a closed loop, not convergence.
A review pass on comment quality found that the development history had leaked out of
docs/VERIFICATION_LOG.md and into the recipes, manifests and env examples: comments narrating
what an earlier revision got wrong, what the original value was, and what a since-removed flag
used to do. That is git history and log material. What belongs next to the code is why the
current shape is correct.
- scripts/evaluate.sh: fixed a regression from the previous commit. Moving cleanup into an
EXIT trap left a bare `wait ${SGLANG_PID}` behind in Step 3, so the SUCCESS path blocked
forever on a server nothing had killed yet. Verified both paths now: success exits 0
immediately, the high-error abort still exits 2, and neither leaks the server.
- Trimmed the post-mortem prose in evaluate.sh's extract_boxed docstring, the semaphore
comment and the pass@k comment to the current-tense reason.
- Collapsed the 6-line uvicorn log-level explanation duplicated verbatim in both recipes to
one line each, with the detail left where it already lived in docs/PORT_NOTES.md.
- Replaced the removed-flag history in the 30B recipe with the reusable trap it illustrates:
miles registers --sglang-* from SGLang's live ServerArgs with ignore_unknown_args, so a
flag SGLang has dropped is accepted silently and does nothing.
- Cut the COLOCATE guard from 16 lines making the same point three times to 2 comment lines
and a 3-line error.
- raycluster.yaml, env_vars.*.example, README: removed the duplicated head-vs-worker
explanation, the four-way-duplicated B300 portability discussion, the "do not parse ray job
list" warning against a usage nobody suggested, and a stale NOTE claiming worker replicas are
not driven by a variable (they are, since ACTOR_NUM_NODES was wired up).
- Rewrote "SPOT IS NOT FREE HERE", which read as a contradiction rather than a warning, and
disambiguated "0.85 (the original hardcode)", which named no identifiable original.
- Unified British spellings to American to match the rest of the repository.
The RayCluster mounts the hf-token Secret without optional: true, so the Ray pods never start without it and report CreateContainerConfigError rather than anything mentioning HuggingFace. That made it a real prerequisite, but the Prerequisites list did not carry it. Adds it as item 8, and gives step 1 the empty-value form for public models, since the pods only need the key to exist. Verified against a live cluster that --from-literal=HF_TOKEN= produces the HF_TOKEN key that secretKeyRef resolves.
…e one it was copied from reward-service.yaml kept slime-reward for the Deployment, the Service and every app label, and env_vars.disaggregated.example pointed RM_URL at that DNS name -- left over from the sibling test case this file was adapted from, while raycluster.yaml had been renamed to miles-ray. Beyond the inconsistency, both test cases claim the same Service name, so deploying the disaggregated overlays of miles and slime into one namespace would collide on miles-reward's predecessor rather than coexisting. Renamed to miles-reward throughout, matching REWARD_IMAGE_NAME which was already miles-reward. Verified that the rendered Service name and the RM_URL host now agree and that the manifest is accepted by a live cluster.
… step 4
Step 8 gave the ls check and the megatron2hf conversion as bare shell commands with no
kubectl exec/cp wrapper, unlike step 4's identical script which explicitly runs it on a GPU
worker pod because the script and /fsx only exist there. A reader following step 8 literally
from their workstation hits convert_checkpoint.sh's cd ${SLIME_DIR} (default /root/miles)
failing under set -euo pipefail, since that path only exists inside the miles image.
COLOCATE named a choice the recipe could not make. --colocate was passed unconditionally, so the recipe refused COLOCATE=false rather than honor it, and the disaggregated layout was unreachable from here. That layout is not a variant of the same run. miles selects the weight sync implementation from --colocate alone: present gives UpdateWeightFromTensor, moving weights by CUDA IPC between processes that share a device; absent gives UpdateWeightFromDistributed, broadcasting them over NCCL, across EFA when the actor and rollout pools sit on different nodes. A user comparing the two placements needs both, and until now this recipe offered one. Emit --colocate from an array so the disaggregated case contributes zero argv tokens, and enforce the sizing rule each layout implies: colocated shares devices, so ROLLOUT_NUM_GPUS must equal the actor GPU count; disaggregated needs separate pools, so actor + rollout must fit CLUSTER_GPUS when the caller states it. Over-subscribing the latter does not fail at submit time -- Ray waits on a placement group that never becomes ready, which reads as a hang -- so refuse while the numbers are still on screen. Reject COLOCATE values other than the two spellings. "True", "1" and "yes" would have fallen through to the disaggregated branch and chosen a layout nobody asked for, which is the failure the previous refusal existed to prevent, pointing the other way. Likewise type-check the GPU counts first: -ne against a non-numeric value errors inside the `if`, which reads as false, so the layout check passes without having run. --sglang-mem-fraction-static becomes SGLANG_MEM_FRACTION, defaulting to the 0.8 this recipe has always used. Disaggregated engines do not share their devices with the trainer and can take more, and a study that varies it needs the knob; the value is resolved once and printed next to the layout so the banner cannot describe a run the argv did not request. A colocated invocation renders a byte-identical argv: 104 tokens, unchanged.
…le "disaggregated" The Known Issue said the root cause was unresolved and that confirming it would mean isolating the MoE path outside the GRPO loop. Did that: served the same checkpoint from SGLang directly, no miles, no Megatron, no GRPO. It reproduces, so nothing in training is involved. Varying one axis at a time locates it in expert parallelism. EP=1 is clean at TP 1, 4 and 8; EP>1 degenerates at every TP tried, and the rate grows with EP (0.594 at EP=2, 0.844 at EP=4 with TP held at 8). The model's own recommended sampling still gives 0.875 and the default auto runner backend still gives 0.844, so neither the temperature nor the triton selection was responsible, which is worth stating because both were plausible and both were wrong. There is no way out inside this test case. EP=1 fixes generation and cannot train: each rank then holds all 128 experts and asks for 108.76 GiB against 139.80 GiB. Raising TP from 8 to 16 moves that by 0.3%, since tensor parallelism does not shard expert weights. The README now says so rather than leaving a reader to discover it. Separately, the README used "disaggregated" for two unrelated things -- splitting the actor and rollout GPU pools, and moving reward scoring to CPU -- and the file named env_vars.disaggregated.example is the second one. Said which is which, and added the 4B COLOCATE=false row to the verification table as UNVERIFIED: the recipe builds that layout now, but only the rendered argv was checked, not a completed run.
Contributor
Author
TODO before merge
|
KeitaW
marked this pull request as draft
August 5, 2026 01:39
Collaborator
|
Thank you! Feel free to open the PR once you clear the remaining TODO items. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
Relates to #1224
Adds a test case for GRPO post-training with miles on Amazon EKS with EFA, as a companion to the existing
3.test_cases/pytorch/slimecase. miles is a fork of slime built for CUDA 13 and Blackwell: SGLang for rollout, Megatron-LM for training, weight sync over Ray actor methods.The reason to carry both is that the two forks diverge exactly where porting is expensive. miles ships on an
nvidia/cudabase rather than NGC, and that difference alone produces three failures that do not occur in the slime case and are not obvious from either repository's documentation.docs/PORT_NOTES.mdrecords them so the next person does not rediscover them.Relates to the slime test case's layout and env-var conventions, which this case mirrors deliberately so the two can be read side by side.
Changes
miles.Dockerfile— takes the upstreamradixark/milesimage (PyTorch 2.11 / CUDA 13.0.1, SGLang and Megatron-LM prebuilt) and adds only the AWS EFA layer: GDRCopy, the EFA installer, and NCCL/EFA runtime defaults. Rebuilding that stack on an NGC base is not viable because of wheel ABI mismatch, so the image is taken as-is. Pinned by digest — see Notes below.kubernetes/— KubeRayRayCluster, an in-cluster BuildKitJobfor clusters with no local Docker daemon, a data-prepPod, and the reward-serviceDeploymentfor the disaggregated overlay. All rendered withenvsubst; every placeholder is defined in anenv_vars.*.example.recipe/run_grpo_qwen3_4b.sh,recipe/run_grpo_qwen3_30b_a3b.sh— dense and MoE GRPO recipes, submitting throughrecipe/launcher/grpo_launch.sh. The 4B recipe builds either GPU layout fromCOLOCATE; the MoE recipe is colocated only, and says so rather than accepting the flag and ignoring it.env_vars.colocated.example,env_vars.disaggregated.example— all tunables, with the reasoning attached to the values that are not free choices (EFA count, memory fractions, the MoE flag set). No account, region, cluster or namespace values: region and account are derived from the caller's AWS credentials at source time.scripts/convert_checkpoint.sh,scripts/evaluate.sh— HF/Megatron conversion and AIME-2024 evaluation.reward_service/+reward_service.Dockerfile— CPU reward service for the disaggregated overlay. Present, not exercised; marked as such.docs/PORT_NOTES.md— the threenvidia/cuda-vs-NGC failures and the porting decisions.docs/EFA_2NODE.md— the 2-node EFA result, the security-group prerequisite that has to be right for it, and the failure signature when it is not.docs/VERIFICATION_LOG.md— what was run, on what, and what came back, read out of the trainer's TensorBoard event files rather than retyped..github/workflows/miles-pin-freshness.yml— see Notes below.Test Plan
Environment:
/fsxp5en.48xlarge(H200 141GB x8)Test commands:
Test Results
Single-node dense baseline, shipped recipe unmodified, Ray job
SUCCEEDED:rollout/raw_rewardrollout/repetition_fracrollout/truncated_ratiotrain/grad_normtrain/train_rollout_logprob_abs_diffperf/step_timeTwo nodes (16 GPU) over EFA, 3 rollout cycles,
SUCCEEDEDin 676s:rollout/raw_rewardactor_train_tflopsperf/step_timeThree optimizer steps evidence a closed loop across the node boundary, not convergence, and the
log says so. Fabric measured separately: NCCL all_reduce busbw 190-257 GB/s with
efa-directand GPUDirect RDMA, no TCP fallback.
What is not verified, listed in the README as its own rows rather than left implicit:
across 4 runs generation falls into a repetition loop and reward stays 0.0
(
repetition_frac0.48-0.70,truncated0.97-0.99, against 0.0 / 0.43-0.58 on denseQwen3-4B). Root cause unresolved; four candidate causes ruled out by measurement. Filed
separately upstream. The
ALTERNATEblock that configures it is commented out with the issuestated inline, so it cannot be enabled without reading it.
layout was run.
pinned base image) but not run.
raycluster.yamlwith the head on a CPU node. That is the shape that ships, butthe validation runs used a head-on-GPU overlay because the validation cluster had no
large-disk CPU node at the time.
save_model()fails inside Megatron's distributed-checkpoint save, soSAVE_INTERVALships above the step count and no save is triggered. This means a default runproduces no checkpoint at all, which the README now states at the top of the conversion
step rather than leaving the reader to discover it after hours of training.
Directory Structure
Mirrors the sibling
3.test_cases/pytorch/slimecase so the two read side by side:Checklist
mainbranch.latest) — base image bydigest, EFA installer and GDRCopy by version, Python packages pinned.
slimecase.