diff --git a/examples/training/nemo-rl/README.md b/examples/training/nemo-rl/README.md index c36eac09a..6aada437a 100644 --- a/examples/training/nemo-rl/README.md +++ b/examples/training/nemo-rl/README.md @@ -32,6 +32,15 @@ NVIDIA Resiliency Extension provides process-level fault tolerance: - **Straggler Detector**: Monitors GPU kernel timing across ranks to detect slow GPUs - **Checkpoint integration**: Saves model state to FSx for resume-on-restart +## Variants + +This directory nests a distinct NeMo-RL example that shares the framework but targets a different +communication mechanism and scale: + +| Variant | Focus | Scale | +|---------|-------|-------| +| [`deepep-v2-efa/`](./deepep-v2-efa) | NeMo-RL (GRPO) + Megatron-LM MoE expert-parallel all-to-all over **DeepEP V2's NCCL-GIN CPU-proxy** on **AWS EFA** (`Qwen3-30B-A3B`); built NGC-from-scratch. See its README for the measured-vs-staged breakdown. | 2× p5.48xlarge (16× H100) or p5en (H200) | + ## Architecture ``` diff --git a/examples/training/nemo-rl/deepep-v2-efa/.dockerignore b/examples/training/nemo-rl/deepep-v2-efa/.dockerignore new file mode 100644 index 000000000..4bc3096b9 --- /dev/null +++ b/examples/training/nemo-rl/deepep-v2-efa/.dockerignore @@ -0,0 +1,11 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +# .gitignore does NOT apply to a Docker build context — without this file the +# filled-in `env_vars` (HF_TOKEN + AWS account id) is uploaded to the daemon, +# and to any remote/CI builder, on every `docker build .` in this directory. +# No COPY in nemo-rl.Dockerfile is broad (each names a specific file), so nothing +# lands in an image layer today; this closes the context-upload exposure and +# guards against a future broad COPY. +env_vars +*.log +__pycache__/ diff --git a/examples/training/nemo-rl/deepep-v2-efa/.gitignore b/examples/training/nemo-rl/deepep-v2-efa/.gitignore new file mode 100644 index 000000000..3f86e687e --- /dev/null +++ b/examples/training/nemo-rl/deepep-v2-efa/.gitignore @@ -0,0 +1,9 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 + +# Local environment file with filled-in secrets/values — never commit +env_vars + +# Run artifacts +*.log +__pycache__/ diff --git a/examples/training/nemo-rl/deepep-v2-efa/README.md b/examples/training/nemo-rl/deepep-v2-efa/README.md new file mode 100644 index 000000000..1bdc03c1f --- /dev/null +++ b/examples/training/nemo-rl/deepep-v2-efa/README.md @@ -0,0 +1,287 @@ + +# NeMo-RL + DeepEP V2 — MoE expert-parallel all-to-all over AWS EFA (NCCL-GIN CPU-proxy) + +This test case wires **[NeMo-RL](https://github.com/NVIDIA-NeMo/RL)** (GRPO post-training) and +**[Megatron-LM](https://github.com/NVIDIA/Megatron-LM)** (the training backend) to **DeepEP V2's +NCCL backend** (`ElasticBuffer`, merged upstream in +[deepseek-ai/DeepEP#605](https://github.com/deepseek-ai/DeepEP/pull/605)) on **AWS EFA**. Unlike +the NVSHMEM-path DeepEP samples (e.g. the SGLang sibling), V2's network traffic rides **NCCL's GIN +(GPU-Initiated Networking) device API**, which on EFA means **aws-ofi-nccl with GIN compiled in**, +running GIN's **CPU-proxy** mode. The image is built NGC-from-scratch from public sources only; the +recipe takes you from image build → static verify → a cross-node dispatch/combine probe → an N-step +Megatron MoE training gate — and documents, honestly, which parts of the full GRPO rollout path +still depend on draft upstream PRs. + +## The mechanism chain + +```text +NeMo-RL (GRPO) ── nemo_rl.models.megatron ── Megatron-core MoE flex dispatcher (fused_a2a) + └─ deep_ep.ElasticBuffer (DeepEP V2, NCCL backend) + └─ NCCL 2.30.4 GIN device API (NCCL_GIN_TYPE=2, CPU-proxy) + └─ aws-ofi-nccl @9c44d34 (ncclGinPlugin, gdrcopy) + └─ libfabric ── EFA (efa-direct, SRD) +``` + +## Measured vs staged (read this first) + +**Measured (2026-05-06, "Wave 28"):** the NeMo-RL 0.5.0rc0 full stack ran E2E on **2× +p5.48xlarge (16× H100)** over EFA with this exact mechanism chain — imports + trainer bring-up, +rollout shape `[64, 8192]` generated in **9.45 s**, Megatron "Shape Y" 3-step training with loss +**26.41 → 24.62**, `ElasticBuffer` confirmed as the active buffer class, and EFA cross-node +hardware TX counters advanced. That run used a privately prebuilt image cascade, which this repo +does not allow as a base — **this folder re-cuts the same substrate NGC-from-scratch from public +sources.** + +**Staged, NOT re-measured:** the image assembly in this folder is **build-staged and has not been +cluster-re-run**; no performance numbers are published from this folder. The **full GRPO +rollout-over-DeepEP path depends on 2 draft upstream PRs** (opt-in image layer, default **OFF** — +see below). What the recipe gates re-verify on the **baseline** image (the 2 opt-in rollout PRs +OFF — it does carry the standardized AWS GIN plugin lineage, incl. the closed-unmerged aws-ofi-nccl +PR#1351, and the `amazon-contributing/DeepEP` fork; see the Pins table): static +substrate asserts (`verify-image.sh`), cross-node `ElasticBuffer` dispatch/combine with an EFA +TX-counter assert (`run-rollout-probe.sh`), and a loss-decreasing Megatron MoE train step on the +stock `alltoall` dispatcher (`train-step.sh`). Never read a build-gate as an E2E pass. + +## Pins (every one justified) + +| Component | Pin | Why | +|---|---|---| +| Base image | `nvcr.io/nvidia/pytorch:26.02-py3` **@sha256:bbc2b67e…** | Same NGC base as the slime RL sibling. Bakes torch 2.11/CUDA 13 with TransformerEngine/apex/flash-attn compiled against that exact ABI — what megatron.core's H100 path needs. Digest-pinned: it is the ABI anchor the whole image builds around, so a silent tag re-push is the most consequential drift possible here. | +| EFA installer | 1.48.0 | The userspace of the measured NCCL-GIN substrate. Bumping is a re-measure event. | +| NCCL | `v2.30.4-1` = commit `1933fdd6` (source build) | GIN device API generation (`nccl_device.h` asserted at build). Commit-pinned, not the bare tag — held to the same moving-ref standard as the other source pins. The NGC base bakes an OLDER NCCL line (2.29.x), so this source-built copy is a deliberately newer line and the `00-nccl-gin.conf` ld.so.conf entry is **load-bearing** — it makes this GIN-capable build win the loader search over the base's baked copy (`verify-image.sh` asserts which one resolves). Source-built so DeepEP has one controlled header/lib root. | +| aws-ofi-nccl | commit `9c44d34` + PR#1351 head `c2e773d` | The GIN CPU-proxy plugin lineage this folder standardises on (same pins as the TensorRT-LLM NcclEP sibling). These SHAs postdate the Wave-28 run, so they are the standardised lineage, not that run's exact pins. Immutable SHAs — `refs/pull/N/head` is a moving ref. | +| gdrcopy | commit `c91ad9f` (= v2.5.2) | GIN **requires** gdrcopy compiled in (trap 4). Commit pin, not tag. | +| DeepEP | `97d8f9bc` (**`amazon-contributing/DeepEP`** fork HEAD) | The AWS EPv2/NCCL-Gin tree. Carries EPv2 (ElasticBuffer + NCCL backend) **and the two former-draft [deepseek-ai/DeepEP#612](https://github.com/deepseek-ai/DeepEP/pull/612) EFA fixes in-code** — the `get_rdma_gbs()` sysfs link-rate fast path (`deep_ep/utils/envs.py`) and the auto-QP overflow clamp (`deep_ep/buffers/elastic.py`) — so **no #612 patch is applied on any flavor**. This is the same fork the repo's [`micro-benchmarks/expert-parallelism`](../../../../micro-benchmarks/expert-parallelism) DeepEP-V2 sample already pins (`setup_deepep_gin.sh`). Full 40-char SHA pinned as `DEEPEP_SHA` in the Dockerfile. | +| Megatron-LM | `19deef67` (main) | The **exact base of draft PR #4632** (ElasticBuffer in the flex dispatcher). | +| NeMo-RL | `46be4e8` | The **exact base (parent commit) of draft PR #2410** — declares `requires-python ">=3.12"`, which the py3.12 NGC base satisfies. `requirements.txt` is generated from this revision. NOT #2411's base `cc75cad` (which bumps `requires-python` to `>=3.13.13` and hard-fails `pip install -e` on this base). The measured Wave-28 evidence ran 0.5.0rc0; re-pinning to a release tag is a re-measure event. | +| GPU arch | `TORCH_CUDA_ARCH_LIST=9.0` (H100/H200) | The only measured arch. Blackwell needs an arch-list override and a re-measure. | + +## The 2 draft upstream PRs (opt-in layer, default OFF) + +Baked only with `--build-arg APPLY_DRAFT_ROLLOUT_PATCHES=1`; commits pinned at immutable SHAs in +[`patches/apply_nemo_rl_patches.py`](patches/apply_nemo_rl_patches.py), applied fail-loud, +self-neutralizing once merged upstream. **The baseline image has zero dependence on them.** +PR states below are as of 2026-08-25 — check them before relying on this table. + +> DeepEP is **not** in this table: the baseline pins the `amazon-contributing/DeepEP` fork, which +> already carries the former draft [deepseek-ai/DeepEP#612](https://github.com/deepseek-ai/DeepEP/pull/612) +> EFA fixes in-code (see the Pins table), so there is no DeepEP patch to opt into on either flavor. + +| PR | State | What it carries | +|---|---|---| +| [NVIDIA/Megatron-LM#4632](https://github.com/NVIDIA/Megatron-LM/pull/4632) | open | DeepEP **V2 ElasticBuffer** support in the MoE flex dispatcher (`fused_a2a.py`) — without it, `--moe-enable-deepep` binds the V1 NVSHMEM `Buffer`, which this NCCL-GIN image intentionally does not build. | +| [NVIDIA-NeMo/RL#2410](https://github.com/NVIDIA-NeMo/RL/pull/2410) | draft, closed unmerged | `LD_LIBRARY_PATH` re-export for OFI plugin discovery in NeMo-RL's own containers, plus the worked 2-node EFA GRPO recipe config (`examples/configs/recipes/llm/aws-efa-grpo-qwen3-30ba3b-2n8g-megatron.yaml`) the full rollout path uses. Applied at its parent commit `46be4e8` (= `NEMO_RL_SHA`), so it lands `--check`-clean. | + +> **Not applied: [NVIDIA-NeMo/RL#2411](https://github.com/NVIDIA-NeMo/RL/pull/2411)** (deep_ep pin bump). Its base is `cc75cad` — 116 commits ahead of `46be4e8`, across a `requires-python` bump to `>=3.13.13` — so it neither applies to this tree nor belongs on this py3.12 base, and it is metadata-only anyway (deep_ep is built from `/opt/DeepEP`, not NeMo-RL's pin). Retained here as a note so the pin history is auditable. + +## The integration traps + +1. **NeMo-RL's pyproject pins torch exactly** (`torch==2.9.0` at the pinned SHA). Letting pip + resolve it would *downgrade* the NGC-baked torch 2.11 and orphan the baked + TransformerEngine/apex/flash-attn ABI. The Dockerfile installs NeMo-RL `--no-deps` and carries + the rest of its dependency set in [`requirements.txt`](requirements.txt) (re-diff on every + `NEMO_RL_SHA` bump). +2. **Stock DeepEP's SM/QP auto-sizers are EFA-blind — the fork fixes them in-code.** On stock + `deepseek-ai/DeepEP` the auto-QP formula (`num_sms*16+1`) overruns aws-ofi-nccl's 128-slot GIN + request ring (hard assert surfaced as `CUDA_ERROR_LAUNCH_FAILED` at the *first* dispatch), and + `get_rdma_gbs()` reads 0 on EFA. The `amazon-contributing/DeepEP` fork this image pins carries + the former-draft #612 fixes structurally — a `get_rdma_gbs()` sysfs link-rate fast path and an + auto-QP overflow clamp — so its auto-sizers are EFA-aware on both flavors. The probe still passes + `num_allocated_qps`/`num_sms`/`num_qps` **explicitly** (`EP_NUM_QPS=2` is the value the p5en + evidence validated) so it is deterministic and auto-sizer-independent; that value survives the + fork's clamp unchanged (`max(2, min(2, max_unordered_gin_qps)) == 2`). +3. **Two NCCLs in one image — the resolved path matters.** The NGC base bakes its own libnccl in a + different directory; if it wins the loader search, you silently run a non-GIN-verified copy. + The image ranks the source build first via `/etc/ld.so.conf.d/00-nccl-gin.conf`, and + `verify-image.sh` asserts **which `libnccl.so.2` the loader actually resolves** (via `dlopen` + + `/proc/self/maps`, not the `ldconfig -p` cache order) *and* its version string. +4. **GIN needs gdrcopy at compile time and gdrdrv at run time.** An aws-ofi-nccl built without + `gdrapi.h` carries "GDRCopy support not available at compile time" and GIN init fails at run + time; the setup script asserts that string is *absent* from the built plugin. At run time, + clusters without a gdrdrv device plugin need `privileged: true` (the unprivileged device cgroup + blocks `open("/dev/gdrdrv")` with EPERM even after CAP_MKNOD — the manifest header documents the + trade-off). +5. **The EFA installer's NGC auto-detect reroutes the plugin install.** On an NGC base the + installer would install the stock `libnccl-ofi-ngc` plugin — which does not carry GIN — and two + plugins on the loader path is a which-one-won guessing game. The Dockerfile passes + `--disable-ngc --disable-build-ngc` and builds the GIN plugin from source. +6. **`flex` on an unpatched image is refused, not degraded.** Megatron's flex dispatcher with + `moe_enable_deepep` needs #4632; on the baseline image `train-step.sh` exits with the reason + instead of failing later inside megatron — and it never silently falls back to a different + dispatcher than the one you asked for. +7. **Opt-in flags use `"${VAR:-default}"` in `env_vars.example`** so a value pre-set on the command + line (`APPLY_DRAFT_ROLLOUT_PATCHES=1 docker build ...`) survives the file being sourced after + it. A hardcoded `export VAR="0"` silently clobbers the build you thought you asked for. +8. **CPU-proxy means CPU is on the data path.** `NCCL_GIN_TYPE=2` runs GIN's proxy threads on the + host cores; the manifest pins requests == limits (Guaranteed QoS) so CFS throttling under node + pressure cannot silently degrade the transport. + +## Runtime requirements (baked in the image ENV, repeated in the manifest and scripts) + +| Env | Why | +|---|---| +| `NCCL_GIN_TYPE=2`, `NCCL_GIN_ENABLE=1` | GIN CPU-proxy — the EFA-viable GIN mode | +| `OFI_NCCL_GIN_GDAKI=0` | GPU-initiated GIN is not the shipped path on EFA | +| `FI_PROVIDER=efa`, `FI_EFA_USE_DEVICE_RDMA=1` | EFA with GPU-direct RDMA | +| `FI_EFA_ENABLE_SHM_TRANSFER=0`, `FI_EFA_FORK_SAFE=1` | no SHM shortcut; fork-safe for the proxy | +| `NCCL_NET_PLUGIN=/opt/aws-ofi-nccl/lib/libnccl-net-ofi.so` | the GIN-capable plugin, explicitly | +| `NCCL_NVLS_ENABLE=0` | prevents NVLS init failures on H100/H200 | +| `DEEP_EP_USE_V2_SHIM=0` | V2-native path, no compatibility shim | + +## Hardware requirements + +2× `p5.48xlarge` (8× H100, 32 EFA NICs/node — the measured topology) or `p5en.48xlarge` +(8× H200, 16 EFA NICs/node; set `EFA_PER_NODE=16`). One CPU node for the Ray head. FSx for +Lustre PVC (full GRPO path only — the recipe gates need no shared storage). 2Mi hugepages +pre-allocated on the GPU nodes (`hugepages-2Mi: 5120Mi` per pod, or the pods sit Pending). + +## Prerequisites + +- An EKS / SageMaker HyperPod EKS cluster with the EFA device plugin, the NVIDIA device plugin, + and the [KubeRay operator](https://docs.ray.io/en/latest/cluster/kubernetes/getting-started.html). +- Docker + network access to `nvcr.io`, `pypi.org`, `github.com`, `efa-installer.amazonaws.com`. +- An image registry you own (ECR); **do not point at anyone's private registry**. + +## Quick Start + +### 1. Configure environment variables + +```bash +cp env_vars.example env_vars # env_vars is gitignored +vim env_vars # REGISTRY/IMAGE/TAG, NAMESPACE, FSX_CLAIM, EFA_PER_NODE, ... +source env_vars +``` + +### 2. Build and push (baseline, and optionally the draft-PR flavor) + +```bash +aws ecr get-login-password --region ${AWS_REGION} | docker login --username AWS --password-stdin ${REGISTRY} +aws ecr create-repository --repository-name ${IMAGE} --region ${AWS_REGION} || true + +# baseline (draft rollout PRs OFF — still carries the AWS GIN lineage + amazon-contributing/DeepEP fork; see Pins) +docker build -f nemo-rl.Dockerfile -t ${FULL_IMAGE} . +# opt-in flavor with the 2 draft PRs baked (use a DISTINCT tag — never overwrite the baseline) +docker build -f nemo-rl.Dockerfile --build-arg APPLY_DRAFT_ROLLOUT_PATCHES=1 \ + -t ${FULL_IMAGE}-draftprs . + +docker push ${FULL_IMAGE} +# push the draft-PR flavor too if you built it (distinct tag — never overwrite the baseline) +docker push ${FULL_IMAGE}-draftprs +``` + +### 3. Gate the image before any cluster deploy + +```bash +recipe/verify-image.sh ${FULL_IMAGE} +# ... ALL CHECKS PASS +``` + +### 4. Deploy the Ray cluster + +```bash +kubectl create namespace ${NAMESPACE} || true +kubectl create secret generic hf-token --from-literal=HF_TOKEN=${HF_TOKEN} -n ${NAMESPACE} +envsubst < kubernetes/raycluster.yaml | kubectl apply -f - +kubectl -n ${NAMESPACE} get pods -w # 1 head + ${NUM_NODES} workers +``` + +### 5. Cross-node rollout probe (minutes, no weights) + +Drives the real `deep_ep.ElasticBuffer` dispatch/combine across the node boundary against a +closed-form oracle, and asserts the EFA hardware TX counters advanced: + +```bash +W0=$(kubectl -n ${NAMESPACE} get pod -l ray.io/node-type=worker -o jsonpath='{.items[0].metadata.name}') +W1=$(kubectl -n ${NAMESPACE} get pod -l ray.io/node-type=worker -o jsonpath='{.items[1].metadata.name}') +W0_IP=$(kubectl -n ${NAMESPACE} get pod ${W0} -o jsonpath='{.status.podIP}') +kubectl -n ${NAMESPACE} exec ${W1} -c ray-worker -- bash -lc \ + "nohup /opt/run-rollout-probe.sh worker ${W0_IP} 1 > /tmp/probe.log 2>&1 &" +kubectl -n ${NAMESPACE} exec ${W0} -c ray-worker -- /opt/run-rollout-probe.sh leader ${W0_IP} +# ... ROLLOUT-PROBE PASS — ElasticBuffer dispatch/combine over EFA verified (tx +NNNB) +``` + +### 6. N-step MoE training gate + +Baseline image → the stock `alltoall` dispatcher (no deep_ep on the data path, still NCCL over +EFA); patched image → `MOE_DISPATCHER=flex` for the DeepEP V2 ElasticBuffer path: + +```bash +kubectl -n ${NAMESPACE} exec ${W1} -c ray-worker -- bash -lc \ + "nohup /opt/train-step.sh worker ${W0_IP} 1 > /tmp/train.log 2>&1 &" +kubectl -n ${NAMESPACE} exec ${W0} -c ray-worker -- /opt/train-step.sh leader ${W0_IP} +# ... TRAIN-STEP-PASS dispatcher=alltoall world=16 ep=16 +# on the -draftprs image, run the flex (DeepEP V2 ElasticBuffer) dispatcher — set +# MOE_DISPATCHER=flex on BOTH nodes (it is one torchrun job across both; if only one +# rank sets flex the ranks disagree on the dispatcher and the collective hangs): +# kubectl -n ${NAMESPACE} exec ${W1} -c ray-worker -- bash -lc \ +# "MOE_DISPATCHER=flex nohup /opt/train-step.sh worker ${W0_IP} 1 > /tmp/train.log 2>&1 &" +# kubectl -n ${NAMESPACE} exec ${W0} -c ray-worker -- \ +# env MOE_DISPATCHER=flex /opt/train-step.sh leader ${W0_IP} +``` + +### 7. Full GRPO run (STAGED — draft-PR image only, not re-measured from this folder) + +The patched image carries the worked 2-node EFA GRPO recipe config from NeMo-RL#2410 at +`/opt/NeMo-RL/examples/configs/recipes/llm/aws-efa-grpo-qwen3-30ba3b-2n8g-megatron.yaml` +(Qwen3-30B-A3B, Megatron backend, `moe_token_dispatcher_type=flex`, `moe_enable_deepep=true`). +Stage the model with `kubernetes/data-prep-pod.yaml`, then launch per +[NeMo-RL's GRPO docs](https://github.com/NVIDIA-NeMo/RL) from the Ray head with that config. +Treat results as your own measurement — this folder publishes none for this path. + +## Known limitations (honest list) + +- **Build-staged, not cluster-re-run.** The NGC-from-scratch assembly here reproduces the measured + Wave-28 mechanism chain from public sources, but this exact image has not itself been re-run on + a cluster. The recipe gates exist so you (or we, next capacity window) can re-verify cheaply. +- **The full rollout path is draft-PR-dependent.** Two upstream PRs, closed-unmerged upstream + (see the table). If upstream supersedes them, the patch layer fails loud or self-neutralizes — + either way the image never ships an ambiguous patch state. +- **Baseline DeepEP gates use explicit SM/QP counts** (trap 2). A probe pass with explicit counts + is deterministic by design; it does not exercise the fork's auto-sizers, and it certifies nothing + about *stock* upstream's EFA-blind auto-sizing (the defect the fork fixes in-code). +- **NCCL topology XML on 32-NIC p5 nodes:** stock NCCL can hit the open issue + [NVIDIA/nccl#2160](https://github.com/NVIDIA/nccl/issues/2160) (`NCCL_TOPO_XML_MAX_NODES=256` + overflow during intra-node XML fusion). If NCCL init fails with a topo-XML error on + p5.48xlarge, rebuild Layer 4 with the define raised (documented one-liner in the issue) — not + baked here because it is a non-upstream one-off. +- **No performance numbers.** Dispatch/combine latency and GRPO throughput on this substrate are + future work; for kernel-level EP benchmarks see + [`micro-benchmarks/expert-parallelism`](../../../../micro-benchmarks/expert-parallelism) — note + that benchmark runs DeepEP V2 on the **EFA-GDA** NCCL-GIN backend (`NCCL_GIN_TYPE=5`), not this + folder's CPU-proxy one (`NCCL_GIN_TYPE=2`). + +## File structure + +```text +deepep-v2-efa/ +├── README.md <- you are here +├── nemo-rl.Dockerfile <- NGC-from-scratch image (baseline + opt-in draft-PR layer) +├── setup_nemo_rl_deepep_efa.sh <- aws-ofi-nccl GIN + DeepEP V2 source builds (in-tree, COPY'd) +├── requirements.txt <- NeMo-RL deps minus the NGC-baked ABI anchors +├── env_vars.example <- copy to env_vars (gitignored), fill in, source +├── patches/ +│ └── apply_nemo_rl_patches.py <- the 2 draft PRs, pinned SHAs, fail-loud, self-neutralizing +├── recipe/ +│ ├── verify-image.sh <- static substrate gate (run before any deploy) +│ ├── run-rollout-probe.sh <- cross-node ElasticBuffer probe + EFA TX-counter assert +│ ├── probe_rollout.py <- the torchrun probe body (oracle-checked dispatch/combine) +│ ├── train-step.sh <- N-step Megatron MoE training gate launcher +│ └── train_moe_step.py <- the torchrun train body (loss finite+decreasing+agreeing) +└── kubernetes/ + ├── raycluster.yaml <- 1 CPU head + N GPU workers (EFA, hugepages, Guaranteed QoS) + └── data-prep-pod.yaml <- stage model/dataset onto FSx (full GRPO path only) +``` + +## References + +- [amazon-contributing/DeepEP](https://github.com/amazon-contributing/DeepEP) — the AWS EPv2/NCCL-Gin + fork this image builds (pinned at `97d8f9bc`); carries the former-draft #612 EFA fixes in-code +- [deepseek-ai/DeepEP](https://github.com/deepseek-ai/DeepEP) — upstream EPv2 / ElasticBuffer (PR#605, merged) +- [aws-ofi-nccl](https://github.com/aws/aws-ofi-nccl) — the GIN-capable NCCL network plugin +- [NeMo-RL](https://github.com/NVIDIA-NeMo/RL) and [Megatron-LM](https://github.com/NVIDIA/Megatron-LM) +- Sibling test cases: [`slime`](../../slime) (RL on HyperPod EKS as a Ray cluster — this folder + mirrors its shape), [`inference/sglang/dsr1-deepep-efa`](../../../inference/sglang/dsr1-deepep-efa) + (the NVSHMEM-path DeepEP serving sample) +- [`micro-benchmarks/expert-parallelism`](../../../../micro-benchmarks/expert-parallelism) — + kernel-level EP benchmarks, including a DeepEP V2 EFA-GDA (`NCCL_GIN_TYPE=5`) benchmark — a + different GIN backend from this folder's CPU-proxy (`NCCL_GIN_TYPE=2`) diff --git a/examples/training/nemo-rl/deepep-v2-efa/env_vars.example b/examples/training/nemo-rl/deepep-v2-efa/env_vars.example new file mode 100644 index 000000000..56952a77d --- /dev/null +++ b/examples/training/nemo-rl/deepep-v2-efa/env_vars.example @@ -0,0 +1,108 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +# ============================================================ +# NeMo-RL + DeepEP V2 over EFA - Environment Variables +# +# Copy this file to env_vars (gitignored) and fill in your values: +# cp env_vars.example env_vars && vim env_vars +# Source it before building or applying manifests: +# source env_vars +# ============================================================ + +# ----- AWS / ECR ----- +# Region and account are derived from your current AWS credentials/config so +# nothing environment-specific is hard-coded. Override AWS_REGION if needed. +export AWS_REGION="${AWS_REGION:-$(aws configure get region)}" +export AWS_ACCOUNT_ID="$(aws sts get-caller-identity --query Account --output text)" +export REGISTRY="${AWS_ACCOUNT_ID}.dkr.ecr.${AWS_REGION}.amazonaws.com/" +export IMAGE="nemo-rl-deepep-efa" +# Immutable tag - never "latest": with imagePullPolicy: IfNotPresent a node +# that cached "latest" silently keeps running the OLD image after a rebuild. +export TAG="v1-20260825" +export FULL_IMAGE="${REGISTRY}${IMAGE}:${TAG}" + +# ----- Opt-in draft-PR image flavor ----- +# 1 = bake the 2 draft upstream PRs (NeMo-RL#2410, Megatron-LM#4632) that the +# full GRPO rollout-over-DeepEP path needs; 0 = the upstream-only baseline. +# (DeepEP needs no patch: the amazon-contributing/DeepEP fork the image pins +# carries the former draft DeepEP#612 fixes in-code, on both flavors.) The ":-" +# default keeps a value pre-set on the command line (APPLY_DRAFT_ROLLOUT_PATCHES=1 +# docker build ...) from being clobbered when this file is sourced afterwards. +export APPLY_DRAFT_ROLLOUT_PATCHES="${APPLY_DRAFT_ROLLOUT_PATCHES:-0}" + +# ----- HuggingFace (full GRPO path only; the recipe gates need no weights) ----- +export HF_TOKEN="" # <-- set your HuggingFace token here if the model is gated + +# ----- Model / shape (the Wave-28 measured shape twin) ----- +# Qwen3-30B-A3B: 128 routed experts, top-k 8, hidden 2048. Any MoE whose +# routed-expert count divides by the EP size works; the probe asserts it. +export MODEL_NAME="Qwen/Qwen3-30B-A3B" +export MODEL_LOCAL="/fsx/models/Qwen3-30B-A3B" + +# ----- Cluster ----- +# HOW THESE REACH THE GATES: sourcing this file sets vars in YOUR shell only — a +# client-side `source` does NOT cross `kubectl exec`; only the pod's own env +# does. So `envsubst < kubernetes/raycluster.yaml` substitutes NUM_NODES / the +# EP_* shape into the worker pod's `env:` block, and the pod carries them. +# run-rollout-probe.sh / train-step.sh read NNODES (not NUM_NODES); the manifest +# bridges the names by emitting `{ name: NNODES, value: "${NUM_NODES}" }`. Set +# NUM_NODES here, and both the RayCluster replica count and the in-pod NNODES the +# launchers read move together. +export NAMESPACE="nemo-rl-deepep" +export FSX_CLAIM="fsx-claim" +export NUM_NODES=2 +export GPUS_PER_NODE=8 +# Per-node EFA NIC count differs by instance type: +# p5.48xlarge (H100): 32 | p5en.48xlarge (H200): 16 +export EFA_PER_NODE=32 +export INSTANCE_TYPE="p5.48xlarge" + +# ----- Probe / train-step shape knobs (defaults mirror the model above) ----- +export EP_EXPERTS=128 +export EP_TOPK=8 +export EP_HIDDEN=2048 +export EP_TOKENS=128 +# Explicit SM/QP counts for the probe. WHY: the amazon-contributing/DeepEP fork +# carries the former #612 EFA fixes in-code (get_rdma_gbs() sysfs link-rate + +# the auto-QP overflow clamp), so its auto-sizers are EFA-aware — but the probe +# pins these anyway so it is deterministic and auto-sizer-independent. 2 QPs is +# the value the p5en evidence validated; it survives the fork's clamp unchanged. +export EP_NUM_SMS=8 +export EP_NUM_QPS=2 + +# ----- EFA / NCCL-GIN transport contract ----- +# Baked into the image ENV and repeated in kubernetes/raycluster.yaml's pod env, +# so a `kubectl exec` shell already inherits them from the container. Exported +# here too only so a LOCAL run (docker run on an EFA host, outside k8s) carries +# the same contract — a client-side `source` does NOT propagate across +# `kubectl exec`. +export FI_PROVIDER=efa +export FI_EFA_USE_DEVICE_RDMA=1 +export FI_EFA_FORK_SAFE=1 +export FI_EFA_ENABLE_SHM_TRANSFER=0 +export NCCL_GIN_TYPE=2 # 2 = CPU-proxy GIN (the EFA-viable path) +export NCCL_GIN_ENABLE=1 +export OFI_NCCL_GIN_GDAKI=0 # GPU-initiated GIN is not the shipped path on EFA +export OFI_NCCL_PROTOCOL=RDMA +export NCCL_NVLS_ENABLE=0 # prevents NVLS init failures on H100/H200 +export NCCL_NET_PLUGIN=/opt/aws-ofi-nccl/lib/libnccl-net-ofi.so +# NOTE: EP_EFA_MAX_QPS / EP_EFA_RDMA_GBS are intentionally absent. They were the +# knobs the old draft DeepEP#612 patch read; the amazon-contributing/DeepEP fork +# resolves both structurally (get_rdma_gbs() sysfs link-rate + the C++ +# _C.{min,max}_unordered_gin_qps clamp), so nothing in deep_ep reads them. The one +# live QP knob is EP_NUM_QPS above (the probe's explicit num_allocated_qps). + +# ----- DeepEP V2 selection ----- +export DEEP_EP_USE_V2_SHIM=0 # V2-native path, no compatibility shim +export HAVE_DEEP_EP_V2=True # rollout bridge feature flag (draft-PR rollout path) + +# ----- NVSHMEM contract (INERT at run time on this image - kept for the rebuild case) ----- +# deep_ep links NVSHMEM as a BUILD-TIME dependency (upstream setup.py links it +# unconditionally; see requirements.txt), but the V2 NCCL-GIN backend does NOT +# USE NVSHMEM at run time - the network path is NCCL-GIN - so these values are +# inert here. If you rebuild with the legacy NVSHMEM backend instead, these are +# the required values on EFA - note IBGDA must be 0 (proxy-based), despite what +# most upstream DeepEP manifests say. +export NVSHMEM_REMOTE_TRANSPORT=libfabric +export NVSHMEM_LIBFABRIC_PROVIDER=efa +export NVSHMEM_IB_ENABLE_IBGDA=0 diff --git a/examples/training/nemo-rl/deepep-v2-efa/kubernetes/data-prep-pod.yaml b/examples/training/nemo-rl/deepep-v2-efa/kubernetes/data-prep-pod.yaml new file mode 100644 index 000000000..37858a82c --- /dev/null +++ b/examples/training/nemo-rl/deepep-v2-efa/kubernetes/data-prep-pod.yaml @@ -0,0 +1,68 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +# Lightweight CPU pod for staging the model + dataset onto FSx (full GRPO path +# only — the recipe gates in recipe/ need NO weights and NO dataset). +# +# Create the HF token Secret first (same Secret raycluster.yaml consumes): +# kubectl create secret generic hf-token \ +# --from-literal=HF_TOKEN=$HF_TOKEN -n ${NAMESPACE} +# +# envsubst < kubernetes/data-prep-pod.yaml | kubectl apply -f - +# kubectl -n ${NAMESPACE} exec -it data-prep -- bash +# # inside: pip install "huggingface_hub[cli]==1.28.0" && \ +# # hf download ${MODEL_NAME} --local-dir ${MODEL_LOCAL} +# # NOTE: the GRPO recipe consumes the hub id (policy.model_name: +# # Qwen/Qwen3-30B-A3B), not ${MODEL_LOCAL} — set HF_HOME on the Ray pods to +# # ${MODEL_LOCAL}'s parent, or override model_name, if you stage to a path. +apiVersion: v1 +kind: Pod +metadata: + name: data-prep + namespace: ${NAMESPACE} + labels: + app: nemo-rl-deepep-data-prep +spec: + # CPU-only staging workload — keep it off GPU nodes so a `sleep infinity` pod + # can never park on (and hold) a p5 node a GPU worker needs. Same label + + # rationale as raycluster.yaml's head group. + affinity: + nodeAffinity: + requiredDuringSchedulingIgnoredDuringExecution: + nodeSelectorTerms: + - matchExpressions: + - key: nvidia.com/gpu.present + operator: NotIn + values: + - "true" + automountServiceAccountToken: false + containers: + - name: data-prep + # Pinned patch release, never a floating tag — same "never latest" + # reasoning env_vars.example states for the training image. + image: python:3.12.14-slim + command: ["sleep", "infinity"] + env: + # Reuse the hf-token Secret (created above); a plaintext env value would + # sit in etcd for the life of this sleep-infinity pod and is visible in + # `kubectl get pod -o yaml`. + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token + key: HF_TOKEN + resources: + requests: + cpu: "4" + memory: "16Gi" + limits: + cpu: "8" + memory: "32Gi" + volumeMounts: + - name: fsx + mountPath: /fsx + volumes: + - name: fsx + persistentVolumeClaim: + claimName: ${FSX_CLAIM} + restartPolicy: Never + terminationGracePeriodSeconds: 30 diff --git a/examples/training/nemo-rl/deepep-v2-efa/kubernetes/raycluster.yaml b/examples/training/nemo-rl/deepep-v2-efa/kubernetes/raycluster.yaml new file mode 100644 index 000000000..dc172ccf9 --- /dev/null +++ b/examples/training/nemo-rl/deepep-v2-efa/kubernetes/raycluster.yaml @@ -0,0 +1,246 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +# NeMo-RL + DeepEP V2 over EFA — RayCluster (1 CPU head + ${NUM_NODES} GPU workers). +# +# Apply with variables substituted (source env_vars first — see env_vars.example): +# envsubst < kubernetes/raycluster.yaml | kubectl apply -f - +# +# TOPOLOGY (honest shape — read before scaling expectations): +# head: Ray GCS + dashboard only, CPU node, no GPUs (slime-sibling pattern). +# workers: one per GPU node. The recipe GATES run in these pods via kubectl exec +# (torchrun leader/worker across the two pods — see README Quick Start): +# kubectl -n ${NAMESPACE} exec -- bash -lc \ +# 'nohup /opt/run-rollout-probe.sh worker 1 > /tmp/probe.log 2>&1 &' +# kubectl -n ${NAMESPACE} exec -- /opt/run-rollout-probe.sh leader +# The FULL NeMo-RL GRPO run additionally needs the opt-in draft-PR image flavor +# (APPLY_DRAFT_ROLLOUT_PATCHES=1) — see README "Measured vs staged". +# +# Image: build nemo-rl.Dockerfile (NGC-from-scratch, public sources only) and push to +# YOUR registry; ${FULL_IMAGE} comes from env_vars. Do NOT point at anyone's private +# registry. +# +# privileged: true is REQUIRED where /dev/gdrdrv is not advertised by a device plugin +# (the unprivileged device cgroup blocks open("/dev/gdrdrv") with EPERM even after +# CAP_MKNOD; without GDRCopy aws-ofi-nccl reports ginType=NONE and the NCCL-GIN +# CPU-proxy path is dead). If your cluster runs a gdrdrv device plugin, drop privileged +# and request the device instead. +# +# Scale seams (per-node EFA NIC count differs by instance type — set in env_vars): +# p5.48xlarge (H100): EFA_PER_NODE=32 | p5en.48xlarge (H200): EFA_PER_NODE=16 +apiVersion: ray.io/v1 +kind: RayCluster +metadata: + name: nemo-rl-deepep + namespace: ${NAMESPACE} + labels: + app: nemo-rl-deepep +spec: + rayVersion: "2.49.2" # matches the ray pin in requirements.txt + enableInTreeAutoscaling: false + headGroupSpec: + rayStartParams: + dashboard-host: "0.0.0.0" + num-gpus: "0" + # Head runs the Ray GCS and dashboard only — no training on head. + template: + metadata: + labels: + ray.io/node-type: head + spec: + # Nothing here talks to the Kubernetes API server — don't mount the + # default ServiceAccount token. + automountServiceAccountToken: false + containers: + - name: ray-head + image: ${FULL_IMAGE} + imagePullPolicy: IfNotPresent + ports: + - containerPort: 6379 # Ray GCS + - containerPort: 8265 # Ray Dashboard + - containerPort: 10001 # Ray Client + env: + # On GPU-operator/HyperPod nodes the NVIDIA runtime auto-injects GPU + # tooling; on a GPU-less head node that injection fails and the + # container cannot start. "void" skips it — correct for the head. + - name: NVIDIA_VISIBLE_DEVICES + value: "void" + - name: RAY_memory_monitor_refresh_ms + value: "0" + # Create with: kubectl create secret generic hf-token \ + # --from-literal=HF_TOKEN=$HF_TOKEN -n ${NAMESPACE} + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token + key: HF_TOKEN + resources: + requests: + cpu: "2" + memory: "8Gi" + limits: + cpu: "4" + memory: "16Gi" + volumeMounts: + - name: fsx + mountPath: /fsx + - name: dshm + mountPath: /dev/shm + volumes: + # /fsx defaults to emptyDir so the recipe GATES (which touch NO shared + # storage — see README) run with zero PVC setup and never sit Pending + # on an unbound claim. For the FULL GRPO path, swap this to the FSx PVC: + # - name: fsx + # persistentVolumeClaim: + # claimName: ${FSX_CLAIM} + - name: fsx + emptyDir: {} + - name: dshm + emptyDir: + medium: Memory + sizeLimit: "16Gi" + # Keep the head off GPU nodes so it never blocks a GPU worker from + # scheduling (label from NVIDIA gpu-feature-discovery — portable). + affinity: + nodeAffinity: + requiredDuringSchedulingIgnoredDuringExecution: + nodeSelectorTerms: + - matchExpressions: + - key: nvidia.com/gpu.present + operator: NotIn + values: + - "true" + restartPolicy: Never + terminationGracePeriodSeconds: 120 + workerGroupSpecs: + - groupName: gpu-workers + replicas: ${NUM_NODES} + minReplicas: ${NUM_NODES} + maxReplicas: ${NUM_NODES} + rayStartParams: + num-gpus: "${GPUS_PER_NODE}" + template: + metadata: + labels: + ray.io/node-type: worker + spec: + # These are privileged root containers running model + training code; + # nothing here talks to the API server, so don't also hand them the + # default ServiceAccount token. + automountServiceAccountToken: false + nodeSelector: + node.kubernetes.io/instance-type: ${INSTANCE_TYPE} + tolerations: + - key: nvidia.com/gpu + operator: Exists + effect: NoSchedule + containers: + - name: ray-worker + image: ${FULL_IMAGE} + imagePullPolicy: IfNotPresent + env: + - name: RAY_memory_monitor_refresh_ms + value: "0" + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token + key: HF_TOKEN + # ---- recipe-gate shape (read INSIDE the pod by the kubectl-exec + # gates) ---- + # A client-side `source env_vars` does NOT cross `kubectl exec`; + # only the pod's own env does. run-rollout-probe.sh / train-step.sh + # read NNODES (not NUM_NODES) and GPUS_PER_NODE; probe_rollout.py / + # train_moe_step.py read the EP_* shape. Without these the gates run + # on hardcoded script defaults regardless of what the operator set. + - { name: NNODES, value: "${NUM_NODES}" } + - { name: GPUS_PER_NODE, value: "${GPUS_PER_NODE}" } + - { name: EP_EXPERTS, value: "${EP_EXPERTS}" } + - { name: EP_TOPK, value: "${EP_TOPK}" } + - { name: EP_HIDDEN, value: "${EP_HIDDEN}" } + - { name: EP_TOKENS, value: "${EP_TOKENS}" } + - { name: EP_NUM_SMS, value: "${EP_NUM_SMS}" } # load-bearing on EFA (trap 2) + - { name: EP_NUM_QPS, value: "${EP_NUM_QPS}" } # load-bearing on EFA (trap 2) + # ---- NCCL-GIN proxy + EFA contract, VERBATIM (baked in the image + # ENV and re-exported by every recipe script; repeated here so + # the manifest alone documents the transport contract) ---- + - { name: NCCL_GIN_TYPE, value: "2" } # 2 = CPU-proxy GIN (the EFA-viable path) + - { name: NCCL_GIN_ENABLE, value: "1" } + - { name: OFI_NCCL_GIN_GDAKI, value: "0" } # GPU-initiated: not the shipped path on EFA + - { name: OFI_NCCL_PROTOCOL, value: "RDMA" } + - { name: NCCL_NET_PLUGIN, value: "/opt/aws-ofi-nccl/lib/libnccl-net-ofi.so" } + - { name: NCCL_NVLS_ENABLE, value: "0" } + - { name: NCCL_DEBUG, value: "WARN" } # INFO to see the efa provider banner proof + - { name: FI_PROVIDER, value: "efa" } + - { name: FI_EFA_USE_DEVICE_RDMA, value: "1" } + - { name: FI_EFA_ENABLE_SHM_TRANSFER, value: "0" } + - { name: FI_EFA_FORK_SAFE, value: "1" } + - { name: RDMAV_FORK_SAFE, value: "1" } + # Exclusion form: governs only NCCL's TCP bootstrap/control socket, + # not the EFA data path. Robust if the primary CNI NIC is not eth0. + - { name: NCCL_SOCKET_IFNAME, value: "^lo,docker,veth" } + # ---- DeepEP V2 knobs ---- + - { name: DEEP_EP_USE_V2_SHIM, value: "0" } + # EP_EFA_MAX_QPS / EP_EFA_RDMA_GBS intentionally absent: they were the + # old draft DeepEP#612 patch knobs; the amazon-contributing/DeepEP fork + # this image pins resolves both structurally (get_rdma_gbs() sysfs + # link-rate + the _C.{min,max}_unordered_gin_qps clamp), so deep_ep reads + # neither. The live QP knob is EP_NUM_QPS above. + - { name: TOKENIZERS_PARALLELISM, value: "false" } + # requests == limits (Guaranteed QoS): NCCL_GIN_TYPE=2 is the CPU-PROXY + # path, so proxy-thread CPU sits on the data path — Burstable QoS + CFS + # throttling under node pressure would degrade network throughput and + # silently taint any number this cluster produces. + resources: + requests: + cpu: "90" + memory: "1024Gi" + nvidia.com/gpu: "${GPUS_PER_NODE}" + vpc.amazonaws.com/efa: "${EFA_PER_NODE}" + hugepages-2Mi: "5120Mi" # needs 2Mi hugepages pre-allocated on the node; without them the pod sits Pending + limits: + cpu: "90" + memory: "1024Gi" + nvidia.com/gpu: "${GPUS_PER_NODE}" + vpc.amazonaws.com/efa: "${EFA_PER_NODE}" + hugepages-2Mi: "5120Mi" + securityContext: + privileged: true # gdrdrv device-cgroup (see header) + capabilities: + add: [IPC_LOCK] # RDMA memory pinning + volumeMounts: + - name: fsx + mountPath: /fsx + - name: dshm + mountPath: /dev/shm + - name: dev-infiniband + mountPath: /dev/infiniband + volumes: + # /fsx defaults to emptyDir (see head group) — the recipe gates need + # no shared storage. Swap to the FSx PVC for the full GRPO path: + # - name: fsx + # persistentVolumeClaim: + # claimName: ${FSX_CLAIM} + - name: fsx + emptyDir: {} + - name: dshm + emptyDir: + medium: Memory + sizeLimit: "128Gi" + - name: dev-infiniband + hostPath: + path: /dev/infiniband + restartPolicy: Never + terminationGracePeriodSeconds: 300 + # One worker per physical node — EFA/GPU are node resources, and same-node + # EFA loopback is silently dropped by SRD anyway, so co-scheduling two + # workers on one node is never right. + affinity: + podAntiAffinity: + requiredDuringSchedulingIgnoredDuringExecution: + - labelSelector: + matchExpressions: + - key: ray.io/node-type + operator: In + values: + - worker + topologyKey: kubernetes.io/hostname diff --git a/examples/training/nemo-rl/deepep-v2-efa/nemo-rl.Dockerfile b/examples/training/nemo-rl/deepep-v2-efa/nemo-rl.Dockerfile new file mode 100644 index 000000000..9c2a36014 --- /dev/null +++ b/examples/training/nemo-rl/deepep-v2-efa/nemo-rl.Dockerfile @@ -0,0 +1,278 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +# +# ============================================================================ +# NeMo-RL + Megatron-LM + DeepEP V2 (NCCL-GIN over AWS EFA) for HyperPod EKS +# ============================================================================ +# +# MoE expert-parallel all-to-all for RL post-training, carried by DeepEP V2's +# NCCL backend (ElasticBuffer) over NCCL-GIN's CPU-proxy and aws-ofi-nccl on +# AWS EFA. Everything is built from PUBLIC sources on the NGC PyTorch base — +# no private registry, no fork URL. Every pin is an immutable SHA with a WHY. +# +# Two image flavors from this one file: +# docker build -f nemo-rl.Dockerfile -t /nemo-rl-deepep-efa: . +# -> BASELINE: upstream-only trees. Gates: imports, ElasticBuffer +# bring-up, cross-node EFA transport probe, non-DeepEP train step. +# docker build --build-arg APPLY_DRAFT_ROLLOUT_PATCHES=1 ... +# -> OPT-IN: additionally bakes 2 DRAFT upstream PRs (see the patches/ +# script) that the full NeMo-RL GRPO rollout-over-DeepEP path needs. +# The baseline image has ZERO dependence on them. + +# Base: same NGC base as the sibling RL test case (slime). Bakes torch 2.11 / +# CUDA 13 with TransformerEngine, apex and flash-attn compiled against that +# exact ABI — which is what megatron.core's H100 path needs; rebuilding any of +# those from PyPI against the baked torch is where images usually go wrong. +# Digest-pinned: the tag is the readable name, the @sha256 is the immutable +# contract — NVIDIA can re-push :26.02-py3, and since it is the ABI anchor the +# whole image builds around, a silent re-push is the most consequential drift +# possible here. Override NGC_PYTORCH_BASE to move it (a re-measure event). +ARG NGC_PYTORCH_BASE=nvcr.io/nvidia/pytorch:26.02-py3@sha256:bbc2b67e2533edd63ff1496bb1ed00a00338cdc1478af6c1a0bf9f4b369977e7 +FROM ${NGC_PYTORCH_BASE} +ARG NGC_PYTORCH_BASE # re-declare: pre-FROM ARGs go out of scope after FROM + +LABEL org.opencontainers.image.description="NeMo-RL + Megatron-LM + DeepEP V2 MoE all-to-all over AWS EFA (NCCL-GIN CPU-proxy)" +LABEL org.opencontainers.image.licenses="MIT-0" +LABEL org.opencontainers.image.source="https://github.com/awslabs/awsome-distributed-ai" + +# ---- pins (every one justified; no floating refs) -------------------------- +# NCCL v2.30.4-1 == commit 1933fdd6: the GIN device API generation the measured +# substrate ran (ships include/nccl_device.h — asserted below). Commit-pinned, +# not the bare tag: held to the same "a tag is a moving ref" standard as the +# gdrcopy/DeepEP/Megatron pins below. The NGC base bakes an OLDER NCCL line +# (2.29.x), so this is deliberately a NEWER line and the 00-nccl-gin.conf +# ld.so.conf entry (Layer 4) is LOAD-BEARING — it makes this GIN-capable source +# build win the loader search over the base's baked copy (verify-image.sh +# asserts which one resolves). Built from source so DeepEP has one controlled +# root of headers + libs. +ARG NCCL_VERSION=1933fdd6360a8bfccaa0166bd71bce363d32e5b6 +# EFA 1.48.0: the userspace of the measured substrate (see README pins table). +# Bumping it is a re-measure event, not a routine bump. +ARG EFA_INSTALLER_VERSION=1.48.0 +# gdrcopy v2.5.2 == commit c91ad9f: commit pin, not tag (a bare tag is a +# moving ref upstream can re-point). GIN REQUIRES gdrapi.h at aws-ofi-nccl +# configure time; without it GIN init fails at run time. +ARG GDRCOPY_SHA=c91ad9f178e5fb729fc5b6dc62a77c3bb364d6c9 +# aws-ofi-nccl @9c44d34 + PR#1351 head c2e773d: the GIN CPU-proxy plugin +# lineage this folder standardises on (same pins as the TensorRT-LLM NcclEP +# sibling sample). Immutable SHAs — refs/pull/N/head is a moving ref. +# NOTE on #1351: it is CLOSED-UNMERGED upstream (adds the +# OFI_NCCL_GDRCOPY_FORCED_PCIE_COPY capability override), so unlike the draft +# PRs in patches/ it will NOT "self-neutralize once merged" — this cherry-pick +# is a PERMANENT baseline carry, not a temporary patch. It is applied on EVERY +# build (baseline included): the "baseline has zero dependence on unmerged PRs" +# statement refers to the four draft PRs in the opt-in layer, not this plugin +# pin. Verified applies clean across the ~3-month gap with both fail-loud asserts +# passing (setup_nemo_rl_deepep_efa.sh). Retire by rebasing onto an aws-ofi-nccl +# release that carries the override, when one ships. +ARG AWS_OFI_NCCL_SHA=9c44d34476f90ddbf4a12d0ac4fc412d46bd8ab4 +ARG AWS_OFI_NCCL_PR=1351 +ARG AWS_OFI_NCCL_PR_SHA=c2e773dfb2c75b765b3415f8ffd1b47e7c239a7b +# DeepEP: the amazon-contributing/DeepEP fork (the AWS EPv2/NCCL-GIN tree) — +# ElasticBuffer + NCCL backend, merged upstream in deepseek-ai/DeepEP#605. Same +# fork the house V2 canonical +# micro-benchmarks/expert-parallelism/deepep-v2-benchmark/setup_deepep_gin.sh +# clones, and the twin of the vllm/deepep-v2-gdaki-efa sibling's pin. It carries +# the EFA delta IN-CODE, including both halves of what was draft +# deepseek-ai/DeepEP#612: the get_rdma_gbs() sysfs link-rate fast path +# (deep_ep/utils/envs.py, provider-agnostic — reads +# /sys/class/infiniband//ports/*/rate) and the auto-QP overflow clamp +# (deep_ep/buffers/elastic.py clamps to _C.{min,max}_unordered_gin_qps). So #612 +# is SUPERSEDED: the fork HEAD is strictly ahead of the old stock-01dc3aa + +# #612-patch pairing, and it addresses the review note that pinning the pre-fix +# upstream fork-point forfeits exactly those AWS fixes — the baseline now carries +# them unconditionally, so #612 drops out of the opt-in layer entirely. Pin an +# IMMUTABLE fork SHA (not the moving `main`) for reproducibility; override +# DEEPEP_SHA to bump. setup_nemo_rl_deepep_efa.sh fail-loud asserts both fix-halves +# are present in the clone. +ARG DEEPEP_SHA=97d8f9bcc1be31e9036db2ab591ef9b9f4e38619 +# Megatron-LM: the exact base of draft PR NVIDIA/Megatron-LM#4632 (DeepEP V2 +# ElasticBuffer support in the flex dispatcher) for --check-clean opt-in. +ARG MEGATRON_LM_SHA=19deef67f910c96c213f33b33b30277be8b94d6d +# NeMo-RL @46be4e8: the exact base of the EFA-recipe draft PR NVIDIA-NeMo/RL#2410 +# (its parent commit) — declares requires-python ">=3.12" and torch==2.9.0, both +# of which the NGC base above satisfies. requirements.txt is generated from THIS +# revision (its version pins match line-for-line). Do NOT bump to #2411's base +# cc75cad: that revision bumped requires-python to ">=3.13.13" and torch to +# 2.10.0, so `pip install -e` (Layer 7) hard-fails the interpreter check on this +# py3.12 base (--no-deps does NOT suppress the Requires-Python floor). #2411 is +# metadata-only for this image (deep_ep is built from /opt/DeepEP, not NeMo-RL's +# pin) and is dropped from the opt-in layer; #2410 applies clean on this base. +# Re-pinning to a release tag is a re-measure event. +ARG NEMO_RL_SHA=46be4e8e2b335722c9af75f84e82ad807dad5bf5 +# 9.0 = Hopper (H100/H200, sm_90) — the only measured arch. Override for +# Blackwell with "9.0;10.0" / matching gencode; that is a re-measure event. +ARG TORCH_CUDA_ARCH_LIST="9.0" +ARG NVCC_GENCODE="-gencode=arch=compute_90,code=sm_90" + +ENV DEBIAN_FRONTEND=noninteractive +SHELL ["/bin/bash", "-c"] + +# ---- Layer 1: system deps; drop the distro/HPC-X verbs + MPI stacks -------- +# Same removal the slime sibling does: the EFA installer below provides +# libfabric + Open MPI, and a leftover HPC-X/UCX on the loader path is a +# classic source of wrong-transport surprises. +# The `|| true` is scoped to the removal ONLY (an already-absent verbs package +# is fine); update+install stay fail-hard, or a transient mirror miss would ship +# an image silently missing cmake/autoconf/libtool and fail 100s of lines later +# inside the gdrcopy/aws-ofi-nccl build with no link to the cause. +RUN apt-get update -y && apt-get install -y --no-install-recommends \ + autoconf automake build-essential cmake curl git jq kmod libtool \ + libhwloc-dev pkg-config \ + && { apt-get remove -y --allow-change-held-packages \ + ibverbs-utils libibverbs-dev libibverbs1 libmlx5-1 || true; } +# OPAL_PREFIX: the NGC base sets it to /opt/hpcx/ompi, which the next line +# deletes — reset it to the Open MPI the EFA installer provides (Layer 2 puts +# /opt/amazon/openmpi on PATH). Nothing in the recipe runs mpirun, so this only +# matters to someone invoking MPI from the image, but a dangling OPAL_PREFIX is +# a latent trap worth closing. +RUN rm -rf /opt/hpcx/ompi /usr/local/mpi /usr/local/ucx && ldconfig +ENV OPAL_PREFIX=/opt/amazon/openmpi + +# ---- Layer 2: AWS EFA userspace (public installer) ------------------------- +# --disable-ngc/--disable-build-ngc: the NGC base trips the installer's NGC +# auto-detect, which would reroute to the libnccl-ofi-ngc plugin path; that +# stock plugin does not carry GIN, and two plugins on the loader path is a +# which-one-won guessing game. We build the GIN plugin from source in Layer 5. +RUN apt-get update -y \ + && curl -fsSL https://efa-installer.amazonaws.com/aws-efa-installer-${EFA_INSTALLER_VERSION}.tar.gz | tar -xzf - -C /tmp \ + && cd /tmp/aws-efa-installer \ + && ./efa_installer.sh -y --skip-kmod --skip-limit-conf --no-verify --disable-ngc --disable-build-ngc \ + && echo "${EFA_INSTALLER_VERSION}" > /opt/efa-installer.version \ + && rm -rf /tmp/aws-efa-installer /var/lib/apt/lists/* +ENV PATH=/opt/amazon/efa/bin:/opt/amazon/openmpi/bin:$PATH +ENV LD_LIBRARY_PATH=/opt/amazon/efa/lib:/opt/amazon/openmpi/lib:${LD_LIBRARY_PATH:-} + +# ---- Layer 3: gdrcopy userspace --------------------------------------------- +# /usr/local prefix so aws-ofi-nccl's configure finds gdrapi.h without flags. +# The matching gdrdrv kernel module must exist on the HOST (see the manifest +# header for the privileged/device-plugin trade-off). +RUN git clone https://github.com/NVIDIA/gdrcopy.git /tmp/gdrcopy \ + && cd /tmp/gdrcopy && git fetch origin ${GDRCOPY_SHA} && git checkout ${GDRCOPY_SHA} \ + && make prefix=/usr/local lib lib_install && ldconfig \ + && rm -rf /tmp/gdrcopy + +# ---- Layer 4: NCCL (GIN-capable, single copy wins) -------------------------- +# The nccl_device.h assert is the point: DeepEP V2's NCCL backend compiles +# against the GIN device API, and a device-header-less NCCL fails only at +# DeepEP build time with a confusing include error. The ld.so.conf entry is +# named 00-* so THIS libnccl.so.2 outranks the base image's baked copy at run +# time — recipe/verify-image.sh asserts which path actually resolves. +ENV NCCL_HOME=/opt/nccl/build +RUN git clone https://github.com/NVIDIA/nccl.git /opt/nccl-src \ + && cd /opt/nccl-src && git checkout ${NCCL_VERSION} \ + && make -j"$(nproc)" src.build BUILDDIR=${NCCL_HOME} CUDA_HOME=/usr/local/cuda NVCC_GENCODE="${NVCC_GENCODE}" \ + && test -f ${NCCL_HOME}/include/nccl_device.h \ + || { echo "ERROR: nccl_device.h missing — ${NCCL_VERSION} is not GIN-capable" >&2; exit 1; } \ + && echo "${NCCL_HOME}/lib" > /etc/ld.so.conf.d/00-nccl-gin.conf && ldconfig \ + && pip3 uninstall -y nvidia-nccl-cu13 nvidia-nccl-cu12 nvidia-nccl 2>/dev/null || true \ + && rm -rf /opt/nccl-src/.git +ENV LD_LIBRARY_PATH=${NCCL_HOME}/lib:${LD_LIBRARY_PATH} + +# ---- Layer 5: aws-ofi-nccl GIN plugin (in-tree script, COPY'd not curled) --- +COPY setup_nemo_rl_deepep_efa.sh /opt/setup_nemo_rl_deepep_efa.sh +RUN chmod +x /opt/setup_nemo_rl_deepep_efa.sh \ + && AWS_OFI_NCCL_SHA=${AWS_OFI_NCCL_SHA} AWS_OFI_NCCL_PR=${AWS_OFI_NCCL_PR} \ + AWS_OFI_NCCL_PR_SHA=${AWS_OFI_NCCL_PR_SHA} NCCL_HOME=${NCCL_HOME} \ + /opt/setup_nemo_rl_deepep_efa.sh ofi +ENV LD_LIBRARY_PATH=/opt/aws-ofi-nccl/lib:${LD_LIBRARY_PATH} +ENV NCCL_NET_PLUGIN=/opt/aws-ofi-nccl/lib/libnccl-net-ofi.so + +# ---- Layer 6: clone the three pinned trees at their SHAs -------------------- +# DeepEP = the amazon-contributing fork (see the DEEPEP_SHA header); Megatron-LM +# and NeMo-RL = upstream. Clones only here; the DeepEP BUILD is deferred to +# Layer 9 so the opt-in patch layer (Layer 8) can land its Megatron/NeMo-RL .py +# edits before anything that imports them runs. (Layer 8 no longer touches +# /opt/DeepEP — the fork carries the former DeepEP#612 .cuh/.py edits in-code — +# but the build stays at Layer 9 to preserve the tested layer order.) +RUN git clone https://github.com/amazon-contributing/DeepEP.git /opt/DeepEP \ + && cd /opt/DeepEP && git fetch origin ${DEEPEP_SHA} && git checkout ${DEEPEP_SHA} \ + && git clone https://github.com/NVIDIA/Megatron-LM.git /opt/Megatron-LM \ + && cd /opt/Megatron-LM && git fetch origin ${MEGATRON_LM_SHA} && git checkout ${MEGATRON_LM_SHA} \ + && git clone https://github.com/NVIDIA-NeMo/RL.git /opt/NeMo-RL \ + && cd /opt/NeMo-RL && git fetch origin ${NEMO_RL_SHA} && git checkout ${NEMO_RL_SHA} + +# ---- Layer 7: NeMo-RL + Megatron-LM python installs ------------------------ +# NeMo-RL's pyproject pins torch exactly (e.g. torch==2.9.0), which would +# DOWNGRADE the NGC-baked torch and orphan the baked TE/apex/flash-attn ABI — +# so --no-deps is mandatory here, with the import-relevant dependency set +# pinned in requirements.txt instead. Megatron-LM stays a PYTHONPATH tree +# (matching the slime sibling's runtime layout) so the opt-in patch layer's +# in-place edits are exactly what executes. +COPY requirements.txt /tmp/requirements.txt +RUN python3 -c 'import sys; assert sys.version_info >= (3, 12), f"NeMo-RL needs python>=3.12, base has {sys.version}"' \ + && pip3 install --no-cache-dir -r /tmp/requirements.txt \ + && pip3 install --no-cache-dir --no-deps -e /opt/NeMo-RL +ENV PYTHONPATH=/opt/Megatron-LM:${PYTHONPATH:-} + +# ---- Layer 7b: make the pip NVSHMEM (linked into deep_ep/_C.so) win at runtime +# deep_ep/_C.so is built against the pip nvidia-nvshmem-cu13 wheel (Layer 7, 3.7.x), +# but torch/lib/libtorch_nvshmem.so — imported before deep_ep — has a NEEDED +# libnvshmem_host.so.3 with a RUNPATH ending in /usr/local/cuda/lib64, where the +# NGC base's OLDER dpkg NVSHMEM (3.4.x, no nvshmem_selected_device_transport) lives. +# RUNPATH OUTRANKS ld.so.cache in the loader search order, so an ld.so.conf.d entry +# does NOT win: the 3.4.x lib loads first by soname and `import deep_ep` then dies +# with `undefined symbol: nvshmem_selected_device_transport, version NVSHMEM`. Only +# LD_LIBRARY_PATH outranks RUNPATH — so symlink the wheel's lib dir to a stable path +# (version-agnostic: no python3.X hardcode) and PREPEND it. Verified: recipe/ +# verify-image.sh deep_ep+ElasticBuffer gate fails without this, passes with it. +RUN ln -sfn "$(python3 -c 'import nvidia.nvshmem; print(nvidia.nvshmem.__path__[0])')/lib" /opt/nvshmem-pip-lib \ + && test -f /opt/nvshmem-pip-lib/libnvshmem_host.so.3 \ + || { echo "ERROR: pip nvshmem lib dir not found — requirements.txt must install nvidia-nvshmem-cu13" >&2; exit 1; } +ENV LD_LIBRARY_PATH=/opt/nvshmem-pip-lib:${LD_LIBRARY_PATH} + +# ---- Layer 8 (OPT-IN, default OFF): the 2 draft upstream PRs ---------------- +# NVIDIA-NeMo/RL#2410, NVIDIA/Megatron-LM#4632 +# — the full GRPO rollout-over-DeepEP path depends on them; the BASELINE image +# does not. (deepseek-ai/DeepEP#612 is no longer here: its two fix-halves are +# carried structurally by the amazon-contributing/DeepEP fork the baseline now +# pins, so it needs no patch.) Commits are pinned inside +# patches/apply_nemo_rl_patches.py at +# immutable SHAs and applied fail-loud (git apply --check first): if a hunk no +# longer applies, the BUILD fails rather than shipping an ambiguous image. +# Retire each entry when its PR merges (the script self-neutralizes). +ARG APPLY_DRAFT_ROLLOUT_PATCHES=0 +COPY patches/apply_nemo_rl_patches.py /opt/patches/apply_nemo_rl_patches.py +RUN if [ "${APPLY_DRAFT_ROLLOUT_PATCHES}" = "1" ]; then \ + python3 /opt/patches/apply_nemo_rl_patches.py \ + --deepep-root /opt/DeepEP --megatron-root /opt/Megatron-LM --nemo-rl-root /opt/NeMo-RL \ + --marker /opt/.draft-rollout-patches-applied; \ + else echo "draft-PR layer skipped (APPLY_DRAFT_ROLLOUT_PATCHES=0 — upstream-only baseline)"; fi + +# ---- Layer 9: build DeepEP V2 (NCCL backend) from the (possibly patched) tree +RUN EP_NCCL_ROOT_DIR=${NCCL_HOME} TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST}" \ + /opt/setup_nemo_rl_deepep_efa.sh deepep + +# ---- Layer 10: recipe scripts (LAST — iteration never invalidates heavy layers) +COPY recipe/verify-image.sh /opt/verify-image.sh +COPY recipe/run-rollout-probe.sh /opt/run-rollout-probe.sh +COPY recipe/probe_rollout.py /opt/probe_rollout.py +COPY recipe/train-step.sh /opt/train-step.sh +COPY recipe/train_moe_step.py /opt/train_moe_step.py +RUN chmod 755 /opt/verify-image.sh /opt/run-rollout-probe.sh /opt/train-step.sh + +# ---- runtime transport contract (also repeated in the manifests) ------------ +ENV FI_PROVIDER=efa \ + FI_EFA_USE_DEVICE_RDMA=1 \ + FI_EFA_FORK_SAFE=1 \ + FI_EFA_ENABLE_SHM_TRANSFER=0 \ + RDMAV_FORK_SAFE=1 \ + NCCL_GIN_TYPE=2 \ + NCCL_GIN_ENABLE=1 \ + OFI_NCCL_GIN_GDAKI=0 \ + OFI_NCCL_PROTOCOL=RDMA \ + NCCL_NVLS_ENABLE=0 \ + NCCL_DEBUG=WARN \ + NCCL_SOCKET_IFNAME=^lo,docker,veth \ + RAY_memory_monitor_refresh_ms=0 \ + TOKENIZERS_PARALLELISM=false \ + DEEP_EP_USE_V2_SHIM=0 +# NOTE: EP_EFA_MAX_QPS / EP_EFA_RDMA_GBS are intentionally absent. They were the +# knobs the old deepseek-ai/DeepEP#612 patch read; the amazon-contributing/DeepEP +# fork resolves both structurally (get_rdma_gbs() sysfs link-rate + the C++ +# _C.{min,max}_unordered_gin_qps clamp), so nothing in deep_ep reads them — a dead +# ENV masquerading as config is worse than its absence. The one live QP knob is +# EP_NUM_QPS (the probe's explicit num_allocated_qps), set in the manifests. + +WORKDIR /opt/NeMo-RL +CMD ["/bin/bash", "-lc", "echo 'gates: /opt/verify-image.sh (in verify mode) | /opt/run-rollout-probe.sh {leader|worker} | /opt/train-step.sh {leader|worker} '; sleep infinity"] diff --git a/examples/training/nemo-rl/deepep-v2-efa/patches/apply_nemo_rl_patches.py b/examples/training/nemo-rl/deepep-v2-efa/patches/apply_nemo_rl_patches.py new file mode 100644 index 000000000..7597a7ee5 --- /dev/null +++ b/examples/training/nemo-rl/deepep-v2-efa/patches/apply_nemo_rl_patches.py @@ -0,0 +1,240 @@ +#!/usr/bin/env python3 +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +"""Opt-in draft-PR layer for the NeMo-RL + DeepEP V2 (EFA) test case. + +The BASELINE image installs three upstream trees as-is (amazon-contributing/DeepEP, +NVIDIA/Megatron-LM, NVIDIA-NeMo/RL) at pinned SHAs and depends on nothing +else. The full GRPO rollout-over-DeepEP path additionally needs two upstream +PRs that are still DRAFT/open — this script bakes them in, and ONLY when the +image is built with ``--build-arg APPLY_DRAFT_ROLLOUT_PATCHES=1``. (A third, +NVIDIA-NeMo/RL#2411, is intentionally excluded — see the #2410 entry below. +DeepEP needs NO patch here: the amazon-contributing/DeepEP fork the baseline +pins carries the former draft deepseek-ai/DeepEP#612 fixes in-code.) + +Design goals (why this shape — mirrors the slime sibling's patch layer): + * Default is upstream. Running this script is opt-in; the baseline image + never executes it, so it can never contaminate the upstream-only flavor. + * Pinned, not floating. Each PR is applied as its individual commits at + IMMUTABLE SHAs fetched from the upstream repo's own commit endpoint + (``https://github.com///commit/.patch``). A bare + ``refs/pull/N/head`` is a moving ref and is never used. + * Fail-loud. Every commit is ``git apply --check``ed before applying; if a + hunk no longer applies against the pinned base, the BUILD fails — an image + whose patch state is ambiguous must not ship. + * Self-neutralizing, at CHANGE granularity. Each PR carries one content + probe per independently-required change (not one probe for the whole PR). + Before applying, ALL of a PR's probes are checked: only if EVERY probe + already passes — the PR merged upstream and the tree pin advanced past all + of it — is the PR skipped and reported. A partially-merged tree (some + probes pass, some don't) is NOT skipped; its commits are applied, and a + commit that then fails --check is a fail-loud build error (the pinned tree + drifted), never a silent skip. When every probe reports already-present, + delete the entry here. (A per-commit reverse-check is deliberately absent: + it false-negatives on multi-commit PRs whose later commits touch the same + hunks, and the all-probes decision above makes it redundant.) + * Post-asserted, on the full list. After a PR applies, EVERY one of its + probes is re-checked before the marker is written — so the marker can never + record a patch state the tree does not actually have (a partially-present + PR fails here), and verify-image.sh can trust the marker it reads. + +Each entry links the upstream PR that will make it unnecessary. + +Usage (from nemo-rl.Dockerfile, Layer 8): + python3 apply_nemo_rl_patches.py \ + --deepep-root /opt/DeepEP --megatron-root /opt/Megatron-LM \ + --nemo-rl-root /opt/NeMo-RL --marker /opt/.draft-rollout-patches-applied +""" +from __future__ import annotations + +import argparse +import shutil +import subprocess +import sys +import urllib.request +from pathlib import Path + +# One entry per draft PR. `commits` are the PR's commits in order, at the +# immutable SHAs they had when this test case was authored (2026-08-25). +# `probes` = a list of (relative path, needle) — ONE per independently-required +# change in the PR (needle None = the file's existence is the assertion). The +# list is deliberately per-CHANGE, not per-PR: a single per-PR probe would let a +# partially-merged upstream tree satisfy the probe while lacking later commits, +# and the marker would then record a patch state the tree does not have. A PR is +# skipped only when EVERY probe passes; after applying, EVERY probe is +# re-asserted before the marker is written. +# +# Pure-cleanup commits (no detectable content addition — e.g. a comment/URL +# removal) intentionally have NO probe: commits are applied atomically in order +# and are never skipped individually, and a PR-level skip fires only when the +# whole PR is already upstream (i.e. merged past the cleanup too), so a cleanup +# commit can be verified by neither presence nor absence without ambiguity. +PATCH_SETS = [ + { + "name": "NVIDIA/Megatron-LM#4632 — moe: DeepEP V2 ElasticBuffer support in the flex dispatcher", + "url": "https://github.com/NVIDIA/Megatron-LM/pull/4632", + "repo": "NVIDIA/Megatron-LM", + "root_arg": "megatron_root", + "commits": [ + "e132d5dd15358940aeb962105e44b402919084c5", # ElasticBuffer support in _DeepepManager + "f5ac3d481a8baca2596f20ae213dee25f87e35bf", # graceful EventOverlap import fallback under V2 + "99b8824ee9c8d26b115e05b3ac563d1ea73b2b6b", # pass num_experts explicitly to V2 backward dispatch + "8056d6d489c73a353d590b8079497fceda4f9aa7", # drop downstream repro URLs + dead conditional (cleanup — no probe) + ], + "probes": [ + ("megatron/core/transformer/moe/fused_a2a.py", "ElasticBuffer"), # commit 1 + ("megatron/core/transformer/moe/fused_a2a.py", "deep_ep.utils.event"), # commit 2 + ("megatron/core/transformer/moe/fused_a2a.py", "_handle_num_experts"), # commit 3 + ], # commit 4 = cleanup, see header + }, + { + "name": "NVIDIA-NeMo/RL#2410 — deps: re-export LD_LIBRARY_PATH for AWS EFA OFI discovery", + "url": "https://github.com/NVIDIA-NeMo/RL/pull/2410", + "repo": "NVIDIA-NeMo/RL", + "root_arg": "nemo_rl_root", + # This PR's commit is cut directly on NEMO_RL_SHA=46be4e8 (its parent), + # so it applies --check-clean on the baseline tree. (#2411, the deep_ep + # pin bump, is intentionally NOT applied here: its base is cc75cad — + # 116 commits ahead of 46be4e8, across the requires-python 3.12->3.13.13 + # bump — so it neither applies here NOR belongs on this py3.12 substrate, + # and it is metadata-only for this image since deep_ep is built from + # /opt/DeepEP, not from NeMo-RL's pin.) + "commits": [ + "7f0f21a7a8d7205d2d741f2bc9cff837462091a5", + ], + # The PR also ships the worked 2-node EFA GRPO recipe config this test + # case's README points at for the full rollout path. + "probes": [ + ("examples/configs/recipes/llm/aws-efa-grpo-qwen3-30ba3b-2n8g-megatron.yaml", None), + ], + }, +] + + +def run(args: list[str], cwd: Path) -> subprocess.CompletedProcess: + return subprocess.run(args, cwd=cwd, capture_output=True, text=True) + + +def fetch_patch(repo: str, sha: str) -> bytes: + url = f"https://github.com/{repo}/commit/{sha}.patch" + with urllib.request.urlopen(url, timeout=60) as resp: + data = resp.read() + if not data.lstrip().startswith(b"From "): + raise RuntimeError(f"{url} did not return a git patch") + return data + + +def apply_commit(root: Path, repo: str, sha: str) -> str: + """Apply one pinned commit into the tree at `root`. Returns a status word. + + No per-commit reverse-check: the whole-PR probe gate in main() decides + skip-vs-apply (a PR is skipped only when EVERY probe already passes). Here we + apply unconditionally and fail loud if the pinned tree drifted under a hunk — + a per-commit reverse-check would false-negative on multi-commit PRs whose + later commits touch the same hunks (the docstring's stated design). + """ + # absolute: git apply runs with cwd=root, so a relative path would resolve + # inside the tree twice + patch_path = (root / f".{sha}.patch").resolve() + patch_path.write_bytes(fetch_patch(repo, sha)) + try: + check = run(["git", "apply", "--check", str(patch_path)], root) + if check.returncode != 0: + raise RuntimeError( + f"{repo} commit {sha} does not apply cleanly:\n{check.stderr}\n" + "The pinned tree moved under the patch — re-pin the tree SHA to the " + "PR's current base, or retire this entry if the PR merged with changes." + ) + applied = run(["git", "apply", str(patch_path)], root) + if applied.returncode != 0: + raise RuntimeError(f"{repo} commit {sha} --check passed but apply failed:\n{applied.stderr}") + return "applied" + finally: + patch_path.unlink(missing_ok=True) + + +def probe_ok(root: Path, probe: tuple[str, str | None]) -> bool: + rel, needle = probe + target = root / rel + if not target.is_file(): + return False + return needle is None or needle in target.read_text(errors="replace") + + +def missing_probes(root: Path, probes: list[tuple[str, str | None]]) -> list[tuple[str, str | None]]: + """The subset of `probes` NOT yet satisfied in the tree at `root`.""" + return [p for p in probes if not probe_ok(root, p)] + + +def describe_probe(probe: tuple[str, str | None]) -> str: + rel, needle = probe + return f"{rel}" if needle is None else f"{needle!r} in {rel}" + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--deepep-root", required=True, type=Path) + parser.add_argument("--megatron-root", required=True, type=Path) + parser.add_argument("--nemo-rl-root", required=True, type=Path) + parser.add_argument("--marker", required=True, type=Path, + help="marker file recording the applied commit SHAs") + args = parser.parse_args() + roots = { + "deepep_root": args.deepep_root, + "megatron_root": args.megatron_root, + "nemo_rl_root": args.nemo_rl_root, + } + + record: list[str] = [] + for pset in PATCH_SETS: + root = roots[pset["root_arg"]] + if not (root / ".git").exists(): + print(f"FATAL: {root} is not a git checkout — cannot apply {pset['name']}", file=sys.stderr) + return 1 + print(f"== {pset['name']} ==") + probes = pset["probes"] + missing = missing_probes(root, probes) + if not missing: + # Every per-change probe already passes — the tree pin advanced past + # ALL of this PR's changes (fully merged upstream). Only then skip; + # a partially-merged tree (some probes missing) is NOT skipped, so it + # can never be recorded as already-present. Delete the entry here. + print(f" all {len(probes)} change-probes already present — skipping; retire this entry") + record.extend(f"{pset['repo']}@{sha} already-present" for sha in pset["commits"]) + continue + if len(missing) != len(probes): + # Partially-present: some changes are in the tree, some are not. Apply + # the PR's commits (git apply --check will fail loud on any hunk that + # the already-present change collides with) rather than silently + # trusting the single-needle hit the old per-PR probe would have. + print(f" partially present ({len(probes) - len(missing)}/{len(probes)} " + f"change-probes hit) — applying to complete: " + f"{', '.join(describe_probe(p) for p in missing)}") + for sha in pset["commits"]: + status = apply_commit(root, pset["repo"], sha) + print(f" {sha[:12]} {status}") + record.append(f"{pset['repo']}@{sha} {status}") + # Post-assert the FULL probe list — a partially-present PR that failed to + # complete, or an apply-succeeded-but-wrong-tree mistake, fails here so + # the marker never records a patch state the tree does not have. + still_missing = missing_probes(root, probes) + if still_missing: + print(f"FATAL: post-assert failed for {pset['name']}: " + f"{'; '.join(describe_probe(p) for p in still_missing)} " + f"absent under {root}", file=sys.stderr) + return 1 + print(f" post-assert OK — all {len(probes)} change-probes present ({pset['url']})") + + # Stale bytecode from the pre-patch install must not shadow the patched + # sources (NeMo-RL is installed -e; Megatron rides PYTHONPATH). + for root in set(roots.values()): + for pycache in root.rglob("__pycache__"): + shutil.rmtree(pycache, ignore_errors=True) + + args.marker.write_text("\n".join(record) + "\n") + print(f"draft-PR layer complete — marker at {args.marker}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/examples/training/nemo-rl/deepep-v2-efa/recipe/probe_rollout.py b/examples/training/nemo-rl/deepep-v2-efa/recipe/probe_rollout.py new file mode 100644 index 000000000..48c9f7900 --- /dev/null +++ b/examples/training/nemo-rl/deepep-v2-efa/recipe/probe_rollout.py @@ -0,0 +1,174 @@ +#!/usr/bin/env python3 +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +"""probe_rollout.py — does DeepEP V2's real dispatch/combine move a +rollout-shaped tensor across the EFA node boundary, and does the combined +result match an oracle? + +This is the transport gate the NeMo-RL rollout path rides: it constructs the +REAL `deep_ep.ElasticBuffer` (the V2 NCCL-GIN buffer, the class the Megatron +flex dispatcher binds to) on a multi-node NCCL group, dispatches a +rollout-shaped activation tensor with random expert routing, runs a local +"expert compute" stand-in with a closed-form oracle, combines, and checks the +combined output per token — no model weights, no checkpoint download, minutes +not hours. Runs under torchrun (see run-rollout-probe.sh) so dispatch/combine +is exercised ACROSS the node boundary, not just instantiated: the defect class +this exists to catch (e.g. the GIN request-ring overflow that the former draft +deepseek-ai/DeepEP#612 addressed) only fires at the first real cross-node +dispatch. + +Note on explicit SM/QP counts: the amazon-contributing/DeepEP fork this image +builds carries the former #612 EFA fixes in-code (the `get_rdma_gbs()` sysfs +link-rate fast path and the auto-QP overflow clamp), so its auto-sizers ARE +EFA-aware. The probe still passes `num_allocated_qps`/`num_sms`/`num_qps` +EXPLICITLY (EP_NUM_QPS=2 is the value the p5en evidence validated) — it makes +the probe deterministic and independent of the auto-sizer, and it survives the +fork's clamp unchanged (max(2, min(2, max_unordered_gin_qps)) == 2). +""" +import datetime +import os +import sys +import traceback + +import torch +import torch.distributed as dist + +RANK = int(os.environ["RANK"]) +WORLD = int(os.environ["WORLD_SIZE"]) +LOCAL = int(os.environ["LOCAL_RANK"]) + +E = int(os.environ.get("EP_EXPERTS", "128")) # routed experts (Qwen3-30B-A3B shape) +TOK = int(os.environ.get("EP_TOKENS", "128")) # tokens per rank +H = int(os.environ.get("EP_HIDDEN", "2048")) # hidden dim +K = int(os.environ.get("EP_TOPK", "8")) # experts per token +NUM_SMS = int(os.environ.get("EP_NUM_SMS", "8")) +NUM_QPS = int(os.environ.get("EP_NUM_QPS", "2")) + + +def log(msg): + print(f"[rank{RANK}] {msg}", flush=True) + + +def main(): + torch.cuda.set_device(LOCAL) + # Explicit timeout so a wedged peer fails fast instead of sitting in the + # default 30-min NCCL timeout (the documented flow runs the worker detached + # under kubectl exec, so a worker-side stall surfaces as a hung leader). + dist.init_process_group( + "nccl", rank=RANK, world_size=WORLD, + timeout=datetime.timedelta(minutes=10), + ) + dev = torch.device("cuda", LOCAL) + assert E % WORLD == 0, f"experts ({E}) must divide by EP size ({WORLD})" + + import deep_ep + from deep_ep import ElasticBuffer + + log(f"probe start world={WORLD} experts={E} tokens={TOK} hidden={H} topk={K} " + f"sms={NUM_SMS} qps={NUM_QPS}") + + # Buffer bring-up alone is a gate: it creates the NCCL comm handle and the + # GIN-backed symmetric buffer across all ranks — if the GIN plugin, the + # gdrcopy handle, or the EFA provider is broken, it dies HERE with a + # transport error, before any dispatch. + buf = ElasticBuffer( + group=dist.group.WORLD, + num_max_tokens_per_rank=TOK, + hidden=H, + num_topk=K, + num_allocated_qps=NUM_QPS, # explicit — see module docstring + explicitly_destroy=True, + ) + log(f"PROBE-BUFFER ElasticBuffer up: bytes={buf.num_bytes} " + f"rdma_ranks={buf.num_rdma_ranks} nvlink_ranks={buf.num_nvlink_ranks}") + + # Rollout-shaped payload with a closed-form oracle: expert e scales its + # tokens by (e+1)/E, weights folded in during local compute, so the + # combined output must equal x * sum_k(((idx_k+1)/E) * w_k) per token — + # checkable without real weights. + # deep_ep._C exports topk_idx_t unconditionally at the pinned SHA — access + # it directly so a future rename fails loudly at import with the real name, + # rather than silently substituting a possibly-wrong index width. + idx_dtype = deep_ep.topk_idx_t + g = torch.Generator(device="cpu").manual_seed(4242 + RANK) + x = (torch.randn(TOK, H, generator=g, dtype=torch.float32) / 8.0).to(dev, torch.bfloat16) + topk_idx = torch.stack([torch.randperm(E, generator=g)[:K] for _ in range(TOK)]).to(dev, idx_dtype) + # keep weights away from 0 so no oracle row degenerates to all-zeros + topk_weights = (torch.rand(TOK, K, generator=g, dtype=torch.float32) * 0.9 + 0.1).to(dev) + + recv_x, recv_topk_idx, recv_topk_weights, handle, _ = buf.dispatch( + x, + topk_idx=topk_idx, + topk_weights=topk_weights, + num_experts=E, + num_sms=NUM_SMS, + num_qps=NUM_QPS, + ) + torch.cuda.synchronize() + log(f"PROBE-DISPATCH recv_x={tuple(recv_x.shape)}") + + # Local expert compute against the oracle. Received routing ids are global + # expert indices (invalid/non-local slots < 0); if a future DeepEP maps + # them to local ids instead, detect and offset rather than misattribute. + num_local = E // WORLD + lo, hi = RANK * num_local, RANK * num_local + num_local + sl = recv_topk_idx.to(torch.int64) + # Rank 0's local experts ARE 0..num_local-1, so global- and local-mapped ids + # are indistinguishable there (sl.max() < num_local always holds) — restrict + # the detector to ranks >= 1 so it can't misfire on a normal run. + if RANK > 0 and sl.numel() > 0 and int(sl.max()) < num_local and E > num_local: + log("PROBE-NOTE recv_topk_idx looks local-mapped; offsetting by rank base") + sl = torch.where(sl >= 0, sl + lo, sl) + moe_out = torch.zeros_like(recv_x, dtype=torch.bfloat16) + for kk in range(sl.shape[1]): + e = sl[:, kk] + m = (e >= lo) & (e < hi) + if not bool(m.any()): + continue + f = ((e[m].to(torch.float32) + 1.0) / E).unsqueeze(1) + w = recv_topk_weights[m, kk].unsqueeze(1).to(torch.float32) + moe_out[m] += (recv_x[m].to(torch.float32) * f * w).to(torch.bfloat16) + + # weights already folded into moe_out — combine just reduces + combined_x, _, _ = buf.combine(moe_out, handle, num_qps=NUM_QPS) + torch.cuda.synchronize() + + fac = ((topk_idx.to(torch.float32) + 1.0) / E) * topk_weights + expect = x.to(torch.float32) * fac.sum(1, keepdim=True) + got = combined_x.to(torch.float32) + # Score each token against ITS OWN scale, not one global maximum: a global + # normaliser averages away exactly the failure this probe exists to catch — + # a handful of badly-wrong low-amplitude rows from a partial dispatch or a + # dropped QP channel. bf16-appropriate absolute floor on the row scale. + row_err = (got - expect).abs().amax(dim=1) + row_scale = expect.abs().amax(dim=1).clamp_min(1e-3) + relmax = (row_err / row_scale).max().item() + nz = int((got.abs().sum(1) > 0).sum()) + ok = (nz == TOK) and (relmax < 0.05) + log(f"PROBE-COMBINE nonzero={nz}/{TOK} relmax={relmax:.4g} {'OK' if ok else 'MISMATCH'}") + + # Every rank must agree — one silently-corrupted rank is a failed run. + flag = torch.tensor([1 if ok else 0], device=dev, dtype=torch.int32) + dist.all_reduce(flag, op=dist.ReduceOp.MIN) + verdict = "PROBE-PASS" if int(flag.item()) == 1 else "PROBE-MISMATCH" + log(f"{verdict} ElasticBuffer dispatch/combine world={WORLD}") + + buf.destroy() + dist.barrier() + log("PROBE-DONE") + return 0 if int(flag.item()) == 1 else 4 + + +if __name__ == "__main__": + rc = 1 + try: + rc = main() + except Exception: + traceback.print_exc() + print(f"[rank{RANK}] PROBE-EXC", flush=True) + finally: + # Tear the group down on every path — a rank that raised must not leave + # its peers blocked in the MIN all_reduce / barrier until the timeout. + if dist.is_initialized(): + dist.destroy_process_group() + sys.exit(rc) diff --git a/examples/training/nemo-rl/deepep-v2-efa/recipe/run-rollout-probe.sh b/examples/training/nemo-rl/deepep-v2-efa/recipe/run-rollout-probe.sh new file mode 100755 index 000000000..5f643fbe7 --- /dev/null +++ b/examples/training/nemo-rl/deepep-v2-efa/recipe/run-rollout-probe.sh @@ -0,0 +1,92 @@ +#!/bin/bash +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. SPDX-License-Identifier: MIT-0 +# run-rollout-probe.sh — prove DeepEP V2's dispatch/combine actually moves bytes over EFA +# BEFORE any model load or GRPO launch. Drives the REAL deep_ep.ElasticBuffer cross-node +# with the SAME NCCL-GIN/EFA env the training uses (probe_rollout.py), asserts the EFA +# hardware TX counters advanced, and exits non-zero on any failure. +# +# leader: run-rollout-probe.sh leader +# worker: run-rollout-probe.sh worker # node-rank 1,2,... for worker nodes +# +# Env (shared with train-step.sh): NNODES (default 2), GPUS_PER_NODE (default 8), +# EP_EXPERTS/EP_TOKENS/EP_HIDDEN/EP_TOPK/EP_NUM_SMS/EP_NUM_QPS (see env_vars.example). +set -uo pipefail + +ROLE="${1:?usage: run-rollout-probe.sh {leader|worker} [node-rank]}" +case "$ROLE" in leader|worker) ;; *) echo "FATAL: unrecognized role '$ROLE' (leader|worker)"; exit 2 ;; esac +LEADER_IP="${2:?need leader ip}" +if [ "$ROLE" = "worker" ]; then NODE_RANK_ARG="${3:?worker requires an explicit node-rank (1,2,...)}"; else NODE_RANK_ARG=0; fi +NNODES="${NNODES:-2}" +GPUS_PER_NODE="${GPUS_PER_NODE:-8}" +PROBE="/opt/probe_rollout.py" +[ -f "$PROBE" ] || { echo "FAIL: $PROBE not in image — rebuild from nemo-rl.Dockerfile"; exit 3; } + +# ---- NCCL-GIN proxy + EFA contract, VERBATIM (baked in the image ENV; re-exported so an +# ad-hoc shell that scrubbed its env still probes the transport under test) ---- +export NCCL_GIN_TYPE=2 NCCL_GIN_ENABLE=1 OFI_NCCL_GIN_GDAKI=0 +export NCCL_NVLS_ENABLE=0 +export FI_PROVIDER=efa FI_EFA_USE_DEVICE_RDMA=1 FI_EFA_ENABLE_SHM_TRANSFER=0 FI_EFA_FORK_SAFE=1 +export OFI_NCCL_PROTOCOL=RDMA +export NCCL_NET_PLUGIN=/opt/aws-ofi-nccl/lib/libnccl-net-ofi.so +# ---- probe shape (defaults = the Qwen3-30B-A3B / Wave-28 shape twin) ---- +export EP_EXPERTS="${EP_EXPERTS:-128}" EP_TOKENS="${EP_TOKENS:-128}" +export EP_HIDDEN="${EP_HIDDEN:-2048}" EP_TOPK="${EP_TOPK:-8}" +# Explicit SM/QP counts: the amazon-contributing/DeepEP fork carries the former #612 EFA +# fixes in-code, so its auto-sizers are EFA-aware — but the probe pins these anyway so it +# is deterministic and auto-sizer-independent. 2 QPs is the value the p5en evidence +# validated, and it survives the fork's clamp unchanged (max(2,min(2,max))==2). +export EP_NUM_SMS="${EP_NUM_SMS:-8}" EP_NUM_QPS="${EP_NUM_QPS:-2}" +export NCCL_DEBUG="${PROBE_NCCL_DEBUG:-INFO}" # INFO so the efa provider banner prints = transport proof + +NODE_RANK="$NODE_RANK_ARG"; [ "$ROLE" = "leader" ] && NODE_RANK=0 +echo "===== DeepEP-V2 rollout probe: role=$ROLE node_rank=$NODE_RANK nnodes=$NNODES gpus/node=$GPUS_PER_NODE leader=$LEADER_IP $(hostname) $(date -u +%FT%TZ) =====" + +# EFA hardware TX bytes across all NICs — the counter delta is the honest "bytes actually +# left this node over EFA" assert (a PASS on SHM/TCP fallback would show delta≈0). +# Sum the byte-counter FAMILIES, not tx_bytes alone: with FI_EFA_USE_DEVICE_RDMA=1 the +# bytes can be accounted as rdma_write_bytes rather than sends, so tx_bytes alone could +# read ~0 on a perfectly healthy RDMA run and false-FAIL the >=1 MiB check below. Missing +# families are skipped (the glob simply finds fewer files), so this is portable across +# kernels that expose different subsets. +efa_tx_total() { + local total=0 v + for f in /sys/class/infiniband/*/ports/1/hw_counters/{tx_bytes,send_bytes,rdma_write_bytes,rdma_read_resp_bytes}; do + [ -f "$f" ] && v=$(cat "$f") && total=$((total + v)) + done + echo "$total" +} +TX_BEFORE=$(efa_tx_total) + +set -o pipefail +# GPUS_PER_NODE torchrun procs per node (one per GPU): NNODES x GPUS_PER_NODE = the EP16 +# shape the Wave-28 measured run used on 2 nodes. +# Wall-clock bound (timeout 900) so a wedged rendezvous returns a verdict instead of +# holding the nodes; the python side also passes an init_process_group timeout. timeout +# exit 124 => treated as a probe failure by the rc check below. +timeout 900 torchrun --nnodes="$NNODES" --nproc-per-node="$GPUS_PER_NODE" --node-rank="$NODE_RANK" \ + --master-addr="$LEADER_IP" --master-port=29501 "$PROBE" 2>&1 | tee /tmp/rollout-probe.$NODE_RANK.log +rc=${PIPESTATUS[0]} + +TX_AFTER=$(efa_tx_total) +TX_DELTA=$((TX_AFTER - TX_BEFORE)) +echo "EFA hw_counters tx_bytes delta on this node: $TX_DELTA" + +# fail-loud contract: gate on the exit code (the probe all-reduces a MIN verdict — one bad +# rank fails every rank), then separately require EFA evidence. +if [ "$rc" -ne 0 ]; then + echo "ROLLOUT-PROBE FAIL (node_rank=$NODE_RANK, torchrun rc=$rc) — see /tmp/rollout-probe.$NODE_RANK.log" + exit 1 +fi +if [ "$NNODES" -gt 1 ] && [ "$TX_DELTA" -lt 1048576 ]; then + # a real cross-node dispatch moves MBs; <1 MiB means the bytes went over SHM/TCP, not EFA + echo "ROLLOUT-PROBE FAIL: probe passed but EFA TX advanced only ${TX_DELTA}B — transport was NOT EFA" + exit 1 +fi +# only provider-specific banners count; a bare "NET/OFI" line also prints for the +# tcp;ofi_rxm fallback, the exact case to rule out +if grep -qiE "efa-direct|Selected Provider is efa" /tmp/rollout-probe.$NODE_RANK.log; then + echo "ROLLOUT-PROBE PASS (node_rank=$NODE_RANK) — ElasticBuffer dispatch/combine over EFA verified (tx +${TX_DELTA}B)" + exit 0 +fi +echo "ROLLOUT-PROBE INCONCLUSIVE: probe passed and counters moved (+${TX_DELTA}B) but no EFA-provider banner in log — confirm transport before trusting" +exit 2 diff --git a/examples/training/nemo-rl/deepep-v2-efa/recipe/train-step.sh b/examples/training/nemo-rl/deepep-v2-efa/recipe/train-step.sh new file mode 100755 index 000000000..57e86e268 --- /dev/null +++ b/examples/training/nemo-rl/deepep-v2-efa/recipe/train-step.sh @@ -0,0 +1,96 @@ +#!/bin/bash +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. SPDX-License-Identifier: MIT-0 +# train-step.sh — the N-step Megatron-core MoE training gate over EFA (the training +# workload's analogue of a served smoke test). Runs recipe/train_moe_step.py under +# torchrun across the nodes with the SAME NCCL-GIN/EFA env the full GRPO run uses, +# and exits non-zero on any failure. +# +# leader: train-step.sh leader +# worker: train-step.sh worker # node-rank 1,2,... for worker nodes +# +# Env: NNODES (default 2), GPUS_PER_NODE (default 8), MOE_DISPATCHER (default alltoall), +# TRAIN_STEPS (default 3), EP_EXPERTS/EP_TOPK/EP_HIDDEN (see env_vars.example). +set -uo pipefail + +ROLE="${1:?usage: train-step.sh leader|worker [node-rank]}" +case "$ROLE" in leader|worker) ;; *) echo "FATAL: unrecognized role '$ROLE' (leader|worker)"; exit 2 ;; esac +LEADER_IP="${2:?need leader ip}" +if [ "$ROLE" = "worker" ]; then NODE_RANK_ARG="${3:?worker requires an explicit node-rank (1,2,...)}"; else NODE_RANK_ARG=0; fi +NNODES="${NNODES:-2}" +GPUS_PER_NODE="${GPUS_PER_NODE:-8}" +DRIVER="/opt/train_moe_step.py" +[ -f "$DRIVER" ] || { echo "FAIL: $DRIVER not in image — rebuild from nemo-rl.Dockerfile"; exit 3; } + +# ---- dispatcher selection guard ---- +# flex (+moe_enable_deepep) is the DeepEP V2 ElasticBuffer path and needs the opt-in +# draft-PR image (Megatron-LM#4632 et al). On an unpatched image, refuse up front with +# the reason, rather than dying later inside megatron with an import error — and never +# silently fall back to a different dispatcher than the one requested. +export MOE_DISPATCHER="${MOE_DISPATCHER:-alltoall}" +if [ "$MOE_DISPATCHER" = "flex" ] && [ ! -f /opt/.draft-rollout-patches-applied ]; then + echo "FATAL: MOE_DISPATCHER=flex (DeepEP V2 ElasticBuffer) needs an image built with" + echo "APPLY_DRAFT_ROLLOUT_PATCHES=1 (bakes Megatron-LM#4632 + NeMo-RL#2410)." + echo "This baseline image is upstream-only — run the default alltoall dispatcher gate," + echo "or rebuild with the opt-in layer." + exit 4 +fi + +# ---- NCCL-GIN proxy + EFA contract, VERBATIM from run-rollout-probe.sh ---- +export NCCL_GIN_TYPE=2 NCCL_GIN_ENABLE=1 OFI_NCCL_GIN_GDAKI=0 +export NCCL_NVLS_ENABLE=0 +export FI_PROVIDER=efa FI_EFA_USE_DEVICE_RDMA=1 FI_EFA_ENABLE_SHM_TRANSFER=0 FI_EFA_FORK_SAFE=1 +export OFI_NCCL_PROTOCOL=RDMA +export NCCL_NET_PLUGIN=/opt/aws-ofi-nccl/lib/libnccl-net-ofi.so +export NCCL_DEBUG="${TRAIN_NCCL_DEBUG:-INFO}" # INFO so the efa provider banner prints = transport proof + +NODE_RANK="$NODE_RANK_ARG"; [ "$ROLE" = "leader" ] && NODE_RANK=0 +echo "===== Megatron MoE train-step: role=$ROLE node_rank=$NODE_RANK nnodes=$NNODES gpus/node=$GPUS_PER_NODE dispatcher=$MOE_DISPATCHER leader=$LEADER_IP $(hostname) $(date -u +%FT%TZ) =====" + +# EFA hardware TX bytes across all NICs — same honest transport assert the rollout probe +# uses: on a multi-node run the MoE dispatcher's cross-node all-to-all (NCCL all-to-all for +# `alltoall`, DeepEP ElasticBuffer for `flex`) MUST move bytes over EFA, so a PASS with the +# counter flat means the traffic fell back to SHM/TCP. Sum the byte-counter FAMILIES, not +# tx_bytes alone: with FI_EFA_USE_DEVICE_RDMA=1 the bytes can be accounted as +# rdma_write_bytes rather than sends, so tx_bytes alone could read ~0 on a healthy RDMA run +# and false-FAIL. Missing families are skipped (the glob finds fewer files). +efa_tx_total() { + local total=0 v + for f in /sys/class/infiniband/*/ports/1/hw_counters/{tx_bytes,send_bytes,rdma_write_bytes,rdma_read_resp_bytes}; do + [ -f "$f" ] && v=$(cat "$f") && total=$((total + v)) + done + echo "$total" +} +TX_BEFORE=$(efa_tx_total) + +set -o pipefail +# Wall-clock bound (timeout 900) so a wedged rendezvous returns a verdict instead of holding +# the nodes; train_moe_step.py also passes an init_process_group timeout. timeout exit 124 => +# treated as a gate failure by the rc check below. +timeout 900 torchrun --nnodes="$NNODES" --nproc-per-node="$GPUS_PER_NODE" --node-rank="$NODE_RANK" \ + --master-addr="$LEADER_IP" --master-port=29502 "$DRIVER" 2>&1 | tee /tmp/train-step.$NODE_RANK.log +rc=${PIPESTATUS[0]} + +TX_AFTER=$(efa_tx_total) +TX_DELTA=$((TX_AFTER - TX_BEFORE)) +echo "EFA hw_counters tx_bytes delta on this node: $TX_DELTA" + +# fail-loud contract, same shape as run-rollout-probe.sh: gate on the exit code (the driver +# all-reduces a MIN verdict — one bad rank fails every rank), then separately require EFA +# evidence, then confirm the provider banner. +if [ "$rc" -ne 0 ]; then + echo "TRAIN-STEP GATE FAIL (node_rank=$NODE_RANK, torchrun rc=$rc) — see /tmp/train-step.$NODE_RANK.log" + exit 1 +fi +if [ "$NNODES" -gt 1 ] && [ "$TX_DELTA" -lt 1048576 ]; then + # a real cross-node MoE all-to-all moves MBs; <1 MiB means the bytes went over SHM/TCP, not EFA + echo "TRAIN-STEP GATE FAIL: step passed but EFA TX advanced only ${TX_DELTA}B — transport was NOT EFA" + exit 1 +fi +# only provider-specific banners count; a bare "NET/OFI" line also prints for the +# tcp;ofi_rxm fallback, the exact case to rule out +if grep -qiE "efa-direct|Selected Provider is efa" /tmp/train-step.$NODE_RANK.log; then + echo "TRAIN-STEP GATE PASS (node_rank=$NODE_RANK, dispatcher=$MOE_DISPATCHER) — MoE step over EFA verified (tx +${TX_DELTA}B)" + exit 0 +fi +echo "TRAIN-STEP GATE INCONCLUSIVE: step passed and counters moved (+${TX_DELTA}B) but no EFA-provider banner in log — confirm transport before trusting" +exit 2 diff --git a/examples/training/nemo-rl/deepep-v2-efa/recipe/train_moe_step.py b/examples/training/nemo-rl/deepep-v2-efa/recipe/train_moe_step.py new file mode 100644 index 000000000..b8e5afd62 --- /dev/null +++ b/examples/training/nemo-rl/deepep-v2-efa/recipe/train_moe_step.py @@ -0,0 +1,248 @@ +#!/usr/bin/env python3 +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +"""train_moe_step.py — N-step Megatron-core MoE training gate over EFA. + +The training analogue of the rollout probe: builds a small expert-parallel GPT +MoE with megatron.core (random weights, synthetic batch — no checkpoint, no +dataset), runs N optimizer steps across all ranks, and gates on: + + 1. loss finite at every step and lower at the last step than the first + (the Wave-28 "Shape Y" oracle shape: loss 26.41 -> 24.62 over 3 steps on + the measured substrate — exact values are seed/substrate-specific and NOT + asserted, only finiteness + decrease are); + 2. cross-rank loss agreement — every rank feeds the SAME seeded batch, so + after the MoE all-to-all round-trip all ranks must compute the SAME loss; + a rank whose dispatched tokens were corrupted in flight disagrees here; + 3. a MIN all-reduced verdict — one bad rank fails every rank. + +Dispatcher selection (MOE_DISPATCHER): + alltoall (default) Megatron's stock all-to-all dispatcher — NCCL all-to-all + over EFA, no deep_ep involvement. This is the BASELINE gate: it + runs on the upstream-only image. + flex Megatron's flex dispatcher with moe_enable_deepep=True — the + DeepEP V2 ElasticBuffer path. Needs the opt-in draft-PR image + (Megatron-LM#4632); train-step.sh refuses it on an unpatched + image rather than failing with a distant import error. + +The loop deliberately runs without megatron's DDP wrapper (to stay independent +of its DDP config API) but reproduces the two things DDP does that this gate +needs, restricted to the REPLICATED (non-expert) params: + - at construction, broadcast them from rank 0 — megatron inits the router + weight and position embeddings under a per-rank-varying RNG context + (measured ~1.5e-1 cross-rank divergence at step 0); + - each step, all-reduce (average) their gradients — with the same batch this + is NOT a no-op for replicated params, because each rank backprops through + DIFFERENT local experts, so the shared router receives different per-rank + gradients and would re-diverge every step (measured spread 5.7e-4 -> 8.9e-3 + over 3 steps without the sync). +Expert params are deliberately left per-rank distinct at init and their +gradients are NOT reduced — that asymmetry is expert parallelism itself. With +replicated params in lock-step and the same batch, gate (2) cross-rank loss +agreement is a true transport test at any step count: a rank whose dispatched +tokens were corrupted in the MoE all-to-all disagrees, a correct one does not. +""" +import datetime +import os +import sys +import traceback + +import torch +import torch.distributed as dist + +RANK = int(os.environ["RANK"]) +WORLD = int(os.environ["WORLD_SIZE"]) +LOCAL = int(os.environ["LOCAL_RANK"]) + +DISPATCHER = os.environ.get("MOE_DISPATCHER", "alltoall") +STEPS = int(os.environ.get("TRAIN_STEPS", "3")) +E = int(os.environ.get("EP_EXPERTS", "128")) +K = int(os.environ.get("EP_TOPK", "8")) +H = int(os.environ.get("EP_HIDDEN", "2048")) +SEQ = int(os.environ.get("TRAIN_SEQ", "256")) +MBS = int(os.environ.get("TRAIN_MBS", "2")) +VOCAB = int(os.environ.get("TRAIN_VOCAB", "8192")) +LR = float(os.environ.get("TRAIN_LR", "1e-3")) + + +def log(msg): + print(f"[rank{RANK}] {msg}", flush=True) + + +def main(): + torch.cuda.set_device(LOCAL) + # Explicit timeout so a wedged peer fails fast instead of holding the nodes + # for the default 30-min NCCL timeout (see run-rollout-probe.sh / the docs: + # the worker runs detached under kubectl exec). + dist.init_process_group( + "nccl", rank=RANK, world_size=WORLD, + timeout=datetime.timedelta(minutes=10), + ) + dev = torch.device("cuda", LOCAL) + # gate 2 (cross-rank loss agreement) is a per-step invariant tracked across + # ALL steps below — needs at least two steps, and losses[-1] < losses[0] + # (the decrease gate) is meaningless with fewer. + assert STEPS >= 2, f"TRAIN_STEPS must be >= 2 (got {STEPS})" + + from megatron.core import parallel_state + from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_spec + from megatron.core.models.gpt.gpt_model import GPTModel + from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed + from megatron.core.transformer.transformer_config import TransformerConfig + + ep_size = int(os.environ.get("EP_SIZE", str(WORLD))) + assert WORLD % ep_size == 0, f"world ({WORLD}) must divide by EP size ({ep_size})" + assert E % ep_size == 0, f"experts ({E}) must divide by EP size ({ep_size})" + + parallel_state.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + expert_model_parallel_size=ep_size, + ) + model_parallel_cuda_manual_seed(123) + log(f"train-step start world={WORLD} ep={ep_size} dispatcher={DISPATCHER} " + f"experts={E} topk={K} hidden={H} steps={STEPS}") + + config = TransformerConfig( + num_layers=2, + hidden_size=H, + num_attention_heads=16, + ffn_hidden_size=4 * H, + moe_ffn_hidden_size=768, # Qwen3-30B-A3B expert width + num_moe_experts=E, + moe_router_topk=K, + moe_token_dispatcher_type=DISPATCHER, + moe_enable_deepep=(DISPATCHER == "flex"), + expert_model_parallel_size=ep_size, + add_bias_linear=False, + # Dropout OFF. This is a transport-DETERMINISM gate: gate 2 asserts every rank + # computes the SAME loss from the SAME seeded batch after the MoE all-to-all + # round-trip. hidden/attention dropout default to 0.1 and inject per-rank-random + # masks (the divergence compounds through weight updates: measured cross-rank + # spread grew 0.07 -> 0.15 -> 0.47 over 3 steps with dropout on), which defeats + # that invariant without indicating any transport fault. moe_input_jitter and + # aux-loss are already off by default, so with dropout off the router decision and + # expert compute are deterministic and a correct all-to-all yields bit-consistent + # cross-rank loss (modulo FP-reduction noise, well under the 1e-2 gate). + hidden_dropout=0.0, + attention_dropout=0.0, + # fp32 throughout — this is a transport-correctness gate, not a fidelity run. + # The megatron `local` layer spec (get_gpt_layer_local_spec, chosen to keep this + # driver free of megatron's TE/DDP wrapper machinery) emits fp32 LayerNorm + # activations; with bf16 params the router's te_general_gemm then sees a + # bf16-weight x fp32-input GEMM, which this NGC base's TransformerEngine cuBLASLt + # build rejects ("unsupported value or parameter") regardless of moe_router_dtype + # (measured: moe_utils RouterGatingLinearFunction always takes the TE path for any + # router_dtype != fp64). fp32 params make every GEMM (fp32,fp32,fp32) — proven to + # run fwd+bwd+opt here — and NCCL/DeepEP all-to-all is dtype-agnostic (DeepEP even + # requires fp32 probs), so fp32 exercises the same transport path a bf16 run would + # while making the cross-rank loss-agreement gate exact (no bf16 rounding near the + # 1e-2 threshold). A bf16 run would instead need the TE layer spec, pulling in the + # megatron wrapper machinery this driver deliberately avoids. + params_dtype=torch.float32, + pipeline_dtype=torch.float32, + use_cpu_initialization=False, + ) + model = GPTModel( + config=config, + transformer_layer_spec=get_gpt_layer_local_spec(num_experts=E, moe_grouped_gemm=False), + vocab_size=VOCAB, + max_sequence_length=SEQ, + ).to(dev) + n_params = sum(p.numel() for p in model.parameters()) + log(f"model up: {n_params/1e6:.1f}M params on this rank") + + # Sync REPLICATED params from rank 0 (what DDP does at construction; this driver + # is deliberately DDP-free — see docstring). Megatron's default RNG tracker seeds + # most replicated params identically, but the router weight and position embeddings + # init under a per-rank-varying RNG context — MEASURED cross-rank divergence 1.5e-1 + # (router.weight) / 1.6e-1 (position_embeddings) at step 0, before any update. Left + # unsynced, ranks route the same batch differently and gate 2 false-fails on a fully + # correct image + transport. Broadcast makes gate 2's invariant hold AND keeps it a + # real transport test (a corrupted all-to-all still diverges the combined output -> + # loss). Expert params (.experts./.local_experts.) are correctly left per-rank + # distinct — that asymmetry IS expert parallelism. + for name, p in model.named_parameters(): + if ".experts." not in name and ".local_experts." not in name: + dist.broadcast(p.data, src=0) + + # Same seeded batch on every rank — see the module docstring for why that + # makes the DDP-less loop sound and turns loss agreement into a gate. + g = torch.Generator(device="cpu").manual_seed(1234) + tokens = torch.randint(0, VOCAB, (MBS, SEQ), generator=g).to(dev) + position_ids = torch.arange(SEQ, device=dev).unsqueeze(0).expand(MBS, -1) + # boolean mask, True = masked (megatron local-spec attention convention) + attention_mask = torch.triu( + torch.ones(SEQ, SEQ, dtype=torch.bool, device=dev), diagonal=1 + ).unsqueeze(0).unsqueeze(0) + labels = torch.roll(tokens, shifts=-1, dims=1) + + # Partition params once: replicated (kept in lock-step across ranks, DDP-style) vs + # expert (per-rank distinct — expert parallelism). Same split as the init broadcast. + replicated_params = [p for n, p in model.named_parameters() + if ".experts." not in n and ".local_experts." not in n] + + optimizer = torch.optim.AdamW(model.parameters(), lr=LR) + losses = [] + max_spread = 0.0 + for step in range(STEPS): + optimizer.zero_grad(set_to_none=True) + # fp32 model (see config) — no autocast needed; every linear sees fp32 x fp32. + per_token_loss = model(tokens, position_ids, attention_mask, labels=labels) + loss = per_token_loss.float().mean() + loss.backward() + # All-reduce (average) gradients of REPLICATED params — the other half of what + # DDP does. The original "same batch => grad all-reduce is a no-op" premise is + # FALSE for these: each rank backprops through DIFFERENT local experts, so the + # shared router/embeddings receive different gradients per rank. Left unsynced the + # router re-diverges after every step (MEASURED: cross-rank spread grew + # 5.7e-4 -> 8.9e-3 over 3 steps without this). Syncing keeps replicated params + # bit-identical across ranks at every step, so gate 2 stays a tight transport test + # regardless of TRAIN_STEPS, and the loss decrease is a true synchronized-train + # decrease. Expert grads are deliberately NOT reduced — that is expert parallelism. + for p in replicated_params: + if p.grad is not None: + dist.all_reduce(p.grad, op=dist.ReduceOp.AVG) + optimizer.step() + losses.append(loss.item()) + # gate 2: cross-rank agreement (same batch => same loss on every rank) + t = torch.tensor([loss.item()], device=dev) + t_min, t_max = t.clone(), t.clone() + dist.all_reduce(t_min, op=dist.ReduceOp.MIN) + dist.all_reduce(t_max, op=dist.ReduceOp.MAX) + spread = (t_max - t_min).item() + max_spread = max(max_spread, spread) + log(f"TRAIN-STEP step={step} loss={loss.item():.4f} cross-rank-spread={spread:.2e}") + + finite = all(l == l and abs(l) != float("inf") for l in losses) + decreasing = losses[-1] < losses[0] + # gate 2 is a PER-STEP invariant — check the worst step, not just the last + # (an early-step transport fault that later clears must still fail the gate). + agree = max_spread < 1e-2 + ok = finite and decreasing and agree + log(f"TRAIN-STEP losses={['%.4f' % l for l in losses]} " + f"finite={finite} decreasing={decreasing} cross-rank-agree={agree}") + + flag = torch.tensor([1 if ok else 0], device=dev, dtype=torch.int32) + dist.all_reduce(flag, op=dist.ReduceOp.MIN) + verdict = "TRAIN-STEP-PASS" if int(flag.item()) == 1 else "TRAIN-STEP-FAIL" + log(f"{verdict} dispatcher={DISPATCHER} world={WORLD} ep={ep_size}") + + dist.barrier() + parallel_state.destroy_model_parallel() + return 0 if int(flag.item()) == 1 else 4 + + +if __name__ == "__main__": + rc = 1 + try: + rc = main() + except Exception: + traceback.print_exc() + print(f"[rank{RANK}] TRAIN-STEP-EXC", flush=True) + finally: + # Tear the group down on every path so a rank that raised does not leave + # its peers blocked in the MIN all_reduce / barrier until the timeout. + if dist.is_initialized(): + dist.destroy_process_group() + sys.exit(rc) diff --git a/examples/training/nemo-rl/deepep-v2-efa/recipe/verify-image.sh b/examples/training/nemo-rl/deepep-v2-efa/recipe/verify-image.sh new file mode 100755 index 000000000..564d04eb3 --- /dev/null +++ b/examples/training/nemo-rl/deepep-v2-efa/recipe/verify-image.sh @@ -0,0 +1,94 @@ +#!/usr/bin/env bash +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. SPDX-License-Identifier: MIT-0 +# Smoke the EFA + DeepEP-V2 + NeMo-RL substrate in the built image BEFORE any cluster +# deploy. Fails loud; there is no unconditional PASS in this file. +# Static image gate: asserts what the image stages (libs, symbols, imports, patch-marker +# consistency). The live cross-node transport proof is run-rollout-probe.sh, not this. +set -euo pipefail +IMG="${1:?usage: verify-image.sh }" + +# EFA device mapping: fi_info -p efa only resolves with the device visible in the +# container. On an EFA host, pass /dev/infiniband through; elsewhere fall back to a +# provider-compiled-in check (fi_info -l) and say so. +DEV_ARGS=() +HAVE_EFA_DEV=0 +if [ -d /dev/infiniband ]; then + DEV_ARGS=(-v /dev/infiniband:/dev/infiniband --device=/dev/infiniband) + HAVE_EFA_DEV=1 +fi + +docker run --rm --gpus all "${DEV_ARGS[@]}" -e HAVE_EFA_DEV="${HAVE_EFA_DEV}" "${IMG}" bash -lc ' + set -euo pipefail + if [ "${HAVE_EFA_DEV}" = "1" ]; then + echo "== fi_info efa (live fabric) ==" + [ "$(/opt/amazon/efa/bin/fi_info -p efa | grep -c "fabric: efa-direct")" -ge 1 ] || { echo "FAIL: no efa-direct"; exit 1; } + else + echo "== fi_info efa (no /dev/infiniband on this host — checking provider is compiled in only) ==" + /opt/amazon/efa/bin/fi_info -l | grep -qi "efa" || { echo "FAIL: efa provider not in libfabric"; exit 1; } + echo " (run this script on an EFA host for the live efa-direct fabric check)" + fi + + echo "== the libnccl the LOADER resolves is the GIN build (path + version) ==" + # Ask the loader, not the ld.so cache. The image sets LD_LIBRARY_PATH to the + # GIN build, and dlopen searches LD_LIBRARY_PATH BEFORE the cache, so + # "ldconfig -p | head -1" (cache order) is not what decides resolution. + # dlopen("libnccl.so.2") reproduces the real loader search (RPATH, + # LD_LIBRARY_PATH, LD_PRELOAD, then cache); /proc/self/maps then reports the + # file actually mapped — the same evidence deep_ep.check_nccl_so() acts on at + # import. This catches a scrubbed LD_LIBRARY_PATH or an LD_PRELOAD shadow, + # which the cache-order check reads right past. + NCCL_SO=$(python3 -c "import ctypes, pathlib; ctypes.CDLL(\"libnccl.so.2\"); print(next(l.split()[-1] for l in pathlib.Path(\"/proc/self/maps\").read_text().splitlines() if \"libnccl.so.2\" in l))") + echo "$NCCL_SO" | grep -q "/opt/nccl/build" || { echo "FAIL: loader resolves libnccl to $NCCL_SO, not the GIN build under /opt/nccl/build (LD_LIBRARY_PATH scrubbed, or a baked libnccl shadows it)"; exit 1; } + [ "$(strings "$NCCL_SO" | grep -c "NCCL version 2.30.4")" -ge 1 ] || { echo "FAIL: $NCCL_SO is not 2.30.4 — wrong NCCL resolved"; exit 1; } + + echo "== aws-ofi-nccl GIN plugin ==" + [ "$(nm -D /opt/aws-ofi-nccl/lib/libnccl-net-ofi.so | grep -c ncclGinPlugin)" -ge 1 ] || { echo "FAIL: no ncclGinPlugin symbol"; exit 1; } + [ "$(strings /opt/aws-ofi-nccl/lib/libnccl-net-ofi.so | grep -c "GDRCopy support not available at compile time")" -eq 0 ] \ + || { echo "FAIL: plugin built without gdrcopy — GIN init will fail at run time"; exit 1; } + + echo "== deep_ep V2 (NCCL backend) import + ElasticBuffer ==" + python3 -c "from deep_ep import ElasticBuffer; import deep_ep; print(\"deep_ep at\", deep_ep.__file__)" \ + || { echo "FAIL: deep_ep import / ElasticBuffer missing — not an EPv2 build?"; exit 1; } + + echo "== deep_ep is the amazon-contributing fork (carries the former DeepEP#612 fixes in-code) ==" + # BASELINE property, asserted on EVERY flavor (patched or not): the image pins the + # amazon-contributing/DeepEP fork, which carries both halves of what was draft + # deepseek-ai/DeepEP#612 structurally — the get_rdma_gbs() sysfs link-rate fast path + # and the auto-QP overflow clamp. So these are NOT gated on the draft-PR marker (as + # the old EP_EFA_MAX_QPS/EP_EFA_RDMA_GBS patch was); a stock-upstream or pre-fix + # checkout lacks them and fails here, not at a distant runtime error. + DEEP_EP_DIR=$(python3 -c "import deep_ep, pathlib; print(pathlib.Path(deep_ep.__file__).parent)") + grep -q unordered_gin_qps "$DEEP_EP_DIR/buffers/elastic.py" \ + || { echo "FAIL: installed deep_ep lacks the #612 auto-QP overflow clamp (unordered_gin_qps) — not the amazon-contributing fork?"; exit 1; } + grep -q _get_sysfs_rdma_gbs "$DEEP_EP_DIR/utils/envs.py" \ + || { echo "FAIL: installed deep_ep lacks the #612 get_rdma_gbs() sysfs link-rate fast path — not the amazon-contributing fork?"; exit 1; } + + echo "== nemo_rl + megatron.core imports ==" + python3 -c "import nemo_rl.algorithms.grpo; import nemo_rl.models.megatron; print(\"nemo_rl OK\")" \ + || { echo "FAIL: nemo_rl imports (algorithms.grpo / models.megatron)"; exit 1; } + python3 -c "import megatron.core; import megatron.core.transformer.moe.fused_a2a; print(\"megatron.core OK, tree:\", megatron.core.__file__)" \ + || { echo "FAIL: megatron.core imports — is /opt/Megatron-LM on PYTHONPATH?"; exit 1; } + + echo "== patch-marker consistency (marker and trees must agree) ==" + # DeepEP is NOT probed here: its former #612 fixes are carried structurally by the + # amazon-contributing fork on EVERY flavor, so they are asserted unconditionally in + # the fork-discriminator gate above — not gated on this draft-PR marker. This block + # covers only the two draft PRs that ARE applied as patches (Megatron-LM#4632, + # NeMo-RL#2410). Probe one needle PER independently-required change, not one per PR: + # a partially-merged tree satisfies a single needle while lacking later commits + # (e.g. ElasticBuffer present but the num_experts backward-dispatch fix absent). This + # mirrors the per-change probe list in patches/apply_nemo_rl_patches.py. + if [ -f /opt/.draft-rollout-patches-applied ]; then + MEGA_A2A=/opt/Megatron-LM/megatron/core/transformer/moe/fused_a2a.py + grep -q ElasticBuffer "$MEGA_A2A" \ + || { echo "FAIL: marker says patched but Megatron-LM#4632 ElasticBuffer support not in the flex dispatcher"; exit 1; } + grep -q _handle_num_experts "$MEGA_A2A" \ + || { echo "FAIL: marker says patched but Megatron-LM#4632 num_experts backward-dispatch fix absent — partially-applied PR"; exit 1; } + test -f /opt/NeMo-RL/examples/configs/recipes/llm/aws-efa-grpo-qwen3-30ba3b-2n8g-megatron.yaml \ + || { echo "FAIL: marker says patched but NeMo-RL#2410 EFA recipe config missing"; exit 1; } + echo " patched image (2 draft PRs baked — full GRPO rollout-over-DeepEP path staged)" + else + echo " unpatched baseline (upstream-only Megatron/NeMo-RL trees — probe/train-step run with explicit SM/QP counts)" + fi + echo "ALL CHECKS PASS" +' diff --git a/examples/training/nemo-rl/deepep-v2-efa/requirements.txt b/examples/training/nemo-rl/deepep-v2-efa/requirements.txt new file mode 100644 index 000000000..bcc40ccf9 --- /dev/null +++ b/examples/training/nemo-rl/deepep-v2-efa/requirements.txt @@ -0,0 +1,84 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 + +# NeMo-RL runtime dependencies for the deepep-v2-efa image. +# +# NeMo-RL itself is installed with --no-deps (nemo-rl.Dockerfile Layer 7) +# because its pyproject pins torch EXACTLY (torch==2.9.0 at NEMO_RL_SHA=46be4e8): +# letting pip resolve that would DOWNGRADE the NGC-baked torch and orphan the +# baked TransformerEngine/apex/flash-attn ABI. This file carries the rest of +# NeMo-RL's [project.dependencies] set, at the specifiers that SHA declares +# EXCEPT for the documented deviations called out inline below, MINUS the +# entries the base image already provides, conflict with this substrate, or +# serve a NeMo-RL feature path this Megatron-GRPO recipe does not exercise: +# - torch, triton, torchvision : NGC-baked (the ABI anchor; never reinstall) +# - setuptools, pip, ninja : present in the base +# - nccl4py, cuda-bindings : NeMo-RL's non-colocated-refit path only; this +# recipe is colocated (GRPO+Megatron), so unused +# - pybase64 : NeMo-RL's SGLang-refit path only (no SGLang here) +# - soundfile : audio multimodal path (not on this text recipe) +# Do NOT read this as "identical to upstream": `pip check` is deliberately NOT a +# build gate, because --no-deps against the NGC-baked torch 2.11 (vs NeMo-RL's +# declared torch==2.9.0) would make pip check false-fail on that intended ABI +# substitution. Re-diff by hand on every NEMO_RL_SHA bump (bottom of file). +# NVSHMEM IS required as a BUILD-TIME dependency, not stripped: upstream DeepEP's +# setup.py links NVSHMEM UNCONDITIONALLY for the deep_ep extension (its own +# `# TODO: make NVSHMEM and legacy optional` is still open at 01dc3aa, at main +# HEAD, and in the amazon-contributing fork), via find_nvshmem_root(optional= +# False) which asserts if NVSHMEM is absent. The NGC base's baked NVSHMEM is the +# dpkg `libnvshmem3-cuda-13` RUNTIME package — it has NO nvshmem.h header and NO +# libnvshmem_device.a, so the Layer-9 deep_ep build would fail at compile/link. +# The pip wheel below ships both (header + static device lib) and resolves +# through find_pkg_root; cu13 to match the base's CUDA line. This mirrors what +# upstream NeMo-RL 46be4e8 lists as `nvidia-nvshmem-cu12 # for deep_ep build`. +# The NCCL-GIN backend does not USE NVSHMEM at run time (the network path is +# nccl.cu), but the extension is still LINKED against it — so it must be present. +# When bumping NEMO_RL_SHA, re-diff this list against the new pyproject. +colored==2.2.3 +ray[default]==2.49.2 +transformers==4.57.1 +wandb>=0.25.0 +numpy +datasets>=4.0.0 +rich +math-verify +accelerate>=0.26 +tensorboard +omegaconf +torchdata +nvidia-ml-py +hydra-core +tiktoken +blobfile +debugpy +nvtx +matplotlib +plotly +sympy>=1.14.0 +pillow>=11.3.0 +num2words>=0.5.14 +mlflow>=3.5.0,<3.6.0 +swanlab +pyzmq + +# NVSHMEM: build-time link dependency of the deep_ep C extension (see header). +# cu13 wheel to match the base's CUDA 13 line; ships nvshmem.h + the static +# libnvshmem_device.a that the dpkg runtime package omits. +nvidia-nvshmem-cu13 + +# megatron.core (PYTHONPATH tree) import-time dependency not in NeMo-RL's set. +einops + +# decord2: imported (as `decord`) by nemo_rl/data/multimodal_utils.py, which the +# GRPO data path pulls in — so `import nemo_rl.algorithms.grpo` +# (recipe/verify-image.sh Gate) is an ImportError without it. Absent from the NGC +# base. NeMo-RL 46be4e8 declares it as `decord2` (the maintained fork that keeps +# the `decord` import name), so match that spelling exactly. +decord2 + +# nvidia-resiliency-ext: import-time dependency of megatron.core's +# dist_checkpointing strategies (strategies/{nvrx,torch}.py, tensor_aware_state_dict.py), +# so `import megatron.core.transformer.moe.fused_a2a` fails without it. The NGC +# base ships 0.5.0; megatron.core at MEGATRON_LM_SHA needs >=0.6.0 (the 0.5.0 +# API is missing symbols it imports) — pin the floor, not an exact version. +nvidia-resiliency-ext>=0.6.0 diff --git a/examples/training/nemo-rl/deepep-v2-efa/setup_nemo_rl_deepep_efa.sh b/examples/training/nemo-rl/deepep-v2-efa/setup_nemo_rl_deepep_efa.sh new file mode 100755 index 000000000..747be2e94 --- /dev/null +++ b/examples/training/nemo-rl/deepep-v2-efa/setup_nemo_rl_deepep_efa.sh @@ -0,0 +1,131 @@ +#!/usr/bin/env bash +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +# +# setup_nemo_rl_deepep_efa.sh — build the two source components that carry +# DeepEP V2's MoE all-to-all onto AWS EFA for the NeMo-RL test case: +# ofi aws-ofi-nccl with GIN (CPU-proxy) compiled in. DeepEP V2's NCCL +# backend drives its network traffic through NCCL's GIN device API, +# and the EFA installer's stock plugin does not carry GIN — hence +# this source build. +# deepep the deep_ep python package (upstream deepseek-ai/DeepEP, EPv2) +# compiled against the GIN-capable NCCL under $NCCL_HOME. Run AFTER +# the optional draft-PR patch layer so patched kernels compile in. +# +# The deepep phase builds the SAME DeepEP V2 NCCL-GIN backend as the repo's +# canonical V2 builder, +# micro-benchmarks/expert-parallelism/deepep-v2-benchmark/setup_deepep_gin.sh. +# This is the NeMo-RL test case's self-contained variant: it ALSO builds the +# aws-ofi-nccl GIN plugin (the ofi phase), and it compiles the same +# amazon-contributing/DeepEP fork the canonical setup_deepep_gin.sh clones (the +# Dockerfile pins it at an immutable fork SHA). The fork carries the EFA delta +# in-code — including both halves of what was draft deepseek-ai/DeepEP#612 — so +# there is no local source patch on the DeepEP tree. +# NOTE on NVSHMEM: deep_ep's _C.so IS build-linked against NVSHMEM — see the +# Dockerfile's Layer 7b, which makes the pip nvidia-nvshmem-cu13 lib win the +# loader search (without it `import deep_ep` dies on an nvshmem undefined +# symbol). NVSHMEM is simply NOT the run-time transport on the V2 NCCL-GIN +# path; the GIN backend carries the traffic. +# Only the V1 NVSHMEM-path setup_deepep_efa.sh copies are vendor-synced (via +# .github/workflows/deepep-vendor-sync.yml); neither this script nor the +# canonical setup_deepep_gin.sh is in that sync. +# +# Runs inside the Docker build (no GPU needed). +set -euo pipefail + +PHASE="${1:?usage: setup_nemo_rl_deepep_efa.sh }" + +# ---- pins (defaults match the Dockerfile ARGs; immutable SHAs only) -------- +AWS_OFI_NCCL_REPO="${AWS_OFI_NCCL_REPO:-https://github.com/aws/aws-ofi-nccl.git}" +AWS_OFI_NCCL_SHA="${AWS_OFI_NCCL_SHA:-9c44d34476f90ddbf4a12d0ac4fc412d46bd8ab4}" # GIN plugin, gdrdrv-2.4 v1-fallback baked +AWS_OFI_NCCL_PR="${AWS_OFI_NCCL_PR:-1351}" # OFI_NCCL_GDRCOPY_FORCED_PCIE_COPY param +AWS_OFI_NCCL_PR_SHA="${AWS_OFI_NCCL_PR_SHA:-c2e773dfb2c75b765b3415f8ffd1b47e7c239a7b}" # IMMUTABLE PR#1351 head +NCCL_HOME="${NCCL_HOME:-/opt/nccl/build}" +DEEPEP_SRC="${DEEPEP_SRC:-/opt/DeepEP}" + +build_ofi() { + echo "== aws-ofi-nccl GIN @ ${AWS_OFI_NCCL_SHA} + PR#${AWS_OFI_NCCL_PR} ==" + git clone "${AWS_OFI_NCCL_REPO}" /opt/aws-ofi-nccl-src + cd /opt/aws-ofi-nccl-src + git config user.email build@local; git config user.name build + git fetch origin "${AWS_OFI_NCCL_SHA}"; git checkout "${AWS_OFI_NCCL_SHA}" + grep -q FALLBACK_V1_FOR_GDRDRV_24 src/nccl_ofi_gdrcopy.cpp # assert the v1-fallback is present (fail-loud) + if [ -n "${AWS_OFI_NCCL_PR_SHA}" ]; then + git fetch origin "${AWS_OFI_NCCL_PR_SHA}" + git cherry-pick "${AWS_OFI_NCCL_PR_SHA}" + grep -q GDRCOPY_FORCED_PCIE_COPY include/nccl_ofi_param.h # assert the param landed (fail-loud) + fi + git rev-parse HEAD > /opt/aws-ofi-nccl.effective.sha + ./autogen.sh + # --with-nccl-headers: the ncclGin* device API headers only >= 2.30.x carries, + # taken from the source-built GIN NCCL under $NCCL_HOME (Dockerfile Layer 4). + ./configure --prefix=/opt/aws-ofi-nccl --with-libfabric=/opt/amazon/efa --with-cuda=/usr/local/cuda \ + --with-nccl-headers="${NCCL_HOME}/include" \ + --enable-cudart-dynamic --enable-platform-aws + make -C src -j"$(nproc)"; make -C src install + test -f /opt/aws-ofi-nccl/lib/libnccl-net-ofi.so + [ "$(nm -D /opt/aws-ofi-nccl/lib/libnccl-net-ofi.so | grep -c ncclGinPlugin)" -ge 1 ] # GIN symbol present (fail-loud) + # GIN needs gdrcopy COMPILED IN (gdrapi.h at configure time) — a gdrapi-less + # build carries this exact runtime-warn string and fails GIN init on first + # use. Assert the string is ABSENT from the built plugin. + [ "$(strings /opt/aws-ofi-nccl/lib/libnccl-net-ofi.so | grep -c 'GDRCopy support not available at compile time')" -eq 0 ] + ldconfig + cd /; rm -rf /opt/aws-ofi-nccl-src + echo "== ofi phase complete: GIN plugin at /opt/aws-ofi-nccl/lib/libnccl-net-ofi.so ==" +} + +build_deepep() { + echo "== DeepEP V2 (NCCL backend) from ${DEEPEP_SRC} against NCCL at ${NCCL_HOME} ==" + [ -d "${DEEPEP_SRC}" ] || { echo "ERROR: ${DEEPEP_SRC} missing — Dockerfile Layer 6 clones it" >&2; exit 1; } + test -f "${NCCL_HOME}/include/nccl_device.h" \ + || { echo "ERROR: ${NCCL_HOME}/include/nccl_device.h missing — NCCL is not GIN-capable" >&2; exit 1; } + # The V2 setup reads EP_NCCL_ROOT_DIR for the NCCL backend's headers/libs; + # assert this tree actually carries the NCCL backend (EPv2) before building, + # so a wrong-SHA checkout fails here and not with a distant include error. + grep -q EP_NCCL_ROOT_DIR "${DEEPEP_SRC}/setup.py" \ + || { echo "ERROR: ${DEEPEP_SRC}/setup.py has no EP_NCCL_ROOT_DIR — not an EPv2 (NCCL backend) tree" >&2; exit 1; } + # Fail-loud: the tree must be the amazon-contributing/DeepEP fork carrying both + # halves of the superseded draft deepseek-ai/DeepEP#612 in-code (a stock-upstream + # or pre-fix checkout lacks them). This is what makes the baseline image carry the + # EFA fixes with no local patch — assert it here, not at a distant runtime failure. + grep -q "_get_sysfs_rdma_gbs" "${DEEPEP_SRC}/deep_ep/utils/envs.py" \ + || { echo "ERROR: ${DEEPEP_SRC} lacks the #612 get_rdma_gbs() sysfs link-rate fast path — not the amazon-contributing fork?" >&2; exit 1; } + grep -q "unordered_gin_qps" "${DEEPEP_SRC}/deep_ep/buffers/elastic.py" \ + || { echo "ERROR: ${DEEPEP_SRC} lacks the #612 auto-QP overflow clamp — not the amazon-contributing fork?" >&2; exit 1; } + cd "${DEEPEP_SRC}" + git rev-parse HEAD > /opt/deepep.effective.sha + export EP_NCCL_ROOT_DIR="${EP_NCCL_ROOT_DIR:-${NCCL_HOME}}" + export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-9.0}" + # --no-deps: deep_ep's metadata must not drag a second torch/NCCL into the + # image (the pip nvidia-nccl wheels are exactly what Layer 4 removed). + pip3 install --no-cache-dir --no-build-isolation --no-deps -v . + # Build sandbox has no GPU, so no import smoke here (recipe/verify-image.sh + # does that on an EFA/GPU host). Assert the compiled extension landed by + # reading the INSTALLED distribution's file manifest (the wheel RECORD, via + # importlib.metadata) — NOT importlib.util.find_spec: this runs with cwd set + # to ${DEEPEP_SRC} (the `cd` above), and Python puts cwd on sys.path, so + # find_spec("deep_ep") would resolve the SOURCE tree's deep_ep/ package — + # which has no compiled .so — and shadow the pip-installed copy, turning this + # into a false-negative that fails the build even though the .so was built and + # installed. metadata.files() is cwd-independent and imports nothing (no GPU + # touched); `cd /` is belt-and-suspenders against any cwd-on-path shadowing. + SO_COUNT=$(cd / && python3 - <<'PY' +import sys +import importlib.metadata as md +try: + files = md.files("deep_ep") +except md.PackageNotFoundError: + sys.exit("deep_ep is not installed (no distribution metadata found)") +so = [f for f in (files or []) if str(f).endswith(".so") and f.locate().is_file()] +print(len(so)) +PY +) + [ "${SO_COUNT}" -ge 1 ] || { echo "ERROR: no compiled deep_ep extension (.so) found after install" >&2; exit 1; } + echo "== deepep phase complete: deep_ep built @ $(cat /opt/deepep.effective.sha) ==" +} + +case "${PHASE}" in + ofi) build_ofi ;; + deepep) build_deepep ;; + *) echo "FATAL: unknown phase '${PHASE}' (ofi|deepep)" >&2; exit 2 ;; +esac