diff --git a/.github/workflows/miles-pin-freshness.yml b/.github/workflows/miles-pin-freshness.yml new file mode 100644 index 000000000..8eee02569 --- /dev/null +++ b/.github/workflows/miles-pin-freshness.yml @@ -0,0 +1,164 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +# +# Fails if an external artifact that 3.test_cases/pytorch/miles pins by tag/version has +# disappeared from its registry/source. This test case needs 8-16 H200 GPUs to run GRPO +# end to end, and no runner in this repo (GitHub-hosted or self-hosted) provides that for +# miles today, so this workflow does NOT build the image and does NOT run training -- it +# only asserts that each pin still resolves. That is a real, current failure mode: the +# base image tag pinned in miles.Dockerfile (MILES_BASE_TAG=dev-202607182122) has already +# been deleted from Docker Hub. +# +# What this catches: +# - the pinned base image digest (MILES_BASE_DIGEST) removed from the registry, and +# separately, the documented tag drifting off that digest (advisory) +# - the EFA installer tarball for EFA_INSTALLER_VERSION removed from the AWS download host +# - the GDRCopy tag for GDRCOPY_VERSION deleted from NVIDIA/gdrcopy on GitHub +# - a reward_service/requirements.txt pin no longer published on PyPI +# +# What this does NOT catch (needs hardware -- see README.md "Verification Status"): +# - whether a replacement image/version still builds, boots, or trains correctly +# - behavioral drift in a dev tag that still resolves (radixark/miles dev-* tags are +# mutable snapshots; "resolves" is not "is the bits this test case was verified against") +# - drift in references with no separate pin, e.g. the exact radixark/Megatron-LM commit +# or SGLang 0.5.16.dev bundled inside MILES_BASE_TAG +# - the top-level requirements.txt, which its own header marks reference-only (miles.Dockerfile +# never installs from it) +# +# A registry/API hiccup (rate limit, 5xx, timeout) is logged as a warning, not a failure -- +# only a confirmed 404 (tag/version/ref genuinely absent) fails the job. + +name: miles pin freshness + +on: + pull_request: + paths: + - "3.test_cases/pytorch/miles/**" + schedule: + # Independent of PR activity: radixark/miles prunes dev-* tags on its own cadence, so a + # pin can go stale with zero activity in this repo. Weekly, Monday morning UTC. + - cron: "17 6 * * 1" + workflow_dispatch: + +permissions: + contents: read + +jobs: + check-pins: + runs-on: ubuntu-latest + timeout-minutes: 10 + defaults: + run: + working-directory: 3.test_cases/pytorch/miles + steps: + - uses: actions/checkout@v4 + + - name: Check the pinned base image digest still resolves + run: | + set -euo pipefail + # The build resolves MILES_BASE_DIGEST, so that is what has to be checked. The tag + # is read only to report which snapshot the digest came from, and to warn if the + # two have drifted apart (the tag moved, or was bumped without the digest). + DIGEST=$(grep -oP '^ARG MILES_BASE_DIGEST=\K\S+' miles.Dockerfile) + TAG=$(grep -oP '^ARG MILES_BASE_TAG=\K\S+' miles.Dockerfile) + echo "pinned digest: ${DIGEST}" + echo "documented tag: ${TAG}" + + TOKEN=$(curl -s --max-time 20 --retry 2 --retry-delay 5 \ + "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"])') || TOKEN="" + if [ -z "${TOKEN}" ]; then + echo "::warning::Could not obtain a Docker Hub token; skipping the digest check." + exit 0 + fi + ACCEPT='application/vnd.oci.image.index.v1+json,application/vnd.docker.distribution.manifest.list.v2+json,application/vnd.docker.distribution.manifest.v2+json' + + CODE=$(curl -s -o /dev/null -w '%{http_code}' --max-time 30 --retry 2 --retry-delay 5 \ + -H "Authorization: Bearer ${TOKEN}" -H "Accept: ${ACCEPT}" \ + "https://registry-1.docker.io/v2/radixark/miles/manifests/${DIGEST}" || echo 000) + if [ "${CODE}" = "404" ]; then + echo "::error file=3.test_cases/pytorch/miles/miles.Dockerfile::The pinned base image digest ${DIGEST} is gone from the registry. Pick a live dated tag from https://hub.docker.com/r/radixark/miles/tags, resolve its digest (see the comment in miles.Dockerfile), re-run the hardware verification matrix in README.md, then bump both MILES_BASE_DIGEST and MILES_BASE_TAG." + exit 1 + elif [ "${CODE}" != "200" ]; then + echo "::warning::Digest check was inconclusive (HTTP ${CODE}); not failing the job." + exit 0 + fi + echo "OK: the pinned digest still resolves." + + # Advisory only: a dev-* tag is mutable, so the tag pointing elsewhere is normal + # and is not a build problem. It does mean the comment is stale. + TAG_DIGEST=$(curl -sI --max-time 30 -H "Authorization: Bearer ${TOKEN}" -H "Accept: ${ACCEPT}" \ + "https://registry-1.docker.io/v2/radixark/miles/manifests/${TAG}" \ + | tr -d '\r' | awk 'tolower($1)=="docker-content-digest:"{print $2}') || TAG_DIGEST="" + if [ -n "${TAG_DIGEST}" ] && [ "${TAG_DIGEST}" != "${DIGEST}" ]; then + echo "::warning::${TAG} now points at ${TAG_DIGEST}, not the pinned ${DIGEST}. The build is unaffected, but the documented tag no longer names the pinned image." + fi + + - name: Check EFA installer tarball is still downloadable + run: | + set -euo pipefail + VER=$(grep -oP '^ARG EFA_INSTALLER_VERSION=\K\S+' miles.Dockerfile) + echo "EFA_INSTALLER_VERSION=${VER}" + URL="https://efa-installer.amazonaws.com/aws-efa-installer-${VER}.tar.gz" + CODE=$(curl -s -o /dev/null -I -w '%{http_code}' --max-time 20 --retry 2 --retry-delay 5 "${URL}" || echo 000) + if [ "${CODE}" = "200" ]; then + echo "OK: EFA installer ${VER} tarball is downloadable." + elif [ "${CODE}" = "404" ]; then + echo "::error file=3.test_cases/pytorch/miles/miles.Dockerfile::${URL} is gone (HTTP 404). Bump EFA_INSTALLER_VERSION to a version still published at https://efa-installer.amazonaws.com/." + exit 1 + else + echo "::warning::EFA installer URL check was inconclusive (HTTP ${CODE}); not failing the job." + fi + + - name: Check GDRCopy tag exists on GitHub + run: | + set -euo pipefail + VER=$(grep -oP '^ARG GDRCOPY_VERSION=\K\S+' miles.Dockerfile) + echo "GDRCOPY_VERSION=${VER}" + URL="https://api.github.com/repos/NVIDIA/gdrcopy/git/refs/tags/${VER}" + CODE=$(curl -s -o /dev/null -w '%{http_code}' --max-time 20 --retry 2 --retry-delay 5 "${URL}" || echo 000) + if [ "${CODE}" = "200" ]; then + echo "OK: NVIDIA/gdrcopy@${VER} exists." + elif [ "${CODE}" = "404" ]; then + echo "::error file=3.test_cases/pytorch/miles/miles.Dockerfile::NVIDIA/gdrcopy tag ${VER} is gone (HTTP 404 from the GitHub API). Bump GDRCOPY_VERSION to a tag that still exists." + exit 1 + else + echo "::warning::GDRCopy tag check was inconclusive (HTTP ${CODE}); not failing the job. The GitHub API is rate-limited for unauthenticated requests (60/hour/IP), which shared GitHub-hosted runner IPs can hit." + fi + + - name: Check reward_service pinned PyPI packages still exist + run: | + set -euo pipefail + python3 - << 'PY' + import re + import sys + import urllib.error + import urllib.request + + fail = False + with open("reward_service/requirements.txt") as f: + for line in f: + line = line.strip() + m = re.match(r'^([A-Za-z0-9_.\-]+)(?:\[[A-Za-z0-9_,.\-]+\])?==([A-Za-z0-9_.\-]+)$', line) + if not m: + continue + pkg, ver = m.group(1), m.group(2) + url = f"https://pypi.org/pypi/{pkg}/{ver}/json" + try: + with urllib.request.urlopen(url, timeout=20) as resp: + code = resp.status + except urllib.error.HTTPError as e: + code = e.code + except Exception as e: + print(f"::warning::PyPI check for {pkg}=={ver} was inconclusive ({e}); not failing the job.") + continue + if code == 200: + print(f"OK: {pkg}=={ver} exists on PyPI.") + elif code == 404: + print(f"::error file=3.test_cases/pytorch/miles/reward_service/requirements.txt::{pkg}=={ver} is gone from PyPI (HTTP 404). Bump the pin in reward_service/requirements.txt.") + fail = True + else: + print(f"::warning::PyPI check for {pkg}=={ver} was inconclusive (HTTP {code}); not failing the job.") + if fail: + sys.exit(1) + PY 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/README.md b/3.test_cases/pytorch/miles/README.md new file mode 100644 index 000000000..3220aecac --- /dev/null +++ b/3.test_cases/pytorch/miles/README.md @@ -0,0 +1,768 @@ +# 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. + +## 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 | | +| +----------------------------------------------------------+ | +| | +| +-----------------------------+ | +| | CPU node (Ray head) | <- pulls the ~18 GB miles image; | +| | no GPU, large ephemeral-fs | needs ample ephemeral-storage | +| +-----------------------------+ | +| | +| 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 | +| **CPU node** | At least one CPU-only node with ample ephemeral-storage (>=150 GiB root volume recommended) for the Ray head, which pulls the ~18 GB miles 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) -- completes without OOM/crash | 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 with explicit expert parallelism | +| Qwen3-30B-A3B MoE GRPO, colocated, 2 nodes (16 GPU) -- produces a usable trained model | Known Issue | across 4 independent runs the model falls into a repetition loop and never reaches an answer (`rollout/repetition_frac` 0.48-0.70, `rollout/truncated` 0.97-0.99, `rollout/raw_reward` 0.0 -- vs. 0.0 / 0.43-0.58 / 0.42-0.55 on dense Qwen3-4B). Root cause found: SGLang's expert parallelism. EP=1 generates correctly at TP 1/4/8; EP>1 degenerates at every TP, worsening with EP. Not fixable here -- EP=1 needs 108.76 GiB per rank to train. See [Known Issues](#known-issues) | +| Qwen3-4B GRPO, **disaggregated** (`COLOCATE=false`), 2 nodes | UNVERIFIED | The recipe builds this layout as of the commit that added the `COLOCATE` branch; before it, `--colocate` was passed unconditionally and `COLOCATE=false` was refused, so the layout was unreachable. Reaching it was verified only to the point of a rendered argv (`--colocate` absent, 104 other tokens unchanged) -- no run has been completed | +| 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 on a CPU node) | UNVERIFIED | the shipped `kubernetes/raycluster.yaml` places the head on a CPU-only node (the recommended production shape); the verification runs above used a head-on-GPU overlay because the validation cluster had no large-disk CPU node available at the time | +| 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 + +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). Note the GPU NodePool's `node-role` label value; + `env_vars`' `GPU_NODE_ROLE` must match it. +2. `kubectl` configured to access the cluster +3. The KubeRay operator installed on the cluster (see step 0 below) +4. A CPU-only node group with ample ephemeral-storage (>=150 GiB root volume) for the Ray + head, which pulls the ~18 GB miles image +5. FSx for Lustre persistent volume claim (`fsx-claim`) available +6. Container registry access (e.g. Amazon ECR) for building/pushing images +7. A Hugging Face account and access token for model downloads +8. 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) + +## 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. This configuration +completes without crashing but currently produces degenerate, zero-reward generation (see +[Known Issues](#known-issues) item 2); treat it as a layout/flag reference, not a working +alternate training configuration, until that is root-caused. + +### 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 - +``` + +### 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 places the Ray head on a CPU-only node (`nodeSelector: node-role: cpu`) +and GPU workers on the GPU node group, each worker declaring a `gpu_node` custom Ray resource +(see [miles-specific requirements](#miles-specific-requirements-found-on-real-hardware) for +why). Ensure the CPU node's root volume has enough headroom to pull 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 + +```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 (completes without crashing, but every +# run to date shows degenerate, zero-reward generation -- see Known Issues item 2 below): +# first uncomment the ALTERNATE block in env_vars, then: +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. **Qwen3-30B-A3B MoE GRPO produces degenerate, zero-reward generation.** The colocated + 16-GPU configuration completes every rollout cycle without crashing (see Verification + Status), but across four independent runs the model falls into a repetition loop and + never reaches an answer: `rollout/repetition_frac` 0.48-0.70 (dense Qwen3-4B: exactly + 0.0), `rollout/truncated` 0.97-0.99 (dense: 0.43-0.58), `rollout/raw_reward` 0.0 (dense: + 0.42-0.55), and `train/mis_ppl_ratio` 1.32-1.78 (dense: 1.0007-1.033). Four candidate + causes were tested and ruled out: raising `--rollout-max-response-len` 8192 -> 16384 + left truncation unchanged at 0.992; `--check-weight-update-equal` passes + (`rollout/weight_version` uniform 1.0, `mixed_version_ratio` 0.0), so trainer/rollout + weights are not diverging; lowering `--rollout-temperature` 1.0 -> 0.6 (the model's own + `generation_config` value) left `repetition_frac` unchanged to 4 decimal places; and the + converted checkpoint's `config.json` / `generation_config.json` match the working dense + model (same `eos_token_id`, `vocab_size`, normal file sizes). + + **The cause is expert parallelism, in SGLang, outside the training loop.** Isolating the + MoE path as suggested above -- serving the same checkpoint from SGLang directly, with no + miles, no Megatron and no GRPO -- reproduces the failure, and varying one axis at a time + locates it. `repetition_frac` over 32 prompts: + + | | EP=1 | EP=2 | EP=4 | + |---|---|---|---| + | TP=1 | 0.000 | (cannot start) | - | + | TP=4 | 0.000 | 0.875 | - | + | TP=8 | 0.000 | 0.594 | 0.844 | + + Tensor parallelism is not involved: EP=1 is clean at TP 1, 4 and 8. Expert parallelism is, + and the rate grows with it. Ruled out along the way: the model's own recommended sampling + (temperature 0.6 / top_p 0.95 / top_k 20) still gives 0.875, and the default `auto` MoE + runner backend still gives 0.844, so neither the temperature nor the `triton` backend + selection is responsible. All 18867 expert weight keys are present in the checkpoint index + with no NaN/Inf tensors, so conversion is not either. + + Dropping to EP=1 fixes generation but cannot train this model: the trainer then needs + 108.76 GiB per rank against 139.80 GiB of capacity, because every rank holds all 128 + experts. Raising TP from 8 to 16 changes that request by 0.3% -- tensor parallelism does not + shard the expert weights, only EP and PP do. + + So treat the 30B MoE recipe as a layout/flag reference (how to fit a 30B MoE actor on + 16x H200 without OOM) rather than a working training configuration, and note that the two + ways out of the repetition -- EP=1, or EP>1 with a fixed SGLang -- are respectively + out of memory and not yet available. + +## 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 +├── 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=CPU node, worker=GPU) +│ ├── 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 +``` + +## 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, completes without crashing -- see +[Known Issues](#known-issues) item 2 for its degenerate-generation problem):** + +``` +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:` | +| 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-training CONTRIBUTING guidelines](https://github.com/awslabs/awsome-distributed-training/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-training](https://github.com/awslabs/awsome-distributed-training) +- [KubeRay Documentation](https://docs.ray.io/en/latest/cluster/kubernetes/index.html) + +## Security + +See [CONTRIBUTING](https://github.com/awslabs/awsome-distributed-training/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..403274d18 --- /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-training +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..4befa7a42 --- /dev/null +++ b/3.test_cases/pytorch/miles/docs/VERIFICATION_LOG.md @@ -0,0 +1,158 @@ +# 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. The 30B MoE configuration in this test case 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. That is exactly what happened +here once, which is why the README now lists "completes without OOM/crash" and "produces a +usable trained model" as separate rows. + +For comparison, on the same cluster and recipe: + +| | dense 4B (this run) | 30B MoE | +|---|---|---| +| `rollout/repetition_frac` | 0.0 | 0.48 to 0.70 | +| `rollout/raw_reward` | 0.516 | 0.0 | +| `rollout/truncated_ratio` | 0.484 | 0.97 to 0.99 | +| exit status | 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 shipped `kubernetes/raycluster.yaml` with the head on a CPU node. The runs above used a + head-on-GPU overlay, because the validation cluster had no large-disk CPU node at the time. + Listed as its own UNVERIFIED row in the README table. +- 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. + +## 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. + +## A correction + +An earlier version of the verification table listed the 30B MoE configuration as simply +"Verified", on the strength of a smoke run that exited 0. That run had reward 0.0 and +repetition 0.96. The job did succeed, and the memory layout claim in that row was and is +true, but "the job completed" was written up as "the configuration works". The table now +separates those two claims, and this log records the metrics that distinguish them. 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..75bc4c708 --- /dev/null +++ b/3.test_cases/pytorch/miles/env_vars.colocated.example @@ -0,0 +1,173 @@ +# 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. +export GPU_NODE_ROLE="gpu-p5en" # e.g. "gpu-b300" on a B300 pool +# 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. +# +# KNOWN ISSUE: this configuration completes every rollout cycle without crashing, but +# across 4 independent runs the model falls into a repetition loop and never reaches an +# answer (reward stays 0.0). Root cause is unresolved (weight-update, response-length, +# sampling-temperature, and checkpoint-corruption causes were all ruled out by +# measurement) -- see README.md Known Issues item 2. Treat this block only as a +# layout/flag reference for fitting a 30B MoE actor on 16x H200 without OOM, not as proof +# the resulting training is useful. +# +# 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..b10ad0638 --- /dev/null +++ b/3.test_cases/pytorch/miles/env_vars.disaggregated.example @@ -0,0 +1,72 @@ +# 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}" +# (or use kubernetes/buildkit-job.yaml with the Dockerfile and tag overridden). +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/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/kubernetes/data-prep-pod.yaml b/3.test_cases/pytorch/miles/kubernetes/data-prep-pod.yaml new file mode 100644 index 000000000..27044d55c --- /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: slime-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..34307657f --- /dev/null +++ b/3.test_cases/pytorch/miles/kubernetes/raycluster.yaml @@ -0,0 +1,171 @@ +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 + 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 + nodeSelector: + node-role: cpu + rayVersion: 2.55.1 + workerGroupSpecs: + # One worker per GPU node, driven by ${ACTOR_NUM_NODES} so the manifest cannot disagree with + # the topology the recipe was given: too few workers leaves the job waiting on placement it + # can never get, with nothing pointing at this file as the cause. podAntiAffinity (below) + # keeps the workers on separate nodes so a multi-node actor spans them. + - groupName: gpu-workers + maxReplicas: ${ACTOR_NUM_NODES} + minReplicas: ${ACTOR_NUM_NODES} + numOfHosts: 1 + rayStartParams: + num-gpus: '8' + resources: '"{\"gpu_node\": 1}"' + replicas: ${ACTOR_NUM_NODES} + 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: + node-role: ${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 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/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/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..91c8e3369 --- /dev/null +++ b/3.test_cases/pytorch/miles/recipe/run_grpo_qwen3_30b_a3b.sh @@ -0,0 +1,237 @@ +#!/usr/bin/env bash +# STATUS: Verified -- completes without OOM/crash (colocated, 2 nodes / 16 GPU H200) -- +# GRPO steps run to completion with --colocate + --use-distributed-optimizer + triton +# MoE runner. Known Issue: across 4 independent runs the resulting generation is +# degenerate (repetition loop, zero reward); do not treat this as a working training +# configuration -- 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 + +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 " 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}" + # 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..22e5b4ec5 --- /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 "${SGLANG_MEM_FRACTION:-0.8}" + # 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/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 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..cdd08ca0c --- /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="SLIME 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 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..59a47e630 --- /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 SLIME-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 " SLIME 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."