From 8243f1fef526565d0eb788cca3b1381252d1ce57 Mon Sep 17 00:00:00 2001 From: littlemex Date: Tue, 18 Aug 2026 19:58:49 +0900 Subject: [PATCH 1/7] feat(miles): container image and in-cluster BuildKit build Take the upstream radixark/miles image (PyTorch 2.11 / CUDA 13, SGLang and Megatron-LM prebuilt) by explicit digest and add only the AWS EFA layer. Rebuilding on an NGC base is not viable (wheel ABI mismatch), so the base is consumed as-is; no floating :latest. A BuildKit Job builds and pushes the image in-cluster for environments with no local Docker daemon. --- 3.test_cases/pytorch/miles/.gitignore | 9 ++ .../miles/kubernetes/buildkit-job.yaml | 55 +++++++ 3.test_cases/pytorch/miles/miles.Dockerfile | 136 ++++++++++++++++++ 3.test_cases/pytorch/miles/requirements.txt | 43 ++++++ 4 files changed, 243 insertions(+) create mode 100644 3.test_cases/pytorch/miles/.gitignore create mode 100644 3.test_cases/pytorch/miles/kubernetes/buildkit-job.yaml create mode 100644 3.test_cases/pytorch/miles/miles.Dockerfile create mode 100644 3.test_cases/pytorch/miles/requirements.txt diff --git a/3.test_cases/pytorch/miles/.gitignore b/3.test_cases/pytorch/miles/.gitignore new file mode 100644 index 000000000..27bc9b754 --- /dev/null +++ b/3.test_cases/pytorch/miles/.gitignore @@ -0,0 +1,9 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 + +# Local environment files with filled-in secrets/values -- never commit +env_vars +env_vars.disaggregated + +# Eval / training artifacts +eval_results/ diff --git a/3.test_cases/pytorch/miles/kubernetes/buildkit-job.yaml b/3.test_cases/pytorch/miles/kubernetes/buildkit-job.yaml new file mode 100644 index 000000000..5c9ec9d5a --- /dev/null +++ b/3.test_cases/pytorch/miles/kubernetes/buildkit-job.yaml @@ -0,0 +1,55 @@ +apiVersion: batch/v1 +kind: Job +metadata: + name: miles-efa-build + namespace: ${NAMESPACE} +spec: + ttlSecondsAfterFinished: 172800 + backoffLimit: 0 + template: + spec: + restartPolicy: Never + # Schedule the build on any node with ample ephemeral-storage (the ~18GB + # base image plus extraction needs ~120GB+ free). buildkit is a pure CPU/IO + # workload and requests NO GPUs, so it does not contend with training. Pin + # it to a node with large local disk via a nodeSelector matching your + # cluster (e.g. a CPU nodegroup with a >=200GB gp3 root, or a GPU node with + # NVMe scratch). Example: + # nodeSelector: + # node.kubernetes.io/instance-type: + containers: + - name: buildkit + image: moby/buildkit:v0.18.2 + securityContext: { privileged: true } + command: + - sh + - -c + - | + set -eux + buildkitd & + for i in $(seq 1 30); do buildctl debug workers >/dev/null 2>&1 && break; sleep 2; done + buildctl build \ + --frontend dockerfile.v0 \ + --local context=/workspace \ + --local dockerfile=/workspace \ + --output type=image,name=${FULL_IMAGE},push=true \ + --progress=plain + env: + - { name: DOCKER_CONFIG, value: /root/.docker } + resources: + requests: { cpu: "16", memory: "64Gi", ephemeral-storage: "120Gi" } + limits: { cpu: "40", memory: "160Gi", ephemeral-storage: "400Gi" } + volumeMounts: + - { name: ctx, mountPath: /workspace } + - { name: docker-config, mountPath: /root/.docker } + volumes: + - name: ctx + configMap: + name: miles-build-context + items: + - { key: Dockerfile, path: Dockerfile } + - name: docker-config + secret: + secretName: ecr-miles-push + items: + - { key: .dockerconfigjson, path: config.json } diff --git a/3.test_cases/pytorch/miles/miles.Dockerfile b/3.test_cases/pytorch/miles/miles.Dockerfile new file mode 100644 index 000000000..b354457ba --- /dev/null +++ b/3.test_cases/pytorch/miles/miles.Dockerfile @@ -0,0 +1,136 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +# ============================================================================ +# miles (radixark/miles) + EFA image for Amazon EKS +# ============================================================================ +# Strategy C (see docs/PORT_NOTES.md): the miles upstream image already bundles +# a matched PyTorch 2.11 / CUDA 13.0.1 stack, the sglang-miles fork, the radixark +# Megatron-LM fork, prebuilt flash_attn / TE / apex wheels, and the B300 sm_103 TE +# FA2 whitelist patch. (Unlike an SGLang-fork HTTP weight-sync path, miles drives +# weight sync through Ray actor methods on the rollout engine -- see README / +# docs/PORT_NOTES.md.) Rebuilding that stack on an NGC base is infeasible (wheel +# ABI mismatch), so we take the miles image as-is and add ONLY the AWS EFA stack. +# +# Pin the base by DIGEST, not by tag. radixark publishes dev-* as mutable snapshots and +# has already deleted a tag this file previously pinned, so a dated tag alone gives neither +# reproducibility nor availability. The tag is kept alongside for readability; the digest is +# what the build resolves. cu13 targets CUDA 13.0.1 (B300 sm_103 / H200 sm_90). +# +# To move to a newer base, resolve its digest first: +# TOK=$(curl -s "https://auth.docker.io/token?service=registry.docker.io\ +#&scope=repository:radixark/miles:pull" | python3 -c "import sys,json;print(json.load(sys.stdin)['token'])") +# curl -sI -H "Authorization: Bearer $TOK" \ +# -H "Accept: application/vnd.oci.image.index.v1+json" \ +# https://registry-1.docker.io/v2/radixark/miles/manifests/ | grep -i docker-content-digest +# ============================================================================ +ARG MILES_BASE_TAG=dev-202607310056 +ARG MILES_BASE_DIGEST=sha256:ca0bb593dd6f4011b444f64d478b72c213e4c70421f4d7f94e593a709562429e +FROM radixark/miles@${MILES_BASE_DIGEST} + +ARG GDRCOPY_VERSION=v2.5.2 +ARG EFA_INSTALLER_VERSION=1.48.0 +# NOTE: these two ARGs are REFERENCE VALUES ONLY -- nothing below installs from them. +# NCCL comes from the miles base image and aws-ofi-nccl is the plugin bundled with the +# EFA installer above; the values here record the versions those components are expected +# to provide (and satisfy the repo's version-gate grep). If you bump the base image or the +# EFA installer, update these to match what those actually ship, or drop them. +ARG NCCL_VERSION=v2.30.4-1 +ARG AWS_OFI_NCCL_VERSION=v1.19.0 + +ENV DEBIAN_FRONTEND=noninteractive +ENV TZ=Etc/UTC + +###################### +# Remove IB libverbs (replaced by the EFA installer below). The miles base is +# lmsysorg/sglang (nvidia/cuda:13.0.1), which ships InfiniBand userspace libs; +# strip them so the EFA installer's libfabric/libibverbs take precedence. +###################### +RUN apt-get update -y \ + && apt-get remove -y --allow-change-held-packages \ + ibverbs-utils libibverbs-dev libibverbs1 libmlx5-1 || true + +RUN DEBIAN_FRONTEND=noninteractive apt-get install -y --allow-unauthenticated \ + autoconf automake build-essential cmake curl gcc gdb git jq kmod libtool \ + openssh-client openssh-server vim \ + && apt-get autoremove -y + +# Permissive SSH for in-cluster MPI/NCCL bootstrap. Port 22 must NOT be exposed +# outside the cluster via a Service. +RUN mkdir -p /var/run/sshd && \ + sed -i 's/[ #]\(.*StrictHostKeyChecking \).*/ \1no/g' /etc/ssh/ssh_config && \ + echo " UserKnownHostsFile /dev/null" >> /etc/ssh/ssh_config && \ + sed -i 's/#\(StrictModes \).*/\1no/g' /etc/ssh/sshd_config && \ + sed -i 's/#PermitRootLogin prohibit-password/PermitRootLogin yes/' /etc/ssh/sshd_config +RUN rm -rf /root/.ssh/ && mkdir -p /root/.ssh/ \ + && ssh-keygen -q -t rsa -N '' -f /root/.ssh/id_rsa \ + && cp /root/.ssh/id_rsa.pub /root/.ssh/authorized_keys \ + && printf "Host *\n StrictHostKeyChecking no\n" >> /root/.ssh/config + +################################################# +## NVIDIA GDRCopy (GPUDirect RDMA copy library) +RUN git clone -b ${GDRCOPY_VERSION} https://github.com/NVIDIA/gdrcopy.git /tmp/gdrcopy \ + && cd /tmp/gdrcopy \ + && make prefix=/opt/gdrcopy install \ + && rm -rf /tmp/gdrcopy +# Physically remove the CUDA forward-compat driver. The slime (NGC) image relied +# on it, but NGC's entrypoint enables compat only when compat >= host driver; +# the nvidia/cuda base under miles has no such guard. This base ships an OLDER +# compat libcuda (580.82.07) than the driver on GPU-operator nodes (e.g. +# 580.159.03), and forward-compat requires compat >= host driver. If anything +# resolves the stale compat libcuda, torch.cuda dies with "Error 803: unsupported +# display driver / cuda driver combination". The host driver already supports +# this image's CUDA 13.0 toolkit, so we delete compat outright (belt-and-suspenders +# vs merely dropping it from LD_LIBRARY_PATH, which a future env edit could undo). +RUN rm -rf /usr/local/cuda/compat /usr/local/cuda-*/compat +ENV LD_LIBRARY_PATH=/opt/gdrcopy/lib:$LD_LIBRARY_PATH +ENV LIBRARY_PATH=/opt/gdrcopy/lib:$LIBRARY_PATH +ENV CPATH=/opt/gdrcopy/include:${CPATH:-} +ENV PATH=/opt/gdrcopy/bin:$PATH + +################################################# +## AWS EFA installer (libfabric + aws-ofi-nccl plugin, no kmod in-container) +RUN cd $HOME \ + && curl -O https://efa-installer.amazonaws.com/aws-efa-installer-${EFA_INSTALLER_VERSION}.tar.gz \ + && tar -xf $HOME/aws-efa-installer-${EFA_INSTALLER_VERSION}.tar.gz \ + && cd aws-efa-installer \ + && ./efa_installer.sh -y -g -d --skip-kmod --skip-limit-conf --no-verify \ + && rm -rf $HOME/aws-efa-installer $HOME/aws-efa-installer-*.tar.gz \ + && rm -rf /var/lib/apt/lists/* + +# EFA / aws-ofi-nccl library paths. Placed AFTER the gdrcopy block so the EFA +# libfabric resolves ahead of any base-image InfiniBand remnants. The host-driver +# dirs are appended at the END: the GPU operator / nvidia-container-toolkit injects +# libcuda.so.1 into /usr/lib64 (AL2023/Bottlerocket) or /usr/lib/x86_64-linux-gnu +# (Ubuntu). torch resolves libcuda via ld.so.cache, but SGLang's server subprocess +# (and Triton/cuda-python loaders) scan LD_LIBRARY_PATH directly and otherwise fail +# with "ImportError: libcuda.so.1: cannot open shared object file". Appending (not +# prepending) means these dirs are only consulted for libraries nothing else +# resolves, so they never shadow the EFA/CUDA-runtime libs above. Non-existent +# dirs in LD_LIBRARY_PATH are harmless. +ENV LD_LIBRARY_PATH=/opt/amazon/efa/lib:/opt/amazon/aws-ofi-nccl/lib:/opt/amazon/ofi-nccl/lib:$LD_LIBRARY_PATH:/usr/lib64:/usr/lib/x86_64-linux-gnu +ENV PATH=/opt/amazon/efa/bin:$PATH +# glibc-path insurance: ensure the host driver dir is in ld.so.cache too. The +# toolkit re-runs ldconfig at container start, so the injected libcuda is picked +# up for the standard glibc resolution path (belt-and-suspenders with the LD above). +RUN echo "/usr/lib64" > /etc/ld.so.conf.d/zz-host-driver.conf \ + && echo "/usr/lib/x86_64-linux-gnu" >> /etc/ld.so.conf.d/zz-host-driver.conf \ + && ldconfig + +##################### +# EFA / NCCL / Ray runtime defaults (match the slime test case). +##################### +ENV FI_PROVIDER="efa" +ENV FI_EFA_USE_DEVICE_RDMA="1" +ENV FI_EFA_FORK_SAFE="1" +ENV NCCL_PROTO="Simple" +ENV NCCL_DEBUG="WARN" +# gloo (Ray's collective for CPU tensors) must bind to the primary NIC. +ENV GLOO_SOCKET_IFNAME="eth0" + +# miles installs itself editable at /root/miles and the Megatron fork at +# /root/Megatron-LM, but does NOT bake PYTHONPATH into the image. The recipe and +# the RayCluster runtime-env set it explicitly; this default covers a bare shell. +ENV PYTHONPATH=/root/miles:/root/Megatron-LM:${PYTHONPATH:-} + +WORKDIR /root/miles +CMD ["/bin/bash"] diff --git a/3.test_cases/pytorch/miles/requirements.txt b/3.test_cases/pytorch/miles/requirements.txt new file mode 100644 index 000000000..ac34fa1f2 --- /dev/null +++ b/3.test_cases/pytorch/miles/requirements.txt @@ -0,0 +1,43 @@ +# STATUS: UNVERIFIED for miles -- inherited from slime. miles bundles its RL deps in the base image (radixark/miles); this file is a reference, not installed by miles.Dockerfile. +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 + +# Reinforcement-learning Python dependencies for the miles 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 +ray[default]==2.55.1 +kernels==0.10.4 +huggingface-hub==1.18.0 +tokenizers==0.22.2 +datasets==4.5.0 +accelerate==1.13.0 +sentencepiece==0.2.1 +protobuf==6.33.5 +math_verify==0.9.0 +qwen_vl_utils==0.0.14 +wandb==0.27.2 +tensorboard==2.20.0 +pylatexenc==2.10 +omegaconf==2.3.0 +aiohttp==3.14.1 From 362a1dea000d60d6a4fd87723fbc0ec8605f9214 Mon Sep 17 00:00:00 2001 From: littlemex Date: Tue, 18 Aug 2026 19:58:50 +0900 Subject: [PATCH 2/7] feat(miles): RayCluster, data prep, and portable env configuration Node placement, EFA device count, and worker replica count are env-driven, not hardcoded. The Ray head co-locates on the GPU pool (CPU_NODE_ROLE defaults to GPU_NODE_ROLE) with a nvidia.com/gpu toleration and runs num-gpus 0: miles control actors import Megatron even at num-gpus 0 and need libcuda, so the head must be on a CUDA-capable node. A karpenter.sh/do-not-disrupt annotation protects the head node from underutilized-consolidation. --- .../pytorch/miles/env_vars.colocated.example | 200 ++++++++++++++++++ .../miles/env_vars.disaggregated.example | 75 +++++++ .../miles/kubernetes/data-prep-pod.yaml | 42 ++++ .../pytorch/miles/kubernetes/raycluster.yaml | 190 +++++++++++++++++ 4 files changed, 507 insertions(+) create mode 100644 3.test_cases/pytorch/miles/env_vars.colocated.example create mode 100644 3.test_cases/pytorch/miles/env_vars.disaggregated.example create mode 100644 3.test_cases/pytorch/miles/kubernetes/data-prep-pod.yaml create mode 100644 3.test_cases/pytorch/miles/kubernetes/raycluster.yaml diff --git a/3.test_cases/pytorch/miles/env_vars.colocated.example b/3.test_cases/pytorch/miles/env_vars.colocated.example new file mode 100644 index 000000000..2ef2e5324 --- /dev/null +++ b/3.test_cases/pytorch/miles/env_vars.colocated.example @@ -0,0 +1,200 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +# ============================================================ +# miles on EKS - Environment (COLOCATED) +# +# Colocated Qwen3-4B GRPO: training (Megatron) and rollout (SGLang) time-share +# the same GPUs. This is the simplest topology and the default reference run. +# +# Copy to env_vars and adjust paths: +# cp env_vars.colocated.example env_vars && vim env_vars +# (HF_TOKEN is provided to the pods via the k8s Secret `hf-token`, not this file.) +# ============================================================ + +# ----- AWS / ECR ----- +# Region and account are derived from your current AWS credentials/config so +# nothing environment-specific is hard-coded. Override AWS_REGION if needed. +export AWS_REGION="${AWS_REGION:-$(aws configure get region)}" +export AWS_ACCOUNT_ID="$(aws sts get-caller-identity --query Account --output text)" +export REGISTRY="${AWS_ACCOUNT_ID}.dkr.ecr.${AWS_REGION}.amazonaws.com/" +export IMAGE="miles-hyperpod" +# Pinned build tag, matching kubernetes/buildkit-job.yaml (avoid ":latest"). +# +# BUMP THIS ON EVERY REBUILD. The Ray pods use imagePullPolicy: IfNotPresent, so pushing a +# new image under a tag a node has already pulled leaves that node running the OLD image, +# with no error and no log line saying so -- and a cluster where some nodes pulled before +# the push and some after is running two different builds in one job. Changing the tag is +# what makes the rollout observable. +export TAG="miles-cuda13-efa-0.1" +export FULL_IMAGE="${REGISTRY}${IMAGE}:${TAG}" + +# ----- HuggingFace ----- +# The pods read HF_TOKEN from the k8s Secret `hf-token` (see raycluster.yaml), so +# it is NOT passed via the Ray runtime-env. Create the Secret once: +# kubectl create secret generic hf-token --from-literal=HF_TOKEN=hf_xxx -n "${NAMESPACE:-default}" + +# ----- Model (Qwen3-4B dense, bf16) ----- +export MODEL_NAME="Qwen/Qwen3-4B" # HF repo id, for the download step + # only; the recipes read MODEL_LOCAL + # and MODEL_DIST, not this +export MODEL_LOCAL="/fsx/models/Qwen3-4B" # SGLang rollout init (--hf-checkpoint) +export MODEL_DIST="/fsx/models/Qwen3-4B_torch_dist" # Megatron training (--ref-load), bf16 +export COLOCATE="true" +export TP_SIZE=1 +export PP_SIZE=1 +export CP_SIZE=1 +export EP_SIZE=1 +export ROLLOUT_NUM_GPUS=8 +export ROLLOUT_GPUS_PER_ENGINE=1 +export ACTOR_NUM_NODES=1 +export ACTOR_GPUS_PER_NODE=8 +export MAX_TOKENS_PER_GPU=8192 +export ROLLOUT_MAX_RESPONSE_LEN=8192 +export MODEL_SCRIPT="qwen3-4B.sh" + +# ----- Cluster ----- +export NAMESPACE="default" +export FSX_CLAIM="fsx-claim" +# Node/GPU topology comes from ACTOR_NUM_NODES / ACTOR_GPUS_PER_NODE / ROLLOUT_NUM_GPUS above; +# ACTOR_NUM_NODES also drives the RayCluster's worker replica count, so the manifest and the +# recipe cannot disagree. +# +# GPU_NODE_ROLE is the `node-role` label on your GPU pool, substituted into +# kubernetes/raycluster.yaml's nodeSelector. Only this value is GPU-generation-specific; see +# the README's Verification Status for what has and has not been run on which hardware. +# Node placement is a label KEY:VALUE pair, and BOTH are cluster-specific. The manifest +# schedules workers on ${GPU_NODE_LABEL_KEY}:${GPU_NODE_ROLE} and the head on +# ${CPU_NODE_LABEL_KEY}:${CPU_NODE_ROLE}. The defaults below (key "node-role") match the +# awsome-distributed-ai Terraform EKS reference; a generic EKS or HyperPod cluster has NO +# "node-role" label. Find yours with `kubectl get nodes --show-labels` -- e.g. the well-known +# `node.kubernetes.io/instance-type` (value "p5en.48xlarge"), an EKS managed-nodegroup label +# `eks.amazonaws.com/nodegroup`, or a HyperPod `sagemaker.amazonaws.com/instance-group-name`. +# A wrong key OR value leaves pods stuck Pending with FailedScheduling, not an obvious error. +export GPU_NODE_LABEL_KEY="node-role" +export GPU_NODE_ROLE="gpu-p5en" # e.g. "gpu-b300" on a B300 pool +# The Ray head must land on a node that has the CUDA driver (libcuda.so.1): miles's control +# actors import Megatron/transformer_engine at startup even though they request num-gpus 0, +# and that import hard-fails on a GPU-less node. So the head defaults to the SAME pool as the +# workers (it runs num-gpus 0 and consumes no GPU, only CPU/disk); the manifest carries a +# toleration for the GPU pool's taint. Do NOT point these at a CPU-only pool. +export CPU_NODE_LABEL_KEY="${GPU_NODE_LABEL_KEY}" +export CPU_NODE_ROLE="${GPU_NODE_ROLE}" + +# WORKER_REPLICAS: number of GPU worker nodes the RayCluster launches. COLOCATE=true shares one +# pool, so this equals ACTOR_NUM_NODES. COLOCATE=false uses SEPARATE actor and rollout GPUs, so +# it must cover both: ceil((ACTOR_NUM_NODES*ACTOR_GPUS_PER_NODE + ROLLOUT_NUM_GPUS)/8). Setting +# it to ACTOR_NUM_NODES for a disaggregated run under-sizes the cluster and the job hangs on +# placement. Derived here so the two layouts cannot silently disagree with the manifest: +if [ "${COLOCATE}" = "true" ]; then + export WORKER_REPLICAS="${ACTOR_NUM_NODES}" +else + export WORKER_REPLICAS=$(( ( ACTOR_NUM_NODES*ACTOR_GPUS_PER_NODE + ROLLOUT_NUM_GPUS + 7 ) / 8 )) +fi +# EFA_PER_NODE: EFA devices each worker pod requests, substituted into +# kubernetes/raycluster.yaml. Set it to your node's FULL allocatable count: +# +# kubectl get node \ +# -o jsonpath='{.status.allocatable.vpc\.amazonaws\.com/efa}' +# +# p5en.48xlarge reports 15 allocatable (16 physical, one held back by the device +# plugin); p5.48xlarge reports 32. Do NOT guess from the instance spec -- read it +# off the node, because the allocatable count is what the scheduler honours. +# +# A partial request is a CORRECTNESS problem, not just a throughput one. When a pod +# claims fewer cards than the node exposes, the device plugin chooses which ones, +# the two pods can end up with different EFA-device-to-NIC-rail mappings, and +# aws-ofi-nccl aborts: +# +# NET/OFI Unexpected number of remote rails for dev N. Expected 1 but got 2 +# ncclInternalError: Internal check failed +# +# This does not appear in an all-reduce smoke test -- all-reduce across the node +# boundary passes with a partial request. It surfaces only on point-to-point +# traffic, i.e. pipeline parallel (PP>1) and MoE expert parallel (EP>1), so an +# EFA validation that only runs all_reduce_perf will report success and the failure +# lands later inside training. Requesting the full count avoids it entirely. +export EFA_PER_NODE=15 + +# ----- Training ----- +export PROMPT_DATA="/fsx/data/dapo-math-17k/dapo-math-17k.jsonl" +export EVAL_DATA="/fsx/data/aime-2024/aime-2024.jsonl" +# Each run MUST get its own CHECKPOINT_DIR: --load otherwise tries to resume a +# scheduler built for a different total-iteration count. train_iters = num_rollout +# here (16 x 8 // 128 = 1 step/rollout), so NUM_ROLLOUT == optimizer steps. +export CHECKPOINT_DIR="/fsx/runs/qwen3-4b/ckpt" +# SAVE_INTERVAL is set beyond NUM_ROLLOUT on purpose: miles's save_model() currently +# fails with a pickle-truncation error (see README Known Issues), so a mid-run save +# would crash the job. Lower this ONLY once that issue is resolved; until then a +# checkpoint-triggering value is a guaranteed failure, so the default must not trigger one. +export SAVE_INTERVAL=1000 +export NUM_ROLLOUT=100 +export ROLLOUT_BATCH_SIZE=16 +export N_SAMPLES_PER_PROMPT=8 +export GLOBAL_BATCH_SIZE=128 +export NUM_STEPS_PER_ROLLOUT=1 +export LEARNING_RATE="1e-6" +export ROLLOUT_TEMPERATURE=1.0 + +# ----- Reward ----- +# Default: in-process rule-based reward (runs on the GPU rollout nodes). +# Built-in types include: deepscaler, dapo, math, f1, gpqa. +export RM_TYPE="deepscaler" + +# ----- Observability (optional) ----- +# TENSORBOARD_DIR is forwarded into the Ray runtime env by the recipe so +# --use-tensorboard writes event files onto shared storage. Extra train.py flags +# can be appended without editing the recipe via EXTRA_TRAIN_ARGS (see the +# EXTRA_TRAIN_ARGS_ARR block in recipe/run_grpo_qwen3_4b.sh). +export TENSORBOARD_DIR="/fsx/tb/qwen3-4b" +export EXTRA_TRAIN_ARGS="--use-tensorboard" + +# ============================================================ +# ALTERNATE: Qwen3-30B-A3B MoE, colocated on 2 nodes (16 GPU) +# ------------------------------------------------------------ +# To run the 30B MoE case instead of 4B, override the model block above with the +# values below (uncomment) and launch recipe/run_grpo_qwen3_30b_a3b.sh. +# +# This block trains cleanly as written: rollout/raw_reward 0.578, repetition 0.0 over +# NUM_ROLLOUT=2 (comparable to the dense 4B run). Here EP_SIZE equals ROLLOUT_GPUS_PER_ENGINE +# (both 2), so the rollout MoE runs pure expert-parallel (moe_tp=1). Setting +# ROLLOUT_GPUS_PER_ENGINE larger than EP_SIZE runs the MoE tensor-parallel AND expert-parallel +# at once (moe_tp>1 and moe_ep>1), which trips a FlashInfer allreduce-fusion bug in this SGLang +# build; the recipe detects that and disables the fusion so those geometries train too (at some +# rollout-throughput cost). See README.md Known Issues item 2. +# +# The actor spans ALL 16 GPU (2 nodes x 8) with --use-distributed-optimizer +# (added by the recipe) so the 30B optimizer state is sharded and fits H200 +# (141GB). Confining the actor to 8 GPU OOMs. Colocated rollout time-shares the +# same 16 GPU. MoE online weight update needs the triton runner + explicit +# --sglang-expert-parallel-size (both added by the recipe from EP_SIZE). +# +# On B300 (288GB HBM/GPU, ~2x H200's 141GB), the disaggregated 8-actor layout +# that OOMs on H200 is expected to fit without --use-distributed-optimizer -- +# UNVERIFIED, not run on this hardware. If you have B300 capacity, that is the +# natural next verification step (see README Verification Status). +# +# export MODEL_NAME="Qwen/Qwen3-30B-A3B" +# export MODEL_LOCAL="/fsx/models/Qwen3-30B-A3B" +# export MODEL_DIST="/fsx/models/Qwen3-30B-A3B_torch_dist" +# export MODEL_SCRIPT="qwen3-30B-A3B.sh" +# export COLOCATE="true" +# export TP_SIZE=2 +# export PP_SIZE=1 +# export CP_SIZE=1 +# export EP_SIZE=2 +# export ACTOR_NUM_NODES=2 +# export ACTOR_GPUS_PER_NODE=8 +# export ROLLOUT_NUM_GPUS=16 +# export ROLLOUT_GPUS_PER_ENGINE=2 +# export MAX_TOKENS_PER_GPU=8192 +# export SGLANG_MEM_FRACTION=0.75 +# export CHECKPOINT_DIR="/fsx/runs/qwen3-30b/ckpt" +# export TENSORBOARD_DIR="/fsx/tb/qwen3-30b" +# NOTE: do NOT set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True for the 30B +# colocated run -- SGLang's torch_memory_saver (used to release rollout memory +# during the training phase) is incompatible with expandable_segments and aborts +# the rollout engine at init. The 4B run tolerates it; the 30B colocated run must +# omit it. CUDA_DEVICE_MAX_CONNECTIONS=1 is required by Megatron whenever TP>1 (or +# CP>1); the 30B recipe sets it in its runtime-env. The 4B recipe does NOT (the +# baseline is TP=1); if you raise the 4B config to TP>1, add it to that recipe's +# runtime-env too or the run asserts at startup. diff --git a/3.test_cases/pytorch/miles/env_vars.disaggregated.example b/3.test_cases/pytorch/miles/env_vars.disaggregated.example new file mode 100644 index 000000000..044ca1e8b --- /dev/null +++ b/3.test_cases/pytorch/miles/env_vars.disaggregated.example @@ -0,0 +1,75 @@ +# STATUS: UNVERIFIED overlay -- disaggregated reward on a CPU pool; not run on miles. +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +# +# DISAGGREGATED workload profile for the miles GRPO sample. +# +# An overlay sourced ON TOP OF the base env_vars. It moves reward scoring to a +# separate CPU instance group (disaggregating it from the GPU rollout) and turns +# up the workload on both tiers. +# +# Usage: +# cp env_vars.colocated.example env_vars && edit secrets (HF_TOKEN) +# cp env_vars.disaggregated.example env_vars.disaggregated +# source env_vars && source env_vars.disaggregated +# +# This profile deliberately increases load on BOTH tiers: +# * GPU tier -- larger rollout batch, more samples/prompt, longer responses +# and bigger global batch => more generation + more optimizer +# work per step on the p5 nodes. +# * Reward tier -- a real reward MODEL (deberta-v3-large) scaled to one replica +# per CPU node, so the CPU pool is the busy, independently-scaled +# tier instead of a millisecond rule-based check. +# +# Assumes: +# * env_vars (Option A: Qwen3-4B, colocated) has already been sourced. +# * A CPU node pool named ${REWARD_NODE_GROUP} already exists, in the same AZ as the GPU +# pool, with no EFA. Create it however your cluster provisions capacity -- a Karpenter +# NodePool, an EKS managed node group, or a SageMaker HyperPod instance group -- then set +# the name below to match. This test case does not create it for you. + +# ----- Heavier GRPO config (GPU tier) ----- +export ROLLOUT_BATCH_SIZE=64 # was 16 -> 4x more prompts per rollout +export N_SAMPLES_PER_PROMPT=16 # was 8 -> 2x more responses per prompt +# GRPO batch-size constraint: ROLLOUT_BATCH_SIZE × N_SAMPLES_PER_PROMPT = GLOBAL_BATCH_SIZE × NUM_STEPS_PER_ROLLOUT +export GLOBAL_BATCH_SIZE=512 # was 128 -> 4x bigger optimizer step +export NUM_STEPS_PER_ROLLOUT=2 # adjusted to satisfy constraint: 64*16 = 512*2 +export ROLLOUT_MAX_RESPONSE_LEN=16384 # was 8192 -> 2x longer generations +export NUM_ROLLOUT=20 # short but heavy run for the demo +export ROLLOUT_TEMPERATURE=1.0 +# 64 prompts x 16 samples = 1024 generations scored per rollout step, which is +# what drives sustained load onto the c5 reward pool. + +# ----- Heavier reward tier (CPU Spot pool, e.g. c5) ----- +export RM_TYPE="remote_rm" +export REWARD_BACKEND="reward_model" +export REWARD_MODEL_NAME="OpenAssistant/reward-model-deberta-v3-large-v2" +# Image for the reward pods. kubernetes/reward-service.yaml substitutes ${REWARD_IMAGE} +# through envsubst, so leaving it unset renders `image: ""` and the Deployment is rejected. +# It is built from reward_service.Dockerfile, which is CPU-only and much smaller than the +# miles GPU image, so it is a separate build: +# docker build -f reward_service.Dockerfile -t "${REWARD_IMAGE}" . +# docker push "${REWARD_IMAGE}" +# On a cluster with no local Docker daemon, build it in-cluster the same way the main image is +# built (see the BuildKit step in README.md), but with a ConfigMap that carries all three build +# inputs -- Dockerfile, reward_service/requirements.txt, and reward_service/app.py -- since +# reward_service.Dockerfile COPYs the latter two from the build context. +export REWARD_IMAGE_NAME="miles-reward" +export REWARD_TAG="v1" # pinned; avoid ":latest" per CONTRIBUTING +export REWARD_IMAGE="${REGISTRY}${REWARD_IMAGE_NAME}:${REWARD_TAG}" +# Name of the CPU node pool hosting the reward pods; it must already exist (see the Assumes +# block at the top of this file -- this test case does not create it for you). +# +# Weigh Spot carefully here. The reward service is stateless, but miles's `remote_rm` client +# has no retry-with-backoff (see the README's Known Issues), so a Spot interruption landing +# mid-rollout can fail the reward call and take the training job with it. Multiple replicas +# behind the Service narrow that window without closing it. Use On-Demand if the run must not +# be interrupted. +export REWARD_NODE_GROUP="reward-spot-c5" +export REWARD_REPLICAS=4 # one per node +# ml.c5.4xlarge is 16 physical vCPU, but this sample provisions the group with +# ThreadsPerCore=1 (hyperthreading off) -> ~8 logical / 7.9 allocatable vCPU on +# EKS. TORCH_NUM_THREADS is sized to fit alongside the pod cpu request below. +# If you provision with ThreadsPerCore=2, raise this toward ~14. +export REWARD_TORCH_THREADS=6 +export RM_URL="http://miles-reward.${NAMESPACE}.svc.cluster.local:8000/score" diff --git a/3.test_cases/pytorch/miles/kubernetes/data-prep-pod.yaml b/3.test_cases/pytorch/miles/kubernetes/data-prep-pod.yaml new file mode 100644 index 000000000..519602612 --- /dev/null +++ b/3.test_cases/pytorch/miles/kubernetes/data-prep-pod.yaml @@ -0,0 +1,42 @@ +# STATUS: UNVERIFIED -- mirrors slime data-prep-pod.yaml; not executed on miles. +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +apiVersion: v1 +kind: Pod +metadata: + name: data-prep + namespace: ${NAMESPACE} + labels: + app: miles-data-prep +spec: + containers: + - name: data-prep + image: python:3.11-slim + command: ["sleep", "infinity"] + env: + - name: HF_TOKEN + # Same Secret the RayCluster reads, so the token is never written into a file + # that could be committed. Create it once: + # kubectl create secret generic hf-token \ + # --from-literal=HF_TOKEN=hf_xxx -n "${NAMESPACE}" + valueFrom: + secretKeyRef: + name: hf-token + key: HF_TOKEN + optional: true + resources: + requests: + cpu: "4" + memory: "16Gi" + limits: + cpu: "8" + memory: "32Gi" + volumeMounts: + - name: fsx + mountPath: /fsx + volumes: + - name: fsx + persistentVolumeClaim: + claimName: ${FSX_CLAIM} + restartPolicy: Never + terminationGracePeriodSeconds: 30 diff --git a/3.test_cases/pytorch/miles/kubernetes/raycluster.yaml b/3.test_cases/pytorch/miles/kubernetes/raycluster.yaml new file mode 100644 index 000000000..5f427d991 --- /dev/null +++ b/3.test_cases/pytorch/miles/kubernetes/raycluster.yaml @@ -0,0 +1,190 @@ +apiVersion: ray.io/v1 +kind: RayCluster +metadata: + labels: + app: miles-grpo + name: miles-ray + namespace: ${NAMESPACE} +spec: + enableInTreeAutoscaling: false + headGroupSpec: + rayStartParams: + dashboard-host: 0.0.0.0 + num-gpus: '0' + template: + metadata: + labels: + ray.io/node-type: head + annotations: + # The head holds the Ray GCS and is a long-lived singleton. On clusters with a + # Karpenter (or similar) consolidator, protect its node from being reclaimed as + # "underutilized" while the head is running. A no-op on clusters without Karpenter. + karpenter.sh/do-not-disrupt: "true" + spec: + containers: + - env: + - name: NVIDIA_VISIBLE_DEVICES + value: void + - name: RAY_memory_monitor_refresh_ms + value: '0' + - name: PYTHONPATH + value: /root/Megatron-LM:/root/miles + - name: HF_TOKEN + valueFrom: + secretKeyRef: + key: HF_TOKEN + name: hf-token + - name: NCCL_DEBUG + value: WARN + image: ${FULL_IMAGE} + imagePullPolicy: IfNotPresent + name: ray-head + ports: + - containerPort: 6379 + protocol: TCP + - containerPort: 8265 + protocol: TCP + - containerPort: 10001 + protocol: TCP + resources: + limits: + cpu: '4' + memory: 16Gi + ephemeral-storage: 60Gi + requests: + cpu: '2' + memory: 8Gi + ephemeral-storage: 40Gi + volumeMounts: + - mountPath: /fsx + name: fsx + - mountPath: /dev/shm + name: dshm + restartPolicy: Never + terminationGracePeriodSeconds: 120 + volumes: + - name: fsx + persistentVolumeClaim: + claimName: ${FSX_CLAIM} + - emptyDir: + medium: Memory + # medium:Memory counts against the pod's memory limit (16Gi above), so + # keep this <= that limit. The head is a CPU coordinator and needs little + # shared memory; the GPU workers (below) get the large 256Gi /dev/shm. + sizeLimit: 8Gi + name: dshm + # The head MUST land on a node that has the CUDA driver (libcuda.so.1): miles's Ray + # control actors (e.g. the rollout manager) import Megatron/transformer_engine at + # startup even though they request num-gpus 0, and that import hard-fails on a + # GPU-less node. Point ${CPU_NODE_ROLE} at a GPU pool's node-role label; the default + # env sets it equal to ${GPU_NODE_ROLE} so the head co-locates on a GPU node (it + # still runs num-gpus 0 and consumes no GPU, only the node's CPU/disk). Do NOT place + # the head on a CPU-only pool. The toleration below lets it schedule onto the (tainted) + # GPU pool. + nodeSelector: + ${CPU_NODE_LABEL_KEY}: ${CPU_NODE_ROLE} + tolerations: + - key: nvidia.com/gpu + operator: Exists + effect: NoSchedule + rayVersion: 2.55.1 + workerGroupSpecs: + # Worker (GPU node) count is ${WORKER_REPLICAS}. For a COLOCATED run this equals + # ACTOR_NUM_NODES; for a DISAGGREGATED run (COLOCATE=false) the actor and rollout GPU pools + # are separate, so it must cover BOTH: ceil((ACTOR_NUM_NODES*ACTOR_GPUS_PER_NODE + + # ROLLOUT_NUM_GPUS) / 8). Setting it to ACTOR_NUM_NODES alone under-sizes a disaggregated run + # and the job waits forever on placement it can never get. env_vars derives it (see the file). + # podAntiAffinity (below) keeps the workers on separate nodes so a multi-node actor spans them. + - groupName: gpu-workers + maxReplicas: ${WORKER_REPLICAS} + minReplicas: ${WORKER_REPLICAS} + numOfHosts: 1 + rayStartParams: + num-gpus: '8' + resources: '"{\"gpu_node\": 1}"' + replicas: ${WORKER_REPLICAS} + template: + metadata: + labels: + ray.io/node-type: worker + spec: + affinity: + podAntiAffinity: + requiredDuringSchedulingIgnoredDuringExecution: + - labelSelector: + matchExpressions: + - key: ray.io/node-type + operator: In + values: + - worker + topologyKey: kubernetes.io/hostname + containers: + - env: + - name: RAY_memory_monitor_refresh_ms + value: '0' + - name: PYTHONPATH + value: /root/Megatron-LM:/root/miles + - name: HF_TOKEN + valueFrom: + secretKeyRef: + key: HF_TOKEN + name: hf-token + - name: FI_PROVIDER + value: efa + - name: FI_EFA_USE_DEVICE_RDMA + value: '1' + - name: FI_EFA_FORK_SAFE + value: '1' + - name: NCCL_PROTO + value: Simple + - name: NCCL_DEBUG + value: WARN + - name: RDMAV_FORK_SAFE + value: '1' + - name: NCCL_SOCKET_IFNAME + value: ^lo + - name: TOKENIZERS_PARALLELISM + value: 'false' + image: ${FULL_IMAGE} + imagePullPolicy: IfNotPresent + name: ray-worker + resources: + # EFA: request the node's FULL allocatable count, not a fraction. See the + # EFA_PER_NODE note in env_vars.colocated.example -- a partial request is a + # correctness problem on point-to-point paths (pipeline / expert parallel), + # not just a throughput one. + limits: + cpu: '96' + memory: 1900Gi + nvidia.com/gpu: '8' + vpc.amazonaws.com/efa: '${EFA_PER_NODE}' + requests: + cpu: '90' + memory: 1800Gi + nvidia.com/gpu: '8' + vpc.amazonaws.com/efa: '${EFA_PER_NODE}' + volumeMounts: + - mountPath: /fsx + name: fsx + - mountPath: /dev/shm + name: dshm + # ${GPU_NODE_ROLE} selects the accelerator pool: e.g. "gpu-p5en" (H200) or + # "gpu-b300" (B300) -- whatever node-role label your cluster's GPU NodePool + # carries (see README Hardware Requirements). Both generations expose the + # same node-role label scheme; only the value differs per cluster. + nodeSelector: + ${GPU_NODE_LABEL_KEY}: ${GPU_NODE_ROLE} + restartPolicy: Never + terminationGracePeriodSeconds: 300 + tolerations: + - effect: NoSchedule + key: nvidia.com/gpu + operator: Exists + volumes: + - name: fsx + persistentVolumeClaim: + claimName: ${FSX_CLAIM} + - emptyDir: + medium: Memory + sizeLimit: 256Gi + name: dshm From b103529d928fb90c3970a860c5502730d1a431ca Mon Sep 17 00:00:00 2001 From: littlemex Date: Tue, 18 Aug 2026 19:58:52 +0900 Subject: [PATCH 3/7] feat(miles): dense and MoE GRPO recipes, plus a head-pod launch helper Qwen3-4B (dense) and Qwen3-30B-A3B (MoE) GRPO recipes via a shared launcher. The MoE recipe derives the SGLang rollout geometry; pure EP (moe_tp=1) and pure TP (moe_ep=1) train cleanly, while combined moe_tp>1 and moe_ep>1 hits a FlashInfer allreduce-fusion bug (drops the moe-tp reduce), so the recipe adds --sglang-enforce-disable-flashinfer-allreduce-fusion for that geometry. run-on-cluster.sh optionally launches a recipe from inside the head pod so the operator needs only kubectl (no local ray CLI, no port-forward); it does not deploy infra. --- .../miles/recipe/launcher/grpo_launch.sh | 69 ++++ .../miles/recipe/run_grpo_qwen3_30b_a3b.sh | 276 +++++++++++++++ .../pytorch/miles/recipe/run_grpo_qwen3_4b.sh | 322 ++++++++++++++++++ 3.test_cases/pytorch/miles/run-on-cluster.sh | 78 +++++ 4 files changed, 745 insertions(+) create mode 100755 3.test_cases/pytorch/miles/recipe/launcher/grpo_launch.sh create mode 100644 3.test_cases/pytorch/miles/recipe/run_grpo_qwen3_30b_a3b.sh create mode 100644 3.test_cases/pytorch/miles/recipe/run_grpo_qwen3_4b.sh create mode 100755 3.test_cases/pytorch/miles/run-on-cluster.sh diff --git a/3.test_cases/pytorch/miles/recipe/launcher/grpo_launch.sh b/3.test_cases/pytorch/miles/recipe/launcher/grpo_launch.sh new file mode 100755 index 000000000..4e6891e1b --- /dev/null +++ b/3.test_cases/pytorch/miles/recipe/launcher/grpo_launch.sh @@ -0,0 +1,69 @@ +#!/usr/bin/env bash +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +# ============================================================ +# miles 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): +# miles's per-model scripts (scripts/models/*.sh, inherited from slime) 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 the 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 model script to source (e.g. qwen3-4B.sh), passed +# through the Ray runtime env by the recipe. +# env SLIME_DIR : framework install dir in the image (default /root/miles; +# variable name kept for parity with the upstream launcher). +# ============================================================ + +set -euo pipefail + +# miles installs itself editable at /root/miles (the variable name is kept as +# SLIME_DIR for parity with the upstream launcher; only the default path moves). +SLIME_DIR="${SLIME_DIR:-/root/miles}" + +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/miles/recipe/run_grpo_qwen3_30b_a3b.sh b/3.test_cases/pytorch/miles/recipe/run_grpo_qwen3_30b_a3b.sh new file mode 100644 index 000000000..bd86ffa8c --- /dev/null +++ b/3.test_cases/pytorch/miles/recipe/run_grpo_qwen3_30b_a3b.sh @@ -0,0 +1,276 @@ +#!/usr/bin/env bash +# STATUS: Verified -- trains cleanly (colocated, 2 nodes / 16 GPU H200). GRPO steps run to +# completion with --colocate + --use-distributed-optimizer + triton MoE runner, and the +# resulting generation is healthy: rollout/raw_reward 0.578, rollout/repetition_frac 0.0 +# over NUM_ROLLOUT=2 (comparable to the dense 4B run), weight_version uniform / mixed 0.0. +# The shipped block runs the rollout MoE in pure expert-parallel (moe_tp=1, i.e. +# EP_SIZE == ROLLOUT_GPUS_PER_ENGINE). Running it tensor-parallel AND expert-parallel at +# once (moe_tp>1 and moe_ep>1) hits a FlashInfer allreduce-fusion bug in this SGLang build +# that drops the moe-tp reduce and yields degenerate, zero-reward text; the recipe detects +# that geometry and disables the fusion (--sglang-enforce-disable-flashinfer-allreduce-fusion) +# so the combined path trains correctly too. See README.md Known Issues item 2. The +# disaggregated actor-8 layout (see the parent slime recipe) needs B300 288GB +# and is UNVERIFIED on H200; colocated 16-GPU with the distributed optimizer is the +# H200-fitting layout and is the one shipped as default here. +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +# ============================================================ +# miles GRPO Training — Qwen3-30B-A3B MoE on HyperPod EKS +# (Colocated Mode: training and rollout time-share the same 16 GPUs) +# +# This configuration runs the 30B MoE model with: +# - actor 2 nodes x 8 GPUs = 16 GPU, TP=2 PP=1 CP=1 EP=2 +# - --use-distributed-optimizer shards the 30B optimizer state across the 16 GPU +# so the static memory fits H200 (141GB); without it the actor OOMs. +# - rollout time-shares the same 16 GPU (COLOCATE=true), 2 GPU per SGLang engine. +# +# Prerequisites: +# - Ray cluster deployed via kubernetes/raycluster.yaml (2 workers for 16 GPU) +# - Model downloaded and converted to torch_dist format +# - Training data on FSx +# - source env_vars with the "ALTERNATE: Qwen3-30B-A3B MoE" block in +# env_vars.colocated.example uncommented +# +# Usage: +# source env_vars # with the 30B MoE block (env_vars.colocated.example) active +# bash recipe/run_grpo_qwen3_30b_a3b.sh +# ============================================================ + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +PROJECT_DIR="$(dirname "${SCRIPT_DIR}")" + +# Source environment if not already loaded. Override the file with ENV_FILE, +# e.g. ENV_FILE=env_vars.disaggregated bash recipe/run_grpo_qwen3_30b_a3b.sh +if [[ -z "${MODEL_LOCAL:-}" ]]; then + ENV_FILE="${ENV_FILE:-${PROJECT_DIR}/env_vars}" + echo "[INFO] Sourcing ${ENV_FILE}..." + source "${ENV_FILE}" +fi + +# Validate required variables (same guard as the 4B recipe; without it a missing +# var trips `set -u` with an opaque error deep in the argv build). +for var in MODEL_LOCAL MODEL_DIST PROMPT_DATA CHECKPOINT_DIR MODEL_SCRIPT RM_TYPE \ + COLOCATE TP_SIZE PP_SIZE CP_SIZE EP_SIZE ACTOR_NUM_NODES ACTOR_GPUS_PER_NODE \ + ROLLOUT_NUM_GPUS ROLLOUT_GPUS_PER_ENGINE NUM_ROLLOUT ROLLOUT_BATCH_SIZE \ + N_SAMPLES_PER_PROMPT GLOBAL_BATCH_SIZE MAX_TOKENS_PER_GPU \ + ROLLOUT_MAX_RESPONSE_LEN ROLLOUT_TEMPERATURE LEARNING_RATE SAVE_INTERVAL EVAL_DATA; do + if [[ -z "${!var:-}" ]]; then + echo "[ERROR] ${var} is not set. Please configure env_vars." + exit 1 + fi +done + +# --colocate is passed unconditionally below, so COLOCATE=false cannot be honored. Refuse +# rather than warn: continuing would run the colocated layout while the env said otherwise. +if [[ "${COLOCATE:-true}" != "true" ]]; then + echo "[ERROR] COLOCATE=${COLOCATE}, but this recipe only builds the colocated layout." >&2 + echo "[ERROR] Disaggregating a 30B actor needs B300-class 288GB HBM and a different" >&2 + echo "[ERROR] rollout layout; it is UNVERIFIED here. Set COLOCATE=true." >&2 + exit 1 +fi + +# SGLang rollout MoE geometry. The recipe passes --rollout-num-gpus-per-engine +# (the engine's tensor-parallel size) and --sglang-expert-parallel-size (EP_SIZE), +# from which SGLang derives moe_ep = EP_SIZE and moe_tp = per-engine GPUs / EP_SIZE. +# EP_SIZE must divide the per-engine GPU count for that split to be integral. +if (( ROLLOUT_GPUS_PER_ENGINE % EP_SIZE != 0 )); then + echo "[ERROR] EP_SIZE=${EP_SIZE} must divide ROLLOUT_GPUS_PER_ENGINE=${ROLLOUT_GPUS_PER_ENGINE}" >&2 + echo "[ERROR] so the SGLang rollout MoE tensor/expert split is integral." >&2 + exit 1 +fi +SGLANG_MOE_TP=$(( ROLLOUT_GPUS_PER_ENGINE / EP_SIZE )) +# When the rollout MoE runs tensor-parallel AND expert-parallel at once (moe_tp>1 and +# moe_ep>1), this SGLang build's FlashInfer allreduce+RMSNorm fusion -- auto-enabled on +# SM90/SM100 -- reduces the post-experts output over the moe-ep group only and drops the +# moe-tp reduce, so each rank keeps a partial sum and generation degenerates to zero +# reward. Disabling the fusion restores the correct two-stage reduce and the combined +# geometry trains cleanly. Pure expert-parallel (moe_tp=1, EP_SIZE == ROLLOUT_GPUS_PER_ENGINE) +# and pure tensor-parallel (moe_ep=1) are unaffected and leave the fusion on. See +# README.md Known Issues item 2. +SGLANG_FUSION_ARGS=() +if (( EP_SIZE > 1 && SGLANG_MOE_TP > 1 )); then + MOE_GEOM_NOTE="(TP+EP; allreduce fusion disabled)" + echo "[INFO] Rollout MoE runs moe_tp=${SGLANG_MOE_TP} x moe_ep=${EP_SIZE} (both >1);" + echo "[INFO] disabling FlashInfer allreduce fusion for this geometry (README Known Issues item 2)." + SGLANG_FUSION_ARGS+=(--sglang-enforce-disable-flashinfer-allreduce-fusion) +elif (( SGLANG_MOE_TP == 1 )); then + MOE_GEOM_NOTE="(pure EP)" +else + MOE_GEOM_NOTE="(pure TP)" +fi + +echo "============================================================" +echo " miles GRPO Training — Qwen3-30B-A3B MoE (Colocated, 16 GPU)" +echo "============================================================" +echo " Model: ${MODEL_LOCAL}" +echo " Megatron ckpt: ${MODEL_DIST}" +echo " Training data: ${PROMPT_DATA}" +echo " Checkpoints: ${CHECKPOINT_DIR}/qwen3-30b-a3b-grpo/" +echo " Actor: ${ACTOR_NUM_NODES} nodes x ${ACTOR_GPUS_PER_NODE} GPUs" +echo " Rollout GPUs: ${ROLLOUT_NUM_GPUS} (${ROLLOUT_GPUS_PER_ENGINE} per engine)" +echo " Parallelism: TP=${TP_SIZE} PP=${PP_SIZE} CP=${CP_SIZE} EP=${EP_SIZE}" +echo " SGLang rollout: moe_ep=${EP_SIZE} moe_tp=${SGLANG_MOE_TP} ${MOE_GEOM_NOTE}" +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 miles 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 miles 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}") +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-url "${RM_URL}") + echo " Reward: remote_rm @ ${RM_URL}" +fi + +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}" + # Colocated on 16 GPU: actor and rollout time-share the same GPUs. Required + # for 30B on H200 -- with all 16 GPU behind the actor, --use-distributed-optimizer + # shards the 30B optimizer state so static memory fits 141GB. Confining the + # actor to 8 GPU (a single node) OOMs (30B static ~121GB/GPU on 8-way). + --colocate + --use-distributed-optimizer + --rollout-num-gpus "${ROLLOUT_NUM_GPUS}" + --rollout-num-gpus-per-engine "${ROLLOUT_GPUS_PER_ENGINE}" + + # 0.75 is the value that completed on H200. Colocated, the rollout engine shares each GPU + # with the actor, so the usual single-tenant fractions (0.8-0.85) leave too little room. + --sglang-mem-fraction-static "${SGLANG_MEM_FRACTION:-0.75}" + # MoE online weight update on SGLang 0.5.12+ requires the triton runner: the + # default flashinfer MoE runner is incompatible with SLIME/miles's in-place + # weight update, so the engine must serve with --sglang-moe-runner-backend + # triton and expert parallelism must be declared to SGLang explicitly. + --sglang-moe-runner-backend triton + --sglang-expert-parallel-size "${EP_SIZE}" + # Disable the FlashInfer allreduce+RMSNorm fusion only when moe_tp>1 and moe_ep>1 + # (empty otherwise); the fusion mis-handles that combined path in this build. + ${SGLANG_FUSION_ARGS[@]+"${SGLANG_FUSION_ARGS[@]}"} + # Lowercase only: this reaches uvicorn's log_level, whose LOG_LEVELS dict has no "WARN" + # key, and the KeyError kills the rollout server before it binds. See docs/PORT_NOTES.md. + --sglang-log-level warning + # Careful when adding --sglang-* flags: miles registers them from SGLang's live + # ServerArgs with ignore_unknown_args, so a flag SGLang has since removed is silently + # accepted and does nothing. Check it against the SGLang version in the base image. +) + +# Optional extra train.py flags injected via EXTRA_TRAIN_ARGS, same mechanism as +# the 4B recipe (e.g. --use-tensorboard). Any flag taking a module path must be a +# DOTTED module path, not "file.py:func" -- load_function does rpartition('.') + +# import_module, so the slash form fails with ModuleNotFoundError at the loss +# forward. The miles package is on PYTHONPATH (/root/miles) below. +EXTRA_TRAIN_ARGS_ARR=() +if [ -n "${EXTRA_TRAIN_ARGS:-}" ]; then + # shellcheck disable=SC2206 + EXTRA_TRAIN_ARGS_ARR=(${EXTRA_TRAIN_ARGS}) + echo " Extra args: ${EXTRA_TRAIN_ARGS}" +fi +# ${arr[@]+...} guards the empty-array + `set -u` case on bash < 4.4 (macOS 3.2). +TRAIN_ARGS+=(${EXTRA_TRAIN_ARGS_ARR[@]+"${EXTRA_TRAIN_ARGS_ARR[@]}"}) + +# 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 miles code itself is already in the image at +# /root/miles (on PYTHONPATH below). MODEL_SCRIPT is forwarded so the launcher +# can source the right model definition. +# +# --entrypoint-resources '{"gpu_node": 0.001}' pins the driver to a GPU worker +# (miles imports mooncake / libcuda at module load; the head is a non-GPU pod). +# CUDA_DEVICE_MAX_CONNECTIONS=1 is required by Megatron for TP>1 (30B is TP=2). +# HF_TOKEN is NOT set here: it is injected into the pod env from the k8s Secret in +# raycluster.yaml, so it never lands in the Ray GCS runtime-env. +echo "[INFO] Submitting Ray job for MoE GRPO training..." + +ray job submit \ + --address="http://127.0.0.1:8265" \ + --entrypoint-resources '{"gpu_node": 0.001}' \ + --working-dir "${SCRIPT_DIR}/launcher" \ + --runtime-env-json="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/Megatron-LM:/root/miles\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"MODEL_SCRIPT\": \"${MODEL_SCRIPT}\", + \"TOKENIZERS_PARALLELISM\": \"false\", + \"NCCL_DEBUG\": \"WARN\", + \"FI_PROVIDER\": \"efa\", + \"FI_EFA_USE_DEVICE_RDMA\": \"1\", + \"TENSORBOARD_DIR\": \"${TENSORBOARD_DIR:-}\" + } + }" \ + -- bash grpo_launch.sh "${TRAIN_ARGS[@]}" + +echo "[INFO] Job submitted. Monitor at http://localhost:8265" diff --git a/3.test_cases/pytorch/miles/recipe/run_grpo_qwen3_4b.sh b/3.test_cases/pytorch/miles/recipe/run_grpo_qwen3_4b.sh new file mode 100644 index 000000000..cdf00e7e4 --- /dev/null +++ b/3.test_cases/pytorch/miles/recipe/run_grpo_qwen3_4b.sh @@ -0,0 +1,322 @@ +#!/usr/bin/env bash +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +# ============================================================ +# miles GRPO Training — Qwen3-4B on HyperPod EKS +# +# Submits a Ray job that runs GRPO with Megatron-LM for training and SGLang for +# rollout. COLOCATE picks how the two share the cluster, and with it how weights +# reach the rollout engines after every training step: +# +# COLOCATE=true one GPU pool, time-shared. Weights move by CUDA IPC. +# COLOCATE=false two GPU pools. Weights move by NCCL broadcast, over EFA when +# the pools are on different nodes. +# +# Both layouts are exercised by this recipe; see the COLOCATE block below for the +# sizing rules each one imposes. +# +# Prerequisites: +# - Ray cluster deployed via kubernetes/raycluster.yaml +# - Model downloaded and converted to torch_dist format +# - Training data on FSx +# - source env_vars before running +# +# Usage: +# source env_vars +# bash recipe/run_grpo_qwen3_4b.sh +# ============================================================ + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +PROJECT_DIR="$(dirname "${SCRIPT_DIR}")" + +# Source environment if not already loaded. Override the file with ENV_FILE, +# e.g. ENV_FILE=env_vars.disaggregated bash recipe/run_grpo_qwen3_4b.sh +if [[ -z "${MODEL_LOCAL:-}" ]]; then + ENV_FILE="${ENV_FILE:-${PROJECT_DIR}/env_vars}" + echo "[INFO] Sourcing ${ENV_FILE}..." + source "${ENV_FILE}" +fi + +# Validate required variables +for var in MODEL_LOCAL MODEL_DIST PROMPT_DATA CHECKPOINT_DIR MODEL_SCRIPT RM_TYPE \ + COLOCATE TP_SIZE PP_SIZE CP_SIZE EP_SIZE ACTOR_NUM_NODES ACTOR_GPUS_PER_NODE \ + ROLLOUT_NUM_GPUS ROLLOUT_GPUS_PER_ENGINE NUM_ROLLOUT ROLLOUT_BATCH_SIZE \ + N_SAMPLES_PER_PROMPT GLOBAL_BATCH_SIZE MAX_TOKENS_PER_GPU ROLLOUT_MAX_RESPONSE_LEN \ + ROLLOUT_TEMPERATURE LEARNING_RATE SAVE_INTERVAL EVAL_DATA; do + if [[ -z "${!var:-}" ]]; then + echo "[ERROR] ${var} is not set. Please configure env_vars." + exit 1 + fi +done + +# COLOCATE selects the weight sync path, which is the whole reason this variable exists. +# miles picks the implementation from --colocate alone (backends/megatron_utils/actor.py): +# +# --colocate present UpdateWeightFromTensor actor and rollout share devices; weights +# move by CUDA IPC handle +# --colocate absent UpdateWeightFromDistributed actor and rollout own separate devices; +# weights move by NCCL broadcast, over EFA +# when the two pools sit on different nodes +# (--update-weight-transfer-mode defaults +# to "broadcast") +# +# Accept only the two spellings. Under `set -u` an unset COLOCATE aborts, but "True", "1" and +# "yes" would otherwise fall through to the disaggregated branch and pick a layout nobody +# asked for -- the same silent disagreement as printing "Colocated: false" while running the +# colocated layout, just pointing the other way. +case "${COLOCATE}" in + true|false) ;; + *) echo "[ERROR] COLOCATE must be exactly 'true' or 'false', got '${COLOCATE}'." >&2 + echo "[ERROR] Values like True/1/yes would silently select a layout you did not ask for." >&2 + exit 1 ;; +esac + +# Arithmetic and -ne on a non-numeric value do not abort the script: the comparison itself +# errors, `if` reads that as false, and the layout check below passes without having run. A +# check that cannot run is worse than no check, because it reads as a pass. $((...)) also +# re-evaluates its operands as expressions, so a non-numeric value is an injection surface. +for _v in ACTOR_NUM_NODES ACTOR_GPUS_PER_NODE ROLLOUT_NUM_GPUS ROLLOUT_GPUS_PER_ENGINE; do + if ! [[ "${!_v}" =~ ^[0-9]+$ ]]; then + echo "[ERROR] ${_v} must be a non-negative integer, got '${!_v}'." >&2 + exit 1 + fi +done +unset _v + +# Engines must tile the rollout pool exactly, in either layout. A remainder leaves GPUs with +# no engine, or asks an engine for a shard that does not exist; neither shows up at submit +# time, so check it here where the numbers are still in view. +if [[ "${ROLLOUT_GPUS_PER_ENGINE}" -le 0 ]] \ + || [[ $((ROLLOUT_NUM_GPUS % ROLLOUT_GPUS_PER_ENGINE)) -ne 0 ]]; then + echo "[ERROR] ROLLOUT_GPUS_PER_ENGINE (${ROLLOUT_GPUS_PER_ENGINE}) must be a positive" >&2 + echo "[ERROR] divisor of ROLLOUT_NUM_GPUS (${ROLLOUT_NUM_GPUS})." >&2 + exit 1 +fi + +# Resolve the KV-pool fraction once. Two readers of the same default drift apart: fix one and +# the banner starts describing a different run than the argv does. Colocated keeps the value +# this recipe has always used, so an existing colocated invocation renders an unchanged argv. +MEM_FRACTION="${SGLANG_MEM_FRACTION:-0.8}" + +# An array, not a string: the disaggregated case must expand to ZERO argv tokens, and an empty +# string would reach argparse as a stray positional. Same property the recipe already relies on +# for EXTRA_TRAIN_ARGS_ARR. +COLOCATE_ARGS=() +[[ "${COLOCATE}" == "true" ]] && COLOCATE_ARGS=(--colocate) + +ACTOR_GPUS=$((ACTOR_NUM_NODES * ACTOR_GPUS_PER_NODE)) +TOTAL_GPUS=$((ACTOR_GPUS + ROLLOUT_NUM_GPUS)) + +# Holds in either layout: the trainer alone cannot exceed the cluster. Checked before the +# per-layout rules so an impossible actor size is reported as such, rather than surfacing as +# whichever layout-specific inequality happens to trip first. +if [[ -n "${CLUSTER_GPUS:-}" ]] && [[ "${ACTOR_GPUS}" -gt "${CLUSTER_GPUS}" ]]; then + echo "[ERROR] actor needs ${ACTOR_NUM_NODES} x ${ACTOR_GPUS_PER_NODE} = ${ACTOR_GPUS} GPUs," >&2 + echo "[ERROR] more than CLUSTER_GPUS=${CLUSTER_GPUS}." >&2 + exit 1 +fi + +if [[ "${COLOCATE}" == "true" ]]; then + # Sharing devices means the rollout count IS the actor count. A mismatch would place + # engines on a different number of GPUs than the trainer holds. + if [[ "${ROLLOUT_NUM_GPUS}" -ne "${ACTOR_GPUS}" ]]; then + echo "[ERROR] COLOCATE=true shares devices, so ROLLOUT_NUM_GPUS (${ROLLOUT_NUM_GPUS})" >&2 + echo "[ERROR] must equal the actor GPU count (${ACTOR_NUM_NODES} x ${ACTOR_GPUS_PER_NODE} = ${ACTOR_GPUS})." >&2 + exit 1 + fi +else + # Separate pools must both fit, and over-subscribing does not fail loudly: Ray waits on a + # placement group that never becomes ready, which reads as a hang rather than as a + # misconfiguration. CLUSTER_GPUS is optional because only the caller knows the cluster; + # when it is set, refuse here instead. Note that actor == rollout is the INTENDED shape on + # a 2-node 8-GPU cluster (8 + 8 = 16), not a mistake to warn about. + if [[ -n "${CLUSTER_GPUS:-}" ]] && [[ "${TOTAL_GPUS}" -gt "${CLUSTER_GPUS}" ]]; then + echo "[ERROR] COLOCATE=false needs actor ${ACTOR_GPUS} + rollout ${ROLLOUT_NUM_GPUS}" >&2 + echo "[ERROR] = ${TOTAL_GPUS} GPUs, more than CLUSTER_GPUS=${CLUSTER_GPUS}." >&2 + echo "[ERROR] Ray would wait forever on an unschedulable placement group." >&2 + exit 1 + fi +fi + +echo "============================================================" +echo " miles GRPO Training — Qwen3-4B" +echo "============================================================" +echo " Model: ${MODEL_LOCAL}" +echo " Megatron ckpt: ${MODEL_DIST}" +echo " Training data: ${PROMPT_DATA}" +echo " Checkpoints: ${CHECKPOINT_DIR}/qwen3-4b-grpo/" +echo " Nodes: ${ACTOR_NUM_NODES} x ${ACTOR_GPUS_PER_NODE} GPUs" +if [[ "${COLOCATE}" == "true" ]]; then + echo " Weight sync: colocated / CUDA IPC (UpdateWeightFromTensor)" + echo " GPUs: ${ACTOR_GPUS} shared by actor and rollout" +else + echo " Weight sync: disaggregated / NCCL broadcast (UpdateWeightFromDistributed)" + echo " GPUs: actor ${ACTOR_GPUS} + rollout ${ROLLOUT_NUM_GPUS} = ${TOTAL_GPUS}" +fi +# The static fraction bounds each engine's KV pool. Colocated must leave room for the trainer +# on the same device; disaggregated owns its GPUs and can take more. Print the value that the +# flag below actually receives, so the log never describes a run that did not happen. +echo " Mem fraction: ${MEM_FRACTION}" +echo " Rollout BS: ${ROLLOUT_BATCH_SIZE} x ${N_SAMPLES_PER_PROMPT}" +echo " Global BS: ${GLOBAL_BATCH_SIZE}" +echo " Num rollouts: ${NUM_ROLLOUT}" +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 miles 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 miles 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}") +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-url "${RM_URL}") + echo " Reward: remote_rm @ ${RM_URL}" +fi + +# Optional extra train.py flags, injected without editing this recipe. Set the +# EXTRA_TRAIN_ARGS env var to a whitespace-separated list of flags in env_vars +# (or on the command line) and they are appended verbatim to the train.py argv. +# This is the supported extension point for flags the baseline recipe does not +# set — e.g. observability (--use-tensorboard) or LR scheduling +# (--lr-decay-style ...). NOTE: any flag that takes a module path (e.g. a custom +# reward or callback) must be a DOTTED module path (load_function does +# rpartition('.') + import_module); the "file.py:func" form fails with +# ModuleNotFoundError. Word-splitting here is intentional so a single env var can +# carry several flags; values containing spaces are not supported (none of the +# intended flags need them). The array keeps each token separate across the Ray +# job boundary, the same shell-safety property the rest of this recipe relies on. +EXTRA_TRAIN_ARGS_ARR=() +if [ -n "${EXTRA_TRAIN_ARGS:-}" ]; then + # shellcheck disable=SC2206 + EXTRA_TRAIN_ARGS_ARR=(${EXTRA_TRAIN_ARGS}) + echo " Extra args: ${EXTRA_TRAIN_ARGS}" +fi + +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_ARGS[@]+"${COLOCATE_ARGS[@]}"} + # Colocated: actor and rollout share the same GPUs, so rollout-num-gpus should + # equal the actor GPU count (ACTOR_NUM_NODES x ACTOR_GPUS_PER_NODE). On 1 node + # this is 8; on 2 nodes 16. Set it explicitly so scaling nodes does not silently + # leave the rollout at a stale GPU count. + --rollout-num-gpus "${ROLLOUT_NUM_GPUS}" + --rollout-num-gpus-per-engine "${ROLLOUT_GPUS_PER_ENGINE}" + + --sglang-mem-fraction-static "${MEM_FRACTION}" + # Lowercase only: this reaches uvicorn's log_level, whose LOG_LEVELS dict has no "WARN" + # key, and the KeyError kills the rollout server before it binds -- training then hangs + # on a health check that never passes. See docs/PORT_NOTES.md. + --sglang-log-level warning + + # Flags injected via the EXTRA_TRAIN_ARGS env var. The shipped default is just + # observability (--use-tensorboard), which does not change the loss; set it to + # empty for a bit-identical baseline, or add flags. See the block above. + # The ${arr[@]+"${arr[@]}"} form expands to nothing (not an "unbound variable" + # error) when the array is empty under `set -u` on bash < 4.4 (e.g. macOS 3.2). + ${EXTRA_TRAIN_ARGS_ARR[@]+"${EXTRA_TRAIN_ARGS_ARR[@]}"} +) + +# 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 miles code itself is already in the image at +# /root/miles (on PYTHONPATH below). MODEL_SCRIPT is forwarded so the launcher +# can source the right model definition. +# +# --entrypoint-resources '{"gpu_node": 0.001}' pins the Ray job DRIVER to a GPU +# worker. miles imports mooncake (libcuda-dependent) at module load, and the head +# is a non-GPU pod, so a head-scheduled driver dies with a libcuda import error. +# The 0.001 fractional request lands the driver on a worker without consuming a +# whole GPU (does not disturb the colocated placement group). HF_TOKEN is NOT set +# here: it is injected into the pod env from the k8s Secret in raycluster.yaml, so +# it never lands in the Ray GCS runtime-env (visible via the dashboard API). +echo "[INFO] Submitting Ray job..." + +ray job submit \ + --address="http://127.0.0.1:8265" \ + --entrypoint-resources '{"gpu_node": 0.001}' \ + --working-dir "${SCRIPT_DIR}/launcher" \ + --runtime-env-json="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/Megatron-LM:/root/miles\", + \"MODEL_SCRIPT\": \"${MODEL_SCRIPT}\", + \"TOKENIZERS_PARALLELISM\": \"false\", + \"NCCL_DEBUG\": \"WARN\", + \"FI_PROVIDER\": \"efa\", + \"FI_EFA_USE_DEVICE_RDMA\": \"1\", + \"TENSORBOARD_DIR\": \"${TENSORBOARD_DIR:-}\" + } + }" \ + -- bash grpo_launch.sh "${TRAIN_ARGS[@]}" + +echo "[INFO] Job submitted. Monitor at http://localhost:8265" diff --git a/3.test_cases/pytorch/miles/run-on-cluster.sh b/3.test_cases/pytorch/miles/run-on-cluster.sh new file mode 100755 index 000000000..3c9c7c1bf --- /dev/null +++ b/3.test_cases/pytorch/miles/run-on-cluster.sh @@ -0,0 +1,78 @@ +#!/usr/bin/env bash +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +# ============================================================ +# Run a miles GRPO recipe FROM the Ray head pod, so the operator's machine needs only +# kubectl (plus AWS auth) -- no local `ray` CLI, no `kubectl port-forward`, no local `envsubst`. +# +# Why this exists: the shipped recipes end in `ray job submit --address http://127.0.0.1:8265`, +# which normally runs on your laptop and therefore needs the ray CLI installed locally, a +# port-forward to the dashboard, and a ray version that matches the cluster (the wheel for the +# cluster's exact ray version may not even exist for your local Python). Running the recipe +# INSIDE the head pod removes all three: ray is already there, 127.0.0.1:8265 is the head's own +# dashboard, and the version always matches. This is a convenience wrapper around the SAME +# recipe -- the local flow in the README Quick Start still works unchanged. +# +# What it does NOT do: it does not create infrastructure, build images, prepare data, or deploy +# manifests. Do those first per the README (steps 0-5): the RayCluster must already be deployed +# with its head Running on a CUDA-capable (GPU) node, and env_vars must be filled in. +# +# Usage: +# ./run-on-cluster.sh [--recipe run_grpo_qwen3_4b.sh] [--env ./env_vars] [--namespace default] [--dry-run] +# ============================================================ +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +NAMESPACE="${NAMESPACE:-default}" +RECIPE="run_grpo_qwen3_4b.sh" +ENV_FILE="${SCRIPT_DIR}/env_vars" +DRY_RUN=0 +REMOTE_DIR="/tmp/miles-run" + +while [[ $# -gt 0 ]]; do + case "$1" in + --recipe) RECIPE="$2"; shift 2;; + --env) ENV_FILE="$2"; shift 2;; + -n|--namespace) NAMESPACE="$2"; shift 2;; + --dry-run) DRY_RUN=1; shift;; + -h|--help) sed -n '2,22p' "${BASH_SOURCE[0]}"; exit 0;; + *) echo "[ERROR] unknown argument: $1" >&2; exit 1;; + esac +done + +# ---- preconditions (fail-fast, with the fix in the message) ---- +command -v kubectl >/dev/null 2>&1 || { echo "[ERROR] kubectl not found on PATH." >&2; exit 1; } +[[ -f "${ENV_FILE}" ]] || { echo "[ERROR] env file not found: ${ENV_FILE}. Copy env_vars.colocated.example to env_vars and fill it in (README step 1)." >&2; exit 1; } +[[ -f "${SCRIPT_DIR}/recipe/${RECIPE}" ]] || { echo "[ERROR] recipe not found: recipe/${RECIPE}" >&2; exit 1; } + +# ---- discover: the Ray head pod ---- +HEAD="$(kubectl -n "${NAMESPACE}" get pods -l ray.io/node-type=head -o jsonpath='{.items[0].metadata.name}' 2>/dev/null || true)" +[[ -n "${HEAD}" ]] || { echo "[ERROR] no Ray head pod (label ray.io/node-type=head) in namespace '${NAMESPACE}'. Deploy the RayCluster first (README step 5)." >&2; exit 1; } +PHASE="$(kubectl -n "${NAMESPACE}" get pod "${HEAD}" -o jsonpath='{.status.phase}' 2>/dev/null || true)" +[[ "${PHASE}" == "Running" ]] || { echo "[ERROR] head pod '${HEAD}' is '${PHASE:-unknown}', not Running. Wait for it (it pulls the ~18 GB image on first start)." >&2; exit 1; } + +echo "[INFO] namespace: ${NAMESPACE}" +echo "[INFO] head pod: ${HEAD}" +echo "[INFO] recipe: recipe/${RECIPE}" +echo "[INFO] env file: ${ENV_FILE}" +echo "[INFO] remote: ${REMOTE_DIR} (recipe/, scripts/, env_vars)" + +if [[ "${DRY_RUN}" == "1" ]]; then + echo "[DRY-RUN] would ship recipe/ scripts/ and '${ENV_FILE}' to ${HEAD}:${REMOTE_DIR}, then run:" + echo " kubectl -n ${NAMESPACE} exec ${HEAD} -- bash -lc 'cd ${REMOTE_DIR} && ENV_FILE=${REMOTE_DIR}/env_vars bash recipe/${RECIPE}'" + exit 0 +fi + +# ---- ship: tar-pipe the test-case files into the head pod ---- +# (kubectl cp is tar under the hood; piping tar ourselves lets us exclude .git and send exactly +# these paths.) The env file is copied to ${REMOTE_DIR}/env_vars regardless of its local name. +kubectl -n "${NAMESPACE}" exec "${HEAD}" -- bash -lc "rm -rf ${REMOTE_DIR} && mkdir -p ${REMOTE_DIR}" +tar cf - -C "${SCRIPT_DIR}" recipe scripts | kubectl -n "${NAMESPACE}" exec -i "${HEAD}" -- tar xf - -C "${REMOTE_DIR}" +tar cf - -C "$(cd "$(dirname "${ENV_FILE}")" && pwd)" "$(basename "${ENV_FILE}")" | kubectl -n "${NAMESPACE}" exec -i "${HEAD}" -- tar xf - -C "${REMOTE_DIR}" +if [[ "$(basename "${ENV_FILE}")" != "env_vars" ]]; then + kubectl -n "${NAMESPACE}" exec "${HEAD}" -- bash -lc "mv -f ${REMOTE_DIR}/$(basename "${ENV_FILE}") ${REMOTE_DIR}/env_vars" +fi + +# ---- run: exec the recipe inside the head pod ---- +echo "[INFO] launching recipe inside ${HEAD}; ray job submit targets the head's own dashboard (127.0.0.1:8265)." +kubectl -n "${NAMESPACE}" exec "${HEAD}" -- bash -lc "cd ${REMOTE_DIR} && ENV_FILE=${REMOTE_DIR}/env_vars bash recipe/${RECIPE}" From bbce9e7d88b7b789d4a63b49b52ac555dcd0cbef Mon Sep 17 00:00:00 2001 From: littlemex Date: Tue, 18 Aug 2026 19:58:54 +0900 Subject: [PATCH 4/7] feat(miles): checkpoint conversion and AIME evaluation HF<->Megatron conversion and AIME-2024 evaluation. pass@k keyed per prompt; boxed-answer extraction counts braces; requests run with bounded concurrency under a generation-sized timeout so a queued timeout is not scored as an incorrect answer. --- .../miles/scripts/convert_checkpoint.sh | 126 ++++++++ .../pytorch/miles/scripts/evaluate.sh | 297 ++++++++++++++++++ 2 files changed, 423 insertions(+) create mode 100644 3.test_cases/pytorch/miles/scripts/convert_checkpoint.sh create mode 100644 3.test_cases/pytorch/miles/scripts/evaluate.sh diff --git a/3.test_cases/pytorch/miles/scripts/convert_checkpoint.sh b/3.test_cases/pytorch/miles/scripts/convert_checkpoint.sh new file mode 100644 index 000000000..35ed19ed0 --- /dev/null +++ b/3.test_cases/pytorch/miles/scripts/convert_checkpoint.sh @@ -0,0 +1,126 @@ +#!/usr/bin/env bash +# STATUS: UNVERIFIED -- mirrors slime scripts/convert_checkpoint.sh; not executed on miles. Related: HF<->Megatron round-trip untested; see README Known Issues (save_model pickle-truncation). +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +# ============================================================ +# Checkpoint Conversion Helper +# +# Converts between HuggingFace and Megatron torch_dist formats. +# +# Usage: +# # HuggingFace -> Megatron (required before training) +# bash scripts/convert_checkpoint.sh hf2megatron \ +# --model-script qwen3-4B.sh \ +# --hf-path /fsx/models/Qwen3-4B \ +# --save-path /fsx/models/Qwen3-4B_torch_dist +# +# # Megatron -> HuggingFace (after training, for evaluation) +# bash scripts/convert_checkpoint.sh megatron2hf \ +# --input-dir /fsx/checkpoints/qwen3-4b-grpo/iter_0060/ \ +# --output-dir /fsx/models/Qwen3-4B-GRPO-step60 \ +# --origin-hf-dir /fsx/models/Qwen3-4B +# ============================================================ + +set -euo pipefail + +DIRECTION="${1:-}" +shift || true + +if [[ -z "${DIRECTION}" ]]; then + echo "Usage: $0 {hf2megatron|megatron2hf} [options]" + exit 1 +fi + +SLIME_DIR="${SLIME_DIR:-/root/miles}" +MEGATRON_DIR="${MEGATRON_DIR:-/root/Megatron-LM}" + +case "${DIRECTION}" in + hf2megatron) + MODEL_SCRIPT="" + HF_PATH="" + SAVE_PATH="" + NUM_GPUS=1 + + while [[ $# -gt 0 ]]; do + case "$1" in + --model-script) MODEL_SCRIPT="$2"; shift 2 ;; + --hf-path) HF_PATH="$2"; shift 2 ;; + --save-path) SAVE_PATH="$2"; shift 2 ;; + --num-gpus) NUM_GPUS="$2"; shift 2 ;; + *) echo "Unknown option: $1"; exit 1 ;; + esac + done + + if [[ -z "${MODEL_SCRIPT}" || -z "${HF_PATH}" || -z "${SAVE_PATH}" ]]; then + echo "Required: --model-script, --hf-path, --save-path" + exit 1 + fi + + echo "[INFO] Converting HuggingFace -> Megatron torch_dist" + echo " Model script: ${MODEL_SCRIPT}" + echo " HF path: ${HF_PATH}" + echo " Save path: ${SAVE_PATH}" + echo " GPUs: ${NUM_GPUS}" + + cd "${SLIME_DIR}" + source "scripts/models/${MODEL_SCRIPT}" + + if [[ ${NUM_GPUS} -gt 1 ]]; then + PYTHONPATH="${MEGATRON_DIR}" torchrun \ + --nproc_per_node="${NUM_GPUS}" \ + tools/convert_hf_to_torch_dist.py \ + "${MODEL_ARGS[@]}" \ + --hf-checkpoint "${HF_PATH}" \ + --save "${SAVE_PATH}" + else + PYTHONPATH="${MEGATRON_DIR}" python3 \ + tools/convert_hf_to_torch_dist.py \ + "${MODEL_ARGS[@]}" \ + --hf-checkpoint "${HF_PATH}" \ + --save "${SAVE_PATH}" + fi + + echo "[INFO] Conversion complete: ${SAVE_PATH}" + ;; + + megatron2hf) + INPUT_DIR="" + OUTPUT_DIR="" + ORIGIN_HF_DIR="" + + while [[ $# -gt 0 ]]; do + case "$1" in + --input-dir) INPUT_DIR="$2"; shift 2 ;; + --output-dir) OUTPUT_DIR="$2"; shift 2 ;; + --origin-hf-dir) ORIGIN_HF_DIR="$2"; shift 2 ;; + *) echo "Unknown option: $1"; exit 1 ;; + esac + done + + if [[ -z "${INPUT_DIR}" || -z "${OUTPUT_DIR}" || -z "${ORIGIN_HF_DIR}" ]]; then + echo "Required: --input-dir, --output-dir, --origin-hf-dir" + exit 1 + fi + + echo "[INFO] Converting Megatron torch_dist -> HuggingFace" + echo " Input: ${INPUT_DIR}" + echo " Output: ${OUTPUT_DIR}" + echo " Origin HF: ${ORIGIN_HF_DIR}" + + cd "${SLIME_DIR}" + + PYTHONPATH="${MEGATRON_DIR}" python3 \ + tools/convert_torch_dist_to_hf.py \ + --input-dir "${INPUT_DIR}" \ + --output-dir "${OUTPUT_DIR}" \ + --origin-hf-dir "${ORIGIN_HF_DIR}" + + echo "[INFO] Conversion complete: ${OUTPUT_DIR}" + ;; + + *) + echo "Unknown direction: ${DIRECTION}" + echo "Usage: $0 {hf2megatron|megatron2hf} [options]" + exit 1 + ;; +esac diff --git a/3.test_cases/pytorch/miles/scripts/evaluate.sh b/3.test_cases/pytorch/miles/scripts/evaluate.sh new file mode 100644 index 000000000..33720ad15 --- /dev/null +++ b/3.test_cases/pytorch/miles/scripts/evaluate.sh @@ -0,0 +1,297 @@ +#!/usr/bin/env bash +# STATUS: UNVERIFIED -- mirrors slime scripts/evaluate.sh; not executed on miles. +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +# ============================================================ +# Evaluation Script for miles-trained Models +# +# Evaluates a HuggingFace-format checkpoint on AIME-2024 using +# SGLang for inference, then scores with the math reward function. +# +# Usage: +# bash scripts/evaluate.sh \ +# --model-path /fsx/models/Qwen3-4B-GRPO-step60 \ +# --eval-data /fsx/data/aime-2024/aime-2024.jsonl \ +# --num-samples 16 \ +# --tp-size 2 \ +# --max-tokens 16384 +# ============================================================ + +set -euo pipefail + +# Defaults +MODEL_PATH="" +EVAL_DATA="/fsx/data/aime-2024/aime-2024.jsonl" +SERVER_PORT="${SERVER_PORT:-30000}" +NUM_SAMPLES=16 +TP_SIZE=2 +MAX_TOKENS=16384 +TEMPERATURE=0.6 +TOP_P=0.95 +OUTPUT_DIR="/fsx/eval_results" + +while [[ $# -gt 0 ]]; do + case "$1" in + --model-path) MODEL_PATH="$2"; shift 2 ;; + --eval-data) EVAL_DATA="$2"; shift 2 ;; + --num-samples) NUM_SAMPLES="$2"; shift 2 ;; + --tp-size) TP_SIZE="$2"; shift 2 ;; + --max-tokens) MAX_TOKENS="$2"; shift 2 ;; + --temperature) TEMPERATURE="$2"; shift 2 ;; + --top-p) TOP_P="$2"; shift 2 ;; + --output-dir) OUTPUT_DIR="$2"; shift 2 ;; + *) echo "Unknown option: $1"; exit 1 ;; + esac +done + +if [[ -z "${MODEL_PATH}" ]]; then + echo "Usage: $0 --model-path [--eval-data ] [--num-samples N] ..." + exit 1 +fi + +MODEL_NAME="$(basename "${MODEL_PATH}")" +TIMESTAMP="$(date +%Y%m%d_%H%M%S)" +RESULT_DIR="${OUTPUT_DIR}/${MODEL_NAME}_${TIMESTAMP}" +mkdir -p "${RESULT_DIR}" + +echo "============================================================" +echo " miles Model Evaluation" +echo "============================================================" +echo " Model: ${MODEL_PATH}" +echo " Eval data: ${EVAL_DATA}" +echo " Samples: ${NUM_SAMPLES} per prompt" +echo " TP size: ${TP_SIZE}" +echo " Max tokens: ${MAX_TOKENS}" +echo " Output: ${RESULT_DIR}" +echo "============================================================" + +# ----- Step 1: Start SGLang server ----- +echo "[INFO] Starting SGLang server (TP=${TP_SIZE})..." +python3 -m sglang.launch_server \ + --model-path "${MODEL_PATH}" \ + --tp "${TP_SIZE}" \ + --host 0.0.0.0 \ + --port "${SERVER_PORT}" \ + --mem-fraction-static "${SGLANG_MEM_FRACTION:-0.85}" \ + --log-level warning & # lowercase: uvicorn KeyErrors on "WARN" + +SGLANG_PID=$! + +# Reap the server on ANY exit, not just the happy path. `set -e` is active, so a non-zero +# exit from the evaluation step below (the high-error-rate abort, or any Python exception) +# would otherwise skip the cleanup at the end of the script and leave a multi-GPU SGLang +# server holding the whole node's memory until someone notices. +cleanup_sglang() { + kill "${SGLANG_PID}" 2>/dev/null || true + wait "${SGLANG_PID}" 2>/dev/null || true +} +trap cleanup_sglang EXIT + +# Wait for server to be ready +SGLANG_STARTUP_TIMEOUT=${SGLANG_STARTUP_TIMEOUT:-300} +echo "[INFO] Waiting for SGLang server to start (timeout=${SGLANG_STARTUP_TIMEOUT}s)..." +for i in $(seq 1 ${SGLANG_STARTUP_TIMEOUT}); do + if curl -s "http://localhost:${SERVER_PORT}/health" > /dev/null 2>&1; then + echo "[INFO] SGLang server ready." + break + fi + if [[ $i -eq ${SGLANG_STARTUP_TIMEOUT} ]]; then + echo "[ERROR] SGLang server failed to start within ${SGLANG_STARTUP_TIMEOUT} seconds." + kill ${SGLANG_PID} 2>/dev/null || true + exit 1 + fi + sleep 1 +done + +# ----- Step 2: Run evaluation ----- +echo "[INFO] Running evaluation..." +# The heredoc below is quoted, so the Python reads these through the environment rather +# than through shell expansion. They must be exported: a plain assignment stays in the +# shell and the Python silently falls back to its defaults, which would produce +# confident-looking results for parameters the caller never asked for. +export EVAL_DATA NUM_SAMPLES MAX_TOKENS TEMPERATURE TOP_P RESULT_DIR SERVER_PORT +export EVAL_MAX_CONCURRENCY="${EVAL_MAX_CONCURRENCY:-32}" +export EVAL_REQUEST_TIMEOUT="${EVAL_REQUEST_TIMEOUT:-1800}" +export EVAL_ERROR_FRACTION_ABORT="${EVAL_ERROR_FRACTION_ABORT:-0.05}" + +python3 - <<'EVAL_SCRIPT' +import json +import sys +import os +import re +import asyncio +import aiohttp + +EVAL_DATA = os.environ.get("EVAL_DATA", "/fsx/data/aime-2024/aime-2024.jsonl") +NUM_SAMPLES = int(os.environ.get("NUM_SAMPLES", "16")) +MAX_TOKENS = int(os.environ.get("MAX_TOKENS", "16384")) +TEMPERATURE = float(os.environ.get("TEMPERATURE", "0.6")) +TOP_P = float(os.environ.get("TOP_P", "0.95")) +RESULT_DIR = os.environ.get("RESULT_DIR", "/fsx/eval_results") +SERVER_URL = f"http://localhost:{os.environ.get('SERVER_PORT', '30000')}/v1/chat/completions" +# Bounded concurrency and a timeout sized for the generation, not for the queue. With +# MAX_TOKENS=16384 a single response can take minutes, so a 300s cap applied to every +# request at once made most of them time out and be tallied as wrong answers. +MAX_CONCURRENCY = int(os.environ.get("EVAL_MAX_CONCURRENCY", "32")) +REQUEST_TIMEOUT = float(os.environ.get("EVAL_REQUEST_TIMEOUT", "1800")) +ERROR_FRACTION_ABORT = float(os.environ.get("EVAL_ERROR_FRACTION_ABORT", "0.05")) + +# Load evaluation prompts +prompts = [] +with open(EVAL_DATA, "r") as f: + for line in f: + item = json.loads(line.strip()) + prompts.append(item) + +print(f"Loaded {len(prompts)} evaluation prompts") + +def extract_boxed(text): + r"""Return the content of the last \boxed{...}, honoring nested braces. + + Scanned with a brace counter rather than a regex: a character class cannot match the + balanced braces in answers like `\boxed{\frac{1}{2}}`. + """ + out = [] + needle = r"\boxed{" + start = text.find(needle) + while start != -1: + i = start + len(needle) + depth = 1 + while i < len(text) and depth: + if text[i] == "{": + depth += 1 + elif text[i] == "}": + depth -= 1 + if depth == 0: + break + i += 1 + if depth == 0: + out.append(text[start + len(needle):i]) + start = text.find(needle, start + len(needle)) + return out[-1].strip() if out else "" + + +async def evaluate_prompt(session, prompt_item, prompt_idx, sample_idx, sem): + """Generate a response and check correctness.""" + messages = [{"role": "user", "content": prompt_item.get("prompt", prompt_item.get("question", ""))}] + + payload = { + "model": "default", + "messages": messages, + "max_tokens": MAX_TOKENS, + "temperature": TEMPERATURE, + "top_p": TOP_P, + } + + try: + # Bound the in-flight requests. Submitting every prompt x sample at once makes most + # requests spend their timeout queued rather than generating, and the tally below + # counts a timeout as a wrong answer -- so accuracy drops for a reason unrelated to + # the model. + async with sem: + async with session.post(SERVER_URL, json=payload, + timeout=aiohttp.ClientTimeout(total=REQUEST_TIMEOUT)) as resp: + result = await resp.json() + response_text = result["choices"][0]["message"]["content"] + + predicted = extract_boxed(response_text) + # A label may be a number in the JSONL, and str.strip() on an int raises + # AttributeError -- which the except below would turn into "incorrect" for + # every sample. + label = str(prompt_item.get("label", prompt_item.get("answer", ""))).strip() + + return { + "prompt_idx": prompt_idx, + "sample_idx": sample_idx, + "predicted": predicted, + "label": label, + "correct": predicted == label if predicted else False, + "response_length": len(response_text), + } + except Exception as e: + return { + "prompt_idx": prompt_idx, + "sample_idx": sample_idx, + "predicted": "", + "label": str(prompt_item.get("label", "")), + "correct": False, + "error": f"{type(e).__name__}: {e}", + } + +async def main(): + results = [] + sem = asyncio.Semaphore(MAX_CONCURRENCY) + async with aiohttp.ClientSession() as session: + tasks = [] + # pass@k is keyed on the loop index, not on a field in the data: AIME-2024 as prepared + # here carries only {"prompt", "label"}, so keying on a missing "idx" would collapse + # every prompt onto one bucket and turn pass@k into "any sample anywhere was right". + for prompt_idx, prompt_item in enumerate(prompts): + for s in range(NUM_SAMPLES): + tasks.append(evaluate_prompt(session, prompt_item, prompt_idx, s, sem)) + + print(f"Evaluating {len(tasks)} total samples " + f"({MAX_CONCURRENCY} concurrent, {REQUEST_TIMEOUT}s timeout each)...") + results = await asyncio.gather(*tasks) + + # Compute metrics + total = len(results) + correct = sum(1 for r in results if r.get("correct", False)) + errors = sum(1 for r in results if r.get("error")) + accuracy = correct / total if total > 0 else 0 + + # Per-prompt pass@k (at least one correct) + from collections import defaultdict + prompt_results = defaultdict(list) + for r in results: + prompt_results[r["prompt_idx"]].append(r.get("correct", False)) + + pass_at_k = sum(1 for prs in prompt_results.values() if any(prs)) / len(prompt_results) if prompt_results else 0 + + print(f"\n{'='*60}") + print(f" Evaluation Results") + print(f"{'='*60}") + print(f" Total samples: {total}") + print(f" Correct: {correct}") + print(f" Errors: {errors}") + print(f" Accuracy: {accuracy:.4f}") + print(f" Pass@{NUM_SAMPLES}: {pass_at_k:.4f}") + print(f" Prompts evaluated:{len(prompt_results)}") + print(f"{'='*60}") + if len(prompt_results) != len(prompts): + print(f" WARNING: {len(prompts)} prompts were loaded but only " + f"{len(prompt_results)} distinct prompt indices appear in the results.") + + # Save results + output_file = os.path.join(RESULT_DIR, "eval_results.json") + with open(output_file, "w") as f: + json.dump({ + "metrics": { + "total_samples": total, + "correct": correct, + "errors": errors, + "accuracy": accuracy, + "pass_at_k": pass_at_k, + "k": NUM_SAMPLES, + "prompts": len(prompt_results), + }, + "results": results, + }, f, indent=2) + print(f" Results saved to: {output_file}") + + # A run where a large share of requests failed has not measured accuracy, it has + # measured the timeout. Exiting non-zero keeps that out of a results table. + if total and errors / total > ERROR_FRACTION_ABORT: + print(f"\nERROR: {errors}/{total} requests failed " + f"(> {ERROR_FRACTION_ABORT:.0%}). These count as incorrect, so the accuracy " + "above understates the model. Raise SGLANG_MEM_FRACTION, lower " + "EVAL_MAX_CONCURRENCY, or raise EVAL_REQUEST_TIMEOUT, then re-run.") + sys.exit(2) + +asyncio.run(main()) +EVAL_SCRIPT + +# ----- Step 3: Cleanup ----- +# Killing and reaping is the EXIT trap's job, so there is deliberately no kill/wait here: a +# bare `wait` would block on a server that is still running and nothing would ever stop it. +echo "[INFO] Evaluation complete; stopping SGLang server." From 09af28dda60bff738e240bf99c0d43c82db6d784 Mon Sep 17 00:00:00 2001 From: littlemex Date: Tue, 18 Aug 2026 19:58:55 +0900 Subject: [PATCH 5/7] feat(miles): optional disaggregated reward service A CPU reward-service Deployment and image for the disaggregated overlay, mirroring the sibling slime test case. Provided as a reference; not exercised in the recorded runs. --- .../miles/kubernetes/reward-service.yaml | 136 ++++++++++++++ .../pytorch/miles/reward_service.Dockerfile | 25 +++ .../pytorch/miles/reward_service/app.py | 174 ++++++++++++++++++ .../miles/reward_service/requirements.txt | 14 ++ 4 files changed, 349 insertions(+) create mode 100644 3.test_cases/pytorch/miles/kubernetes/reward-service.yaml create mode 100644 3.test_cases/pytorch/miles/reward_service.Dockerfile create mode 100644 3.test_cases/pytorch/miles/reward_service/app.py create mode 100644 3.test_cases/pytorch/miles/reward_service/requirements.txt diff --git a/3.test_cases/pytorch/miles/kubernetes/reward-service.yaml b/3.test_cases/pytorch/miles/kubernetes/reward-service.yaml new file mode 100644 index 000000000..49d38e70a --- /dev/null +++ b/3.test_cases/pytorch/miles/kubernetes/reward-service.yaml @@ -0,0 +1,136 @@ +# STATUS: UNVERIFIED -- mirrors slime reward-service.yaml (remote_rm on CPU pool); not deployed on miles. miles remote_rm lacks slime's retry. +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +# +# Remote reward service for miles GRPO (rm_type=remote_rm). +# +# Runs on a same-AZ CPU instance group -- NOT on the GPU nodes. This frees the +# scarce GPU rollout/training actors from CPU-bound reward scoring and lets the +# reward layer scale independently (bump `replicas`). The reward RPC is +# low-bandwidth HTTP, so this group does NOT need EFA. +apiVersion: apps/v1 +kind: Deployment +metadata: + name: miles-reward + namespace: ${NAMESPACE} + labels: + app: miles-reward +spec: + replicas: ${REWARD_REPLICAS} + # Roll new pods before tearing old ones down so reward capacity never drops to + # zero -- important on Spot, where nodes can be reclaimed at any time. + strategy: + type: RollingUpdate + rollingUpdate: + maxUnavailable: 1 + maxSurge: 1 + selector: + matchLabels: + app: miles-reward + template: + metadata: + labels: + app: miles-reward + spec: + # Pin to a specific HyperPod CPU instance group (e.g. the c5 compute pool) + # via REWARD_NODE_GROUP, and additionally guard against GPU nodes with the + # gpu-feature-discovery label so the reward pods never land on p5 nodes. + nodeSelector: + sagemaker.amazonaws.com/instance-group-name: ${REWARD_NODE_GROUP} + affinity: + nodeAffinity: + requiredDuringSchedulingIgnoredDuringExecution: + nodeSelectorTerms: + - matchExpressions: + - key: nvidia.com/gpu.present + operator: NotIn + values: + - "true" + topologySpreadConstraints: + # Spread replicas across the CPU nodes so the reward tier uses the whole + # pool rather than stacking on one node. On Spot this also limits the + # blast radius if a single node is reclaimed. + - maxSkew: 1 + topologyKey: kubernetes.io/hostname + whenUnsatisfiable: ScheduleAnyway + labelSelector: + matchLabels: + app: miles-reward + # The reward server is stateless; this is enough time to drain in-flight + # /score requests when HyperPod gives a Spot interruption notice (it taints + # the node and gracefully evicts within terminationGracePeriodSeconds). + terminationGracePeriodSeconds: 30 + # HyperPod taints interrupted Spot nodes; tolerate briefly so in-flight + # requests can finish before the pod is evicted. + tolerations: + - key: "sagemaker.amazonaws.com/node-health-status" + operator: "Exists" + effect: "NoExecute" + tolerationSeconds: 30 + containers: + - name: reward + image: ${REWARD_IMAGE} + # Image tag is immutable/versioned (REWARD_TAG), so avoid re-pulling + # on every pod start (fast Spot recovery, air-gap friendly). + imagePullPolicy: IfNotPresent + env: + - name: REWARD_BACKEND + value: "${REWARD_BACKEND}" + - name: REWARD_MODEL_NAME + value: "${REWARD_MODEL_NAME}" + - name: TORCH_NUM_THREADS + value: "${REWARD_TORCH_THREADS}" + # Create the secret with: + # kubectl create secret generic hf-token --from-literal=HF_TOKEN=$HF_TOKEN -n $NAMESPACE + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token + key: HF_TOKEN + ports: + - containerPort: 8000 + resources: + requests: + cpu: "6" + memory: "12Gi" + limits: + cpu: "8" + memory: "16Gi" + # startupProbe gates liveness until the reward model has finished its + # one-time CPU download+load (cold HF cache on a freshly reclaimed Spot + # node can take minutes). Without it, livenessProbe would race the cold + # load and crashloop during the very Spot churn this design rides out. + startupProbe: + httpGet: + path: /health + port: 8000 + periodSeconds: 10 + failureThreshold: 30 # up to ~5 min for a cold model download/load + readinessProbe: + httpGet: + path: /health + port: 8000 + initialDelaySeconds: 20 + periodSeconds: 10 + livenessProbe: + httpGet: + path: /health + port: 8000 + initialDelaySeconds: 60 + periodSeconds: 30 +--- +apiVersion: v1 +kind: Service +metadata: + name: miles-reward + namespace: ${NAMESPACE} + labels: + app: miles-reward +spec: + selector: + app: miles-reward + ports: + - protocol: TCP + port: 8000 + targetPort: 8000 + type: ClusterIP diff --git a/3.test_cases/pytorch/miles/reward_service.Dockerfile b/3.test_cases/pytorch/miles/reward_service.Dockerfile new file mode 100644 index 000000000..89419923a --- /dev/null +++ b/3.test_cases/pytorch/miles/reward_service.Dockerfile @@ -0,0 +1,25 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +# +# CPU-only image for the miles remote reward service. Deliberately lightweight +# (no CUDA / NGC base) so it schedules on cheap CPU instances and starts fast. +FROM python:3.12-slim + +ENV PYTHONUNBUFFERED=1 \ + HF_HOME=/opt/hf-cache \ + TOKENIZERS_PARALLELISM=false + +WORKDIR /opt/reward_service + +# Build tools needed by sentencepiece/tokenizers wheels on slim. +RUN apt-get update && apt-get install -y --no-install-recommends \ + build-essential \ + && rm -rf /var/lib/apt/lists/* + +COPY reward_service/requirements.txt /tmp/requirements.txt +RUN pip install --no-cache-dir -r /tmp/requirements.txt + +COPY reward_service/app.py /opt/reward_service/app.py + +EXPOSE 8000 +CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000", "--workers", "1"] diff --git a/3.test_cases/pytorch/miles/reward_service/app.py b/3.test_cases/pytorch/miles/reward_service/app.py new file mode 100644 index 000000000..d01699b7a --- /dev/null +++ b/3.test_cases/pytorch/miles/reward_service/app.py @@ -0,0 +1,174 @@ +# STATUS: UNVERIFIED -- UNVERIFIED on miles -- reward payload contract matches slime, but the remote_rm path was not exercised with miles. +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +""" +Remote Reward Service for miles GRPO training. + +This is a CPU-hosted HTTP reward server that implements the contract expected by +SLIME's ``remote_rm`` hook (slime/rollout/rm_hub/__init__.py): + + POST {RM_URL} + body: {"prompt": str, "response": str, "label": str | null} + returns: a bare JSON number (the scalar reward), which SLIME assigns + directly to ``sample.reward``. + +Why run this off the GPU nodes? + The GPU rollout engines (SGLang) and trainers (Megatron) are the expensive, + scarce resource. Scoring should not steal their CPU. A reward *model* (a + small sequence classifier) or any heavy verifier (code-exec, RAG, unit + tests) is CPU/IO-bound and latency-tolerant, so it belongs on a cheap, + independently-scalable CPU instance group in the same AZ. The reward RPC is + low-bandwidth HTTP and does NOT use EFA/RDMA. + +Backends (select with REWARD_BACKEND): + - "reward_model" (default): a HuggingFace AutoModelForSequenceClassification + that scores (prompt, response) pairs on CPU. This is the case where the + CPU offload is a genuine throughput win. + - "math_verify": rule-based LaTeX/sympy verification against ``label``; + useful as a zero-dependency fallback / for math datasets. +""" + +import logging +import os +from typing import Optional + +from fastapi import FastAPI +from pydantic import BaseModel + +logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") +logger = logging.getLogger("reward_service") + +REWARD_BACKEND = os.environ.get("REWARD_BACKEND", "reward_model").strip() +REWARD_MODEL_NAME = os.environ.get( + "REWARD_MODEL_NAME", "OpenAssistant/reward-model-deberta-v3-large-v2" +) +# Cap intra-op threads so a single replica does not monopolise the node; scale +# horizontally with replicas instead. +TORCH_NUM_THREADS = int(os.environ.get("TORCH_NUM_THREADS", "4")) +MAX_LENGTH = int(os.environ.get("REWARD_MAX_LENGTH", "2048")) + + +class ScoreRequest(BaseModel): + prompt: str | list = "" + response: str = "" + label: Optional[str] = None + + +# --------------------------------------------------------------------------- # +# Pluggable scorer backends +# --------------------------------------------------------------------------- # +class Scorer: + """Backend interface: score a single (prompt, response, label) -> float.""" + + def score(self, prompt: str, response: str, label: Optional[str]) -> float: + raise NotImplementedError + + +class RewardModelScorer(Scorer): + """HuggingFace sequence-classifier reward model running on CPU.""" + + def __init__(self, model_name: str): + import torch + from transformers import AutoModelForSequenceClassification, AutoTokenizer + + torch.set_num_threads(TORCH_NUM_THREADS) + self.torch = torch + logger.info("Loading reward model %s on CPU ...", model_name) + self.tokenizer = AutoTokenizer.from_pretrained(model_name) + self.model = AutoModelForSequenceClassification.from_pretrained( + model_name, torch_dtype=torch.float32 + ) + self.model.eval() + logger.info("Reward model loaded.") + + def _to_text(self, prompt) -> str: + if isinstance(prompt, list): + # chat-format prompt: concatenate message contents + return "\n".join( + m.get("content", "") for m in prompt if isinstance(m, dict) + ) + return prompt or "" + + def score(self, prompt, response: str, label: Optional[str]) -> float: + prompt_text = self._to_text(prompt) + with self.torch.no_grad(): + inputs = self.tokenizer( + prompt_text, + response, + return_tensors="pt", + truncation=True, + max_length=MAX_LENGTH, + ) + logits = self.model(**inputs).logits + # Single-logit reward models output a scalar score; multi-class + # models -> take the positive/last class logit. + score = logits.squeeze(-1) if logits.shape[-1] == 1 else logits[..., -1] + return float(score.reshape(-1)[0].item()) + + +class MathVerifyScorer(Scorer): + """Rule-based math verification (LaTeX \\boxed{} + sympy) against label.""" + + def __init__(self): + from math_verify import parse, verify + + self._parse = parse + self._verify = verify + + def score(self, prompt, response: str, label: Optional[str]) -> float: + if not label: + return 0.0 + try: + # parsing_timeout=None: math_verify's default timeout uses + # signal.alarm(), which only works on the main thread. FastAPI runs + # sync handlers in a threadpool, so disable the signal-based timeout. + gold = self._parse( + label if "\\boxed" in str(label) else f"\\boxed{{{label}}}", + parsing_timeout=None, + ) + pred = self._parse(response, parsing_timeout=None) + return 1.0 if self._verify(gold, pred, timeout_seconds=None) else 0.0 + except Exception as e: # noqa: BLE001 - verifier must never crash the service + logger.warning("math_verify scoring error: %s", e) + return 0.0 + + +def _build_scorer() -> Scorer: + if REWARD_BACKEND == "math_verify": + logger.info("Using math_verify backend.") + return MathVerifyScorer() + logger.info("Using reward_model backend: %s", REWARD_MODEL_NAME) + return RewardModelScorer(REWARD_MODEL_NAME) + + +app = FastAPI(title="miles Remote Reward Service") +_scorer: Optional[Scorer] = None + + +@app.on_event("startup") +def _startup(): + global _scorer + _scorer = _build_scorer() + + +@app.get("/health") +def health(): + return {"status": "ok", "backend": REWARD_BACKEND} + + +@app.post("/score") +def score(req: ScoreRequest): + # SLIME's remote_rm assigns the returned JSON value directly to + # sample.reward, so we return a bare float. + # + # Never raise: SLIME's remote_rm client calls resp.raise_for_status(), + # retries up to 10x with backoff, then re-raises; batched_async_rm gathers + # without return_exceptions, so a single failed /score (e.g. an odd + # truncation, unexpected logits shape, or CPU OOM in the reward model) would + # take down the whole rollout step. Guard the handler and return a neutral + # reward instead, mirroring the rule-based backend's defensive behavior. + try: + return _scorer.score(req.prompt, req.response, req.label) + except Exception as e: # noqa: BLE001 - must never 500 the rollout loop + logger.warning("scoring error (returning neutral reward 0.0): %s", e) + return 0.0 diff --git a/3.test_cases/pytorch/miles/reward_service/requirements.txt b/3.test_cases/pytorch/miles/reward_service/requirements.txt new file mode 100644 index 000000000..6abfb1782 --- /dev/null +++ b/3.test_cases/pytorch/miles/reward_service/requirements.txt @@ -0,0 +1,14 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +# +# Pinned dependencies for the CPU reward service. CPU-only torch keeps the +# image small (no CUDA). Versions are pinned for reproducible builds. +--extra-index-url https://download.pytorch.org/whl/cpu +torch==2.5.1 +transformers==4.46.3 +fastapi==0.115.5 +uvicorn[standard]==0.32.1 +pydantic==2.10.2 +math_verify==0.9.0 +sentencepiece==0.2.1 +protobuf==5.29.6 From ec59e119427384c3ce4d51788f09895b9345db20 Mon Sep 17 00:00:00 2001 From: littlemex Date: Tue, 18 Aug 2026 19:58:57 +0900 Subject: [PATCH 6/7] docs(miles): README, porting notes, EFA and verification logs README with prerequisites, walkthrough (incl. in-cluster BuildKit and an optional head-pod launch path), verification status, and known issues. The Ray head co-locates on the GPU pool (num-gpus 0) because miles control actors import Megatron and need libcuda; a CPU-only head is not safe. VERIFICATION_LOG records what ran, on what hardware, the metrics, the 30B MoE FlashInfer-fusion root cause, and why the head must run on a CUDA-capable node. --- 3.test_cases/pytorch/miles/README.md | 883 ++++++++++++++++++ 3.test_cases/pytorch/miles/docs/EFA_2NODE.md | 82 ++ 3.test_cases/pytorch/miles/docs/PORT_NOTES.md | 115 +++ .../pytorch/miles/docs/VERIFICATION_LOG.md | 322 +++++++ 4 files changed, 1402 insertions(+) create mode 100644 3.test_cases/pytorch/miles/README.md create mode 100644 3.test_cases/pytorch/miles/docs/EFA_2NODE.md create mode 100644 3.test_cases/pytorch/miles/docs/PORT_NOTES.md create mode 100644 3.test_cases/pytorch/miles/docs/VERIFICATION_LOG.md diff --git a/3.test_cases/pytorch/miles/README.md b/3.test_cases/pytorch/miles/README.md new file mode 100644 index 000000000..836962fd8 --- /dev/null +++ b/3.test_cases/pytorch/miles/README.md @@ -0,0 +1,883 @@ +# Reinforcement Learning with miles on Amazon SageMaker HyperPod EKS + +This test case runs GRPO post-training with [**miles**](https://github.com/radixark/miles) on +[Amazon SageMaker HyperPod](https://aws.amazon.com/sagemaker/hyperpod/) with Amazon EKS +orchestration. It mirrors the sibling [`3.test_cases/pytorch/slime/`](../slime/) test case +end to end -- building a container image, preparing data, deploying a multi-node Ray cluster, +converting model checkpoints, and launching GRPO training on NVIDIA GPUs interconnected with +Elastic Fabric Adapter (EFA) -- but targets **miles**, a fork of SLIME built for CUDA 13 and +NVIDIA Blackwell as first-class hardware. + +The hardware-verified runs recorded here were done on a plain Amazon EKS cluster with KubeRay, +EFA, and FSx for Lustre. The same manifests are expected to run unchanged on a SageMaker +HyperPod EKS cluster -- which adds deep health checks and automatic node replacement on top of +the same EKS/KubeRay/EFA stack -- but that was not separately validated. + +## Introduction + +[**miles**](https://github.com/radixark/miles) is a direct fork of +[**SLIME**](https://github.com/THUDM/slime) (an LLM post-training framework for RL scaling). +It inherits SLIME's core design -- integrating two best-of-breed backends under Ray +orchestration: + +- **[SGLang](https://github.com/sgl-project/sglang)** for high-throughput rollout generation + (inference), providing RadixAttention, continuous batching, and tensor parallelism. +- **[Megatron-LM](https://github.com/NVIDIA/Megatron-LM)** (a radixark fork) for scalable + distributed training, with support for TP, PP, CP, EP, and ZeRO-style sharding. + +**Ray** manages resource orchestration, supporting both **colocated** (training and rollout +share the same GPUs, time-sliced) and **disaggregated** (separate GPU pools connected by +NCCL/EFA weight sync) deployment topologies -- the same two topologies SLIME supports. + +Note that "disaggregated" appears here in two unrelated senses, and the files use both. +`COLOCATE=false` splits the **actor and rollout GPU pools**, which is what selects the weight +sync implementation. `env_vars.disaggregated.example` does something else entirely: it moves +**reward scoring** to a CPU pool and leaves the GPU layout alone. The two are independent and +can be combined. + +What miles changes relative to upstream SLIME is the platform target: it ships a matched +PyTorch 2.11 / CUDA 13.0.1 stack with prebuilt `flash_attn` / Transformer Engine / `apex` +wheels and a Blackwell (sm_103) Transformer Engine FA2 whitelist patch, and it rewrites the +rollout/training weight-sync path from HTTP endpoints on an SGLang fork to direct Ray actor +methods. See [Why miles](#why-miles-vs-slime) below. + +**Amazon SageMaker HyperPod** provides purpose-built infrastructure for distributed model +training with deep health checks, automatic node replacement, and managed Kubernetes (EKS) +integration. Combined with FSx for Lustre shared storage and EFA networking, HyperPod +delivers the resilient, high-performance fabric that large-scale RL workloads demand. + +## Why miles vs SLIME + +| Aspect | SLIME | miles | +|--------|-------|-------| +| Target CUDA / GPU generation | NGC-image CUDA stack, Hopper-first | CUDA 13.0.1 native, NVIDIA Blackwell (sm_103) first-class | +| Base image | NGC PyTorch (nightly ABI) | `radixark/miles` (matched PyTorch 2.11 + cu130 stable ABI) | +| Rollout <-> training weight sync | HTTP endpoints on an SGLang fork | Ray actor methods (`begin_weight_update` / `pull_weights`) on the rollout engine directly | +| Framework install path (in image) | `/opt/slime` | `/root/miles` (editable install) | +| Megatron fork | NVIDIA/Megatron-LM (SLIME-compatible commit) | radixark/Megatron-LM fork | +| SGLang version pinned | 0.5.12.post1 | 0.5.16.dev | +| `train.py` CLI / GRPO flags | -- | Compatible with SLIME's (same names, same semantics) | +| Relationship | Upstream | Direct fork of SLIME (diverged at commit `fcce96ca0`, 2025-10-05) | + +Because miles is a fork rather than an independent reimplementation, the `train.py` CLI used +by this test case's recipes is drop-in compatible with SLIME's -- the [Quick +Start](#quick-start) below tracks the sibling slime test case step for step. + +## Architecture + +The deployment uses a colocated architecture where training and rollout share the same GPU +pool, connected by EFA networking and FSx for Lustre shared storage. (A disaggregated +topology -- separate GPU pools for training and rollout -- is also supported by miles/Ray but +is not the hardware-verified path documented here; see [Verification Status](#verification-status).) + +``` ++-----------------------------------------------------------------------+ +| Amazon SageMaker HyperPod EKS Cluster | +| | +| +-----------------------------+ +-----------------------------+ | +| | Node 1 (GPU worker) | | Node 2 (GPU worker) | | +| | 8x H-series GPU | EFA | | 8x H-series GPU | EFA | | +| +-----------------------------+ +-----------------------------+ | +| | | | | | +| +--------------EFA (GPU<->GPU RDMA)-----------+ | +| | | | +| +----------------------------------------------------------+ | +| | FSx for Lustre, mounted as PVC `fsx-claim` | | +| | /fsx/models /fsx/data /fsx/runs | | +| +----------------------------------------------------------+ | +| | +| +-----------------------------+ | +| | Ray head (num-gpus 0) | <- co-located on the GPU pool; | +| | on a GPU node (needs CUDA) | needs libcuda for Megatron impt | +| +-----------------------------+ | +| | +| Kubernetes Resources: | +| - KubeRay operator | +| - EFA device plugin, NVIDIA device plugin (kube-system) | +| - FSx CSI driver (kube-system) | ++-----------------------------------------------------------------------+ +``` + +### miles Internal Loop + +``` + +------------------+ + | Data Buffer | + | (prompt queue + | + | rollout cache) | + +--------+---------+ + | + +--------------+--------------+ + | | + v v + +-------------------+ +-------------------+ + | Rollout | | Training | + | (SGLang engines, | | (Megatron-LM | + | Ray actors) | <----> | TP/PP/CP/EP) | + | | weight | | + | - RadixAttention | sync | - GRPO | + | - Cont. batching | (Ray | - Dynamic batch | + | - TP per engine | actor | - Gradient ckpt | + | | calls) | | + +-------------------+ +-------------------+ +``` + +1. **Data Buffer** manages prompts, dispatches them for rollout, and stores generated samples + with rewards. +2. **Rollout** runs SGLang engines as Ray actors, generating responses and scoring them via a + reward function. Unlike SLIME, weight-sync calls (`begin_weight_update` / `pull_weights`) + go directly to Ray actor methods on the rollout engine rather than HTTP endpoints. +3. **Training** reads batches from the Data Buffer, computes GRPO advantages, and updates the + policy via Megatron-LM. Updated weights are synced back to the rollout engines. + +## Hardware Requirements + +miles's own image already builds for Hopper and Blackwell alike (CUDA 13.0.1, with the +sm_103 Transformer Engine FA2 whitelist patch applied on top of the standard sm_90 build), +and nothing in this test case's recipes, env files, or manifests hard-codes a GPU generation +or CUDA/SM version -- the only per-cluster knob is which accelerator NodePool to target +(`GPU_NODE_ROLE` in `env_vars`, substituted into `kubernetes/raycluster.yaml`'s +`nodeSelector`). The table below is the hardware this was actually run on; a B300 row is +included as the expected-compatible target (see Verification Status for what is and is not +confirmed on real hardware). + +| Component | H200 (verified) | B300 (expected-compatible, unverified) | +|-----------|------------------|------------------------------------------| +| **Instance type** | p5en.48xlarge | p6-b300.48xlarge | +| **GPUs per node** | 8x NVIDIA H200 (141 GB HBM) | 8x NVIDIA B300 (288 GB HBM) | +| **EFA per node** | 16 physical, **15 allocatable** on EKS | 16 physical, verify allocatable | +| **Storage** | FSx for Lustre, mounted as PVC `fsx-claim` | same | +| **Kubernetes** | EKS, KubeRay operator | same | +| **Ray head placement** | The head co-locates on the GPU pool (`num-gpus 0`); it must be on a CUDA-capable node (miles control actors import Megatron even at `num-gpus 0`). Give the GPU node's root volume >=150 GiB for the ~18 GB image | same | + +B300's ~2x larger HBM is expected to relax the memory-driven choices made for H200 in +`env_vars.colocated.example` (e.g. the 30B MoE case needing `--use-distributed-optimizer` and +a colocated 16-GPU layout to avoid OOM at 8 GPU -- see the ALTERNATE block in that file); it +should not require code changes, only re-tuning env values once measured on B300. + +### Verification Status + +This test case is presented as a GRPO reference, not a research study -- the table below +states plainly what has been run on real hardware and what has not, so you know exactly what +you are inheriting. + +Every row below is backed by a recorded run: see [docs/VERIFICATION_LOG.md](./docs/VERIFICATION_LOG.md) for the submission commands, Ray job ids, the flags each +job actually received, and the metric values read back from the trainer's event files. + +| Component | Status | Environment / Note | +|-----------|--------|--------------------| +| Qwen3-4B GRPO, colocated, 1 node (8 GPU) | Verified | p5en.48xlarge, H200x8 | +| Qwen3-4B GRPO, colocated, 2 nodes (16 GPU) over EFA, 3 rollout cycles | Verified | 2x p5en.48xlarge; all 3 cycles SUCCEEDED in 676s with `raw_reward` 0.477/0.523/0.492 -- metrics in [docs/VERIFICATION_LOG.md](./docs/VERIFICATION_LOG.md). Fabric separately measured at NCCL all_reduce busbw 190-257 GB/s over EFA (`efa-direct` + GPUDirect RDMA, no TCP fallback), see [docs/EFA_2NODE.md](./docs/EFA_2NODE.md) | +| Qwen3-30B-A3B MoE GRPO, colocated, 2 nodes (16 GPU) -- trains cleanly | Verified | actor spans all 16 GPU with `--use-distributed-optimizer` (required -- confining the actor to 8 GPU OOMs) and the SGLang `triton` MoE runner backend. The shipped block runs the rollout MoE pure expert-parallel (`moe_tp=1`, `EP_SIZE = ROLLOUT_GPUS_PER_ENGINE = 2`): `rollout/raw_reward` 0.578, `rollout/repetition_frac` 0.0, `weight_version` uniform / `mixed_version_ratio` 0.0 -- comparable to the dense 4B run. Running the MoE tensor-parallel AND expert-parallel at once (`moe_tp>1` and `moe_ep>1`) hits a FlashInfer allreduce-fusion bug in this build, so the recipe disables that fusion for such geometries and they train too -- see [Known Issues](#known-issues) item 2 | +| Qwen3-4B GRPO, **disaggregated** (`COLOCATE=false`), 2 nodes | Verified | run on 2x p5en (actor 8 + rollout 8): Ray job SUCCEEDED, `raw_reward` 0.52, `repetition_frac` 0.0, `weight_version` 2 with `mixed_version_ratio` 0.0 -- i.e. weights synced across the node boundary over NCCL/EFA (`UpdateWeightFromDistributed`). A disaggregated run needs GPU nodes for BOTH the actor and rollout pools, so `WORKER_REPLICAS` (derived in `env_vars` as `ceil((ACTOR_NUM_NODES*ACTOR_GPUS_PER_NODE + ROLLOUT_NUM_GPUS)/8)`, here 2) drives the RayCluster worker count rather than `ACTOR_NUM_NODES` alone (see [docs/VERIFICATION_LOG.md](./docs/VERIFICATION_LOG.md)) | +| Qwen3-30B-A3B MoE GRPO, **disaggregated** | UNVERIFIED | the disaggregated actor-8 layout needs B300-class 288GB HBM; only the colocated 16-GPU layout was run, on H200 | +| Any workload on B300 (p6-b300.48xlarge) | UNVERIFIED | expected-compatible (miles's base image already targets CUDA 13 / sm_103; nothing here is GPU-generation-specific) but not run on this hardware -- see Hardware Requirements | +| RayCluster manifest as shipped (head co-located on the GPU pool) | Verified | the shipped `kubernetes/raycluster.yaml` schedules the head onto the GPU pool (`num-gpus 0`, with a `nvidia.com/gpu` toleration) alongside the worker: dense Qwen3-4B GRPO SUCCEEDED with `raw_reward` 0.53 and `repetition_frac` 0.0. The head must be on a CUDA-capable node -- miles control actors import Megatron even at `num-gpus 0`, so a CPU-only head node fails with `libcuda.so.1: cannot open shared object file`. See [docs/VERIFICATION_LOG.md](./docs/VERIFICATION_LOG.md) | +| Disaggregated reward service (`remote_rm` on a CPU pool) | UNVERIFIED | `reward_service/`, `kubernetes/reward-service.yaml` are present and mirror the sibling slime test case, but were not deployed/exercised on miles | +| Checkpoint save-back / long-run checkpointing | Known Issue | `save_model()` fails with a pickle-truncation error in Megatron's distributed checkpoint save; see [Known Issues](#known-issues) | + +## Prerequisites + +This test case does not create cluster infrastructure; it deploys onto a cluster that already +provides the pieces below. The manifests reference them by label/name, so if any is missing the +failure is a scheduling or mount error, not an obvious message. Confirm each before starting. + +1. An Amazon SageMaker HyperPod cluster with EKS orchestration and GPU instance groups + (p5en.48xlarge/H200 -- verified -- or p6-b300.48xlarge/B300 -- expected-compatible, + see Hardware Requirements -- with EFA). +2. **Node placement labels.** `kubernetes/raycluster.yaml` schedules the GPU workers on + `${GPU_NODE_LABEL_KEY}: ${GPU_NODE_ROLE}` and the Ray head on + `${CPU_NODE_LABEL_KEY}: ${CPU_NODE_ROLE}`. There is no universal `node-role` label; both the + key and the value are cluster-specific. You can either point the four env vars at a label your + nodes already carry -- e.g. `GPU_NODE_LABEL_KEY=node.kubernetes.io/instance-type`, + `GPU_NODE_ROLE=p5en.48xlarge`, or on SageMaker HyperPod the + `sagemaker.amazonaws.com/instance-group-name` of your GPU group -- or add your own label: + `kubectl label node node-role=gpu` (then `GPU_NODE_LABEL_KEY=node-role`, + `GPU_NODE_ROLE=gpu`). If you have no dedicated CPU pool, set the CPU_ vars equal to the GPU + ones (the head runs `num-gpus 0`). A wrong key OR value leaves pods `Pending` with + `FailedScheduling`, not an obvious error. +3. **Cluster add-ons that advertise the scheduled resources**, all of which the manifests + request and none of which this test case installs: + - the NVIDIA device plugin (or GPU Operator, or a GPU AMI that bundles it) advertising + `nvidia.com/gpu`; + - the AWS EFA Kubernetes device plugin (`aws-efa-k8s-device-plugin`) advertising + `vpc.amazonaws.com/efa` -- read the allocatable count off a node for `EFA_PER_NODE` + (`kubectl get node -o jsonpath='{.status.allocatable.vpc\.amazonaws\.com/efa}'`); + - the FSx for Lustre CSI driver, bound to the `fsx-claim` PVC below. +4. `kubectl` and `helm` configured to access the cluster. +5. The KubeRay operator installed (see step 0 below). +6. The Ray head co-locates on the GPU pool (it runs `num-gpus 0` and uses no GPU, only CPU/disk). + It must be on a CUDA-capable node: miles control actors import Megatron/`transformer_engine` + even at `num-gpus 0`, which loads `libcuda.so.1` at import, so a CPU-only head node fails. + Ensure the GPU node's root volume has ample ephemeral-storage (>=150 GiB) for the ~18 GB image. +7. An FSx for Lustre `PersistentVolumeClaim` named `fsx-claim`, RWX, mounted at `/fsx`. +8. **EFA security group (multi-node / `COLOCATE=false` only).** The EFA node security group must + allow all traffic to itself on BOTH ingress AND egress (self-referencing). EFA's OS-bypass + SRD traffic is not ordinary IP, so a CIDR-only egress rule does not authorize it and NCCL + over EFA fails with `Unreachable remote` / `Unexpected number of remote rails`. See + [docs/EFA_2NODE.md](./docs/EFA_2NODE.md). +9. Container registry access (e.g. Amazon ECR) for building/pushing images. +10. A Hugging Face account and access token for model downloads. +11. A Kubernetes Secret named `hf-token` in the target namespace. `kubernetes/raycluster.yaml` + mounts it without `optional: true`, so the Ray pods will not start without it; if your model + is public, create it with an empty value (step 1 has the command). + +### Cluster assumptions and portability + +`kubernetes/raycluster.yaml` was validated on an EKS cluster meeting the Prerequisites above +(a `p5en.48xlarge` / H200 pool with the EFA and FSx add-ons) and carries a few values sized for +that node; on a different cluster, adjust these before deploying or the failure is a silent +`Pending`/`ImagePullBackOff`/OOM rather than a message: + +- **EFA is requested unconditionally, including the single-node colocated run.** Without the EFA + device plugin (or on an instance type without EFA), the worker is `Unschedulable`. For a + single-node run on a non-EFA cluster, delete the two `vpc.amazonaws.com/efa` lines from the + worker `resources`; multi-node runs require EFA (verified) and the self-referencing SG above. +- **Worker resources are sized for `p5en.48xlarge` (2 TiB RAM):** cpu 90/96, memory 1800/1900 Gi, + and a 256 Gi memory-backed `/dev/shm` (which counts against the pod memory limit). On a smaller + GPU node, lower these below the node's allocatable or the worker will not schedule. +- **Every node pulls the ~18 GB image**, not just the head: give GPU node root volumes >=100 GiB + free as well (the head needs >=150 GiB, prerequisite 6). +- **Registry auth is assumed to come from the node IAM role (ECR).** For a private or + cross-account registry, add `imagePullSecrets` to the pod specs. +- **`fsx-claim` must be `ReadWriteMany` and exist in the target `${NAMESPACE}`.** A PVC is + namespaced, so one created in `default` is invisible to a run in another namespace; an + `ReadWriteOnce` (e.g. EBS) claim fails to mount across a multi-node run. +- **The disaggregated reward-service overlay** (`kubernetes/reward-service.yaml`, UNVERIFIED) + selects nodes by the HyperPod label `sagemaker.amazonaws.com/instance-group-name`, which a + non-HyperPod cluster does not have; edit its `nodeSelector` for your CPU pool if you deploy it. +- **One worker pod consumes a whole 8-GPU node** (`nvidia.com/gpu: 8`, `num-gpus: '8'`). On + nodes with a different GPU count, change these together with `ACTOR_GPUS_PER_NODE` and the + actor/rollout GPU split; the manifest is not a fractional-GPU RayCluster. +- **If you point `CPU_NODE_ROLE` at the GPU pool** (no dedicated CPU pool), the head also needs + a toleration for that pool's taint (commonly `nvidia.com/gpu:NoSchedule`), or it stays Pending + even though the label matches. Add it to the head pod spec. +- **`/fsx` must already hold the model and data before you launch.** This test case does not + download or convert during training: `MODEL_LOCAL` (HF checkpoint), `MODEL_DIST` (Megatron + `torch_dist`, from step 4), `PROMPT_DATA`, and `EVAL_DATA` must all exist. Quick Start steps 3 + and 4 create them. + +Preflight (run after `source env_vars`, before deploying) -- turns a silent `Pending`/mount +failure into an early, named error: + +```bash +: "${NAMESPACE:?}" "${FULL_IMAGE:?}" "${FSX_CLAIM:?}" "${GPU_NODE_ROLE:?}" "${CPU_NODE_ROLE:?}" "${EFA_PER_NODE:?}" "${WORKER_REPLICAS:?}" +kubectl get nodes -l "${GPU_NODE_LABEL_KEY}=${GPU_NODE_ROLE}" -o name # GPU pool exists? +kubectl get pvc "${FSX_CLAIM}" -n "${NAMESPACE}" # PVC Bound, RWX, this ns? +kubectl get secret hf-token -n "${NAMESPACE}" # hf-token present? +kubectl get crd rayclusters.ray.io # KubeRay installed? +# and on /fsx (from a pod that mounts it): +# ls -ld "$MODEL_LOCAL" "$MODEL_DIST"; test -s "$PROMPT_DATA"; test -s "$EVAL_DATA" +``` + +## Quick Start + +### 0. Install the KubeRay Operator (one-time per cluster) + +The Ray cluster is managed by the KubeRay operator. If it is not already present +(`kubectl get crd rayclusters.ray.io`), install it with Helm: + +```bash +helm repo add kuberay https://ray-project.github.io/kuberay-helm/ +helm repo update +helm install kuberay-operator kuberay/kuberay-operator \ + --version 1.4.2 --namespace kuberay-operator --create-namespace +``` + +### 1. Configure Environment Variables + +```bash +cd 3.test_cases/pytorch/miles +cp env_vars.colocated.example env_vars +# Edit env_vars with your cluster-specific values +source env_vars +``` + +Key variables (see `env_vars.colocated.example` for the full annotated file): + +```bash +# AWS / ECR (region and account are auto-derived from your AWS credentials) +export AWS_REGION="${AWS_REGION:-$(aws configure get region)}" +export AWS_ACCOUNT_ID="$(aws sts get-caller-identity --query Account --output text)" +export REGISTRY="${AWS_ACCOUNT_ID}.dkr.ecr.${AWS_REGION}.amazonaws.com/" +export IMAGE="miles-hyperpod" +export TAG="miles-cuda13-efa-0.1" # pinned image tag (avoid :latest) +export FULL_IMAGE="${REGISTRY}${IMAGE}:${TAG}" + +# Model (Qwen3-4B dense, bf16, colocated -- the verified default) +export MODEL_NAME="Qwen/Qwen3-4B" +export MODEL_LOCAL="/fsx/models/Qwen3-4B" # SGLang rollout init +export MODEL_DIST="/fsx/models/Qwen3-4B_torch_dist" # Megatron training (torch_dist) +export COLOCATE="true" +export ACTOR_NUM_NODES=1 +export ACTOR_GPUS_PER_NODE=8 +export ROLLOUT_NUM_GPUS=8 + +# Cluster +export NAMESPACE="default" +export FSX_CLAIM="fsx-claim" +export GPU_NODE_ROLE="gpu-p5en" # node-role label on your GPU pool +export EFA_PER_NODE=15 # the node's FULL allocatable EFA count -- see env_vars +``` + +`ACTOR_NUM_NODES` drives the RayCluster's worker replica count as well as the actor layout, so +one variable sets both and the manifest cannot disagree with the recipe. Raise it to 2 for the +2-node (16 GPU) configurations. + +The pods read `HF_TOKEN` from a Kubernetes Secret (not from `env_vars` / the Ray runtime-env), +so it never lands in the Ray dashboard's job-submission metadata. Create it once per +namespace. This Secret is **required**, not optional: `kubernetes/raycluster.yaml` references +it without `optional: true`, so if it is absent the Ray pods never start and report +`CreateContainerConfigError` rather than anything about HuggingFace. If the model you point at +is public and you have no token, create the Secret with an empty value -- the pods only need +the key to exist: + +```bash +kubectl create secret generic hf-token \ + --from-literal=HF_TOKEN=hf_xxx \ + -n "${NAMESPACE}" + +# Public model and no token? The pods only need the key to exist: +# kubectl create secret generic hf-token --from-literal=HF_TOKEN= -n "${NAMESPACE}" +``` + +`env_vars.colocated.example` also ships a commented-out **ALTERNATE** block for +Qwen3-30B-A3B MoE, colocated on 2 nodes (16 GPU) -- uncomment it (and the matching recipe in +step 7) to run the MoE configuration instead of the 4B dense default. It trains cleanly as +shipped (`rollout/raw_reward` 0.578, `rollout/repetition_frac` 0.0), running the rollout MoE +pure expert-parallel (`moe_tp=1`, `EP_SIZE = ROLLOUT_GPUS_PER_ENGINE`). Geometries that run it +tensor-parallel and expert-parallel at once (`moe_tp>1` and `moe_ep>1`) hit a FlashInfer +allreduce-fusion bug in this build; the recipe disables that fusion for them so they train too +(see [Known Issues](#known-issues) item 2). + +### 2. Build and Push the Container Image + +The image takes `radixark/miles:` (which already bundles miles, SGLang, +Megatron-LM, and the matched PyTorch 2.11 / CUDA 13.0.1 stack) as its base and adds **only** +the AWS EFA networking layer (`miles.Dockerfile`) -- GDRCopy, the EFA installer, and NCCL/EFA +runtime defaults. Pin the base by digest (`MILES_BASE_DIGEST` build arg, with `MILES_BASE_TAG` kept alongside for readability) and pin the resulting image +tag to dated/versioned values; never use `:latest`. + +```bash +# Authenticate to ECR +aws ecr get-login-password --region ${AWS_REGION} | \ + docker login --username AWS --password-stdin ${REGISTRY} + +# Create repository (first time only) +aws ecr create-repository --repository-name ${IMAGE} --region ${AWS_REGION} || true + +# Build image (build context is this test-case directory) +docker build -t ${FULL_IMAGE} -f miles.Dockerfile . +docker push ${FULL_IMAGE} +``` + +If the cluster has no node with local Docker access, `kubernetes/buildkit-job.yaml` builds +and pushes the image **in-cluster** with a rootless buildkit Job instead (CPU/IO-only, no GPU +required) -- edit its `nodeSelector` to a node with >=120 GiB free ephemeral-storage, then +create the ConfigMap and Secret it mounts and apply it through `envsubst` (like +`raycluster.yaml` above) so `${AWS_ACCOUNT_ID}`/`${AWS_REGION}` resolve: + +```bash +kubectl create configmap miles-build-context \ + --from-file=Dockerfile=miles.Dockerfile -n "${NAMESPACE}" +kubectl create secret docker-registry ecr-miles-push \ + --docker-server="${REGISTRY}" --docker-username=AWS \ + --docker-password="$(aws ecr get-login-password --region "${AWS_REGION}")" \ + -n "${NAMESPACE}" +envsubst < kubernetes/buildkit-job.yaml | kubectl apply -f - +kubectl -n "${NAMESPACE}" logs -f job/miles-efa-build +``` + +### 3. Download and Prepare the Model + +```bash +# Create a data-prep pod +envsubst < kubernetes/data-prep-pod.yaml | kubectl apply -f - +kubectl exec -it data-prep -- bash + +# Inside the pod: +pip install huggingface_hub +# HF_TOKEN comes from the hf-token Secret. If you created that Secret with an empty +# value because the model is public, skip this login. +huggingface-cli login --token "${HF_TOKEN}" + +# Download model +huggingface-cli download Qwen/Qwen3-4B --local-dir /fsx/models/Qwen3-4B + +# Download training + evaluation datasets +huggingface-cli download --repo-type dataset zhuzilin/dapo-math-17k \ + --local-dir /fsx/data/dapo-math-17k +huggingface-cli download --repo-type dataset zhuzilin/aime-2024 \ + --local-dir /fsx/data/aime-2024 +``` + +### 4. Convert Model Weights to Megatron Format + +miles's Megatron training backend requires weights in `torch_dist` format. Use the bundled +conversion helper, which sources the model script (`scripts/models/.sh`) from the +miles install at `/root/miles`: + +Run this **inside a pod built from the miles image**, not on your workstation: the script reads +the model scripts from `/root/miles`, writes to `/fsx`, and uses `torchrun` when `--num-gpus` is +passed. A **GPU worker** of the RayCluster is such a pod, so deploy step 5 first and come back +here, or run the same command as a one-off Job on the GPU node pool. It must not be the head: +the head runs with `num-gpus 0` and `NVIDIA_VISIBLE_DEVICES=void`, where importing CUDA fails +(see `docs/PORT_NOTES.md`). + +```bash +# Must be a GPU worker, not the head -- see above. +W=$(kubectl get pod -n "${NAMESPACE}" -l ray.io/node-type=worker -o name | head -1) +W=${W#pod/} +kubectl cp scripts/convert_checkpoint.sh "${NAMESPACE}/${W}:/tmp/convert_checkpoint.sh" +kubectl exec -n "${NAMESPACE}" "${W}" -- \ + bash /tmp/convert_checkpoint.sh hf2megatron \ + --model-script qwen3-4B.sh \ + --hf-path /fsx/models/Qwen3-4B \ + --save-path /fsx/models/Qwen3-4B_torch_dist +``` + +For larger models, pass `--num-gpus 8` to parallelize the conversion with `torchrun`: + +```bash +bash scripts/convert_checkpoint.sh hf2megatron \ + --model-script qwen3-30B-A3B.sh \ + --hf-path /fsx/models/Qwen3-30B-A3B \ + --save-path /fsx/models/Qwen3-30B-A3B_torch_dist \ + --num-gpus 8 +``` + +### 5. Deploy the Ray Cluster + +```bash +# Substitute environment variables into the manifest +source env_vars +envsubst < kubernetes/raycluster.yaml | kubectl apply -f - + +# Watch pods come up (1 head + workers) +kubectl get pods -w -l ray.io/is-ray-node=yes + +# Port-forward the Ray dashboard (the recipes submit to 127.0.0.1:8265) +kubectl port-forward -n "${NAMESPACE}" svc/miles-ray-head-svc 8265:8265 & +``` + +The shipped manifest schedules the Ray head onto the GPU pool (`${CPU_NODE_ROLE}` defaults to +`${GPU_NODE_ROLE}`, with a `nvidia.com/gpu` toleration; the head runs `num-gpus 0`) and GPU +workers on `node-role: ${GPU_NODE_ROLE}`, each worker declaring a `gpu_node` custom Ray resource +(see [miles-specific requirements](#miles-specific-requirements-found-on-real-hardware) for +why). The head must be on a CUDA-capable node -- a CPU-only head fails with `libcuda.so.1` +because miles control actors import Megatron even at `num-gpus 0`. Ensure the GPU node's root +volume has headroom for the ~18 GB image (150 GiB or more recommended). + +### 6. Configure the Reward + +This sample scores rollouts in one of two ways -- pick one in `env_vars`: + +- **Built-in rule-based (default, verified):** `RM_TYPE="deepscaler"` (also `dapo`, `math`, + `f1`, `gpqa`). Runs in-process on the rollout actors; no extra setup. This is the reward + path used for every hardware-verified run in this test case. +- **Remote reward service on a CPU pool (`RM_TYPE="remote_rm"` + `RM_URL`):** offloads scoring + to a separate CPU instance group via `reward_service/` and + `kubernetes/reward-service.yaml`, mirroring the sibling slime test case. This path is + **UNVERIFIED on miles** -- the manifests are present but have not been deployed/exercised + here, and miles's `remote_rm` client lacks slime's retry-with-backoff behavior. See + `env_vars.disaggregated.example` for the overlay. + +No file copy is needed for the default path -- `deepscaler` is built into miles. + +### 7. Launch GRPO Training + +The recipes end in `ray job submit --address http://127.0.0.1:8265`, so run them from a machine +that has the Ray CLI whose version matches the cluster (`pip install "ray=="`) +with the dashboard port-forwarded (step 5). If you cannot install a matching Ray CLI locally -- +e.g. no wheel exists for your local Python version -- or want to skip the port-forward, use the +`./run-on-cluster.sh` helper instead, which ships the recipe into the head pod and runs it there +(Ray is already present, `127.0.0.1:8265` is the head's own dashboard, and the version always +matches). `./run-on-cluster.sh --dry-run` prints what it will do; `./run-on-cluster.sh --recipe +run_grpo_qwen3_30b_a3b.sh` runs the MoE recipe. It only launches the recipe -- deploy the +RayCluster (step 5) first. + +```bash +# Qwen3-4B, colocated (verified 1-node and 2-node paths): +bash recipe/run_grpo_qwen3_4b.sh + +# Qwen3-30B-A3B MoE, colocated on 2 nodes / 16 GPU (trains cleanly: reward 0.578, +# repetition 0.0 -- see Known Issues item 2 for the one rollout-geometry constraint): +# uncomment the ALTERNATE block in env_vars, then launch. The shipped block runs the rollout +# MoE pure expert-parallel (moe_tp=1); for moe_tp>1 & moe_ep>1 the recipe disables the FlashInfer fusion. +bash recipe/run_grpo_qwen3_30b_a3b.sh +``` + +Monitor training: + +```bash +# Ray dashboard (after port-forward) +open http://localhost:8265 + +# Follow Ray job logs, using the submission id the recipe printed when it submitted +# ("raysubmit_..."). `ray job list` renders a table for humans; parse it at your own risk. +ray job logs --address http://localhost:8265 --follow + +# Monitor GPU utilization. This has to be a WORKER: the head runs with num-gpus 0 and +# NVIDIA_VISIBLE_DEVICES=void, so nvidia-smi there reports no devices. +W=$(kubectl get pod -n "${NAMESPACE}" -l ray.io/node-type=worker -o name | head -1) +kubectl exec -n "${NAMESPACE}" "${W#pod/}" -- nvidia-smi +``` + +Both recipes build the `train.py` argv as a bash array and submit it through +`recipe/launcher/grpo_launch.sh`, which sources the model script (`scripts/models/.sh`) +and expands `MODEL_ARGS` in the same shell -- avoiding a shell-escaping trap where an outer +`ray job submit -- bash -c "..."` string would expand the array before it is defined. The job +is submitted with `--entrypoint-resources '{"gpu_node": 0.001}'` so the Ray driver lands on a +GPU worker rather than the (non-GPU) head; see +[miles-specific requirements](#miles-specific-requirements-found-on-real-hardware) for why. + +### 8. Convert Checkpoints Back to HuggingFace Format + +**Read this before running the training step, not after.** With the shipped values there is +nothing here to convert. `SAVE_INTERVAL=1000` sits deliberately above `NUM_ROLLOUT=100` so +that no save is ever triggered, because `save_model()` currently fails inside Megatron's +distributed-checkpoint save (Known Issues item 1) and a triggered save would end the run. The +consequence is that a run can finish `SUCCEEDED` after hours of training and leave **no +checkpoint at all** -- the trained weights are gone, and this step is where you would find +that out. The `iter_0060/` path below is therefore an illustration of the command's shape, +not a directory the default configuration produces. + +This runs on the same GPU-worker pod as step 4, for the same reason: the script and `/fsx` +only exist there, not on your workstation. + +Check first, and expect it to be empty on a default run: + +```bash +W=$(kubectl get pod -n "${NAMESPACE}" -l ray.io/node-type=worker -o name | head -1) +W=${W#pod/} +kubectl exec -n "${NAMESPACE}" "${W}" -- ls /fsx/runs/qwen3-4b/ckpt/qwen3-4b-grpo/ +``` + +```bash +kubectl cp scripts/convert_checkpoint.sh "${NAMESPACE}/${W}:/tmp/convert_checkpoint.sh" +kubectl exec -n "${NAMESPACE}" "${W}" -- \ + bash /tmp/convert_checkpoint.sh megatron2hf \ + --input-dir /fsx/runs/qwen3-4b/ckpt/qwen3-4b-grpo/iter_0060/ \ + --output-dir /fsx/models/Qwen3-4B-GRPO-step60 \ + --origin-hf-dir /fsx/models/Qwen3-4B +``` + +Lowering `SAVE_INTERVAL` to produce a checkpoint is what hits the save bug, so treat this +sample as a measurement and instrumentation harness rather than a route to trained weights +until that issue is resolved -- see [Known Issues](#known-issues). + +## miles-specific requirements (found on real hardware) + +Three issues surface on the miles base image (`nvidia/cuda`) that do **not** occur on +SLIME's NGC base. All three are fixed in this test case's `miles.Dockerfile` and manifests; +they are called out here because carrying an EFA layer or launch pattern from a different +base image verbatim will silently reintroduce them. + +1. **CUDA compat shadowing kills `torch.cuda` (Error 803).** The miles base ships an older + CUDA forward-compat `libcuda` than the actual node driver. CUDA forward-compat requires + `compat >= host driver`; when the older bundled compat library is preferred via + `LD_LIBRARY_PATH`, `torch.cuda.is_available()` dies with `Error 803: unsupported display + driver / cuda driver combination`, and the SGLang engine fails at `get_device()`. + `miles.Dockerfile` removes `/usr/local/cuda*/compat` outright and lets the host driver + resolve instead. (SLIME's NGC base does not hit this because NGC's entrypoint only enables + compat when `compat >= host driver`.) + +2. **`libcuda.so.1` not found inside the SGLang subprocess.** With compat removed, `torch` + still resolves `libcuda` via `ld.so.cache`, but the SGLang server subprocess (and + Triton/cuda-python loaders, which scan `LD_LIBRARY_PATH` directly rather than the cache) + fail with `ImportError: libcuda.so.1: cannot open shared object file`. Fix: append the + driver-injection directories (`/usr/lib64`, `/usr/lib/x86_64-linux-gnu`) to + `LD_LIBRARY_PATH` and register them with `ldconfig` -- appended, not prepended, so they are + only consulted for libraries nothing else resolves. + +3. **The Ray job driver dies on the head (no libcuda).** miles's training actor imports + `mooncake` (P2P weight transfer, libcuda-dependent) at module load. The Ray job driver + runs wherever the job is submitted to land; on a non-GPU head pod, the import dies with + the same `libcuda.so.1` error. Fix: the GPU worker group declares a `gpu_node` custom Ray + resource (`rayStartParams.resources: '{"gpu_node": 1}'`), and the recipe submits with + `ray job submit --entrypoint-resources '{"gpu_node": 0.001}'`, landing the driver on a GPU + worker without consuming a full GPU slot from the colocated placement group. + +Full evidence and root-cause detail: [docs/PORT_NOTES.md](./docs/PORT_NOTES.md). + +## Known Issues + +1. **`save_model()` fails with `_pickle.UnpicklingError: pickle data was truncated`.** This + occurs in Megatron's distributed checkpoint save (`gather_object`) after a training step + completes; the GRPO loop itself is unaffected, but any run that must checkpoint -- + including the megatron2hf conversion in step 8 above -- is currently blocked. Workaround + for short smoke runs: set `SAVE_INTERVAL` beyond the total step count so no save is + triggered. + +2. **The 30B MoE rollout degenerates only when SGLang runs the MoE tensor-parallel and + expert-parallel at the same time (`moe_tp>1` and `moe_ep>1`).** With the shipped rollout + geometry (pure expert-parallel, `moe_tp=1`) the 30B MoE trains cleanly -- `rollout/raw_reward` + 0.578, `rollout/repetition_frac` 0.0 -- so this is a rollout-configuration constraint, not a + model or a general "SGLang expert parallelism" problem. The recipe enforces it (below). + + SGLang derives the rollout MoE geometry as `moe_ep = --sglang-expert-parallel-size` (EP_SIZE) + and `moe_tp = --rollout-num-gpus-per-engine / EP_SIZE`. Serving the converted checkpoint + directly from SGLang -- no miles, no Megatron, no GRPO -- and sweeping the engine's + tensor-parallel size (TP) against EP isolates it. `repetition_frac` over 32 prompts, with + the resulting `moe_tp` = TP/EP annotated: + + | engine TP \\ EP | EP=1 (`moe_ep`=1) | EP=2 | EP=4 | EP=8 | + |---|---|---|---|---| + | TP=1 | 0.000 (`moe_tp`=1) | (cannot start) | - | - | + | TP=4 | 0.000 (`moe_tp`=4) | 0.875 (`moe_tp`=2) | - | - | + | TP=8 | 0.000 (`moe_tp`=8) | 0.594 (`moe_tp`=4) | 0.844 (`moe_tp`=2) | ~0.0 (`moe_tp`=1) | + + Read by `moe_tp`, the pattern is exact: every clean cell has `moe_tp=1` (the whole EP=8 + column, pure expert-parallel) or `moe_ep=1` (the whole EP=1 column, pure tensor-parallel); + every degenerate cell has both `moe_tp>1` and `moe_ep>1`. The earlier reading -- "EP>1 + degenerates" -- came from a sweep whose EP>1 cells all happened to have TP>EP, i.e. + `moe_tp>1`; the pure expert-parallel case (EP=TP, the EP=8 column) was not in it. A stock + `sglang.Engine` at `tp_size=8` confirms the missing column: `ep_size=8` (`moe_tp=1`) + generates coherent text (4-gram repetition 0.009), while `ep_size=4` and `ep_size=2` + (`moe_tp` 2 and 4) collapse to "7. 7. 7...", "1010...", ",,,,". Ruled out along the way: the + model's own recommended sampling (temperature 0.6 / top_p 0.95 / top_k 20) does not change a + degenerate cell, the `auto` vs `triton` MoE runner backend does not either, and all 18867 + expert weight keys are present in the checkpoint index with no NaN/Inf, so sampling, backend + selection and conversion are not involved. `--check-weight-update-equal` passes + (`weight_version` uniform, `mixed_version_ratio` 0.0), so trainer/rollout weights are not + diverging. + + Root cause (`0.5.16.dev` in the image): the FlashInfer allreduce+RMSNorm fusion. On SM90/SM100 + it is auto-enabled for Qwen3-MoE (`tp_size>1`, no dp-attention, `moe_a2a=none`) without regard + to `moe_ep`/`moe_tp`. With the fusion on, both post-experts all-reduces in + `models/qwen3_moe.py` `forward_normal` are skipped and deferred to the next layer's fused + `layernorm.forward_with_allreduce_fusion`. That fused reduce + (`layernorm.py::_forward_with_allreduce_fusion`, `flashinfer_comm_fusion.py`) selects its group + with `if moe_ep_size>1: use moe_ep_group else: use moe_tp_group` -- assuming the two are mutually + exclusive. When both are `>1` it reduces over the moe-ep group only and never reduces the moe-tp + group, so each rank keeps a partial sum over the intermediate dimension and generation collapses + from layer 0. Pure expert-parallel (`moe_ep=tp`) and pure tensor-parallel (`moe_tp=tp`) are + unaffected because `_MOE_EP`/`_MOE_TP` alias the full TP group, so either branch reduces over all + ranks. Setting `enforce_disable_flashinfer_allreduce_fusion` on the same combined config restores + the correct two-stage reduce and generation is clean (verified). The MoE weight sharding, the + dispatcher's `local_expert_mapping`, and the moe-ep/moe-tp process groups are all correct; the + defect is purely in the fused-reduce group selection. See + [docs/VERIFICATION_LOG.md](./docs/VERIFICATION_LOG.md) "30B MoE root cause". + + The recipe handles it automatically: `run_grpo_qwen3_30b_a3b.sh` computes `moe_tp` from + `ROLLOUT_GPUS_PER_ENGINE / EP_SIZE`, and when both `moe_tp>1` and `moe_ep>1` it adds + `--sglang-enforce-disable-flashinfer-allreduce-fusion` so that geometry trains correctly too; + pure expert-parallel (the shipped default) and pure tensor-parallel leave the fusion on. Losing + the fusion costs some rollout throughput but not correctness. The underlying SGLang bug is a + candidate for an upstream fix (generalise the fused-reduce group to the full TP group when + `moe_tp>1` and `moe_ep>1`). + +## File Structure + +``` +miles/ # 3.test_cases/pytorch/miles +├── README.md # This documentation +├── .gitignore +├── env_vars.colocated.example # Base config: Qwen3-4B colocated (+ 30B MoE ALTERNATE block) +├── env_vars.disaggregated.example # Overlay: reward model on a CPU pool + heavier GRPO +├── run-on-cluster.sh # Optional: run a recipe from the head pod (no local ray/port-forward) +├── miles.Dockerfile # radixark/miles base + EFA layer +├── requirements.txt # Reference Python deps (miles bundles these in the base image) +├── reward_service.Dockerfile # CPU-only image for the remote reward service +├── reward_service/ +│ ├── app.py # FastAPI reward server (reward_model / math_verify) +│ └── requirements.txt # Pinned CPU deps (no CUDA) +├── kubernetes/ +│ ├── buildkit-job.yaml # In-cluster image build -> registry (CPU-only) +│ ├── raycluster.yaml # KubeRay cluster manifest (head co-located on GPU pool, num-gpus 0) +│ ├── reward-service.yaml # CPU reward service Deployment + Service +│ └── data-prep-pod.yaml # Utility pod for data preparation +├── recipe/ +│ ├── run_grpo_qwen3_4b.sh # GRPO submit script (Qwen3-4B, colocated) +│ ├── run_grpo_qwen3_30b_a3b.sh # GRPO submit script (Qwen3-30B-A3B MoE, colocated 16 GPU) +│ └── 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 +└── docs/ + ├── EFA_2NODE.md # 2-node EFA/NCCL verification record + ├── PORT_NOTES.md # miles-specific porting notes and the 3 hardware pitfalls + └── VERIFICATION_LOG.md # runs, job ids, flags, metrics (source of Verification Status) +``` + +## Training Configuration Deep Dive + +### GRPO (Group Relative Policy Optimization) + +GRPO is a critic-free RL algorithm that estimates advantages by comparing rewards within a +group of responses generated for the same prompt, eliminating the need for a separate value +model. + +| Parameter | Description | Recipe Setting (4B) | +|-----------|-------------|----------------------| +| `--advantage-estimator grpo` | Use GRPO advantage estimation | GRPO | +| `--rollout-batch-size` | Prompts per rollout | 16 | +| `--n-samples-per-prompt` | Responses generated per prompt | 8 | +| `--global-batch-size` | Samples per optimizer step | 128 | +| `--num-steps-per-rollout` | Optimizer steps per rollout cycle | 1 | +| `--eps-clip` / `--eps-clip-high` | PPO-style clipping bounds | 0.2 / 0.28 | +| `--kl-loss-coef` | KL penalty coefficient | 0.0 | +| `--entropy-coef` | Entropy bonus coefficient | 0.0 | + +The constraint `rollout_batch_size * n_samples_per_prompt == global_batch_size * +num_steps_per_rollout` (16 * 8 = 128 * 1) must always hold; the recipes validate every +required variable is set before submitting. + +### Parallelism Strategy + +**Qwen3-4B (colocated, verified on 1 and 2 nodes):** + +``` +Tensor Parallel (TP) = 1 +Pipeline Parallel (PP) = 1 +Context Parallel (CP) = 1 +Expert Parallel (EP) = 1 +Colocated: rollout-num-gpus == actor GPU count (8 on 1 node, 16 on 2 nodes) +``` + +**Qwen3-30B-A3B MoE (colocated on 2 nodes / 16 GPU, trains cleanly -- reward 0.578, +repetition 0.0; see [Known Issues](#known-issues) item 2 for the rollout-geometry constraint):** + +``` +Tensor Parallel (TP) = 2 +Pipeline Parallel (PP) = 1 +Context Parallel (CP) = 1 +Expert Parallel (EP) = 2 +Expert Tensor Parallel = 1 +--use-distributed-optimizer # required: shards the 30B optimizer state across all + # 16 GPU so static memory fits H200 (141 GB); confining + # the actor to 8 GPU (one node) OOMs +--colocate # rollout time-shares the same 16 GPU +--rollout-num-gpus-per-engine 2 +--sglang-moe-runner-backend triton # MoE online weight update requires the triton + # runner; flashinfer is incompatible with the + # in-place weight update on SGLang 0.5.12+ +--sglang-expert-parallel-size 2 +``` + +### Dynamic Batching + +```bash +--use-dynamic-batch-size +--max-tokens-per-gpu 8192 # Per-GPU token budget per micro-batch +``` + +Variable-length responses (math solutions can range from under 100 to 8,000+ tokens) make +dynamic batching materially more efficient than a fixed micro-batch size. + +## Reward Function + +1. **Built-in rule-based rewards (default, verified).** miles ships several rule-based + reward types -- `deepscaler`, `dapo`, `math`, `f1`, `gpqa` -- selected via `--rm-type`. The + math types extract the `\boxed{...}` answer and grade it against the label using + `math_verify` (LaTeX/sympy equivalence). The bundled DAPO-math recipe defaults to + `RM_TYPE="deepscaler"` and needs no extra setup. + +2. **Remote reward service on a CPU pool (UNVERIFIED on miles).** Set `RM_TYPE="remote_rm"` + and `RM_URL` to offload scoring to a separate CPU instance group -- useful for a heavier + reward such as a reward model, code execution, or RAG lookups. The service + (`reward_service/app.py`) exposes a `reward_model` backend (HuggingFace sequence + classifier on CPU) and a `math_verify` backend behind `POST /score`, plus `GET /health`. + This path mirrors the sibling slime test case's manifests but has not been deployed or + exercised on miles; miles's `remote_rm` client also lacks slime's retry-with-backoff, so + treat it as a starting point rather than a validated path. + +## Software Versions + +| Component | Version | +|-----------|---------| +| miles | `radixark/miles` (fork point `fcce96ca0`, 2025-10-05) | +| Base image | `radixark/miles:dev-202607310056` @ `sha256:ca0bb593dd6f4011b444f64d478b72c213e4c70421f4d7f94e593a709562429e` | +| SGLang | 0.5.16.dev | +| Megatron-LM | radixark fork (miles-compatible) | +| Ray | 2.55.1 | +| CUDA | 13.0.1 | +| PyTorch | 2.11 | +| EFA installer | 1.48.0 | +| GDRCopy | v2.5.2 | + +Pin the image tag to a dated/versioned value (e.g. `miles-cuda13-efa-0.1`); never use +`:latest`, per the [awsome-distributed-ai CONTRIBUTING guidelines](https://github.com/awslabs/awsome-distributed-ai/blob/main/CONTRIBUTING.md). + +## Troubleshooting + +**Pod stuck in `Pending` state** +```bash +kubectl describe pod +# Check for resource constraints -- GPU/EFA/memory/ephemeral-storage requests may +# exceed node capacity. The Ray head in particular needs enough ephemeral-storage +# to pull the ~18 GB image; a too-small root volume causes an Evict mid-pull rather +# than a clean Pending. +``` + +**Ray workers fail to connect to head node** +```bash +kubectl get svc miles-ray-head-svc +kubectl exec -- nslookup miles-ray-head-svc +# Ensure RAY_memory_monitor_refresh_ms=0 is set (prevents OOM kills during init) +``` + +**NCCL/EFA initialization errors** +```bash +kubectl exec -- fi_info -p efa +kubectl exec -- env | grep NCCL +# Ensure FI_PROVIDER=efa and FI_EFA_USE_DEVICE_RDMA=1 are set +# If NCCL selects efa but every transfer times out (Error 15, "Unreachable remote"), +# check that the EFA security group has self-referencing all-traffic on BOTH +# ingress AND egress -- EFA's OS-bypass SRD traffic is not ordinary IP traffic, so +# a CIDR-only egress rule does not authorize it. See docs/EFA_2NODE.md. +``` + +**`ImportError: libcuda.so.1: cannot open shared object file`** +- Confirm the Ray job driver landed on a GPU worker (`--entrypoint-resources + '{"gpu_node": 0.001}'`) and not the head. +- Confirm `LD_LIBRARY_PATH` includes the driver-injection directories + (`/usr/lib64`, `/usr/lib/x86_64-linux-gnu`) appended at the end. +- See [miles-specific requirements](#miles-specific-requirements-found-on-real-hardware). + +**`torch.cuda.is_available()` is `False` / `Error 803`** +- The image's CUDA forward-compat library is older than the node driver. Confirm + `/usr/local/cuda*/compat` has been removed (as `miles.Dockerfile` does) rather than + merely dropped from `LD_LIBRARY_PATH`. + +**SGLang fails to start (CUDA OOM) in colocated mode** +- In colocated mode, SGLang launches after Megatron occupies GPU memory. Reduce + `--sglang-mem-fraction-static` (0.8 for 4B; 0.75 was the value that completed for the 30B + MoE colocated run -- the original 0.85 hardcode left too little room). + +**MoE online weight update fails on SGLang 0.5.12+** +- Ensure `--sglang-moe-runner-backend triton` and `--sglang-expert-parallel-size` are both + set explicitly; the default flashinfer MoE runner is incompatible with the in-place + weight update used by GRPO training. + +**Weight conversion fails** +- Ensure `PYTHONPATH` includes the Megatron-LM directory (`/root/Megatron-LM`). +- Verify model config parameters match (`--rotary-base`, `--vocab-size`, etc.). +- For MoE models, ensure `--expert-model-parallel-size` is set correctly. + +**Checkpoint save fails with `_pickle.UnpicklingError: pickle data was truncated`** +- Known upstream issue in Megatron's distributed checkpoint save path; see + [Known Issues](#known-issues). Raise `SAVE_INTERVAL` beyond the run length to avoid + triggering a save on short smoke runs. + +## References + +- [miles GitHub Repository](https://github.com/radixark/miles) +- [SLIME GitHub Repository](https://github.com/THUDM/slime) +- [SLIME Blog: An SGLang-Native Post-Training Framework for RL Scaling](https://lmsys.org/blog/2025-07-09-slime/) +- [SGLang Project](https://github.com/sgl-project/sglang) +- [Megatron-LM](https://github.com/NVIDIA/Megatron-LM) +- [GRPO Paper (DeepSeek-R1)](https://arxiv.org/abs/2402.03300) +- [Amazon SageMaker HyperPod Documentation](https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-hyperpod.html) +- [Sibling test case: 3.test_cases/pytorch/slime/](../slime/) +- [awsome-distributed-ai](https://github.com/awslabs/awsome-distributed-ai) +- [KubeRay Documentation](https://docs.ray.io/en/latest/cluster/kubernetes/index.html) + +## Security + +See [CONTRIBUTING](https://github.com/awslabs/awsome-distributed-ai/blob/main/CONTRIBUTING.md) for more information. + +## License + +This sample code is made available under the MIT-0 license. See the LICENSE file. diff --git a/3.test_cases/pytorch/miles/docs/EFA_2NODE.md b/3.test_cases/pytorch/miles/docs/EFA_2NODE.md new file mode 100644 index 000000000..aa7e0672e --- /dev/null +++ b/3.test_cases/pytorch/miles/docs/EFA_2NODE.md @@ -0,0 +1,82 @@ +# 2-node EFA verification (p5en.48xlarge x2) + +Hardware verification that multi-node NCCL runs over EFA on this cluster, using two +Capacity Block p5en.48xlarge nodes (8x H200 each, us-east-2a, same UltraCluster). No +placement group is used or needed: a Capacity Block already colocates its nodes within one +UltraCluster spine, and `PlacementGroupArn` on the reservation is null, so layering a +self-made cluster placement group would only risk "capacity reserved but PG-unsatisfiable". + +## Result: EFA works across 2 nodes + +`torchrun --nnodes 2 --nproc_per_node 8` running a NCCL all_reduce (16 GPUs), with +`NCCL_DEBUG=INFO NCCL_DEBUG_SUBSYS=INIT,NET`: + +- Both nodes: `NET/OFI Selected provider is efa, fabric is efa-direct (found 8 nics)`. +- Inter-node channels: `Channel .. : 3[3] -> 14[6] [send] via NET/Libfabric/3/GDRDMA`, + `GPU Direct RDMA Enabled`. Zero `NET/Socket` lines (no TCP fallback). +- Bandwidth (algbw / busbw, all_reduce, 16 GPUs): + + | message | algbw | busbw | + | --- | --- | --- | + | 64 MB | 101.5 | 190.3 GB/s | + | 256 MB | 100.5 | 188.5 GB/s | + | 1024 MB | 127.2 | 238.4 GB/s | + | 4096 MB | 135.8 | 254.7 GB/s | + | 8192 MB | 137.1 | 257.1 GB/s | + +busbw 190-257 GB/s is far above the ~10 GB/s a TCP path would give, confirming EFA. + +**These numbers were taken with a partial EFA request, which is no longer what the manifest +does -- and that turned out to matter for correctness, not just throughput.** The measurement +above ran with the pods requesting `vpc.amazonaws.com/efa: 8` of the node's allocatable 15, +roughly half the fabric, and an earlier version of this document described requesting them all +as "a throughput improvement, not a correctness issue". That was wrong. + +Claiming fewer cards than the node exposes lets the device plugin choose which ones, and the +two pods can end up with different EFA-device-to-NIC-rail mappings. `aws-ofi-nccl` then aborts +with `NET/OFI Unexpected number of remote rails for dev N. Expected 1 but got 2` followed by +`ncclInternalError`. Crucially it does **not** appear in the all_reduce test above: TP16 +all-reduce across the node boundary passed with 8 cards, while TP8 x PP2 on the same cluster +failed within minutes, because the pipeline stage boundary uses point-to-point send/recv. +MoE expert-parallel all-to-all is P2P-like and is affected the same way. + +So `EFA_PER_NODE` now defaults to the node's full allocatable count (15 on p5en.48xlarge, +32 on p5.48xlarge) and `kubernetes/raycluster.yaml` takes it from there. Read the count off the +node rather than the instance spec: + +```bash +kubectl get node -o jsonpath='{.status.allocatable.vpc\.amazonaws\.com/efa}' +``` + +The throughput point still stands on its own -- the full request is also expected to exceed +~300 GB/s -- but an EFA validation that only exercises all-reduce will report success on a +configuration that breaks later inside training. Include a P2P path. + +## Root cause fixed: EFA security-group needs self-referencing egress + +The first 2-node run failed even though NCCL selected the efa provider and established +GPUDirect RDMA channels: every data transfer timed out with +`NET/OFI ... Error 15 (Unreachable remote (never received a response))`. + +Cause: the EFA security group had self-referencing all-traffic on **ingress** but only +`0.0.0.0/0` on **egress**. EFA's OS-bypass SRD traffic is not IP traffic, so a CIDR egress +rule does not authorize it -- EFA SGs need self-referencing all-traffic on BOTH directions. +Bootstrap (TCP over eth0) and NCCL init succeeded under the asymmetric rule, which is why +the failure only appeared at the first SRD data transfer. + +Fix: add a self-referencing security-group rule allowing all protocols on both ingress and +egress (e.g. `aws_security_group_rule` with `source_security_group_id` == the SG's own ID, +protocol `-1`, for both directions) to the EFA-capable nodes' security group. Single-node EFA +was unaffected because intra-node traffic does not traverse the SG. + +## How to reproduce the check + +1. Two EFA nodes in the same subnet with the EFA SG (self-ref ingress AND egress). +2. sshd not required -- use torchrun with `--master_addr `; NCCL data path is + EFA regardless of the launcher. +3. On each pod: `FI_PROVIDER=efa NCCL_DEBUG=INFO NCCL_DEBUG_SUBSYS=INIT,NET + NCCL_SOCKET_IFNAME=eth0 torchrun --nnodes 2 --node_rank <0|1> --nproc_per_node 8 + --master_addr --master_port 29500 bench.py`. +4. Pass: `Selected provider is efa` on both nodes, inter-node `[send] via NET/Libfabric`, + no `NET/Socket`, busbw > 100 GB/s. Fail (TCP): busbw single-digit GB/s, or Error 15 + (check the SG self-ref egress rule first). diff --git a/3.test_cases/pytorch/miles/docs/PORT_NOTES.md b/3.test_cases/pytorch/miles/docs/PORT_NOTES.md new file mode 100644 index 000000000..d892a985d --- /dev/null +++ b/3.test_cases/pytorch/miles/docs/PORT_NOTES.md @@ -0,0 +1,115 @@ +# Porting notes (slime -> miles) + +Record of porting the sibling [slime test case](../../slime/) to +[miles](https://github.com/radixark/miles), slime's direct fork. miles diverged from slime +at fork point `fcce96ca0` (2025-10-05) and rewrote the train loop sync -> async, but the +`train.py` CLI is compatible, so the same GRPO recipe runs with only path/branding changes +(verified on hardware). + +## Why miles + +- miles targets CUDA 13.0.1 / PyTorch 2.11 / Blackwell (sm_103) as first-class. Where the + slime image reached sm_103 via an NGC base plus hand patches, the miles image already + applies the sm_103 Transformer Engine FA2 whitelist patch. +- miles has features slime lacks (fp8 rollout, fully-async, true-on-policy) -- out of scope + here, but future material. + +## Docker strategy C: miles official image + EFA layer + +miles's dependencies (the sglang-miles fork, radixark/Megatron-LM fork, prebuilt +flash_attn/TE/apex wheels) are built for the PyTorch 2.11 stable + cu130 ABI and do not +load on an NGC (nightly ABI) base. So the image takes `radixark/miles:` as the +base and adds only the AWS EFA stack (`miles.Dockerfile`): + +- remove IB libverbs (let the EFA installer's libfabric win) +- gdrcopy (GPUDirect RDMA) +- EFA installer 1.48.0 (`--skip-kmod`, libfabric + aws-ofi-nccl plugin) +- NCCL/EFA runtime env (`FI_PROVIDER=efa` etc.), `GLOO_SOCKET_IFNAME=eth0` +- `PYTHONPATH=/root/miles:/root/Megatron-LM` (miles does not bake this into the image) + +Pin by digest (`MILES_BASE_DIGEST`), never `:latest` and not by tag alone: radixark +publishes `dev-*` as mutable snapshots and has already deleted a tag this test case once +pinned. The dated tag is kept next to the digest for readability. Per awsome-distributed-ai +CONTRIBUTING. + +## slime -> miles changes (recipe / launcher) + +| Item | slime | miles | +| --- | --- | --- | +| framework install dir | `/opt/slime` | `/root/miles` (editable install) | +| Megatron | `/opt/Megatron-LM` | `/root/Megatron-LM` (radixark fork) | +| recipe runtime-env PYTHONPATH | `/opt/Megatron-LM` | `/root/Megatron-LM:/root/miles` | +| launcher `SLIME_DIR` default | `/opt/slime` | `/root/miles` (variable name kept) | +| model script | `scripts/models/qwen3-4B.sh` | same (path compatible) | + +The `train.py` flags are unchanged, so the GRPO recipe argv is identical to slime's apart +from the paths above. SGLang passthrough flags (`--sglang-*`) work the same way on miles +(ServerArgs auto-expose via parse_known_args). + +## Pitfalls found on real hardware + +All three occur on the miles base (nvidia/cuda) and NOT on slime's NGC base, and are fixed +in `miles.Dockerfile` / the manifests. + +### 1. CUDA compat shadowing -> torch.cuda dies (Error 803) + +Carrying slime.Dockerfile's EFA layer verbatim also carries +`ENV LD_LIBRARY_PATH=...:/usr/local/cuda/compat:...`, which is fatal on the miles base: + +- The miles base bundles a CUDA forward-compat libcuda `580.82.07`. +- The validation node's host driver is `580.159.03` (nvidia-smi shows "CUDA Version: 13.0"). +- CUDA forward-compat requires compat >= host driver. When the older compat (580.82.07) is + preferred via LD_LIBRARY_PATH, torch.cuda dies with `RuntimeError: ... Error 803: system + has unsupported display driver / cuda driver combination`, and the SGLang engine fails at + `get_device()` ("No accelerator available"). +- Verified: dropping compat from LD_LIBRARY_PATH restores `torch.cuda.is_available() == True`. +- Fix: `miles.Dockerfile` deletes `/usr/local/cuda*/compat` outright and uses the host + driver (`/usr/lib64/libcuda.so`), which already supports the image's CUDA 13.0 toolkit. + slime (NGC) could use compat because NGC's entrypoint enables it only when compat >= host. + +### 2. libcuda.so.1 not found in the SGLang subprocess + +With compat gone, torch still works (ld.so.cache resolves libcuda), but the SGLang server +subprocess (and Triton / cuda-python loaders, which scan LD_LIBRARY_PATH directly rather +than the cache) fail with `ImportError: libcuda.so.1: cannot open shared object file`. Fix: +append the driver-injection dirs (`/usr/lib64`, `/usr/lib/x86_64-linux-gnu`) to +LD_LIBRARY_PATH and register them with `ldconfig`. Appending (not prepending) means they +are consulted only for libraries nothing else resolves. + +### 3. Ray job driver dies on the head (no libcuda) + +miles's `MegatronTrainRayActor` imports `mooncake` (P2P weight transfer, libcuda-dependent) +at module load. The Ray job driver runs on the head, which is a non-GPU pod with no libcuda +injected, so the driver dies with the same `ImportError: libcuda.so.1`. Fix: declare a +`gpu_node` custom resource on the worker (`rayStartParams.resources: '{"gpu_node": 1}'`) +and submit with `ray job submit --entrypoint-resources '{"gpu_node": 0.001}'`, which places +the driver on a GPU worker without consuming a GPU logical count (so it does not conflict +with the colocated 8-GPU placement group). Note: `begin_weight_update` / `pull_weights` are +Ray actor methods on miles's rollout engine (`sglang_engine.py`), not HTTP endpoints on an +SGLang fork. + +## GPU-generation portability (H200 / B300) + +Everything in `miles.Dockerfile`, the recipes, and the manifests is GPU-generation-agnostic: +the base image tag (`MILES_BASE_TAG`) already carries both the sm_90 (Hopper/H200) and +sm_103 (Blackwell/B300) builds internally, and no recipe/env file hard-codes a CUDA or SM +version. The one cluster-specific knob is `GPU_NODE_ROLE` in `env_vars` (substituted into +`kubernetes/raycluster.yaml`'s worker `nodeSelector`), which just needs to match whichever +accelerator NodePool -- H200 or B300 -- you point it at. This test case is hardware-verified +only on H200 (p5en.48xlarge); the analysis above is why B300 (p6-b300.48xlarge) is expected +to work unmodified, not a claim that it has been run. + +## Known Issue + +- `save_model()` fails with `_pickle.UnpicklingError: pickle data was truncated` in + Megatron's distributed checkpoint save (`gather_object`), independent of the GRPO loop. + Blocks the HF<->Megatron round-trip and long checkpointing runs. File upstream on miles. + +## Residual patches (of slime's 7, what remains on miles) + +- `--sglang-log-level warning` (lowercase; avoids the uvicorn KeyError; recipe-only). +- GPU-less Ray driver Megatron `validate_args` CUDA probe (may surface for 30B MoE; may + already be fixed in the radixark fork -- UNVERIFIED). + +The numpy<2 pin, torch_memory_saver preload `.so` selection, and manual mbridge pin are all +resolved by the miles official image. diff --git a/3.test_cases/pytorch/miles/docs/VERIFICATION_LOG.md b/3.test_cases/pytorch/miles/docs/VERIFICATION_LOG.md new file mode 100644 index 000000000..089d61ff6 --- /dev/null +++ b/3.test_cases/pytorch/miles/docs/VERIFICATION_LOG.md @@ -0,0 +1,322 @@ +# Verification log + +What was run, on what, and what came back. Every metric below was read out of the trainer's +own TensorBoard event files rather than retyped from a terminal, and every run named here has +an event file behind it. + +This file exists because an earlier version of the README's verification table contained a +row that was wrong, and a log is what would have caught it. See "A correction" at the end. + +## Environment + +| | | +|---|---| +| Cluster | Amazon EKS, KubeRay, 2 nodes | +| Instance | `p5en.48xlarge` (H200 141GB x8 per node) | +| Interconnect | EFA between nodes, GPUDirect RDMA | +| Shared storage | FSx for Lustre, PERSISTENT_2, mounted at `/fsx` | +| Base image | `radixark/miles@sha256:ca0bb593dd6f4011b444f64d478b72c213e4c70421f4d7f94e593a709562429e` (tag `dev-202607310056`, cu13) plus the AWS EFA layer built by `miles.Dockerfile` | +| Model | Qwen3-4B, converted to Megatron format with `scripts/convert_checkpoint.sh hf2megatron` | +| Data | DAPO-Math-17k for training prompts, AIME-2024 for eval, `deepscaler` reward | + +## Run: the shipped 4B recipe, unmodified + +This is the reader's path: the shipped `env_vars.colocated.example` with only the +cluster-specific values filled in (model and data paths, checkpoint and TensorBoard +directories, namespace, FSx claim) and `NUM_ROLLOUT=1` to keep it short. No changes to the +recipe itself. + +``` +ENV_FILE=env_upstream bash recipe/run_grpo_qwen3_4b.sh +``` + +Submitted as Ray job `raysubmit_GdnWJy6KJsJDmsF2`, terminal status `SUCCEEDED`. + +The flags the job actually received, from the job log rather than from the recipe source: + +``` +bash grpo_launch.sh \ + --hf-checkpoint /fsx/models/Qwen3-4B \ + --ref-load /fsx/models/Qwen3-4B_torch_dist \ + --load /fsx/runs/upstream_verify/ckpt/qwen3-4b-grpo/ \ + --save /fsx/runs/upstream_verify/ckpt/qwen3-4b-grpo/ \ + --save-interval 1000 \ + --prompt-data /fsx/data/dapo-math-17k/dapo-math-17k.jsonl \ + --input-key prompt --label-key label ... +``` + +Results, read back from `/fsx/tb/upstream_verify` (78 scalar tags recorded): + +| metric | value | reading | +|---|---|---| +| `rollout/raw_reward` | 0.5156 | the model answers correctly about half the time, which is the expected range for 4B on this data | +| `rollout/repetition_frac` | 0.0 | no degenerate generation | +| `rollout/truncated_ratio` | 0.4844 | about half the responses hit the 8192-token cap, normal for this prompt set | +| `train/grad_norm` | 0.6479 | finite and unremarkable | +| `perf/step_time` | 273.3 s | one full rollout-plus-train cycle | + +Logs kept in full: 2,242 lines of job log and 2,276 lines of submission log. + +## Why reward and repetition are the metrics to check, not the exit code + +`rollout/repetition_frac` and `rollout/raw_reward` are what separate a run that works from a +run that merely finishes. A misconfigured 30B MoE run -- one whose SGLang rollout runs the +MoE tensor-parallel and expert-parallel at once (`moe_tp>1` and `moe_ep>1`) -- exits 0, prints +an unremarkable loss, and produces nothing usable: repetition 0.48 to 0.70, reward pinned at +0.0. If you verify by exit code you will record it as working -- which is exactly what happened +once, before the metric was checked. The same 30B model with the shipped rollout geometry +(`moe_tp=1`, pure expert-parallel) trains cleanly: reward 0.578, repetition 0.0. The metric, +not the exit code, is what tells the two apart -- and what pointed to the rollout geometry as +the cause. See "30B MoE root cause" below. + +For comparison, on the same cluster and recipe: + +| | dense 4B | 30B MoE, `moe_tp=1` (shipped) | 30B MoE, `moe_tp>1` (misconfigured) | +|---|---|---|---| +| `rollout/repetition_frac` | 0.0 | 0.0 | 0.48 to 0.70 | +| `rollout/raw_reward` | 0.516 | 0.578 | 0.0 | +| `rollout/truncated_ratio` | 0.484 | 0.42 | 0.97 to 0.99 | +| exit status | SUCCEEDED | SUCCEEDED | SUCCEEDED | + +## Two-node EFA + +The 2-node rows in the README's table come from earlier runs on this same cluster with the +worker replica count at 2. NCCL all_reduce busbw measured 190 to 257 GB/s with `efa-direct` +and GPUDirect RDMA and no TCP fallback; `docs/EFA_2NODE.md` has the security-group +prerequisite and the failure signature when it is missing. + +`docs/EFA_2NODE.md` is a fabric measurement, though, not a training one, so on its own it does +not support the "3 rollout cycles" part of that row. The GRPO run behind it, Qwen3-4B dense +colocated across 2 nodes (actor 2x8, rollout sharing the same 16 GPU): + +| rollout | `rollout/raw_reward` | `actor_train_tflops` | `perf/step_time` | +|---|---|---|---| +| 0 | 0.477 | 100.4 | 119.2s | +| 1 | 0.523 | 245.4 | 73.1s | +| 2 | 0.492 | 235.7 | 81.6s | + +All three cycles completed, job `SUCCEEDED` in 676s. Weight sync, rollout, reference log-probs +and the Megatron backward all crossed the node boundary over EFA. `ppo_kl` stayed 0.0 at +dropout 0, matching the single-node arm. Reward moving in the 0.48-0.52 band over three steps +is not a training result -- three optimizer steps show nothing about convergence -- it is +evidence that the loop closes end to end on two nodes, which is what that row claims. + +One detail worth repeating from that document, because it cost real time: requesting fewer +EFA devices than the instance exposes lets the device plugin pick different cards on each +node, and NCCL then fails with `NET/OFI Unexpected number of remote rails`. It only shows up +on point-to-point traffic (pipeline or expert parallel), so an all-reduce smoke test passes +and the problem surfaces later. + +## What this log does not cover + +- The disaggregated reward service. Present in the repository, never deployed. +- Any B300 hardware. +- Checkpoint save-back. `save_model()` fails with a pickle-truncation error inside Megatron's + distributed checkpoint save; `SAVE_INTERVAL` ships above the step count so a run does not + trigger it. Known Issues item 1. + +## Ray head must run on a GPU (CUDA-capable) node, not a CPU-only node + +An earlier version of this test case placed the Ray head on a CPU-only node (to keep it off +the expensive GPU pool). That shape is not safe for miles and the shipped manifest no longer +uses it. The reason is in the framework, not the recipe: miles's Ray control actors are +created with `num_gpus=0` (e.g. `create_rollout_manager` in `miles/ray/placement_group.py` +does `RolloutManager.options(num_cpus=1, num_gpus=0)`), but their module import pulls in +Megatron / `transformer_engine`, whose shared library `dlopen`s `libcuda.so.1` at import time. +Ray is then free to place such a zero-GPU actor on any node with a spare CPU -- including a +CPU-only head node -- where the import hard-fails: + +``` +OSError: libcuda.so.1: cannot open shared object file: No such file or directory + ... File ".../transformer_engine/common/__init__.py", ... _load_core_library() + ray.exceptions.ActorDiedError: RolloutManager.__init__() ... (TemporaryActor, ip=) +``` + +So any CPU-only node in the miles Ray cluster is a latent hazard: whether a run survives +depends on whether Ray happens to place the Megatron-importing actor on a GPU worker instead. +The earlier "head-on-CPU SUCCEEDED (reward 0.531)" run was that placement luck, not a +guarantee. The fix is to keep every node in the Ray cluster CUDA-capable: the shipped +`raycluster.yaml` now schedules the head onto the GPU pool (`CPU_NODE_ROLE` defaults to +`GPU_NODE_ROLE`) with a `nvidia.com/gpu` toleration; the head still runs `num-gpus 0` and +consumes no GPU, so on a colocated run it simply co-locates on a worker's GPU node at no extra +cost, and the ~18 GB image is already cached there. This also removes the head from the CPU +Karpenter pool, sidestepping the "underutilized" consolidation churn that pool's 30s policy +caused. The verified metrics for the colocated dense 4B run (reward 0.53, repetition 0.0, +weight_version uniform) are unchanged; only the head's node placement changed. + +Note (upstream): the sibling `slime` test case ships the same head-on-CPU shape and the same +`num_gpus=0` control-actor pattern, so it shares this latent hazard; worth raising upstream. + +## Second review round: defects found by reading, and what was re-run to confirm + +A later review pass went through the reader's path again rather than re-running the same job, +and found six defects that no amount of re-running would have surfaced, because the affected +paths either are not on the default path or fail silently. They are listed here because the +fixes changed shipped files, and a log that only records successes is not much of a log. + +| what was wrong | why it was silent | +|---|---| +| `kubernetes/reward-service.yaml` referenced `${REWARD_IMAGE}`, which no env file defined | `envsubst` renders an undefined variable as the empty string, so the Deployment went out with `image: ""` | +| `raycluster.yaml` hardcoded `replicas: 1` with a comment telling the reader to edit it for 2-node runs | a 2-node config (the 30B MoE block sets `ACTOR_NUM_NODES=2`) then started 8 GPU worth of workers and the job waited on placement rather than erroring. Now driven by `${ACTOR_NUM_NODES}` | +| `raycluster.yaml` requested `vpc.amazonaws.com/efa: 8` of the node's 15 allocatable | a partial request breaks point-to-point NCCL (see the rails note above) but passes an all-reduce smoke test. Now `${EFA_PER_NODE}` | +| `scripts/evaluate.sh` keyed pass@k on `prompt_item.get("idx", 0)` | the prepared AIME-2024 file has only `prompt` and `label`, so all 30 prompts collapsed onto key 0 and pass@k became "any of the 480 samples was correct". Reproduced: for a model correct on 1 prompt of 30, the old code reports 1.0000 where the truth is 0.0333 | +| `scripts/evaluate.sh` extracted answers with `\\boxed\{([^}]*)\}` | the character class stops at the first brace, so `\boxed{\frac{1}{2}}` yielded `\frac{1` and was scored wrong. Now a brace-counting scan | +| `scripts/evaluate.sh` submitted every request at once under a 300s timeout | with `MAX_TOKENS=16384` most requests time out in the queue, and the handler counts a timeout as an incorrect answer, so accuracy sags for a reason unrelated to the model. Now bounded concurrency, a generation-sized timeout, and a non-zero exit when the error rate exceeds 5% | + +The README's hardware table also said "EFA per node: 16 devices" while EKS reports 15 +allocatable on `p5en.48xlarge`; `docs/EFA_2NODE.md` still described a partial EFA request as a +throughput trade-off rather than a correctness problem; and step 4's prose offered "exec into +its head" for the checkpoint conversion two lines above a code block that correctly targets a +worker. All three are corrected. + +What was re-verified on hardware after these changes, rather than assumed: + +- All four manifests were rendered through `envsubst` from the shipped example env and + validated with `kubectl apply --dry-run=server` against a live EKS cluster. All four are + accepted, including `reward-service.yaml`, which previously could not be. +- `raycluster.yaml` was rendered at both `ACTOR_NUM_NODES=1` and `=2` and validated at each; + `replicas`, `minReplicas` and `maxReplicas` track the value, and `vpc.amazonaws.com/efa` + resolves to 15, matching what the node reports as allocatable. +- The `\boxed{}` extractor and the pass@k indexing were checked against the actual + `aime-2024.jsonl` on the cluster, which is where the missing `idx` field was confirmed. +- `bash -n` passes on every script, and the Python inside the `evaluate.sh` heredoc compiles. + +## slime-parity comparison (model x topology coverage) + +The sibling `3.test_cases/pytorch/slime` README lists a "Supported Model Sizes" matrix +(Qwen3-4B colocated, GLM-Z1-9B colocated, Qwen3-30B-A3B disaggregated, Qwen2.5-72B +disaggregated) but publishes no measured results for it. The runs below cover that matrix on +miles, `NUM_ROLLOUT=2`, reading metrics from the trainer's event files. Rows marked +(shipped recipe) use the recipes in this test case unmodified; rows marked (campaign config) +used an authored model script or a `--colocate`-conditional recipe variant and are reported as +findings, not as shipped support. + +Hardware shorthand used in the HW column: + +- **P5ENx1** = 1x `p5en.48xlarge` (8x H200, 141 GiB each). +- **P5ENx2** = 2x `p5en.48xlarge` (16x H200). +- **P6B300x2** = 2x `p6-b300.48xlarge` (16x B300, 288 GiB each) -- the expected-compatible target + for cases that do not fit H200; not run here. + +The HW column names the configuration a row was actually run on, so a "does not fit" is scoped to +that hardware rather than read as a property of the model. + +| Model | Layout | HW | Result | reward | repetition | notes | +|---|---|---|---|---|---|---| +| Qwen3-4B dense | colocated (shipped recipe) | P5ENx1 | SUCCEEDED | 0.53 | 0.0 | head co-located on GPU pool, above | +| Qwen3-4B dense | disaggregated `COLOCATE=false` (shipped recipe) | P5ENx2 | SUCCEEDED | 0.52 | 0.0 | weight sync over NCCL/EFA (`weight_version` 2, mixed 0.0); worker `replicas` must cover actor+rollout GPUs | +| GLM-Z1-9B dense | colocated TP2 (campaign config) | P5ENx1 | SUCCEEDED | 0.68 | 0.0 | TP>1 requires `CUDA_DEVICE_MAX_CONNECTIONS=1` in the Ray runtime-env | +| Qwen3-30B-A3B MoE | colocated, `moe_tp=1` pure EP (shipped recipe) | P5ENx2 | SUCCEEDED | 0.578 | 0.0 | shipped 30B block, EP_SIZE=ROLLOUT_GPUS_PER_ENGINE=2 so the rollout MoE is pure expert-parallel; trains cleanly, comparable to dense 4B. `--use-distributed-optimizer` shards the 30B optimizer state to fit H200 | +| Qwen3-30B-A3B MoE | SGLang `moe_tp>1` and `moe_ep>1`, fusion left on (campaign config) | P5ENx2 | completes, degenerate | 0.0 | 0.56 to 0.80 | root cause, below: a FlashInfer allreduce-fusion bug drops the moe-tp reduce in this combined path. Reproduced disaggregated (per-engine 4 / EP 2 -> moe_tp=2) and in stock serving | +| Qwen3-30B-A3B MoE | colocated, `moe_tp=2` x `moe_ep=2`, fusion disabled (recipe auto) | P5ENx2 | SUCCEEDED | 0.555 | 0.0 | per-engine 4 / EP 2 -> moe_tp=2; the recipe adds `--sglang-enforce-disable-flashinfer-allreduce-fusion` automatically and the combined geometry trains cleanly, comparable to pure EP (0.578) | +| Qwen2.5-72B dense | disaggregated TP4 PP2, actor8+rollout8 (campaign config) | P5ENx2 | OOM | -- | -- | did not fit on P5ENx2: ~144 GiB/GPU (8-way shard, DP=1 so Adam cannot shard) vs 141 GiB. Expected to fit **P6B300x2** (288 GiB) or an H200 layout with optimizer sharding (DP>1 / TP8 / offload) -- not run. `rms_norm_eps` is 1e-6, not 1e-5 | + +Takeaways: the disaggregated (`COLOCATE=false`) weight-sync path works on miles and is now +verified, not just argv-rendered. The dense 4B and GLM-Z1-9B cases train cleanly, and the 30B +MoE trains cleanly too once the SGLang rollout runs pure expert-parallel (`moe_tp=1`): reward +0.578, no repetition, comparable to the dense 4B run. What degenerates is not "the 30B MoE" or +"SGLang expert parallelism" -- pure expert-parallel (`moe_ep=8`) and pure tensor-parallel +(`moe_ep=1`) both generate cleanly -- but specifically the combined path where the rollout MoE +runs tensor-parallel AND expert-parallel at once (`moe_tp>1` and `moe_ep>1`). The earlier +"colocated 30B degenerates" reading conflated the Megatron TP/EP labels with the SGLang rollout +geometry; the run that degenerated had the rollout at `moe_tp>1`, and the shipped colocated +block (`moe_tp=1`) does not. See "30B MoE root cause" below. +The 72B dense case did not fit the 16-GPU disaggregated layout on H200 in the configurations +tried (DP=1 leaves the optimizer unshardable; DP>1 / TP8 / offload were not attempted) -- which +is consistent with slime listing 72B as a config without measured evidence. + +On the comparison with slime specifically: slime's "Supported Model Sizes" table lists +Qwen3-30B-A3B and Qwen2.5-72B as parallelism configurations (TP/PP and rollout/training GPU +counts), but ships no runnable env for them and reports no reward/success metric, so it is not +evidence that either trains -- the 30B is verified here (reward 0.578) and unverified on slime, +and the 72B is unverified on both: + +- The 72B layout slime tabulates (TP4 PP2, training on 8 GPUs, so 8-way sharding with DP=1) needs + roughly 18 GiB weights + 18 GiB grads + ~108 GiB Adam state = ~144 GiB per GPU with a naive + distributed optimizer that cannot shard at DP=1. That exceeds the H200's 141 GiB here (hence the + OOM) and is well above the H100 80 GiB in slime's own table -- i.e. the tabulated layout does not + fit either card as written, which is why "slime does it on p5" has no measured run behind it. + Fitting 72B GRPO needs optimizer sharding (DP>1, i.e. more actor GPUs / the colocated-16 layout), + heavier model parallel (TP8), or CPU/optimizer offload -- none of which were attempted here. + +## 30B MoE root cause + +The 30B MoE degeneration is not "the 30B model" and not "SGLang expert parallelism". It is a +single, narrow condition in the SGLang build shipped in the miles image: the rollout MoE +corrupts its own output when it runs **tensor-parallel and expert-parallel at the same time** +(`moe_tp>1` and `moe_ep>1`). Either axis alone is fine. + +SGLang derives the rollout MoE geometry from two recipe flags: +`moe_ep = --sglang-expert-parallel-size` (EP_SIZE) and +`moe_tp = --rollout-num-gpus-per-engine / EP_SIZE`. So the trigger is set by the ratio of the +per-engine GPU count to EP_SIZE, not by the Megatron TP/EP -- which is why labelling the runs +by Megatron TP/EP hid it. + +GRPO, same 30B model and cluster, `NUM_ROLLOUT=2`, metrics from the trainer's event files: + +| rollout geometry | `moe_tp` x `moe_ep` | reward | repetition | +|---|---|---|---| +| colocated, per-engine 2, EP 2 (shipped) | 1 x 2 (pure EP) | 0.578 | 0.0 | +| colocated, per-engine 2, EP 1 | 2 x 1 (pure TP) | 0.531 | 0.0 | +| disaggregated, per-engine 4, EP 2 | 2 x 2 (combined) | 0.0 | 0.56 | + +Reproduced without any RL or weight-update code, in a stock `sglang.Engine` on the same +`/fsx/models/Qwen3-30B-A3B`, `temperature=0.0`, on math prompts (4-gram repetition score): + +| `tp_size` | `ep_size` | `moe_tp` x `moe_ep` | mean repetition | sample output | +|---|---|---|---|---| +| 8 | 8 | 1 x 8 (pure EP) | 0.009 | coherent ("...Okay, so I need to solve the equation 3x + 7 = 22...") | +| 8 | 4 | 2 x 4 (combined) | 0.807 | " 7. 7. 7. 7..." | +| 8 | 2 | 4 x 2 (combined) | 0.327 | ",,,,, and and and", "10101010...", "aaaa..." | + +That the stock engine reproduces it rules out the miles RL path (weight sync, the on-policy +topk branch, Megatron) as the cause; it is in SGLang's serving path. + +What the code shows (SGLang `0.5.16.dev` in the image): + +- `models/qwen3_moe.py` `forward_normal` runs two post-experts all-reduces -- an expert-parallel + one over the moe-ep group (guarded by `self.ep_size = moe_ep_size > 1`) and a tensor-parallel + one over the moe-tp group (guarded by `self.tp_size = moe_tp_size > 1`). With `moe_tp=1` the + second is skipped; with `moe_ep=1` the first is skipped; only the combined case runs both. +- `should_skip_post_experts_all_reduce` returns `False` for both paths in this configuration + (triton runner, no dp-attention, no flashinfer/deepep A2A, no reduce-scatter), so neither + reduce is being dropped -- a missing reduce is not the cause. +- the standard dispatcher's `local_expert_mapping` is built from `moe_ep_rank` only, and the + moe-tp ranks inside one expert-parallel group correctly share it, so the global->local expert + mapping is not itself wrong. + +The `forward_normal` reduces, the moe-ep/moe-tp process groups (orthogonal by construction: for +`tp=8, ep=2` the ep groups are `{0,4},{1,5},{2,6},{3,7}` and the tp groups `{0,1,2,3},{4,5,6,7}`), +the weight sharding and the dispatcher's `local_expert_mapping` are all correct. The defect is in +the **FlashInfer allreduce+RMSNorm fusion**. On SM90/SM100 it is auto-enabled for Qwen3-MoE +(`tp_size>1`, no dp-attention, `moe_a2a=none`) regardless of `moe_ep`/`moe_tp`. With it on, both +post-experts reduces in `forward_normal` are skipped and deferred to the next layer's fused +`layernorm.forward_with_allreduce_fusion`; that fused reduce +(`layernorm.py::_forward_with_allreduce_fusion`, `flashinfer_comm_fusion.py`) picks its group with +`if moe_ep_size>1: moe_ep_group else: moe_tp_group`, treating the two axes as mutually exclusive. +When both are `>1` it reduces over the moe-ep group only and never over the moe-tp group, so each +rank keeps a partial sum over the intermediate dimension -- a total collapse from layer 0. Pure EP +and pure TP escape it because `_MOE_EP`/`_MOE_TP` alias the full TP group, so either branch reduces +over all ranks. Causal proof: rerunning the same combined config (`tp=8, ep=2, moe_tp=2`, triton) +with `enforce_disable_flashinfer_allreduce_fusion=True` generates cleanly (4-gram repetition 0.006, +same as pure EP). This is a genuine SGLang bug; a minimal stock-`sglang.Engine` reproducer is +captured for an upstream report. + +The recipe's response is to disable the fusion for the affected geometry rather than forbid it: +`run_grpo_qwen3_30b_a3b.sh` computes `moe_tp = ROLLOUT_GPUS_PER_ENGINE / EP_SIZE`, and when both +`moe_tp>1` and `moe_ep>1` it adds `--sglang-enforce-disable-flashinfer-allreduce-fusion` so the +combined geometry trains correctly too; pure EP (the shipped default) and pure TP keep the fusion. + +## A correction + +This log has now corrected the 30B MoE row twice, and both corrections are worth keeping +visible. First, an early table listed the configuration as simply "Verified" on the strength +of a smoke run that exited 0 while producing reward 0.0 and repetition 0.96 -- "the job +completed" written up as "the configuration works". The table was split to separate those +claims. Second, the follow-up reading -- "the 30B MoE degenerates, cause suspected in SGLang +expert parallelism" -- was itself too broad: it generalised from runs that happened to have +the rollout at `moe_tp>1`, and it labelled runs by Megatron TP/EP, which is not what sets the +SGLang rollout geometry. Measuring the model across the actual rollout geometries showed the +shipped colocated block trains cleanly (reward 0.578) and isolated the real trigger. The +lesson both times is the same: read the metric, name the exact variable, and do not let a +plausible summary outrun the measurement. From fea8451920a582f86d5c926ebe5580df9874bae3 Mon Sep 17 00:00:00 2001 From: littlemex Date: Wed, 19 Aug 2026 08:09:35 +0900 Subject: [PATCH 7/7] docs(miles): rewrite README to mirror the slime sibling; drop docs/ and internal narrative - Rewrite README.md to mirror the sibling slime test case structure, fold in the measurements, remove development/campaign narrative and env-specific details, reduce parenthetical asides, and add a slime-vs-miles comparison with commit-pinned upstream links. - Remove docs/ (EFA_2NODE, PORT_NOTES, VERIFICATION_LOG); content is folded into the README or dropped as internal. Recipes now reference the README. - convert_checkpoint.sh: expose the conversion GPU count as CONVERT_NUM_GPUS (default 1, still overridable by --num-gpus) instead of a hardcoded value. - env_vars examples: correct the SAVE_INTERVAL rationale (the end-of-run save succeeds; only reload and megatron2hf are untested), clarify node-role labels, and drop cluster-provisioner jargon. --- 3.test_cases/pytorch/miles/README.md | 946 ++++-------------- 3.test_cases/pytorch/miles/docs/EFA_2NODE.md | 82 -- 3.test_cases/pytorch/miles/docs/PORT_NOTES.md | 115 --- .../pytorch/miles/docs/VERIFICATION_LOG.md | 322 ------ .../pytorch/miles/env_vars.colocated.example | 15 +- .../miles/env_vars.disaggregated.example | 3 +- 3.test_cases/pytorch/miles/miles.Dockerfile | 12 +- .../miles/recipe/run_grpo_qwen3_30b_a3b.sh | 2 +- .../pytorch/miles/recipe/run_grpo_qwen3_4b.sh | 2 +- .../miles/scripts/convert_checkpoint.sh | 4 +- 10 files changed, 235 insertions(+), 1268 deletions(-) delete mode 100644 3.test_cases/pytorch/miles/docs/EFA_2NODE.md delete mode 100644 3.test_cases/pytorch/miles/docs/PORT_NOTES.md delete mode 100644 3.test_cases/pytorch/miles/docs/VERIFICATION_LOG.md diff --git a/3.test_cases/pytorch/miles/README.md b/3.test_cases/pytorch/miles/README.md index 836962fd8..810267671 100644 --- a/3.test_cases/pytorch/miles/README.md +++ b/3.test_cases/pytorch/miles/README.md @@ -1,281 +1,149 @@ # Reinforcement Learning with miles on Amazon SageMaker HyperPod EKS -This test case runs GRPO post-training with [**miles**](https://github.com/radixark/miles) on -[Amazon SageMaker HyperPod](https://aws.amazon.com/sagemaker/hyperpod/) with Amazon EKS -orchestration. It mirrors the sibling [`3.test_cases/pytorch/slime/`](../slime/) test case -end to end -- building a container image, preparing data, deploying a multi-node Ray cluster, -converting model checkpoints, and launching GRPO training on NVIDIA GPUs interconnected with -Elastic Fabric Adapter (EFA) -- but targets **miles**, a fork of SLIME built for CUDA 13 and -NVIDIA Blackwell as first-class hardware. - -The hardware-verified runs recorded here were done on a plain Amazon EKS cluster with KubeRay, -EFA, and FSx for Lustre. The same manifests are expected to run unchanged on a SageMaker -HyperPod EKS cluster -- which adds deep health checks and automatic node replacement on top of -the same EKS/KubeRay/EFA stack -- but that was not separately validated. +This test case runs GRPO post-training with [**miles**](https://github.com/radixark/miles) on [Amazon SageMaker HyperPod](https://aws.amazon.com/sagemaker/hyperpod/) with Amazon EKS orchestration, or on a plain Amazon EKS cluster. It mirrors the sibling [`3.test_cases/pytorch/slime/`](../slime/) test case: build a container image, prepare data, deploy a multi-node Ray cluster, convert model checkpoints, and launch GRPO training on NVIDIA GPUs interconnected with Elastic Fabric Adapter networking, EFA. miles is a fork of SLIME built for CUDA 13 and NVIDIA Blackwell as first-class hardware. ## Introduction -[**miles**](https://github.com/radixark/miles) is a direct fork of -[**SLIME**](https://github.com/THUDM/slime) (an LLM post-training framework for RL scaling). -It inherits SLIME's core design -- integrating two best-of-breed backends under Ray -orchestration: +[**miles**](https://github.com/radixark/miles) is a fork of [**SLIME**](https://github.com/THUDM/slime), an LLM post-training framework for RL scaling. It keeps SLIME's design of two backends under Ray orchestration: -- **[SGLang](https://github.com/sgl-project/sglang)** for high-throughput rollout generation - (inference), providing RadixAttention, continuous batching, and tensor parallelism. -- **[Megatron-LM](https://github.com/NVIDIA/Megatron-LM)** (a radixark fork) for scalable - distributed training, with support for TP, PP, CP, EP, and ZeRO-style sharding. +- **[SGLang](https://github.com/sgl-project/sglang)** for high-throughput rollout generation, providing RadixAttention, continuous batching, and tensor parallelism. +- **[Megatron-LM](https://github.com/NVIDIA/Megatron-LM)**, a radixark fork, for scalable distributed training with TP, PP, CP, EP, and ZeRO-style sharding. -**Ray** manages resource orchestration, supporting both **colocated** (training and rollout -share the same GPUs, time-sliced) and **disaggregated** (separate GPU pools connected by -NCCL/EFA weight sync) deployment topologies -- the same two topologies SLIME supports. +**Ray** manages resource orchestration and supports two deployment topologies, the same two SLIME supports. In colocated mode the training actors and the rollout engines share one GPU pool. In disaggregated mode, selected with `COLOCATE=false`, the training actors and the rollout engines run on separate GPU pools and synchronize weights over NCCL/EFA. A third, independent choice moves reward scoring to CPU nodes and is configured by `env_vars.disaggregated.example`; it can be combined with either GPU topology. -Note that "disaggregated" appears here in two unrelated senses, and the files use both. -`COLOCATE=false` splits the **actor and rollout GPU pools**, which is what selects the weight -sync implementation. `env_vars.disaggregated.example` does something else entirely: it moves -**reward scoring** to a CPU pool and leaves the GPU layout alone. The two are independent and -can be combined. +### Why miles on HyperPod? -What miles changes relative to upstream SLIME is the platform target: it ships a matched -PyTorch 2.11 / CUDA 13.0.1 stack with prebuilt `flash_attn` / Transformer Engine / `apex` -wheels and a Blackwell (sm_103) Transformer Engine FA2 whitelist patch, and it rewrites the -rollout/training weight-sync path from HTTP endpoints on an SGLang fork to direct Ray actor -methods. See [Why miles](#why-miles-vs-slime) below. +This test case adds a miles reference for GRPO post-training and keeps it structurally close to the sibling slime case, so both can be used the same way; it is not a recommendation of one over the other. The table below records how miles differs from slime, with a source link for each row. The slime column is the sibling AWS test case `3.test_cases/pytorch/slime` at commit [`63becdd3`](https://github.com/awslabs/awsome-distributed-ai/tree/63becdd38d0a047a4ec79acad893a6652df1bcbe/3.test_cases/pytorch/slime), on THUDM/slime `v0.2.4`; the miles column is [`radixark/miles`](https://github.com/radixark/miles/tree/fc04f666c08aebd72b241cef586a3939fdd6fa8e) at commit `fc04f66`. -**Amazon SageMaker HyperPod** provides purpose-built infrastructure for distributed model -training with deep health checks, automatic node replacement, and managed Kubernetes (EKS) -integration. Combined with FSx for Lustre shared storage and EFA networking, HyperPod -delivers the resilient, high-performance fabric that large-scale RL workloads demand. - -## Why miles vs SLIME - -| Aspect | SLIME | miles | +| Aspect | slime — sibling test case, THUDM/slime v0.2.4 | miles — radixark/miles | |--------|-------|-------| -| Target CUDA / GPU generation | NGC-image CUDA stack, Hopper-first | CUDA 13.0.1 native, NVIDIA Blackwell (sm_103) first-class | -| Base image | NGC PyTorch (nightly ABI) | `radixark/miles` (matched PyTorch 2.11 + cu130 stable ABI) | -| Rollout <-> training weight sync | HTTP endpoints on an SGLang fork | Ray actor methods (`begin_weight_update` / `pull_weights`) on the rollout engine directly | -| Framework install path (in image) | `/opt/slime` | `/root/miles` (editable install) | -| Megatron fork | NVIDIA/Megatron-LM (SLIME-compatible commit) | radixark/Megatron-LM fork | -| SGLang version pinned | 0.5.12.post1 | 0.5.16.dev | -| `train.py` CLI / GRPO flags | -- | Compatible with SLIME's (same names, same semantics) | -| Relationship | Upstream | Direct fork of SLIME (diverged at commit `fcce96ca0`, 2025-10-05) | - -Because miles is a fork rather than an independent reimplementation, the `train.py` CLI used -by this test case's recipes is drop-in compatible with SLIME's -- the [Quick -Start](#quick-start) below tracks the sibling slime test case step for step. +| Base image | NGC PyTorch container [`nvcr.io/nvidia/pytorch:26.02-py3`](https://github.com/awslabs/awsome-distributed-ai/blob/63becdd38d0a047a4ec79acad893a6652df1bcbe/3.test_cases/pytorch/slime/slime.Dockerfile#L8) | [`radixark/miles`](https://github.com/radixark/miles/blob/fc04f666c08aebd72b241cef586a3939fdd6fa8e/docker/Dockerfile#L14-L23), built on `lmsysorg/sglang:v0.5.16` with [torch 2.11.0 / CUDA 13.0.1](https://github.com/sgl-project/sglang/blob/fdebc938f7f4d16fe6b9f55dcd9a767cf0899ea1/docker/Dockerfile#L1) | +| SGLang | [`0.5.12.post1`](https://github.com/awslabs/awsome-distributed-ai/blob/63becdd38d0a047a4ec79acad893a6652df1bcbe/3.test_cases/pytorch/slime/slime.Dockerfile#L20) | the [`sglang-miles`](https://github.com/radixark/miles/blob/fc04f666c08aebd72b241cef586a3939fdd6fa8e/docker/Dockerfile#L14) branch, a development build based on `v0.5.16` | +| Megatron fork | [NVIDIA/Megatron-LM](https://github.com/awslabs/awsome-distributed-ai/blob/63becdd38d0a047a4ec79acad893a6652df1bcbe/3.test_cases/pytorch/slime/slime.Dockerfile#L18) `@3714d81` | [radixark/Megatron-LM](https://github.com/radixark/miles/blob/fc04f666c08aebd72b241cef586a3939fdd6fa8e/docker/Dockerfile#L22-L23) `miles-main` | +| Framework install path | [`/opt/slime`](https://github.com/awslabs/awsome-distributed-ai/blob/63becdd38d0a047a4ec79acad893a6652df1bcbe/3.test_cases/pytorch/slime/slime.Dockerfile#L189) | [`/root/miles`](https://github.com/radixark/miles/blob/fc04f666c08aebd72b241cef586a3939fdd6fa8e/docker/Dockerfile#L238) | +| Weight-sync path | `SGLangEngine(RayActor)` methods called over Ray, each issuing an HTTP request to its local SGLang server | same design; miles adds a two-phase [`begin_weight_update` / `end_weight_update`](https://github.com/radixark/miles/blob/fc04f666c08aebd72b241cef586a3939fdd6fa8e/miles/backends/sglang_utils/sglang_engine.py#L636) session on top | +| Relationship | upstream | fork of SLIME, [shared ancestor commit `fcce96ca0`](https://github.com/radixark/miles/commit/fcce96ca06e47e92f685db91c2a7327fff095906), dated 2025-10-06 | -## Architecture +The `train.py` CLI and GRPO flags are compatible with SLIME's, so the recipes here are drop-in and the [Quick Start](#quick-start) tracks the sibling test case step for step. -The deployment uses a colocated architecture where training and rollout share the same GPU -pool, connected by EFA networking and FSx for Lustre shared storage. (A disaggregated -topology -- separate GPU pools for training and rollout -- is also supported by miles/Ray but -is not the hardware-verified path documented here; see [Verification Status](#verification-status).) +## Architecture ``` -+-----------------------------------------------------------------------+ -| Amazon SageMaker HyperPod EKS Cluster | -| | -| +-----------------------------+ +-----------------------------+ | -| | Node 1 (GPU worker) | | Node 2 (GPU worker) | | -| | 8x H-series GPU | EFA | | 8x H-series GPU | EFA | | -| +-----------------------------+ +-----------------------------+ | -| | | | | | -| +--------------EFA (GPU<->GPU RDMA)-----------+ | -| | | | -| +----------------------------------------------------------+ | -| | FSx for Lustre, mounted as PVC `fsx-claim` | | -| | /fsx/models /fsx/data /fsx/runs | | -| +----------------------------------------------------------+ | -| | -| +-----------------------------+ | -| | Ray head (num-gpus 0) | <- co-located on the GPU pool; | -| | on a GPU node (needs CUDA) | needs libcuda for Megatron impt | -| +-----------------------------+ | -| | -| Kubernetes Resources: | -| - KubeRay operator | -| - EFA device plugin, NVIDIA device plugin (kube-system) | -| - FSx CSI driver (kube-system) | -+-----------------------------------------------------------------------+ ++---------------------------------------------------------------------+ +| Amazon EKS Cluster, HyperPod-compatible | +| 2x p5en.48xlarge, EKS orchestration | +| | +| +----------------------------+ +----------------------------+ | +| | Node 1: p5en.48xlarge | | Node 2: p5en.48xlarge | | +| | 8x H200 141GB 16x EFA | | 8x H200 141GB 16x EFA | | +| | ~2TB RAM 192 vCPU | | ~2TB RAM 192 vCPU | | +| +----------------------------+ +----------------------------+ | +| | | | +| +----- EFA 3200 Gbps/node, GPU<->GPU RDMA -----+ | +| | +| +---------------------------------------------------------+ | +| | FSx for Lustre, RWX, mounted at /fsx | | +| | /fsx/models /fsx/data /fsx/runs | | +| +---------------------------------------------------------+ | +| | +| Kubernetes Resources: | +| - KubeRay operator, kuberay-operator namespace | +| - EFA device plugin, kube-system | +| - FSx CSI driver, kube-system | +| - NVIDIA device plugin, kube-system | ++---------------------------------------------------------------------+ ``` -### miles Internal Loop +The diagram shows the 2-node layout. A single node is also supported for the dense 4B case; the multi-node layout adds the second node for MoE and disaggregated runs. + +## miles Internal Architecture ``` - +------------------+ - | Data Buffer | - | (prompt queue + | - | rollout cache) | - +--------+---------+ - | - +--------------+--------------+ - | | - v v - +-------------------+ +-------------------+ - | Rollout | | Training | - | (SGLang engines, | | (Megatron-LM | - | Ray actors) | <----> | TP/PP/CP/EP) | - | | weight | | - | - RadixAttention | sync | - GRPO | - | - Cont. batching | (Ray | - Dynamic batch | - | - TP per engine | actor | - Gradient ckpt | - | | calls) | | - +-------------------+ +-------------------+ + +----------------------+ + | Data Buffer | + | prompt queue and | + | rollout cache | + +----------+-----------+ + | + +----------------+-----------------+ + | | + v v + +-----------------------+ +-----------------------+ + | Rollout | | Training | + | SGLang Ray actors | | Megatron-LM | + | RadixAttention | | TP / PP / CP / EP | + | continuous batching | | GRPO, dynamic batch | + +-----------------------+ +-----------------------+ + ^ ^ + | weight sync | + +----------------------------------+ ``` -1. **Data Buffer** manages prompts, dispatches them for rollout, and stores generated samples - with rewards. -2. **Rollout** runs SGLang engines as Ray actors, generating responses and scoring them via a - reward function. Unlike SLIME, weight-sync calls (`begin_weight_update` / `pull_weights`) - go directly to Ray actor methods on the rollout engine rather than HTTP endpoints. -3. **Training** reads batches from the Data Buffer, computes GRPO advantages, and updates the - policy via Megatron-LM. Updated weights are synced back to the rollout engines. +The Data Buffer manages prompts, dispatches them for rollout, and stores generated samples with rewards. Rollout runs SGLang engines as Ray actors, generating responses and scoring them with a reward function. Training reads batches from the buffer, computes GRPO advantages, updates the policy with Megatron-LM, and syncs the updated weights back to the rollout engines. ## Hardware Requirements -miles's own image already builds for Hopper and Blackwell alike (CUDA 13.0.1, with the -sm_103 Transformer Engine FA2 whitelist patch applied on top of the standard sm_90 build), -and nothing in this test case's recipes, env files, or manifests hard-codes a GPU generation -or CUDA/SM version -- the only per-cluster knob is which accelerator NodePool to target -(`GPU_NODE_ROLE` in `env_vars`, substituted into `kubernetes/raycluster.yaml`'s -`nodeSelector`). The table below is the hardware this was actually run on; a B300 row is -included as the expected-compatible target (see Verification Status for what is and is not -confirmed on real hardware). - -| Component | H200 (verified) | B300 (expected-compatible, unverified) | -|-----------|------------------|------------------------------------------| -| **Instance type** | p5en.48xlarge | p6-b300.48xlarge | -| **GPUs per node** | 8x NVIDIA H200 (141 GB HBM) | 8x NVIDIA B300 (288 GB HBM) | -| **EFA per node** | 16 physical, **15 allocatable** on EKS | 16 physical, verify allocatable | -| **Storage** | FSx for Lustre, mounted as PVC `fsx-claim` | same | -| **Kubernetes** | EKS, KubeRay operator | same | -| **Ray head placement** | The head co-locates on the GPU pool (`num-gpus 0`); it must be on a CUDA-capable node (miles control actors import Megatron even at `num-gpus 0`). Give the GPU node's root volume >=150 GiB for the ~18 GB image | same | - -B300's ~2x larger HBM is expected to relax the memory-driven choices made for H200 in -`env_vars.colocated.example` (e.g. the 30B MoE case needing `--use-distributed-optimizer` and -a colocated 16-GPU layout to avoid OOM at 8 GPU -- see the ALTERNATE block in that file); it -should not require code changes, only re-tuning env values once measured on B300. - -### Verification Status - -This test case is presented as a GRPO reference, not a research study -- the table below -states plainly what has been run on real hardware and what has not, so you know exactly what -you are inheriting. - -Every row below is backed by a recorded run: see [docs/VERIFICATION_LOG.md](./docs/VERIFICATION_LOG.md) for the submission commands, Ray job ids, the flags each -job actually received, and the metric values read back from the trainer's event files. - -| Component | Status | Environment / Note | -|-----------|--------|--------------------| -| Qwen3-4B GRPO, colocated, 1 node (8 GPU) | Verified | p5en.48xlarge, H200x8 | -| Qwen3-4B GRPO, colocated, 2 nodes (16 GPU) over EFA, 3 rollout cycles | Verified | 2x p5en.48xlarge; all 3 cycles SUCCEEDED in 676s with `raw_reward` 0.477/0.523/0.492 -- metrics in [docs/VERIFICATION_LOG.md](./docs/VERIFICATION_LOG.md). Fabric separately measured at NCCL all_reduce busbw 190-257 GB/s over EFA (`efa-direct` + GPUDirect RDMA, no TCP fallback), see [docs/EFA_2NODE.md](./docs/EFA_2NODE.md) | -| Qwen3-30B-A3B MoE GRPO, colocated, 2 nodes (16 GPU) -- trains cleanly | Verified | actor spans all 16 GPU with `--use-distributed-optimizer` (required -- confining the actor to 8 GPU OOMs) and the SGLang `triton` MoE runner backend. The shipped block runs the rollout MoE pure expert-parallel (`moe_tp=1`, `EP_SIZE = ROLLOUT_GPUS_PER_ENGINE = 2`): `rollout/raw_reward` 0.578, `rollout/repetition_frac` 0.0, `weight_version` uniform / `mixed_version_ratio` 0.0 -- comparable to the dense 4B run. Running the MoE tensor-parallel AND expert-parallel at once (`moe_tp>1` and `moe_ep>1`) hits a FlashInfer allreduce-fusion bug in this build, so the recipe disables that fusion for such geometries and they train too -- see [Known Issues](#known-issues) item 2 | -| Qwen3-4B GRPO, **disaggregated** (`COLOCATE=false`), 2 nodes | Verified | run on 2x p5en (actor 8 + rollout 8): Ray job SUCCEEDED, `raw_reward` 0.52, `repetition_frac` 0.0, `weight_version` 2 with `mixed_version_ratio` 0.0 -- i.e. weights synced across the node boundary over NCCL/EFA (`UpdateWeightFromDistributed`). A disaggregated run needs GPU nodes for BOTH the actor and rollout pools, so `WORKER_REPLICAS` (derived in `env_vars` as `ceil((ACTOR_NUM_NODES*ACTOR_GPUS_PER_NODE + ROLLOUT_NUM_GPUS)/8)`, here 2) drives the RayCluster worker count rather than `ACTOR_NUM_NODES` alone (see [docs/VERIFICATION_LOG.md](./docs/VERIFICATION_LOG.md)) | -| Qwen3-30B-A3B MoE GRPO, **disaggregated** | UNVERIFIED | the disaggregated actor-8 layout needs B300-class 288GB HBM; only the colocated 16-GPU layout was run, on H200 | -| Any workload on B300 (p6-b300.48xlarge) | UNVERIFIED | expected-compatible (miles's base image already targets CUDA 13 / sm_103; nothing here is GPU-generation-specific) but not run on this hardware -- see Hardware Requirements | -| RayCluster manifest as shipped (head co-located on the GPU pool) | Verified | the shipped `kubernetes/raycluster.yaml` schedules the head onto the GPU pool (`num-gpus 0`, with a `nvidia.com/gpu` toleration) alongside the worker: dense Qwen3-4B GRPO SUCCEEDED with `raw_reward` 0.53 and `repetition_frac` 0.0. The head must be on a CUDA-capable node -- miles control actors import Megatron even at `num-gpus 0`, so a CPU-only head node fails with `libcuda.so.1: cannot open shared object file`. See [docs/VERIFICATION_LOG.md](./docs/VERIFICATION_LOG.md) | -| Disaggregated reward service (`remote_rm` on a CPU pool) | UNVERIFIED | `reward_service/`, `kubernetes/reward-service.yaml` are present and mirror the sibling slime test case, but were not deployed/exercised on miles | -| Checkpoint save-back / long-run checkpointing | Known Issue | `save_model()` fails with a pickle-truncation error in Megatron's distributed checkpoint save; see [Known Issues](#known-issues) | +This test case is validated on the following configuration: + +| Component | Specification | +|-----------|---------------| +| **Instance type** | p5en.48xlarge | +| **Nodes** | 1 for single-node, or 2 for multi-node and MoE | +| **GPUs per node** | 8x NVIDIA H200 141GB | +| **Total GPUs** | 8 on one node, 16 on two nodes | +| **GPU memory** | 1,128 GB per node aggregate | +| **Host RAM per node** | ~2 TB | +| **EFA per node** | 16 devices, 15 allocatable on EKS | +| **Storage** | FSx for Lustre, RWX, mounted at `/fsx` | +| **Kubernetes** | EKS, KubeRay operator | + +Other instance types are expected to work with resource-value retuning. p6-b300.48xlarge, a Blackwell instance, is expected-compatible because miles's base image targets CUDA 13 / sm_103 and nothing here hard-codes a GPU generation, but it has not been verified. + +## Supported Model Sizes + +The TP and PP columns describe Megatron training-side parallelism, which is distinct from the SGLang rollout MoE geometry `moe_tp` / `moe_ep`. The shipped 30B MoE recipe runs the rollout pure expert-parallel, `moe_tp=1`; see [Known Issues](#known-issues). + +| Model | Parameters | Topology | TP | PP | Rollout GPUs | Training GPUs | +|-------|-----------|----------|----|----|-------------|---------------| +| Qwen3-4B | 4B Dense | Colocated | 1 | 1 | 8, shared | 8, shared | +| GLM-Z1-9B | 9B Dense | Colocated | 2 | 1 | 16, shared | 16, shared | +| Qwen3-30B-A3B | 30B MoE | Colocated | 2 | 1 | 16, shared | 16, shared | +| Qwen2.5-72B * | 72B Dense | Disaggregated | 4 | 2 | 8 | 8 | + +\* Qwen2.5-72B does not fit this 16-GPU H200 layout; see the Validation table and Known Issues. + +## Validation + +Each configuration below was launched on 2x p5en.48xlarge and confirmed to complete within the listed wall time. Reward and repetition are the last scalar from the trainer's TensorBoard event files; they indicate the loop closes and generation is healthy, not convergence. Known limitations are the 30B MoE rollout degenerating in one specific parallelism geometry and Qwen2.5-72B not fitting the 16-GPU H200 layout; both are in [Known Issues](#known-issues). + +| Config | reward | repetition | wall time | +|--------|--------|------------|-----------| +| Qwen3-4B dense, colocated 1 node | 0.531 | 0.0 | ~13 min | +| Qwen3-4B dense, disaggregated 2 nodes | 0.523 | 0.0 | ~12 min | +| GLM-Z1-9B dense, colocated TP2 | 0.680 | 0.0 | ~13 min | +| Qwen3-30B-A3B MoE, colocated pure EP, `moe_tp=1` | 0.578 | 0.0 | ~21 min | +| Qwen3-30B-A3B MoE, colocated pure TP, `moe_ep=1` | 0.531 | 0.0 | ~25 min | +| Qwen3-30B-A3B MoE, disaggregated pure EP | 0.65 | 0.0 | ~18 min | +| Qwen3-30B-A3B MoE, colocated combined `moe_tp=2` x `moe_ep=2` | 0.555 | 0.0 | ~21 min | +| Qwen3-30B-A3B MoE, combined geometry, FlashInfer fusion on | 0.0 | 0.56 | ~20 min, degenerate output | +| Qwen2.5-72B dense, disaggregated | OOM | -- | did not fit 16x H200 | ## Prerequisites -This test case does not create cluster infrastructure; it deploys onto a cluster that already -provides the pieces below. The manifests reference them by label/name, so if any is missing the -failure is a scheduling or mount error, not an obvious message. Confirm each before starting. - -1. An Amazon SageMaker HyperPod cluster with EKS orchestration and GPU instance groups - (p5en.48xlarge/H200 -- verified -- or p6-b300.48xlarge/B300 -- expected-compatible, - see Hardware Requirements -- with EFA). -2. **Node placement labels.** `kubernetes/raycluster.yaml` schedules the GPU workers on - `${GPU_NODE_LABEL_KEY}: ${GPU_NODE_ROLE}` and the Ray head on - `${CPU_NODE_LABEL_KEY}: ${CPU_NODE_ROLE}`. There is no universal `node-role` label; both the - key and the value are cluster-specific. You can either point the four env vars at a label your - nodes already carry -- e.g. `GPU_NODE_LABEL_KEY=node.kubernetes.io/instance-type`, - `GPU_NODE_ROLE=p5en.48xlarge`, or on SageMaker HyperPod the - `sagemaker.amazonaws.com/instance-group-name` of your GPU group -- or add your own label: - `kubectl label node node-role=gpu` (then `GPU_NODE_LABEL_KEY=node-role`, - `GPU_NODE_ROLE=gpu`). If you have no dedicated CPU pool, set the CPU_ vars equal to the GPU - ones (the head runs `num-gpus 0`). A wrong key OR value leaves pods `Pending` with - `FailedScheduling`, not an obvious error. -3. **Cluster add-ons that advertise the scheduled resources**, all of which the manifests - request and none of which this test case installs: - - the NVIDIA device plugin (or GPU Operator, or a GPU AMI that bundles it) advertising - `nvidia.com/gpu`; - - the AWS EFA Kubernetes device plugin (`aws-efa-k8s-device-plugin`) advertising - `vpc.amazonaws.com/efa` -- read the allocatable count off a node for `EFA_PER_NODE` - (`kubectl get node -o jsonpath='{.status.allocatable.vpc\.amazonaws\.com/efa}'`); - - the FSx for Lustre CSI driver, bound to the `fsx-claim` PVC below. -4. `kubectl` and `helm` configured to access the cluster. -5. The KubeRay operator installed (see step 0 below). -6. The Ray head co-locates on the GPU pool (it runs `num-gpus 0` and uses no GPU, only CPU/disk). - It must be on a CUDA-capable node: miles control actors import Megatron/`transformer_engine` - even at `num-gpus 0`, which loads `libcuda.so.1` at import, so a CPU-only head node fails. - Ensure the GPU node's root volume has ample ephemeral-storage (>=150 GiB) for the ~18 GB image. -7. An FSx for Lustre `PersistentVolumeClaim` named `fsx-claim`, RWX, mounted at `/fsx`. -8. **EFA security group (multi-node / `COLOCATE=false` only).** The EFA node security group must - allow all traffic to itself on BOTH ingress AND egress (self-referencing). EFA's OS-bypass - SRD traffic is not ordinary IP, so a CIDR-only egress rule does not authorize it and NCCL - over EFA fails with `Unreachable remote` / `Unexpected number of remote rails`. See - [docs/EFA_2NODE.md](./docs/EFA_2NODE.md). -9. Container registry access (e.g. Amazon ECR) for building/pushing images. -10. A Hugging Face account and access token for model downloads. -11. A Kubernetes Secret named `hf-token` in the target namespace. `kubernetes/raycluster.yaml` - mounts it without `optional: true`, so the Ray pods will not start without it; if your model - is public, create it with an empty value (step 1 has the command). - -### Cluster assumptions and portability - -`kubernetes/raycluster.yaml` was validated on an EKS cluster meeting the Prerequisites above -(a `p5en.48xlarge` / H200 pool with the EFA and FSx add-ons) and carries a few values sized for -that node; on a different cluster, adjust these before deploying or the failure is a silent -`Pending`/`ImagePullBackOff`/OOM rather than a message: - -- **EFA is requested unconditionally, including the single-node colocated run.** Without the EFA - device plugin (or on an instance type without EFA), the worker is `Unschedulable`. For a - single-node run on a non-EFA cluster, delete the two `vpc.amazonaws.com/efa` lines from the - worker `resources`; multi-node runs require EFA (verified) and the self-referencing SG above. -- **Worker resources are sized for `p5en.48xlarge` (2 TiB RAM):** cpu 90/96, memory 1800/1900 Gi, - and a 256 Gi memory-backed `/dev/shm` (which counts against the pod memory limit). On a smaller - GPU node, lower these below the node's allocatable or the worker will not schedule. -- **Every node pulls the ~18 GB image**, not just the head: give GPU node root volumes >=100 GiB - free as well (the head needs >=150 GiB, prerequisite 6). -- **Registry auth is assumed to come from the node IAM role (ECR).** For a private or - cross-account registry, add `imagePullSecrets` to the pod specs. -- **`fsx-claim` must be `ReadWriteMany` and exist in the target `${NAMESPACE}`.** A PVC is - namespaced, so one created in `default` is invisible to a run in another namespace; an - `ReadWriteOnce` (e.g. EBS) claim fails to mount across a multi-node run. -- **The disaggregated reward-service overlay** (`kubernetes/reward-service.yaml`, UNVERIFIED) - selects nodes by the HyperPod label `sagemaker.amazonaws.com/instance-group-name`, which a - non-HyperPod cluster does not have; edit its `nodeSelector` for your CPU pool if you deploy it. -- **One worker pod consumes a whole 8-GPU node** (`nvidia.com/gpu: 8`, `num-gpus: '8'`). On - nodes with a different GPU count, change these together with `ACTOR_GPUS_PER_NODE` and the - actor/rollout GPU split; the manifest is not a fractional-GPU RayCluster. -- **If you point `CPU_NODE_ROLE` at the GPU pool** (no dedicated CPU pool), the head also needs - a toleration for that pool's taint (commonly `nvidia.com/gpu:NoSchedule`), or it stays Pending - even though the label matches. Add it to the head pod spec. -- **`/fsx` must already hold the model and data before you launch.** This test case does not - download or convert during training: `MODEL_LOCAL` (HF checkpoint), `MODEL_DIST` (Megatron - `torch_dist`, from step 4), `PROMPT_DATA`, and `EVAL_DATA` must all exist. Quick Start steps 3 - and 4 create them. - -Preflight (run after `source env_vars`, before deploying) -- turns a silent `Pending`/mount -failure into an early, named error: +1. A SageMaker HyperPod EKS cluster, or a plain EKS cluster, with a p5en.48xlarge GPU instance group and EFA. Validated on p5en.48xlarge; other instance types may need resource-value retuning. +2. An FSx for Lustre `PersistentVolumeClaim` mounted at `/fsx`. The claim name is set through `FSX_CLAIM` and defaults to `fsx-claim`, matching the sibling slime test case. +3. Amazon ECR access for building and pushing the image. +4. A Hugging Face account and access token for model downloads. -```bash -: "${NAMESPACE:?}" "${FULL_IMAGE:?}" "${FSX_CLAIM:?}" "${GPU_NODE_ROLE:?}" "${CPU_NODE_ROLE:?}" "${EFA_PER_NODE:?}" "${WORKER_REPLICAS:?}" -kubectl get nodes -l "${GPU_NODE_LABEL_KEY}=${GPU_NODE_ROLE}" -o name # GPU pool exists? -kubectl get pvc "${FSX_CLAIM}" -n "${NAMESPACE}" # PVC Bound, RWX, this ns? -kubectl get secret hf-token -n "${NAMESPACE}" # hf-token present? -kubectl get crd rayclusters.ray.io # KubeRay installed? -# and on /fsx (from a pod that mounts it): -# ls -ld "$MODEL_LOCAL" "$MODEL_DIST"; test -s "$PROMPT_DATA"; test -s "$EVAL_DATA" -``` +The KubeRay operator is installed in step 0 below. ## Quick Start -### 0. Install the KubeRay Operator (one-time per cluster) +The default path is Qwen3-4B dense, colocated on one node: it builds the image, prepares the model and data, deploys a Ray cluster, and runs a short GRPO loop. To run the 30B MoE case instead, uncomment the `ALTERNATE` block in `env_vars.colocated.example` and launch the MoE recipe in step 7. + +### 0. Install the KubeRay Operator, once per cluster -The Ray cluster is managed by the KubeRay operator. If it is not already present -(`kubectl get crd rayclusters.ray.io`), install it with Helm: +If `kubectl get crd rayclusters.ray.io` returns nothing, install it with Helm: ```bash helm repo add kuberay https://ray-project.github.io/kuberay-helm/ @@ -289,594 +157,212 @@ helm install kuberay-operator kuberay/kuberay-operator \ ```bash cd 3.test_cases/pytorch/miles cp env_vars.colocated.example env_vars -# Edit env_vars with your cluster-specific values +# Edit env_vars for your cluster, then: source env_vars ``` -Key variables (see `env_vars.colocated.example` for the full annotated file): +`env_vars` includes `NAMESPACE`, which must already exist. The Hugging Face token is read from a Kubernetes Secret named `hf-token`, wired into the pods with `secretKeyRef`, so create it once per namespace. This is the same flow as the sibling slime test case: ```bash -# AWS / ECR (region and account are auto-derived from your AWS credentials) -export AWS_REGION="${AWS_REGION:-$(aws configure get region)}" -export AWS_ACCOUNT_ID="$(aws sts get-caller-identity --query Account --output text)" -export REGISTRY="${AWS_ACCOUNT_ID}.dkr.ecr.${AWS_REGION}.amazonaws.com/" -export IMAGE="miles-hyperpod" -export TAG="miles-cuda13-efa-0.1" # pinned image tag (avoid :latest) -export FULL_IMAGE="${REGISTRY}${IMAGE}:${TAG}" - -# Model (Qwen3-4B dense, bf16, colocated -- the verified default) -export MODEL_NAME="Qwen/Qwen3-4B" -export MODEL_LOCAL="/fsx/models/Qwen3-4B" # SGLang rollout init -export MODEL_DIST="/fsx/models/Qwen3-4B_torch_dist" # Megatron training (torch_dist) -export COLOCATE="true" -export ACTOR_NUM_NODES=1 -export ACTOR_GPUS_PER_NODE=8 -export ROLLOUT_NUM_GPUS=8 - -# Cluster -export NAMESPACE="default" -export FSX_CLAIM="fsx-claim" -export GPU_NODE_ROLE="gpu-p5en" # node-role label on your GPU pool -export EFA_PER_NODE=15 # the node's FULL allocatable EFA count -- see env_vars +kubectl create secret generic hf-token --from-literal=HF_TOKEN=hf_xxx -n "${NAMESPACE}" +# Public model with no token: create it with an empty value so the key exists. ``` -`ACTOR_NUM_NODES` drives the RayCluster's worker replica count as well as the actor layout, so -one variable sets both and the manifest cannot disagree with the recipe. Raise it to 2 for the -2-node (16 GPU) configurations. - -The pods read `HF_TOKEN` from a Kubernetes Secret (not from `env_vars` / the Ray runtime-env), -so it never lands in the Ray dashboard's job-submission metadata. Create it once per -namespace. This Secret is **required**, not optional: `kubernetes/raycluster.yaml` references -it without `optional: true`, so if it is absent the Ray pods never start and report -`CreateContainerConfigError` rather than anything about HuggingFace. If the model you point at -is public and you have no token, create the Secret with an empty value -- the pods only need -the key to exist: - -```bash -kubectl create secret generic hf-token \ - --from-literal=HF_TOKEN=hf_xxx \ - -n "${NAMESPACE}" - -# Public model and no token? The pods only need the key to exist: -# kubectl create secret generic hf-token --from-literal=HF_TOKEN= -n "${NAMESPACE}" -``` - -`env_vars.colocated.example` also ships a commented-out **ALTERNATE** block for -Qwen3-30B-A3B MoE, colocated on 2 nodes (16 GPU) -- uncomment it (and the matching recipe in -step 7) to run the MoE configuration instead of the 4B dense default. It trains cleanly as -shipped (`rollout/raw_reward` 0.578, `rollout/repetition_frac` 0.0), running the rollout MoE -pure expert-parallel (`moe_tp=1`, `EP_SIZE = ROLLOUT_GPUS_PER_ENGINE`). Geometries that run it -tensor-parallel and expert-parallel at once (`moe_tp>1` and `moe_ep>1`) hit a FlashInfer -allreduce-fusion bug in this build; the recipe disables that fusion for them so they train too -(see [Known Issues](#known-issues) item 2). - ### 2. Build and Push the Container Image -The image takes `radixark/miles:` (which already bundles miles, SGLang, -Megatron-LM, and the matched PyTorch 2.11 / CUDA 13.0.1 stack) as its base and adds **only** -the AWS EFA networking layer (`miles.Dockerfile`) -- GDRCopy, the EFA installer, and NCCL/EFA -runtime defaults. Pin the base by digest (`MILES_BASE_DIGEST` build arg, with `MILES_BASE_TAG` kept alongside for readability) and pin the resulting image -tag to dated/versioned values; never use `:latest`. +The image takes `radixark/miles` as its base and adds only the AWS EFA layer. The base is pinned by `sha256` digest in `miles.Dockerfile`. ```bash -# Authenticate to ECR -aws ecr get-login-password --region ${AWS_REGION} | \ - docker login --username AWS --password-stdin ${REGISTRY} - -# Create repository (first time only) +aws ecr get-login-password --region ${AWS_REGION} | docker login --username AWS --password-stdin ${REGISTRY} aws ecr create-repository --repository-name ${IMAGE} --region ${AWS_REGION} || true - -# Build image (build context is this test-case directory) docker build -t ${FULL_IMAGE} -f miles.Dockerfile . docker push ${FULL_IMAGE} ``` -If the cluster has no node with local Docker access, `kubernetes/buildkit-job.yaml` builds -and pushes the image **in-cluster** with a rootless buildkit Job instead (CPU/IO-only, no GPU -required) -- edit its `nodeSelector` to a node with >=120 GiB free ephemeral-storage, then -create the ConfigMap and Secret it mounts and apply it through `envsubst` (like -`raycluster.yaml` above) so `${AWS_ACCOUNT_ID}`/`${AWS_REGION}` resolve: +If no node has a local Docker daemon, `kubernetes/buildkit-job.yaml` builds and pushes the image in-cluster with a rootless BuildKit Job: ```bash -kubectl create configmap miles-build-context \ - --from-file=Dockerfile=miles.Dockerfile -n "${NAMESPACE}" +kubectl create configmap miles-build-context --from-file=Dockerfile=miles.Dockerfile -n "${NAMESPACE}" kubectl create secret docker-registry ecr-miles-push \ - --docker-server="${REGISTRY}" --docker-username=AWS \ - --docker-password="$(aws ecr get-login-password --region "${AWS_REGION}")" \ - -n "${NAMESPACE}" + --docker-server="${REGISTRY}" --docker-username=AWS \ + --docker-password="$(aws ecr get-login-password --region ${AWS_REGION})" -n "${NAMESPACE}" envsubst < kubernetes/buildkit-job.yaml | kubectl apply -f - -kubectl -n "${NAMESPACE}" logs -f job/miles-efa-build ``` ### 3. Download and Prepare the Model +As in the sibling slime test case, this uses a data-prep pod and `huggingface-cli` inside it: + ```bash -# Create a data-prep pod envsubst < kubernetes/data-prep-pod.yaml | kubectl apply -f - -kubectl exec -it data-prep -- bash - +kubectl wait --for=condition=Ready pod/data-prep -n "${NAMESPACE}" --timeout=300s +kubectl exec -it data-prep -n "${NAMESPACE}" -- bash # Inside the pod: -pip install huggingface_hub -# HF_TOKEN comes from the hf-token Secret. If you created that Secret with an empty -# value because the model is public, skip this login. -huggingface-cli login --token "${HF_TOKEN}" - -# Download model +pip install -U "huggingface_hub[cli]" huggingface-cli download Qwen/Qwen3-4B --local-dir /fsx/models/Qwen3-4B - -# Download training + evaluation datasets -huggingface-cli download --repo-type dataset zhuzilin/dapo-math-17k \ - --local-dir /fsx/data/dapo-math-17k -huggingface-cli download --repo-type dataset zhuzilin/aime-2024 \ - --local-dir /fsx/data/aime-2024 +huggingface-cli download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /fsx/data/dapo-math-17k +huggingface-cli download --repo-type dataset zhuzilin/aime-2024 --local-dir /fsx/data/aime-2024 ``` -### 4. Convert Model Weights to Megatron Format - -miles's Megatron training backend requires weights in `torch_dist` format. Use the bundled -conversion helper, which sources the model script (`scripts/models/.sh`) from the -miles install at `/root/miles`: - -Run this **inside a pod built from the miles image**, not on your workstation: the script reads -the model scripts from `/root/miles`, writes to `/fsx`, and uses `torchrun` when `--num-gpus` is -passed. A **GPU worker** of the RayCluster is such a pod, so deploy step 5 first and come back -here, or run the same command as a one-off Job on the GPU node pool. It must not be the head: -the head runs with `num-gpus 0` and `NVIDIA_VISIBLE_DEVICES=void`, where importing CUDA fails -(see `docs/PORT_NOTES.md`). +### 4. Deploy the Ray Cluster ```bash -# Must be a GPU worker, not the head -- see above. -W=$(kubectl get pod -n "${NAMESPACE}" -l ray.io/node-type=worker -o name | head -1) -W=${W#pod/} -kubectl cp scripts/convert_checkpoint.sh "${NAMESPACE}/${W}:/tmp/convert_checkpoint.sh" -kubectl exec -n "${NAMESPACE}" "${W}" -- \ - bash /tmp/convert_checkpoint.sh hf2megatron \ - --model-script qwen3-4B.sh \ - --hf-path /fsx/models/Qwen3-4B \ - --save-path /fsx/models/Qwen3-4B_torch_dist +source env_vars +envsubst < kubernetes/raycluster.yaml | kubectl apply -f - +kubectl get pods -w -n "${NAMESPACE}" -l ray.io/is-ray-node=yes # Ctrl-C once head and workers are Running +kubectl port-forward -n "${NAMESPACE}" svc/miles-ray-head-svc 8265:8265 & ``` -For larger models, pass `--num-gpus 8` to parallelize the conversion with `torchrun`: +The manifest schedules GPU workers on the GPU pool and the Ray head on the same pool. The head uses no GPU but must sit on a CUDA-capable node; see [Known Issues](#known-issues). -```bash -bash scripts/convert_checkpoint.sh hf2megatron \ - --model-script qwen3-30B-A3B.sh \ - --hf-path /fsx/models/Qwen3-30B-A3B \ - --save-path /fsx/models/Qwen3-30B-A3B_torch_dist \ - --num-gpus 8 -``` +### 5. Convert Model Weights to Megatron Format -### 5. Deploy the Ray Cluster +miles's Megatron backend needs weights in `torch_dist` format. Run this on a GPU worker pod, which now exists after step 4. The conversion runs single-process by default; set `CONVERT_NUM_GPUS` or pass `--num-gpus` to parallelize it with `torchrun`: ```bash -# Substitute environment variables into the manifest -source env_vars -envsubst < kubernetes/raycluster.yaml | kubectl apply -f - - -# Watch pods come up (1 head + workers) -kubectl get pods -w -l ray.io/is-ray-node=yes - -# Port-forward the Ray dashboard (the recipes submit to 127.0.0.1:8265) -kubectl port-forward -n "${NAMESPACE}" svc/miles-ray-head-svc 8265:8265 & +W=$(kubectl get pod -n "${NAMESPACE}" -l ray.io/node-type=worker -o name | head -1); W=${W#pod/} +kubectl cp scripts/convert_checkpoint.sh "${NAMESPACE}/${W}:/tmp/convert_checkpoint.sh" +kubectl exec -n "${NAMESPACE}" "${W}" -- bash /tmp/convert_checkpoint.sh hf2megatron \ + --model-script qwen3-4B.sh --hf-path /fsx/models/Qwen3-4B --save-path /fsx/models/Qwen3-4B_torch_dist ``` -The shipped manifest schedules the Ray head onto the GPU pool (`${CPU_NODE_ROLE}` defaults to -`${GPU_NODE_ROLE}`, with a `nvidia.com/gpu` toleration; the head runs `num-gpus 0`) and GPU -workers on `node-role: ${GPU_NODE_ROLE}`, each worker declaring a `gpu_node` custom Ray resource -(see [miles-specific requirements](#miles-specific-requirements-found-on-real-hardware) for -why). The head must be on a CUDA-capable node -- a CPU-only head fails with `libcuda.so.1` -because miles control actors import Megatron even at `num-gpus 0`. Ensure the GPU node's root -volume has headroom for the ~18 GB image (150 GiB or more recommended). - ### 6. Configure the Reward -This sample scores rollouts in one of two ways -- pick one in `env_vars`: - -- **Built-in rule-based (default, verified):** `RM_TYPE="deepscaler"` (also `dapo`, `math`, - `f1`, `gpqa`). Runs in-process on the rollout actors; no extra setup. This is the reward - path used for every hardware-verified run in this test case. -- **Remote reward service on a CPU pool (`RM_TYPE="remote_rm"` + `RM_URL`):** offloads scoring - to a separate CPU instance group via `reward_service/` and - `kubernetes/reward-service.yaml`, mirroring the sibling slime test case. This path is - **UNVERIFIED on miles** -- the manifests are present but have not been deployed/exercised - here, and miles's `remote_rm` client lacks slime's retry-with-backoff behavior. See - `env_vars.disaggregated.example` for the overlay. - -No file copy is needed for the default path -- `deepscaler` is built into miles. +Pick a reward strategy in `env_vars`. The default is the built-in rule-based reward `RM_TYPE="deepscaler"`, with `dapo`, `math`, `f1`, and `gpqa` also available. It runs in-process on the rollout actors and needs no extra setup. A remote reward service on a CPU pool is available with `RM_TYPE="remote_rm"` and `RM_URL`, mirroring the sibling slime test case; for that path, deploy `kubernetes/reward-service.yaml` first. ### 7. Launch GRPO Training -The recipes end in `ray job submit --address http://127.0.0.1:8265`, so run them from a machine -that has the Ray CLI whose version matches the cluster (`pip install "ray=="`) -with the dashboard port-forwarded (step 5). If you cannot install a matching Ray CLI locally -- -e.g. no wheel exists for your local Python version -- or want to skip the port-forward, use the -`./run-on-cluster.sh` helper instead, which ships the recipe into the head pod and runs it there -(Ray is already present, `127.0.0.1:8265` is the head's own dashboard, and the version always -matches). `./run-on-cluster.sh --dry-run` prints what it will do; `./run-on-cluster.sh --recipe -run_grpo_qwen3_30b_a3b.sh` runs the MoE recipe. It only launches the recipe -- deploy the -RayCluster (step 5) first. +Run the recipe from a machine with a matching Ray CLI and the dashboard port-forwarded, as in step 4, or use `./run-on-cluster.sh` to launch from inside the head pod with only `kubectl`. ```bash -# Qwen3-4B, colocated (verified 1-node and 2-node paths): +# Qwen3-4B, colocated: bash recipe/run_grpo_qwen3_4b.sh - -# Qwen3-30B-A3B MoE, colocated on 2 nodes / 16 GPU (trains cleanly: reward 0.578, -# repetition 0.0 -- see Known Issues item 2 for the one rollout-geometry constraint): -# uncomment the ALTERNATE block in env_vars, then launch. The shipped block runs the rollout -# MoE pure expert-parallel (moe_tp=1); for moe_tp>1 & moe_ep>1 the recipe disables the FlashInfer fusion. +# Qwen3-30B-A3B MoE, colocated on 2 nodes: uncomment the ALTERNATE block in env_vars first. bash recipe/run_grpo_qwen3_30b_a3b.sh ``` -Monitor training: - -```bash -# Ray dashboard (after port-forward) -open http://localhost:8265 - -# Follow Ray job logs, using the submission id the recipe printed when it submitted -# ("raysubmit_..."). `ray job list` renders a table for humans; parse it at your own risk. -ray job logs --address http://localhost:8265 --follow - -# Monitor GPU utilization. This has to be a WORKER: the head runs with num-gpus 0 and -# NVIDIA_VISIBLE_DEVICES=void, so nvidia-smi there reports no devices. -W=$(kubectl get pod -n "${NAMESPACE}" -l ray.io/node-type=worker -o name | head -1) -kubectl exec -n "${NAMESPACE}" "${W#pod/}" -- nvidia-smi -``` - -Both recipes build the `train.py` argv as a bash array and submit it through -`recipe/launcher/grpo_launch.sh`, which sources the model script (`scripts/models/.sh`) -and expands `MODEL_ARGS` in the same shell -- avoiding a shell-escaping trap where an outer -`ray job submit -- bash -c "..."` string would expand the array before it is defined. The job -is submitted with `--entrypoint-resources '{"gpu_node": 0.001}'` so the Ray driver lands on a -GPU worker rather than the (non-GPU) head; see -[miles-specific requirements](#miles-specific-requirements-found-on-real-hardware) for why. +Monitor with the Ray dashboard at `http://localhost:8265`, or follow the job log with `ray job logs --address http://localhost:8265 --follow`. ### 8. Convert Checkpoints Back to HuggingFace Format -**Read this before running the training step, not after.** With the shipped values there is -nothing here to convert. `SAVE_INTERVAL=1000` sits deliberately above `NUM_ROLLOUT=100` so -that no save is ever triggered, because `save_model()` currently fails inside Megatron's -distributed-checkpoint save (Known Issues item 1) and a triggered save would end the run. The -consequence is that a run can finish `SUCCEEDED` after hours of training and leave **no -checkpoint at all** -- the trained weights are gone, and this step is where you would find -that out. The `iter_0060/` path below is therefore an illustration of the command's shape, -not a directory the default configuration produces. - -This runs on the same GPU-worker pod as step 4, for the same reason: the script and `/fsx` -only exist there, not on your workstation. - -Check first, and expect it to be empty on a default run: - -```bash -W=$(kubectl get pod -n "${NAMESPACE}" -l ray.io/node-type=worker -o name | head -1) -W=${W#pod/} -kubectl exec -n "${NAMESPACE}" "${W}" -- ls /fsx/runs/qwen3-4b/ckpt/qwen3-4b-grpo/ -``` +A run writes a Megatron `torch_dist` checkpoint to `CHECKPOINT_DIR` at the end of training. To evaluate the trained weights outside miles, convert them back to HuggingFace format on a GPU worker pod: ```bash +W=$(kubectl get pod -n "${NAMESPACE}" -l ray.io/node-type=worker -o name | head -1); W=${W#pod/} kubectl cp scripts/convert_checkpoint.sh "${NAMESPACE}/${W}:/tmp/convert_checkpoint.sh" -kubectl exec -n "${NAMESPACE}" "${W}" -- \ - bash /tmp/convert_checkpoint.sh megatron2hf \ - --input-dir /fsx/runs/qwen3-4b/ckpt/qwen3-4b-grpo/iter_0060/ \ - --output-dir /fsx/models/Qwen3-4B-GRPO-step60 \ - --origin-hf-dir /fsx/models/Qwen3-4B +# Replace iter_NNNN with the actual iteration: +# kubectl exec -n "${NAMESPACE}" "${W}" -- ls /fsx/runs/qwen3-4b/ckpt/qwen3-4b-grpo/ +kubectl exec -n "${NAMESPACE}" "${W}" -- bash /tmp/convert_checkpoint.sh megatron2hf \ + --input-dir /fsx/runs/qwen3-4b/ckpt/qwen3-4b-grpo/iter_NNNN/ \ + --output-dir /fsx/models/Qwen3-4B-GRPO --origin-hf-dir /fsx/models/Qwen3-4B ``` -Lowering `SAVE_INTERVAL` to produce a checkpoint is what hits the save bug, so treat this -sample as a measurement and instrumentation harness rather than a route to trained weights -until that issue is resolved -- see [Known Issues](#known-issues). - -## miles-specific requirements (found on real hardware) - -Three issues surface on the miles base image (`nvidia/cuda`) that do **not** occur on -SLIME's NGC base. All three are fixed in this test case's `miles.Dockerfile` and manifests; -they are called out here because carrying an EFA layer or launch pattern from a different -base image verbatim will silently reintroduce them. - -1. **CUDA compat shadowing kills `torch.cuda` (Error 803).** The miles base ships an older - CUDA forward-compat `libcuda` than the actual node driver. CUDA forward-compat requires - `compat >= host driver`; when the older bundled compat library is preferred via - `LD_LIBRARY_PATH`, `torch.cuda.is_available()` dies with `Error 803: unsupported display - driver / cuda driver combination`, and the SGLang engine fails at `get_device()`. - `miles.Dockerfile` removes `/usr/local/cuda*/compat` outright and lets the host driver - resolve instead. (SLIME's NGC base does not hit this because NGC's entrypoint only enables - compat when `compat >= host driver`.) - -2. **`libcuda.so.1` not found inside the SGLang subprocess.** With compat removed, `torch` - still resolves `libcuda` via `ld.so.cache`, but the SGLang server subprocess (and - Triton/cuda-python loaders, which scan `LD_LIBRARY_PATH` directly rather than the cache) - fail with `ImportError: libcuda.so.1: cannot open shared object file`. Fix: append the - driver-injection directories (`/usr/lib64`, `/usr/lib/x86_64-linux-gnu`) to - `LD_LIBRARY_PATH` and register them with `ldconfig` -- appended, not prepended, so they are - only consulted for libraries nothing else resolves. - -3. **The Ray job driver dies on the head (no libcuda).** miles's training actor imports - `mooncake` (P2P weight transfer, libcuda-dependent) at module load. The Ray job driver - runs wherever the job is submitted to land; on a non-GPU head pod, the import dies with - the same `libcuda.so.1` error. Fix: the GPU worker group declares a `gpu_node` custom Ray - resource (`rayStartParams.resources: '{"gpu_node": 1}'`), and the recipe submits with - `ray job submit --entrypoint-resources '{"gpu_node": 0.001}'`, landing the driver on a GPU - worker without consuming a full GPU slot from the colocated placement group. - -Full evidence and root-cause detail: [docs/PORT_NOTES.md](./docs/PORT_NOTES.md). +## miles-specific requirements -## Known Issues +Three items surface on the miles base image that do not occur on the sibling slime NGC base. All three are handled in `miles.Dockerfile` and the manifests; the inline comments at each fix carry the detail. -1. **`save_model()` fails with `_pickle.UnpicklingError: pickle data was truncated`.** This - occurs in Megatron's distributed checkpoint save (`gather_object`) after a training step - completes; the GRPO loop itself is unaffected, but any run that must checkpoint -- - including the megatron2hf conversion in step 8 above -- is currently blocked. Workaround - for short smoke runs: set `SAVE_INTERVAL` beyond the total step count so no save is - triggered. - -2. **The 30B MoE rollout degenerates only when SGLang runs the MoE tensor-parallel and - expert-parallel at the same time (`moe_tp>1` and `moe_ep>1`).** With the shipped rollout - geometry (pure expert-parallel, `moe_tp=1`) the 30B MoE trains cleanly -- `rollout/raw_reward` - 0.578, `rollout/repetition_frac` 0.0 -- so this is a rollout-configuration constraint, not a - model or a general "SGLang expert parallelism" problem. The recipe enforces it (below). - - SGLang derives the rollout MoE geometry as `moe_ep = --sglang-expert-parallel-size` (EP_SIZE) - and `moe_tp = --rollout-num-gpus-per-engine / EP_SIZE`. Serving the converted checkpoint - directly from SGLang -- no miles, no Megatron, no GRPO -- and sweeping the engine's - tensor-parallel size (TP) against EP isolates it. `repetition_frac` over 32 prompts, with - the resulting `moe_tp` = TP/EP annotated: - - | engine TP \\ EP | EP=1 (`moe_ep`=1) | EP=2 | EP=4 | EP=8 | - |---|---|---|---|---| - | TP=1 | 0.000 (`moe_tp`=1) | (cannot start) | - | - | - | TP=4 | 0.000 (`moe_tp`=4) | 0.875 (`moe_tp`=2) | - | - | - | TP=8 | 0.000 (`moe_tp`=8) | 0.594 (`moe_tp`=4) | 0.844 (`moe_tp`=2) | ~0.0 (`moe_tp`=1) | - - Read by `moe_tp`, the pattern is exact: every clean cell has `moe_tp=1` (the whole EP=8 - column, pure expert-parallel) or `moe_ep=1` (the whole EP=1 column, pure tensor-parallel); - every degenerate cell has both `moe_tp>1` and `moe_ep>1`. The earlier reading -- "EP>1 - degenerates" -- came from a sweep whose EP>1 cells all happened to have TP>EP, i.e. - `moe_tp>1`; the pure expert-parallel case (EP=TP, the EP=8 column) was not in it. A stock - `sglang.Engine` at `tp_size=8` confirms the missing column: `ep_size=8` (`moe_tp=1`) - generates coherent text (4-gram repetition 0.009), while `ep_size=4` and `ep_size=2` - (`moe_tp` 2 and 4) collapse to "7. 7. 7...", "1010...", ",,,,". Ruled out along the way: the - model's own recommended sampling (temperature 0.6 / top_p 0.95 / top_k 20) does not change a - degenerate cell, the `auto` vs `triton` MoE runner backend does not either, and all 18867 - expert weight keys are present in the checkpoint index with no NaN/Inf, so sampling, backend - selection and conversion are not involved. `--check-weight-update-equal` passes - (`weight_version` uniform, `mixed_version_ratio` 0.0), so trainer/rollout weights are not - diverging. - - Root cause (`0.5.16.dev` in the image): the FlashInfer allreduce+RMSNorm fusion. On SM90/SM100 - it is auto-enabled for Qwen3-MoE (`tp_size>1`, no dp-attention, `moe_a2a=none`) without regard - to `moe_ep`/`moe_tp`. With the fusion on, both post-experts all-reduces in - `models/qwen3_moe.py` `forward_normal` are skipped and deferred to the next layer's fused - `layernorm.forward_with_allreduce_fusion`. That fused reduce - (`layernorm.py::_forward_with_allreduce_fusion`, `flashinfer_comm_fusion.py`) selects its group - with `if moe_ep_size>1: use moe_ep_group else: use moe_tp_group` -- assuming the two are mutually - exclusive. When both are `>1` it reduces over the moe-ep group only and never reduces the moe-tp - group, so each rank keeps a partial sum over the intermediate dimension and generation collapses - from layer 0. Pure expert-parallel (`moe_ep=tp`) and pure tensor-parallel (`moe_tp=tp`) are - unaffected because `_MOE_EP`/`_MOE_TP` alias the full TP group, so either branch reduces over all - ranks. Setting `enforce_disable_flashinfer_allreduce_fusion` on the same combined config restores - the correct two-stage reduce and generation is clean (verified). The MoE weight sharding, the - dispatcher's `local_expert_mapping`, and the moe-ep/moe-tp process groups are all correct; the - defect is purely in the fused-reduce group selection. See - [docs/VERIFICATION_LOG.md](./docs/VERIFICATION_LOG.md) "30B MoE root cause". - - The recipe handles it automatically: `run_grpo_qwen3_30b_a3b.sh` computes `moe_tp` from - `ROLLOUT_GPUS_PER_ENGINE / EP_SIZE`, and when both `moe_tp>1` and `moe_ep>1` it adds - `--sglang-enforce-disable-flashinfer-allreduce-fusion` so that geometry trains correctly too; - pure expert-parallel (the shipped default) and pure tensor-parallel leave the fusion on. Losing - the fusion costs some rollout throughput but not correctness. The underlying SGLang bug is a - candidate for an upstream fix (generalise the fused-reduce group to the full TP group when - `moe_tp>1` and `moe_ep>1`). +- The base image's CUDA forward-compat `libcuda` is older than the node driver, so `miles.Dockerfile` removes `/usr/local/cuda*/compat` and uses the host driver. +- The SGLang subprocess resolves `libcuda.so.1` only after the driver-injection directories are appended to `LD_LIBRARY_PATH` and registered with `ldconfig`. +- miles's Ray control actors import Megatron at startup and link `libcuda.so.1`, so they must land on a CUDA-capable node. The manifest keeps every Ray node on the GPU pool and declares a `gpu_node` custom resource so the job driver lands on a GPU worker. -## File Structure - -``` -miles/ # 3.test_cases/pytorch/miles -├── README.md # This documentation -├── .gitignore -├── env_vars.colocated.example # Base config: Qwen3-4B colocated (+ 30B MoE ALTERNATE block) -├── env_vars.disaggregated.example # Overlay: reward model on a CPU pool + heavier GRPO -├── run-on-cluster.sh # Optional: run a recipe from the head pod (no local ray/port-forward) -├── miles.Dockerfile # radixark/miles base + EFA layer -├── requirements.txt # Reference Python deps (miles bundles these in the base image) -├── reward_service.Dockerfile # CPU-only image for the remote reward service -├── reward_service/ -│ ├── app.py # FastAPI reward server (reward_model / math_verify) -│ └── requirements.txt # Pinned CPU deps (no CUDA) -├── kubernetes/ -│ ├── buildkit-job.yaml # In-cluster image build -> registry (CPU-only) -│ ├── raycluster.yaml # KubeRay cluster manifest (head co-located on GPU pool, num-gpus 0) -│ ├── reward-service.yaml # CPU reward service Deployment + Service -│ └── data-prep-pod.yaml # Utility pod for data preparation -├── recipe/ -│ ├── run_grpo_qwen3_4b.sh # GRPO submit script (Qwen3-4B, colocated) -│ ├── run_grpo_qwen3_30b_a3b.sh # GRPO submit script (Qwen3-30B-A3B MoE, colocated 16 GPU) -│ └── 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 -└── docs/ - ├── EFA_2NODE.md # 2-node EFA/NCCL verification record - ├── PORT_NOTES.md # miles-specific porting notes and the 3 hardware pitfalls - └── VERIFICATION_LOG.md # runs, job ids, flags, metrics (source of Verification Status) -``` +miles adds the sm_103 Transformer Engine FA2 whitelist patch that the sibling slime image reached differently, because the miles base is not the NGC image. The residual patches slime carries for CUDA 13 and Blackwell are tracked in [awsome-distributed-ai issue #1163](https://github.com/awslabs/awsome-distributed-ai/issues/1163); miles applies the equivalent set through its own image build. -## Training Configuration Deep Dive - -### GRPO (Group Relative Policy Optimization) - -GRPO is a critic-free RL algorithm that estimates advantages by comparing rewards within a -group of responses generated for the same prompt, eliminating the need for a separate value -model. +## Known Issues -| Parameter | Description | Recipe Setting (4B) | -|-----------|-------------|----------------------| -| `--advantage-estimator grpo` | Use GRPO advantage estimation | GRPO | -| `--rollout-batch-size` | Prompts per rollout | 16 | -| `--n-samples-per-prompt` | Responses generated per prompt | 8 | -| `--global-batch-size` | Samples per optimizer step | 128 | -| `--num-steps-per-rollout` | Optimizer steps per rollout cycle | 1 | -| `--eps-clip` / `--eps-clip-high` | PPO-style clipping bounds | 0.2 / 0.28 | -| `--kl-loss-coef` | KL penalty coefficient | 0.0 | -| `--entropy-coef` | Entropy bonus coefficient | 0.0 | +1. **The 30B MoE rollout degenerates when SGLang runs the MoE tensor-parallel and expert-parallel at the same time**, that is when `moe_tp>1` and `moe_ep>1`. Pure expert-parallel with `moe_tp=1`, the shipped default, and pure tensor-parallel with `moe_ep=1` both train cleanly. This is an upstream SGLang bug in the FlashInfer allreduce+RMSNorm fusion, tracked and being fixed upstream in sgl-project/sglang PRs [#32963](https://github.com/sgl-project/sglang/pull/32963), [#32511](https://github.com/sgl-project/sglang/pull/32511), and [#32012](https://github.com/sgl-project/sglang/pull/32012). Disabling the fusion with `enforce_disable_flashinfer_allreduce_fusion` restores clean generation, with 4-gram repetition at 0.006 on par with pure EP, and the recipe applies this automatically for the combined geometry. The workaround stays correct after upstream ships the fix; once you move to a fixed build you can drop the flag to regain the fusion's throughput. -The constraint `rollout_batch_size * n_samples_per_prompt == global_batch_size * -num_steps_per_rollout` (16 * 8 = 128 * 1) must always hold; the recipes validate every -required variable is set before submitting. +2. **Qwen2.5-72B does not fit the 16-GPU H200 layout.** The disaggregated TP4 PP2 configuration runs out of memory on 2x p5en.48xlarge. It needs a larger cluster or optimizer and activation offload, neither of which has been run here. -### Parallelism Strategy +3. **The Ray head must run on a CUDA-capable node.** miles's control actors link `libcuda.so.1` at import even with `num-gpus 0`, so the manifest places the head on the GPU pool rather than a CPU node. Requiring a GPU-capable node for a zero-GPU control process is an upstream limitation in miles, not a constraint introduced by this test case. -**Qwen3-4B (colocated, verified on 1 and 2 nodes):** +## Reward Function -``` -Tensor Parallel (TP) = 1 -Pipeline Parallel (PP) = 1 -Context Parallel (CP) = 1 -Expert Parallel (EP) = 1 -Colocated: rollout-num-gpus == actor GPU count (8 on 1 node, 16 on 2 nodes) -``` +miles ships rule-based reward types `deepscaler`, `dapo`, `math`, `f1`, and `gpqa`, selected with `--rm-type`. The math types extract the `\boxed{...}` answer and grade it with `math_verify`. A remote reward service on a CPU pool is available with `RM_TYPE="remote_rm"`, mirroring the sibling slime test case. -**Qwen3-30B-A3B MoE (colocated on 2 nodes / 16 GPU, trains cleanly -- reward 0.578, -repetition 0.0; see [Known Issues](#known-issues) item 2 for the rollout-geometry constraint):** +## File Structure ``` -Tensor Parallel (TP) = 2 -Pipeline Parallel (PP) = 1 -Context Parallel (CP) = 1 -Expert Parallel (EP) = 2 -Expert Tensor Parallel = 1 ---use-distributed-optimizer # required: shards the 30B optimizer state across all - # 16 GPU so static memory fits H200 (141 GB); confining - # the actor to 8 GPU (one node) OOMs ---colocate # rollout time-shares the same 16 GPU ---rollout-num-gpus-per-engine 2 ---sglang-moe-runner-backend triton # MoE online weight update requires the triton - # runner; flashinfer is incompatible with the - # in-place weight update on SGLang 0.5.12+ ---sglang-expert-parallel-size 2 -``` - -### Dynamic Batching - -```bash ---use-dynamic-batch-size ---max-tokens-per-gpu 8192 # Per-GPU token budget per micro-batch +miles/ +├── README.md +├── env_vars.colocated.example # Qwen3-4B colocated, plus a 30B MoE ALTERNATE block +├── env_vars.disaggregated.example # overlay: reward model on a CPU pool +├── run-on-cluster.sh # launch a recipe from inside the head pod +├── miles.Dockerfile # radixark/miles base + AWS EFA layer +├── reward_service.Dockerfile +├── reward_service/ # FastAPI reward app, CPU +├── kubernetes/ +│ ├── raycluster.yaml # KubeRay cluster manifest +│ ├── buildkit-job.yaml # in-cluster image build +│ ├── data-prep-pod.yaml # model and data download +│ └── reward-service.yaml # CPU reward service +├── recipe/ # GRPO recipes and launcher +└── scripts/ # convert_checkpoint.sh, evaluate.sh ``` -Variable-length responses (math solutions can range from under 100 to 8,000+ tokens) make -dynamic batching materially more efficient than a fixed micro-batch size. - -## Reward Function - -1. **Built-in rule-based rewards (default, verified).** miles ships several rule-based - reward types -- `deepscaler`, `dapo`, `math`, `f1`, `gpqa` -- selected via `--rm-type`. The - math types extract the `\boxed{...}` answer and grade it against the label using - `math_verify` (LaTeX/sympy equivalence). The bundled DAPO-math recipe defaults to - `RM_TYPE="deepscaler"` and needs no extra setup. - -2. **Remote reward service on a CPU pool (UNVERIFIED on miles).** Set `RM_TYPE="remote_rm"` - and `RM_URL` to offload scoring to a separate CPU instance group -- useful for a heavier - reward such as a reward model, code execution, or RAG lookups. The service - (`reward_service/app.py`) exposes a `reward_model` backend (HuggingFace sequence - classifier on CPU) and a `math_verify` backend behind `POST /score`, plus `GET /health`. - This path mirrors the sibling slime test case's manifests but has not been deployed or - exercised on miles; miles's `remote_rm` client also lacks slime's retry-with-backoff, so - treat it as a starting point rather than a validated path. - ## Software Versions | Component | Version | |-----------|---------| -| miles | `radixark/miles` (fork point `fcce96ca0`, 2025-10-05) | -| Base image | `radixark/miles:dev-202607310056` @ `sha256:ca0bb593dd6f4011b444f64d478b72c213e4c70421f4d7f94e593a709562429e` | -| SGLang | 0.5.16.dev | -| Megatron-LM | radixark fork (miles-compatible) | +| miles | `radixark/miles` at commit `fc04f66` | +| Base image | `radixark/miles`, pinned by `sha256` digest in `miles.Dockerfile` | +| SGLang | `sglang-miles` branch, based on `v0.5.16` | +| Megatron-LM | radixark fork, `miles-main` | | Ray | 2.55.1 | | CUDA | 13.0.1 | -| PyTorch | 2.11 | +| PyTorch | 2.11.0 | | EFA installer | 1.48.0 | | GDRCopy | v2.5.2 | -Pin the image tag to a dated/versioned value (e.g. `miles-cuda13-efa-0.1`); never use -`:latest`, per the [awsome-distributed-ai CONTRIBUTING guidelines](https://github.com/awslabs/awsome-distributed-ai/blob/main/CONTRIBUTING.md). - ## Troubleshooting -**Pod stuck in `Pending` state** +**Pod stuck in `Pending`** ```bash -kubectl describe pod -# Check for resource constraints -- GPU/EFA/memory/ephemeral-storage requests may -# exceed node capacity. The Ray head in particular needs enough ephemeral-storage -# to pull the ~18 GB image; a too-small root volume causes an Evict mid-pull rather -# than a clean Pending. +kubectl describe pod +# Check GPU/EFA/memory/ephemeral-storage requests against node capacity. +# The head needs enough ephemeral-storage to pull the ~18 GB image. ``` -**Ray workers fail to connect to head node** +**Ray workers cannot connect to the head** ```bash -kubectl get svc miles-ray-head-svc -kubectl exec -- nslookup miles-ray-head-svc -# Ensure RAY_memory_monitor_refresh_ms=0 is set (prevents OOM kills during init) +kubectl get svc miles-ray-head-svc -n "${NAMESPACE}" +kubectl exec -n "${NAMESPACE}" -- nslookup miles-ray-head-svc ``` +Separately, if workers are killed mid-run rather than failing to connect, ensure `RAY_memory_monitor_refresh_ms=0` is set to disable the memory monitor. -**NCCL/EFA initialization errors** +**NCCL/EFA errors** ```bash kubectl exec -- fi_info -p efa kubectl exec -- env | grep NCCL -# Ensure FI_PROVIDER=efa and FI_EFA_USE_DEVICE_RDMA=1 are set -# If NCCL selects efa but every transfer times out (Error 15, "Unreachable remote"), -# check that the EFA security group has self-referencing all-traffic on BOTH -# ingress AND egress -- EFA's OS-bypass SRD traffic is not ordinary IP traffic, so -# a CIDR-only egress rule does not authorize it. See docs/EFA_2NODE.md. +# Ensure FI_PROVIDER=efa. For multi-node, the EFA nodes' security group must +# allow all traffic to itself on both ingress and egress, since EFA SRD is not IP. +``` + +**`ImportError: libcuda.so.1`** — the Ray job driver or a control actor landed on a node without the driver; see miles-specific requirements. + +**`torch.cuda.is_available()` is `False` or `Error 803`** +```bash +kubectl exec -- ls /usr/local/cuda*/compat # expect: not found +# The base image's CUDA forward-compat library was older than the node driver; +# miles.Dockerfile removes it. ``` -**`ImportError: libcuda.so.1: cannot open shared object file`** -- Confirm the Ray job driver landed on a GPU worker (`--entrypoint-resources - '{"gpu_node": 0.001}'`) and not the head. -- Confirm `LD_LIBRARY_PATH` includes the driver-injection directories - (`/usr/lib64`, `/usr/lib/x86_64-linux-gnu`) appended at the end. -- See [miles-specific requirements](#miles-specific-requirements-found-on-real-hardware). - -**`torch.cuda.is_available()` is `False` / `Error 803`** -- The image's CUDA forward-compat library is older than the node driver. Confirm - `/usr/local/cuda*/compat` has been removed (as `miles.Dockerfile` does) rather than - merely dropped from `LD_LIBRARY_PATH`. - -**SGLang fails to start (CUDA OOM) in colocated mode** -- In colocated mode, SGLang launches after Megatron occupies GPU memory. Reduce - `--sglang-mem-fraction-static` (0.8 for 4B; 0.75 was the value that completed for the 30B - MoE colocated run -- the original 0.85 hardcode left too little room). - -**MoE online weight update fails on SGLang 0.5.12+** -- Ensure `--sglang-moe-runner-backend triton` and `--sglang-expert-parallel-size` are both - set explicitly; the default flashinfer MoE runner is incompatible with the in-place - weight update used by GRPO training. - -**Weight conversion fails** -- Ensure `PYTHONPATH` includes the Megatron-LM directory (`/root/Megatron-LM`). -- Verify model config parameters match (`--rotary-base`, `--vocab-size`, etc.). -- For MoE models, ensure `--expert-model-parallel-size` is set correctly. - -**Checkpoint save fails with `_pickle.UnpicklingError: pickle data was truncated`** -- Known upstream issue in Megatron's distributed checkpoint save path; see - [Known Issues](#known-issues). Raise `SAVE_INTERVAL` beyond the run length to avoid - triggering a save on short smoke runs. +**SGLang out-of-memory in colocated mode** — lower `--sglang-mem-fraction-static` to 0.8 for 4B or 0.75 for the 30B MoE colocated run. ## References -- [miles GitHub Repository](https://github.com/radixark/miles) -- [SLIME GitHub Repository](https://github.com/THUDM/slime) -- [SLIME Blog: An SGLang-Native Post-Training Framework for RL Scaling](https://lmsys.org/blog/2025-07-09-slime/) -- [SGLang Project](https://github.com/sgl-project/sglang) +- [miles](https://github.com/radixark/miles) +- [SLIME](https://github.com/THUDM/slime) +- [SGLang](https://github.com/sgl-project/sglang) - [Megatron-LM](https://github.com/NVIDIA/Megatron-LM) -- [GRPO Paper (DeepSeek-R1)](https://arxiv.org/abs/2402.03300) -- [Amazon SageMaker HyperPod Documentation](https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-hyperpod.html) +- [GRPO paper](https://arxiv.org/abs/2402.03300) +- [Amazon SageMaker HyperPod documentation](https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-hyperpod.html) +- [KubeRay documentation](https://docs.ray.io/en/latest/cluster/kubernetes/index.html) - [Sibling test case: 3.test_cases/pytorch/slime/](../slime/) - [awsome-distributed-ai](https://github.com/awslabs/awsome-distributed-ai) -- [KubeRay Documentation](https://docs.ray.io/en/latest/cluster/kubernetes/index.html) ## Security -See [CONTRIBUTING](https://github.com/awslabs/awsome-distributed-ai/blob/main/CONTRIBUTING.md) for more information. +See [CONTRIBUTING](https://github.com/awslabs/awsome-distributed-ai/blob/main/CONTRIBUTING.md). ## License diff --git a/3.test_cases/pytorch/miles/docs/EFA_2NODE.md b/3.test_cases/pytorch/miles/docs/EFA_2NODE.md deleted file mode 100644 index aa7e0672e..000000000 --- a/3.test_cases/pytorch/miles/docs/EFA_2NODE.md +++ /dev/null @@ -1,82 +0,0 @@ -# 2-node EFA verification (p5en.48xlarge x2) - -Hardware verification that multi-node NCCL runs over EFA on this cluster, using two -Capacity Block p5en.48xlarge nodes (8x H200 each, us-east-2a, same UltraCluster). No -placement group is used or needed: a Capacity Block already colocates its nodes within one -UltraCluster spine, and `PlacementGroupArn` on the reservation is null, so layering a -self-made cluster placement group would only risk "capacity reserved but PG-unsatisfiable". - -## Result: EFA works across 2 nodes - -`torchrun --nnodes 2 --nproc_per_node 8` running a NCCL all_reduce (16 GPUs), with -`NCCL_DEBUG=INFO NCCL_DEBUG_SUBSYS=INIT,NET`: - -- Both nodes: `NET/OFI Selected provider is efa, fabric is efa-direct (found 8 nics)`. -- Inter-node channels: `Channel .. : 3[3] -> 14[6] [send] via NET/Libfabric/3/GDRDMA`, - `GPU Direct RDMA Enabled`. Zero `NET/Socket` lines (no TCP fallback). -- Bandwidth (algbw / busbw, all_reduce, 16 GPUs): - - | message | algbw | busbw | - | --- | --- | --- | - | 64 MB | 101.5 | 190.3 GB/s | - | 256 MB | 100.5 | 188.5 GB/s | - | 1024 MB | 127.2 | 238.4 GB/s | - | 4096 MB | 135.8 | 254.7 GB/s | - | 8192 MB | 137.1 | 257.1 GB/s | - -busbw 190-257 GB/s is far above the ~10 GB/s a TCP path would give, confirming EFA. - -**These numbers were taken with a partial EFA request, which is no longer what the manifest -does -- and that turned out to matter for correctness, not just throughput.** The measurement -above ran with the pods requesting `vpc.amazonaws.com/efa: 8` of the node's allocatable 15, -roughly half the fabric, and an earlier version of this document described requesting them all -as "a throughput improvement, not a correctness issue". That was wrong. - -Claiming fewer cards than the node exposes lets the device plugin choose which ones, and the -two pods can end up with different EFA-device-to-NIC-rail mappings. `aws-ofi-nccl` then aborts -with `NET/OFI Unexpected number of remote rails for dev N. Expected 1 but got 2` followed by -`ncclInternalError`. Crucially it does **not** appear in the all_reduce test above: TP16 -all-reduce across the node boundary passed with 8 cards, while TP8 x PP2 on the same cluster -failed within minutes, because the pipeline stage boundary uses point-to-point send/recv. -MoE expert-parallel all-to-all is P2P-like and is affected the same way. - -So `EFA_PER_NODE` now defaults to the node's full allocatable count (15 on p5en.48xlarge, -32 on p5.48xlarge) and `kubernetes/raycluster.yaml` takes it from there. Read the count off the -node rather than the instance spec: - -```bash -kubectl get node -o jsonpath='{.status.allocatable.vpc\.amazonaws\.com/efa}' -``` - -The throughput point still stands on its own -- the full request is also expected to exceed -~300 GB/s -- but an EFA validation that only exercises all-reduce will report success on a -configuration that breaks later inside training. Include a P2P path. - -## Root cause fixed: EFA security-group needs self-referencing egress - -The first 2-node run failed even though NCCL selected the efa provider and established -GPUDirect RDMA channels: every data transfer timed out with -`NET/OFI ... Error 15 (Unreachable remote (never received a response))`. - -Cause: the EFA security group had self-referencing all-traffic on **ingress** but only -`0.0.0.0/0` on **egress**. EFA's OS-bypass SRD traffic is not IP traffic, so a CIDR egress -rule does not authorize it -- EFA SGs need self-referencing all-traffic on BOTH directions. -Bootstrap (TCP over eth0) and NCCL init succeeded under the asymmetric rule, which is why -the failure only appeared at the first SRD data transfer. - -Fix: add a self-referencing security-group rule allowing all protocols on both ingress and -egress (e.g. `aws_security_group_rule` with `source_security_group_id` == the SG's own ID, -protocol `-1`, for both directions) to the EFA-capable nodes' security group. Single-node EFA -was unaffected because intra-node traffic does not traverse the SG. - -## How to reproduce the check - -1. Two EFA nodes in the same subnet with the EFA SG (self-ref ingress AND egress). -2. sshd not required -- use torchrun with `--master_addr `; NCCL data path is - EFA regardless of the launcher. -3. On each pod: `FI_PROVIDER=efa NCCL_DEBUG=INFO NCCL_DEBUG_SUBSYS=INIT,NET - NCCL_SOCKET_IFNAME=eth0 torchrun --nnodes 2 --node_rank <0|1> --nproc_per_node 8 - --master_addr --master_port 29500 bench.py`. -4. Pass: `Selected provider is efa` on both nodes, inter-node `[send] via NET/Libfabric`, - no `NET/Socket`, busbw > 100 GB/s. Fail (TCP): busbw single-digit GB/s, or Error 15 - (check the SG self-ref egress rule first). diff --git a/3.test_cases/pytorch/miles/docs/PORT_NOTES.md b/3.test_cases/pytorch/miles/docs/PORT_NOTES.md deleted file mode 100644 index d892a985d..000000000 --- a/3.test_cases/pytorch/miles/docs/PORT_NOTES.md +++ /dev/null @@ -1,115 +0,0 @@ -# Porting notes (slime -> miles) - -Record of porting the sibling [slime test case](../../slime/) to -[miles](https://github.com/radixark/miles), slime's direct fork. miles diverged from slime -at fork point `fcce96ca0` (2025-10-05) and rewrote the train loop sync -> async, but the -`train.py` CLI is compatible, so the same GRPO recipe runs with only path/branding changes -(verified on hardware). - -## Why miles - -- miles targets CUDA 13.0.1 / PyTorch 2.11 / Blackwell (sm_103) as first-class. Where the - slime image reached sm_103 via an NGC base plus hand patches, the miles image already - applies the sm_103 Transformer Engine FA2 whitelist patch. -- miles has features slime lacks (fp8 rollout, fully-async, true-on-policy) -- out of scope - here, but future material. - -## Docker strategy C: miles official image + EFA layer - -miles's dependencies (the sglang-miles fork, radixark/Megatron-LM fork, prebuilt -flash_attn/TE/apex wheels) are built for the PyTorch 2.11 stable + cu130 ABI and do not -load on an NGC (nightly ABI) base. So the image takes `radixark/miles:` as the -base and adds only the AWS EFA stack (`miles.Dockerfile`): - -- remove IB libverbs (let the EFA installer's libfabric win) -- gdrcopy (GPUDirect RDMA) -- EFA installer 1.48.0 (`--skip-kmod`, libfabric + aws-ofi-nccl plugin) -- NCCL/EFA runtime env (`FI_PROVIDER=efa` etc.), `GLOO_SOCKET_IFNAME=eth0` -- `PYTHONPATH=/root/miles:/root/Megatron-LM` (miles does not bake this into the image) - -Pin by digest (`MILES_BASE_DIGEST`), never `:latest` and not by tag alone: radixark -publishes `dev-*` as mutable snapshots and has already deleted a tag this test case once -pinned. The dated tag is kept next to the digest for readability. Per awsome-distributed-ai -CONTRIBUTING. - -## slime -> miles changes (recipe / launcher) - -| Item | slime | miles | -| --- | --- | --- | -| framework install dir | `/opt/slime` | `/root/miles` (editable install) | -| Megatron | `/opt/Megatron-LM` | `/root/Megatron-LM` (radixark fork) | -| recipe runtime-env PYTHONPATH | `/opt/Megatron-LM` | `/root/Megatron-LM:/root/miles` | -| launcher `SLIME_DIR` default | `/opt/slime` | `/root/miles` (variable name kept) | -| model script | `scripts/models/qwen3-4B.sh` | same (path compatible) | - -The `train.py` flags are unchanged, so the GRPO recipe argv is identical to slime's apart -from the paths above. SGLang passthrough flags (`--sglang-*`) work the same way on miles -(ServerArgs auto-expose via parse_known_args). - -## Pitfalls found on real hardware - -All three occur on the miles base (nvidia/cuda) and NOT on slime's NGC base, and are fixed -in `miles.Dockerfile` / the manifests. - -### 1. CUDA compat shadowing -> torch.cuda dies (Error 803) - -Carrying slime.Dockerfile's EFA layer verbatim also carries -`ENV LD_LIBRARY_PATH=...:/usr/local/cuda/compat:...`, which is fatal on the miles base: - -- The miles base bundles a CUDA forward-compat libcuda `580.82.07`. -- The validation node's host driver is `580.159.03` (nvidia-smi shows "CUDA Version: 13.0"). -- CUDA forward-compat requires compat >= host driver. When the older compat (580.82.07) is - preferred via LD_LIBRARY_PATH, torch.cuda dies with `RuntimeError: ... Error 803: system - has unsupported display driver / cuda driver combination`, and the SGLang engine fails at - `get_device()` ("No accelerator available"). -- Verified: dropping compat from LD_LIBRARY_PATH restores `torch.cuda.is_available() == True`. -- Fix: `miles.Dockerfile` deletes `/usr/local/cuda*/compat` outright and uses the host - driver (`/usr/lib64/libcuda.so`), which already supports the image's CUDA 13.0 toolkit. - slime (NGC) could use compat because NGC's entrypoint enables it only when compat >= host. - -### 2. libcuda.so.1 not found in the SGLang subprocess - -With compat gone, torch still works (ld.so.cache resolves libcuda), but the SGLang server -subprocess (and Triton / cuda-python loaders, which scan LD_LIBRARY_PATH directly rather -than the cache) fail with `ImportError: libcuda.so.1: cannot open shared object file`. Fix: -append the driver-injection dirs (`/usr/lib64`, `/usr/lib/x86_64-linux-gnu`) to -LD_LIBRARY_PATH and register them with `ldconfig`. Appending (not prepending) means they -are consulted only for libraries nothing else resolves. - -### 3. Ray job driver dies on the head (no libcuda) - -miles's `MegatronTrainRayActor` imports `mooncake` (P2P weight transfer, libcuda-dependent) -at module load. The Ray job driver runs on the head, which is a non-GPU pod with no libcuda -injected, so the driver dies with the same `ImportError: libcuda.so.1`. Fix: declare a -`gpu_node` custom resource on the worker (`rayStartParams.resources: '{"gpu_node": 1}'`) -and submit with `ray job submit --entrypoint-resources '{"gpu_node": 0.001}'`, which places -the driver on a GPU worker without consuming a GPU logical count (so it does not conflict -with the colocated 8-GPU placement group). Note: `begin_weight_update` / `pull_weights` are -Ray actor methods on miles's rollout engine (`sglang_engine.py`), not HTTP endpoints on an -SGLang fork. - -## GPU-generation portability (H200 / B300) - -Everything in `miles.Dockerfile`, the recipes, and the manifests is GPU-generation-agnostic: -the base image tag (`MILES_BASE_TAG`) already carries both the sm_90 (Hopper/H200) and -sm_103 (Blackwell/B300) builds internally, and no recipe/env file hard-codes a CUDA or SM -version. The one cluster-specific knob is `GPU_NODE_ROLE` in `env_vars` (substituted into -`kubernetes/raycluster.yaml`'s worker `nodeSelector`), which just needs to match whichever -accelerator NodePool -- H200 or B300 -- you point it at. This test case is hardware-verified -only on H200 (p5en.48xlarge); the analysis above is why B300 (p6-b300.48xlarge) is expected -to work unmodified, not a claim that it has been run. - -## Known Issue - -- `save_model()` fails with `_pickle.UnpicklingError: pickle data was truncated` in - Megatron's distributed checkpoint save (`gather_object`), independent of the GRPO loop. - Blocks the HF<->Megatron round-trip and long checkpointing runs. File upstream on miles. - -## Residual patches (of slime's 7, what remains on miles) - -- `--sglang-log-level warning` (lowercase; avoids the uvicorn KeyError; recipe-only). -- GPU-less Ray driver Megatron `validate_args` CUDA probe (may surface for 30B MoE; may - already be fixed in the radixark fork -- UNVERIFIED). - -The numpy<2 pin, torch_memory_saver preload `.so` selection, and manual mbridge pin are all -resolved by the miles official image. diff --git a/3.test_cases/pytorch/miles/docs/VERIFICATION_LOG.md b/3.test_cases/pytorch/miles/docs/VERIFICATION_LOG.md deleted file mode 100644 index 089d61ff6..000000000 --- a/3.test_cases/pytorch/miles/docs/VERIFICATION_LOG.md +++ /dev/null @@ -1,322 +0,0 @@ -# Verification log - -What was run, on what, and what came back. Every metric below was read out of the trainer's -own TensorBoard event files rather than retyped from a terminal, and every run named here has -an event file behind it. - -This file exists because an earlier version of the README's verification table contained a -row that was wrong, and a log is what would have caught it. See "A correction" at the end. - -## Environment - -| | | -|---|---| -| Cluster | Amazon EKS, KubeRay, 2 nodes | -| Instance | `p5en.48xlarge` (H200 141GB x8 per node) | -| Interconnect | EFA between nodes, GPUDirect RDMA | -| Shared storage | FSx for Lustre, PERSISTENT_2, mounted at `/fsx` | -| Base image | `radixark/miles@sha256:ca0bb593dd6f4011b444f64d478b72c213e4c70421f4d7f94e593a709562429e` (tag `dev-202607310056`, cu13) plus the AWS EFA layer built by `miles.Dockerfile` | -| Model | Qwen3-4B, converted to Megatron format with `scripts/convert_checkpoint.sh hf2megatron` | -| Data | DAPO-Math-17k for training prompts, AIME-2024 for eval, `deepscaler` reward | - -## Run: the shipped 4B recipe, unmodified - -This is the reader's path: the shipped `env_vars.colocated.example` with only the -cluster-specific values filled in (model and data paths, checkpoint and TensorBoard -directories, namespace, FSx claim) and `NUM_ROLLOUT=1` to keep it short. No changes to the -recipe itself. - -``` -ENV_FILE=env_upstream bash recipe/run_grpo_qwen3_4b.sh -``` - -Submitted as Ray job `raysubmit_GdnWJy6KJsJDmsF2`, terminal status `SUCCEEDED`. - -The flags the job actually received, from the job log rather than from the recipe source: - -``` -bash grpo_launch.sh \ - --hf-checkpoint /fsx/models/Qwen3-4B \ - --ref-load /fsx/models/Qwen3-4B_torch_dist \ - --load /fsx/runs/upstream_verify/ckpt/qwen3-4b-grpo/ \ - --save /fsx/runs/upstream_verify/ckpt/qwen3-4b-grpo/ \ - --save-interval 1000 \ - --prompt-data /fsx/data/dapo-math-17k/dapo-math-17k.jsonl \ - --input-key prompt --label-key label ... -``` - -Results, read back from `/fsx/tb/upstream_verify` (78 scalar tags recorded): - -| metric | value | reading | -|---|---|---| -| `rollout/raw_reward` | 0.5156 | the model answers correctly about half the time, which is the expected range for 4B on this data | -| `rollout/repetition_frac` | 0.0 | no degenerate generation | -| `rollout/truncated_ratio` | 0.4844 | about half the responses hit the 8192-token cap, normal for this prompt set | -| `train/grad_norm` | 0.6479 | finite and unremarkable | -| `perf/step_time` | 273.3 s | one full rollout-plus-train cycle | - -Logs kept in full: 2,242 lines of job log and 2,276 lines of submission log. - -## Why reward and repetition are the metrics to check, not the exit code - -`rollout/repetition_frac` and `rollout/raw_reward` are what separate a run that works from a -run that merely finishes. A misconfigured 30B MoE run -- one whose SGLang rollout runs the -MoE tensor-parallel and expert-parallel at once (`moe_tp>1` and `moe_ep>1`) -- exits 0, prints -an unremarkable loss, and produces nothing usable: repetition 0.48 to 0.70, reward pinned at -0.0. If you verify by exit code you will record it as working -- which is exactly what happened -once, before the metric was checked. The same 30B model with the shipped rollout geometry -(`moe_tp=1`, pure expert-parallel) trains cleanly: reward 0.578, repetition 0.0. The metric, -not the exit code, is what tells the two apart -- and what pointed to the rollout geometry as -the cause. See "30B MoE root cause" below. - -For comparison, on the same cluster and recipe: - -| | dense 4B | 30B MoE, `moe_tp=1` (shipped) | 30B MoE, `moe_tp>1` (misconfigured) | -|---|---|---|---| -| `rollout/repetition_frac` | 0.0 | 0.0 | 0.48 to 0.70 | -| `rollout/raw_reward` | 0.516 | 0.578 | 0.0 | -| `rollout/truncated_ratio` | 0.484 | 0.42 | 0.97 to 0.99 | -| exit status | SUCCEEDED | SUCCEEDED | SUCCEEDED | - -## Two-node EFA - -The 2-node rows in the README's table come from earlier runs on this same cluster with the -worker replica count at 2. NCCL all_reduce busbw measured 190 to 257 GB/s with `efa-direct` -and GPUDirect RDMA and no TCP fallback; `docs/EFA_2NODE.md` has the security-group -prerequisite and the failure signature when it is missing. - -`docs/EFA_2NODE.md` is a fabric measurement, though, not a training one, so on its own it does -not support the "3 rollout cycles" part of that row. The GRPO run behind it, Qwen3-4B dense -colocated across 2 nodes (actor 2x8, rollout sharing the same 16 GPU): - -| rollout | `rollout/raw_reward` | `actor_train_tflops` | `perf/step_time` | -|---|---|---|---| -| 0 | 0.477 | 100.4 | 119.2s | -| 1 | 0.523 | 245.4 | 73.1s | -| 2 | 0.492 | 235.7 | 81.6s | - -All three cycles completed, job `SUCCEEDED` in 676s. Weight sync, rollout, reference log-probs -and the Megatron backward all crossed the node boundary over EFA. `ppo_kl` stayed 0.0 at -dropout 0, matching the single-node arm. Reward moving in the 0.48-0.52 band over three steps -is not a training result -- three optimizer steps show nothing about convergence -- it is -evidence that the loop closes end to end on two nodes, which is what that row claims. - -One detail worth repeating from that document, because it cost real time: requesting fewer -EFA devices than the instance exposes lets the device plugin pick different cards on each -node, and NCCL then fails with `NET/OFI Unexpected number of remote rails`. It only shows up -on point-to-point traffic (pipeline or expert parallel), so an all-reduce smoke test passes -and the problem surfaces later. - -## What this log does not cover - -- The disaggregated reward service. Present in the repository, never deployed. -- Any B300 hardware. -- Checkpoint save-back. `save_model()` fails with a pickle-truncation error inside Megatron's - distributed checkpoint save; `SAVE_INTERVAL` ships above the step count so a run does not - trigger it. Known Issues item 1. - -## Ray head must run on a GPU (CUDA-capable) node, not a CPU-only node - -An earlier version of this test case placed the Ray head on a CPU-only node (to keep it off -the expensive GPU pool). That shape is not safe for miles and the shipped manifest no longer -uses it. The reason is in the framework, not the recipe: miles's Ray control actors are -created with `num_gpus=0` (e.g. `create_rollout_manager` in `miles/ray/placement_group.py` -does `RolloutManager.options(num_cpus=1, num_gpus=0)`), but their module import pulls in -Megatron / `transformer_engine`, whose shared library `dlopen`s `libcuda.so.1` at import time. -Ray is then free to place such a zero-GPU actor on any node with a spare CPU -- including a -CPU-only head node -- where the import hard-fails: - -``` -OSError: libcuda.so.1: cannot open shared object file: No such file or directory - ... File ".../transformer_engine/common/__init__.py", ... _load_core_library() - ray.exceptions.ActorDiedError: RolloutManager.__init__() ... (TemporaryActor, ip=) -``` - -So any CPU-only node in the miles Ray cluster is a latent hazard: whether a run survives -depends on whether Ray happens to place the Megatron-importing actor on a GPU worker instead. -The earlier "head-on-CPU SUCCEEDED (reward 0.531)" run was that placement luck, not a -guarantee. The fix is to keep every node in the Ray cluster CUDA-capable: the shipped -`raycluster.yaml` now schedules the head onto the GPU pool (`CPU_NODE_ROLE` defaults to -`GPU_NODE_ROLE`) with a `nvidia.com/gpu` toleration; the head still runs `num-gpus 0` and -consumes no GPU, so on a colocated run it simply co-locates on a worker's GPU node at no extra -cost, and the ~18 GB image is already cached there. This also removes the head from the CPU -Karpenter pool, sidestepping the "underutilized" consolidation churn that pool's 30s policy -caused. The verified metrics for the colocated dense 4B run (reward 0.53, repetition 0.0, -weight_version uniform) are unchanged; only the head's node placement changed. - -Note (upstream): the sibling `slime` test case ships the same head-on-CPU shape and the same -`num_gpus=0` control-actor pattern, so it shares this latent hazard; worth raising upstream. - -## Second review round: defects found by reading, and what was re-run to confirm - -A later review pass went through the reader's path again rather than re-running the same job, -and found six defects that no amount of re-running would have surfaced, because the affected -paths either are not on the default path or fail silently. They are listed here because the -fixes changed shipped files, and a log that only records successes is not much of a log. - -| what was wrong | why it was silent | -|---|---| -| `kubernetes/reward-service.yaml` referenced `${REWARD_IMAGE}`, which no env file defined | `envsubst` renders an undefined variable as the empty string, so the Deployment went out with `image: ""` | -| `raycluster.yaml` hardcoded `replicas: 1` with a comment telling the reader to edit it for 2-node runs | a 2-node config (the 30B MoE block sets `ACTOR_NUM_NODES=2`) then started 8 GPU worth of workers and the job waited on placement rather than erroring. Now driven by `${ACTOR_NUM_NODES}` | -| `raycluster.yaml` requested `vpc.amazonaws.com/efa: 8` of the node's 15 allocatable | a partial request breaks point-to-point NCCL (see the rails note above) but passes an all-reduce smoke test. Now `${EFA_PER_NODE}` | -| `scripts/evaluate.sh` keyed pass@k on `prompt_item.get("idx", 0)` | the prepared AIME-2024 file has only `prompt` and `label`, so all 30 prompts collapsed onto key 0 and pass@k became "any of the 480 samples was correct". Reproduced: for a model correct on 1 prompt of 30, the old code reports 1.0000 where the truth is 0.0333 | -| `scripts/evaluate.sh` extracted answers with `\\boxed\{([^}]*)\}` | the character class stops at the first brace, so `\boxed{\frac{1}{2}}` yielded `\frac{1` and was scored wrong. Now a brace-counting scan | -| `scripts/evaluate.sh` submitted every request at once under a 300s timeout | with `MAX_TOKENS=16384` most requests time out in the queue, and the handler counts a timeout as an incorrect answer, so accuracy sags for a reason unrelated to the model. Now bounded concurrency, a generation-sized timeout, and a non-zero exit when the error rate exceeds 5% | - -The README's hardware table also said "EFA per node: 16 devices" while EKS reports 15 -allocatable on `p5en.48xlarge`; `docs/EFA_2NODE.md` still described a partial EFA request as a -throughput trade-off rather than a correctness problem; and step 4's prose offered "exec into -its head" for the checkpoint conversion two lines above a code block that correctly targets a -worker. All three are corrected. - -What was re-verified on hardware after these changes, rather than assumed: - -- All four manifests were rendered through `envsubst` from the shipped example env and - validated with `kubectl apply --dry-run=server` against a live EKS cluster. All four are - accepted, including `reward-service.yaml`, which previously could not be. -- `raycluster.yaml` was rendered at both `ACTOR_NUM_NODES=1` and `=2` and validated at each; - `replicas`, `minReplicas` and `maxReplicas` track the value, and `vpc.amazonaws.com/efa` - resolves to 15, matching what the node reports as allocatable. -- The `\boxed{}` extractor and the pass@k indexing were checked against the actual - `aime-2024.jsonl` on the cluster, which is where the missing `idx` field was confirmed. -- `bash -n` passes on every script, and the Python inside the `evaluate.sh` heredoc compiles. - -## slime-parity comparison (model x topology coverage) - -The sibling `3.test_cases/pytorch/slime` README lists a "Supported Model Sizes" matrix -(Qwen3-4B colocated, GLM-Z1-9B colocated, Qwen3-30B-A3B disaggregated, Qwen2.5-72B -disaggregated) but publishes no measured results for it. The runs below cover that matrix on -miles, `NUM_ROLLOUT=2`, reading metrics from the trainer's event files. Rows marked -(shipped recipe) use the recipes in this test case unmodified; rows marked (campaign config) -used an authored model script or a `--colocate`-conditional recipe variant and are reported as -findings, not as shipped support. - -Hardware shorthand used in the HW column: - -- **P5ENx1** = 1x `p5en.48xlarge` (8x H200, 141 GiB each). -- **P5ENx2** = 2x `p5en.48xlarge` (16x H200). -- **P6B300x2** = 2x `p6-b300.48xlarge` (16x B300, 288 GiB each) -- the expected-compatible target - for cases that do not fit H200; not run here. - -The HW column names the configuration a row was actually run on, so a "does not fit" is scoped to -that hardware rather than read as a property of the model. - -| Model | Layout | HW | Result | reward | repetition | notes | -|---|---|---|---|---|---|---| -| Qwen3-4B dense | colocated (shipped recipe) | P5ENx1 | SUCCEEDED | 0.53 | 0.0 | head co-located on GPU pool, above | -| Qwen3-4B dense | disaggregated `COLOCATE=false` (shipped recipe) | P5ENx2 | SUCCEEDED | 0.52 | 0.0 | weight sync over NCCL/EFA (`weight_version` 2, mixed 0.0); worker `replicas` must cover actor+rollout GPUs | -| GLM-Z1-9B dense | colocated TP2 (campaign config) | P5ENx1 | SUCCEEDED | 0.68 | 0.0 | TP>1 requires `CUDA_DEVICE_MAX_CONNECTIONS=1` in the Ray runtime-env | -| Qwen3-30B-A3B MoE | colocated, `moe_tp=1` pure EP (shipped recipe) | P5ENx2 | SUCCEEDED | 0.578 | 0.0 | shipped 30B block, EP_SIZE=ROLLOUT_GPUS_PER_ENGINE=2 so the rollout MoE is pure expert-parallel; trains cleanly, comparable to dense 4B. `--use-distributed-optimizer` shards the 30B optimizer state to fit H200 | -| Qwen3-30B-A3B MoE | SGLang `moe_tp>1` and `moe_ep>1`, fusion left on (campaign config) | P5ENx2 | completes, degenerate | 0.0 | 0.56 to 0.80 | root cause, below: a FlashInfer allreduce-fusion bug drops the moe-tp reduce in this combined path. Reproduced disaggregated (per-engine 4 / EP 2 -> moe_tp=2) and in stock serving | -| Qwen3-30B-A3B MoE | colocated, `moe_tp=2` x `moe_ep=2`, fusion disabled (recipe auto) | P5ENx2 | SUCCEEDED | 0.555 | 0.0 | per-engine 4 / EP 2 -> moe_tp=2; the recipe adds `--sglang-enforce-disable-flashinfer-allreduce-fusion` automatically and the combined geometry trains cleanly, comparable to pure EP (0.578) | -| Qwen2.5-72B dense | disaggregated TP4 PP2, actor8+rollout8 (campaign config) | P5ENx2 | OOM | -- | -- | did not fit on P5ENx2: ~144 GiB/GPU (8-way shard, DP=1 so Adam cannot shard) vs 141 GiB. Expected to fit **P6B300x2** (288 GiB) or an H200 layout with optimizer sharding (DP>1 / TP8 / offload) -- not run. `rms_norm_eps` is 1e-6, not 1e-5 | - -Takeaways: the disaggregated (`COLOCATE=false`) weight-sync path works on miles and is now -verified, not just argv-rendered. The dense 4B and GLM-Z1-9B cases train cleanly, and the 30B -MoE trains cleanly too once the SGLang rollout runs pure expert-parallel (`moe_tp=1`): reward -0.578, no repetition, comparable to the dense 4B run. What degenerates is not "the 30B MoE" or -"SGLang expert parallelism" -- pure expert-parallel (`moe_ep=8`) and pure tensor-parallel -(`moe_ep=1`) both generate cleanly -- but specifically the combined path where the rollout MoE -runs tensor-parallel AND expert-parallel at once (`moe_tp>1` and `moe_ep>1`). The earlier -"colocated 30B degenerates" reading conflated the Megatron TP/EP labels with the SGLang rollout -geometry; the run that degenerated had the rollout at `moe_tp>1`, and the shipped colocated -block (`moe_tp=1`) does not. See "30B MoE root cause" below. -The 72B dense case did not fit the 16-GPU disaggregated layout on H200 in the configurations -tried (DP=1 leaves the optimizer unshardable; DP>1 / TP8 / offload were not attempted) -- which -is consistent with slime listing 72B as a config without measured evidence. - -On the comparison with slime specifically: slime's "Supported Model Sizes" table lists -Qwen3-30B-A3B and Qwen2.5-72B as parallelism configurations (TP/PP and rollout/training GPU -counts), but ships no runnable env for them and reports no reward/success metric, so it is not -evidence that either trains -- the 30B is verified here (reward 0.578) and unverified on slime, -and the 72B is unverified on both: - -- The 72B layout slime tabulates (TP4 PP2, training on 8 GPUs, so 8-way sharding with DP=1) needs - roughly 18 GiB weights + 18 GiB grads + ~108 GiB Adam state = ~144 GiB per GPU with a naive - distributed optimizer that cannot shard at DP=1. That exceeds the H200's 141 GiB here (hence the - OOM) and is well above the H100 80 GiB in slime's own table -- i.e. the tabulated layout does not - fit either card as written, which is why "slime does it on p5" has no measured run behind it. - Fitting 72B GRPO needs optimizer sharding (DP>1, i.e. more actor GPUs / the colocated-16 layout), - heavier model parallel (TP8), or CPU/optimizer offload -- none of which were attempted here. - -## 30B MoE root cause - -The 30B MoE degeneration is not "the 30B model" and not "SGLang expert parallelism". It is a -single, narrow condition in the SGLang build shipped in the miles image: the rollout MoE -corrupts its own output when it runs **tensor-parallel and expert-parallel at the same time** -(`moe_tp>1` and `moe_ep>1`). Either axis alone is fine. - -SGLang derives the rollout MoE geometry from two recipe flags: -`moe_ep = --sglang-expert-parallel-size` (EP_SIZE) and -`moe_tp = --rollout-num-gpus-per-engine / EP_SIZE`. So the trigger is set by the ratio of the -per-engine GPU count to EP_SIZE, not by the Megatron TP/EP -- which is why labelling the runs -by Megatron TP/EP hid it. - -GRPO, same 30B model and cluster, `NUM_ROLLOUT=2`, metrics from the trainer's event files: - -| rollout geometry | `moe_tp` x `moe_ep` | reward | repetition | -|---|---|---|---| -| colocated, per-engine 2, EP 2 (shipped) | 1 x 2 (pure EP) | 0.578 | 0.0 | -| colocated, per-engine 2, EP 1 | 2 x 1 (pure TP) | 0.531 | 0.0 | -| disaggregated, per-engine 4, EP 2 | 2 x 2 (combined) | 0.0 | 0.56 | - -Reproduced without any RL or weight-update code, in a stock `sglang.Engine` on the same -`/fsx/models/Qwen3-30B-A3B`, `temperature=0.0`, on math prompts (4-gram repetition score): - -| `tp_size` | `ep_size` | `moe_tp` x `moe_ep` | mean repetition | sample output | -|---|---|---|---|---| -| 8 | 8 | 1 x 8 (pure EP) | 0.009 | coherent ("...Okay, so I need to solve the equation 3x + 7 = 22...") | -| 8 | 4 | 2 x 4 (combined) | 0.807 | " 7. 7. 7. 7..." | -| 8 | 2 | 4 x 2 (combined) | 0.327 | ",,,,, and and and", "10101010...", "aaaa..." | - -That the stock engine reproduces it rules out the miles RL path (weight sync, the on-policy -topk branch, Megatron) as the cause; it is in SGLang's serving path. - -What the code shows (SGLang `0.5.16.dev` in the image): - -- `models/qwen3_moe.py` `forward_normal` runs two post-experts all-reduces -- an expert-parallel - one over the moe-ep group (guarded by `self.ep_size = moe_ep_size > 1`) and a tensor-parallel - one over the moe-tp group (guarded by `self.tp_size = moe_tp_size > 1`). With `moe_tp=1` the - second is skipped; with `moe_ep=1` the first is skipped; only the combined case runs both. -- `should_skip_post_experts_all_reduce` returns `False` for both paths in this configuration - (triton runner, no dp-attention, no flashinfer/deepep A2A, no reduce-scatter), so neither - reduce is being dropped -- a missing reduce is not the cause. -- the standard dispatcher's `local_expert_mapping` is built from `moe_ep_rank` only, and the - moe-tp ranks inside one expert-parallel group correctly share it, so the global->local expert - mapping is not itself wrong. - -The `forward_normal` reduces, the moe-ep/moe-tp process groups (orthogonal by construction: for -`tp=8, ep=2` the ep groups are `{0,4},{1,5},{2,6},{3,7}` and the tp groups `{0,1,2,3},{4,5,6,7}`), -the weight sharding and the dispatcher's `local_expert_mapping` are all correct. The defect is in -the **FlashInfer allreduce+RMSNorm fusion**. On SM90/SM100 it is auto-enabled for Qwen3-MoE -(`tp_size>1`, no dp-attention, `moe_a2a=none`) regardless of `moe_ep`/`moe_tp`. With it on, both -post-experts reduces in `forward_normal` are skipped and deferred to the next layer's fused -`layernorm.forward_with_allreduce_fusion`; that fused reduce -(`layernorm.py::_forward_with_allreduce_fusion`, `flashinfer_comm_fusion.py`) picks its group with -`if moe_ep_size>1: moe_ep_group else: moe_tp_group`, treating the two axes as mutually exclusive. -When both are `>1` it reduces over the moe-ep group only and never over the moe-tp group, so each -rank keeps a partial sum over the intermediate dimension -- a total collapse from layer 0. Pure EP -and pure TP escape it because `_MOE_EP`/`_MOE_TP` alias the full TP group, so either branch reduces -over all ranks. Causal proof: rerunning the same combined config (`tp=8, ep=2, moe_tp=2`, triton) -with `enforce_disable_flashinfer_allreduce_fusion=True` generates cleanly (4-gram repetition 0.006, -same as pure EP). This is a genuine SGLang bug; a minimal stock-`sglang.Engine` reproducer is -captured for an upstream report. - -The recipe's response is to disable the fusion for the affected geometry rather than forbid it: -`run_grpo_qwen3_30b_a3b.sh` computes `moe_tp = ROLLOUT_GPUS_PER_ENGINE / EP_SIZE`, and when both -`moe_tp>1` and `moe_ep>1` it adds `--sglang-enforce-disable-flashinfer-allreduce-fusion` so the -combined geometry trains correctly too; pure EP (the shipped default) and pure TP keep the fusion. - -## A correction - -This log has now corrected the 30B MoE row twice, and both corrections are worth keeping -visible. First, an early table listed the configuration as simply "Verified" on the strength -of a smoke run that exited 0 while producing reward 0.0 and repetition 0.96 -- "the job -completed" written up as "the configuration works". The table was split to separate those -claims. Second, the follow-up reading -- "the 30B MoE degenerates, cause suspected in SGLang -expert parallelism" -- was itself too broad: it generalised from runs that happened to have -the rollout at `moe_tp>1`, and it labelled runs by Megatron TP/EP, which is not what sets the -SGLang rollout geometry. Measuring the model across the actual rollout geometries showed the -shipped colocated block trains cleanly (reward 0.578) and isolated the real trigger. The -lesson both times is the same: read the metric, name the exact variable, and do not let a -plausible summary outrun the measurement. diff --git a/3.test_cases/pytorch/miles/env_vars.colocated.example b/3.test_cases/pytorch/miles/env_vars.colocated.example index 2ef2e5324..cf78fb10f 100644 --- a/3.test_cases/pytorch/miles/env_vars.colocated.example +++ b/3.test_cases/pytorch/miles/env_vars.colocated.example @@ -64,9 +64,10 @@ export FSX_CLAIM="fsx-claim" # the README's Verification Status for what has and has not been run on which hardware. # Node placement is a label KEY:VALUE pair, and BOTH are cluster-specific. The manifest # schedules workers on ${GPU_NODE_LABEL_KEY}:${GPU_NODE_ROLE} and the head on -# ${CPU_NODE_LABEL_KEY}:${CPU_NODE_ROLE}. The defaults below (key "node-role") match the -# awsome-distributed-ai Terraform EKS reference; a generic EKS or HyperPod cluster has NO -# "node-role" label. Find yours with `kubectl get nodes --show-labels` -- e.g. the well-known +# ${CPU_NODE_LABEL_KEY}:${CPU_NODE_ROLE}. The defaults below use a custom key "node-role", +# which a generic EKS or HyperPod cluster does NOT carry until you add it +# (`kubectl label node node-role=gpu`). Find an existing label with +# `kubectl get nodes --show-labels` -- e.g. the well-known # `node.kubernetes.io/instance-type` (value "p5en.48xlarge"), an EKS managed-nodegroup label # `eks.amazonaws.com/nodegroup`, or a HyperPod `sagemaker.amazonaws.com/instance-group-name`. # A wrong key OR value leaves pods stuck Pending with FailedScheduling, not an obvious error. @@ -122,10 +123,10 @@ export EVAL_DATA="/fsx/data/aime-2024/aime-2024.jsonl" # scheduler built for a different total-iteration count. train_iters = num_rollout # here (16 x 8 // 128 = 1 step/rollout), so NUM_ROLLOUT == optimizer steps. export CHECKPOINT_DIR="/fsx/runs/qwen3-4b/ckpt" -# SAVE_INTERVAL is set beyond NUM_ROLLOUT on purpose: miles's save_model() currently -# fails with a pickle-truncation error (see README Known Issues), so a mid-run save -# would crash the job. Lower this ONLY once that issue is resolved; until then a -# checkpoint-triggering value is a guaranteed failure, so the default must not trigger one. +# SAVE_INTERVAL is kept at or above NUM_ROLLOUT so short smoke runs do not pay the +# ~270s in-loop save cost on every rollout; the single end-of-run save still fires and +# succeeded on every verification run (a torch_dist checkpoint is written to CHECKPOINT_DIR). +# Reloading that checkpoint and the megatron2hf back-conversion are untested (README Known Issues). export SAVE_INTERVAL=1000 export NUM_ROLLOUT=100 export ROLLOUT_BATCH_SIZE=16 diff --git a/3.test_cases/pytorch/miles/env_vars.disaggregated.example b/3.test_cases/pytorch/miles/env_vars.disaggregated.example index 044ca1e8b..f8cb8051a 100644 --- a/3.test_cases/pytorch/miles/env_vars.disaggregated.example +++ b/3.test_cases/pytorch/miles/env_vars.disaggregated.example @@ -24,8 +24,7 @@ # Assumes: # * env_vars (Option A: Qwen3-4B, colocated) has already been sourced. # * A CPU node pool named ${REWARD_NODE_GROUP} already exists, in the same AZ as the GPU -# pool, with no EFA. Create it however your cluster provisions capacity -- a Karpenter -# NodePool, an EKS managed node group, or a SageMaker HyperPod instance group -- then set +# pool, with no EFA. Create it however your cluster provisions capacity, then set # the name below to match. This test case does not create it for you. # ----- Heavier GRPO config (GPU tier) ----- diff --git a/3.test_cases/pytorch/miles/miles.Dockerfile b/3.test_cases/pytorch/miles/miles.Dockerfile index b354457ba..407742eea 100644 --- a/3.test_cases/pytorch/miles/miles.Dockerfile +++ b/3.test_cases/pytorch/miles/miles.Dockerfile @@ -3,13 +3,11 @@ # ============================================================================ # miles (radixark/miles) + EFA image for Amazon EKS # ============================================================================ -# Strategy C (see docs/PORT_NOTES.md): the miles upstream image already bundles -# a matched PyTorch 2.11 / CUDA 13.0.1 stack, the sglang-miles fork, the radixark -# Megatron-LM fork, prebuilt flash_attn / TE / apex wheels, and the B300 sm_103 TE -# FA2 whitelist patch. (Unlike an SGLang-fork HTTP weight-sync path, miles drives -# weight sync through Ray actor methods on the rollout engine -- see README / -# docs/PORT_NOTES.md.) Rebuilding that stack on an NGC base is infeasible (wheel -# ABI mismatch), so we take the miles image as-is and add ONLY the AWS EFA stack. +# The miles upstream image already bundles a matched PyTorch 2.11 / CUDA 13.0.1 +# stack, the sglang-miles fork, the radixark Megatron-LM fork, prebuilt +# flash_attn / TE / apex wheels, and the sm_103 TE FA2 whitelist patch. Rebuilding +# that stack on an NGC base is infeasible due to wheel ABI mismatch, so we take the +# miles image as-is and add ONLY the AWS EFA stack. See README. # # Pin the base by DIGEST, not by tag. radixark publishes dev-* as mutable snapshots and # has already deleted a tag this file previously pinned, so a dated tag alone gives neither diff --git a/3.test_cases/pytorch/miles/recipe/run_grpo_qwen3_30b_a3b.sh b/3.test_cases/pytorch/miles/recipe/run_grpo_qwen3_30b_a3b.sh index bd86ffa8c..5cee933c8 100644 --- a/3.test_cases/pytorch/miles/recipe/run_grpo_qwen3_30b_a3b.sh +++ b/3.test_cases/pytorch/miles/recipe/run_grpo_qwen3_30b_a3b.sh @@ -219,7 +219,7 @@ TRAIN_ARGS=( # (empty otherwise); the fusion mis-handles that combined path in this build. ${SGLANG_FUSION_ARGS[@]+"${SGLANG_FUSION_ARGS[@]}"} # Lowercase only: this reaches uvicorn's log_level, whose LOG_LEVELS dict has no "WARN" - # key, and the KeyError kills the rollout server before it binds. See docs/PORT_NOTES.md. + # key, and the KeyError kills the rollout server before it binds. See README (miles-specific requirements). --sglang-log-level warning # Careful when adding --sglang-* flags: miles registers them from SGLang's live # ServerArgs with ignore_unknown_args, so a flag SGLang has since removed is silently diff --git a/3.test_cases/pytorch/miles/recipe/run_grpo_qwen3_4b.sh b/3.test_cases/pytorch/miles/recipe/run_grpo_qwen3_4b.sh index cdf00e7e4..5f024a851 100644 --- a/3.test_cases/pytorch/miles/recipe/run_grpo_qwen3_4b.sh +++ b/3.test_cases/pytorch/miles/recipe/run_grpo_qwen3_4b.sh @@ -274,7 +274,7 @@ TRAIN_ARGS=( --sglang-mem-fraction-static "${MEM_FRACTION}" # Lowercase only: this reaches uvicorn's log_level, whose LOG_LEVELS dict has no "WARN" # key, and the KeyError kills the rollout server before it binds -- training then hangs - # on a health check that never passes. See docs/PORT_NOTES.md. + # on a health check that never passes. See README (miles-specific requirements). --sglang-log-level warning # Flags injected via the EXTRA_TRAIN_ARGS env var. The shipped default is just diff --git a/3.test_cases/pytorch/miles/scripts/convert_checkpoint.sh b/3.test_cases/pytorch/miles/scripts/convert_checkpoint.sh index 35ed19ed0..02ff2de93 100644 --- a/3.test_cases/pytorch/miles/scripts/convert_checkpoint.sh +++ b/3.test_cases/pytorch/miles/scripts/convert_checkpoint.sh @@ -39,7 +39,9 @@ case "${DIRECTION}" in MODEL_SCRIPT="" HF_PATH="" SAVE_PATH="" - NUM_GPUS=1 + # Single process by default (validated); set CONVERT_NUM_GPUS or pass --num-gpus to + # parallelize the conversion with torchrun. + NUM_GPUS="${CONVERT_NUM_GPUS:-1}" while [[ $# -gt 0 ]]; do case "$1" in