diff --git a/.claude/skills/mad-slurm-multinode/SKILL.md b/.claude/skills/mad-slurm-multinode/SKILL.md new file mode 100644 index 0000000..cfb1a82 --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/SKILL.md @@ -0,0 +1,329 @@ +--- +name: mad-slurm-multinode +description: Deploy and run madengine performance tests on an unprepared SLURM cluster from scratch. Use when the user wants to perf-test a model via a madengine run engine (any engine that has a template under assets/manifests/), set up a fresh node (clone MAD + madengine, conda/miniforge env, pip install), build a mad.env, pick/adapt a run manifest, and launch a multi-node SLURM run. Covers the cluster archetypes documented in references/cluster-types.md (such as CX7/Mellanox-RoCE, AMD-AINIC/Pollara, and Broadcom-Thor2-RoCE), required interactive inputs (node host, work dir, data dir, MAD_DOCKER_BUILDS, SLURM partition/account/qos/reservation/nodelist), and the cluster-specific communication-backend env vars (RCCL/NCCL over the right transport, e.g. RDMA) that decide whether the right transport is exercised. +compatibility: Requires git, docker, conda/miniforge, a SLURM CLI (sbatch/sinfo), rocm-smi or nvidia-smi, and internet access for cloning and image pulls. +metadata: + author: mkuznet1 (mikhail.kuznetsov@amd.com) + version: "0.1" +--- + +# mad-slurm-multinode + +Bring up madengine on a fresh SLURM node and run a Llama-3.1 perf test +end to end: prerequisites -> clone -> conda env -> install -> mad.env -> +manifest -> `madengine run`. Covers the cluster archetypes documented in +[references/cluster-types.md](references/cluster-types.md). + +> **Skill paths — read first.** When this skill activates you are given the +> absolute path to its directory (the folder containing this `SKILL.md`). All +> `scripts/`, `references/`, and `assets/` paths below are relative to THAT +> directory, not to `$WORKDIR`. The workflow `cd`s into `$WORKDIR`, so set the +> skill dir once and reference files through it: +> +> ```bash +> export SKILL_DIR="" +> ``` +> +> Then use `"$SKILL_DIR/scripts/..."`, `"$SKILL_DIR/assets/..."`, etc. Markdown +> links like `references/foo.md` are also relative to `$SKILL_DIR`. + +## When this fires + +- "Run a perf test with madengine on ``" / "set up madengine on this cluster". +- A Llama-3.1 perf workload that has a template under `assets/manifests/`, on SLURM. +- A node that has nothing installed yet (no conda, no repos, no rundir). + +Out of scope: +- RCCL build-vs-build validation delta campaigns (build image A vs B and + compare the transport) — that is a separate workflow, out of scope here. +- Accumulating many iterations across multiple SLURM allocations (large + multi-allocation campaigns) — out of scope here. +- Single-GPU local (non-SLURM) smoke tests -> just `madengine run --tags ...` directly. + +## Responsibilities — who does what + +The skill stays in its lane and follows the steps rather than re-deriving the +workload. The split is: + +- **mad-slurm-multinode (this skill)** configures one run: brings up the node + (repos, conda env), writes `mad.env` and the manifest (paths, SLURM selectors, + transport vars), launches `madengine run`, and reads the perf CSV. +- **madengine** orchestrates: renders the sbatch script from the manifest, + ensures the docker image per node (load tar / build / pull), runs the + multi-node job, and aggregates results. +- **MAD** holds the Dockerfiles and the per-model `run.sh`. Inside the + container, `run.sh` runs the actual workload **and acquires/prepares the + dataset and model weights** (data prep, model download, etc.). The + skill points it at the right paths and launches; staging datasets or + downloading models by hand is `run.sh`'s job, not a pre-step the skill adds. + +**Do not fix madengine or MAD code.** When a failure is rooted in the madengine +orchestrator or the MAD repo (a Dockerfile, `run.sh`, model code, an +orchestration bug), the skill reports it to the user with the evidence and +stops — it does not patch those repos unless the user explicitly asks. The +skill's own edits stay confined to `rundir/` (`mad.env`, the manifest) and the +user-provided inputs. A known per-workload workaround that only sets env/manifest +values (see [references/gotchas.md](references/gotchas.md)) is applied in the +manifest, not by editing madengine/MAD. + +## Required inputs — ask before exploring + +The first action of a run, before any other tool call (Shell, Read, +transcript/past-session search, or web search), is to list the required inputs +that are missing or not clearly identifiable and resolve them in a single +`AskQuestion` round. A missing or unrecognized input is a question to the user, +not a research task — inferring it from prior sessions, `agent-transcripts`, +repo history, or the web is a known failure mode of this skill and is skipped. +Tool calls begin once the required inputs are in hand. + +These come from the user, and the skill asks for them rather than assuming. +The first four have no safe defaults, so the skill asks; a value already in the +conversation gets reused. + +1. **Compute node hostname or IP** to run on (e.g. the login/jump node you SSH + to) — where the bootstrap, build, and `sbatch` submission happen. +2. **Working directory** `$WORKDIR` — the root that holds the cloned repos + (`MAD`, `madengine`), the `rundir/` (the filled `mad.env`, the + filled manifest, `slurm_output/`, and run logs), and the perf-result CSVs. + Picking it explicitly keeps one run's artifacts together and reusable. +3. **Cluster archetype** — one of the archetypes in + [references/cluster-types.md](references/cluster-types.md). You may run + `scripts/detect_cluster_env.sh` to *propose* it, but confirm with the user. +4. **Data root** (`MAD_DATAHOME` + caches) and **`MAD_DOCKER_BUILDS` dir** — + state what each holds so the user can point them at the right place: + - the **data root** holds datasets, tokenizer, and model weights, plus the + `HF_HOME`/`TORCH_HOME`/pip caches the run reads and writes; + - **`MAD_DOCKER_BUILDS`** is the shared image-tar cache: rank 0 saves the + built image there and every worker loads the same tar, so it lives on + shared FS visible to every node (see Gotchas). +5. **SLURM specifics, on demand**: `partition`, `account`, `qos`, + `reservation`, `nodelist`/`exclude`, node count. These are cluster-private + and are not stored in this skill — requested per run. +6. **Branch per repo** for `MAD` and `madengine`. No branch is + assumed — the skill asks which branch to use for each repo and `git switch`es + to it (see "Repo + branch rule"). + +The HF token file `~/.huggingface/token` provides gated Llama-3.1 access; +`mad.env` reads it into `MAD_SECRETS_HFTOKEN` and madengine forwards it to the +container (see [references/manifests.md](references/manifests.md) "Secrets"), so +confirm it exists before launch. + +**Probe one node before asking about cluster shape.** The cluster-shape values +(GPU arch, HCA list, GID index, management iface) are discoverable on the node +itself, so the skill allocates one node from the user-provided partition and +runs the probe there rather than interrogating the user: + +```bash +srun -p [--reservation ] [--nodelist ] -N1 \ + bash "$SKILL_DIR/scripts/detect_cluster_env.sh" +``` + +Only the cluster-private selectors (partition / account / qos / reservation / +nodelist) come from the user; the probe reports the rest. The same holds for +`scripts/preflight.sh`, which reflects whichever node it runs on — for +compute-node values it runs through `srun` on an allocated node. + +Defaults you MAY assume unless told otherwise (state them when you do): +- conda env name `madenv`, Python 3.12. + +**Repo + branch rule (every run):** the repo and the branch for both +`MAD` and `madengine` are confirmed with the user before cloning or +switching. No branch name is assumed — not a default, not one inferred from a +prior session — the branch is asked and then `git switch`ed to. An existing +checkout is left as-is rather than switched silently: when the directory already +exists, the skill shows its current branch +(`git -C branch --show-current`) and asks before changing it. + +## Workflow + +Track progress with this checklist: + +- [ ] Step 0: preflight (`scripts/preflight.sh`) +- [ ] Step 1: clone repos (idempotent) +- [ ] Step 2: conda env (idempotent) +- [ ] Step 3: `pip install -e ./madengine` +- [ ] Step 4: rundir + mad.env (filled + sourced) +- [ ] Step 5: pick + adapt manifest +- [ ] Step 6: launch + collect results + +Full step detail (commands, idempotency guards) is in +[references/deploy-bootstrap.md](references/deploy-bootstrap.md). The short +version: + +Each step below ends with a **Guard** — the enumerated conditions are the only +branches to consider. If a guard matches, take its action; if none match, the +step passed and the next step follows. This keeps the path deterministic and +avoids re-analyzing a step that already succeeded. + +### Run logs + +Every long command (`madengine run`, docker build, conda/pip install) streams +its stdout+stderr to `/.cursor.logs//-.log` +via `tee`, which keeps a run reproducible and inspectable after the fact +(several bugs here surfaced only in these logs). Kinds map to subdirectories +(`.cursor.logs/build/`, `.cursor.logs/run/`, `.cursor.logs/install/`); a +filename carries a UTC timestamp and a short slug and holds no secrets. Example: + +```bash +mkdir -p "$WORKDIR/rundir/.cursor.logs/run" +madengine run --manifest-file run_manifest_.json --live-output \ + -o perf_.csv \ + 2>&1 | tee "$WORKDIR/rundir/.cursor.logs/run/$(date -u +%Y%m%dT%H%M%SZ)-primus-8b.log" +``` + +The log files are `*.log`, which `.gitignore` already ignores, so the +`.cursor.logs/` contents never land in git. + +### Step 0 — Preflight + +```bash +bash "$SKILL_DIR/scripts/preflight.sh" +``` + +Checks Python >= 3.10, docker (present + daemon reachable), conda/miniforge, +SLURM CLI (`sinfo`/`sbatch`), and a GPU SMI (`rocm-smi`/`nvidia-smi`). Fix +any FAIL before continuing. If docker or SLURM is missing this is not a valid +target node — stop and tell the user. + +**Guard:** +- **If Step 0 reports any FAIL** → stop, report which check failed, do not continue. +- **If docker or SLURM is missing** → this node is not a valid target; tell the user and stop. + +### Steps 1-3 — Bootstrap (only do what's missing) + +```bash +cd "$WORKDIR" +# Repo URLs are confirmed with the user; clone only if missing. +[ -d MAD ] || git clone --recursive +[ -d madengine ] || git clone --recursive +# Branch is ASKED per repo (no default). For an existing checkout, show the +# current branch first and ask before switching: +# git -C MAD branch --show-current +# git -C madengine branch --show-current +( cd MAD && git switch "" && git submodule update --init --recursive ) +( cd madengine && git switch "" && git submodule update --init --recursive ) + +# conda: install miniforge only if `conda` is absent (see deploy-bootstrap.md) +conda env list | grep -q '^madenv ' || conda create -y -n madenv python=3.12 +conda activate madenv +pip install -e ./madengine +``` + +After cloning or switching, the chosen manifest's `dockerfile` and the model +`scripts`/`run.sh` resolve under the `MAD` checkout (later +`$MODEL_DIR`). A missing path stops the run with a report (likely a wrong +branch, uninitialized submodules, or a different layout) rather than a silent +search — ask the user for the correct branch. + +**Guard:** +- **If a clone, `git switch`, or submodule init fails** → stop, report, ask the user for the correct repo/branch; do not guess another branch. +- **If `pip install -e ./madengine` fails** → stop and report; a madengine bug is not patched here (see Responsibilities). +- **If the manifest's `dockerfile`/`run.sh` do not resolve under `$MODEL_DIR`** → stop and report (likely wrong branch or uninitialized submodules); ask the user. + +### Step 4 — rundir + mad.env + +```bash +cd "$WORKDIR" && mkdir -p rundir && cd rundir +# copy the mad.env template for the archetype (one file per archetype under +# assets/mad.env/; see references/cluster-types.md). The template stays unedited: +cp "$SKILL_DIR/assets/mad.env/mad.env..template" ./mad.env +# FILL every placeholder, then: +source mad.env +``` + +Every `` placeholder is resolved before `source`. The cluster-specific +values (`MAD_SYSTEM_GPU_ARCHITECTURE`, `NCCL_SOCKET_IFNAME`, +`NCCL_IB_GID_INDEX`, and the per-manifest `NCCL_IB_HCA` list) are confirmed +against the actual node — defaults in the template are archetype-typical, not +guaranteed. +Run `bash "$SKILL_DIR/scripts/detect_cluster_env.sh"` to propose them and use +the matrix in [references/cluster-types.md](references/cluster-types.md) to +interpret. + +**Guard:** +- **If a `` value is unknown** → ask the user; do not infer it from a prior session, the repo, or the web. +- **If `source mad.env` errors, or `MODEL_DIR`/`MAD_DOCKER_BUILDS` come back empty** → stop and report before continuing. + +### Step 5 — Manifest + +Copy the matching template from `assets/manifests/` into `rundir/`, then adapt. +That directory holds one `*.template.json` per workload/size; pick the one named +for the requested workload. + +```bash +cp "$SKILL_DIR/assets/manifests/.template.json" run_manifest_.json +# fill the manifest, then statically validate it (GPU-free; reads $MODEL_DIR): +bash "$SKILL_DIR/scripts/validate_manifest.sh" run_manifest_.json +``` + +`validate_manifest.sh` is a read-only static check: JSON validity, leftover +`` placeholders, `NCCL_IB_HCA` set and equal in both env blocks, +network-interface consistency, `slurm.nodes == distributed.nnodes` (+ nodelist +cardinality), no stray `MAD_SECRETS_HFTOKEN`, AINIC transport-var symmetry, and +(with `mad.env` sourced) that the `dockerfile`/`run.sh` resolve under +`$MODEL_DIR`. It exits non-zero on any FAIL. + +What the run sets per run (cluster-private, kept out of the templates): +`deployment_config.slurm.{partition,account,qos,reservation,nodelist,exclude,nodes}`, +node count consistency (`slurm.nodes` == `distributed.nnodes`), the +`NCCL_IB_HCA` device list, and host paths. Field-by-field guidance: +[references/manifests.md](references/manifests.md). + +The manifest's `dockerfile` and the model `scripts`/`run.sh` resolve under +`$MODEL_DIR` (`[ -f "$MODEL_DIR/" ] && [ -f "$MODEL_DIR/" ]`). +A missing path stops the run with a report rather than a silent search. + +**Guard:** +- **If no template matches the requested engine** → ask the user; do not invent a manifest. +- **If `scripts/validate_manifest.sh` reports any FAIL** → fix it before launching (it covers JSON validity, leftover placeholders, `NCCL_IB_HCA`, iface consistency, `nodes == nnodes`, a stray HF token, and dockerfile/run.sh resolution under `$MODEL_DIR`). + +### Step 6 — Launch + results + +```bash +cd "$WORKDIR/rundir" +source mad.env +# -o writes the aggregated perf to a per-run CSV (instead of the default +# perf.csv), so each workload keeps its own result and parallel runs never +# clobber a shared file: +madengine run --manifest-file run_manifest_.json --live-output \ + -o perf_.csv +``` + +Pre-building or pulling the image is unnecessary. With `local_image: true`, +madengine ensures the image itself on each node: if it is not already present +locally it loads it from the `MAD_DOCKER_BUILDS` tar, otherwise builds it from +the manifest's `dockerfile` (and falls back to `docker pull`), then `docker +save`s it to `MAD_DOCKER_BUILDS` so workers load the same tar. Implications: +the first run is slower (it builds once on rank 0), `MAD_DOCKER_BUILDS` lives on +shared FS, and the manifest's `dockerfile` resolves (it lives in the cloned +`MAD`) or the `docker_image` tag is pullable. Details: +[references/launch-and-results.md](references/launch-and-results.md). + +Perf lands in the `-o` file (`perf_.csv`) and the per-model +`multiple_results` CSV. Multi-node aggregation and how to read the result are in +the same file. + +**Guard:** +- **If `madengine run` fails** → read the captured `.cursor.logs/run/...` log and triage with [references/launch-and-results.md](references/launch-and-results.md). +- **If the cause is config/input the skill owns** (mad.env, manifest, paths, transport vars) → fix it in `rundir/` and re-run. +- **If the cause is madengine or MAD code** → stop and report to the user; do not patch those repos unless asked (see Responsibilities). + +## Gotchas + +Cross-cutting and per-workload pitfalls — mad.env sourcing, `MAD_DOCKER_BUILDS` +on shared FS, HF-token handling, `docker_mounts` direction, `NCCL_IB_HCA` +per-cluster, AINIC transport vars, perf-CSV aggregation, and per-workload notes +(sglang_disagg and primus_megatron training) — are in +[references/gotchas.md](references/gotchas.md), read before a run. + +## Examples + +Sanitized end-to-end walkthroughs live in [examples/](examples/); the run +templates they build on live in `assets/` (`assets/mad.env/`, +`assets/manifests/`). New workloads and clusters are added there, not inlined +here. + + + diff --git a/.claude/skills/mad-slurm-multinode/assets/mad.env/mad.env.amd-ainic.template b/.claude/skills/mad-slurm-multinode/assets/mad.env/mad.env.amd-ainic.template new file mode 100644 index 0000000..49c2f97 --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/assets/mad.env/mad.env.amd-ainic.template @@ -0,0 +1,75 @@ +# mad.env — AMD-AINIC / Pollara archetype (AMD AINIC rdma* HCAs, gfx950) +# Copy to /rundir/mad.env, resolve every , then `source mad.env`. +# Do NOT commit a filled copy into this skill — it carries cluster-private paths. +# +# AINIC specifics vs CX7: gfx950, eno0, GID 1, ionic drivers, RCCL_AINIC_ROCE. +# Confirm against the node (see references/cluster-types.md + +# scripts/detect_cluster_env.sh). + +# --- secrets (file-based, never inline a literal hf_... token) --- +# This file is meant to be `source`d, so we warn (not exit) on a missing/empty +# token: an exit here would kill the caller's shell instead of just this file. +if [[ ! -s ~/.huggingface/token ]]; then + echo "WARNING: ~/.huggingface/token is missing or empty; MAD_SECRETS_HFTOKEN" >&2 + echo " will be empty and HF downloads will fail later. Run" >&2 + echo " 'huggingface-cli login' or place a token at that path." >&2 +fi +export MAD_SECRETS_HFTOKEN=$(cat ~/.huggingface/token 2>/dev/null) + +# --- system --- +export MAD_SYSTEM_GPU_ARCHITECTURE=gfx950 # AINIC/MI355-class default; confirm on node +export MAD_VERBOSE_CONFIG=true + +# --- MAD source (point at the MAD checkout) --- +export MAD_SETUP_MODEL_DIR=false +export MODEL_DIR=/ # e.g. /MAD/ + +# --- shared dir convenience (set once, reuse below) --- +SHARE_DIR= # shared FS visible to every node, e.g. /shared/ + +# --- data / cache roots --- +export MAD_DATAHOME=$SHARE_DIR/data/models +export HF_HOME=$SHARE_DIR/data/cache/huggingface +export TORCH_HOME=$SHARE_DIR/data/cache/torch +export XDG_CACHE_HOME=$SHARE_DIR/data/cache/xdg +export PIP_CACHE_DIR=$SHARE_DIR/data/cache/pip +export TRITON_CACHE_DIR=$SHARE_DIR/data/cache/triton + +# --- image tar cache: MUST be on shared FS visible to every compute node --- +export MAD_DOCKER_BUILDS=$SHARE_DIR/mad_docker_builds + +# --- HF helper aliases --- +export TRANSFORMERS_CACHE=$HF_HOME +export HUGGINGFACE_HUB_CACHE=$HF_HOME/hub + +# --- MAD metadata --- +export MAD_DEPLOYMENT_TYPE=slurm +export BUILD_NUMBER=${BUILD_NUMBER:-0} + +# --- AINIC RoCE / NCCL transport (Pollara, vendor 0x1dd8, exposed as rdma0..7) --- +# NO mlx5 devices. NCCL_IB_HCA (the rdma* list) is set PER-MANIFEST. +export NCCL_IB_DISABLE=0 +export NCCL_SOCKET_IFNAME=eno0 # management iface; confirm with `ip link` +export GLOO_SOCKET_IFNAME=eno0 +export NCCL_IB_GID_INDEX=1 # RoCEv2 GID on AINIC; confirm with show_gids + +# --- AINIC transport selection (also required in BOTH manifest env blocks) --- +# Without these RCCL silently falls back to verbs/sockets and certifies the +# wrong codepath. +export RCCL_AINIC_ROCE=1 +export RDMAV_DRIVERS=ionic +export IBV_DRIVERS=ionic + +# --- AINIC driver mounts (manifest-level, documented here for reference) --- +# The ionic verbs driver lives on the host and must be bind-mounted into every +# container. Add to built_models.additional_docker_run_options: +# -v /usr/lib/x86_64-linux-gnu/libionic.so:/usr/lib/x86_64-linux-gnu/libionic.so:ro +# -v /usr/lib/x86_64-linux-gnu/libionic.so.1:/usr/lib/x86_64-linux-gnu/libionic.so.1:ro +# -v /etc/libibverbs.d:/etc/libibverbs.d:ro +# And to context.docker_mounts: +# "/etc/libibverbs.d": "/etc/libibverbs.d" +# "/usr/lib/x86_64-linux-gnu/libibverbs": "/usr/lib/x86_64-linux-gnu/libibverbs" +# See references/cluster-types.md "AINIC driver mounts" for the full rationale. + +# Workload-specific data paths are NOT set here — they live in the workload +# manifest, which keeps this file cluster/cache-only and reusable across workloads. diff --git a/.claude/skills/mad-slurm-multinode/assets/mad.env/mad.env.cx7-roce.template b/.claude/skills/mad-slurm-multinode/assets/mad.env/mad.env.cx7-roce.template new file mode 100644 index 0000000..8866266 --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/assets/mad.env/mad.env.cx7-roce.template @@ -0,0 +1,55 @@ +# mad.env — CX7 / Mellanox-RoCE archetype (Mellanox mlx5 HCAs, gfx942) +# Copy to /rundir/mad.env, resolve every , then `source mad.env`. +# Do NOT commit a filled copy into this skill — it carries cluster-private paths. +# +# Verify the cluster-specific values against the actual node (see +# references/cluster-types.md and scripts/detect_cluster_env.sh): +# MAD_SYSTEM_GPU_ARCHITECTURE, NCCL_SOCKET_IFNAME, NCCL_IB_GID_INDEX + +# --- secrets (file-based, never inline a literal hf_... token) --- +# This file is meant to be `source`d, so we warn (not exit) on a missing/empty +# token: an exit here would kill the caller's shell instead of just this file. +if [[ ! -s ~/.huggingface/token ]]; then + echo "WARNING: ~/.huggingface/token is missing or empty; MAD_SECRETS_HFTOKEN" >&2 + echo " will be empty and HF downloads will fail later. Run" >&2 + echo " 'huggingface-cli login' or place a token at that path." >&2 +fi +export MAD_SECRETS_HFTOKEN=$(cat ~/.huggingface/token 2>/dev/null) + +# --- system --- +export MAD_SYSTEM_GPU_ARCHITECTURE=gfx942 # CX7/MI300-class default; confirm on node +export MAD_VERBOSE_CONFIG=true + +# --- MAD source (point at the MAD checkout) --- +export MAD_SETUP_MODEL_DIR=false +export MODEL_DIR=/ # e.g. /MAD/ + +# --- data / cache roots: keep large artifacts off /home, on shared FS --- +export MAD_DATAHOME=/models +export HF_HOME=/cache/huggingface +export TORCH_HOME=/cache/torch +export XDG_CACHE_HOME=/cache/xdg +export PIP_CACHE_DIR=/cache/pip +export TRITON_CACHE_DIR=/cache/triton + +# --- image tar cache: MUST be on shared FS visible to every compute node --- +export MAD_DOCKER_BUILDS= + +# --- HF helper aliases --- +export TRANSFORMERS_CACHE=$HF_HOME +export HUGGINGFACE_HUB_CACHE=$HF_HOME/hub + +# --- MAD metadata --- +export MAD_DEPLOYMENT_TYPE=slurm +export BUILD_NUMBER=${BUILD_NUMBER:-0} + +# --- RoCE / NCCL transport (CX7: Mellanox mlx5 HCAs, RoCEv2 over rdma*) --- +# Bootstrap/control plane is TCP over the management iface; data plane is RDMA +# via NCCL_IB_HCA which is set PER-MANIFEST (device list differs per node). +export NCCL_IB_DISABLE=0 +export NCCL_SOCKET_IFNAME=eth0 # management iface; confirm with `ip link` +export GLOO_SOCKET_IFNAME=eth0 +export NCCL_IB_GID_INDEX=3 # RoCEv2 GID; confirm with show_gids + +# Workload-specific data paths are NOT set here — they live in the workload +# manifest, which keeps this file cluster/cache-only and reusable across workloads. diff --git a/.claude/skills/mad-slurm-multinode/assets/mad.env/mad.env.thor2-bnxt.template b/.claude/skills/mad-slurm-multinode/assets/mad.env/mad.env.thor2-bnxt.template new file mode 100644 index 0000000..92053d3 --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/assets/mad.env/mad.env.thor2-bnxt.template @@ -0,0 +1,57 @@ +# mad.env — Broadcom Thor2 / RoCE archetype (bnxt_re HCAs, gfx950 / MI355X-class) +# Copy to /rundir/mad.env, resolve every , then `source mad.env`. +# Do NOT commit a filled copy into this skill — it carries cluster-private paths. +# +# Verify the cluster-specific values against the actual node (see +# references/cluster-types.md and scripts/detect_cluster_env.sh): +# MAD_SYSTEM_GPU_ARCHITECTURE, NCCL_SOCKET_IFNAME, NCCL_IB_GID_INDEX + +# --- secrets (file-based, never inline a literal hf_... token) --- +# This file is meant to be `source`d, so we warn (not exit) on a missing/empty +# token: an exit here would kill the caller's shell instead of just this file. +if [[ ! -s ~/.huggingface/token ]]; then + echo "WARNING: ~/.huggingface/token is missing or empty; MAD_SECRETS_HFTOKEN" >&2 + echo " will be empty and HF downloads will fail later. Run" >&2 + echo " 'huggingface-cli login' or place a token at that path." >&2 +fi +export MAD_SECRETS_HFTOKEN=$(cat ~/.huggingface/token 2>/dev/null) + +# --- system --- +export MAD_SYSTEM_GPU_ARCHITECTURE=gfx950 # MI355X-class default; confirm on node (gfx942 = MI300X) +export MAD_VERBOSE_CONFIG=true + +# --- MAD source (point at the MAD checkout) --- +export MAD_SETUP_MODEL_DIR=false +export MODEL_DIR=/ # e.g. /MAD/ + +# --- data / cache roots: keep large artifacts off /home, on shared FS --- +export MAD_DATAHOME=/models +export HF_HOME=/cache/huggingface +export TORCH_HOME=/cache/torch +export XDG_CACHE_HOME=/cache/xdg +export PIP_CACHE_DIR=/cache/pip +export TRITON_CACHE_DIR=/cache/triton + +# --- image tar cache: MUST be on shared FS visible to every compute node --- +export MAD_DOCKER_BUILDS= + +# --- HF helper aliases --- +export TRANSFORMERS_CACHE=$HF_HOME +export HUGGINGFACE_HUB_CACHE=$HF_HOME/hub + +# --- MAD metadata --- +export MAD_DEPLOYMENT_TYPE=slurm +export BUILD_NUMBER=${BUILD_NUMBER:-0} + +# --- RoCE / NCCL transport (Broadcom Thor2 bnxt_re HCAs, RoCEv2) --- +# Bootstrap/control plane is TCP over the management iface (fenic0 here); data +# plane is RDMA via NCCL_IB_HCA which is set PER-MANIFEST (bnxt_re device list). +export NCCL_IB_DISABLE=0 +export NCCL_SOCKET_IFNAME=fenic0 # management iface; confirm with `ip -br link` +export GLOO_SOCKET_IFNAME=fenic0 +export NCCL_IB_GID_INDEX=3 # RoCEv2 GID; confirm with show_gids +export RDMAV_DRIVERS=bnxt_re # Broadcom Thor2 ibverbs provider +export IBV_DRIVERS=bnxt_re + +# Workload-specific data paths are NOT set here — they live in the workload +# manifest, which keeps this file cluster/cache-only and reusable across workloads. diff --git a/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-70b.template.json b/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-70b.template.json new file mode 100644 index 0000000..4c1aeac --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-70b.template.json @@ -0,0 +1,111 @@ +{ + "built_images": { + "rocm-primus-llama31-70b": { + "model": "primus_pyt_megatron_lm_train_llama-3.1-70b_scaleout", + "docker_image": "rocm/primus:v26.4->", + "dockerfile": "docker/primus_megatron_train_rccl_overlay.ubuntu.amd.Dockerfile", + "base_docker": "rocm/primus:v26.4", + "build_duration": 0, + "local_image": true, + "registry_image": null, + "registry": null, + "gpu_vendor": "AMD" + } + }, + "built_models": { + "rocm-primus-llama31-70b": { + "name": "primus_pyt_megatron_lm_train_llama-3.1-70b_scaleout", + "url": "", + "dockerfile": "docker/primus_megatron_train_rccl_overlay.ubuntu.amd.Dockerfile", + "scripts": "scripts/primus/megatron-lm/run.sh", + "n_gpus": "-1", + "owner": "mad.support@amd.com", + "training_precision": "", + "multiple_results": "perf_primus-megatron-Llama-3.1-70B.csv", + "tags": ["pyt", "pretrain", "llama-3.1-70b", "training"], + "timeout": -1, + "args": "--model_repo primus_pyt_megatron_lm_train_llama-3.1-70b", + "additional_docker_run_options": "--privileged --group-add render --shm-size 64G --device=/dev/infiniband --cap-add IPC_LOCK --ulimit memlock=-1 -v /sys:/sys:ro -v /run/udev:/run/udev:ro" + } + }, + "context": { + "docker_env_vars": { + "NCCL_DEBUG": "INFO", + "NCCL_DEBUG_SUBSYS": "INIT,NET", + "NCCL_IB_DISABLE": "0", + "NCCL_NET": "IB", + "NCCL_IB_HCA": "", + "NCCL_IB_GID_INDEX": "", + "NCCL_SOCKET_IFNAME": "", + "GLOO_SOCKET_IFNAME": "", + "RDMAV_DRIVERS": "", + "IBV_DRIVERS": "", + "RCCL_AINIC_ROCE": "", + "NCCL_NET_PLUGIN": "none", + "IBV_SHOW_WARNINGS": "1" + }, + "docker_mounts": { + "/dev/infiniband": "/dev/infiniband" + }, + "docker_build_arg": {}, + "gpu_vendor": "AMD", + "guest_os": "UBUNTU", + "docker_gpus": "0,1,2,3,4,5,6,7" + }, + "credentials_required": [], + "summary": { + "successful_builds": [], + "failed_builds": [], + "total_build_time": 0, + "successful_pushes": [], + "failed_pushes": [] + }, + "deployment_config": { + "target": "slurm", + "slurm": { + "partition": "", + "account": "", + "qos": "", + "nodes": 2, + "gpus_per_node": 8, + "time": "12:00:00", + "output_dir": "./slurm_output", + "exclusive": true, + "network_interface": "" + }, + "distributed": { + "launcher": "torchrun", + "backend": "nccl", + "port": 29500, + "nnodes": 2, + "nproc_per_node": 8 + }, + "env_vars": { + "HF_HOME": "${HF_HOME}", + "TORCH_HOME": "${TORCH_HOME}", + "XDG_CACHE_HOME": "${XDG_CACHE_HOME}", + "PIP_CACHE_DIR": "${PIP_CACHE_DIR}", + "MAD_DATAHOME": "${MAD_DATAHOME}", + "NCCL_DEBUG": "INFO", + "NCCL_DEBUG_SUBSYS": "INIT,NET", + "NCCL_IB_DISABLE": "0", + "NCCL_NET": "IB", + "NCCL_SOCKET_IFNAME": "", + "GLOO_SOCKET_IFNAME": "", + "NCCL_IB_GID_INDEX": "", + "NCCL_IB_HCA": "", + "RDMAV_DRIVERS": "", + "IBV_DRIVERS": "", + "RCCL_AINIC_ROCE": "", + "NCCL_NET_PLUGIN": "none", + "NCCL_TIMEOUT": "900", + "TORCH_NCCL_ASYNC_ERROR_HANDLING": "1", + "TORCH_NCCL_HIGH_PRIORITY": "1", + "OMP_NUM_THREADS": "8", + "MIOPEN_FIND_MODE": "1", + "MIOPEN_USER_DB_PATH": "${XDG_CACHE_HOME}/miopen" + }, + "debug": false, + "docker_gpus": "0,1,2,3,4,5,6,7" + } +} diff --git a/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-8b.template.json b/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-8b.template.json new file mode 100644 index 0000000..0d93b7a --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/assets/manifests/primus_llama-3.1-8b.template.json @@ -0,0 +1,111 @@ +{ + "built_images": { + "rocm-primus-llama31-8b": { + "model": "primus_pyt_megatron_lm_train_llama-3.1-8b_scaleout", + "docker_image": "rocm/primus:v26.4->", + "dockerfile": "docker/primus_megatron_train_rccl_overlay.ubuntu.amd.Dockerfile", + "base_docker": "rocm/primus:v26.4", + "build_duration": 0, + "local_image": true, + "registry_image": null, + "registry": null, + "gpu_vendor": "AMD" + } + }, + "built_models": { + "rocm-primus-llama31-8b": { + "name": "primus_pyt_megatron_lm_train_llama-3.1-8b_scaleout", + "url": "", + "dockerfile": "docker/primus_megatron_train_rccl_overlay.ubuntu.amd.Dockerfile", + "scripts": "scripts/primus/megatron-lm/run.sh", + "n_gpus": "-1", + "owner": "mad.support@amd.com", + "training_precision": "", + "multiple_results": "perf_primus-megatron-Llama-3.1-8B.csv", + "tags": ["pyt", "pretrain", "llama-3.1-8b", "training"], + "timeout": -1, + "args": "--model_repo primus_pyt_megatron_lm_train_llama-3.1-8b", + "additional_docker_run_options": "--privileged --group-add render --shm-size 64G --device=/dev/infiniband --cap-add IPC_LOCK --ulimit memlock=-1 -v /sys:/sys:ro -v /run/udev:/run/udev:ro" + } + }, + "context": { + "docker_env_vars": { + "NCCL_DEBUG": "INFO", + "NCCL_DEBUG_SUBSYS": "INIT,NET", + "NCCL_IB_DISABLE": "0", + "NCCL_NET": "IB", + "NCCL_IB_HCA": "", + "NCCL_IB_GID_INDEX": "", + "NCCL_SOCKET_IFNAME": "", + "GLOO_SOCKET_IFNAME": "", + "RDMAV_DRIVERS": "", + "IBV_DRIVERS": "", + "RCCL_AINIC_ROCE": "", + "NCCL_NET_PLUGIN": "none", + "IBV_SHOW_WARNINGS": "1" + }, + "docker_mounts": { + "/dev/infiniband": "/dev/infiniband" + }, + "docker_build_arg": {}, + "gpu_vendor": "AMD", + "guest_os": "UBUNTU", + "docker_gpus": "0,1,2,3,4,5,6,7" + }, + "credentials_required": [], + "summary": { + "successful_builds": [], + "failed_builds": [], + "total_build_time": 0, + "successful_pushes": [], + "failed_pushes": [] + }, + "deployment_config": { + "target": "slurm", + "slurm": { + "partition": "", + "account": "", + "qos": "", + "nodes": 2, + "gpus_per_node": 8, + "time": "12:00:00", + "output_dir": "./slurm_output", + "exclusive": true, + "network_interface": "" + }, + "distributed": { + "launcher": "torchrun", + "backend": "nccl", + "port": 29500, + "nnodes": 2, + "nproc_per_node": 8 + }, + "env_vars": { + "HF_HOME": "${HF_HOME}", + "TORCH_HOME": "${TORCH_HOME}", + "XDG_CACHE_HOME": "${XDG_CACHE_HOME}", + "PIP_CACHE_DIR": "${PIP_CACHE_DIR}", + "MAD_DATAHOME": "${MAD_DATAHOME}", + "NCCL_DEBUG": "INFO", + "NCCL_DEBUG_SUBSYS": "INIT,NET", + "NCCL_IB_DISABLE": "0", + "NCCL_NET": "IB", + "NCCL_SOCKET_IFNAME": "", + "GLOO_SOCKET_IFNAME": "", + "NCCL_IB_GID_INDEX": "", + "NCCL_IB_HCA": "", + "RDMAV_DRIVERS": "", + "IBV_DRIVERS": "", + "RCCL_AINIC_ROCE": "", + "NCCL_NET_PLUGIN": "none", + "NCCL_TIMEOUT": "900", + "TORCH_NCCL_ASYNC_ERROR_HANDLING": "1", + "TORCH_NCCL_HIGH_PRIORITY": "1", + "OMP_NUM_THREADS": "8", + "MIOPEN_FIND_MODE": "1", + "MIOPEN_USER_DB_PATH": "${XDG_CACHE_HOME}/miopen" + }, + "debug": false, + "docker_gpus": "0,1,2,3,4,5,6,7" + } +} diff --git a/.claude/skills/mad-slurm-multinode/assets/manifests/sglang_disagg_deepseek-r1.template.json b/.claude/skills/mad-slurm-multinode/assets/manifests/sglang_disagg_deepseek-r1.template.json new file mode 100644 index 0000000..63a2bb1 --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/assets/manifests/sglang_disagg_deepseek-r1.template.json @@ -0,0 +1,151 @@ +{ + "built_images": { + "sglang-disagg-deepseek-r1": { + "model": "sglang_disagg_deepseek-r1", + "docker_image": "smifix-mori--rixl--mooncake-gfx942 (append -rdma62 when built with --build-arg ENABLE_RDMA62=1)>", + "dockerfile": "docker/sglang_disagg_inference_full_overlay.ubuntu.amd.Dockerfile", + "base_docker": "lmsysorg/sglang:v0.5.12.post1-rocm720-mi30x", + "build_duration": 0, + "local_image": true, + "registry_image": null, + "registry": null, + "gpu_vendor": "AMD" + } + }, + "built_models": { + "sglang-disagg-deepseek-r1": { + "name": "sglang_disagg_deepseek-r1", + "url": "", + "dockerfile": "docker/sglang_disagg_inference_full_overlay.ubuntu.amd.Dockerfile", + "scripts": "scripts/sglang_disagg/run.sh", + "data": "", + "n_gpus": "-1", + "owner": "mad.support@amd.com", + "training_precision": "", + "multiple_results": "perf_sglang-disagg-DeepSeek-R1.csv", + "tags": ["pyt", "sglang", "disagg", "inference", "deepseek-r1"], + "timeout": 21600, + "args": "--model-name DeepSeek-R1 --model-path deepseek-ai/DeepSeek-R1-0528) and peers wait on MODEL_PATH/.stage_done> --run-mori 1 --dp-mode 1 --xp 2 --yd 2 --gpus-per-node 8 --kv-transfer-backend nixl", + "additional_docker_run_options": "--privileged --group-add render --shm-size 64G --device=/dev/infiniband --cap-add IPC_LOCK --ulimit memlock=-1 --ulimit nofile=1048576:1048576 -v /sys:/sys:ro -v /sys/class/infiniband:/sys/class/infiniband:ro -v /run/udev:/run/udev:ro -v : -v ${SLURM_SUBMIT_DIR}/slurm_output/run_logs:/run_logs -e SLURM_JOB_ID" + } + }, + "context": { + "skip_perf_collection": false, + "docker_env_vars": { + "MODEL_NAME": "DeepSeek-R1", + "MODEL_PATH": "", + "RUN_MORI": "1", + "DP_MODE": "1", + "xP": "2", + "yD": "2", + "GPUS_PER_NODE": "8", + "KV_TRANSFER_BACKEND": "nixl", + "MORI_JIT_CACHE_DIR": "/tmp/mori_jit", + "MORI_SHMEM_HEAP_SIZE": "17179869184", + "RDMA_CORE_CACHE": "", + "NCCL_DEBUG": "WARN", + "NCCL_DEBUG_SUBSYS": "INIT,NET", + "NCCL_IB_DISABLE": "0", + "NCCL_IB_HCA": "", + "IB_DEVICES": "", + "MORI_RDMA_DEVICES": "", + "NCCL_SOCKET_IFNAME": "", + "GLOO_SOCKET_IFNAME": "", + "MORI_SOCKET_IFNAME": "", + "NCCL_IB_GID_INDEX": "", + "MORI_IB_GID_INDEX": "", + "RDMAV_DRIVERS": "", + "IBV_DRIVERS": "", + "TORCH_NCCL_HIGH_PRIORITY": "1", + "GPU_MAX_HW_QUEUES": "2", + "TORCH_NCCL_ASYNC_ERROR_HANDLING": "1", + "NCCL_TIMEOUT": "900", + "HSA_ENABLE_SDMA": "0", + "HSA_FORCE_FINE_GRAIN_PCIE": "1", + "OMP_NUM_THREADS": "8", + "MIOPEN_FIND_MODE": "1", + "SMOKE": "0", + "BENCHMARK_ITR": "1", + "BENCHMARK_COMBINATIONS": "1024/1024 8192/1024", + "BENCHMARK_CONCURRENCY_LEVELS": "8 16 32 64", + "BENCHMARK_PRECHECK": "0", + "BENCHMARK_FAIL_FAST": "1" + }, + "docker_mounts": { + "/dev/infiniband": "/dev/infiniband", + "/sys/class/infiniband": "/sys/class/infiniband", + "": "" + }, + "docker_build_arg": {}, + "gpu_vendor": "AMD", + "guest_os": "UBUNTU", + "docker_gpus": "0,1,2,3,4,5,6,7" + }, + "credentials_required": [], + "summary": { + "successful_builds": [], + "failed_builds": [], + "total_build_time": 0, + "successful_pushes": [], + "failed_pushes": [] + }, + "deployment_config": { + "target": "slurm", + "slurm": { + "partition": "", + "account": "", + "qos": "", + "nodes": 4, + "gpus_per_node": 8, + "time": "02:30:00", + "output_dir": "./slurm_output", + "exclusive": true, + "network_interface": "" + }, + "distributed": { + "launcher": "sglang-disagg", + "backend": "nccl", + "port": 29500, + "nnodes": 4, + "nproc_per_node": 8, + "sglang_disagg": { + "prefill_nodes": 2, + "decode_nodes": 2 + } + }, + "env_vars": { + "MODEL_NAME": "DeepSeek-R1", + "MODEL_PATH": "", + "RUN_MORI": "1", + "DP_MODE": "1", + "xP": "2", + "yD": "2", + "GPUS_PER_NODE": "8", + "KV_TRANSFER_BACKEND": "nixl", + "MORI_JIT_CACHE_DIR": "/tmp/mori_jit", + "MORI_SHMEM_HEAP_SIZE": "17179869184", + "RDMA_CORE_CACHE": "", + "NCCL_DEBUG": "WARN", + "NCCL_IB_DISABLE": "0", + "NCCL_IB_HCA": "", + "IB_DEVICES": "", + "MORI_RDMA_DEVICES": "", + "NCCL_SOCKET_IFNAME": "", + "GLOO_SOCKET_IFNAME": "", + "MORI_SOCKET_IFNAME": "", + "NCCL_IB_GID_INDEX": "", + "MORI_IB_GID_INDEX": "", + "RDMAV_DRIVERS": "", + "IBV_DRIVERS": "", + "NCCL_TIMEOUT": "900", + "OMP_NUM_THREADS": "8", + "BENCHMARK_COMBINATIONS": "1024/1024 8192/1024", + "BENCHMARK_CONCURRENCY_LEVELS": "8 16 32 64", + "BENCHMARK_FAIL_FAST": "1" + }, + "monitor": true, + "debug": false, + "docker_gpus": "0,1,2,3,4,5,6,7", + "gpus_per_node": 8 + } +} diff --git a/.claude/skills/mad-slurm-multinode/examples/amd-ainic-walkthrough.md b/.claude/skills/mad-slurm-multinode/examples/amd-ainic-walkthrough.md new file mode 100644 index 0000000..880b1ba --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/examples/amd-ainic-walkthrough.md @@ -0,0 +1,93 @@ +# Example walkthrough — Primus Llama-3.1-8B on an AMD-AINIC / Pollara cluster + +> Sanitized end-to-end walkthrough: no real node names, reservations, or tokens. +> Replace every `<...>` with your values. + +Scenario: a 4-node Primus 8B run on a reserved AINIC partition. + +## 0. Inputs gathered from the requester + +| Input | Example value | +|--------------------|----------------------------------------| +| Compute/jump node | `` | +| `$WORKDIR` | `/run` | +| Archetype | AMD-AINIC / Pollara | +| Shared root | `` (e.g. /shared/) | +| `MAD_DOCKER_BUILDS`| `/mad_docker_builds` | +| SLURM | partition ``, reservation ``, nodelist `,,,` | +| Nodes | 4 | + +## 1. Preflight + bootstrap + +```bash +export SKILL_DIR="" +bash "$SKILL_DIR/scripts/preflight.sh" # expect gfx950, docker ok, sbatch ok +bash "$SKILL_DIR/scripts/detect_cluster_env.sh" # expect rdma0..7, ionic, GID ~1, eno0 +cd /run +[ -d MAD ] || git clone --recursive +( cd MAD && git switch "" && git submodule update --init --recursive ) +[ -d madengine ] || git clone --recursive +( cd madengine && git switch "" && git submodule update --init --recursive ) +conda env list | grep -q '^madenv ' || conda create -y -n madenv python=3.12 +conda activate madenv +pip install -e ./madengine +``` + +## 2. rundir + mad.env (from mad.env.amd-ainic.template) + +Filled values (illustrative): +```bash +export MAD_SYSTEM_GPU_ARCHITECTURE=gfx950 +SHARE_DIR= +export MODEL_DIR=$SHARE_DIR/MAD/ +export MAD_DOCKER_BUILDS=$SHARE_DIR/mad_docker_builds +export NCCL_SOCKET_IFNAME=eno0 +export NCCL_IB_GID_INDEX=1 +export RCCL_AINIC_ROCE=1 +export RDMAV_DRIVERS=ionic +export IBV_DRIVERS=ionic +``` +```bash +cd /run && mkdir -p rundir && cd rundir +cp "$SKILL_DIR/assets/mad.env/mad.env.amd-ainic.template" ./mad.env +# edit placeholders ... +source mad.env +``` + +## 3. Manifest (from primus_llama-3.1-8b.template.json) + +Edits applied to `run_manifest_primus_8b_ainic.json`: +- `slurm.partition=`, `reservation=`, + `nodelist=,,,`, `nodes=4` (remove `account` if unused) +- `distributed.nnodes=4` +- `NCCL_IB_HCA=rdma0:1,rdma1:1,rdma2:1,rdma3:1,rdma4:1,rdma5:1,rdma6:1,rdma7:1` + in BOTH env blocks +- `NCCL_SOCKET_IFNAME=eno0`, `GLOO_SOCKET_IFNAME=eno0`, `NCCL_IB_GID_INDEX=1` +- `RDMAV_DRIVERS=ionic`, `IBV_DRIVERS=ionic`, `RCCL_AINIC_ROCE=1` in BOTH blocks +- `slurm.network_interface=eno0` +- `docker_image=rocm/primus:v26.2-rccl--` + +For a custom RCCL drop you can also set `context.docker_build_arg` with +`BUILD_GPU_TARGETS=gfx950`, `RCCL_REPO`, `RCCL_COMMIT` — this rebuilds RCCL via +the `rccl_overlay` Dockerfile shipped in MAD. +```bash +python3 -m json.tool run_manifest_primus_8b_ainic.json >/dev/null && echo "valid" +``` + +## 4. Launch + read result + +```bash +source mad.env +madengine run --manifest-file run_manifest_primus_8b_ainic.json --live-output +squeue -u $USER +``` +Confirm the AINIC/ionic transport was selected in the rank-0 RCCL debug log +before trusting the number (otherwise it fell back to verbs/sockets — see +references/cluster-types.md). Aggregated perf in `rundir/perf.csv`. + +## Notes + +- 4 nodes minimum is recommended for AINIC to exercise multi-rail collectives. +- Hold the same node set across compared runs (`nodelist`) for apples-to-apples. +- Deep AINIC transport validation / RCCL build-vs-build deltas are a separate + workflow, out of scope for this perf-run skill. diff --git a/.claude/skills/mad-slurm-multinode/examples/cx7-roce-walkthrough.md b/.claude/skills/mad-slurm-multinode/examples/cx7-roce-walkthrough.md new file mode 100644 index 0000000..03effd2 --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/examples/cx7-roce-walkthrough.md @@ -0,0 +1,84 @@ +# Example walkthrough — Primus Llama-3.1-8B on a CX7 / Mellanox-RoCE cluster + +> Sanitized end-to-end walkthrough: no real node names, SLURM queues, accounts, +> or tokens. Replace every `<...>` with your values. + +Scenario: bring up a fresh CX7 login node and run a 2-node Primus 8B perf test. + +## 0. Inputs gathered from the requester + +| Input | Example value | +|--------------------|----------------------------------------| +| Compute/login node | `` | +| `$WORKDIR` | `~/source/run1` | +| Archetype | CX7 / Mellanox-RoCE | +| Data root | `` (shared FS) | +| `MAD_DOCKER_BUILDS`| `/mad_docker_builds` | +| SLURM | partition ``, account ``, qos `` | +| Nodes | 2 | + +## 1. Preflight + bootstrap + +```bash +export SKILL_DIR="" +bash "$SKILL_DIR/scripts/preflight.sh" # expect gfx942, docker ok, sbatch ok +cd ~/source/run1 +[ -d MAD ] || git clone --recursive +( cd MAD && git switch "" && git submodule update --init --recursive ) +[ -d madengine ] || git clone --recursive +( cd madengine && git switch "" && git submodule update --init --recursive ) +conda env list | grep -q '^madenv ' || conda create -y -n madenv python=3.12 +conda activate madenv +pip install -e ./madengine +``` + +## 2. rundir + mad.env (from mad.env.cx7-roce.template) + +Filled values (illustrative): +```bash +export MAD_SYSTEM_GPU_ARCHITECTURE=gfx942 +export MODEL_DIR=~/source/run1/MAD/ +export MAD_DATAHOME=/models +export MAD_DOCKER_BUILDS=/mad_docker_builds +export NCCL_SOCKET_IFNAME=eth0 +export NCCL_IB_GID_INDEX=3 +``` +```bash +cd ~/source/run1 && mkdir -p rundir && cd rundir +cp "$SKILL_DIR/assets/mad.env/mad.env.cx7-roce.template" ./mad.env +# edit placeholders ... +source mad.env +[ -d "$MODEL_DIR/scripts" ] && echo "MODEL_DIR ok" +``` + +## 3. Manifest (from primus_llama-3.1-8b.template.json) + +Edits applied to `run_manifest_primus_8b.json`: +- `slurm.partition=`, `account=`, `qos=`, `nodes=2` +- `distributed.nnodes=2` +- `NCCL_IB_HCA=mlx5_0:1,mlx5_1:1` in BOTH `context.docker_env_vars` and + `deployment_config.env_vars` +- `NCCL_SOCKET_IFNAME=eth0`, `GLOO_SOCKET_IFNAME=eth0`, `NCCL_IB_GID_INDEX=3` +- `RDMAV_DRIVERS=mlx5`, `IBV_DRIVERS=mlx5`, deleted the `RCCL_AINIC_ROCE` key +- `docker_image=rocm/primus:v26.2-rccl--` +```bash +python3 -m json.tool run_manifest_primus_8b.json >/dev/null && echo "manifest valid" +``` + +## 4. Launch + read result + +```bash +source mad.env +madengine run --manifest-file run_manifest_primus_8b.json --live-output +squeue -u $USER # watch the 2-node job +``` +Result lands in `rundir/perf.csv` (aggregated) and +`perf_primus-megatron-Llama-3.1-8B.csv`. Report tok/s/gpu and TFLOPS/gpu from +the aggregated CSV (node_0's local CSV can read empty — see +references/launch-and-results.md). + +## Notes + +- 70B: same flow with `primus_llama-3.1-70b.template.json`; bump walltime/nodes. +- Data staging (download + preprocessing) is handled by the MAD `run.sh`; the + skill just fills the workload's data mounts and points the env vars at them. diff --git a/.claude/skills/mad-slurm-multinode/examples/thor2-bnxt-walkthrough.md b/.claude/skills/mad-slurm-multinode/examples/thor2-bnxt-walkthrough.md new file mode 100644 index 0000000..b368243 --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/examples/thor2-bnxt-walkthrough.md @@ -0,0 +1,103 @@ +# Example walkthrough — Primus Llama-3.1-70B on a Broadcom-Thor2-RoCE cluster + +> Sanitized end-to-end walkthrough: no real node names, SLURM queues, accounts, +> or tokens. Replace every `<...>` with your values. + +Scenario: bring up a fresh Broadcom-Thor2-RoCE login node and run a 16-node Primus 70B perf +test over RoCEv2 on Broadcom Thor2 (`bnxt_re`) NICs. + +## 0. Inputs gathered from the requester + +| Input | Example value | +|--------------------|----------------------------------------| +| Compute/login node | `` | +| `$WORKDIR` | `` (shared NFS) | +| Archetype | Broadcom-Thor2-RoCE | +| Data root | `` (shared NFS) | +| `MAD_DOCKER_BUILDS`| `/mad_docker_builds` (shared NFS)| +| SLURM | partition `` (no account/qos) | +| Nodes | 16 (8 GPU/node) | + +## 1. Preflight + bootstrap + +```bash +export SKILL_DIR="" +bash "$SKILL_DIR/scripts/preflight.sh" # expect gfx950, docker ok, sbatch ok +cd +[ -d MAD ] || git clone --recursive +( cd MAD && git switch "" && git submodule update --init --recursive ) +[ -d madengine ] || git clone https://github.com/ROCm/madengine --recursive +( cd madengine && git switch develop && git submodule update --init --recursive ) # PR #142 merged +conda env list | grep -q '^madenv ' || conda create -y -n madenv python=3.12 +conda activate madenv +pip install -e ./madengine +``` + +Confirm the data-plane HCAs and GID on an allocated node before filling the env: +```bash +srun -p -N1 bash "$SKILL_DIR/scripts/detect_cluster_env.sh" +# expect: gfx950, bnxt_re0..7, mgmt iface fenic0, RoCEv2 GID 3 +``` + +## 2. rundir + mad.env (from mad.env.thor2-bnxt.template) + +Filled values (illustrative): +```bash +export MAD_SYSTEM_GPU_ARCHITECTURE=gfx950 +export MODEL_DIR=/MAD/ +export MAD_DATAHOME=/models +export MAD_DOCKER_BUILDS=/mad_docker_builds +export NCCL_SOCKET_IFNAME=fenic0 +export NCCL_IB_GID_INDEX=3 +export RDMAV_DRIVERS=bnxt_re +export IBV_DRIVERS=bnxt_re +``` +```bash +cd && mkdir -p rundir && cd rundir +cp "$SKILL_DIR/assets/mad.env/mad.env.thor2-bnxt.template" ./mad.env +# edit placeholders ... +source mad.env +[ -d "$MODEL_DIR/scripts" ] && echo "MODEL_DIR ok" +``` + +## 3. Manifest (from primus_llama-3.1-70b.template.json) + +Edits applied to `run_manifest_primus_70b.json`: +- `slurm.partition=`, removed `account`/`qos` keys (cluster has none), + `nodes=16`, and **added `"skip_gpus_directive": true`** (this cluster's SLURM rejects + the `--gpus-per-node` directive — see references/cluster-types.md) +- `distributed.nnodes=16` +- `NCCL_IB_HCA=bnxt_re0:1,bnxt_re1:1,bnxt_re2:1,bnxt_re3:1,bnxt_re4:1,bnxt_re5:1,bnxt_re6:1,bnxt_re7:1` + in BOTH `context.docker_env_vars` and `deployment_config.env_vars` +- `NCCL_SOCKET_IFNAME=fenic0`, `GLOO_SOCKET_IFNAME=fenic0`, `NCCL_IB_GID_INDEX=3` +- `RDMAV_DRIVERS=bnxt_re`, `IBV_DRIVERS=bnxt_re`, deleted the `RCCL_AINIC_ROCE` key +- added the bnxt_re tuning vars to both env blocks (QPS/inline/merge/adaptive/ + gdr/dmabuf — see references/cluster-types.md "Broadcom-Thor2 specifics") +- added `"/etc/libibverbs.d": "/etc/libibverbs.d"` to `context.docker_mounts` +- `docker_image=rocm/primus:v26.3-rccl--`, `base_docker=rocm/primus:v26.3` +```bash +bash "$SKILL_DIR/scripts/validate_manifest.sh" run_manifest_primus_70b.json +``` + +## 4. Launch + read result + +```bash +source mad.env +madengine run --manifest-file run_manifest_primus_70b.json --live-output \ + -o perf_primus_70b.csv +squeue -u $USER # watch the 16-node job (PR #142 auto-selects healthy nodes) +``` +Result lands in `rundir/perf_primus_70b.csv` (aggregated) and +`perf_primus-megatron-Llama-3.1-70B.csv`. Report tok/s/gpu and TFLOPS/gpu from +the aggregated CSV (a worker node's local CSV can read header-only — see +references/launch-and-results.md). + +## Notes + +- PR #142 (merged in madengine `develop`) runs a GPU health check and + auto-selects healthy nodes; leave `nodelist` unset to use it, or pin nodes by + filling `slurm.nodelist` (cardinality must equal `nodes`). +- With `local_image: true` + `MAD_DOCKER_BUILDS` on shared NFS, rank 0 saves the + image tar once and every worker loads it — no manual `docker load` per node. +- The Primus 70B recipe may sweep BF16+FP8; a BF16 cold-start flake can mark the + run FAILURE and skip perf collection. Re-run for a clean perf CSV. diff --git a/.claude/skills/mad-slurm-multinode/references/cluster-types.md b/.claude/skills/mad-slurm-multinode/references/cluster-types.md new file mode 100644 index 0000000..a819d40 --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/references/cluster-types.md @@ -0,0 +1,125 @@ +# Cluster archetypes and the env vars that differ + +Keywords: CX7 Mellanox mlx5 RoCE, AMD AINIC Pollara ionic rdma, Broadcom +Thor2 bnxt_re RoCE, gfx942 gfx950, NCCL_IB_HCA, NCCL_SOCKET_IFNAME, +NCCL_IB_GID_INDEX, RCCL_AINIC_ROCE, RDMAV_DRIVERS, IBV_DRIVERS, eth0 eno0 fenic0, +show_gids, ibv_devices, skip_gpus_directive + +Three archetypes seen in practice. The differences below decide whether the run +exercises the intended RDMA transport. Values are archetype-*typical* and are +confirmed on the actual node (`scripts/detect_cluster_env.sh`). + +## Difference matrix + +| Variable | CX7 / Mellanox-RoCE | AMD-AINIC / Pollara | Broadcom-Thor2-RoCE | +|------------------------------|--------------------------------------|----------------------------------------|----------------------------------------| +| `MAD_SYSTEM_GPU_ARCHITECTURE`| `gfx942` | `gfx950` | `gfx950` (MI355X; `gfx942` = MI300X) | +| HCA hardware | Mellanox CX7 `mlx5_*` | AMD AINIC `rdma0..7` (vendor 0x1dd8) | Broadcom Thor2 `bnxt_re0..7` | +| `NCCL_IB_HCA` | `mlx5_0:1,mlx5_1:1` (or 8-rail mlx5 list) | `rdma0:1,rdma1:1,...,rdma7:1` | `bnxt_re0:1,bnxt_re1:1,...,bnxt_re7:1` | +| `NCCL_SOCKET_IFNAME` | `eth0` | `eno0` | `fenic0` | +| `GLOO_SOCKET_IFNAME` | `eth0` | `eno0` | `fenic0` | +| `NCCL_IB_GID_INDEX` | `3` | `1` | `3` | +| `RDMAV_DRIVERS` / `IBV_DRIVERS` | `mlx5` | `ionic` | `bnxt_re` | +| `RCCL_AINIC_ROCE` | (unset) | `1` (required) | (unset) | +| Shared FS paths | shared inference FS + off-/home scratch | shared team FS (mounted on all nodes) | shared NFS work/data root on all nodes | +| SLURM selection | `partition`+`account`+`qos`, sometimes `exclude` | usually `partition`+`reservation`+`nodelist` | `partition`, `exclude`/`nodelist`; `skip_gpus_directive` (see below) | +| AINIC driver mounts | (not needed) | add to every manifest's `built_models.additional_docker_run_options` and `context.docker_mounts` — see **AINIC driver mounts** below | (not needed) | + +## How to confirm each value on the node + +These values live on the node, so the probe runs there before the user is +asked. The login/jump node and the compute nodes can differ, so compute-node +values come from running `scripts/detect_cluster_env.sh` through `srun` on an +allocated node (`srun -p [--reservation ] [--nodelist ] +-N1 bash scripts/detect_cluster_env.sh`); only the cluster-private selectors +(partition / account / qos / reservation / nodelist) come from the user. + + +- GPU arch: `rocm-smi --showhw` or `rocminfo | grep gfx` -> the `gfx9xx` target. +- IB/RDMA HCAs: `ibv_devices` (lists `mlx5_*` or `rdma*`). The `NCCL_IB_HCA` + list should enumerate the GPU-attached HCAs (`:1` = port 1). +- Management iface: `ip -br link` / `ip -br addr` -> the routable host iface + used for bootstrap (`eth0` vs `eno0`). This is not the data-plane device. +- RoCEv2 GID index: `show_gids` (or read + `/sys/class/infiniband//ports/1/gid_attrs/types/*`); pick the RoCEv2 + (v2) GID index. CX7 commonly 3, AINIC commonly 1. +- Driver: if `ibv_devices` shows `mlx5_*` -> `mlx5`; if `rdma*` (ionic) -> + `ionic` and you are on AINIC (set `RCCL_AINIC_ROCE=1`). + +## AINIC driver mounts + +The ionic verbs driver lives on the host (`/usr/lib/x86_64-linux-gnu/libionic.so*`, +`/etc/libibverbs.d/`) and must be bind-mounted into every container so that +ibverbs inside the container can find and load the ionic provider. + +**`built_models.additional_docker_run_options`** — add these `-v` flags: +``` +-v /usr/lib/x86_64-linux-gnu/libionic.so:/usr/lib/x86_64-linux-gnu/libionic.so:ro +-v /usr/lib/x86_64-linux-gnu/libionic.so.1:/usr/lib/x86_64-linux-gnu/libionic.so.1:ro +-v /etc/libibverbs.d:/etc/libibverbs.d:ro +``` + +**`context.docker_mounts`** — add: +```json +"/etc/libibverbs.d": "/etc/libibverbs.d", +"/usr/lib/x86_64-linux-gnu/libibverbs": "/usr/lib/x86_64-linux-gnu/libibverbs" +``` + +Without these mounts `libibverbs` inside the container finds no provider, reports +0 IB devices, and RCCL silently falls back to TCP sockets. + +## Broadcom-Thor2 specifics + +Broadcom-Thor2 nodes carry Broadcom Thor2 NICs exposed as `bnxt_re0..7` ibverbs devices +over RoCEv2. Two things differ from the CX7/AINIC archetypes: + +**1. `skip_gpus_directive` in the manifest's `deployment_config.slurm`.** This cluster's +SLURM rejects the generated `--gpus-per-node` directive (the partition does not +advertise GPU GRES the way the default sbatch template assumes), so the job is +refused before launch. Set `"skip_gpus_directive": true` so madengine emits the +sbatch script without that directive and relies on `exclusive`/`nproc_per_node` +instead. Symptom when missing: sbatch is rejected with an invalid `--gpus`/GRES +error and no job id is returned. + +**2. bnxt_re transport tuning (set in BOTH manifest env blocks).** The Broadcom +provider needs a conservative QP/feature profile to run RoCEv2 reliably; the +defaults can hang or fall back. Validated set: + +```json +"NCCL_IB_QPS_PER_CONNECTION": "1", +"NCCL_IB_USE_INLINE": "0", +"NCCL_IB_MERGE_NICS": "0", +"NCCL_IB_SPLIT_DATA_ON_QPS": "0", +"NCCL_IB_ADAPTIVE_ROUTING": "0", +"NCCL_GDR_FLUSH_DISABLE": "1", +"NCCL_DMABUF_ENABLE": "0", +"NCCL_IB_ROCE_VERSION_NUM": "2" +``` + +Also point ibverbs at the host provider inside the container (the image carries +the bnxt_re provider, but the host `/etc/libibverbs.d` is mounted to be safe): +`LIBIBVERBS_DRIVER_PATH=/usr/lib/x86_64-linux-gnu/libibverbs` and add +`"/etc/libibverbs.d": "/etc/libibverbs.d"` to `context.docker_mounts`. Confirm +the device list with `ibv_devices` (expect `bnxt_re*`) and the RoCEv2 GID with +`show_gids` (commonly index 3). + +## MIOpen first-run compile time on gfx950 (AINIC) + +On gfx950 (MI355X) MIOpen compiles GPU kernels on first use, which can take +**20–40 minutes**. This affects any workload that starts MIOpen (inference, training). + +- Set `BARRIER_TIMEOUT_S` high enough (≥ `7200`) for the first run. +- Mount a **persistent** MIOpen cache dir so subsequent runs skip compilation: + - env var: `MIOPEN_USER_DB_PATH` → in-container path (e.g. `/miopen_cache`) + - mount: container `/miopen_cache` → host shared-FS dir +- Once the cache is warm, `BARRIER_TIMEOUT_S` can be reduced to a few minutes. + +## Why this matters + +If `NCCL_IB_HCA` names devices that do not exist on the node, RCCL/NCCL +initializes zero NICs and either aborts (`Failed to initialize any NET +plugin`) or silently falls back to TCP sockets — making the perf number a +measurement of the wrong path. On AINIC, omitting `RCCL_AINIC_ROCE=1` / +`RDMAV_DRIVERS=ionic` has the same effect (fallback to verbs/sockets). These +appear in BOTH `context.docker_env_vars` and `deployment_config.env_vars` of the +manifest (and in `mad.env` host env as a belt-and-suspenders default). diff --git a/.claude/skills/mad-slurm-multinode/references/deploy-bootstrap.md b/.claude/skills/mad-slurm-multinode/references/deploy-bootstrap.md new file mode 100644 index 0000000..25bdf7e --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/references/deploy-bootstrap.md @@ -0,0 +1,143 @@ +# Deploy bootstrap (Steps 0-4) + +Keywords: clone MAD madengine, miniforge conda madenv, pip install -e, +rundir, source mad.env, idempotent setup, git switch, python 3.12, docker check + +Detailed, idempotent bring-up of a fresh SLURM node. Run each step only if its +result is missing — re-running a completed step stays safe. Substitute +`$WORKDIR` with the user-provided working directory. `$SKILL_DIR` is the +absolute path to this skill's directory (the folder with `SKILL.md`); export it +once because the steps below `cd` into `$WORKDIR`, so bare `scripts/...` paths +would not resolve. + +## Step 0 — Preflight + +```bash +bash "$SKILL_DIR/scripts/preflight.sh" +``` + +Hard requirements (FAIL -> stop, this is not a valid target node): +- `docker` present and `docker info` succeeds (daemon reachable, user in group). +- A SLURM client: `sinfo` and `sbatch` on PATH. +- `git` (needed to clone/switch the repos in Step 1). `preflight.sh` treats a + missing `git` as a hard FAIL. + +Soft requirements (warn, can be installed in later steps): +- Python >= 3.10 (madengine needs 3.10+; we create a 3.12 conda env anyway). +- conda/miniforge (installed in Step 2 if absent). +- A GPU SMI (`rocm-smi` for AMD, `nvidia-smi` for NVIDIA) for arch detection. +- HF token at `~/.huggingface/token` (required at run time for gated Llama-3.1). + +## Step 1 — Clone repos (idempotent) + +```bash +cd "$WORKDIR" + +# MAD source -> MODEL_DIR. Repo URL is confirmed with the user; clone if missing. +if [ ! -d MAD ]; then + git clone --recursive +fi +# Branch is ASKED (no default). Switch, then sync submodules. +( cd MAD && git switch "" && git submodule update --init --recursive ) + +# madengine. Repo URL is confirmed with the user; clone if missing. +if [ ! -d madengine ]; then + git clone --recursive +fi +# Branch is ASKED (no default). +( cd madengine && git switch "" && git submodule update --init --recursive ) +``` + +Confirm the repo and branch with the user before cloning or switching. No branch +name is assumed (not a default, not one inferred from a prior session); the +branch is asked per repo and then `git switch`ed to. When the dirs already +exist, the skill shows the current branch +(`git -C branch --show-current`) and asks before changing it rather than +switching silently. An existing checkout is left intact, since the user may have +local edits. + +After cloning or switching, the assets the chosen manifest references resolve +under the `MAD` checkout (this becomes `$MODEL_DIR`). Confirm the +manifest's `dockerfile` and the model `scripts`/`run.sh` exist: + +```bash +# example for primus_llama-3.1-8b — adjust to the chosen manifest's fields +( cd MAD \ + && [ -f docker/primus_megatron_train_rccl_overlay.ubuntu.amd.Dockerfile ] \ + && [ -f scripts/primus/megatron-lm/run.sh ] \ + || echo "STOP: manifest dockerfile/run.sh not found under MAD" ) +``` + +A missing path stops the run with a report rather than a silent search. The +usual causes are a wrong branch, uninitialized submodules, or a different repo +layout, so the next step is to ask the user for the correct branch. + +## Step 2 — conda env (idempotent) + +```bash +# an existing conda/miniforge may just be missing from PATH — source it first +if ! command -v conda >/dev/null 2>&1; then + for cp in "$HOME/miniforge3" "$HOME/miniconda3" "$HOME/mambaforge" \ + "/opt/conda" "${CONDA_PREFIX:-}" "${MAMBA_ROOT_PREFIX:-}"; do + if [ -n "$cp" ] && [ -f "$cp/etc/profile.d/conda.sh" ]; then + source "$cp/etc/profile.d/conda.sh"; break + fi + done +fi + +# install miniforge only if conda is still missing +if ! command -v conda >/dev/null 2>&1; then + wget -q https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-x86_64.sh + bash Miniforge3-Linux-x86_64.sh -b -p "$HOME/miniforge3" + source "$HOME/miniforge3/etc/profile.d/conda.sh" +fi + +conda env list | grep -q '^madenv ' || conda create -y -n madenv python=3.12 +conda activate madenv +``` + +Notes: +- An already-installed conda/miniforge that is simply off PATH gets sourced from + its common install prefix (`$HOME/miniforge3`, `$HOME/miniconda3`, + `$HOME/mambaforge`, `/opt/conda`, `$CONDA_PREFIX`, `$MAMBA_ROOT_PREFIX`) + rather than reinstalled, which avoids duplicate installs. `preflight.sh` + reports the same find. +- `-b` runs the miniforge installer unattended; `-p` sets the prefix. +- If conda was just installed in this shell, `source .../conda.sh` (or open a + new shell) before `conda activate`. + +## Step 3 — Install madengine + +```bash +cd "$WORKDIR" +pip install -e ./madengine +madengine --help >/dev/null && echo "madengine OK" +``` + +Editable install so the user's branch changes are picked up without reinstall. + +## Step 4 — rundir + mad.env + +```bash +cd "$WORKDIR" +mkdir -p rundir && cd rundir + +# copy the archetype template (see cluster-types.md to choose): +# CX7/Mellanox-RoCE -> mad.env.cx7-roce.template +# AMD-AINIC/Pollara -> mad.env.amd-ainic.template +cp "$SKILL_DIR/assets/mad.env/mad.env..template" ./mad.env +# resolve every , confirm node-specific values, then: +source mad.env +``` + +After `source mad.env`, sanity-check the critical exports in the SAME shell: + +```bash +echo "MODEL_DIR=$MODEL_DIR" +echo "MAD_DOCKER_BUILDS=$MAD_DOCKER_BUILDS" +echo "MAD_SYSTEM_GPU_ARCHITECTURE=$MAD_SYSTEM_GPU_ARCHITECTURE" +[ -d "$MODEL_DIR/scripts" ] || echo "WARNING: MODEL_DIR/scripts missing — run.sh will not be found" +``` + +The HF token file (`~/.huggingface/token`) exists and is valid — Llama-3.1 +is gated, so a missing file is created before sourcing. diff --git a/.claude/skills/mad-slurm-multinode/references/gotchas.md b/.claude/skills/mad-slurm-multinode/references/gotchas.md new file mode 100644 index 0000000..5601676 --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/references/gotchas.md @@ -0,0 +1,283 @@ +# Gotchas — cross-cutting and per-workload + +Keywords: source mad.env MODEL_DIR run-script, MAD_DOCKER_BUILDS shared storage, +MAD_SECRETS_HFTOKEN HF 401 single-quoted, docker_mounts container_path host_path -v, +run_logs shared NFS mount, SLURM_SUBMIT_DIR sbatch cwd rundir, WORKDIR not +exported shlex.quote docker_mounts additional_docker_run_options, +RCCL_AINIC_ROCE RDMAV_DRIVERS ionic, +NCCL_IB_HCA mlx5 rdma per-cluster, perf.csv login-node aggregation, +slurm.nodes distributed.nnodes nodelist world size, +heterogeneous nodes NCCL_SOCKET_IFNAME GLOO_SOCKET_IFNAME network_interface, +routable interface eth0 eth1 IPv6 link-local fe80 gloo connect timeout subnet + +Cross-cutting and per-workload pitfalls observed in real runs. SKILL.md links +here; this file is read before a run. + +## Cross-cutting + +- **`source mad.env` in the same shell before `madengine run`.** `MODEL_DIR` is + exported there; without it the container's `scripts/` directory resolves empty + and the run dies with a missing run-script error. +- **`MAD_DOCKER_BUILDS` lives on shared storage** visible to every compute + node. The login/build node saves the image tar there; workers load it from + the same path. A node-local path makes workers rebuild (or fail). +- **The HF token lives in `mad.env`, not in the manifest** — a stray + `MAD_SECRETS_HFTOKEN` key in the manifest causes HF 401s (madengine renders it + single-quoted, so the container gets the literal `${MAD_SECRETS_HFTOKEN}`). + Full explanation in [references/manifests.md](manifests.md) + ("Secrets (HF token)"). +- **`docker_mounts` format is `{container_path: host_path}`.** madengine + renders each entry as `-v :`. Swapping the two + sides (a tempting mistake when the paths differ) hands Docker a non-existent + host path, which it silently creates as an empty directory, and the workload + fails with a "path does not exist" error. The direction is worth re-checking + whenever a host data directory maps to a container-internal path that differs + from the host path. See + [references/manifests.md](manifests.md) ("docker_mounts direction"). +- **Prefer `${SLURM_SUBMIT_DIR}` over `$WORKDIR`/relative paths in + `additional_docker_run_options` host mounts — it is a real, always-set SLURM + var that already equals the shared `rundir`.** `$WORKDIR` is a doc-only + convention (nothing in mad.env/madengine actually exports it), and a + relative host path (e.g. `./slurm_output/run_logs`) resolves against + whatever CWD the per-node `madengine run` process happens to have. SLURM + itself sets `SLURM_SUBMIT_DIR` to the directory `sbatch` was invoked from + (`deployment/slurm.py` runs `sbatch` without a `cwd=` override, and Step 6 + always launches from `$WORKDIR/rundir`), and that value is inherited + end-to-end — through the generated job script, `srun` (no `--export` + restriction), and the final `docker run` (`additional_docker_run_options` + is concatenated unquoted and `console.sh()` runs with `env=None`, i.e. full + inherited environment) — so `${SLURM_SUBMIT_DIR}` reliably expands to + `rundir` at every node with zero manual filling. This only works for + `additional_docker_run_options`, NOT `docker_mounts`: madengine renders + `docker_mounts` values through `shlex.quote()`, which single-quotes the + whole string and blocks `$VAR`/`${VAR}` expansion outright — don't duplicate + a shared-path mount there. The `sglang_disagg_deepseek-r1` template's + `/run_logs` mount (`-v ${SLURM_SUBMIT_DIR}/slurm_output/run_logs:/run_logs`) + is the reference example: `scripts/sglang_disagg/run.sh` requires `/run_logs` + and fails fast (no `/tmp` fallback) if it is missing or not writable, but + only checks writability, not actual cross-node sharedness — a node-local dir + that happens to exist would pass that check and silently break the + cross-node readiness rendezvous (see the `sglang_disagg` section below). +- **`NCCL_IB_HCA` is per-cluster, not portable.** CX7 uses `mlx5_*`; AINIC + uses `rdma0..7`. Copying an mlx5 list onto an AINIC node (or vice versa) + inits zero NICs. Verify on the node. +- **AINIC transport vars appear in BOTH manifest blocks.** + `RCCL_AINIC_ROCE=1`, `RDMAV_DRIVERS=ionic`, `IBV_DRIVERS=ionic` go in + `context.docker_env_vars` AND `deployment_config.env_vars`. Missing from + either, RCCL falls back to verbs/sockets and the run measures the + wrong transport. +- **Multi-node perf CSV: trust the login-node aggregation.** Throughput is + often printed only on the last global rank, so node-0's local CSV can look + empty even on a healthy run. Read the aggregated `perf.csv` in `rundir`, + not a single node's file. +- **`slurm.nodes` equals `distributed.nnodes`** and matches `--nodelist` + cardinality, or sbatch/torchrun disagree on world size. +- **Node environments can be heterogeneous across a cluster — don't trust a + single detect probe.** The interface that carries the routable control-plane + IP is not guaranteed to have the same name on every node (e.g. one node routes + on `eth0` while another routes on `eth1`), and a given named interface may hold + only a non-routable IPv6 link-local address (`fe80::...`) on some nodes. If you + pin `NCCL_SOCKET_IFNAME`/`GLOO_SOCKET_IFNAME`/`network_interface` by name based + on one probe node, the bootstrap can silently fail on the actually-allocated + nodes. Typical symptom: torchrun rendezvous and `world_size` form fine (that + goes over the hostname's routable IP), but `initializing torch distributed` + hangs and gloo times out connecting peers over `fe80::` link-local addresses. + Mitigation: verify the routable interface on the *allocated* nodes (not just + the probe node), and prefer selecting the interface by its routable subnet + rather than hard-coding an interface name. This is an environment + (cluster-provisioning) inconsistency, not a workload bug — flag it to the + cluster owner if a uniform-environment guarantee is expected. + +## sglang_disagg + +Keywords: launcher sglang-disagg run.sh xP yD RUN_MORI DP_MODE KV_TRANSFER_BACKEND +nixl mooncake, detokenizer hang health check No response from detokenizer, +MoRI overlay #366 inter-node decode freeze, RCCL overlay rsmi_init libtorch_hip +undefined symbol torch broken, rocm720 base librocm_smi64, patchelf add-needed +DT_NEEDED smifix, single full-overlay Dockerfile no base chaining, ENABLE_RDMA62 +rdma-core v62 optional stage, mooncake baked launcher build-layer removed runtime pip, +self-discover node IPs rendezvous IP_SYNC_TIMEOUT SGLANG_NODE_IPS not forwarded, +per-node docker load shared tar, perf CSV rank0 BENCHMARK_FAIL_FAST, circuit +breaker prefill workers Service Unavailable BENCHMARK_POINT_RETRIES transient +sweep retry. + +- **The launcher is `sglang-disagg` with `scripts/sglang_disagg/run.sh` as the + entrypoint** (PR #142 native launcher). `run.sh` reads topology from + `--xp/--yd` (or `xP`/`yD` env) and `--kv-transfer-backend` (`mori`/`mooncake`/ + `nixl`), then always execs the unified `sglang_disagg_mori_io_ep.sh` (it + branches internally on `KV_TRANSFER_BACKEND`; the older Mooncake/NIXL-only + `sglang_disagg_server.sh` was retired as a functional subset). There is no + `slurm_multi` wrapper — one `madengine run` is launched per node and the + roles (proxy / prefill x xP / decode x yD) are derived from `NODE_RANK` + the + ordered IP list. Do not reintroduce a `*_mn` slurm wrapper; it duplicates + topology logic and drifts from `run.sh`. + +- **A newer RCCL build drops the `rocm_smi` `DT_NEEDED` and breaks + `import torch` — re-add it with `patchelf`.** The RCCL stage of the single + full-overlay Dockerfile + (`docker/sglang_disagg_inference_full_overlay.ubuntu.amd.Dockerfile` — base + SGLang + RCCL + MoRI + NIXL/Mooncake KV-transfer + Mooncake pip in one file, no + base-image chaining) uses a rocm-systems RCCL (e.g. develop `78e8ba0`) on top of + the rocm720 sglang base and prepends the overlay lib dir to `LD_LIBRARY_PATH`, so + the overlay `librccl.so` resolves before the base one. That build links `librccl` + *without* the `DT_NEEDED librocm_smi64.so` the base librccl carried, so + `rocm_smi` is never transitively loaded and every later stage (and the run) dies + on `import torch` with `libtorch_hip.so: undefined symbol: rsmi_init`. Fix + (already baked into the Dockerfile): re-add the dependency on the overlay librccl + during the RCCL stage — + `patchelf --add-needed librocm_smi64.so. "$(readlink -f .../librccl.so)"` — + then sanity-check `python3 -c "import torch"` *with the overlay librccl resolved + first* (the exact case that regressed). With the smifix the full overlay + **base -> RCCL (+smifix) -> MoRI -> NIXL/Mooncake KV-transfer -> Mooncake pip** + runs green (validated 4-node DeepSeek-R1, 56/56 sweep points). If you don't need + a specific RCCL version, skip the RCCL stage and the base RCCL also works; if you + keep it, apply the smifix. The Dockerfile sanity-checks `import torch` after each + stage — a failure at the NIXL stage usually means an earlier stage (RCCL) is the + real culprit, not NIXL. + +- **Keep runtime installs out of the launcher — bake them into the image.** The + launcher (`scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh`) used to + `pip install` py-spy/flask/pyyaml and build rdma-core v62 at job start; that + mutated the container runtime env and silently corrupted the Python libs (the + model flags then stopped parsing and the server came up with defaults). Those pip + deps and the Mooncake KV backend are now baked into the full-overlay Dockerfile; + rdma-core v62 is an optional stage in the SAME Dockerfile, gated by + `--build-arg ENABLE_RDMA62=1` (for hosts whose RDMA stack needs a newer + libibverbs/librdmacm/libmlx5 than the base ships — no host `libibverbs` + bind-mounts needed; unset/default skips the stage entirely, mirroring + `docker/primus_megatron_train_rccl_overlay...Dockerfile`'s `RDMA_CORE_VERSION` + gate). There is no separate rdma-core-variant Dockerfile anymore — it was a + ~92%-identical copy of the base overlay and was merged in to cut duplication. + Do not reintroduce a runtime build-layer in the launcher. + +- **The MoRI overlay (#366, e.g. commit `a14e6992`) is the one that fixes the + mid-decode inter-node freeze.** (Commit hashes here and below — RCCL, MoRI, + NIXL — are illustrative pins valid when written; they can drift after a squash, + so track the PR/branch as the source of truth, not the SHA.) Without it, multi-node decode hangs partway + through generation (observed ~token 31 during warmup): the MoE expert-parallel + all-to-all stalls, the decode worker stops responding, and the proxy reports + `No response from detokenizer` / the benchmark stream aborts + (`ClientPayloadError` / `TransferEncodingError`). It is `RUN_MORI=1` + + `MORI_RDMA_DEVICES`/`MORI_SOCKET_IFNAME`/`MORI_IB_GID_INDEX` that route this + transport — set them alongside the NCCL/IB vars. + +- **`No response from detokenizer` on idle decode ranks can be a false + positive.** In a `yD>1` decode pool, idle decode ranks can trip the + detokenizer health check even when generation is healthy (sglang upstream + #20756). A *recent* sglang base fixes this; pin a base new enough to include it + (v0.5.12.post1+ worked) rather than chasing it at the transport layer. A true + hang (above) and this false positive look similar in the proxy log — confirm by + checking whether the *active* decode rank is still emitting tokens + (`py-spy dump` / `gstack` on the decode PID) before blaming the network. + +- **`run.sh` self-discovers the rank-ordered node IPs in-container — do NOT set + `--ipaddrs` / `IPADDRS`.** madengine's container runner does not forward + `SGLANG_NODE_IPS`, so a manifest that relied on it fell back to + `IPADDRS=localhost` and `DP_MODE=1` failed (roles dialed loopback and never + connected). `run.sh` now derives the rank-ordered IP list from the forwarded + `MASTER_ADDR`/`NODE_RANK`/`NNODES`: the SLURM nodelist when `scontrol`/`getent` + exist, else a rank-0 TCP rendezvous on `MASTER_ADDR`. `NNODES` falls back to `xP+yD` (proxy/router is + co-located on rank 0, no separate node). Leave `IPADDRS` / `--ipaddrs` unset in + the manifest; only override them for a reproduction launched outside madengine. + Raise `IP_SYNC_TIMEOUT` (default 1800s) if image-load skew delays peers. + +- **Let madengine distribute the image — don't hand-roll per-node `docker + load`.** For manifest-driven runs madengine already fans the image out on every + node via `container_runner._ensure_local_image_available()` (inspect → `docker + load` the `MAD_DOCKER_BUILDS/.tar` → build → pull); rank 0 `docker save`s + the tar, a TCP barrier syncs, then workers load it (see `launch-and-results.md` + and ROCm/madengine@d28f2f5). So with `local_image: true` + `MAD_DOCKER_BUILDS` + on shared FS the ~20 GB overlay image is distributed for you — no manual + `docker save`/`docker load` step is required. A manual per-node `docker load` + is only relevant for the standalone holder/`srun` reproduction that launches the + container *outside* madengine's dispatcher. + +- **Perf CSV is produced only on rank 0; make the benchmark fail-fast.** Set + `MAD_COLLECT_METRICS=true`/`MAD_SKIP_PERF_COLLECTION=false` on rank 0 only, and + `BENCHMARK_FAIL_FAST=1` so a zero-result sweep aborts instead of writing an + empty `perf_sglang-disagg-DeepSeek-R1.csv`. The sweep is + `BENCHMARK_COMBINATIONS` (isl/osl) x `BENCHMARK_CONCURRENCY_LEVELS`; read the + rank-0 aggregated CSV, not a non-zero node's local file. + +- **A single transient gateway circuit-open should not nuke the whole sweep — + retry points.** Occasionally a point fails at warmup with + `Service Unavailable ... No available prefill workers (all circuits open or + unhealthy)` even though the servers recover seconds later: the sgl-model-gateway + circuit breaker opened on a transient prefill detok hiccup. Under + `BENCHMARK_FAIL_FAST=1` that one blip otherwise aborts a ~50-minute run. The + sweep retries each point `BENCHMARK_POINT_RETRIES` times (default 2) with a + `BENCHMARK_RETRY_COOLDOWN_SECONDS` (default 45s) pause so the breaker can + re-close; a point that fails *every* attempt still fails the run, so real + regressions are not masked. Distinguish this transient from a true hard detok + freeze (heartbeat frozen for minutes, never recovering) by checking + `last_heartbeat` in `/run_logs//prefill_NODE*.log` / + `decode_NODE*.log` — a hard freeze needs a real fix (e.g. the MoRI overlay), + not a retry. + +## primus_megatron (training) + +Keywords: primus megatron scaleout MoE training, primus_turbo DeepEP token dispatcher, +use_turbo_deepep moe_enable_deepep flex alltoall, moe_shared_expert_overlap assert, +print_rank_last throughput last global rank, multi-node perf collection rank-0. + +- **DeepEP is opt-in and OFF by default (dispatcher = `alltoall` over RCCL).** + The baseline MoE token dispatcher is Megatron `alltoall` (a `torch.distributed` + all-to-all over the EP group, carried by RCCL on the RoCE/IB fabric) — *not* + MoRI. To switch to the Primus-Turbo **DeepEP** dispatcher, set + `PRIMUS_USE_DEEPEP=1` in the manifest env (a primus scaleout manifest can carry + it defaulting to `"0"`), which makes + `scripts/primus/megatron-lm/primus_megatron-lm_benchmark_report.sh` + (invoked by the model `run.sh`) pass `--use_turbo_deepep true` to Primus. Primus + then (`_is_turbo_deepep_enabled`) auto-sets `moe_enable_deepep=True` and + `moe_token_dispatcher_type='flex'` and swaps in + `PrimusTurboDeepEPTokenDispatcher`. DeepEP lives inside `primus_turbo` (module + `primus_turbo.pytorch.deep_ep`), not as a standalone `deep_ep` package. + +- **DeepEP requires `tensor_model_parallel_size == 1` + `enable_primus_turbo`, and + is incompatible with `moe_shared_expert_overlap`.** With DeepEP on, Primus ROCm + `validate_args` asserts + `AssertionError: DeepEP not support moe_shared_expert_overlap, please set + moe_shared_expert_overlap=False` (DeepSeek-V3 enables shared-expert overlap by + default). So enabling DeepEP must also pass `--moe_shared_expert_overlap false` + (done automatically alongside `--use_turbo_deepep true`). Measured 4-node x + 8-GPU DeepSeek-V3 proxy config (MI300X, BF16, seq 4096): DeepEP ~295-299 TFLOP/s/GPU, + ~13.1-13.3k tokens/s/GPU vs alltoall ~270 / ~12.05k (~+9%). + +- **Throughput is printed on the global last rank (`print_rank_last`), not rank 0 + — read the aggregated perf, not node 0's local CSV.** On a multi-node run the + Megatron/Primus training_log line (`TFLOP/s/GPU`, `tokens/s/GPU`) is emitted by + the last global rank, which usually lands on the *last* node, while madengine's + designated metric collector is rank 0 (node 0). madengine's login-node + `collect_results` handles this (it picks the richest per-node `multiple_results` + CSV via `_select_best_multiple_results_csv`, and treats an empty local perf on a + `skip_perf_collection` SLURM node as SUCCESS). Use a madengine that has this + multi-node aggregation (present in ROCm/madengine `develop`); otherwise the run + trains fine but is mis-reported as FAILED with an empty `perf.csv`. + +- **Per-model global batch size must be normalized to the world size, or Megatron + aborts at startup on non-standard node counts.** The config `global_batch_size` + is tuned for a single 8-GPU node; at any node count where + `GBS % (MBS * world_size) != 0` (e.g. a 3-node/24-GPU run: `512 % (4*24) != 0`) + Megatron aborts with `global batch size ... is not divisible by micro batch + size ... times data parallel size`. In our consolidated + `scripts/primus/megatron-lm/primus_megatron-lm_benchmark_report.sh` this is + handled generically for *every* model branch: `scaleout_gbs_override` reads + MBS/GBS from the run's config YAML and, when `NUM_GPUS > 8`, rounds GBS up to + the next multiple of `MBS * NUM_GPUS` (via `normalize_global_batch_size`) and + passes it as an explicit `--global_batch_size` override; single-node runs emit + no override so behavior there is byte-for-byte unchanged. Any new per-model + branch that hardcodes GBS should route it through `scaleout_gbs_override` to + stay node-count-portable. (Upstream fix for the 8B branch: + [mkuznet1/MAD#1](https://github.com/mkuznet1/MAD/pull/1) — our tree already + generalizes it, so no per-branch change is needed here.) + +- **`rocm/primus:v26.4` auto-loads `librccl-anp.so`, which deadlocks a bundled + RCCL overlay — set `NCCL_NET_PLUGIN=none`.** The v26.4 base image ships an + environment default of `NCCL_NET_PLUGIN=librccl-anp.so` (the ANP net plugin). + When the image carries a *bundled* RCCL overlay (the whole point of the + `rccl_overlay` Dockerfile), that plugin is incompatible with the overlay + `librccl` and RCCL init hangs at the first collective — the run never starts and + eventually times out. Set `NCCL_NET_PLUGIN=none` in the manifest env (both + `context.docker_env_vars` and `deployment_config.env_vars`, like the other + transport vars) to disable the plugin and let the bundled `librccl` drive the + IB/RoCE net path directly. The primus template ships this key set to `none`. diff --git a/.claude/skills/mad-slurm-multinode/references/launch-and-results.md b/.claude/skills/mad-slurm-multinode/references/launch-and-results.md new file mode 100644 index 0000000..df1bb58 --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/references/launch-and-results.md @@ -0,0 +1,104 @@ +# Launch and results + +Keywords: madengine run --manifest-file --live-output, perf.csv perf_super, +multiple_results, slurm_output node_N.out, multi-node aggregation, tok/s/gpu, +TFLOPS, sbatch squeue, MODEL_DIR scripts missing + +## Launch + +From `/rundir`, in a shell where `mad.env` has been sourced: + +```bash +cd "$WORKDIR/rundir" +source mad.env # same shell that runs madengine +madengine run --manifest-file run_manifest_.json --live-output +``` + +Useful flags (`madengine run --help` for the full list): +- `--manifest-file/-m ` — run an existing manifest (skips the separate + build phase; the image is still ensured at run time, see below). +- `--live-output/-l` — stream container/job output in real time. +- `--output/-o ` — performance output CSV (default `perf.csv`). +- `--summary-output/-s ` — JSON run summary. +- `--keep-alive` — leave containers up after the run (debugging). +- `--keep-model-dir` — keep the staged model dir (debugging). +- `--verbose/-v` — verbose logging. + +`madengine` renders the sbatch script from the `deployment_config` and submits +the SLURM job itself. Watch the job with `squeue -u $USER`; per-node +stdout/stderr land under `slurm_output/` (e.g. `node_0.out`, `node_1.out`). + +## Where the image comes from (madengine handles it — no pre-build needed) + +For a `local_image: true` manifest, the execution phase calls +`_ensure_local_image_available()` on each node (`container_runner.py`), so no +pre-build or pre-pull is needed. Per node the logic is: + +1. `docker image inspect ` — already present? Then nothing to do. +2. On the primary node (NODE_RANK 0), if the image is missing: + - if a tar exists at `MAD_DOCKER_BUILDS/.tar` -> `docker load` it; + - else **build it from the manifest `dockerfile`** (`docker build -f --pull ...`), + and if that build fails, **fall back to `docker pull `**; + - if `MAD_DOCKER_BUILDS` is set and the tar is absent -> `docker save` the + image into that tar. +3. A TCP barrier syncs all nodes; then worker nodes (rank > 0) `docker load` + the tar (when `MAD_DOCKER_BUILDS` is set) or build/pull independently. + +Implications for the agent: +- **No manual build/pull step.** Just run `madengine run -m`. +- **First run is slower** when the image doesn't exist yet — rank 0 builds it + once, then everyone reuses the tar. +- **`MAD_DOCKER_BUILDS` lives on shared FS** so rank 0's tar is visible to + workers (otherwise every node rebuilds/pulls). This is the same gotcha as in + `SKILL.md`. +- **The build path needs the manifest's `dockerfile` to resolve** — it lives in + the cloned `MAD` (the dockerfile path is relative to the MAD repo). + If the dockerfile is missing/unresolvable, madengine falls back to + `docker pull `, which then requires that tag to exist in a + reachable registry. +- Keep the manifest's `docker_image` tag meaningful: it is both the build tag + and the pull fallback tag, and it names the tar in `MAD_DOCKER_BUILDS`. + +## Where results land + +- `perf.csv` (or `--output` path) in `rundir` — the aggregated performance row. +- `perf_super*.csv` / `perf_entry*.csv` — intermediate per-entry files. +- The model's `multiple_results` CSV (e.g. + `perf_primus-megatron-Llama-3.1-8B.csv`) — per-model detail. +- `slurm_output/node_*.out` — raw logs (throughput banner, NCCL/RCCL transport + lines, tracebacks). + +## Multi-node aggregation gotcha + +Throughput is frequently emitted only by the LAST global rank (often the last +node), not node 0. So node 0's local perf CSV can be empty even on a fully +healthy run. The login-node aggregated `perf.csv` in `rundir` is the source of +truth over any single node's file. A node-0-only "empty perf" check is a false +negative — the aggregation across nodes picks the non-empty result. + +## Reading the number + +- Primus: report tok/s/gpu and TFLOPS/gpu (from the aggregated CSV / rank-last + `.out` throughput banner). +- sglang_disagg: request throughput / latency from the disagg benchmark CSV. + +## Quick failure triage + +- Run dies immediately with a missing run-script / empty `scripts/` -> + `MODEL_DIR` not exported. Re-`source mad.env` in the launching shell and + confirm `[ -d "$MODEL_DIR/scripts" ]`. +- Workers rebuild the image or fail to find it -> `MAD_DOCKER_BUILDS` is not on + shared storage visible to every node. Point it at shared FS and re-run. +- All ranks exit in the first minutes with `Failed to initialize any NET + plugin` / RCCL falls back to sockets -> wrong `NCCL_IB_HCA` for the node, or + AINIC transport vars missing in one of the two env blocks. See + cluster-types.md. +- Data/gated-model error -> HF token missing/expired (`~/.huggingface/token`). + A rendered `docker run` showing `MAD_SECRETS_HFTOKEN='${MAD_SECRETS_HFTOKEN}'` + means a manifest wrongly redeclared the token; remove that key and re-source + `mad.env` (see manifests.md "Secrets"). +- "path does not exist" / an unexpectedly empty mounted dir -> swapped + `docker_mounts`. Entries are keyed `{container_path: host_path}`, so the host + path is the value (see manifests.md "docker_mounts direction"). +- `MODEL_PATH is missing` -> wrong model host path; the model + resolves at `$MODEL_DIR/$MODEL_NAME` inside the container. diff --git a/.claude/skills/mad-slurm-multinode/references/manifests.md b/.claude/skills/mad-slurm-multinode/references/manifests.md new file mode 100644 index 0000000..b164f12 --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/references/manifests.md @@ -0,0 +1,151 @@ +# Manifest anatomy and per-run adaptation + +Keywords: run_manifest json, built_images local_image, docker_env_vars, +deployment_config slurm partition account qos reservation nodelist exclude, +distributed nnodes nproc_per_node launcher torchrun, multiple_results, +docker_mounts, ${VAR} expansion, MAD_SECRETS_HFTOKEN + +A run manifest is a build+run description madengine executes with +`madengine run --manifest-file `. The templates in `assets/manifests/` +are sanitized — fill them per run and write the result into `/rundir/`. +A filled, cluster-private manifest stays in the rundir, not back in the skill. + +## Top-level keys + +- `built_images` — the image(s) to run. `docker_image` is the local tag. + `local_image: true` means the image is treated as locally provided: at run + time madengine ensures it per node (load from `MAD_DOCKER_BUILDS` tar, else + build from `dockerfile`, else `docker pull`) — see launch-and-results.md + ("Where the image comes from"). `dockerfile` is the path used for that build. +- `built_models` — model/run metadata: `scripts` (run.sh path, resolved under + `MODEL_DIR`), `args` (`--model_repo ...`), `multiple_results` (the per-model + perf CSV name to look for), `additional_docker_run_options` (privileged, + shm-size, device mounts, ulimits). +- `context` — `docker_env_vars` (env inside the container), + `docker_mounts` (bind mounts, keyed `{container_path: host_path}`; madengine + renders each entry as `-v :`), `docker_build_arg`, + `docker_gpus`. +- `deployment_config` — `slurm` block, `distributed` block, `env_vars` + (env on the SLURM job / host side). + +## Verify assets resolve under MODEL_DIR + +The `built_images.*.dockerfile` and `built_models.*.scripts` (the `run.sh`) +paths are relative to `$MODEL_DIR` — the cloned `MAD` checkout. After +cloning or switching its branch, both resolve there +(`[ -f "$MODEL_DIR/" ] && [ -f "$MODEL_DIR/" ]`). A missing +path stops the run with a report (a wrong branch, uninitialized submodules, or a +different layout) rather than a silent search; the next step is to ask the user +for the correct branch. + +## Secrets (HF token) + +The HF token comes from `mad.env`, not the manifest. `mad.env` exports +`MAD_SECRETS_HFTOKEN=$(cat ~/.huggingface/token)`, and madengine forwards every +`MAD_SECRETS_*` from the sourced host env into the container's env +(`core/context.py` reads `MAD_SECRETS*` from `os.environ`; host-env values take +priority over manifest entries). A manifest that also sets +`"MAD_SECRETS_HFTOKEN": "${MAD_SECRETS_HFTOKEN}"` triggers an HF 401: madengine +renders that entry single-quoted in `docker run`, so the container receives the +literal string instead of the token. The templates carry no token key, and +keeping it that way is the fix. + +## Fill checklist (every run) + +Cluster-private — request from the user, keep out of the skill: +- `deployment_config.slurm.partition` +- `deployment_config.slurm.account` (omit the key if the cluster has no accounts) +- `deployment_config.slurm.qos` +- `deployment_config.slurm.reservation` (AINIC reserved campaigns) — **NOTE**: madengine + 2.0.0.post53 silently drops this field from the generated sbatch script. After submit, fix + with `scontrol update jobid= Reservation=`. Alternatively, `export + SBATCH_RESERVATION=` in the shell before `madengine run`. +- `deployment_config.slurm.nodelist` and/or `exclude` (specific node names) + +Cluster-shape — confirm on the node (see cluster-types.md): +- `NCCL_IB_HCA` (in BOTH env blocks) — the mlx5/rdma device list +- `NCCL_SOCKET_IFNAME` / `GLOO_SOCKET_IFNAME` / `slurm.network_interface` — the + same iface value in all three places (`context.docker_env_vars`, `slurm`, and + `deployment_config.env_vars`); a mismatch splits control-plane routing. +- `NCCL_IB_GID_INDEX` +- `RDMAV_DRIVERS` / `IBV_DRIVERS` (+ `RCCL_AINIC_ROCE=1` on AINIC, delete on CX7) +- `MAD_SYSTEM_GPU_ARCHITECTURE` / `BUILD_GPU_TARGETS` + +Run-shape — the user's run choice: +- `slurm.nodes` equals `distributed.nnodes` equals the cardinality of + `nodelist` (if used). A mismatch -> sbatch/torchrun disagree on world size. +- `slurm.time` walltime, `gpus_per_node`, `nproc_per_node`. +- `docker_image` tag (the RCCL/build tag you intend to test). + +Paths — from `mad.env` via `${VAR}` expansion where possible +(`${HF_HOME}`, `${MAD_DATAHOME}`, ...), explicit `` otherwise. +madengine expands `${VAR}` from the sourced host env, so keep mad.env and the +manifest consistent. Any workload-specific data paths are explicit +`/...` in the workload manifest itself, not exported +from mad.env (which stays cluster/cache-only). + +## Placeholder convention in templates + +JSON has no comments, so placeholders carry inline guidance: +- `""` — a transport/cluster value; + replace the whole string with the archetype value from the matrix in + [cluster-types.md](cluster-types.md). The templates are cluster-agnostic, so + the per-archetype values live only in that matrix, not inline in the manifest. +- `"RCCL_AINIC_ROCE": ""` — set per archetype + (AINIC: `"1"`); delete the key entirely on an archetype that does not need it. +- `${VAR}` — leave as-is; madengine expands from the sourced env. + +Re-validate after editing with the static checker (GPU-free, includes JSON +validity plus the fill checklist above): +`bash "$SKILL_DIR/scripts/validate_manifest.sh" .json` (source +`mad.env` first so it can also resolve the dockerfile/run.sh under `$MODEL_DIR`). + +## docker_mounts direction + +`docker_mounts` is keyed `{container_path: host_path}` and madengine renders +each entry as `-v :` (`execution/container_runner.py` +emits `-v {value}:{key}`; `orchestration/run_orchestrator.py` reads the dict as +`{container_path: host_path}`). The key is the path inside the container; the +value is the path on the host. Swapping the two sides points Docker at a +non-existent host path, which Docker then creates as an empty directory, and the +workload fails with a missing-path error (for example a `MODEL_PATH is missing` +error). When a host data path and its container mount point +differ, the host path belongs on the value side. + +## Per-workload notes + +- **Primus** (`primus_pyt_megatron_lm_train_llama-3.1-{8b,70b}_scaleout`): uses the + `rccl_overlay` Dockerfile; `launcher: torchrun`, `nproc_per_node: 8`. + `multiple_results` = `perf_primus-megatron-Llama-3.1-{8B,70B}.csv`. 70B is + heavier — give it more walltime / more nodes. +- **sglang_disagg** (`sglang_disagg_deepseek-r1`): disaggregated + prefill/decode serving of DeepSeek-R1 on SGLang. + `launcher: sglang-disagg`, `scripts/sglang_disagg/run.sh` entrypoint, + `nproc_per_node: 8`, typically 4 nodes (e.g. xP=2 prefill, yD=2 decode; the + proxy/load-balancer rides rank 0). Extra knobs: `xP`, `yD`, `RUN_MORI` (MoE + expert-parallel all-to-all via MoRI — keep `1` for multi-node), `DP_MODE`, + `KV_TRANSFER_BACKEND` (`nixl` or `mooncake`), `GPUS_PER_NODE`, the `MORI_*` + transport vars (mirror the `NCCL_IB_*` data-plane values), and the benchmark + sweep (`BENCHMARK_COMBINATIONS`, `BENCHMARK_CONCURRENCY_LEVELS`, + `BENCHMARK_FAIL_FAST`). `run.sh` self-discovers the rank-ordered node IP list + in-container (SLURM nodelist, else a rank-0 TCP rendezvous on the forwarded + `MASTER_ADDR`), so do NOT set `IPADDRS` / `--ipaddrs`: madengine does not + forward `SGLANG_NODE_IPS`, and a stale hand-filled list breaks role addressing + (tune `IP_SYNC_TIMEOUT`, default 1800s, for image-load skew instead). The image + is a single full-overlay build + (`docker/sglang_disagg_inference_full_overlay.ubuntu.amd.Dockerfile`): base + SGLang + RCCL + MoRI + NIXL/Mooncake KV-transfer + Mooncake pip in one + Dockerfile (no base-image chaining). For hosts whose RDMA stack needs a + newer rdma-core than the base ships, build the same Dockerfile with + `--build-arg ENABLE_RDMA62=1` (bakes rdma-core v62 in, so no host `libibverbs` + bind-mounts are needed; unset/default skips that stage). The full overlay + re-adds the `rocm_smi` `DT_NEEDED` + (`patchelf --add-needed librocm_smi64.so.`) so a newer RCCL on the rocm720 + base does not break `import torch` (see [gotchas.md](gotchas.md#sglang_disagg)). + Perf lands in `perf_sglang-disagg-DeepSeek-R1.csv`, collected on rank 0 only. + +## Scaling a manifest (e.g. 2-node -> 4-node) + +Change `slurm.nodes`, `distributed.nnodes`, and the `nodelist` cardinality +together. Everything else (env, mounts, image) stays the same. Real examples +exist at 2 and 4 nodes for primus and sglang_disagg. diff --git a/.claude/skills/mad-slurm-multinode/scripts/detect_cluster_env.sh b/.claude/skills/mad-slurm-multinode/scripts/detect_cluster_env.sh new file mode 100755 index 0000000..47e5f03 --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/scripts/detect_cluster_env.sh @@ -0,0 +1,106 @@ +#!/usr/bin/env bash +# Propose cluster-specific values for mad.env / manifests by inspecting the node. +# Read-only. The agent confirms these with the user before writing configs; +# they are best-effort guesses, not authoritative. +# +# This reflects the node it runs on. The login/jump node and the compute nodes +# can differ (HCAs, GPU arch, interfaces), so for compute-node values run it on +# an allocated node, e.g.: +# srun -p [--reservation ] [--nodelist ] -N1 \ +# bash detect_cluster_env.sh +set -u + +echo "== mad-slurm-multinode cluster env detection (proposals only) ==" + +# --- GPU architecture --- +arch="" +if command -v rocminfo >/dev/null 2>&1; then + arch=$(rocminfo 2>/dev/null | grep -m1 -oE 'gfx[0-9a-f]+') +elif command -v rocm-smi >/dev/null 2>&1; then + arch=$(rocm-smi --showhw 2>/dev/null | grep -m1 -oE 'gfx[0-9a-f]+') +fi +echo "MAD_SYSTEM_GPU_ARCHITECTURE : ${arch:-}" + +# --- RDMA / IB HCAs --- +hcas="" +if command -v ibv_devices >/dev/null 2>&1; then + hcas=$(ibv_devices 2>/dev/null | awk 'NR>2 {print $1}' | grep -E '^(mlx5_|rdma|bnxt_re)' | sort) +fi +if [ -z "$hcas" ] && [ -d /sys/class/infiniband ]; then + hcas=$(ls /sys/class/infiniband 2>/dev/null | grep -E '^(mlx5_|rdma|bnxt_re)' | sort) +fi +if [ -n "$hcas" ]; then + list=$(echo "$hcas" | sed 's/$/:1/' | paste -sd, -) + echo "NCCL_IB_HCA (all ports :1) : $list" + if echo "$hcas" | grep -q '^rdma'; then + fam="AINIC (ionic) -> RDMAV_DRIVERS=ionic, IBV_DRIVERS=ionic, RCCL_AINIC_ROCE=1, GID likely 1" + elif echo "$hcas" | grep -q '^mlx5'; then + fam="CX7/Mellanox (mlx5) -> RDMAV_DRIVERS=mlx5, IBV_DRIVERS=mlx5, GID likely 3" + elif echo "$hcas" | grep -q '^bnxt_re'; then + fam="Broadcom Thor2 (bnxt_re) -> RDMAV_DRIVERS=bnxt_re, IBV_DRIVERS=bnxt_re, GID: check show_gids" + else + fam="unknown vendor -> archetype not auto-classified; confirm with the user (see references/cluster-types.md)" + fi + echo "archetype guess : $fam" + echo " (note: NCCL_IB_HCA usually lists only the GPU-attached HCAs; trim as needed)" +else + echo "NCCL_IB_HCA : " +fi + +# --- HCA <-> netdev / NUMA affinity (management vs data-plane) --- +# Some HCAs back a routed management iface (e.g. mlx5_1->eth0, mlx5_6->eth1) +# rather than a GPU rail. NCCL_IB_HCA should list only the data-plane HCAs, so +# split them by their backing netdev: eth*/eno*/bond* = management (skip), +# rdma*/ib*/none = data-plane (use). Confirm against GPU<->NIC affinity below. +echo "HCA -> netdev / NUMA affinity:" +if command -v ibdev2netdev >/dev/null 2>&1; then + ibdev2netdev 2>/dev/null | sed 's/^/ /' +fi +mgmt_hcas=""; data_hcas="" +if [ -d /sys/class/infiniband ]; then + for d in /sys/class/infiniband/*; do + [ -e "$d" ] || continue + dev=$(basename "$d") + numa=$(cat "$d/device/numa_node" 2>/dev/null) + nd=$(ls "$d/device/net" 2>/dev/null | paste -sd, -) + printf ' %-10s numa=%-3s netdev=%s\n' "$dev" "${numa:-?}" "${nd:-none}" + if echo "$nd" | grep -qE '^(eth|eno|ens|enp|em|bond)'; then + mgmt_hcas="$mgmt_hcas $dev" + else + data_hcas="$data_hcas $dev" + fi + done +fi +if [ -n "$data_hcas" ]; then + dlist=$(echo $data_hcas | tr ' ' '\n' | sed 's/$/:1/' | paste -sd, -) + echo " management HCAs (skip) :${mgmt_hcas:- none}" + echo " data-plane HCAs (use) :$data_hcas" + echo " NCCL_IB_HCA (data-plane :1) : $dlist" + echo " (heuristic: a HCA backing a routed eth*/bond iface is management;" + echo " rdma*/ib* rails attached to the GPUs are data-plane — confirm by affinity)" +fi + +# --- GPU PCIe bus / NUMA (to match data-plane rails to GPUs) --- +if command -v rocm-smi >/dev/null 2>&1; then + echo "GPU PCIe bus (rocm-smi --showbus):" + rocm-smi --showbus 2>/dev/null | grep -iE 'GPU|bus' | sed 's/^/ /' | head -20 +fi + +# --- management interface --- +echo "candidate NCCL_SOCKET_IFNAME :" +if command -v ip >/dev/null 2>&1; then + ip -br addr 2>/dev/null | awk '$2=="UP" && $1!="lo" {print " "$1" "$3}' +else + echo " " +fi +echo " (pick the routable management iface, e.g. eth0 on CX7, eno0 on AINIC)" + +# --- RoCEv2 GID index --- +if command -v show_gids >/dev/null 2>&1; then + echo "RoCEv2 GID candidates (show_gids, look for v2):" + show_gids 2>/dev/null | grep -iE 'v2|RoCE' | head -10 | sed 's/^/ /' +else + echo "NCCL_IB_GID_INDEX : " +fi + +echo "== confirm all of the above with the user before writing mad.env/manifest ==" diff --git a/.claude/skills/mad-slurm-multinode/scripts/preflight.sh b/.claude/skills/mad-slurm-multinode/scripts/preflight.sh new file mode 100755 index 0000000..63a6892 --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/scripts/preflight.sh @@ -0,0 +1,90 @@ +#!/usr/bin/env bash +# Step 0 preflight for mad-slurm-multinode. +# Prints a PASS/WARN/FAIL table of prerequisites. Exits non-zero if any HARD +# requirement (docker daemon, SLURM client) is missing. +set -u + +pass=0; warn=0; fail=0 +row() { printf ' %-6s %-22s %s\n' "$1" "$2" "$3"; } +ok() { row "PASS" "$1" "$2"; pass=$((pass+1)); } +wn() { row "WARN" "$1" "$2"; warn=$((warn+1)); } +bad() { row "FAIL" "$1" "$2"; fail=$((fail+1)); } + +echo "== mad-slurm-multinode preflight ==" + +# Python >= 3.10 (soft: a conda env is created later anyway) +if command -v python3 >/dev/null 2>&1; then + pv=$(python3 -c 'import sys;print("%d.%d"%sys.version_info[:2])' 2>/dev/null) + if python3 -c 'import sys;exit(0 if sys.version_info[:2]>=(3,10) else 1)' 2>/dev/null; then + ok "python3" "$pv (>= 3.10)" + else + wn "python3" "$pv (< 3.10; conda env 3.12 will be created)" + fi +else + wn "python3" "not found (conda env 3.12 will be created)" +fi + +# Docker (HARD): present + daemon reachable +if command -v docker >/dev/null 2>&1; then + if docker info >/dev/null 2>&1; then + ok "docker" "daemon reachable" + else + bad "docker" "present but 'docker info' failed (daemon/perms?)" + fi +else + bad "docker" "not found" +fi + +# SLURM client (HARD) +if command -v sbatch >/dev/null 2>&1 && command -v sinfo >/dev/null 2>&1; then + ok "slurm" "$(sinfo --version 2>/dev/null | head -1)" +else + bad "slurm" "sbatch/sinfo not found" +fi + +# conda / miniforge (soft: installed in Step 2 if missing) +if command -v conda >/dev/null 2>&1; then + ok "conda" "$(conda --version 2>/dev/null)" +else + # conda is not on PATH; an existing install may just need sourcing + cphit="" + for cp in "$HOME/miniforge3" "$HOME/miniconda3" "$HOME/mambaforge" \ + "/opt/conda" "${CONDA_PREFIX:-}" "${MAMBA_ROOT_PREFIX:-}"; do + if [ -n "$cp" ] && [ -f "$cp/etc/profile.d/conda.sh" ]; then cphit="$cp"; break; fi + done + if [ -n "$cphit" ]; then + wn "conda" "found at $cphit, not on PATH — source $cphit/etc/profile.d/conda.sh" + else + wn "conda" "not found (miniforge installed in Step 2)" + fi +fi + +# GPU SMI (soft: used for arch detection) +if command -v rocm-smi >/dev/null 2>&1; then + ok "gpu-smi" "rocm-smi (AMD)" +elif command -v nvidia-smi >/dev/null 2>&1; then + ok "gpu-smi" "nvidia-smi (NVIDIA)" +else + wn "gpu-smi" "neither rocm-smi nor nvidia-smi found" +fi + +# git (HARD: needed to clone/switch repos) +if command -v git >/dev/null 2>&1; then + ok "git" "$(git --version 2>/dev/null)" +else + bad "git" "not found" +fi + +# HF token file (soft: required at run time for gated Llama-3.1) +if [ -s "$HOME/.huggingface/token" ]; then + ok "hf-token" "~/.huggingface/token present" +else + wn "hf-token" "~/.huggingface/token missing (needed before launch)" +fi + +echo "== summary: $pass PASS, $warn WARN, $fail FAIL ==" +if [ "$fail" -gt 0 ]; then + echo "Hard requirements failed. This node is not ready for madengine SLURM runs." + exit 1 +fi +exit 0 diff --git a/.claude/skills/mad-slurm-multinode/scripts/validate_manifest.sh b/.claude/skills/mad-slurm-multinode/scripts/validate_manifest.sh new file mode 100644 index 0000000..f156dd6 --- /dev/null +++ b/.claude/skills/mad-slurm-multinode/scripts/validate_manifest.sh @@ -0,0 +1,147 @@ +#!/usr/bin/env bash +# GPU-free static validator for a filled run manifest. +# Usage: bash validate_manifest.sh +# Reads $MODEL_DIR (export it via `source mad.env`) to also resolve the +# dockerfile / run.sh paths. Read-only; prints a PASS/WARN/FAIL table and exits +# non-zero if any FAIL. No GPU, no docker, no network needed. +set -u + +manifest="${1:-}" +if [ -z "$manifest" ]; then + echo "usage: $0 " >&2 + exit 2 +fi +if ! command -v python3 >/dev/null 2>&1; then + echo "python3 required for validate_manifest.sh" >&2 + exit 2 +fi + +exec python3 - "$manifest" <<'PY' +import json, os, re, sys + +path = sys.argv[1] +model_dir = os.environ.get("MODEL_DIR", "") + +passes = warns = fails = 0 +def ok(m): + global passes; passes += 1; print(f" PASS {m}") +def warn(m): + global warns; warns += 1; print(f" WARN {m}") +def fail(m): + global fails; fails += 1; print(f" FAIL {m}") + +print(f"== mad-slurm-multinode manifest validation: {path} ==") + +if not os.path.isfile(path): + print(f" FAIL manifest file not found: {path}") + print("== 0 PASS, 0 WARN, 1 FAIL ==") + sys.exit(1) + +raw = open(path, encoding="utf-8", errors="replace").read() + +# 1) JSON validity +try: + m = json.loads(raw) + ok("valid JSON") +except Exception as e: + print(f" FAIL invalid JSON: {e}") + print("== 0 PASS, 0 WARN, 1 FAIL ==") + sys.exit(1) + +# 2) no unfilled placeholders +ph = re.findall(r']*>', raw) +if ph: + fail(f"{len(ph)} unfilled placeholder(s) remain, e.g. {ph[0]!r}") +else: + ok("no placeholders left") + +ctx = (m.get("context") or {}).get("docker_env_vars") or {} +dep = m.get("deployment_config") or {} +slurm = dep.get("slurm") or {} +dist = dep.get("distributed") or {} +denv = dep.get("env_vars") or {} + +def is_ph(v): + return isinstance(v, str) and (" 1: + fail(f"network interface mismatch across blocks: {present}") +elif vals: + ok(f"network interface consistent ({vals[0]})") + +# 5) node count consistency +nodes, nnodes = slurm.get("nodes"), dist.get("nnodes") +if nodes is not None and nnodes is not None: + if nodes != nnodes: + fail(f"slurm.nodes ({nodes}) != distributed.nnodes ({nnodes})") + else: + ok(f"slurm.nodes == distributed.nnodes ({nodes})") +nl = slurm.get("nodelist") +if isinstance(nl, str) and nl and not is_ph(nl): + if "[" in nl: + warn(f"nodelist uses a bracket range; cardinality not auto-checked ({nl})") + elif nodes is not None: + count = len([x for x in nl.split(",") if x.strip()]) + if count != nodes: + fail(f"nodelist has {count} node(s) but slurm.nodes={nodes}") + else: + ok(f"nodelist cardinality matches slurm.nodes ({count})") + +# 6) stray HF token in the manifest -> HF 401 +if "MAD_SECRETS_HFTOKEN" in ctx or "MAD_SECRETS_HFTOKEN" in denv: + fail("MAD_SECRETS_HFTOKEN is declared in the manifest -> HF 401; " + "remove it (the token comes from mad.env)") +else: + ok("no MAD_SECRETS_HFTOKEN key in manifest") + +# 7) AINIC transport vars must be symmetric across both env blocks +for k in ("RCCL_AINIC_ROCE", "RDMAV_DRIVERS", "IBV_DRIVERS"): + inc = (k in ctx) and not is_ph(ctx.get(k)) + ind = (k in denv) and not is_ph(denv.get(k)) + if inc != ind: + fail(f"{k} set in only one env block (context={inc}, env_vars={ind}); " + "transport vars must be in BOTH") + +# 8) dockerfile / run.sh resolve under MODEL_DIR (if available) +asset_paths = [] +for v in (m.get("built_images") or {}).values(): + if isinstance(v, dict) and v.get("dockerfile"): + asset_paths.append(("dockerfile", v["dockerfile"])) +for v in (m.get("built_models") or {}).values(): + if isinstance(v, dict) and v.get("scripts"): + asset_paths.append(("scripts", v["scripts"])) +if model_dir and os.path.isdir(model_dir): + for kind, p in asset_paths: + full = os.path.join(model_dir, p) + if os.path.exists(full): + ok(f"{kind} resolves under MODEL_DIR: {p}") + else: + fail(f"{kind} not found under MODEL_DIR: {p}") +else: + warn(f"MODEL_DIR not set/usable ('{model_dir}') -> " + "dockerfile/run.sh path checks skipped") + +print(f"== validate_manifest: {passes} PASS, {warns} WARN, {fails} FAIL ==") +sys.exit(1 if fails else 0) +PY diff --git a/.gitignore b/.gitignore index 1285852..b7c59cf 100644 --- a/.gitignore +++ b/.gitignore @@ -72,3 +72,6 @@ venv.bak/ *.log *.out *.html + +# Skill manifest templates are checked in despite *.json above +!.claude/skills/mad-slurm-multinode/assets/manifests/*.template.json diff --git a/docker/primus_megatron_train_rccl_overlay.ubuntu.amd.Dockerfile b/docker/primus_megatron_train_rccl_overlay.ubuntu.amd.Dockerfile new file mode 100644 index 0000000..113413a --- /dev/null +++ b/docker/primus_megatron_train_rccl_overlay.ubuntu.amd.Dockerfile @@ -0,0 +1,327 @@ +# CONTEXT {'gpu_vendor': 'AMD', 'guest_os': 'UBUNTU'} +# +# Primus megatron-lm training image with a candidate RCCL built from source. +# +# base : rocm/primus:v26.4 (ROCm 7.2.x, torch 2.10, python3.12) +# + : RCCL built from ${RCCL_REPO}@${RCCL_BRANCH}[/${RCCL_COMMIT}] for +# ${BUILD_GPU_TARGETS} (gfx942=MI300X/MI325X, gfx950=MI350X/MI355X) +# + : candidate librccl installed over EVERY location torch resolves +# + : optional rdma-core ${RDMA_CORE_VERSION} from source (Broadcom Thor2 +# bnxt_re EFAULT fix on some fabrics; empty = keep base rdma-core) +# +# WHY install the candidate librccl into BOTH the system rocm lib dir AND +# torch/lib (not just LD_LIBRARY_PATH=/opt/rccl/lib): +# the primus v26.x base ships NO bundled librccl under torch/lib and +# libtorch_hip.so links the SYSTEM /opt/rocm/lib/librccl.so.1 (verified via ldd +# on v26.3). LD_LIBRARY_PATH alone is not enough, so we overwrite the system +# librccl and drop a copy into torch/lib (RPATH=$ORIGIN) as belt-and-suspenders. +# This is also safe on earlier v26.2/v26.3 bases. +# +# The git-clone source keeps .git, so RCCL's git_version.cmake bakes the real +# commit into the binary as a separate rcclGitHash string ":" +# (e.g. "HEAD:c67fbe4" for a detached checkout) — no tarball/hash-override hack +# required. The final gate confirms one of those baked short hashes is a prefix +# of the recorded RCCL_BUILT_SHA. +# +# IMAGE LAYOUT: RCCL is compiled in a throwaway `rccl-builder` stage and only +# the installed ${RCCL_INSTALL_DIR} tree is COPYed into the final image. This +# keeps the full RCCL source/build tree AND the heavy build toolchain +# (build-essential, cmake, ninja) out of the shipped image. +# +ARG BASE_DOCKER=rocm/primus:v26.4 + +# ---- shared base: apt mirror hygiene (so both stages resolve the same) ------ +FROM ${BASE_DOCKER} AS apt-base +SHELL ["/bin/bash", "-o", "pipefail", "-c"] +USER root +# Force https mirrors + ensure the universe/multiverse components are enabled so +# build deps resolve consistently in both stages. +RUN set -e && \ + sed -i 's|http://archive.ubuntu.com/ubuntu|https://archive.ubuntu.com/ubuntu|g; s|http://security.ubuntu.com/ubuntu|https://security.ubuntu.com/ubuntu|g' /etc/apt/sources.list 2>/dev/null || true && \ + if [ -d /etc/apt/sources.list.d ]; then \ + sed -i 's|http://archive.ubuntu.com/ubuntu|https://archive.ubuntu.com/ubuntu|g; s|http://security.ubuntu.com/ubuntu|https://security.ubuntu.com/ubuntu|g' /etc/apt/sources.list.d/*.list 2>/dev/null || true; \ + sed -i 's|http://archive.ubuntu.com/ubuntu|https://archive.ubuntu.com/ubuntu|g; s|http://security.ubuntu.com/ubuntu|https://security.ubuntu.com/ubuntu|g' /etc/apt/sources.list.d/*.sources 2>/dev/null || true; \ + sed -i -E 's/^Components:.*/Components: main restricted universe multiverse/g' /etc/apt/sources.list.d/*.sources 2>/dev/null || true; \ + fi + +# ---- Stage 1 (throwaway): build + verify candidate RCCL --------------------- +# Everything heavy (full RCCL clone, submodules, build tree, gcc/cmake/ninja) +# lives here and is discarded; only ${RCCL_INSTALL_DIR} is carried forward. +FROM apt-base AS rccl-builder + +ARG RCCL_REPO=https://github.com/ROCm/rocm-systems.git +ARG RCCL_BRANCH=develop +ARG RCCL_COMMIT= +ARG RCCL_INSTALL_DIR=/opt/rccl +# gfx942 = MI300X / MI325X ; gfx950 = MI350X / MI355X +ARG BUILD_GPU_TARGETS=gfx942 + +RUN mkdir -p "${RCCL_INSTALL_DIR}" +WORKDIR /opt + +RUN apt-get -o Acquire::ForceIPv4=true \ + -o Acquire::Retries=5 \ + -o Acquire::http::Timeout=30 \ + -o Acquire::https::Timeout=30 \ + update && \ + DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + git \ + ca-certificates \ + cmake \ + ninja-build \ + pkg-config \ + make \ + patch \ + build-essential \ + libatomic1 \ + patchelf && \ + rm -rf /var/lib/apt/lists/* + +# clone the candidate RCCL and record the exact built commit +RUN if [[ -n "${RCCL_COMMIT}" ]]; then \ + git clone "${RCCL_REPO}" /tmp/rccl; \ + else \ + git clone --depth 1 --branch "${RCCL_BRANCH}" "${RCCL_REPO}" /tmp/rccl; \ + fi && \ + cd /tmp/rccl && \ + if [[ -n "${RCCL_BRANCH}" ]]; then git checkout "${RCCL_BRANCH}"; fi && \ + if [[ -n "${RCCL_COMMIT}" ]]; then git checkout "${RCCL_COMMIT}"; fi && \ + if [[ -d projects/rccl ]]; then \ + cd projects/rccl && git submodule update --init --recursive; \ + echo "/tmp/rccl/projects/rccl" > /tmp/BLD_RCCL_HOME.txt; \ + else \ + git submodule update --init --recursive; \ + echo "/tmp/rccl" > /tmp/BLD_RCCL_HOME.txt; \ + fi && \ + git -C "$(cat /tmp/BLD_RCCL_HOME.txt)" rev-parse HEAD > "${RCCL_INSTALL_DIR}/RCCL_BUILT_SHA" && \ + echo "RCCL_BUILT_SHA=$(cat ${RCCL_INSTALL_DIR}/RCCL_BUILT_SHA)" + +# v26.4 compat: ROCm ships as pip wheels and the amdclang++ wrapper in +# _rocm_sdk_devel/bin/ resolves its clang++/clang-23 helpers relative to its own +# directory -- but those binaries live ONLY in _rocm_sdk_devel/lib/llvm/bin/. +# RCCL's device-code compile invokes bin/amdclang++ directly and dies with +# "amdclang++: binary '.../_rocm_sdk_devel/bin/clang++' does not exist". Symlink +# the helpers into bin/ so the wrapper resolves. No-op on v26.3 / non-wheel bases. +# (Fixing CC/CXX is NOT enough: --offload-device-only still calls bin/amdclang++.) +RUN SDK="$(ls -d /opt/venv/lib/python*/site-packages/_rocm_sdk_devel 2>/dev/null || true)" && \ + if [ -n "$SDK" ] && [ ! -e "$SDK/bin/clang++" ] && [ -d "$SDK/lib/llvm/bin" ]; then \ + for b in clang clang++ clang-23; do ln -sf ../lib/llvm/bin/$b "$SDK/bin/$b"; done; \ + echo "[v26.4-fix] symlinked clang/clang++/clang-23 into $SDK/bin"; \ + fi + +RUN set -e && \ + BLD_RCCL_HOME=$(cat /tmp/BLD_RCCL_HOME.txt) && \ + cd "${BLD_RCCL_HOME}" && \ + ./install.sh --amdgpu_targets="${BUILD_GPU_TARGETS}" --prefix="${RCCL_INSTALL_DIR}" + +# ---- build-verification gate: artifact exists + 0 undefined GDAKI/atomic ---- +# Fail first if the built librccl is missing (otherwise `nm` would error, the +# counts would default to 0, and the gate would wrongly pass). The `|| true` on +# the grep counts is only to tolerate "no match" (grep exit 1) under -o pipefail. +RUN set -e; \ + LIB="${RCCL_INSTALL_DIR}/lib/librccl.so"; \ + echo "=== librccl.so ==="; ls -laL "${RCCL_INSTALL_DIR}/lib/" | grep -i rccl || true; \ + [ -e "$LIB" ] || { echo "BUILD GATE FAIL: $LIB is missing"; exit 1; }; \ + nd=$(nm -D --undefined-only "$LIB" | grep -cE "[Gg]daki|GDAKI" || true); \ + na=$(nm -D --undefined-only "$LIB" | grep -E "__atomic_" | grep -vc "@LIBATOMIC" || true); \ + echo "gdaki_undefined=${nd:-0} unversioned_atomic_undefined=${na:-0}"; \ + if [ "${nd:-0}" -ne 0 ] || [ "${na:-0}" -ne 0 ]; then echo "BUILD GATE FAIL: undefined symbols present"; exit 1; fi; \ + echo "GATE PASS (built RCCL @ $(cat ${RCCL_INSTALL_DIR}/RCCL_BUILT_SHA))" + +# ---- final image ------------------------------------------------------------ +FROM apt-base + +ARG BASE_DOCKER +ARG RCCL_REPO=https://github.com/ROCm/rocm-systems.git +ARG RCCL_BRANCH=develop +ARG RCCL_INSTALL_DIR=/opt/rccl +# gfx942 = MI300X / MI325X ; gfx950 = MI350X / MI355X +ARG BUILD_GPU_TARGETS=gfx942 +# Optional rdma-core from source (e.g. 63.0 for the Broadcom Thor2 bnxt_re +# EFAULT fix). Empty string keeps the base image's rdma-core untouched. +ARG RDMA_CORE_VERSION= + +ENV WORKSPACE_DIR=/workspace +ENV RCCL_HOME=${RCCL_INSTALL_DIR} +# Prepend the candidate RCCL dir but KEEP the base image's paths (ROCm libs). +# The candidate librccl is also installed over the system path + ldconfig, so +# this is belt-and-suspenders rather than the sole resolution mechanism. +ENV LD_LIBRARY_PATH=${RCCL_INSTALL_DIR}/lib:${LD_LIBRARY_PATH} +ENV NCCL_DEBUG=WARN + +RUN mkdir -p "${WORKSPACE_DIR}" +WORKDIR /opt + +# Bring in ONLY the installed candidate RCCL tree (no source/build artifacts). +COPY --from=rccl-builder ${RCCL_INSTALL_DIR} ${RCCL_INSTALL_DIR} + +# Minimal tools needed to splice the candidate librccl into the image and to run +# the final ELF verification: patchelf (--add-needed), binutils (objdump / nm / +# strings / readelf) and libatomic1 (runtime dep of the candidate librccl). +RUN apt-get -o Acquire::ForceIPv4=true -o Acquire::Retries=5 update && \ + DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + patchelf binutils libatomic1 && \ + rm -rf /var/lib/apt/lists/* + +# ---- Stage 2: install candidate librccl over EVERY location torch resolves -- +# Overwrite the system /opt/rocm librccl AND drop a copy into torch/lib so the +# candidate always wins regardless of how libtorch_hip is linked. Also +# add-needed librocm_smi64 (folded-in rsmifix) so `import torch` resolves +# rsmi_init on bases where the overlay librccl lacks it in NEEDED. +# add_needed() adds librocm_smi64 to NEEDED idempotently (skips if already +# present, so re-running never appends duplicates) and FAILS the build if +# patchelf cannot apply it -- `import torch` depends on this resolving rsmi_init. +RUN set -e && \ + add_needed() { \ + local lib="$1"; \ + if objdump -p "$lib" | grep -qE "NEEDED[[:space:]]+librocm_smi64"; then \ + echo " add_needed: already present in $lib"; return 0; \ + fi; \ + patchelf --add-needed librocm_smi64.so.1 "$lib"; \ + objdump -p "$lib" | grep -qE "NEEDED[[:space:]]+librocm_smi64" || { echo "PATCHELF FAILED on $lib"; exit 1; }; \ + echo " add_needed: applied to $lib"; \ + }; \ + src=$(ls -L "${RCCL_INSTALL_DIR}/lib/librccl.so.1.0" 2>/dev/null || ls -L "${RCCL_INSTALL_DIR}/lib/librccl.so") && \ + echo "candidate librccl: $src" && \ + add_needed "$src" && \ + if [ -e "/opt/rocm/lib/librccl.so.1" ]; then \ + ROCM_LIB=$(dirname "$(readlink -f /opt/rocm/lib/librccl.so.1)"); \ + echo "system rocm lib dir: $ROCM_LIB"; \ + ls -la "$ROCM_LIB"/librccl.so* 2>/dev/null || true; \ + for f in "$ROCM_LIB"/librccl.so*; do \ + [ -f "$f" ] && [ ! -L "$f" ] && { echo "overwriting real file: $f"; cp -fL "$src" "$f"; add_needed "$f"; }; \ + done; \ + ldconfig; \ + echo "after:"; ls -laL "$ROCM_LIB"/librccl.so.1 2>/dev/null || true; \ + else \ + echo "[rccl-overlay] /opt/rocm/lib/librccl.so.1 not present -- targeted /opt/rocm/lib overwrite skipped (global sweep covers it)"; \ + fi && \ + TORCH_LIB=$(ls -d /opt/venv/lib/python*/site-packages/torch/lib 2>/dev/null | head -1) && \ + [ -n "$TORCH_LIB" ] && [ -d "$TORCH_LIB" ] && echo "torch lib dir: $TORCH_LIB" && \ + for f in "$TORCH_LIB"/librccl.so*; do \ + if [ -f "$f" ] && [ ! -L "$f" ]; then cp -fL "$src" "$f"; add_needed "$f"; fi; \ + done && \ + cp -fL "$src" "$TORCH_LIB/librccl.so.1.0" && \ + ln -sf librccl.so.1.0 "$TORCH_LIB/librccl.so.1" && \ + ln -sf librccl.so.1 "$TORCH_LIB/librccl.so" && \ + add_needed "$TORCH_LIB/librccl.so.1.0" && \ + if [ ! -e "$TORCH_LIB/librocm_smi64.so" ]; then \ + smi_src=$(ls /opt/rocm/lib/librocm_smi64.so.1 2>/dev/null || \ + ls /opt/venv/lib/python*/site-packages/_rocm_sdk_libraries/lib/librocm_smi64.so.1 2>/dev/null || \ + find /opt -name 'librocm_smi64.so.1' -not -type l 2>/dev/null | head -1 || true); \ + if [ -n "$smi_src" ]; then \ + cp -v -L "$smi_src" "$TORCH_LIB/librocm_smi64.so"; \ + ln -sfv librocm_smi64.so "$TORCH_LIB/librocm_smi64.so.1"; \ + echo "librocm_smi64 copied from: $smi_src"; \ + else \ + echo "[rccl-overlay] librocm_smi64.so.1 not found -- skipping copy (add-needed already applied to candidate)"; \ + fi; \ + fi && \ + echo "=== global sweep: overwrite EVERY real librccl on disk with candidate ===" && \ + canon_src=$(readlink -f "$src") && \ + for f in $(find / -name 'librccl.so*' -not -type l 2>/dev/null); do \ + [ "$(readlink -f "$f")" = "$canon_src" ] && continue; \ + echo " overwrite: $f"; cp -fL "$src" "$f"; add_needed "$f"; \ + done && \ + ldconfig +# WHY the global sweep: v26.4 splits ROCm libs into _rocm_sdk_libraries/lib +# (runtime .so) and _rocm_sdk_devel/lib (dev). torch maps librccl.so.1 from +# _rocm_sdk_libraries at RUNTIME (soname wins once loaded), which the targeted +# /opt/rocm + torch/lib overwrites above do NOT cover -- so the candidate is +# built and shipped but the STOCK base librccl is what actually runs. Overwriting +# every non-symlink librccl on disk makes whichever copy the loader maps the +# candidate, independent of ROCm's per-version layout churn. + +# ---- Stage 3 (optional): rdma-core from source ------------------------------ +# Builds only when RDMA_CORE_VERSION is non-empty (e.g. 63.0). Replaces the +# distro rdma-core to pick up the Broadcom Thor2 bnxt_re EFAULT fix. The build +# toolchain it needs is installed here (not in the base image) so the default +# image (RDMA_CORE_VERSION empty) never pays for it. +# +# ORDERING IS LOAD-BEARING: the source rdma-core is installed by REPLACING the +# distro libibverbs/librdmacm packages via `dpkg -r --force-all`. That leaves +# still-installed dependents (libopenmpi3t64, libucx0, rdmacm-utils, ...) with +# dangling deps, so ANY apt command after that point aborts with "Unmet +# dependencies". Therefore every apt operation (install + the docs-tooling +# purge/autoremove) MUST run BEFORE the dpkg removal, while the dpkg state is +# still consistent. Only dpkg + `ninja install` may run afterwards. +RUN if [[ -n "${RDMA_CORE_VERSION}" ]]; then \ + set -e; \ + apt-get -o Acquire::ForceIPv4=true -o Acquire::Retries=5 update; \ + DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + git ca-certificates cmake ninja-build build-essential pkg-config make \ + libnl-3-dev libnl-route-3-dev libudev-dev libssl-dev libsystemd-dev \ + python3-docutils pandoc; \ + rm -rf /var/lib/apt/lists/*; \ + git clone --depth 1 --branch "v${RDMA_CORE_VERSION}" https://github.com/linux-rdma/rdma-core.git /tmp/rdma-core; \ + cd /tmp/rdma-core && git log -1 --format='%H %s'; \ + mkdir build && cd build; \ + cmake -GNinja \ + -DCMAKE_INSTALL_PREFIX=/usr \ + -DCMAKE_INSTALL_LIBDIR=lib/x86_64-linux-gnu \ + -DCMAKE_INSTALL_SYSCONFDIR=/etc \ + -DCMAKE_INSTALL_RUNDIR=/run \ + -DCMAKE_BUILD_TYPE=Release \ + -DNO_PYVERBS=1 \ + ..; \ + ninja -j"$(nproc)"; \ + DEBIAN_FRONTEND=noninteractive apt-get purge -y python3-docutils pandoc; \ + DEBIAN_FRONTEND=noninteractive apt-get autoremove -y; \ + rm -rf /var/lib/apt/lists/*; \ + to_remove=""; \ + for p in ibverbs-providers libibverbs1 libibverbs-dev ibverbs-utils \ + librdmacm1 librdmacm-dev rdma-core libibumad3 libibmad5 \ + infiniband-diags; do \ + if dpkg-query -W -f='${Status}' "$p" 2>/dev/null | grep -q "install ok installed"; then \ + to_remove="$to_remove $p"; \ + fi; \ + done; \ + if [[ -n "$to_remove" ]]; then \ + echo "removing distro rdma packages:$to_remove"; \ + dpkg -r --force-all $to_remove; \ + fi; \ + ninja install && ldconfig; \ + cd / && rm -rf /tmp/rdma-core; \ + else \ + echo "RDMA_CORE_VERSION empty — keeping base image rdma-core"; \ + fi + +# ---- final verification: import torch + every librccl == candidate ---------- +RUN set -e && \ + if [[ -n "${RDMA_CORE_VERSION}" ]]; then \ + echo "=== rdma-core ===" && readelf -d /usr/lib/x86_64-linux-gnu/libibverbs.so.1 | grep SONAME && \ + { ls /usr/lib/x86_64-linux-gnu/libibverbs/libbnxt_re-rdmav*.so 2>/dev/null | head || true; }; \ + fi && \ + echo "=== import torch (must succeed) ===" && \ + if [ -x /opt/venv/bin/python3 ]; then PY=/opt/venv/bin/python3; else PY=$(command -v python3); fi && \ + echo "using python: $PY" && \ + "$PY" -c "import torch; print('torch', torch.__version__, 'hip', torch.version.hip)" && \ + echo "=== ldd libtorch_hip rccl resolution ===" && \ + TL=$(ls -d /opt/venv/lib/python*/site-packages/torch/lib | head -1) && \ + ldd "$TL/libtorch_hip.so" | grep -iE "rccl|smi" && \ + BUILT_SHA=$(cat "${RCCL_INSTALL_DIR}/RCCL_BUILT_SHA") && \ + echo "=== ASSERT: EVERY librccl on disk is the candidate built @ ${BUILT_SHA} ===" && \ + for p in $(find / -name 'librccl.so*' -not -type l 2>/dev/null); do \ + v=$(strings "$p" 2>/dev/null | grep -m1 -oE "RCCL version 2\.[0-9]+\.[0-9]+"); \ + echo "$v" | grep -q "RCCL version 2." || { echo "ASSERT FAIL: $p is not RCCL 2.x ($p -> ${v:-NONE})"; exit 1; }; \ + h=""; \ + for cand in $(strings "$p" 2>/dev/null | grep -oE "[A-Za-z0-9_./-]+:[0-9a-f]{7,40}" | sed 's/.*://' | sort -u); do \ + case "$BUILT_SHA" in "$cand"*) h="$cand"; break;; esac; \ + done; \ + echo " $p -> ${v:-NONE} / baked=${h:-NONE}"; \ + [ -n "$h" ] || { echo "ASSERT FAIL: $p has no baked git commit matching built SHA $BUILT_SHA (upstream 'Unknown' fallback?)"; exit 1; }; \ + done && \ + echo "=== candidate RCCL commit: ${BUILT_SHA} ===" + +# primus_base is derived from BASE_DOCKER so the label cannot drift from the +# actual base when BASE_DOCKER is overridden. +LABEL primus_base="${BASE_DOCKER}" +LABEL rccl_repo="${RCCL_REPO}" +LABEL rccl_branch="${RCCL_BRANCH}" +LABEL build_gpu_targets="${BUILD_GPU_TARGETS}" +LABEL rdma_core_version="${RDMA_CORE_VERSION}" + +WORKDIR ${WORKSPACE_DIR} + +# Record final Python environment for posterity. +RUN pip3 list diff --git a/docker/sglang_disagg_inference_full_overlay.ubuntu.amd.Dockerfile b/docker/sglang_disagg_inference_full_overlay.ubuntu.amd.Dockerfile new file mode 100644 index 0000000..ad51cd9 --- /dev/null +++ b/docker/sglang_disagg_inference_full_overlay.ubuntu.amd.Dockerfile @@ -0,0 +1,285 @@ +# CONTEXT {'gpu_vendor': 'AMD', 'guest_os': 'UBUNTU'} +############################################################################### +# SGLang Disaggregated P/D — MERGED full overlay (RCCL + MoRI + KV-transfer) +# +# Single-Dockerfile equivalent of the three chained overlays: +# base sglang -> RCCL overlay (+smifix) -> MoRI overlay -> KV-transfer (RIXL) +# Built in one `docker build` step so madengine's single-dockerfile build path +# produces the whole stack at once (no base_docker chaining between overlays). +# +# Pins (override via --build-arg): +# RCCL : ROCm/rocm-systems develop @ RCCL_COMMIT (default 78e8ba0) + smifix +# MoRI : ROCm/mori @ MORI_COMMIT (default a14e6992, includes #366) +# NIXL : ROCm/RIXL @ RIXL_COMMIT (default f33a5599) + ROCm/ucx @ UCX_COMMIT +# (the AMD NIXL implementation; exposed to SGLang as `nixl` via alias). +# NOTE: this is NOT ai-dynamo/nixl. The rocm KV transport for +# KV_TRANSFER_BACKEND=nixl is ROCm/RIXL; recipe mirrors the proven +# scripts/kvcache_transfer_bench/Dockerfile. +# +# Hosts whose RDMA stack needs a newer rdma-core than the base image ships: +# pass --build-arg ENABLE_RDMA62=1 to additionally build + install rdma-core +# v62 from source (newer libibverbs/librdmacm/libmlx5; no host libibverbs +# bind-mounts needed). Default (unset) keeps the image-default rdma-core and +# skips that stage entirely. +############################################################################### +ARG BASE_DOCKER=lmsysorg/sglang:v0.5.12.post1-rocm720-mi30x +FROM $BASE_DOCKER + +SHELL ["/bin/bash", "-o", "pipefail", "-c"] +USER root + +############################################################################### +# 1) RCCL overlay — rebuild RCCL from source so the RCCL under test wins over +# the base image's librccl. (mirrors sglang_disagg_inference_rccl_overlay) +############################################################################### +ARG RCCL_REPO=https://github.com/ROCm/rocm-systems.git +ARG RCCL_BRANCH=develop +ARG RCCL_COMMIT=78e8ba0 +ARG RCCL_INSTALL_DIR=/opt/rccl +ARG BUILD_GPU_TARGETS=gfx942 + +ENV RCCL_HOME=${RCCL_INSTALL_DIR} +# Prepend the overlay RCCL so it wins over the base image's librccl. +ENV LD_LIBRARY_PATH=${RCCL_INSTALL_DIR}/lib:${LD_LIBRARY_PATH} + +RUN mkdir -p "${RCCL_INSTALL_DIR}" +WORKDIR /opt + +RUN sed -i 's|http://|https://|g' /etc/apt/sources.list 2>/dev/null || true && \ + apt-get -o Acquire::Retries=5 update && \ + DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + git cmake ninja-build pkg-config make patch && \ + rm -rf /var/lib/apt/lists/* + +RUN if [[ -n "${RCCL_COMMIT}" ]]; then \ + git clone "${RCCL_REPO}" /tmp/rccl; \ + else \ + git clone --depth 1 --branch "${RCCL_BRANCH}" "${RCCL_REPO}" /tmp/rccl; \ + fi && \ + cd /tmp/rccl && \ + if [[ -n "${RCCL_BRANCH}" ]]; then git checkout "${RCCL_BRANCH}" || true; fi && \ + if [[ -n "${RCCL_COMMIT}" ]]; then git checkout "${RCCL_COMMIT}"; fi && \ + if [[ -d projects/rccl ]]; then \ + cd projects/rccl && git submodule update --init --recursive; \ + echo "/tmp/rccl/projects/rccl" > /tmp/BLD_RCCL_HOME.txt; \ + else \ + git submodule update --init --recursive; \ + echo "/tmp/rccl" > /tmp/BLD_RCCL_HOME.txt; \ + fi + +RUN set -e && \ + BLD_RCCL_HOME=$(cat /tmp/BLD_RCCL_HOME.txt) && \ + cd "${BLD_RCCL_HOME}" && \ + ./install.sh --amdgpu_targets="${BUILD_GPU_TARGETS}" --prefix="${RCCL_INSTALL_DIR}" + +RUN rm -rf /tmp/rccl /tmp/BLD_RCCL_HOME.txt + +# Re-add the rocm_smi dependency the rocm-systems RCCL build drops (smifix). +# Newer rocm-systems RCCL builds librccl WITHOUT a DT_NEEDED on librocm_smi64.so; +# torch's libtorch_hip.so relies on librccl to transitively pull in rocm_smi, so +# without this `import torch` dies with "undefined symbol: rsmi_init". +RUN set -e; \ + command -v patchelf >/dev/null 2>&1 || pip install --no-cache-dir patchelf >/dev/null 2>&1 \ + || { apt-get -o Acquire::Retries=5 update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends patchelf && rm -rf /var/lib/apt/lists/*; }; \ + librccl="$(readlink -f ${RCCL_INSTALL_DIR}/lib/librccl.so)"; \ + smi_soname="$(basename "$(ls /opt/rocm*/lib/librocm_smi64.so.[0-9]* 2>/dev/null | grep -E 'librocm_smi64\.so\.[0-9]+$' | head -1)")"; \ + test -n "${smi_soname}" || { echo "ROCM_SMI_SONAME_NOT_FOUND"; exit 1; }; \ + if readelf -d "${librccl}" 2>/dev/null | grep -q "NEEDED.*${smi_soname}"; then \ + echo "RCCL_SMI_NEEDED_ALREADY_PRESENT ${smi_soname}"; \ + else \ + patchelf --add-needed "${smi_soname}" "${librccl}" && echo "RCCL_SMI_NEEDED_ADDED ${smi_soname} -> ${librccl}"; \ + fi + +# Sanity: overlay librccl present AND torch imports with overlay librccl first. +RUN set -e; \ + test -e "${RCCL_INSTALL_DIR}/lib/librccl.so" || { echo "RCCL_OVERLAY_MISSING"; exit 1; }; \ + echo "RCCL_OVERLAY_OK $(ls -l ${RCCL_INSTALL_DIR}/lib/librccl.so*)"; \ + python3 -c "import torch; print('RCCL_OVERLAY_TORCH_OK', torch.__version__)" + +RUN pip list 2>/dev/null | grep -iE "sglang|torch" || true + +############################################################################### +# 2) MoRI overlay — substitutable MoE EP all-to-all (dispatch/combine, IBGDA). +# Default ROCm/mori @ a14e6992 includes #366 (internode decode hang fix). +############################################################################### +ARG MORI_REPO=https://github.com/ROCm/mori.git +ARG MORI_BRANCH=main +ARG MORI_COMMIT=a14e6992ffa95478e83127fe2672afff2840856f +ARG MORI_WHEEL_URL= +ARG MORI_GPU_ARCHS=gfx942 +ARG MORI_VERSION=1.2.0 +ARG MORI_SRC_DIR=/sgl-workspace/mori + +ENV MORI_GPU_ARCHS=${MORI_GPU_ARCHS} \ + MORI_SKIP_PRECOMPILE=1 \ + CMAKE_BUILD_TYPE=Release \ + SETUPTOOLS_SCM_PRETEND_VERSION=${MORI_VERSION} + +WORKDIR /sgl-workspace + +RUN sed -i 's|http://|https://|g' /etc/apt/sources.list 2>/dev/null || true && \ + apt-get -o Acquire::Retries=5 update && \ + DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + git cmake ninja-build pkg-config make patch && \ + rm -rf /var/lib/apt/lists/* + +RUN set -e; \ + if [[ -n "${MORI_WHEEL_URL}" ]]; then \ + echo "[mori-overlay] installing prebuilt wheel: ${MORI_WHEEL_URL}"; \ + pip install --no-cache-dir --force-reinstall "${MORI_WHEEL_URL}"; \ + else \ + echo "[mori-overlay] source build ${MORI_REPO}@${MORI_BRANCH}${MORI_COMMIT:+ (${MORI_COMMIT})} archs=${MORI_GPU_ARCHS}"; \ + rm -rf "${MORI_SRC_DIR}"; \ + if [[ -n "${MORI_COMMIT}" ]]; then \ + git clone "${MORI_REPO}" "${MORI_SRC_DIR}"; \ + cd "${MORI_SRC_DIR}" && git checkout "${MORI_COMMIT}"; \ + else \ + git clone --depth 1 --branch "${MORI_BRANCH}" "${MORI_REPO}" "${MORI_SRC_DIR}"; \ + fi; \ + cd "${MORI_SRC_DIR}" && git submodule update --init --recursive || true; \ + pip install --no-cache-dir --force-reinstall .; \ + fi + +# Sanity: MoRI must import and report the overlay version. +RUN python3 -c "import mori, importlib.metadata as m; \ +print('MORI_OVERLAY_OK', getattr(mori,'__version__', m.version('amd_mori')))" \ + || { echo 'MORI_OVERLAY_IMPORT_FAILED'; exit 1; } + +############################################################################### +# 3) KV-transfer overlay — ROCm UCX + ROCm/RIXL (the AMD "nixl" KV transport). +# Mirrors the proven recipe in scripts/kvcache_transfer_bench/Dockerfile. +# SGLang's KV_TRANSFER_BACKEND=nixl imports `nixl`; we alias rixl -> nixl. +############################################################################### +ARG ROCM_PATH=/opt/rocm +ARG KV_WORKSPACE=/sgl-workspace +ARG UCX_REPO=https://github.com/ROCm/ucx.git +ARG UCX_COMMIT=da3fac2a +ARG RIXL_REPO=https://github.com/ROCm/RIXL.git +ARG RIXL_COMMIT=f33a5599 + +ENV UCX_HOME=${KV_WORKSPACE}/ucx +ENV RIXL_HOME=${KV_WORKSPACE}/rixl +ENV PATH=${UCX_HOME}/bin:${PATH} +ENV LD_LIBRARY_PATH=${RIXL_HOME}/lib/x86_64-linux-gnu:${UCX_HOME}/lib:${LD_LIBRARY_PATH} + +WORKDIR ${KV_WORKSPACE} + +RUN sed -i 's|http://|https://|g' /etc/apt/sources.list 2>/dev/null || true && \ + apt-get -o Acquire::Retries=5 update && \ + DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + git cmake ninja-build pkg-config make patch autoconf automake libtool \ + cython3 libaio-dev libibverbs-dev librdmacm-dev libpci-dev \ + libgflags-dev libgoogle-glog-dev && \ + rm -rf /var/lib/apt/lists/* + +# meson/ninja/pybind11 build frontends for the RIXL meson build + wheel. +RUN pip install --no-cache-dir -U meson ninja pybind11 pyyaml build wheel + +# --- ROCm UCX --------------------------------------------------------------- +RUN set -e; \ + git clone "${UCX_REPO}" "${UCX_HOME}.src" && cd "${UCX_HOME}.src" && \ + git checkout "${UCX_COMMIT}" && \ + ./autogen.sh && mkdir -p build && cd build && \ + ../configure --prefix="${UCX_HOME}" --enable-shared --disable-static \ + --disable-doxygen-doc --enable-optimizations --enable-devel-headers \ + --with-rocm="${ROCM_PATH}" --with-verbs --with-dm --enable-mt && \ + make -j"$(nproc)" && make install && \ + cd "${KV_WORKSPACE}" && rm -rf "${UCX_HOME}.src" + +# --- ROCm/RIXL (the AMD NIXL implementation) -------------------------------- +RUN set -e; \ + git clone "${RIXL_REPO}" "${RIXL_HOME}" && cd "${RIXL_HOME}" && \ + git checkout "${RIXL_COMMIT}" && (git submodule update --init --recursive || true); \ + meson setup build --prefix="${RIXL_HOME}" -Ducx_path="${UCX_HOME}" -Drocm_path="${ROCM_PATH}" && \ + cd build && ninja && ninja install + +# Install the meson-built RIXL python package into site-packages directly. +# The upstream contrib/build-wheel.sh requires uv + py3.12 + auditwheel, which +# the py3.10 sglang base lacks; `ninja install` already produced the +# cpython-310 bindings under the RIXL prefix, so copy them into site-packages +# and alias `nixl` -> `rixl` for SGLang's KV_TRANSFER_BACKEND=nixl import path. +RUN set -e; \ + echo "=== RIXL python install tree ==="; \ + find "${RIXL_HOME}" -type d \( -name site-packages -o -name dist-packages \) | sed 's/^/PYDIR: /'; \ + pysp="$(find "${RIXL_HOME}" -type d \( -name site-packages -o -name dist-packages \) | head -1)"; \ + test -n "${pysp}" || { echo "RIXL_PY_INSTALL_NOT_FOUND"; find "${RIXL_HOME}" -name '_bindings*.so' -o -name '*.py' | head; exit 1; }; \ + echo "RIXL_PY_SRC=${pysp}"; ls -la "${pysp}"; \ + SP="$(python3 -c 'import site; print(site.getsitepackages()[0])')"; \ + copied=0; \ + for d in "${pysp}"/*/; do \ + name="$(basename "$d")"; \ + case "$name" in *dist-info|__pycache__) continue;; esac; \ + rm -rf "${SP}/${name}"; cp -a "$d" "${SP}/${name}"; echo "INSTALLED_PY_PKG ${name} -> ${SP}/${name}"; copied=1; \ + done; \ + test "$copied" = 1 || { echo "NO_RIXL_PKG_COPIED"; exit 1; }; \ + echo "import rixl, sys; sys.modules['nixl'] = rixl" > "${SP}/nixl_alias.pth"; \ + echo "NIXL_ALIAS_WRITTEN ${SP}/nixl_alias.pth" + +# Sanity: nixl (== rixl) import must succeed; report version + agent symbol. +RUN set -e; \ + python3 -c "import nixl; print('NIXL_IMPORT_OK', getattr(nixl,'__file__','?'))"; \ + python3 -c "import rixl, importlib.metadata as m; print('RIXL_VERSION', m.version('rixl'))" || true; \ + python3 -c "from nixl._api import nixl_agent, nixl_agent_config; print('NIXL_AGENT_OK')" \ + || echo "NIXL_AGENT_API_DIFFERS (verify sglang nixl import path at runtime)" + +############################################################################### +# 4) Runtime python deps + Mooncake KV-transfer backend. +# These were formerly built/pip-installed by the launcher at job start +# (scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh); baked into the image so +# the launcher no longer mutates the runtime environment. +############################################################################### +# Runtime python deps formerly pip-installed by the launcher. +RUN pip install --no-cache-dir py-spy pyyaml pandas \ + && pip install --no-cache-dir --ignore-installed --force-reinstall flask + +# Mooncake KV-transfer backend (KV_TRANSFER_BACKEND=mooncake). Mirrors +# docker/sglang_disagg_inference_kvtransfer_overlay.ubuntu.amd.Dockerfile:62-74. +# Default: pip wheel pin; set MOONCAKE_COMMIT for a source build. +ARG MOONCAKE_VERSION=0.3.6.post1 +ARG MOONCAKE_REPO=https://github.com/kvcache-ai/Mooncake.git +ARG MOONCAKE_COMMIT= +RUN set -e; \ + if [[ -n "${MOONCAKE_COMMIT}" ]]; then \ + echo "[full-overlay] Mooncake source ${MOONCAKE_REPO}@${MOONCAKE_COMMIT}"; \ + git clone "${MOONCAKE_REPO}" /tmp/mooncake && cd /tmp/mooncake && \ + git checkout "${MOONCAKE_COMMIT}" && (git submodule update --init --recursive || true); \ + pip install --no-cache-dir --force-reinstall . && rm -rf /tmp/mooncake; \ + elif [[ -n "${MOONCAKE_VERSION}" ]]; then \ + echo "[full-overlay] Mooncake pip mooncake-transfer-engine==${MOONCAKE_VERSION}"; \ + pip install --no-cache-dir --force-reinstall "mooncake-transfer-engine==${MOONCAKE_VERSION}"; \ + else \ + echo "[full-overlay] Mooncake: keeping base image version"; \ + fi + +# Sanity: mooncake must import alongside nixl/rixl (non-fatal — module path may vary). +RUN python3 -c "import mooncake; print('MOONCAKE_OVERLAY_OK')" \ + || echo "MOONCAKE_IMPORT_DIFFERS (verify sglang mooncake import path at runtime)" + +############################################################################### +# 5) Optional: build + install rdma-core v62 from source, for hosts whose RDMA +# stack needs a newer libibverbs/librdmacm/libmlx5 than the base image +# ships. Formerly built at job start by the launcher; baked here (gated by +# ENABLE_RDMA62, empty by default) so the runtime env is not mutated and +# the default build never pays for it. Mirrors the RDMA_CORE_VERSION-gated +# stage in docker/primus_megatron_train_rccl_overlay.ubuntu.amd.Dockerfile. +############################################################################### +ARG ENABLE_RDMA62= +ARG RDMA_VER=v62.0 +RUN if [[ -n "${ENABLE_RDMA62}" ]]; then \ + set -e; \ + git clone --branch "${RDMA_VER}" --depth 1 https://github.com/linux-rdma/rdma-core.git /tmp/rdma-core && \ + cd /tmp/rdma-core && mkdir -p build && cd build && \ + cmake -GNinja -DCMAKE_INSTALL_PREFIX=/usr -DNO_MAN_PAGES=1 .. && \ + ninja && ninja install && ldconfig && \ + rm -rf /tmp/rdma-core; \ + else \ + echo "ENABLE_RDMA62 unset — keeping base image rdma-core"; \ + fi + +# Sanity: rebuilt libibverbs present (only when ENABLE_RDMA62 was set) and +# python still imports (linker not corrupted). +RUN set -e; \ + if [[ -n "${ENABLE_RDMA62}" ]]; then \ + ls -l /usr/lib/libibverbs.so* 2>/dev/null || ls -l /usr/lib/*/libibverbs.so* 2>/dev/null || true; \ + fi; \ + python3 -c "import torch; print(\"RDMA62_TORCH_OK\" if \"${ENABLE_RDMA62}\" else \"RDMA62_SKIPPED_TORCH_OK\", torch.__version__)" diff --git a/models.json b/models.json index af914f4..4ce979e 100644 --- a/models.json +++ b/models.json @@ -1977,6 +1977,29 @@ "args": "--model_repo deepseek-ai/DeepSeek-R1-Distill-Qwen-32B --test_option latency --num_gpu 8 --datatype bfloat16 --dataset random --batch_size 1,8,32 --lat_input_output_len '128:128;128:1024;1024:128;1024:1024'" }, + { + "name": "sglang_disagg_deepseek-r1", + "url": "", + "dockerfile": "docker/sglang_disagg_inference_full_overlay", + "scripts": "scripts/sglang_disagg/run.sh", + "n_gpus": "-1", + "owner": "mad.support@amd.com", + "training_precision": "", + "multiple_results": "perf_sglang-disagg-DeepSeek-R1.csv", + "tags": [ + "pyt", + "sglang", + "inference", + "disaggregated", + "deepseek-r1", + "moe", + "ep-internode", + "rccl-tp", + "mori-a2a" + ], + "timeout": 14400, + "args": "" + }, { "name": "pyt_hy_video", "url": "", @@ -2660,6 +2683,66 @@ "args": "--workload z_image", "multiple_results": "results.csv" }, + { + "name": "primus_pyt_megatron_lm_train_llama-3.1-8b_scaleout", + "url": "", + "dockerfile": "docker/primus_megatron_train_rccl_overlay", + "scripts": "scripts/primus/megatron-lm/run.sh", + "n_gpus": "-1", + "owner": "mad.support@amd.com", + "training_precision": "", + "multiple_results": "perf_primus-megatron-Llama-3.1-8B.csv", + "tags": [ + "pyt", + "pretrain", + "llama3", + "training", + "primus", + "scaleout" + ], + "timeout": -1, + "args": "--model_repo primus_pyt_megatron_lm_train_llama-3.1-8b" + }, + { + "name": "primus_pyt_megatron_lm_train_llama-3.1-70b_scaleout", + "url": "", + "dockerfile": "docker/primus_megatron_train_rccl_overlay", + "scripts": "scripts/primus/megatron-lm/run.sh", + "n_gpus": "-1", + "owner": "mad.support@amd.com", + "training_precision": "", + "multiple_results": "perf_primus-megatron-Llama-3.1-70B.csv", + "tags": [ + "pyt", + "pretrain", + "llama3", + "training", + "primus", + "scaleout" + ], + "timeout": -1, + "args": "--model_repo primus_pyt_megatron_lm_train_llama-3.1-70b" + }, + { + "name": "primus_pyt_megatron_lm_train_llama-3.1-405b_scaleout", + "url": "", + "dockerfile": "docker/primus_megatron_train_rccl_overlay", + "scripts": "scripts/primus/megatron-lm/run.sh", + "n_gpus": "-1", + "owner": "mad.support@amd.com", + "training_precision": "", + "multiple_results": "perf_primus-megatron-Llama-3.1-405B.csv", + "tags": [ + "pyt", + "pretrain", + "llama3", + "training", + "primus", + "scaleout" + ], + "timeout": -1, + "args": "--model_repo primus_pyt_megatron_lm_train_llama-3.1-405b" + }, { "name": "pyt_sglang_disagg_mori_io_llama-3.1-8b", "url": "", diff --git a/scripts/primus/megatron-lm/primus_megatron-lm_benchmark_report.py b/scripts/primus/megatron-lm/primus_megatron-lm_benchmark_report.py index d1d1f3b..f75bb89 100644 --- a/scripts/primus/megatron-lm/primus_megatron-lm_benchmark_report.py +++ b/scripts/primus/megatron-lm/primus_megatron-lm_benchmark_report.py @@ -114,7 +114,25 @@ def find_match(file_path, search_string): print(f"Warning: No matches found for '{search_string}' pattern") return None +def find_match_running_avg(file_path, search_string): + """Return the running-average value (the number AFTER the slash) from the + LAST logged iteration, e.g. '1885.6' from 'tokens/s/GPU): 2355.1/1885.6'. + + Primus logs throughput as 'instant/running_avg'. The instantaneous value + has large run-to-run jitter (MoE routing imbalance, collective stragglers, + kernel warmup), so the trailing running average is a steadier figure, + especially for multi-node scaleout runs. Returns None when no match.""" + with open(file_path, 'r') as f: + content = f.read() + pattern = fr"{re.escape(search_string)}\):\s*\d+\.?\d*/(\d+\.?\d*)" + matches = re.findall(pattern, content) + if matches: + print(f"Found {len(matches)} running-avg matches for '{search_string}', using last: {matches[-1]}") + return matches[-1] + return None + if args.model == "Llama-3.1-8B" or args.model == "Llama-3.1-70B" or \ + args.model == "Llama-3.1-405B" or \ args.model == "Llama-2-7B" or args.model == "Llama-2-70B" or \ args.model == "Mixtral-8x7B" or args.model == "Mixtral-8x22B-proxy" or \ args.model == "DeepSeek-V2-lite" or args.model == "DeepSeek-V3-proxy" or \ @@ -135,9 +153,23 @@ def find_match(file_path, search_string): # Write data for metrics that are found (don't require both to be present) # Skip tokens/s/GPU for Qwen-3-32B if tok_per_s_per_gpu is not None and args.model != "Qwen-3-32B": - data.append({'model': args.model, 'performance': tok_per_s_per_gpu, 'metric': 'tok_per_s_per_gpu', 'mode': args.mode, 'precision': args.precision, 'batch_size': args.batch_size, 'global_batch_size': args.global_batch_size, 'posttrain_type': args.posttrain_type, 'seq_len': args.seq_len, 'device': args.device, 'num_gpus': args.num_gpus}) + data.append({'model': args.model, 'performance': tok_per_s_per_gpu, 'metric': 'tok_per_s_per_gpu', 'mode': args.mode, 'precision': args.precision, 'batch_size': args.batch_size, 'global_batch_size': args.global_batch_size, 'posttrain_type': args.posttrain_type, 'seq_len': args.seq_len, 'device': args.device, 'num_gpus': args.num_gpus, 'run': ''}) if TFLOPS_per_gpu is not None: - data.append({'model': args.model, 'performance': TFLOPS_per_gpu, 'metric': 'TFLOPS_per_gpu', 'mode': args.mode, 'precision': args.precision, 'batch_size': args.batch_size, 'global_batch_size': args.global_batch_size, 'posttrain_type': args.posttrain_type, 'seq_len': args.seq_len, 'device': args.device, 'num_gpus': args.num_gpus}) + data.append({'model': args.model, 'performance': TFLOPS_per_gpu, 'metric': 'TFLOPS_per_gpu', 'mode': args.mode, 'precision': args.precision, 'batch_size': args.batch_size, 'global_batch_size': args.global_batch_size, 'posttrain_type': args.posttrain_type, 'seq_len': args.seq_len, 'device': args.device, 'num_gpus': args.num_gpus, 'run': ''}) + + # Additionally emit the trailing running-average throughput (value after the + # slash in the last logged iteration). Kept as separate, distinctly-named + # rows (metric suffix '_avg', run='avg') so the instantaneous metrics above + # are preserved and never confused with the steady-state average. This is a + # steadier figure for multi-node scaleout runs where per-iteration jitter is + # largest. Skipped for Qwen-3-32B, matching the tokens/s/GPU handling above. + if args.model != "Qwen-3-32B": + tok_avg = find_match_running_avg(input_file, "tokens/s/GPU") + if tok_avg is not None: + data.append({'model': args.model, 'performance': tok_avg, 'metric': 'tok_per_s_per_gpu_avg', 'mode': args.mode, 'precision': args.precision, 'batch_size': args.batch_size, 'global_batch_size': args.global_batch_size, 'posttrain_type': args.posttrain_type, 'seq_len': args.seq_len, 'device': args.device, 'num_gpus': args.num_gpus, 'run': 'avg'}) + tflops_avg = find_match_running_avg(input_file, "TFLOP/s/GPU") + if tflops_avg is not None: + data.append({'model': args.model, 'performance': tflops_avg, 'metric': 'TFLOPS_per_gpu_avg', 'mode': args.mode, 'precision': args.precision, 'batch_size': args.batch_size, 'global_batch_size': args.global_batch_size, 'posttrain_type': args.posttrain_type, 'seq_len': args.seq_len, 'device': args.device, 'num_gpus': args.num_gpus, 'run': 'avg'}) if not os.path.exists(output_file) or os.stat(output_file).st_size == 0: mode = 'w' # Write if file doesn't exist or is empty @@ -146,7 +178,7 @@ def find_match(file_path, search_string): with open(output_file, mode=mode, newline='') as file: print("Preparing to write performance data...") print("Data: ", data) - writer = csv.DictWriter(file, fieldnames=['model','performance','metric','mode','precision','batch_size','global_batch_size','posttrain_type','seq_len','device','num_gpus']) + writer = csv.DictWriter(file, fieldnames=['model','performance','metric','mode','precision','batch_size','global_batch_size','posttrain_type','seq_len','device','num_gpus','run']) if mode == 'w': writer.writeheader() writer.writerows(data) diff --git a/scripts/primus/megatron-lm/primus_megatron-lm_benchmark_report.sh b/scripts/primus/megatron-lm/primus_megatron-lm_benchmark_report.sh index a4604f0..fd6e891 100755 --- a/scripts/primus/megatron-lm/primus_megatron-lm_benchmark_report.sh +++ b/scripts/primus/megatron-lm/primus_megatron-lm_benchmark_report.sh @@ -140,8 +140,17 @@ export PATH="${ROCM_BIN}:${PATH}" # Run rocminfo and grep for "AMD Instinct" DEVICE=$("${ROCM_BIN}/rocminfo" | grep "AMD Instinct" | head -n1 | awk '{print $5}') echo "DEVICE found: $DEVICE" +# Fall back to amd-smi when the rocminfo output format changes or rocminfo is +# unavailable (some newer ROCm/scaleout stacks). +if [[ -z "$DEVICE" && -x "${ROCM_BIN}/amd-smi" ]]; then + DEVICE=$("${ROCM_BIN}/amd-smi" 2>/dev/null | awk '/AMD Instinct/ {print $5; exit}') +fi if [[ -z "$DEVICE" || ("$DEVICE" != "MI300X" && "$DEVICE" != "MI355X") ]]; then ARCH=$("${ROCM_BIN}/rocminfo" | grep -o 'gfx942\|gfx950' | head -n 1 | tr -d '[:space:]') + # Fall back to the arch injected by madengine when rocminfo cannot report it. + if [[ -z "$ARCH" && -n "${MAD_SYSTEM_GPU_ARCHITECTURE:-}" ]]; then + ARCH=$(printf '%s' "${MAD_SYSTEM_GPU_ARCHITECTURE}" | tr -d '[:space:]') + fi case "$ARCH" in "gfx942") DEVICE="MI300X" ;; "gfx950") DEVICE="MI355X" ;; @@ -161,13 +170,83 @@ else CONFIG_DEVICE="$DEVICE" fi -# Set common environment variables -export NNODES=1 -export CPUS_PER_TASK=128 +# Set common environment variables. +# Keep compatibility with SLURM/madengine distributed env and only fall back to +# single-node defaults: single-node runs behave exactly as before, while +# multi-node scaleout runs pick up the injected topology. +NNODES="${NNODES:-1}" +export NNODES +export CPUS_PER_TASK="${CPUS_PER_TASK:-128}" export HSA_NO_SCRATCH_RECLAIM=1 export NVTE_CK_USES_BWD_V3=1 -GPUS_PER_NODE=8 # default to 8 GPUs per node -NUM_GPUS=8 +GPUS_PER_NODE="${GPUS_PER_NODE:-${NPROC_PER_NODE:-8}}" +if ! [[ "$NNODES" =~ ^[0-9]+$ ]]; then NNODES=1; fi +if ! [[ "$GPUS_PER_NODE" =~ ^[0-9]+$ ]]; then GPUS_PER_NODE=8; fi +NUM_GPUS=$((NNODES * GPUS_PER_NODE)) +echo "[INFO] Distributed params: NNODES=$NNODES GPUS_PER_NODE=$GPUS_PER_NODE NUM_GPUS=$NUM_GPUS" + +# Round global batch size up to a multiple of micro_batch_size * world_size so +# the global batch stays divisible across an arbitrary number of ranks. Note: +# this uses total world size (NUM_GPUS), not the actual Megatron data-parallel +# degree, which can be smaller than world size when TP/PP/CP/EP > 1. +normalize_global_batch_size() { + local mbs="$1" + local gbs="$2" + local world_size="$3" + local unit=$((mbs * world_size)) + if (( unit <= 0 )); then + echo "$gbs" + return + fi + if (( gbs % unit == 0 )); then + echo "$gbs" + return + fi + echo "$(( ((gbs + unit - 1) / unit) * unit ))" +} + +# Pure helper: the possibly-adjusted global batch size for NUM_GPUS>8 (world +# size scaleout), or the original gbs unchanged for single-node runs / values +# that don't parse as plain integers. Callers must reassign their own GBS +# variable to this so the value used for the actual run (and reported to the +# perf CSV) always agree -- see scaleout_gbs_override below. +effective_global_batch_size() { + local mbs="$1" + local gbs="$2" + if (( NUM_GPUS > 8 )) && [[ "$mbs" =~ ^[0-9]+$ && "$gbs" =~ ^[0-9]+$ ]]; then + normalize_global_batch_size "$mbs" "$gbs" "$NUM_GPUS" + else + echo "$gbs" + fi +} + +# Emit the extra CLI flags needed to keep a run valid on multiple nodes: when +# NUM_GPUS > 8 the config global_batch_size (tuned for a single 8-GPU node) is +# renormalized against total world size and passed as an override. Prints +# nothing for single-node runs, so behavior there is byte-for-byte identical +# to before. Callers should also run `GBS=$(effective_global_batch_size ...)` +# so the reported GBS in the perf CSV matches the value actually used. +scaleout_gbs_override() { + local mbs="$1" + local gbs="$2" + local adjusted + adjusted=$(effective_global_batch_size "$mbs" "$gbs") + if [[ "$adjusted" != "$gbs" ]]; then + echo "[INFO] Adjusted global batch size for distributed run: ${gbs} -> ${adjusted} (MBS=${mbs}, NUM_GPUS=${NUM_GPUS})" >&2 + printf -- '--global_batch_size %s' "$adjusted" + fi +} + +# Launch a single Primus pretrain run. Every rank must execute an identical +# command so the model shape stays consistent across nodes; a mismatch makes +# multi-node ranks rendezvous with different shapes and deadlocks the collectives. +run_primus() { + local config="$1"; shift + bash runner/primus-cli direct \ + --log_file "/tmp/primus_$MODEL_REPO.log" \ + -- train pretrain \ + --config "$config" "$@" 2>&1 | tee "$TRAIN_LOG" +} cd /workspace/Primus @@ -180,26 +259,19 @@ if [ "$MODEL_REPO" == "Llama-3.1-8B" ]; then MBS=$(grep -E '^\s*micro_batch_size:' $EXP | head -n1 | awk '{print $2}' | tr -d '\r') GBS=$(grep -E '^\s*global_batch_size:' $EXP | head -n1 | awk '{print $2}' | tr -d '\r') echo "[INFO] Extracted MBS=$MBS, GBS=$GBS from config: $EXP" + GBS_OVERRIDE=$(scaleout_gbs_override "$MBS" "$GBS") + GBS=$(effective_global_batch_size "$MBS" "$GBS") if [[ "$DEVICE" == "MI355X" || "$DEVICE" == "MI350X" ]]; then if [[ "$DATATYPE" == "MXFP4" ]]; then - NVTE_USE_CAST_TRANSPOSE_TRITON=0 bash runner/primus-cli direct \ - --log_file /tmp/primus_$MODEL_REPO.log \ - -- train pretrain \ - --config $EXP 2>&1 | tee -a $TRAIN_LOG + NVTE_USE_CAST_TRANSPOSE_TRITON=0 run_primus "$EXP" $GBS_OVERRIDE else - bash runner/primus-cli direct \ - --log_file /tmp/primus_$MODEL_REPO.log \ - -- train pretrain \ - --config $EXP 2>&1 | tee -a $TRAIN_LOG + run_primus "$EXP" $GBS_OVERRIDE fi elif [[ "$DEVICE" == "MI300X" || "$DEVICE" == "MI325X" ]]; then if [[ "$DATATYPE" == "MXFP8" || "$DATATYPE" == "MXFP4" ]]; then echo "Error: Datatype $DATATYPE is not supported for $MODEL_REPO on $DEVICE. Only supported on MI355X/MI350X." else - bash runner/primus-cli direct \ - --log_file /tmp/primus_$MODEL_REPO.log \ - -- train pretrain \ - --config $EXP 2>&1 | tee -a $TRAIN_LOG + run_primus "$EXP" $GBS_OVERRIDE fi fi if [ -f "$TRAIN_LOG" ]; then @@ -221,26 +293,19 @@ elif [ "$MODEL_REPO" == "Llama-3.1-70B" ]; then MBS=$(grep -E '^\s*micro_batch_size:' $EXP | head -n1 | awk '{print $2}' | tr -d '\r') GBS=$(grep -E '^\s*global_batch_size:' $EXP | head -n1 | awk '{print $2}' | tr -d '\r') echo "[INFO] Extracted MBS=$MBS, GBS=$GBS from config: $EXP" + GBS_OVERRIDE=$(scaleout_gbs_override "$MBS" "$GBS") + GBS=$(effective_global_batch_size "$MBS" "$GBS") if [[ "$DEVICE" == "MI355X" || "$DEVICE" == "MI350X" ]]; then if [ "$DATATYPE" == "FP8" ]; then - bash runner/primus-cli direct \ - --log_file /tmp/primus_$MODEL_REPO.log \ - -- train pretrain \ - --config $EXP 2>&1 | tee -a $TRAIN_LOG + run_primus "$EXP" $GBS_OVERRIDE elif [ "$DATATYPE" == "BF16" ]; then - bash runner/primus-cli direct \ - --log_file /tmp/primus_$MODEL_REPO.log \ - -- train pretrain \ - --config $EXP 2>&1 | tee -a $TRAIN_LOG + run_primus "$EXP" $GBS_OVERRIDE fi elif [[ "$DEVICE" == "MI300X" || "$DEVICE" == "MI325X" ]]; then if [ "$DATATYPE" == "FP8" ]; then echo "Error: Datatype FP8 is not supported for $MODEL_REPO on $DEVICE. Only BF16 is supported." elif [ "$DATATYPE" == "BF16" ]; then - bash runner/primus-cli direct \ - --log_file /tmp/primus_$MODEL_REPO.log \ - -- train pretrain \ - --config $EXP 2>&1 | tee -a $TRAIN_LOG + run_primus "$EXP" $GBS_OVERRIDE fi fi if [ -f "$TRAIN_LOG" ]; then @@ -287,6 +352,40 @@ elif [ "$MODEL_REPO" == "Llama-3.1-70B-proxy" ]; then echo "$MODEL_REPO,,TFLOPS_per_gpu,$MODE,$DATATYPE,$MBS,$GBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG fi +elif [ "$MODEL_REPO" == "Llama-3.1-405B" ]; then + echo "[INFO] $MODEL_REPO TRAINING" # BF16 training only + export EXP=examples/megatron/configs/$CONFIG_DEVICE/llama3.1_405B-$DATATYPE-pretrain.yaml + rm -f "$TRAIN_LOG" + + SEQ_LEN=8192 + if [[ -f "$EXP" ]]; then + MBS=$(grep -E '^\s*micro_batch_size:' $EXP | head -n1 | awk '{print $2}' | tr -d '\r') + GBS=$(grep -E '^\s*global_batch_size:' $EXP | head -n1 | awk '{print $2}' | tr -d '\r') + echo "[INFO] Extracted MBS=$MBS, GBS=$GBS from config: $EXP" + fi + if [[ "$DATATYPE" == "FP8" ]]; then + echo "Error: Datatype FP8 is not supported for $MODEL_REPO. Only BF16 is supported." + elif [[ ! -f "$EXP" ]]; then + echo "Error: Config file not found: $EXP" + echo "Hint: add llama3.1_405B-$DATATYPE-pretrain.yaml in Primus configs." + else + GBS_OVERRIDE=$(scaleout_gbs_override "$MBS" "$GBS") + GBS=$(effective_global_batch_size "$MBS" "$GBS") + run_primus "$EXP" $GBS_OVERRIDE + fi + MBS="${MBS:-NA}" + GBS="${GBS:-NA}" + if [ -f "$TRAIN_LOG" ]; then + echo "[INFO] Benchmarking" + python3 $perf_script --model $MODEL_REPO --input $TRAIN_LOG --output $PERF_LOG --mode $MODE --precision $DATATYPE --batch_size $MBS --global_batch_size $GBS --seq_len $SEQ_LEN --device $DEVICE --num_gpus $NUM_GPUS + rm $TRAIN_LOG + else + echo "[INFO] Training log not found - configuration not supported." + echo "model,performance,metric,mode,precision,batch_size,global_batch_size,seq_len,device,num_gpus" > $PERF_LOG + echo "$MODEL_REPO,,tok_per_s_per_gpu,$MODE,$DATATYPE,$MBS,$GBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG + echo "$MODEL_REPO,,TFLOPS_per_gpu,$MODE,$DATATYPE,$MBS,$GBS,$SEQ_LEN,$DEVICE,$NUM_GPUS" >> $PERF_LOG + fi + elif [ "$MODEL_REPO" == "Llama-3.3-70B" ]; then echo "[INFO] $MODEL_REPO TRAINING" export EXP=examples/megatron/configs/$CONFIG_DEVICE/llama3.3_70B-$DATATYPE-pretrain.yaml @@ -455,6 +554,15 @@ elif [ "$MODEL_REPO" == "DeepSeek-V3-proxy" ]; then SEQ_LEN=4096 # Set MBS/GBS through command line for proxy models echo "[INFO] Extracted MBS=$MBS, GBS=$GBS from config: $EXP" + # DeepEP (Primus-Turbo) is opt-in via PRIMUS_USE_DEEPEP=1 and applies to every + # supported device/precision branch below. DeepEP is incompatible with + # moe_shared_expert_overlap (Primus ROCm validate_args asserts it), so disable + # that overlap when enabling DeepEP. Default (unset/0) keeps alltoall over RCCL. + DEEPEP_ARGS="" + if [[ "${PRIMUS_USE_DEEPEP:-0}" == "1" ]]; then + DEEPEP_ARGS="--use_turbo_deepep true --moe_shared_expert_overlap false" + echo "[INFO] DeepEP enabled via PRIMUS_USE_DEEPEP=1 -> ${DEEPEP_ARGS}" + fi if [[ "$DEVICE" == "MI355X" || "$DEVICE" == "MI350X" ]]; then if [ "$DATATYPE" == "FP8" ]; then echo "Error: Datatype FP8 is not supported for $MODEL_REPO on $DEVICE. Only BF16 is supported." @@ -478,7 +586,7 @@ elif [ "$MODEL_REPO" == "DeepSeek-V3-proxy" ]; then --recompute_layer_ids null \ --overlap_grad_reduce false \ --overlap_param_gather false \ - --gradient_accumulation_fusion false 2>&1 | tee $TRAIN_LOG + --gradient_accumulation_fusion false $DEEPEP_ARGS 2>&1 | tee $TRAIN_LOG fi elif [[ "$DEVICE" == "MI300X" || "$DEVICE" == "MI325X" ]]; then if [ "$DATATYPE" == "FP8" ]; then @@ -489,7 +597,7 @@ elif [ "$MODEL_REPO" == "DeepSeek-V3-proxy" ]; then bash runner/primus-cli direct \ --log_file /tmp/primus_$MODEL_REPO.log \ -- train pretrain \ - --config $EXP --num_layers 3 --moe_layer_freq 1 --micro_batch_size $MBS --global_batch_size $GBS 2>&1 | tee $TRAIN_LOG + --config $EXP --num_layers 3 --moe_layer_freq 1 --micro_batch_size $MBS --global_batch_size $GBS $DEEPEP_ARGS 2>&1 | tee $TRAIN_LOG fi fi if [ -f "$TRAIN_LOG" ]; then diff --git a/scripts/primus/megatron-lm/run.sh b/scripts/primus/megatron-lm/run.sh index 6d8267f..afe4254 100755 --- a/scripts/primus/megatron-lm/run.sh +++ b/scripts/primus/megatron-lm/run.sh @@ -40,6 +40,8 @@ if [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_llama-3.1-8b" ]]; then model="Llama-3.1-8B" elif [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_llama-3.1-70b" ]]; then model="Llama-3.1-70B" +elif [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_llama-3.1-405b" ]]; then + model="Llama-3.1-405B" elif [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_llama-3.1-70b-proxy" ]]; then model="Llama-3.1-70B-proxy" elif [[ "$MODEL_REPO" == "primus_pyt_megatron_lm_train_llama-3.3-70b" ]]; then @@ -114,8 +116,17 @@ export PATH="${ROCM_BIN}:${PATH}" # Detect device DEVICE=$("${ROCM_BIN}/rocminfo" | grep "AMD Instinct" | head -n1 | awk '{print $5}') +# Fall back to amd-smi when the rocminfo output format changes or rocminfo is +# unavailable (some newer ROCm/scaleout stacks). +if [[ -z "$DEVICE" && -x "${ROCM_BIN}/amd-smi" ]]; then + DEVICE=$("${ROCM_BIN}/amd-smi" 2>/dev/null | awk '/AMD Instinct/ {print $5; exit}') +fi if [ -z "$DEVICE" ]; then ARCH=$("${ROCM_BIN}/rocminfo" | grep -o 'gfx942\|gfx950' | head -n 1 | tr -d '[:space:]') + # Fall back to the arch injected by madengine when rocminfo cannot report it. + if [[ -z "$ARCH" && -n "${MAD_SYSTEM_GPU_ARCHITECTURE:-}" ]]; then + ARCH=$(printf '%s' "${MAD_SYSTEM_GPU_ARCHITECTURE}" | tr -d '[:space:]') + fi case "$ARCH" in "gfx942") DEVICE="MI300X" ;; "gfx950") DEVICE="MI355X" ;; diff --git a/scripts/sglang_disagg/README.MD b/scripts/sglang_disagg/README.MD index 0540b19..c16722f 100644 --- a/scripts/sglang_disagg/README.MD +++ b/scripts/sglang_disagg/README.MD @@ -54,12 +54,56 @@ docker build --build-arg MORI_COMMIT= -t sglang_disagg_pd_image -f sglang_d | File | Description | |------|-------------| -| `run_xPyD_models.slurm` | SLURM script to launch docker containers on all nodes via `sbatch` | -| `sglang_disagg_mori_io_ep.sh` | Container entrypoint — starts prefill/decode servers, proxy, and benchmark | +| `run.sh` | **madengine entrypoint** (models.json `scripts` target). Bridges madengine's `sglang-disagg` launcher env to the launchers below | +| `run_xPyD_models.slurm` | Standalone SLURM script to launch docker containers on all nodes via `sbatch` | +| `sglang_disagg_mori_io_ep.sh` | Unified container entrypoint (MoRI EP / Mooncake / NIXL, selected via `KV_TRANSFER_BACKEND`) — starts prefill/decode servers, proxy, and benchmark | | `models.yaml` | Model-specific CLI flags for all supported models | | `mori_ep_env.sh` | RDMA/NCCL/Gloo environment variables | | `benchmark_xPyD.sh` | Concurrency sweep benchmark using sglang bench_serving | | `benchmark_parser.py` | Log parser for CONCURRENCY benchmark logs | +| `parse_to_csv.py` | Log parser for CONCURRENCY benchmark logs (emits summary CSV and madengine perf CSV) | + +## Running via madengine (`madengine run`) + +This workload is registered in `models.json` (e.g. `sglang_disagg_deepseek-r1`) and +runs through madengine's built-in **`sglang-disagg`** launcher. The launcher runs one +container per node and exports the cluster topology; `run.sh` maps it to the launchers +above: + +| madengine env | bridged to | +|---------------|-----------| +| `SGLANG_NODE_RANK` | `NODE_RANK` (0=proxy, 1..xP=prefill, rest=decode) | +| `SGLANG_DISAGG_PREFILL_NODES` / `..._DECODE_NODES` | `xP` / `yD` | +| `SGLANG_NODE_IPS` | `IPADDRS` (+ `MASTER_ADDR` = rank-0 IP) | +| `SGLANG_TP_SIZE` | `IO_EP_TP_SIZE` / per-server `--tp-size` | + +Model + transport come from deployment `env_vars` (or `args`): `MODEL_NAME` (short key +in `models.yaml`), `MODEL_PATH` (mounted weights), `RUN_MORI` (1=MoRI EP default, +0=Mooncake/NIXL), `DP_MODE` (1=DP-attention EP for DeepSeek-V3/R1 inter-node). + +Deploy with an additional-context config carrying `slurm` + `distributed.launcher: +"sglang-disagg"` (+ `distributed.sglang_disagg.{prefill_nodes,decode_nodes}` for a +custom split) + `env_vars`; see the Confluence package for a ready 5-node EP16 manifest. + +## Docker overlays (substitutable component versions) + +The base image (`docker/sglang_disagg_inference.ubuntu.amd.Dockerfile`) bakes RCCL, +MoRI and Mooncake/NIXL. To test *specific* versions of these components without +rebuilding the base, use the merged full overlay, which layers RCCL, MoRI and the +KV-transfer (Mooncake/NIXL) steps on top of `BASE_DOCKER` in a single image +(`base -> rccl -> mori -> kvtransfer`): + +| Overlay Dockerfile | Substitutes | Key build args | +|--------------------|-------------|----------------| +| `docker/sglang_disagg_inference_full_overlay.ubuntu.amd.Dockerfile` | RCCL + MoRI + Mooncake/NIXL (merged) | `RCCL_REPO/RCCL_BRANCH/RCCL_COMMIT`, `MORI_REPO/MORI_BRANCH/MORI_COMMIT` or `MORI_WHEEL_URL`, `MOONCAKE_VERSION` / `NIXL_VERSION`, `ENABLE_RDMA62` (set to `1` on hosts needing a newer rdma-core than the base image ships, to additionally build rdma-core v62 from source) | + +> **RCCL note:** a newer rocm-systems RCCL build links `librccl` without +> the `librocm_smi64.so` `DT_NEEDED` the base carried. Since the overlay is first +> on `LD_LIBRARY_PATH`, this would break `import torch` +> (`libtorch_hip.so: undefined symbol: rsmi_init`). The overlay Dockerfile +> re-adds the dependency (`patchelf --add-needed`) and sanity-checks `import +> torch` with the overlay librccl resolved first, so the full +> `base -> rccl -> mori -> kvtransfer` chain stays importable. ## Quick Start @@ -72,6 +116,7 @@ export xP=1 export yD=1 export MODEL_NAME=Llama-3.1-8B-Instruct export RUN_MORI=1 # MoRI (default). Set RUN_MORI=0 for Mooncake (KV_TRANSFER_BACKEND=mooncake) +export KV_TRANSFER_BACKEND=mori # Validated/default. Use nixl only with an image that provides nixl._api. # num_nodes = xP + yD sbatch -N 2 -n 2 --nodelist= run_xPyD_models.slurm @@ -115,7 +160,7 @@ Logs are written to `${LOG_PATH}/${SLURM_JOB_ID}/`: Parse benchmark results: ```bash -python3 benchmark_parser.py /benchmark_XXX_CONCURRENCY.log +python3 parse_to_csv.py /benchmark_XXX_CONCURRENCY.log -o results.csv ``` Smoke test from the proxy node: diff --git a/scripts/sglang_disagg/benchmark_xPyD.sh b/scripts/sglang_disagg/benchmark_xPyD.sh index 05e6e3c..c28d18f 100644 --- a/scripts/sglang_disagg/benchmark_xPyD.sh +++ b/scripts/sglang_disagg/benchmark_xPyD.sh @@ -1,61 +1,142 @@ #!/bin/bash +set -uo pipefail timestamp=$(date "+%Y%m%d_%H%M%S") -LOG="/run_logs/${SLURM_JOB_ID}/benchmark_${SLURM_JOB_ID}_${timestamp}_xP${xP}_yD${yD}_$MODEL_NAME" -echo "==== Benchmark Serving Concurrency Sweep Test ${LOG} ===== " -echo "UTC Time: $(TZ=UTC date '+%Y-%m-%d %H:%M:%S %Z')" | tee -a ${LOG}_CONCURRENCY.log >/dev/null -echo "PST Time: $(TZ=America/Los_Angeles date '+%Y-%m-%d %H:%M:%S %Z')" | tee -a ${LOG}_CONCURRENCY.log >/dev/null - -sleep 60 -echo "Test run:" | tee -a ${LOG}_CONCURRENCY.log >/dev/null -python3 -m sglang.bench_serving \ - --model $MODEL_PATH \ - --backend sglang \ - --host 127.0.0.1 \ - --port 2322 \ - --dataset-name random \ - --random-input 1024 \ - --random-output 1024\ - --random-range-ratio 1.0 \ - --max-concurrency 512 \ - --num-prompt 1024 \ - --pd-separated \ - 2>&1 | tee -a ${LOG}_CONCURRENCY.log >/dev/null -echo "" -CON="8 16 32 64 128 256 512" -# ISL/OSL combinations — override via BENCHMARK_COMBINATIONS env var (space-separated, e.g. "1024/1024 8192/1024") +RUN_LOG_JOB_ID="${SLURM_JOB_ID:-0}" +# Honor RUN_LOG_DIR exported by run.sh (single source of truth for the log +# location, including its /run_logs-vs-fallback decision). Only derive the +# default when this script is invoked standalone without run.sh in the env. +RUN_LOG_DIR="${RUN_LOG_DIR:-/run_logs/${RUN_LOG_JOB_ID}}" +mkdir -p "$RUN_LOG_DIR" 2>/dev/null || true + +LOG="${RUN_LOG_DIR}/benchmark_${RUN_LOG_JOB_ID}_${timestamp}_xP${xP}_yD${yD}_${MODEL_NAME}" +LOG_FILE="${LOG}_CONCURRENCY.log" + +BENCHMARK_ITR="${BENCHMARK_ITR:-1}" +if ! [[ "$BENCHMARK_ITR" =~ ^[0-9]+$ ]] || [[ "$BENCHMARK_ITR" -lt 1 ]]; then + echo "ERROR: BENCHMARK_ITR must be a positive integer, got '${BENCHMARK_ITR}'" >&2 + exit 1 +fi + +CON="${BENCHMARK_CONCURRENCY_LEVELS:-8 16 32 64 128 256 512}" +# ISL/OSL combinations; override with e.g. BENCHMARK_COMBINATIONS="1024/1024 8192/1024". IFS=' ' read -ra COMBINATIONS <<< "${BENCHMARK_COMBINATIONS:-1024/1024 8192/1024}" -echo "Benchmarking iterations: $BENCHMARK_ITR" | tee -a ${LOG}_CONCURRENCY.log >/dev/null -for i in {1..$BENCHMARK_ITR}; do - sleep 60 - echo "RUNNING: the benchserving script for iter: $i" | tee -a ${LOG}_CONCURRENCY.log >/dev/null + +# --- Per-point resilience (kept from #328) ----------------------------------- +# A single transient gateway circuit-open / server recovery hiccup should not +# nuke a multi-hour sweep. Retry the whole point after a cooldown; a persistently +# failing point still fails the run (fail-fast preserved). +# Set BENCHMARK_POINT_RETRIES=0 to restore single-attempt behavior. +BENCHMARK_FAIL_FAST="${BENCHMARK_FAIL_FAST:-1}" +BENCHMARK_POINT_RETRIES="${BENCHMARK_POINT_RETRIES:-2}" +BENCHMARK_RETRY_COOLDOWN_SECONDS="${BENCHMARK_RETRY_COOLDOWN_SECONDS:-45}" + +log_msg() { + echo "$*" | tee -a "$LOG_FILE" >/dev/null +} + +run_serving_bench() { + local isl="$1" + local osl="$2" + local con="$3" + local prompts="$4" + + # Extra context on a separate INFO line; the parseable marker below stays in + # the exact shape parse_to_csv.py expects: "RUNNING: prompts isl X osl Y con Z". + log_msg "INFO: model=${MODEL_NAME} xP=${xP} yD=${yD} job=${RUN_LOG_JOB_ID} prompts=${prompts}" + log_msg "RUNNING: prompts isl ${isl} osl ${osl} con ${con}" + python3 -m sglang.bench_serving \ + --model "$MODEL_PATH" \ + --backend sglang \ + --host 127.0.0.1 \ + --port 2322 \ + --dataset-name random \ + --random-input "$isl" \ + --random-output "$osl" \ + --random-range-ratio 1.0 \ + --max-concurrency "$con" \ + --num-prompt "$prompts" \ + --pd-separated \ + 2>&1 | tee -a "$LOG_FILE" >/dev/null +} + +log_msg "==== Benchmark Serving Concurrency Sweep Test ${LOG} =====" +log_msg "UTC Time: $(TZ=UTC date '+%Y-%m-%d %H:%M:%S %Z')" +log_msg "PST Time: $(TZ=America/Los_Angeles date '+%Y-%m-%d %H:%M:%S %Z')" +log_msg "Benchmarking iterations: ${BENCHMARK_ITR}" +log_msg "Benchmarking combinations: ${COMBINATIONS[*]}" +log_msg "Benchmarking concurrency levels: ${CON}" + +benchmark_failures=0 + +# Optional warmup / precheck. Runs BEFORE the "iter: 1" marker, so parse_to_csv.py +# (which ignores everything prior to that marker) never counts it as a result. +if [[ "${BENCHMARK_PRECHECK:-0}" == "1" ]]; then + sleep "${BENCHMARK_START_DELAY_SECONDS:-10}" + log_msg "Test run:" + if ! run_serving_bench \ + "${BENCHMARK_PRECHECK_ISL:-128}" \ + "${BENCHMARK_PRECHECK_OSL:-16}" \ + "${BENCHMARK_PRECHECK_CONCURRENCY:-1}" \ + "${BENCHMARK_PRECHECK_PROMPTS:-1}"; then + benchmark_failures=$((benchmark_failures + 1)) + log_msg "ERROR: benchmark precheck failed" + [[ "$BENCHMARK_FAIL_FAST" == "1" ]] && exit 1 + fi +fi + +for ((i = 1; i <= BENCHMARK_ITR; i++)); do + sleep "${BENCHMARK_ITERATION_DELAY_SECONDS:-60}" + log_msg "RUNNING: the benchserving script for iter: $i" for combo in "${COMBINATIONS[@]}"; do - IFS="/" read -r isl osl <<< "$combo" - for con in $CON; do - p_con=$(($con * 2)) - if [ "$p_con" -lt 16 ]; then - p_con=16 - fi - echo "RUNNING: prompts $prompts isl $isl osl $osl con $con" | tee -a ${LOG}_CONCURRENCY.log >/dev/null - python3 -m sglang.bench_serving \ - --model $MODEL_PATH \ - --backend sglang \ - --host 127.0.0.1 \ - --port 2322 \ - --dataset-name random \ - --random-input $isl \ - --random-output $osl \ - --random-range-ratio 1.0 \ - --max-concurrency $con \ - --num-prompt $p_con \ - --pd-separated \ - 2>&1 | tee -a ${LOG}_CONCURRENCY.log >/dev/null - - sleep 10 + IFS="/" read -r isl osl <<< "$combo" + for con in $CON; do + p_con=$((con * 2)) + if [[ "$p_con" -lt 16 ]]; then + p_con=16 + fi + total_attempts=$((BENCHMARK_POINT_RETRIES + 1)) + point_ok=0 + for ((attempt = 1; attempt <= total_attempts; attempt++)); do + if run_serving_bench "$isl" "$osl" "$con" "$p_con"; then + point_ok=1 + break + fi + if [[ "$attempt" -lt "$total_attempts" ]]; then + log_msg "WARN: bench_serving failed (attempt ${attempt}/${total_attempts}) for iter=${i} isl=${isl} osl=${osl} con=${con}; retrying after ${BENCHMARK_RETRY_COOLDOWN_SECONDS}s (transient circuit-open / server recovery)" + sleep "${BENCHMARK_RETRY_COOLDOWN_SECONDS}" + fi + done + if [[ "$point_ok" -ne 1 ]]; then + benchmark_failures=$((benchmark_failures + 1)) + log_msg "ERROR: bench_serving failed after ${total_attempts} attempt(s) for iter=${i} isl=${isl} osl=${osl} con=${con}" + [[ "$BENCHMARK_FAIL_FAST" == "1" ]] && break 3 + fi + + sleep "${BENCHMARK_BETWEEN_RUN_SECONDS:-10}" done done done -python3 parse_to_csv.py ${LOG}_CONCURRENCY.log -o ${LOG}_CONCURRENCY.csv \ - --perf-csv /run_logs/${SLURM_JOB_ID}/perf.csv \ - --model-name "${MODEL_NAME}" \ - 2>&1 | tee -a ${LOG}_CONCURRENCY.log >/dev/null + +log_msg "==== Benchmark Serving Concurrency End Time ${LOG} =====" +log_msg "UTC Time: $(TZ=UTC date '+%Y-%m-%d %H:%M:%S %Z')" +log_msg "PST Time: $(TZ=America/Los_Angeles date '+%Y-%m-%d %H:%M:%S %Z')" + +# --- Finalization: aligned with develop (parse_to_csv.py + MAD_OUTPUT_CSV) ---- +# Write the madengine multiple_results CSV to the container cwd under the name +# madengine harvests (MAD_OUTPUT_CSV, e.g. perf_sglang-disagg-.csv); keep +# a copy under /run_logs for post-mortem inspection. +PERF_CSV="${MAD_OUTPUT_CSV:-perf_sglang-disagg-${MODEL_NAME}.csv}" +if ! python3 parse_to_csv.py "${LOG}_CONCURRENCY.log" -o "${LOG}_CONCURRENCY.csv" \ + --perf-csv "${PERF_CSV}" \ + --model-name "${MODEL_NAME}" \ + 2>&1 | tee -a "$LOG_FILE" >/dev/null; then + benchmark_failures=$((benchmark_failures + 1)) + log_msg "ERROR: parse_to_csv.py failed to produce ${PERF_CSV}" +fi +cp -f "${PERF_CSV}" "${RUN_LOG_DIR}/" 2>/dev/null || true + +if [[ "$benchmark_failures" -ne 0 ]]; then + log_msg "ERROR: benchmark completed with ${benchmark_failures} failure(s)" + exit 1 +fi diff --git a/scripts/sglang_disagg/ip_rendezvous.py b/scripts/sglang_disagg/ip_rendezvous.py new file mode 100755 index 0000000..a124ad7 --- /dev/null +++ b/scripts/sglang_disagg/ip_rendezvous.py @@ -0,0 +1,150 @@ +#!/usr/bin/env python3 +"""In-container node-IP rendezvous for disaggregated P/D launchers. + +madengine forwards MASTER_ADDR/NODE_RANK/NNODES into each container but NOT the +rank-ordered node IP list. This helper rebuilds that list with stdlib only: + + rank 0 -> tiny TCP server: collects {rank: ip} from peers, then broadcasts + the comma-separated, rank-ordered IP list back to everyone. + rank>0 -> connects to MASTER_ADDR, reports its own IP, then blocks until + rank 0 broadcasts the full list. + +On success the rank-ordered "ip0,ip1,..." list is printed to stdout (exit 0); +on failure an empty line is printed (exit 2), so callers can use +`out="$(ip_rendezvous.py ... || true)"` and test for emptiness. + +Usage: + ip_rendezvous.py +Env: + IP_SYNC_TIMEOUT total rendezvous budget in seconds (default 1800). +""" +import json +import os +import socket +import sys +import time + + +def recv_line(conn): + buf = b"" + while b"\n" not in buf: + chunk = conn.recv(4096) + if not chunk: + break + buf += chunk + return buf.split(b"\n", 1)[0].decode("utf-8", "replace") if buf else None + + +def serve(nnodes, host_ip, port, token, deadline): + """rank-0 path: collect peer IPs, broadcast the rank-ordered list.""" + by_rank = {0: host_ip} + # rank -> connection. Workers reconnect while we wait for slow nodes, so we + # keep only the latest socket per rank (closing the stale one) to avoid + # broadcasting over half-open sockets. + conns = {} + srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + srv.bind(("0.0.0.0", port)) + srv.listen(max(8, nnodes + 2)) + srv.settimeout(1.0) + try: + while len(by_rank) < nnodes and time.time() < deadline: + try: + conn, _ = srv.accept() + except socket.timeout: + continue + conn.settimeout(5.0) + line = recv_line(conn) + if not line: + conn.close() + continue + try: + msg = json.loads(line) + except Exception: + conn.close() + continue + if msg.get("token") != token: + conn.close() + continue + wrank = int(msg.get("rank", -1)) + wip = str(msg.get("ip", "")).strip() + # rank 0 is the server itself (by_rank[0] is already seeded with our + # own host_ip above); only accept reports from peers, rank 1..nnodes-1, + # so a buggy/spoofed rank-0 message can't clobber our own entry. + if 1 <= wrank < nnodes and wip: + old = conns.get(wrank) + if old is not None: + try: + old.close() + except Exception: + pass + by_rank[wrank] = wip + conns[wrank] = conn + else: + conn.close() + if len(by_rank) != nnodes: + return None + ipaddrs = ",".join(by_rank[i] for i in range(nnodes)) + payload = (ipaddrs + "\n").encode() + # Best-effort broadcast: a single dead/stale peer socket must not abort + # the whole rendezvous, so swallow per-conn send errors. + for conn in conns.values(): + try: + conn.sendall(payload) + except Exception: + pass + finally: + try: + conn.close() + except Exception: + pass + return ipaddrs + finally: + srv.close() + + +def report(rank, host_ip, master_addr, port, token, deadline): + """rank>0 path: report own IP, block on recv until rank-0 broadcasts.""" + payload = (json.dumps({"token": token, "rank": rank, "ip": host_ip}) + "\n").encode() + while time.time() < deadline: + try: + sock = socket.create_connection((master_addr, port), timeout=3.0) + sock.sendall(payload) + # rank-0 broadcasts once, only after every node has reported, which + # can be far longer than a few seconds when peers load the image at + # different times. Block on recv for the remaining deadline so we + # stay connected and never miss that one-shot broadcast. + remaining = max(1.0, deadline - time.time()) + sock.settimeout(remaining) + line = recv_line(sock) + sock.close() + if line: + return line.strip() + except Exception: + time.sleep(1.0) + return None + + +def main(argv): + rank = int(argv[1]) + nnodes = int(argv[2]) + host_ip = argv[3] + master_addr = argv[4] + port = int(argv[5]) + job_id = argv[6] + token = f"JOB{job_id}" + timeout_s = int(os.environ.get("IP_SYNC_TIMEOUT", "1800")) + deadline = time.time() + timeout_s + if rank == 0: + result = serve(nnodes, host_ip, port, token, deadline) + else: + result = report(rank, host_ip, master_addr, port, token, deadline) + if result: + print(result) + return 0 + print("") + return 2 + + +if __name__ == "__main__": + sys.exit(main(sys.argv)) diff --git a/scripts/sglang_disagg/parse_to_csv.py b/scripts/sglang_disagg/parse_to_csv.py index 8d4d65d..4fe659d 100644 --- a/scripts/sglang_disagg/parse_to_csv.py +++ b/scripts/sglang_disagg/parse_to_csv.py @@ -1,170 +1,255 @@ #!/usr/bin/env python3 """ Parse SGLang benchmark log file and save results to CSV. -Extracts: Concurrency, Input tokens, Output tokens, Total Token throughput (tok/s) -For each configuration, takes the MAX Total Token throughput across all iterations. + +Extracts a full metric set per (isl, osl, concurrency) configuration: + - Request throughput (req/s) + - Input / Output / Total token throughput (tok/s) + - Mean E2E latency (ms), Mean TTFT (ms), Mean ITL (ms) + +For each configuration the row from the iteration with the MAX total token +throughput is retained (all metrics come from that single best run, so they +stay mutually consistent). """ import re import csv from pathlib import Path -from typing import List, Dict, Tuple -from collections import defaultdict +from typing import Dict, Tuple + +# Metric key -> madengine metric name. Order defines CSV row order per config. +METRIC_FIELDS = [ + ("request_throughput_req_s", "request_throughput_req_s"), + ("input_token_throughput_tok_s", "input_token_throughput_tok_s"), + ("output_token_throughput_tok_s", "output_token_throughput_tok_s"), + ("total_token_throughput_tok_s", "total_token_throughput_tok_s"), + ("mean_e2e_latency_ms", "mean_e2e_latency_ms"), + ("mean_ttft_ms", "mean_ttft_ms"), + ("mean_itl_ms", "mean_itl_ms"), +] + + +def _extract(pattern: str, text: str): + """Return the first capture group as float (commas stripped), else None.""" + m = re.search(pattern, text) + return float(m.group(1).replace(",", "")) if m else None def parse_benchmark_log(log_file: str) -> Dict[Tuple[int, int, int], Dict]: - """Parse benchmark log file and extract results, keeping max throughput per configuration.""" - results = defaultdict(lambda: {'concurrency': None, 'input_tokens': None, - 'output_tokens': None, 'max_throughput': 0.0}) - - with open(log_file, 'r') as f: + """Parse benchmark log, keeping the best-throughput run per configuration.""" + results: Dict[Tuple[int, int, int], Dict] = {} + + with open(log_file, "r") as f: content = f.read() - - # Find the start of the first iteration (ignore warmup) - first_iter_match = re.search(r'RUNNING: the benchserving script for iter: 1', content) + + # Find the start of the first iteration (ignore warmup). + first_iter_match = re.search(r"RUNNING: the benchserving script for iter: 1", content) if not first_iter_match: print("Warning: No iteration 1 found. Processing entire file.") start_pos = 0 else: start_pos = first_iter_match.start() - - # Process only from first iteration onwards + content = content[start_pos:] - - # Split by benchmark result sections - sections = re.split(r'============ Serving Benchmark Result ============', content) - - current_input_seq_len = None - current_output_seq_len = None - current_concurrency = None - - for i, section in enumerate(sections[1:], 1): # Skip first empty section - # Look for configuration in previous sections (from RUNNING line) - if i > 1: - prev_section = sections[i-1] - - # Extract config: prompts isl osl con - config_match = re.search(r'RUNNING: prompts\s+isl\s+(\d+)\s+osl\s+(\d+)\s+con\s+(\d+)', prev_section) - if config_match: - current_input_seq_len = int(config_match.group(1)) - current_output_seq_len = int(config_match.group(2)) - current_concurrency = int(config_match.group(3)) - - # Extract Total token throughput (tok/s) from benchmark result section - throughput_match = re.search(r'Total token throughput \(tok/s\):\s+([\d.]+)', section) - throughput = float(throughput_match.group(1)) if throughput_match else None - - # Only process if we have a valid configuration from RUNNING line and throughput - if current_input_seq_len and current_output_seq_len and current_concurrency and throughput is not None: - config_key = (current_input_seq_len, current_output_seq_len, current_concurrency) - - # Update results for this configuration - # Always use values from RUNNING line (isl, osl, con) - results[config_key]['concurrency'] = current_concurrency - results[config_key]['input_tokens'] = current_input_seq_len - results[config_key]['output_tokens'] = current_output_seq_len - - # Keep the maximum throughput - if throughput > results[config_key]['max_throughput']: - results[config_key]['max_throughput'] = throughput - + + # Split by benchmark result sections; config lives in the preceding section. + sections = re.split(r"============ Serving Benchmark Result ============", content) + + current_isl = None + current_osl = None + current_con = None + + for i, section in enumerate(sections[1:], 1): # Skip first (pre-first-result) chunk + prev_section = sections[i - 1] + + # Config emitted by benchmark_xPyD.sh: "RUNNING: prompts isl X osl Y con Z" + config_match = re.search( + r"RUNNING: prompts\s+isl\s+(\d+)\s+osl\s+(\d+)\s+con\s+(\d+)", prev_section + ) + if config_match: + current_isl = int(config_match.group(1)) + current_osl = int(config_match.group(2)) + current_con = int(config_match.group(3)) + + if current_isl is None or current_osl is None or current_con is None: + continue + + total_throughput = _extract(r"Total token throughput \(tok/s\):\s+([\d,\.]+)", section) + if total_throughput is None: + continue + + row = { + "isl": current_isl, + "osl": current_osl, + "con": current_con, + "request_throughput_req_s": _extract( + r"Request throughput \(req/s\):\s+([\d,\.]+)", section + ), + "input_token_throughput_tok_s": _extract( + r"Input token throughput \(tok/s\):\s+([\d,\.]+)", section + ), + "output_token_throughput_tok_s": _extract( + r"Output token throughput \(tok/s\):\s+([\d,\.]+)", section + ), + "total_token_throughput_tok_s": total_throughput, + "mean_e2e_latency_ms": _extract(r"Mean E2E Latency \(ms\):\s+([\d,\.]+)", section), + "mean_ttft_ms": _extract(r"Mean TTFT \(ms\):\s+([\d,\.]+)", section), + "mean_itl_ms": _extract(r"Mean ITL \(ms\):\s+([\d,\.]+)", section), + } + + config_key = (current_isl, current_osl, current_con) + best = results.get(config_key) + if best is None or total_throughput > best["total_token_throughput_tok_s"]: + results[config_key] = row + return results +def _sorted_results(results: Dict[Tuple[int, int, int], Dict]): + # Sort by concurrency, then isl, then osl. + return [ + results[key] + for key in sorted(results.keys(), key=lambda k: (k[2], k[0], k[1])) + ] + + def save_to_csv(results: Dict[Tuple[int, int, int], Dict], output_file: str): - """Save results to CSV file with specified columns.""" + """Save a wide per-configuration summary CSV (all metrics).""" if not results: print("No results to save.") return - - # Define column order - fieldnames = ['Concurrency', 'Input tokens', 'Output tokens', 'Total Token throughput (tok/s)'] - - with open(output_file, 'w', newline='') as f: + + fieldnames = [ + "Concurrency", + "Input tokens", + "Output tokens", + "Request throughput (req/s)", + "Input token throughput (tok/s)", + "Output token throughput (tok/s)", + "Total Token throughput (tok/s)", + "Mean E2E Latency (ms)", + "Mean TTFT (ms)", + "Mean ITL (ms)", + ] + + def _fmt(value): + return f"{value:.2f}" if value is not None else "" + + with open(output_file, "w", newline="") as f: writer = csv.DictWriter(f, fieldnames=fieldnames) writer.writeheader() - - # Sort by concurrency, then input tokens, then output tokens - for (input_tokens, output_tokens, concurrency), data in sorted(results.items(), - key=lambda x: (x[0][2], x[0][0], x[0][1])): - row = { - 'Concurrency': data['concurrency'], - 'Input tokens': data['input_tokens'], - 'Output tokens': data['output_tokens'], - 'Total Token throughput (tok/s)': f"{data['max_throughput']:.2f}" - } - writer.writerow(row) - + for data in _sorted_results(results): + writer.writerow( + { + "Concurrency": data["con"], + "Input tokens": data["isl"], + "Output tokens": data["osl"], + "Request throughput (req/s)": _fmt(data["request_throughput_req_s"]), + "Input token throughput (tok/s)": _fmt(data["input_token_throughput_tok_s"]), + "Output token throughput (tok/s)": _fmt(data["output_token_throughput_tok_s"]), + "Total Token throughput (tok/s)": _fmt(data["total_token_throughput_tok_s"]), + "Mean E2E Latency (ms)": _fmt(data["mean_e2e_latency_ms"]), + "Mean TTFT (ms)": _fmt(data["mean_ttft_ms"]), + "Mean ITL (ms)": _fmt(data["mean_itl_ms"]), + } + ) + print(f"Saved {len(results)} benchmark configurations to {output_file}") def _get_run_metadata(pipeline: str = "sglang"): """Collect run metadata from environment variables.""" import os - xP = os.environ.get('xP', '1') - yD = os.environ.get('yD', '1') - dp_mode = os.environ.get('DP_MODE', '0') - run_mori = os.environ.get('RUN_MORI', '0') - gpus = os.environ.get('GPUS_PER_NODE', '8') - - # Determine backend tag - if dp_mode == '1': - backend = 'mori_dp' - elif run_mori == '1': - backend = 'mori_io' + + xP = os.environ.get("xP", "1") + yD = os.environ.get("yD", "1") + dp_mode = os.environ.get("DP_MODE", "0") + run_mori = os.environ.get("RUN_MORI", "1") + kv_backend = os.environ.get("KV_TRANSFER_BACKEND", "mooncake").lower() + gpus = os.environ.get("GPUS_PER_NODE", "8") + + # Determine backend tag. RUN_MORI is the launcher's own switch and takes + # priority; DP_MODE only distinguishes mori_dp vs mori_io *within* the MoRI + # path. When MoRI is off, tag by the actual KV transfer backend so a + # Mooncake/NIXL run with DP_MODE=1 is not mislabelled as mori_dp. + if run_mori == "1": + backend = "mori_dp" if dp_mode == "1" else "mori_io" + elif kv_backend in ("mooncake", "nixl"): + backend = kv_backend else: - backend = 'mooncake' + backend = "mooncake" return { - 'pipeline': pipeline, - 'deployment_type': f'disagg_{xP}P{yD}D', - 'tags': f'{pipeline}_disagg,{backend}', - 'n_gpus': str(int(xP) * int(gpus) + int(yD) * int(gpus)), - 'nnodes': str(int(xP) + int(yD)), - 'gpus_per_node': gpus, - 'docker_image': os.environ.get('DOCKER_IMAGE_NAME', ''), - 'machine_name': os.environ.get('SLURM_JOB_NODELIST', ''), - 'launcher': 'slurm_multi', - 'gpu_architecture': 'gfx942', + "pipeline": pipeline, + "deployment_type": f"disagg_{xP}P{yD}D", + "tags": f"{pipeline}_disagg,{backend}", + "n_gpus": str(int(xP) * int(gpus) + int(yD) * int(gpus)), + "nnodes": str(int(xP) + int(yD)), + "gpus_per_node": gpus, + "docker_image": os.environ.get("DOCKER_IMAGE_NAME", ""), + "machine_name": os.environ.get("SLURM_JOB_NODELIST", ""), + # This script only runs behind madengine's native sglang-disagg launcher + # (scripts/sglang_disagg/run.sh), so the launcher name is fixed. The GPU + # architecture varies by cluster; madengine exports it, with a gfx942 + # fallback for standalone/manual invocations. + "launcher": "sglang-disagg", + "gpu_architecture": os.environ.get("MAD_SYSTEM_GPU_ARCHITECTURE", "gfx942"), } -def save_perf_csv(results: Dict[Tuple[int, int, int], Dict], output_file: str, - model_name: str = "", pipeline: str = "sglang"): - """Save results in madengine perf.csv format.""" +def save_perf_csv( + results: Dict[Tuple[int, int, int], Dict], + output_file: str, + model_name: str = "", + pipeline: str = "sglang", +): + """Save results in madengine perf.csv format (one row per config x metric).""" if not results: print("No results to save to perf.csv.") return + import os + meta = _get_run_metadata(pipeline) + xP = os.environ.get("xP", "1") + yD = os.environ.get("yD", "1") fieldnames = [ - 'model', 'n_gpus', 'nnodes', 'gpus_per_node', 'training_precision', - 'pipeline', 'args', 'tags', 'docker_file', 'base_docker', 'docker_sha', - 'docker_image', 'git_commit', 'machine_name', 'deployment_type', 'launcher', - 'gpu_architecture', 'performance', 'metric', 'relative_change', 'status', - 'build_duration', 'test_duration', 'dataname', 'data_provider_type', - 'data_size', 'data_download_duration', 'build_number', - 'additional_docker_run_options', + "model", "n_gpus", "nnodes", "gpus_per_node", "training_precision", + "pipeline", "args", "tags", "docker_file", "base_docker", "docker_sha", + "docker_image", "git_commit", "machine_name", "deployment_type", "launcher", + "gpu_architecture", "performance", "metric", "relative_change", "status", + "build_duration", "test_duration", "dataname", "data_provider_type", + "data_size", "data_download_duration", "build_number", + "additional_docker_run_options", ] - with open(output_file, 'w', newline='') as f: + rows_written = 0 + with open(output_file, "w", newline="") as f: writer = csv.DictWriter(f, fieldnames=fieldnames) writer.writeheader() - for (input_tokens, output_tokens, concurrency), data in sorted( - results.items(), key=lambda x: (x[0][2], x[0][0], x[0][1]) - ): - row = { - 'model': model_name, - 'performance': f"{data['max_throughput']:.2f}", - 'metric': f"tok/s (isl={data['input_tokens']} osl={data['output_tokens']} con={data['concurrency']})", - 'status': 'SUCCESS', - } - row.update(meta) - writer.writerow(row) + for data in _sorted_results(results): + config_key = ( + f"{xP}p{yD}d_isl{data['isl']}_osl{data['osl']}_con{data['con']}" + ) + for source_key, metric_name in METRIC_FIELDS: + value = data.get(source_key) + if value is None: + continue + row = { + "model": config_key, + "performance": f"{value:.2f}", + "metric": metric_name, + "status": "SUCCESS", + } + row.update(meta) + # model_name (catalog key) is prepended by madengine; keep the + # per-config key in `model` so metrics stay disambiguated. + writer.writerow(row) + rows_written += 1 - print(f"Saved {len(results)} rows to perf.csv: {output_file}") + print(f"Saved {rows_written} rows to perf.csv: {output_file}") def main(): @@ -172,40 +257,42 @@ def main(): import sys import argparse - parser = argparse.ArgumentParser(description='Parse SGLang benchmark log file and save results to CSV') - parser.add_argument('log_file', type=str, help='Path to benchmark log file') - parser.add_argument('-o', '--output', type=str, help='Output CSV file name (default: _results.csv)') - parser.add_argument('--perf-csv', type=str, help='Also generate madengine perf.csv at this path') - parser.add_argument('--model-name', type=str, default='', help='Model name for perf.csv') + parser = argparse.ArgumentParser( + description="Parse SGLang benchmark log file and save results to CSV" + ) + parser.add_argument("log_file", type=str, help="Path to benchmark log file") + parser.add_argument( + "-o", "--output", type=str, + help="Output CSV file name (default: _results.csv)", + ) + parser.add_argument( + "--perf-csv", type=str, help="Also generate madengine perf.csv at this path" + ) + parser.add_argument("--model-name", type=str, default="", help="Model name for perf.csv") args = parser.parse_args() log_file = args.log_file - # Check if file exists if not Path(log_file).exists(): print(f"Error: Log file not found: {log_file}") sys.exit(1) print(f"Parsing log file: {log_file}") - # Parse the log file results = parse_benchmark_log(log_file) if not results: print("No benchmark results found in log file.") - return + sys.exit(1) - # Generate output filename if args.output: output_file = args.output else: - output_file = Path(log_file).stem + '_results.csv' + output_file = Path(log_file).stem + "_results.csv" - # Save to CSV save_to_csv(results, output_file) - # Save madengine perf.csv if requested if args.perf_csv: save_perf_csv(results, args.perf_csv, args.model_name) @@ -214,5 +301,5 @@ def main(): print(f" Output file: {output_file}") -if __name__ == '__main__': +if __name__ == "__main__": main() diff --git a/scripts/sglang_disagg/run.sh b/scripts/sglang_disagg/run.sh new file mode 100755 index 0000000..0b2b3d1 --- /dev/null +++ b/scripts/sglang_disagg/run.sh @@ -0,0 +1,221 @@ +#!/bin/bash +# ============================================================================= +# madengine-native entrypoint for SGLang Disaggregated P/D inference. +# +# This is the `scripts` target referenced by models.json (sglang_disagg_*). +# madengine's built-in `sglang-disagg` SLURM/K8s launcher +# (src/madengine/deployment/slurm.py::_generate_sglang_disagg_command) runs ONE +# container per node and exports the cluster topology as env vars. This script +# bridges those into the env contract of the proven MAD-private disagg +# launchers and execs the right one (MoRI EP, or Mooncake/NIXL). +# +# madengine -> proven env mapping: +# SGLANG_NODE_RANK -> NODE_RANK (0=proxy, 1..xP=prefill, rest=decode) +# SGLANG_DISAGG_PREFILL_NODES -> xP +# SGLANG_DISAGG_DECODE_NODES -> yD +# SGLANG_NODE_IPS -> IPADDRS (comma-separated, rank order) +# SGLANG_TP_SIZE -> IO_EP_TP_SIZE / per-server --tp-size +# MASTER_PORT -> MASTER_PORT +# +# Model identity + transport come from deployment env_vars / models.json args: +# MODEL_NAME short catalog key in models.yaml (e.g. DeepSeek-R1) [required] +# MODEL_PATH weights dir mounted into the container [required] +# RUN_MORI 1 = MoRI EP all-to-all (default), 0 = Mooncake/NIXL +# KV_TRANSFER_BACKEND prefill->decode KV transfer backend (mori/mooncake/nixl) +# DP_MODE 1 = DP-attention EP (DeepSeek-V3/R1 inter-node), 0 = TP +# ============================================================================= +set -uo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +# --- parse optional CLI args (models.json "args"), env takes precedence ------ +while [[ $# -gt 0 ]]; do + case "$1" in + --model-name|--model_name|--model_repo) MODEL_NAME="${MODEL_NAME:-$2}"; shift 2;; + --model-path|--model_path) MODEL_PATH="${MODEL_PATH:-$2}"; shift 2;; + --run-mori) RUN_MORI="${RUN_MORI:-$2}"; shift 2;; + --dp-mode) DP_MODE="${DP_MODE:-$2}"; shift 2;; + --xp|--xP|--prefill-nodes) xP="${xP:-$2}"; shift 2;; + --yd|--yD|--decode-nodes) yD="${yD:-$2}"; shift 2;; + --ipaddrs|--node-ips) IPADDRS="${IPADDRS:-$2}"; shift 2;; + --master-port) MASTER_PORT="${MASTER_PORT:-$2}"; shift 2;; + --gpus-per-node) GPUS_PER_NODE="${GPUS_PER_NODE:-$2}"; shift 2;; + --kv-transfer-backend|--transfer-backend) KV_TRANSFER_BACKEND="${KV_TRANSFER_BACKEND:-$2}"; shift 2;; + *) shift;; + esac +done + +# --- bridge madengine sglang-disagg topology -> proven launcher env --------- +export NODE_RANK="${NODE_RANK:-${SGLANG_NODE_RANK:-${SLURM_PROCID:-0}}}" +# Persist run.sh output to the shared /run_logs mount so failures are +# diagnosable from the submission node (madengine discards model stdout). +# The dir is created from inside the container (root); make it world-writable +# so an ordinary user outside the container can clean it up afterwards. +# +# /run_logs MUST be a shared (NFS-backed) mount passed through the docker env: +# the disaggregated launchers write per-node readiness logs +# (prefill_NODE*/decode_NODE*/proxy_NODE*) to /run_logs/${SLURM_JOB_ID} and the +# rank-0 proxy greps them across nodes. Disaggregated P/D always spans >=2 nodes +# (rank-0 prefill/proxy + at least one decode node), so a node-local fallback +# would silently break that cross-node rendezvous. Require /run_logs and fail +# fast if it is missing or not writable. +RUN_LOG_DIR="/run_logs/${SLURM_JOB_ID:-local}" +if ! mkdir -p "${RUN_LOG_DIR}" 2>/dev/null || [[ ! -w "${RUN_LOG_DIR}" ]]; then + echo "FATAL: /run_logs is not writable. It must be a shared (NFS-backed)" >&2 + echo " mount passed through the docker env: the disaggregated launchers" >&2 + echo " write per-node readiness logs to /run_logs/\${SLURM_JOB_ID} and the" >&2 + echo " rank-0 proxy greps them across nodes. Mount it and re-run." >&2 + exit 1 +fi +chmod 1777 /run_logs 2>/dev/null || true +chmod -R 0777 "${RUN_LOG_DIR}" 2>/dev/null || true +export RUN_LOG_DIR +exec > >(tee -a "${RUN_LOG_DIR}/run_sh_rank${NODE_RANK}.log") 2>&1 +export xP="${xP:-${SGLANG_DISAGG_PREFILL_NODES:-1}}" +export yD="${yD:-${SGLANG_DISAGG_DECODE_NODES:-1}}" +export NNODES="${NNODES:-${SGLANG_DISAGG_TOTAL_NODES:-${SLURM_JOB_NUM_NODES:-$((xP + yD))}}}" +export MASTER_PORT="${MASTER_PORT:-23731}" +export IO_EP_TP_SIZE="${IO_EP_TP_SIZE:-${SGLANG_TP_SIZE:-8}}" +# madengine forwards MASTER_ADDR (rank-0 host) into the container; use it as the +# rendezvous address for in-container IP discovery below. +export MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}" + +# IPADDRS = comma-separated node IPs in rank order. It may be supplied +# explicitly via env/arg or SGLANG_NODE_IPS. madengine's container runner does +# NOT forward SGLANG_NODE_IPS into the container, so when it is absent we +# self-discover the node IPs in-container using only the forwarded +# MASTER_ADDR/NODE_RANK/NNODES, mirroring scripts/vllm_dissag/run.sh. +export IPADDRS="${IPADDRS:-${SGLANG_NODE_IPS:-}}" + +# Prefer the routable transport interface (NCCL_SOCKET_IFNAME, eth0 here); +# hostname -I order is not deterministic and may surface link-local addrs. +# NCCL_SOCKET_IFNAME may be a comma-separated list (see set_env_vars.sh), so +# take only the first interface before handing it to `ip addr show`. +_HOST_IFACE="${NCCL_SOCKET_IFNAME:-eth0}"; _HOST_IFACE="${_HOST_IFACE%%,*}" +HOST_IP="$(ip -4 -o addr show "${_HOST_IFACE}" 2>/dev/null | awk '{print $4}' | cut -d/ -f1 | head -n1)" +[[ -z "${HOST_IP}" ]] && HOST_IP="$(hostname -I | awk '{print $1}')" +# Derive the rendezvous port from the job id. Use a wide modulo so two +# concurrent jobs whose ids happen to share a rank-0 host are very unlikely to +# collide (a narrow 10000-wide range collides whenever ids differ by a multiple +# of 10000). The value is deterministic across nodes (depends only on the job +# id, which is identical on every node). +IP_SYNC_PORT="${IP_SYNC_PORT:-$((20000 + (${SLURM_JOB_ID:-0} % 40000)))}" + +# rank-0 acts as a tiny TCP rendezvous server: peers connect to MASTER_ADDR and +# report their host IP; rank-0 aggregates the rank-ordered list and broadcasts. +# The implementation lives in the sibling ip_rendezvous.py so it can be unit tested. +_tcp_discover_ipaddrs() { + python3 "$SCRIPT_DIR/ip_rendezvous.py" \ + "${NODE_RANK}" "${NNODES}" "${HOST_IP}" "${MASTER_ADDR}" "${IP_SYNC_PORT}" "${SLURM_JOB_ID:-0}" +} + +# Method 1: SLURM nodelist (only if scontrol/getent exist inside the container). +if [[ -z "${IPADDRS}" && -n "${SLURM_JOB_NODELIST:-}" ]] && command -v scontrol >/dev/null 2>&1 && command -v getent >/dev/null 2>&1; then + _ips=""; _cnt=0 + while IFS= read -r _n; do + [[ -z "${_n}" ]] && continue + _ip="$(getent ahostsv4 "${_n}" | awk 'NR==1{print $1}')" + [[ -n "${_ip}" ]] && { _ips="${_ips:+${_ips},}${_ip}"; _cnt=$((_cnt + 1)); } + done < <(scontrol show hostnames "${SLURM_JOB_NODELIST}") + [[ "${_cnt}" == "${NNODES}" && -n "${_ips}" ]] && IPADDRS="${_ips}" +fi + +# Method 2: TCP rendezvous via forwarded MASTER_ADDR (primary path). +if [[ -z "${IPADDRS}" ]]; then + _tcp="$(_tcp_discover_ipaddrs || true)" + [[ -n "${_tcp}" ]] && IPADDRS="${_tcp}" +fi + +if [[ -z "${IPADDRS}" ]]; then + echo "ERROR: unable to determine IPADDRS (node IP list);" \ + "MASTER_ADDR=${MASTER_ADDR} NNODES=${NNODES} NODE_RANK=${NODE_RANK}." >&2 + exit 1 +fi +export IPADDRS +# rank-0 IP is the master/load-balancer host. +export MASTER_ADDR="$(echo "$IPADDRS" | cut -d',' -f1)" + +# --- model + transport defaults -------------------------------------------- +export MODEL_NAME="${MODEL_NAME:-}" +export MODEL_PATH="${MODEL_PATH:-}" +export RUN_MORI="${RUN_MORI:-1}" +export DP_MODE="${DP_MODE:-0}" +# Default the KV transfer backend from RUN_MORI so the non-MoRI path does not +# silently inherit an invalid `mori` backend (mooncake/NIXL use mooncake). +if [[ "$RUN_MORI" == "1" ]]; then + export KV_TRANSFER_BACKEND="${KV_TRANSFER_BACKEND:-mori}" +else + export KV_TRANSFER_BACKEND="${KV_TRANSFER_BACKEND:-mooncake}" +fi + +# Helper scripts (socket_barrier.py, set_env_vars.sh, benchmark_xPyD.sh, ...) +# live alongside this script; point the Mooncake-path lookups at them. +export MOONCAKE_COOKBOOK_PATH="${MOONCAKE_COOKBOOK_PATH:-$SCRIPT_DIR}" +mkdir -p "${RUN_LOG_DIR}" 2>/dev/null || true + +if [[ -z "$MODEL_NAME" || -z "$MODEL_PATH" ]]; then + echo "ERROR: MODEL_NAME and MODEL_PATH must be set (via deployment env_vars or args)." >&2 + echo " MODEL_NAME='$MODEL_NAME' MODEL_PATH='$MODEL_PATH'" >&2 + exit 1 +fi + +# --- stage model weights if missing (rank-0 gated, NFS-shared MODEL_PATH) ----- +# run.sh historically assumed pre-staged weights. Stage from HF when MODEL_PATH +# is incomplete. Only NODE_RANK 0 downloads; peers wait on a sentinel. +# Treat weights as complete only if our own download sentinel is present, OR a +# config.json plus at least one weight shard exists (covers manually pre-staged +# dirs). A partial earlier download (config.json present, shards missing, no +# sentinel) must NOT be mistaken for a complete stage. +_weights_complete() { + [[ -f "${MODEL_PATH}/.stage_done" ]] && return 0 + [[ -f "${MODEL_PATH}/config.json" ]] || return 1 + compgen -G "${MODEL_PATH}/*.safetensors" >/dev/null 2>&1 && return 0 + compgen -G "${MODEL_PATH}/*.bin" >/dev/null 2>&1 && return 0 + return 1 +} +if ! _weights_complete; then + _MODEL_REPO="${MODEL_REPO:-}" + if [[ -z "${_MODEL_REPO}" ]]; then + case "${MODEL_NAME}" in + DeepSeek-R1) _MODEL_REPO="deepseek-ai/DeepSeek-R1-0528" ;; + Qwen3-Next-80B) _MODEL_REPO="Qwen/Qwen3-Next-80B-A3B-Instruct" ;; + *) echo "ERROR: weights missing at ${MODEL_PATH}; set MODEL_REPO for ${MODEL_NAME}" >&2; exit 1 ;; + esac + fi + _SENTINEL="${MODEL_PATH}/.stage_done" + if [[ "${NODE_RANK:-0}" == "0" ]]; then + echo "[stage_weights] NODE_RANK=0 downloading ${_MODEL_REPO} -> ${MODEL_PATH}" + mkdir -p "${MODEL_PATH}" + # huggingface-cli is deprecated and refuses to run on newer huggingface_hub + # (it prints "use `hf` instead" and exits non-zero); prefer the `hf` CLI and + # fall back to huggingface-cli only on older images that lack `hf`. + if command -v hf >/dev/null 2>&1; then + _HF_DL=(hf download) + else + _HF_DL=(huggingface-cli download) + fi + HF_TOKEN="${MAD_SECRETS_HFTOKEN:-${HF_TOKEN:-}}" "${_HF_DL[@]}" "${_MODEL_REPO}" --local-dir "${MODEL_PATH}" || { echo "[stage_weights] download failed" >&2; exit 1; } + touch "${_SENTINEL}" + else + echo "[stage_weights] NODE_RANK=${NODE_RANK} waiting for weights at ${MODEL_PATH}" + for _i in $(seq 1 720); do + _weights_complete && break + sleep 10 + done + _weights_complete || { echo "[stage_weights] timed out waiting for weights" >&2; exit 1; } + fi +fi + + +echo "==============================================================" +echo " SGLang Disaggregated (madengine bridge)" +echo " MODEL_NAME=$MODEL_NAME MODEL_PATH=$MODEL_PATH" +echo " NODE_RANK=$NODE_RANK xP=$xP yD=$yD TP=$IO_EP_TP_SIZE" +echo " RUN_MORI=$RUN_MORI DP_MODE=$DP_MODE KV_TRANSFER_BACKEND=$KV_TRANSFER_BACKEND" +echo " MASTER_ADDR=$MASTER_ADDR IPADDRS=$IPADDRS" +echo "==============================================================" + +# sglang_disagg_mori_io_ep.sh is a functional superset of the retired +# sglang_disagg_server.sh (per-model config from models.yaml, mori/mooncake/nixl +# via KV_TRANSFER_BACKEND, DP_MODE support) — route both RUN_MORI values there. +exec bash "$SCRIPT_DIR/sglang_disagg_mori_io_ep.sh" diff --git a/scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh b/scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh index f224493..acf8220 100755 --- a/scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh +++ b/scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh @@ -76,53 +76,8 @@ BARRIER_PORT="${BARRIER_PORT:-4342}" # Dependencies and Environment Setup # ============================================================================= -pip install py-spy -pip install --ignore-installed --force-reinstall flask -pip install pyyaml - - -host_ip=$(ip route get 1.1.1.1 | awk '/src/ {print $7}') -host_name=$(hostname) - -if [[ "$PARALLEL_MODE" != "dp" && "$PARALLEL_MODE" != "tp" ]]; then - echo "ERROR: PARALLEL_MODE must be 'dp' or 'tp' (got: ${PARALLEL_MODE})" - exit 1 -fi - -# ============================================================================= -# Parallelism Settings -# ============================================================================= - -# DP_MODE=0: CLI --tp-size is IO_EP_TP_SIZE (default 8) on every worker; PREFILL_EP_SIZE/DECODE_EP_SIZE -# still scale with xP/yD×GPUS_PER_NODE for MoRI env (not passed as CLI unless DP_MODE=1). -# DP_MODE=1: --tp-size scales with cluster; --dp-size/--ep-size on CLI (same total degree as Nnodes×GPUS_PER_NODE). +# === Model-Specific Configuration from YAML === GPUS_PER_NODE="${GPUS_PER_NODE:-8}" -GENERIC_TP_SIZE="${GENERIC_TP_SIZE:-8}" - -if [[ "$DP_MODE" == "1" ]]; then - PREFILL_TP_SIZE=$((xP * GPUS_PER_NODE)) - DECODE_TP_SIZE=$((yD * GPUS_PER_NODE)) -else - PREFILL_TP_SIZE="${GENERIC_TP_SIZE}" - DECODE_TP_SIZE="${GENERIC_TP_SIZE}" -fi - - -if [[ "$DP_MODE" == "1" ]]; then - PREFILL_EP_SIZE=$((xP * GPUS_PER_NODE)) - DECODE_EP_SIZE=$((yD * GPUS_PER_NODE)) - PREFILL_DP_SIZE=$((xP * GPUS_PER_NODE)) - DECODE_DP_SIZE=$((yD * GPUS_PER_NODE)) - export PREFILL_DP_SIZE DECODE_DP_SIZE PREFILL_EP_SIZE DECODE_EP_SIZE -else - unset PREFILL_DP_SIZE DECODE_DP_SIZE PREFILL_EP_SIZE DECODE_EP_SIZE 2>/dev/null || true -fi -export PREFILL_TP_SIZE DECODE_TP_SIZE - -# ============================================================================= -# Model-Specific Configuration from YAML -# ============================================================================= - MODELS_YAML="${MODELS_YAML:-${SCRIPT_DIR}/models.yaml}" if [[ ! -f "$MODELS_YAML" ]]; then @@ -130,6 +85,20 @@ if [[ ! -f "$MODELS_YAML" ]]; then exit 1 fi +if ! python3 -c "import yaml" >/dev/null 2>&1; then + echo "[stage_deps] PyYAML not found; installing at runtime (expected on the" >&2 + echo " base, non-overlay image) ..." >&2 + pip install --no-cache-dir --quiet pyyaml || true +fi +if ! python3 -c "import yaml" >/dev/null 2>&1; then + echo "ERROR: PyYAML is not installed and could not be installed at runtime," >&2 + echo " but is required to parse ${MODELS_YAML}. Use an image built" >&2 + echo " from docker/sglang_disagg_inference_full_overlay*.Dockerfile" >&2 + echo " (which installs pyyaml at build time), or ensure network" >&2 + echo " access for 'pip install pyyaml'." >&2 + exit 1 +fi + export MODELS_YAML MODEL_NAME PARALLEL_MODE xP GPUS_PER_NODE eval "$(python3 - <<'PY' import os @@ -144,11 +113,20 @@ mode = os.environ["PARALLEL_MODE"] xP = int(os.environ.get("xP", "1")) gpus_per_node = int(os.environ.get("GPUS_PER_NODE", "8")) + +def q(v): + return shlex.quote(str(v if v is not None else "")) + + with open(config_path, "r", encoding="utf-8") as f: models = yaml.safe_load(f) or {} if model_name not in models: - print(f'echo "ERROR: Model {model_name} not found in {config_path}"; exit 1') + # Quote the interpolated values: this string is eval'd by the shell, so an + # unescaped model_name / config_path could break the launcher or inject + # shell tokens. + msg = f"ERROR: Model {model_name} not found in {config_path}" + print(f"echo {q(msg)}; exit 1") sys.exit(0) cfg = models[model_name] or {} @@ -156,10 +134,6 @@ prefill = cfg.get("prefill", {}) or {} decode = cfg.get("decode", {}) or {} -def q(v): - return shlex.quote(str(v if v is not None else "")) - - prefill_flags = prefill.get(mode, "") or "" exports = { @@ -175,6 +149,49 @@ for key, value in exports.items(): PY )" +host_ip=$(ip route get 1.1.1.1 | awk '/src/ {print $7}') +host_name=$(hostname) + +if [[ "$PARALLEL_MODE" != "dp" && "$PARALLEL_MODE" != "tp" ]]; then + echo "ERROR: PARALLEL_MODE must be 'dp' or 'tp' (got: ${PARALLEL_MODE})" + exit 1 +fi + +# ============================================================================= +# Parallelism Settings +# ============================================================================= + +# DP_MODE=0: CLI --tp-size is IO_EP_TP_SIZE (default 8) on every worker; PREFILL_EP_SIZE/DECODE_EP_SIZE +# still scale with xP/yD×GPUS_PER_NODE for MoRI env (not passed as CLI unless DP_MODE=1). +# DP_MODE=1: --tp-size scales with cluster; --dp-size/--ep-size on CLI (same total degree as Nnodes×GPUS_PER_NODE). +GPUS_PER_NODE="${GPUS_PER_NODE:-8}" +GENERIC_TP_SIZE="${GENERIC_TP_SIZE:-8}" + +if [[ "$DP_MODE" == "1" ]]; then + PREFILL_TP_SIZE=$((xP * GPUS_PER_NODE)) + DECODE_TP_SIZE=$((yD * GPUS_PER_NODE)) +else + PREFILL_TP_SIZE="${GENERIC_TP_SIZE}" + DECODE_TP_SIZE="${GENERIC_TP_SIZE}" +fi + + +if [[ "$DP_MODE" == "1" ]]; then + PREFILL_EP_SIZE=$((xP * GPUS_PER_NODE)) + DECODE_EP_SIZE=$((yD * GPUS_PER_NODE)) + PREFILL_DP_SIZE=$((xP * GPUS_PER_NODE)) + DECODE_DP_SIZE=$((yD * GPUS_PER_NODE)) + export PREFILL_DP_SIZE DECODE_DP_SIZE PREFILL_EP_SIZE DECODE_EP_SIZE +else + unset PREFILL_DP_SIZE DECODE_DP_SIZE PREFILL_EP_SIZE DECODE_EP_SIZE 2>/dev/null || true +fi +export PREFILL_TP_SIZE DECODE_TP_SIZE + +# ============================================================================= +# Model-Specific Configuration from YAML +# ============================================================================= + + PREFILL_MODEL_CONFIG="${MODEL_BASE_FLAGS} ${MODEL_MODE_FLAGS} ${MODEL_PREFILL_FLAGS} ${MODEL_EXPERIMENTAL_FLAGS}" DECODE_MODEL_CONFIG="${MODEL_BASE_FLAGS} ${MODEL_MODE_FLAGS} ${MODEL_DECODE_FLAGS} ${MODEL_EXPERIMENTAL_FLAGS}" echo "Using model-specific configuration for: $MODEL_NAME (mode=${PARALLEL_MODE})" @@ -184,9 +201,20 @@ export PREFILL_MODEL_CONFIG DECODE_MODEL_CONFIG MODEL_EXPERIMENTAL_FLAGS # shellcheck disable=SC1091 source "${SCRIPT_DIR}/mori_ep_env.sh" -# KV transfer backend: default mori, switchable to mooncake (Mooncake). +# KV transfer backend: default mori, switchable to mooncake or nixl. # Kept out of models.yaml so model config is backend-agnostic. _TRANSFER_BACKEND="${KV_TRANSFER_BACKEND:-mori}" + +# _TRANSFER_BACKEND is interpolated into an eval'd launch command below; restrict +# it to known-good values to avoid invalid backends and shell-token injection. +case "$_TRANSFER_BACKEND" in + mori|mooncake|nixl) ;; + *) + echo "ERROR: unsupported KV_TRANSFER_BACKEND='$_TRANSFER_BACKEND' (expected: mori|mooncake|nixl)" >&2 + exit 1 + ;; +esac + PREFILL_MODEL_CONFIG+=" --disaggregation-transfer-backend ${_TRANSFER_BACKEND}" DECODE_MODEL_CONFIG+=" --disaggregation-transfer-backend ${_TRANSFER_BACKEND}" export PREFILL_MODEL_CONFIG DECODE_MODEL_CONFIG @@ -534,15 +562,17 @@ PY echo "" fi + benchmark_status=0 if [[ "${SKIP_BENCHMARK:-0}" != "1" ]] && [[ -n "${MOONCAKE_COOKBOOK_PATH:-}" ]]; then if [[ -f "${MOONCAKE_COOKBOOK_PATH}/benchmark_xPyD.sh" ]]; then echo "Running ${MOONCAKE_COOKBOOK_PATH}/benchmark_xPyD.sh" ( cd "${MOONCAKE_COOKBOOK_PATH}" || exit 1 bash benchmark_xPyD.sh - ) + ) || benchmark_status=$? else echo "WARN: benchmark_xPyD.sh not found under MOONCAKE_COOKBOOK_PATH=${MOONCAKE_COOKBOOK_PATH}" >&2 + benchmark_status=1 fi fi @@ -552,6 +582,11 @@ PY echo "Killing the co-located prefill server (pid=${_node0_prefill_pid})" kill "${_node0_prefill_pid}" + if [[ "${benchmark_status}" -ne 0 ]]; then + echo "ERROR: benchmark failed with status ${benchmark_status}" >&2 + exit "${benchmark_status}" + fi + elif [[ "$NODE_RANK" -ge 1 && "$NODE_RANK" -lt "$xP" ]]; then echo "${host_name}:${host_ip} is Prefill Node (Model: ${MODEL_NAME:-default})" # NODE_RANK 0..xP-1 map directly to PREFILL_NODE_RANK 0..xP-1 (proxy co-located on NODE_RANK=0). diff --git a/scripts/sglang_disagg/sglang_disagg_server.sh b/scripts/sglang_disagg/sglang_disagg_server.sh deleted file mode 100755 index 261dcdd..0000000 --- a/scripts/sglang_disagg/sglang_disagg_server.sh +++ /dev/null @@ -1,275 +0,0 @@ -#!/bin/bash -# SGLang Disaggregated Server Launcher with Model-Specific Configurations -# ============================================================================= - -# ============================================================================= -# Environment Configuration -# ============================================================================= - -MASTER_ADDR="${MASTER_ADDR:-localhost}" -MASTER_PORT="${MASTER_PORT:-23731}" -NODE_RANK="${NODE_RANK:-0}" -MODEL_PATH=$MODEL_PATH -MODEL_NAME="${MODEL_NAME:-}" -xP="${xP:-1}" -yD="${yD:-1}" -IPADDRS="${IPADDRS:-localhost}" - -# ============================================================================= -# Dependencies and Environment Setup -# ============================================================================= - -pip install py-spy -pip install --ignore-installed --force-reinstall flask - -source $MOONCAKE_COOKBOOK_PATH/set_env_vars.sh -#trap 'echo "Error occurred. Cleaning up..."; exit 0' ERR - -host_ip=$(hostname -I | awk '{print $1}') -host_name=$(hostname) - -# ============================================================================= -# Model-Specific Configuration Maps -# ============================================================================= - -declare -A MODEL_PREFILL_CONFIGS=( - ["Qwen3-32B"]="--tp-size 8" - ["Mixtral-8x7B-v0.1"]="--tp-size 8" - ["Llama-3.1-8B-Instruct"]="--tp-size 8" - ["Llama-3.1-405B-Instruct-FP8-KV"]="--tp-size 8" - ["amd-Llama-3.3-70B-Instruct-FP8-KV"]="--tp-size 8" - ["DeepSeek-V3"]="--tp-size 8" -) - -declare -A MODEL_DECODE_CONFIGS=( - ["Qwen3-32B"]="--tp-size 8" - ["Mixtral-8x7B-v0.1"]="--tp-size 8" - ["Llama-3.1-8B-Instruct"]="--tp-size 8" - ["Llama-3.1-405B-Instruct-FP8-KV"]="--tp-size 8" - ["amd-Llama-3.3-70B-Instruct-FP8-KV"]="--tp-size 8" - ["DeepSeek-V3"]="--tp-size 8" -) - -# ============================================================================= -# Configuration Selection Functions -# ============================================================================= - -get_model_config() { - local mode="$1" - local model_name="$2" - - if [[ "$mode" == "prefill" ]]; then - if [[ -n "${MODEL_PREFILL_CONFIGS[$model_name]}" ]]; then - echo "${MODEL_PREFILL_CONFIGS[$model_name]}" - else - echo "--tp-size 4" - fi - elif [[ "$mode" == "decode" ]]; then - if [[ -n "${MODEL_DECODE_CONFIGS[$model_name]}" ]]; then - echo "${MODEL_DECODE_CONFIGS[$model_name]}" - else - echo "--tp-size 4" - fi - fi -} - -if [[ -z "$MODEL_NAME" ]]; then - echo "Warning: MODEL_NAME not set, using default configurations" - PREFILL_MODEL_CONFIG="--tp-size 4" - DECODE_MODEL_CONFIG="--tp-size 4" -else - PREFILL_MODEL_CONFIG=$(get_model_config "prefill" "$MODEL_NAME") - DECODE_MODEL_CONFIG=$(get_model_config "decode" "$MODEL_NAME") - echo "Using model-specific configuration for: $MODEL_NAME" -fi - -# ============================================================================= -# Container Synchronization -# ============================================================================= - -echo "Waiting at the container creation barrier on $host_name" -python $MOONCAKE_COOKBOOK_PATH/socket_barrier.py \ - --local-ip ${host_ip} \ - --local-port 4342 \ - --enable-port \ - --node-ips ${IPADDRS} \ - --node-ports 4342 - -# ============================================================================= -# Cluster Topology Configuration -# ============================================================================= - -IFS=',' read -ra IP_ARRAY <<< "$IPADDRS" - -PREFILL_ARGS="" -DECODE_ARGS="" - -for ((i=1; i<=$xP && i<${#IP_ARRAY[@]}; i++)); do - PREFILL_ARGS+=" --prefill http://${IP_ARRAY[$i]}:2322 " -done - -for ((i=$xP+1; i<${#IP_ARRAY[@]}; i++)); do - DECODE_ARGS+=" --decode http://${IP_ARRAY[$i]}:2322 " -done - -# ============================================================================= -# Node Role Assignment and Server Launch -# ============================================================================= - -if [ "$NODE_RANK" -eq 0 ]; then - echo "NODE INFO =======================================" - echo "================================================" - echo "Node List : ${SLURM_JOB_NODELIST}" - echo "Node IPs : ${IPADDRS}" - echo "Model Name : ${MODEL_NAME:-'Not specified'}" - echo "================================================" - - echo "CLUSTER INFO ====================================" - echo "================================================" - echo "${host_name}:${host_ip} is Proxy Node" - echo "${PREFILL_ARGS} are Proxy's Prefill" - echo "${DECODE_ARGS} are Proxy's Decode" - echo "================================================" - echo "Proxy server is waiting for prefill and decode nodes to be ready ..." \ - | tee /run_logs/${SLURM_JOB_ID}/proxy_NODE${NODE_RANK}.log >/dev/null - sleep 20 - TIMEOUT_SECONDS=4000 - SLEEP_SECONDS=10 - SEARCH_SIGNAL="The server is fired up and ready to roll!" - SECONDS=0 - for ((i=1; i<=$xP && i<${#IP_ARRAY[@]}; i++)); do - LOG_FILE=/run_logs/${SLURM_JOB_ID}/prefill_NODE${i}.log - #wait until prefill nodes get ready - until grep -q "${SEARCH_SIGNAL}" "${LOG_FILE}"; do - if [ $SECONDS -ge $TIMEOUT_SECONDS ]; then - echo "Awaited ${SECONDS} seconds. Timeout reached. Signal not found in prefill ${i} file" \ - | tee -a /run_logs/${SLURM_JOB_ID}/proxy_NODE${NODE_RANK}.log >/dev/null - - fi - sleep $SLEEP_SECONDS - SECONDS=$(( SECONDS + SLEEP_SECONDS)) - done - done - - for ((i=$xP+1; i<${#IP_ARRAY[@]}; i++)); do - LOG_FILE=/run_logs/${SLURM_JOB_ID}/decode_NODE${i}.log - #wait until decode nodes get ready - until grep -q "${SEARCH_SIGNAL}" "${LOG_FILE}"; do - if [ $SECONDS -ge $TIMEOUT_SECONDS ]; then - echo "Awaited ${SECONDS} seconds. Timeout reached. Signal not found in decode ${i} file" \ - | tee -a /run_logs/${SLURM_JOB_ID}/proxy_NODE${NODE_RANK}.log >/dev/null - fi - sleep $SLEEP_SECONDS - SECONDS=$(( SECONDS + SLEEP_SECONDS)) - done - done - - sleep 10 - - python -m sglang_router.launch_router \ - --pd-disaggregation \ - ${PREFILL_ARGS} \ - ${DECODE_ARGS} \ - --host 0.0.0.0 \ - --port 2322 \ - 2>&1 | tee -a /run_logs/${SLURM_JOB_ID}/proxy_NODE${NODE_RANK}.log >/dev/null & - - proxy_pid=$! - - echo "Waiting for all prefill and decode servers to be up . . ." - python $MOONCAKE_COOKBOOK_PATH/socket_barrier.py \ - --node-ips ${IPADDRS} \ - --node-ports 2322 - - echo "Proxy Server Ready for benchmarking on ${host_name}:${host_ip}" - - sleep 10 - cd /opt/mooncake-cookbook - bash /opt/mooncake-cookbook/benchmark_xPyD.sh - - echo "Killing the proxy server" - kill $proxy_pid - -elif [ "$NODE_RANK" -gt 0 ] && [ "$NODE_RANK" -le "$xP" ]; then - echo "${host_name}:${host_ip} is Prefill Node (Model: ${MODEL_NAME:-'default'})" - echo "Using prefill config: $PREFILL_MODEL_CONFIG" - - PREFILL_CMD="MC_TE_METRIC=true python3 -m sglang.launch_server \ - --model-path $MODEL_PATH \ - --disaggregation-mode prefill \ - --disaggregation-ib-device ${IBDEVICES} \ - --host ${host_ip} \ - --port 2322 \ - --stream-output \ - --trust-remote-code \ - --disaggregation-transfer-backend mooncake" - - if [[ -n "$PREFILL_MODEL_CONFIG" ]]; then - PREFILL_CMD="$PREFILL_CMD $PREFILL_MODEL_CONFIG" - fi - - eval "$PREFILL_CMD" \ - 2>&1 | tee /run_logs/${SLURM_JOB_ID}/prefill_NODE${NODE_RANK}.log >/dev/null & - - prefill_pid=$! - - echo "Waiting for proxy server to be up..." - python $MOONCAKE_COOKBOOK_PATH/socket_barrier.py \ - --node-ips ${MASTER_ADDR} \ - --node-ports 2322 - - echo "Waiting until proxy server closes..." - python $MOONCAKE_COOKBOOK_PATH/socket_wait.py \ - --remote-ip ${MASTER_ADDR} \ - --remote-port 2322 - - echo "Killing the prefill server" - kill $prefill_pid - -else - echo "${host_name}:${host_ip} is Decode Node (Model: ${MODEL_NAME:-'default'})" - echo "Using decode config: $DECODE_MODEL_CONFIG" - - DECODE_CMD="python3 -m sglang.launch_server \ - --model-path ${MODEL_PATH} \ - --disaggregation-mode decode \ - --disaggregation-ib-device ${IBDEVICES} \ - --host ${host_ip} \ - --port 2322 \ - --stream-output \ - --trust-remote-code \ - --disaggregation-transfer-backend mooncake" - - if [[ -n "$DECODE_MODEL_CONFIG" ]]; then - DECODE_CMD="$DECODE_CMD $DECODE_MODEL_CONFIG" - fi - - eval "$DECODE_CMD" \ - 2>&1 | tee /run_logs/${SLURM_JOB_ID}/decode_NODE${NODE_RANK}.log >/dev/null & - - decode_pid=$! - - echo "Waiting for proxy server to be up..." - python $MOONCAKE_COOKBOOK_PATH/socket_barrier.py \ - --node-ips ${MASTER_ADDR} \ - --node-ports 2322 - - echo "Waiting until proxy server closes..." - python $MOONCAKE_COOKBOOK_PATH/socket_wait.py \ - --remote-ip ${MASTER_ADDR} \ - --remote-port 2322 - - echo "Killing the decode server" - kill $decode_pid - -fi - -# ============================================================================= -# Cleanup -# ============================================================================= - -echo "Killing the etcd server" -kill $etcd_pid - -echo "Script completed successfully" -exit 0