From ce6c118d6fee726b462f46806fbd79581f7ee17d Mon Sep 17 00:00:00 2001 From: Anton Alexander Date: Tue, 25 Aug 2026 07:07:38 +0000 Subject: [PATCH 01/17] feat(deepep-efa): NeMo-RL GRPO + Megatron Shape-Y MoE all-to-all over EFA (DeepEP-V2 NCCL-GIN) - new test case 3.test_cases/pytorch/nemo-rl/deepep-v2-efa (mirrors slime sibling shape) - NGC-from-scratch Dockerfile, DeepEP-V2/NCCL-Gin/EFA from public source; opt-in default-OFF draft-PR layer - recipe verify+rollout-probe+train-step, RayCluster 2-node manifest, engine-index README - honest measured (Wave-28 2xp5 H100) vs staged (rc image build-staged, draft-PR-dependent) Signed-off-by: Anton Alexander --- 3.test_cases/pytorch/nemo-rl/README.md | 23 ++ .../pytorch/nemo-rl/deepep-v2-efa/.gitignore | 9 + .../pytorch/nemo-rl/deepep-v2-efa/README.md | 266 ++++++++++++++++++ .../nemo-rl/deepep-v2-efa/env_vars.example | 93 ++++++ .../kubernetes/data-prep-pod.yaml | 40 +++ .../deepep-v2-efa/kubernetes/raycluster.yaml | 214 ++++++++++++++ .../nemo-rl/deepep-v2-efa/nemo-rl.Dockerfile | 206 ++++++++++++++ .../patches/apply_nemo_rl_patches.py | 202 +++++++++++++ .../deepep-v2-efa/recipe/probe_rollout.py | 147 ++++++++++ .../deepep-v2-efa/recipe/run-rollout-probe.sh | 82 ++++++ .../deepep-v2-efa/recipe/train-step.sh | 59 ++++ .../deepep-v2-efa/recipe/train_moe_step.py | 160 +++++++++++ .../deepep-v2-efa/recipe/verify-image.sh | 69 +++++ .../nemo-rl/deepep-v2-efa/requirements.txt | 47 ++++ .../deepep-v2-efa/setup_nemo_rl_deepep_efa.sh | 100 +++++++ 15 files changed, 1717 insertions(+) create mode 100644 3.test_cases/pytorch/nemo-rl/README.md create mode 100644 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/.gitignore create mode 100644 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/README.md create mode 100644 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/env_vars.example create mode 100644 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/kubernetes/data-prep-pod.yaml create mode 100644 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/kubernetes/raycluster.yaml create mode 100644 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/nemo-rl.Dockerfile create mode 100644 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/patches/apply_nemo_rl_patches.py create mode 100644 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/probe_rollout.py create mode 100644 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/run-rollout-probe.sh create mode 100644 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train-step.sh create mode 100644 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train_moe_step.py create mode 100644 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/verify-image.sh create mode 100644 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/requirements.txt create mode 100644 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/setup_nemo_rl_deepep_efa.sh diff --git a/3.test_cases/pytorch/nemo-rl/README.md b/3.test_cases/pytorch/nemo-rl/README.md new file mode 100644 index 000000000..afc030ea8 --- /dev/null +++ b/3.test_cases/pytorch/nemo-rl/README.md @@ -0,0 +1,23 @@ + + +# NeMo-RL test cases + +[NeMo-RL](https://github.com/NVIDIA-NeMo/RL) is NVIDIA's scalable post-training library (GRPO, +DPO, SFT) for models from 1 GPU to thousands, with Megatron-core and DTensor training backends. +The samples in this directory deploy NeMo-RL on AWS with high-performance EFA networking and +expert-parallel MoE all-to-all. + +## Available test cases + +| Test case | Orchestrator | Description | +| --- | --- | --- | +| [`deepep-v2-efa`](./deepep-v2-efa) | Kubernetes (Ray cluster, 2-node) | GRPO post-training with MoE expert-parallel dispatch/combine via **DeepEP V2's NCCL backend** (`ElasticBuffer`) over **AWS EFA**, using aws-ofi-nccl with **GIN** (GPU-Initiated Networking) CPU-proxy. Image built NGC-from-scratch from public sources; the recipe runs image build → static verify → cross-node rollout probe → N-step Megatron MoE training gate. Mechanism chain measured on 2× p5.48xlarge (H100); this folder's image assembly is build-staged with an opt-in draft-PR layer for the full rollout path (honest measured-vs-staged breakdown in its README). | + +For RL post-training with a different stack (SLIME + SGLang) on the same HyperPod-EKS Ray-cluster +pattern, see [`3.test_cases/pytorch/slime`](../slime). For kernel-level expert-parallelism +dispatch/combine benchmarks over EFA — including a DeepEP V2 benchmark on the same NCCL-GIN +substrate this test case uses — see +[`micro-benchmarks/expert-parallelism`](../../../micro-benchmarks/expert-parallelism). diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/.gitignore b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/.gitignore new file mode 100644 index 000000000..3f86e687e --- /dev/null +++ b/3.test_cases/pytorch/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/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/README.md b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/README.md new file mode 100644 index 000000000..21ebe49a0 --- /dev/null +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/README.md @@ -0,0 +1,266 @@ + +# 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 + +``` +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 4 draft upstream PRs** (opt-in image layer, default **OFF** — +see below). What the recipe gates re-verify on the **baseline (upstream-only)** image: 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` | 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. | +| EFA installer | 1.48.0 | The userspace of the measured NCCL-GIN substrate. Bumping is a re-measure event. | +| NCCL | `v2.30.4-1` (source build) | GIN device API generation (`nccl_device.h` asserted at build). Same NCCL line the NGC base bakes — the ld.so.conf override introduces no drift; source-built so DeepEP has one controlled header/lib root. | +| aws-ofi-nccl | commit `9c44d34` + PR#1351 head `c2e773d` | The GIN CPU-proxy plugin pins of the measured substrate (same as the TensorRT-LLM NcclEP sibling). 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 | `01dc3aa` (upstream `deepseek-ai/DeepEP` main) | Carries EPv2 (ElasticBuffer + NCCL backend) and is the **exact base of draft PR #612**, so the opt-in layer applies `--check`-clean. Stock upstream — NOT a fork. | +| Megatron-LM | `19deef67` (main) | The **exact base of draft PR #4632** (ElasticBuffer in the flex dispatcher). | +| NeMo-RL | `cc75cad` (main) | The **exact shared base of draft PRs #2410 + #2411**. 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 4 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. + +| PR | State | What it carries | +|---|---|---| +| [deepseek-ai/DeepEP#612](https://github.com/deepseek-ai/DeepEP/pull/612) | open | EFA awareness: auto-QP capped at `EP_EFA_MAX_QPS` (aws-ofi-nccl's GIN request ring is 128 slots; upstream's auto formula overruns it → `CUDA_ERROR_LAUNCH_FAILED` at first dispatch), `get_rdma_gbs()` EFA fast path (SM auto-sizing), dispatch scaleout-interval tuning. | +| [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. | +| [NVIDIA-NeMo/RL#2411](https://github.com/NVIDIA-NeMo/RL/pull/2411) | draft, closed unmerged | Bumps NeMo-RL's `deep_ep` pin to the V2 merge commit (metadata-only for this image — deep_ep is built from `/opt/DeepEP` — applied for tree self-consistency). | + +## 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. **Upstream DeepEP's SM/QP auto-sizers are EFA-blind** at the pinned SHA: auto-QP + (`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 probe passes `num_allocated_qps`/`num_sms`/`num_qps` **explicitly** (`EP_NUM_QPS=2` is the + value the #612 evidence validated on p5en) so the baseline image stays probeable; the opt-in + layer (#612) fixes the auto-sizers. +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 path** `libnccl.so.2` resolves to *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 | +| `OFI_NCCL_GIN_MAX_REQUESTS=512` | GIN request-ring depth of the measured substrate | +| `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 | +| `EP_EFA_MAX_QPS=2`, `EP_EFA_RDMA_GBS=25.0` | read only by the **patched** deep_ep (#612); inert on baseline | + +## 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 (upstream-only) +docker build -f nemo-rl.Dockerfile -t ${FULL_IMAGE} . +# opt-in flavor with the 4 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} +``` + +### 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: +# 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.** Four upstream PRs, one of them with two entries + 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 + does not certify upstream's auto-sizing on EFA — that certification is exactly PR #612. +- **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 on the same NCCL-GIN substrate see + [`micro-benchmarks/expert-parallelism`](../../../../micro-benchmarks/expert-parallelism). + +## File structure + +``` +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 4 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 + +- [DeepEP](https://github.com/deepseek-ai/DeepEP) — EPv2 / ElasticBuffer (PR#605) +- [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), [`sglang/dsr1-deepep-efa`](../../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 (NCCL GIN) benchmark diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/env_vars.example b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/env_vars.example new file mode 100644 index 000000000..8e8018c77 --- /dev/null +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/env_vars.example @@ -0,0 +1,93 @@ +# 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 4 draft upstream PRs (NeMo-RL#2410/#2411, Megatron-LM#4632, +# DeepEP#612) that the full GRPO rollout-over-DeepEP path needs; 0 = the +# upstream-only baseline. 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 ----- +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: upstream DeepEP's auto-sizers are +# EFA-blind at the pinned SHA (auto-QP overruns aws-ofi-nccl's 128-slot GIN +# ring; get_rdma_gbs() returns 0) - exactly what draft PR DeepEP#612 fixes. +# Explicit values keep the BASELINE image probeable; the patched image also +# works with the auto-sizers. +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; exported +# here too so ad-hoc shells (kubectl exec) carry the same contract. +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_GIN_MAX_REQUESTS=512 +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 +# Read only by the PATCHED deep_ep (DeepEP#612); inert on the baseline image. +export EP_EFA_MAX_QPS=2 +export EP_EFA_RDMA_GBS=25.0 + +# ----- 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 NeMo-RL#2411 path) + +# ----- NVSHMEM contract (INERT on this image - kept for the rebuild case) ----- +# This image links NO NVSHMEM (the V2 NCCL-GIN backend replaces it). 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/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/kubernetes/data-prep-pod.yaml b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/kubernetes/data-prep-pod.yaml new file mode 100644 index 000000000..2b5b42821 --- /dev/null +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/kubernetes/data-prep-pod.yaml @@ -0,0 +1,40 @@ +# 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). +# +# envsubst < kubernetes/data-prep-pod.yaml | kubectl apply -f - +# kubectl -n ${NAMESPACE} exec -it data-prep -- bash +# # inside: pip install "huggingface_hub[cli]" && \ +# # hf download ${MODEL_NAME} --local-dir ${MODEL_LOCAL} +apiVersion: v1 +kind: Pod +metadata: + name: data-prep + namespace: ${NAMESPACE} + labels: + app: nemo-rl-deepep-data-prep +spec: + containers: + - name: data-prep + image: python:3.12-slim + command: ["sleep", "infinity"] + env: + - name: HF_TOKEN + value: "" # Set before applying, or pass via kubectl set env + resources: + requests: + cpu: "4" + memory: "16Gi" + limits: + cpu: "8" + memory: "32Gi" + volumeMounts: + - name: fsx + mountPath: /fsx + volumes: + - name: fsx + persistentVolumeClaim: + claimName: ${FSX_CLAIM} + restartPolicy: Never + terminationGracePeriodSeconds: 30 diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/kubernetes/raycluster.yaml b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/kubernetes/raycluster.yaml new file mode 100644 index 000000000..724d164b9 --- /dev/null +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/kubernetes/raycluster.yaml @@ -0,0 +1,214 @@ +# 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: + 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: + - name: fsx + persistentVolumeClaim: + claimName: ${FSX_CLAIM} + - 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: + 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 + # ---- 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_GIN_MAX_REQUESTS, value: "512" } + - { 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" } + # Read only by the PATCHED deep_ep (DeepEP#612); inert on baseline. + - { name: EP_EFA_MAX_QPS, value: "2" } + - { name: EP_EFA_RDMA_GBS, value: "25.0" } + - { 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: + - name: fsx + persistentVolumeClaim: + claimName: ${FSX_CLAIM} + - 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/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/nemo-rl.Dockerfile b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/nemo-rl.Dockerfile new file mode 100644 index 000000000..75354af13 --- /dev/null +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/nemo-rl.Dockerfile @@ -0,0 +1,206 @@ +# 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 4 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. +ARG NGC_PYTORCH_BASE=nvcr.io/nvidia/pytorch:26.02-py3 +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: the GIN device API generation the measured substrate ran +# (ships include/nccl_device.h — asserted below). Same NCCL line the NGC base +# bakes, so the ld.so.conf override below introduces no version drift; built +# from source so the DeepEP build has one controlled root of headers + libs. +ARG NCCL_VERSION=v2.30.4-1 +# 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 pins +# of the measured NCCL-GIN substrate (same pins as the TensorRT-LLM NcclEP +# sibling sample). Immutable SHAs — refs/pull/N/head is a moving ref. +ARG AWS_OFI_NCCL_SHA=9c44d34476f90ddbf4a12d0ac4fc412d46bd8ab4 +ARG AWS_OFI_NCCL_PR=1351 +ARG AWS_OFI_NCCL_PR_SHA=c2e773dfb2c75b765b3415f8ffd1b47e7c239a7b +# DeepEP: upstream deepseek-ai/DeepEP main @01dc3aa — carries EPv2 (the +# ElasticBuffer + NCCL backend, merged upstream in PR#605) and is the EXACT +# base of draft PR deepseek-ai/DeepEP#612, so the opt-in layer applies +# --check-clean. NOT the amazon-contributing fork: baseline is stock upstream. +ARG DEEPEP_SHA=01dc3aaac82068020353dce2c302e38153c0bfaa +# 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: the exact shared base of draft PRs NVIDIA-NeMo/RL#2410 + #2411. +# The measured Wave-28 evidence ran 0.5.0rc0 (see README, measured-vs-staged); +# re-pinning to a release tag is a re-measure event. +ARG NEMO_RL_SHA=cc75cadfe061301bd121306f1648957583760528 +# 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. +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 +RUN rm -rf /opt/hpcx/ompi /usr/local/mpi /usr/local/ucx && ldconfig + +# ---- 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 upstream trees at their pinned SHAs ---------- +# Clones only here; the DeepEP BUILD is deferred to Layer 9 so the opt-in +# patch layer (Layer 8) can land its .cuh/.py edits BEFORE kernels compile. +RUN git clone https://github.com/deepseek-ai/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 8 (OPT-IN, default OFF): the 4 draft upstream PRs ---------------- +# NVIDIA-NeMo/RL#2410 + #2411, NVIDIA/Megatron-LM#4632, deepseek-ai/DeepEP#612 +# — the full GRPO rollout-over-DeepEP path depends on them; the BASELINE image +# does not. 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_GIN_MAX_REQUESTS=512 \ + 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 +# Read only by the PATCHED deep_ep (DeepEP#612 adds the EFA awareness); inert +# on the baseline image. Kept here so both flavors run with one manifest. +ENV EP_EFA_MAX_QPS=2 \ + EP_EFA_RDMA_GBS=25.0 + +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/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/patches/apply_nemo_rl_patches.py b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/patches/apply_nemo_rl_patches.py new file mode 100644 index 000000000..1672845d9 --- /dev/null +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/patches/apply_nemo_rl_patches.py @@ -0,0 +1,202 @@ +#!/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 (deepseek-ai/DeepEP, +NVIDIA/Megatron-LM, NVIDIA-NeMo/RL) at pinned SHAs and depends on nothing +else. The full GRPO rollout-over-DeepEP path additionally needs four 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``. + +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. Before applying anything, each PR's content probe is + checked: if it already passes — because the PR merged upstream and the + tree pin advanced past it — the whole PR is skipped and reported (a + per-commit reverse-check would false-negative on multi-commit PRs whose + later commits touch the same hunks). When a PR reports already-present, + delete its entry here. + * Post-asserted. After each PR a content probe (needle in file) confirms the + intended change is really in the tree, catching apply-succeeded-but-wrong- + tree mistakes. + +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). +# `probe` = (relative path, needle) asserted present after the PR applies; +# needle None = the file's existence is the assertion. +PATCH_SETS = [ + { + "name": "deepseek-ai/DeepEP#612 — aws-efa: QP cap, get_rdma_gbs fast path, scaleout interval", + "url": "https://github.com/deepseek-ai/DeepEP/pull/612", + "repo": "deepseek-ai/DeepEP", + "root_arg": "deepep_root", + "commits": [ + "4eddba396006b8e4aa1a3a9a505396020aba4ef7", # cap auto-QP at 2 on EFA (128-slot GIN ring) + "922a1fa7c0cd3ef047c0919638a87f9a2360346b", # EFA fast path in get_rdma_gbs (SM auto-sizing) + "28d1f7fb173f728be51632ce0026fea23243e350", # dispatch kScaleoutUpdateInterval 6 -> 16 + ], + "probe": ("deep_ep/buffers/elastic.py", "EP_EFA_MAX_QPS"), + }, + { + "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 + ], + "probe": ("megatron/core/transformer/moe/fused_a2a.py", "ElasticBuffer"), + }, + { + "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", + "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. + "probe": ("examples/configs/recipes/llm/aws-efa-grpo-qwen3-30ba3b-2n8g-megatron.yaml", None), + }, + { + "name": "NVIDIA-NeMo/RL#2411 — deps: bump deep_ep pin to the V2 merge commit", + "url": "https://github.com/NVIDIA-NeMo/RL/pull/2411", + "repo": "NVIDIA-NeMo/RL", + "root_arg": "nemo_rl_root", + "commits": [ + "711147f4401a3f532cbcdb6cb6b7e00e2023569e", + ], + # Metadata-only for this image (deep_ep is built from /opt/DeepEP, not + # from NeMo-RL's pin), but applied so the tree is self-consistent. + "probe": ("pyproject.toml", "b306af0"), + }, +] + + +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.""" + # 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: + # Already present? (PR merged upstream and the pin moved past it.) + if run(["git", "apply", "--reverse", "--check", str(patch_path)], root).returncode == 0: + return "already-present-upstream" + 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 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']} ==") + if probe_ok(root, pset["probe"]): + # PR-level neutralization: the content probe already passes, so the + # tree pin advanced past this PR's merge — delete its entry here. + print(f" already present in the tree (probe satisfied) — skipping; retire this entry") + record.extend(f"{pset['repo']}@{sha} already-present" for sha in pset["commits"]) + continue + for sha in pset["commits"]: + status = apply_commit(root, pset["repo"], sha) + print(f" {sha[:12]} {status}") + record.append(f"{pset['repo']}@{sha} {status}") + if not probe_ok(root, pset["probe"]): + rel, needle = pset["probe"] + print(f"FATAL: post-assert failed for {pset['name']}: " + f"{'missing file' if needle is None else f'needle {needle!r} absent in'} {root / rel}", + file=sys.stderr) + return 1 + print(f" post-assert OK ({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/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/probe_rollout.py b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/probe_rollout.py new file mode 100644 index 000000000..feef8dc32 --- /dev/null +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/probe_rollout.py @@ -0,0 +1,147 @@ +#!/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 draft PR +deepseek-ai/DeepEP#612 fixes) only fires at the first real cross-node dispatch. + +Baseline-image note: upstream DeepEP's SM/QP auto-sizers are EFA-blind at the +pinned SHA (auto-QP overruns aws-ofi-nccl's 128-slot GIN request ring; +`get_rdma_gbs()` reads 0 on EFA) — the exact gaps DeepEP#612 fixes. The probe +therefore passes `num_allocated_qps`/`num_sms`/`num_qps` EXPLICITLY +(EP_NUM_QPS=2 is the value the #612 evidence validated on p5en), which keeps +the unpatched baseline probeable; on a patched image the auto-sizers also work. +""" +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) + dist.init_process_group("nccl", rank=RANK, world_size=WORLD) + 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. + idx_dtype = getattr(deep_ep, "topk_idx_t", torch.int64) + 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) + if 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) + relmax = ((got - expect).abs().max() / expect.abs().max().clamp_min(1e-6)).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__": + try: + sys.exit(main()) + except Exception: + traceback.print_exc() + print(f"[rank{RANK}] PROBE-EXC", flush=True) + sys.exit(1) diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/run-rollout-probe.sh b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/run-rollout-probe.sh new file mode 100644 index 000000000..fd5c5437e --- /dev/null +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/run-rollout-probe.sh @@ -0,0 +1,82 @@ +#!/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 OFI_NCCL_GIN_MAX_REQUESTS=512 +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: upstream DeepEP's auto-sizers are EFA-blind at the pinned SHA +# (DeepEP#612 is the fix; opt-in layer). 2 QPs is the value the #612 evidence validated. +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). +efa_tx_total() { + local total=0 v + for f in /sys/class/infiniband/*/ports/1/hw_counters/tx_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. +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/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train-step.sh b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train-step.sh new file mode 100644 index 000000000..70b337677 --- /dev/null +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train-step.sh @@ -0,0 +1,59 @@ +#!/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 + DeepEP#612 + NeMo-RL#2410/#2411)." + 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 OFI_NCCL_GIN_MAX_REQUESTS=512 +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:-WARN}" + +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) =====" + +set -o pipefail +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]} + +if [ "$rc" -eq 0 ]; then + echo "TRAIN-STEP GATE PASS (node_rank=$NODE_RANK, dispatcher=$MOE_DISPATCHER)" + exit 0 +fi +echo "TRAIN-STEP GATE FAIL (node_rank=$NODE_RANK, torchrun rc=$rc) — see /tmp/train-step.$NODE_RANK.log" +exit 1 diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train_moe_step.py b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train_moe_step.py new file mode 100644 index 000000000..5fffa6c60 --- /dev/null +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train_moe_step.py @@ -0,0 +1,160 @@ +#!/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. + +Every rank sees the same synthetic batch, so data-parallel gradient +all-reduce would be a mathematical no-op — the loop deliberately runs without +a DDP wrapper to stay independent of megatron's DDP config API, and the +cross-rank loss-agreement gate (2) is what makes that sound. +""" +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) + dist.init_process_group("nccl", rank=RANK, world_size=WORLD) + dev = torch.device("cuda", LOCAL) + + 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, + params_dtype=torch.bfloat16, + pipeline_dtype=torch.bfloat16, + bf16=True, + 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") + + # 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) + + optimizer = torch.optim.AdamW(model.parameters(), lr=LR) + losses = [] + for step in range(STEPS): + optimizer.zero_grad(set_to_none=True) + per_token_loss = model(tokens, position_ids, attention_mask, labels=labels) + loss = per_token_loss.float().mean() + loss.backward() + 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() + 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] + agree = 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__": + try: + sys.exit(main()) + except Exception: + traceback.print_exc() + print(f"[rank{RANK}] TRAIN-STEP-EXC", flush=True) + sys.exit(1) diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/verify-image.sh b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/verify-image.sh new file mode 100644 index 000000000..6760ba928 --- /dev/null +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/verify-image.sh @@ -0,0 +1,69 @@ +#!/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 "== single GIN-capable libnccl wins (path + version string — the NGC base bakes its own NCCL in a DIFFERENT dir, so path matters) ==" + NCCL_SO=$(ldconfig -p | grep "libnccl.so.2 " | head -1 | awk "{print \$NF}") + echo "$NCCL_SO" | grep -q "/opt/nccl/build" || { echo "FAIL: baked libnccl shadows the GIN build ($NCCL_SO)"; 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 "== 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/site-packages must agree) ==" + if [ -f /opt/.draft-rollout-patches-applied ]; then + DEEP_EP_DIR=$(python3 -c "import deep_ep, pathlib; print(pathlib.Path(deep_ep.__file__).parent)") + grep -q EP_EFA_MAX_QPS "$DEEP_EP_DIR/buffers/elastic.py" \ + || { echo "FAIL: marker says patched but DeepEP#612 EFA cap not in installed deep_ep"; exit 1; } + grep -q ElasticBuffer /opt/Megatron-LM/megatron/core/transformer/moe/fused_a2a.py \ + || { echo "FAIL: marker says patched but Megatron-LM#4632 not in the flex dispatcher"; 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 (4 draft PRs baked — full GRPO rollout-over-DeepEP path staged)" + else + DEEP_EP_DIR=$(python3 -c "import deep_ep, pathlib; print(pathlib.Path(deep_ep.__file__).parent)") + if grep -q EP_EFA_MAX_QPS "$DEEP_EP_DIR/buffers/elastic.py" 2>/dev/null; then + echo "FAIL: no marker but the DeepEP#612 patch IS present — ambiguous patch state"; exit 1 + fi + echo " unpatched baseline (upstream-only trees — probe/train-step run with explicit SM/QP counts)" + fi + echo "ALL CHECKS PASS" +' diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/requirements.txt b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/requirements.txt new file mode 100644 index 000000000..4f910d018 --- /dev/null +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/requirements.txt @@ -0,0 +1,47 @@ +# 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 the pinned SHA): +# 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 same specifiers the pinned SHA +# declares, MINUS the entries the base image already provides or that conflict +# with this test case's substrate: +# - torch, triton, torchvision : NGC-baked (the ABI anchor; never reinstall) +# - setuptools, pip, ninja : present in the base +# - nvidia-nvshmem-cu12 : cu12 wheel on a cu13 base, and unused here — +# this image's deep_ep is the NCCL-GIN backend +# (no NVSHMEM linked at all) +# 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 +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 + +# megatron.core (PYTHONPATH tree) import-time dependency not in NeMo-RL's set. +einops diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/setup_nemo_rl_deepep_efa.sh b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/setup_nemo_rl_deepep_efa.sh new file mode 100644 index 000000000..fdbcbbc0a --- /dev/null +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/setup_nemo_rl_deepep_efa.sh @@ -0,0 +1,100 @@ +#!/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. +# +# Distinct from the NVSHMEM-path setup_deepep_efa.sh (vendor-synced via +# .github/workflows/deepep-vendor-sync.yml) — this script builds the NCCL-GIN +# V2 backend, links NO NVSHMEM, and is intentionally NOT vendor-synced. +# +# Runs inside the Docker build (no GPU needed). +set -euo pipefail + +PHASE="${1:?usage: setup_nemo_rl_deepep_efa.sh {ofi|deepep}}" + +# ---- 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; } + 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. + SO_COUNT=$(python3 - <<'PY' +import glob, importlib.util, os, sys +spec = importlib.util.find_spec("deep_ep") +if spec is None or not spec.submodule_search_locations: + sys.exit("deep_ep not installed") +root = list(spec.submodule_search_locations)[0] +print(len(glob.glob(os.path.join(os.path.dirname(root), "deep_ep*", "**", "*.so"), recursive=True) + + glob.glob(os.path.join(root, "**", "*.so"), recursive=True))) +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 From 6c14c2e288e0de8041a20cf2add76914e49f9a2d Mon Sep 17 00:00:00 2001 From: Anton Alexander Date: Wed, 26 Aug 2026 05:05:30 +0000 Subject: [PATCH 02/17] fix(deepep-efa): NeMo-RL image builds + passes its own verify-image gates on p5en Run-to-green on a live p5en (8xH200, EFA): the image failed to build on the pinned NeMo-RL SHA, and once buildable it failed its own recipe/verify-image.sh on two counts. All fixed and verified against a pristine-rootfs in-pod replay of Dockerfile layers 1-9 from these exact files -> full 6/6 gate suite now PASS (efa-direct x4 rails, GIN-capable NCCL 2.30.4 at /opt/nccl/build, ncclGinPlugin + gdrcopy-compiled-in, deep_ep V2 ElasticBuffer import, nemo_rl + megatron.core imports, clean upstream-only baseline marker). cgk p5en, 2026-08-26. Build fix -- NEMO_RL_SHA cc75cad -> 46be4e8 (+ drop draft PR #2411): cc75cad bumped requires-python to ">=3.13.13" and torch to 2.10.0, so `pip install -e /opt/NeMo-RL` (Layer 7) hard-fails the interpreter check on this py3.12 NGC base (--no-deps does NOT suppress the Requires-Python floor). 46be4e8 is the parent commit of draft PR NVIDIA-NeMo/RL#2410 (the EFA recipe), declares requires-python ">=3.12", and is the revision requirements.txt is generated from. #2411 (a deep_ep pin bump based on cc75cad) is therefore dropped from the opt-in layer -- it neither applies to the 46be4e8 tree nor belongs on py3.12, and is metadata-only here (deep_ep builds from /opt/DeepEP, not NeMo-RL's pin). The opt-in layer is now 3 draft PRs, not 4; docs, the patches script, env_vars.example, and the recipe scripts are reconciled to match. Build fix -- restore nvidia-nvshmem-cu13 in requirements.txt: upstream DeepEP's setup.py links NVSHMEM UNCONDITIONALLY (find_nvshmem_root (optional=False) asserts if absent -- its "make NVSHMEM optional" TODO is still open at 01dc3aa). The NGC base's baked NVSHMEM is the dpkg runtime package with no nvshmem.h and no libnvshmem_device.a, so the Layer-9 deep_ep build fails at link without the pip wheel (which ships both). cu13 to match the base's CUDA line; this mirrors NeMo-RL 46be4e8's own `nvidia-nvshmem-cu12 # for deep_ep build`. Gate 4 -- deep_ep import (nvshmem runtime ABI), new Layer 7b: `from deep_ep import ElasticBuffer` died with `undefined symbol: nvshmem_selected_device_transport, version NVSHMEM`. deep_ep/_C.so is correctly built against the pip nvidia-nvshmem-cu13 wheel (3.7.x), but torch/lib/libtorch_nvshmem.so -- imported first -- has a NEEDED libnvshmem_host.so.3 whose RUNPATH ends in /usr/local/cuda/lib64, where the NGC base's OLDER dpkg NVSHMEM (3.4.x, lacking that symbol) lives. RUNPATH outranks ld.so.cache, so the stale copy wins by soname before _C.so can pull the right one -- which is why an ld.so.conf.d entry (Layer 4's house style) does NOT fix it. Fix (Layer 7b): symlink the wheel's lib dir to a stable, version-agnostic /opt/nvshmem-pip-lib and prepend it to LD_LIBRARY_PATH (only LD_LIBRARY_PATH outranks RUNPATH). Fail-loud if the wheel dir is absent. Gate 5 -- nemo_rl / megatron.core imports (missing requirements): `import nemo_rl.algorithms.grpo` and `import megatron.core.transformer.moe.fused_a2a` both ImportError'd. - decord: imported by nemo_rl/data/multimodal_utils.py (GRPO data path); absent from the NGC base, an extra in NeMo-RL's own deps. - nvidia-resiliency-ext>=0.6.0: import-time dep of megatron.core's dist_checkpointing strategies; the base ships 0.5.0, which is too old. setup script -- deep_ep .so assert made cwd-independent: the post-install assert used importlib.util.find_spec("deep_ep") while cwd was the DeepEP source tree; Python puts cwd on sys.path, so find_spec resolved the SOURCE package (no compiled .so) and shadowed the pip-installed copy, a false negative that failed the build even when the .so was built. Rewrote to read the installed distribution's file manifest (importlib.metadata.files), which is cwd-independent and imports nothing. Verified: pristine-rootfs in-pod replay of layers 1-9 from these exact files, then recipe/verify-image.sh -> ALL IMAGE GATES PASS. cgk p5en, 2026-08-26. Signed-off-by: Anton Alexander --- .../pytorch/nemo-rl/deepep-v2-efa/README.md | 21 ++++++----- .../nemo-rl/deepep-v2-efa/env_vars.example | 4 +- .../nemo-rl/deepep-v2-efa/nemo-rl.Dockerfile | 37 +++++++++++++++---- .../patches/apply_nemo_rl_patches.py | 24 +++++------- .../deepep-v2-efa/recipe/train-step.sh | 2 +- .../deepep-v2-efa/recipe/verify-image.sh | 2 +- .../nemo-rl/deepep-v2-efa/requirements.txt | 37 ++++++++++++++++--- .../deepep-v2-efa/setup_nemo_rl_deepep_efa.sh | 29 ++++++++++----- 8 files changed, 106 insertions(+), 50 deletions(-) diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/README.md b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/README.md index 21ebe49a0..53df830f3 100644 --- a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/README.md +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/README.md @@ -34,7 +34,7 @@ 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 4 draft upstream PRs** (opt-in image layer, default **OFF** — +rollout-over-DeepEP path depends on 3 draft upstream PRs** (opt-in image layer, default **OFF** — see below). What the recipe gates re-verify on the **baseline (upstream-only)** image: 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 @@ -51,10 +51,10 @@ stock `alltoall` dispatcher (`train-step.sh`). Never read a build-gate as an E2E | gdrcopy | commit `c91ad9f` (= v2.5.2) | GIN **requires** gdrcopy compiled in (trap 4). Commit pin, not tag. | | DeepEP | `01dc3aa` (upstream `deepseek-ai/DeepEP` main) | Carries EPv2 (ElasticBuffer + NCCL backend) and is the **exact base of draft PR #612**, so the opt-in layer applies `--check`-clean. Stock upstream — NOT a fork. | | Megatron-LM | `19deef67` (main) | The **exact base of draft PR #4632** (ElasticBuffer in the flex dispatcher). | -| NeMo-RL | `cc75cad` (main) | The **exact shared base of draft PRs #2410 + #2411**. The measured Wave-28 evidence ran 0.5.0rc0; re-pinning to a release tag is a re-measure event. | +| 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 4 draft upstream PRs (opt-in layer, default OFF) +## The 3 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, @@ -65,8 +65,9 @@ PR states below are as of 2026-08-25 — check them before relying on this table |---|---|---| | [deepseek-ai/DeepEP#612](https://github.com/deepseek-ai/DeepEP/pull/612) | open | EFA awareness: auto-QP capped at `EP_EFA_MAX_QPS` (aws-ofi-nccl's GIN request ring is 128 slots; upstream's auto formula overruns it → `CUDA_ERROR_LAUNCH_FAILED` at first dispatch), `get_rdma_gbs()` EFA fast path (SM auto-sizing), dispatch scaleout-interval tuning. | | [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. | -| [NVIDIA-NeMo/RL#2411](https://github.com/NVIDIA-NeMo/RL/pull/2411) | draft, closed unmerged | Bumps NeMo-RL's `deep_ep` pin to the V2 merge commit (metadata-only for this image — deep_ep is built from `/opt/DeepEP` — applied for tree self-consistency). | +| [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 @@ -152,7 +153,7 @@ aws ecr create-repository --repository-name ${IMAGE} --region ${AWS_REGION} || t # baseline (upstream-only) docker build -f nemo-rl.Dockerfile -t ${FULL_IMAGE} . -# opt-in flavor with the 4 draft PRs baked (use a DISTINCT tag — never overwrite the baseline) +# opt-in flavor with the 3 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 . @@ -218,9 +219,9 @@ Treat results as your own measurement — this folder publishes none for this pa - **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.** Four upstream PRs, one of them with two entries - 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. +- **The full rollout path is draft-PR-dependent.** Three 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 does not certify upstream's auto-sizing on EFA — that certification is exactly PR #612. - **NCCL topology XML on 32-NIC p5 nodes:** stock NCCL can hit the open issue @@ -242,7 +243,7 @@ deepep-v2-efa/ ├── 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 4 draft PRs, pinned SHAs, fail-loud, self-neutralizing +│ └── apply_nemo_rl_patches.py <- the 3 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 diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/env_vars.example b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/env_vars.example index 8e8018c77..72fdee08e 100644 --- a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/env_vars.example +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/env_vars.example @@ -22,7 +22,7 @@ export TAG="v1-20260825" export FULL_IMAGE="${REGISTRY}${IMAGE}:${TAG}" # ----- Opt-in draft-PR image flavor ----- -# 1 = bake the 4 draft upstream PRs (NeMo-RL#2410/#2411, Megatron-LM#4632, +# 1 = bake the 3 draft upstream PRs (NeMo-RL#2410, Megatron-LM#4632, # DeepEP#612) that the full GRPO rollout-over-DeepEP path needs; 0 = the # upstream-only baseline. The ":-" default keeps a value pre-set on the # command line (APPLY_DRAFT_ROLLOUT_PATCHES=1 docker build ...) from being @@ -81,7 +81,7 @@ export EP_EFA_RDMA_GBS=25.0 # ----- 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 NeMo-RL#2411 path) +export HAVE_DEEP_EP_V2=True # rollout bridge feature flag (draft-PR rollout path) # ----- NVSHMEM contract (INERT on this image - kept for the rebuild case) ----- # This image links NO NVSHMEM (the V2 NCCL-GIN backend replaces it). If you diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/nemo-rl.Dockerfile b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/nemo-rl.Dockerfile index 75354af13..c21d611d7 100644 --- a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/nemo-rl.Dockerfile +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/nemo-rl.Dockerfile @@ -15,7 +15,7 @@ # -> 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 4 DRAFT upstream PRs (see the patches/ +# -> OPT-IN: additionally bakes 3 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. @@ -58,10 +58,17 @@ ARG DEEPEP_SHA=01dc3aaac82068020353dce2c302e38153c0bfaa # 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: the exact shared base of draft PRs NVIDIA-NeMo/RL#2410 + #2411. -# The measured Wave-28 evidence ran 0.5.0rc0 (see README, measured-vs-staged); -# re-pinning to a release tag is a re-measure event. -ARG NEMO_RL_SHA=cc75cadfe061301bd121306f1648957583760528 +# 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" @@ -153,8 +160,24 @@ RUN python3 -c 'import sys; assert sys.version_info >= (3, 12), f"NeMo-RL needs && pip3 install --no-cache-dir --no-deps -e /opt/NeMo-RL ENV PYTHONPATH=/opt/Megatron-LM:${PYTHONPATH:-} -# ---- Layer 8 (OPT-IN, default OFF): the 4 draft upstream PRs ---------------- -# NVIDIA-NeMo/RL#2410 + #2411, NVIDIA/Megatron-LM#4632, deepseek-ai/DeepEP#612 +# ---- 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 3 draft upstream PRs ---------------- +# NVIDIA-NeMo/RL#2410, NVIDIA/Megatron-LM#4632, deepseek-ai/DeepEP#612 # — the full GRPO rollout-over-DeepEP path depends on them; the BASELINE image # does not. 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 diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/patches/apply_nemo_rl_patches.py b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/patches/apply_nemo_rl_patches.py index 1672845d9..bbd846d96 100644 --- a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/patches/apply_nemo_rl_patches.py +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/patches/apply_nemo_rl_patches.py @@ -5,9 +5,10 @@ The BASELINE image installs three upstream trees as-is (deepseek-ai/DeepEP, NVIDIA/Megatron-LM, NVIDIA-NeMo/RL) at pinned SHAs and depends on nothing -else. The full GRPO rollout-over-DeepEP path additionally needs four upstream +else. The full GRPO rollout-over-DeepEP path additionally needs three 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``. +image is built with ``--build-arg APPLY_DRAFT_ROLLOUT_PATCHES=1``. (A fourth, +NVIDIA-NeMo/RL#2411, is intentionally excluded — see the #2410 entry below.) Design goals (why this shape — mirrors the slime sibling's patch layer): * Default is upstream. Running this script is opt-in; the baseline image @@ -80,6 +81,13 @@ "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", ], @@ -87,18 +95,6 @@ # case's README points at for the full rollout path. "probe": ("examples/configs/recipes/llm/aws-efa-grpo-qwen3-30ba3b-2n8g-megatron.yaml", None), }, - { - "name": "NVIDIA-NeMo/RL#2411 — deps: bump deep_ep pin to the V2 merge commit", - "url": "https://github.com/NVIDIA-NeMo/RL/pull/2411", - "repo": "NVIDIA-NeMo/RL", - "root_arg": "nemo_rl_root", - "commits": [ - "711147f4401a3f532cbcdb6cb6b7e00e2023569e", - ], - # Metadata-only for this image (deep_ep is built from /opt/DeepEP, not - # from NeMo-RL's pin), but applied so the tree is self-consistent. - "probe": ("pyproject.toml", "b306af0"), - }, ] diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train-step.sh b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train-step.sh index 70b337677..2632fc8f8 100644 --- a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train-step.sh +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train-step.sh @@ -29,7 +29,7 @@ DRIVER="/opt/train_moe_step.py" 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 + DeepEP#612 + NeMo-RL#2410/#2411)." + echo "APPLY_DRAFT_ROLLOUT_PATCHES=1 (bakes Megatron-LM#4632 + DeepEP#612 + 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 diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/verify-image.sh b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/verify-image.sh index 6760ba928..5dcb23137 100644 --- a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/verify-image.sh +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/verify-image.sh @@ -57,7 +57,7 @@ docker run --rm --gpus all "${DEV_ARGS[@]}" -e HAVE_EFA_DEV="${HAVE_EFA_DEV}" "$ || { echo "FAIL: marker says patched but Megatron-LM#4632 not in the flex dispatcher"; 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 (4 draft PRs baked — full GRPO rollout-over-DeepEP path staged)" + echo " patched image (3 draft PRs baked — full GRPO rollout-over-DeepEP path staged)" else DEEP_EP_DIR=$(python3 -c "import deep_ep, pathlib; print(pathlib.Path(deep_ep.__file__).parent)") if grep -q EP_EFA_MAX_QPS "$DEEP_EP_DIR/buffers/elastic.py" 2>/dev/null; then diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/requirements.txt b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/requirements.txt index 4f910d018..868fa1bdc 100644 --- a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/requirements.txt +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/requirements.txt @@ -4,17 +4,26 @@ # 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 the pinned SHA): +# 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 same specifiers the pinned SHA +# NeMo-RL's [project.dependencies] set, at the same specifiers that SHA # declares, MINUS the entries the base image already provides or that conflict # with this test case's substrate: # - torch, triton, torchvision : NGC-baked (the ABI anchor; never reinstall) # - setuptools, pip, ninja : present in the base -# - nvidia-nvshmem-cu12 : cu12 wheel on a cu13 base, and unused here — -# this image's deep_ep is the NCCL-GIN backend -# (no NVSHMEM linked at all) +# 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 @@ -43,5 +52,23 @@ 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 + +# decord: imported 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; not in NeMo-RL's own +# [project.dependencies] at 46be4e8 (it's an extra), so it must be listed here. +decord + +# 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/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/setup_nemo_rl_deepep_efa.sh b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/setup_nemo_rl_deepep_efa.sh index fdbcbbc0a..a9625e505 100644 --- a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/setup_nemo_rl_deepep_efa.sh +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/setup_nemo_rl_deepep_efa.sh @@ -19,7 +19,7 @@ # Runs inside the Docker build (no GPU needed). set -euo pipefail -PHASE="${1:?usage: setup_nemo_rl_deepep_efa.sh {ofi|deepep}}" +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}" @@ -78,15 +78,24 @@ build_deepep() { # 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. - SO_COUNT=$(python3 - <<'PY' -import glob, importlib.util, os, sys -spec = importlib.util.find_spec("deep_ep") -if spec is None or not spec.submodule_search_locations: - sys.exit("deep_ep not installed") -root = list(spec.submodule_search_locations)[0] -print(len(glob.glob(os.path.join(os.path.dirname(root), "deep_ep*", "**", "*.so"), recursive=True) - + glob.glob(os.path.join(root, "**", "*.so"), recursive=True))) + # 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; } From 33ee8e7e51ac98628ebb7a3062f18ccee25067ce Mon Sep 17 00:00:00 2001 From: Anton Alexander Date: Wed, 26 Aug 2026 06:19:14 +0000 Subject: [PATCH 03/17] =?UTF-8?q?fix(nemo-rl/deepep-efa):=20train-step=20M?= =?UTF-8?q?oE=20gate=20=E2=80=94=20sync=20replicated=20params=20(DDP-free)?= =?UTF-8?q?,=20+=20shell=20bug?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Megatron-core MoE train-step gate (recipe/train_moe_step.py) now PASSES at full 8-GPU scale on a live p5en (8xH200, EFA): loss 9.43 -> 6.33 -> 4.17 -> 3.43 over real optimizer steps with the alltoall dispatcher (NCCL all-to-all over EFA), cross-rank loss agreement to ~1e-5. Two correctness fixes, both measured. train_moe_step.py -- gate 2 (cross-rank loss agreement) false-failed on a fully correct image + transport because the DDP-free loop never synced the REPLICATED (non-expert) params: - init: megatron seeds most replicated params identically via its RNG tracker, but router.weight and position_embeddings init under a per-rank-varying RNG context -- MEASURED cross-rank divergence 1.52e-1 / 1.64e-1 at step 0, before any update (diag all-reduce MAX-MIN, split expert vs replicated by param name; qkv/proj/layernorm/word_embeddings were already 0.0). Fix: broadcast the non-expert params from rank 0 after construction (what DDP does at build). - per step: the old docstring claimed "same batch => grad all-reduce is a no-op". FALSE for replicated params under EP -- each rank backprops through DIFFERENT local experts, so the shared router gets different per-rank grads and re-diverges every step (spread grew 5.7e-4 -> 8.9e-3 over 3 steps with init-broadcast alone). Fix: all-reduce(AVG) replicated-param grads each step. Result: spread FLAT ~1e-5 across 5 steps -> the gate is a real transport test at any step count, not a rubber stamp. Expert params/grads are deliberately left per-rank distinct -- that asymmetry IS expert parallelism. Also: fp32 params throughout (not bf16). The `local` layer spec emits fp32 LayerNorm activations; with bf16 params the router's te_general_gemm sees a bf16-weight x fp32-input GEMM that this NGC base's TE cuBLASLt build rejects ("unsupported value or parameter") regardless of moe_router_dtype. fp32 makes every GEMM (fp32,fp32,fp32) and keeps the all-to-all transport identical (NCCL/DeepEP all-to-all is dtype-agnostic; DeepEP requires fp32 probs). train-step.sh -- the usage string `${1:?usage: ... {leader|worker} ...}` has a literal `}` inside the `${1:?word}` expansion, which terminates the parameter expansion early; $ROLE became the literal "leader [node-rank]}" and the leader invocation aborted with "unrecognized role". Removed the interior brace from the message. (bash-parse-only fix; no behavior change on the happy path other than making the leader path actually run.) Verified: single-node 8xH200 over EFA on cgk p5en, 2026-08-26. Signed-off-by: Anton Alexander --- .../deepep-v2-efa/recipe/train-step.sh | 2 +- .../deepep-v2-efa/recipe/train_moe_step.py | 81 +++++++++++++++++-- 2 files changed, 75 insertions(+), 8 deletions(-) diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train-step.sh b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train-step.sh index 2632fc8f8..6ff86241f 100644 --- a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train-step.sh +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train-step.sh @@ -12,7 +12,7 @@ # 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]}" +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 diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train_moe_step.py b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train_moe_step.py index 5fffa6c60..2d34aa21d 100644 --- a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train_moe_step.py +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train_moe_step.py @@ -25,10 +25,22 @@ (Megatron-LM#4632); train-step.sh refuses it on an unpatched image rather than failing with a distant import error. -Every rank sees the same synthetic batch, so data-parallel gradient -all-reduce would be a mathematical no-op — the loop deliberately runs without -a DDP wrapper to stay independent of megatron's DDP config API, and the -cross-rank loss-agreement gate (2) is what makes that sound. +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 os import sys @@ -92,9 +104,32 @@ def main(): moe_enable_deepep=(DISPATCHER == "flex"), expert_model_parallel_size=ep_size, add_bias_linear=False, - params_dtype=torch.bfloat16, - pipeline_dtype=torch.bfloat16, - bf16=True, + # 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( @@ -106,6 +141,20 @@ def main(): 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) @@ -117,13 +166,31 @@ def main(): ).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 = [] 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) From 4052250f5d12015af159be65d44ad97a16a41f98 Mon Sep 17 00:00:00 2001 From: Anton Alexander Date: Wed, 26 Aug 2026 17:31:45 +0000 Subject: [PATCH 04/17] =?UTF-8?q?fix(nemo-rl/deepep-efa):=20security=20?= =?UTF-8?q?=E2=80=94=20.dockerignore=20+=20data-prep=20pod=20hardening?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses KeitaW review threads on PR #1242: - Add .dockerignore beside nemo-rl.Dockerfile (env_vars, *.log, __pycache__/). .gitignore does not apply to a Docker build context, so the filled-in env_vars (HF_TOKEN + AWS account id) was uploaded to the daemon / any remote builder on every `docker build .`. No broad COPY exists so nothing reached an image layer; this closes the context-upload exposure. (thread PRRT…M3nJ) - data-prep-pod.yaml: * HF_TOKEN plaintext env -> secretKeyRef from the same hf-token Secret raycluster.yaml already consumes (a literal "" could not work anyway, and a plaintext value would sit in etcd for the sleep-infinity pod's life). (thread PRRT…M3nY) * pin python:3.12-slim -> python:3.12.14-slim (never a floating tag — the "never latest" rule env_vars.example states). (thread PRRT…M3nU) * add the head's nvidia.com/gpu.present NotIn ["true"] nodeAffinity so this CPU-only sleep-infinity pod cannot park on and hold a GPU node; pin the header's `pip install "huggingface_hub[cli]==1.28.0"`; add automountServiceAccountToken:false (nothing here talks to the API server). (thread PRRT…M3nh) Signed-off-by: Anton Alexander --- .../nemo-rl/deepep-v2-efa/.dockerignore | 11 ++++++ .../kubernetes/data-prep-pod.yaml | 34 +++++++++++++++++-- 2 files changed, 42 insertions(+), 3 deletions(-) create mode 100644 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/.dockerignore diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/.dockerignore b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/.dockerignore new file mode 100644 index 000000000..4bc3096b9 --- /dev/null +++ b/3.test_cases/pytorch/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/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/kubernetes/data-prep-pod.yaml b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/kubernetes/data-prep-pod.yaml index 2b5b42821..37858a82c 100644 --- a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/kubernetes/data-prep-pod.yaml +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/kubernetes/data-prep-pod.yaml @@ -3,10 +3,17 @@ # 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]" && \ +# # 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: @@ -15,13 +22,34 @@ metadata: 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 - image: python:3.12-slim + # 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 - value: "" # Set before applying, or pass via kubectl set env + valueFrom: + secretKeyRef: + name: hf-token + key: HF_TOKEN resources: requests: cpu: "4" From 5df9552e5214be281ffc838e7db273b408e4faae Mon Sep 17 00:00:00 2001 From: Anton Alexander Date: Wed, 26 Aug 2026 17:35:23 +0000 Subject: [PATCH 05/17] =?UTF-8?q?fix(nemo-rl/deepep-efa):=20raycluster=20m?= =?UTF-8?q?anifest=20=E2=80=94=20pod=20env,=20FSx=20default,=20SA=20token?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses KeitaW review threads on PR #1242: - Put the recipe-gate shape in the worker pod env: NNODES (the name the scripts read — NOT NUM_NODES, which only raycluster's replicas consume) + GPUS_PER_NODE + EP_EXPERTS/EP_TOPK/EP_HIDDEN/EP_TOKENS/EP_NUM_SMS/EP_NUM_QPS. A client-side `source env_vars` does not cross `kubectl exec`; without these the gates ran on hardcoded script defaults. (thread PRRT…M3YG) - Default /fsx to emptyDir on both head and workers (PVC form left documented inline for the full GRPO swap). The recipe gates touch no shared storage, so the advertised cheap entry point no longer sits Pending on an unbound claim. (thread PRRT…M3YL) - Add automountServiceAccountToken:false to both pod specs — nothing here talks to the API server, and the workers are privileged root containers running model/training code. (thread PRRT…M3nj) - Drop OFI_NCCL_GIN_MAX_REQUESTS from the worker env: no such parameter exists in the pinned aws-ofi-nccl (the GIN params are gin_cq_process_max_iter, gin_gdaki, gin_strong_signal, gdrcopy_forced_pcie_copy). It was inert and contradicted trap 2's 128-slot ring. (thread PRRT…M3R0) Signed-off-by: Anton Alexander --- .../deepep-v2-efa/kubernetes/raycluster.yaml | 40 ++++++++++++++++--- 1 file changed, 35 insertions(+), 5 deletions(-) diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/kubernetes/raycluster.yaml b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/kubernetes/raycluster.yaml index 724d164b9..291b475c5 100644 --- a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/kubernetes/raycluster.yaml +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/kubernetes/raycluster.yaml @@ -47,6 +47,9 @@ spec: 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} @@ -83,9 +86,14 @@ spec: - 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 - persistentVolumeClaim: - claimName: ${FSX_CLAIM} + emptyDir: {} - name: dshm emptyDir: medium: Memory @@ -115,6 +123,10 @@ spec: 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: @@ -133,13 +145,27 @@ spec: 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_GIN_MAX_REQUESTS, value: "512" } - { 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" } @@ -187,9 +213,13 @@ spec: - 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 - persistentVolumeClaim: - claimName: ${FSX_CLAIM} + emptyDir: {} - name: dshm emptyDir: medium: Memory From c455bb18c2c88bc1a70c89e2ddaf67d8dd75c56d Mon Sep 17 00:00:00 2001 From: Anton Alexander Date: Wed, 26 Aug 2026 17:37:59 +0000 Subject: [PATCH 06/17] =?UTF-8?q?fix(nemo-rl/deepep-efa):=20recipe=20probe?= =?UTF-8?q?s=20=E2=80=94=20per-token=20scoring,=20all-step=20gate,=20teard?= =?UTF-8?q?own?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses KeitaW review threads on PR #1242: probe_rollout.py: - Score combine per token against its OWN scale (row_err/row_scale, 1e-3 floor) instead of one global max — a global normaliser averages away the badly-wrong low-amplitude rows a partial dispatch / dropped QP channel produces, exactly what the probe exists to catch. (thread PRRT…M3fS) - deep_ep.topk_idx_t direct access instead of getattr(...,torch.int64): the attribute is exported unconditionally at the pin, so the default was unreachable and would silently mask a future rename. (thread PRRT…M3fx) - Restrict the local-mapped detector to RANK>0: on rank 0 the local experts ARE 0..num_local-1, so it fired every normal run and printed a misleading note. (thread PRRT…M3fs) train_moe_step.py: - Track max_spread across ALL steps for gate 2 (was reading the final step only), so an early transport fault that later clears still fails; assert STEPS>=2. (thread PRRT…M3fb) both: - Pass timeout=10min to init_process_group and destroy the group in a finally so a raised rank fails fast instead of holding peers for the 30-min NCCL default (the worker runs detached under kubectl exec). (thread PRRT…M3f3) Signed-off-by: Anton Alexander --- .../deepep-v2-efa/recipe/probe_rollout.py | 37 ++++++++++++++++--- .../deepep-v2-efa/recipe/train_moe_step.py | 29 +++++++++++++-- 2 files changed, 56 insertions(+), 10 deletions(-) diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/probe_rollout.py b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/probe_rollout.py index feef8dc32..4d30b4779 100644 --- a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/probe_rollout.py +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/probe_rollout.py @@ -23,6 +23,7 @@ (EP_NUM_QPS=2 is the value the #612 evidence validated on p5en), which keeps the unpatched baseline probeable; on a patched image the auto-sizers also work. """ +import datetime import os import sys import traceback @@ -48,7 +49,13 @@ def log(msg): def main(): torch.cuda.set_device(LOCAL) - dist.init_process_group("nccl", rank=RANK, world_size=WORLD) + # 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})" @@ -77,7 +84,10 @@ def main(): # 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. - idx_dtype = getattr(deep_ep, "topk_idx_t", torch.int64) + # 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) @@ -101,7 +111,10 @@ def main(): num_local = E // WORLD lo, hi = RANK * num_local, RANK * num_local + num_local sl = recv_topk_idx.to(torch.int64) - if sl.numel() > 0 and int(sl.max()) < num_local and E > num_local: + # 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) @@ -121,7 +134,13 @@ def main(): 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) - relmax = ((got - expect).abs().max() / expect.abs().max().clamp_min(1e-6)).item() + # 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'}") @@ -139,9 +158,15 @@ def main(): if __name__ == "__main__": + rc = 1 try: - sys.exit(main()) + rc = main() except Exception: traceback.print_exc() print(f"[rank{RANK}] PROBE-EXC", flush=True) - sys.exit(1) + 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/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train_moe_step.py b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train_moe_step.py index 2d34aa21d..b8e5afd62 100644 --- a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train_moe_step.py +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train_moe_step.py @@ -42,6 +42,7 @@ 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 @@ -70,8 +71,18 @@ def log(msg): def main(): torch.cuda.set_device(LOCAL) - dist.init_process_group("nccl", rank=RANK, world_size=WORLD) + # 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 @@ -173,6 +184,7 @@ def main(): 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. @@ -199,11 +211,14 @@ def main(): 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] - agree = spread < 1e-2 + # 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}") @@ -219,9 +234,15 @@ def main(): if __name__ == "__main__": + rc = 1 try: - sys.exit(main()) + rc = main() except Exception: traceback.print_exc() print(f"[rank{RANK}] TRAIN-STEP-EXC", flush=True) - sys.exit(1) + 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) From 9e841d8a24fbcd15ffe5606502585b0a2e1f1042 Mon Sep 17 00:00:00 2001 From: Anton Alexander Date: Wed, 26 Aug 2026 17:43:44 +0000 Subject: [PATCH 07/17] =?UTF-8?q?fix(nemo-rl/deepep-efa):=20recipe=20shell?= =?UTF-8?q?s=20=E2=80=94=20EFA=20counter=20families,=20timeout,=20train-st?= =?UTF-8?q?ep=20transport=20gate?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - run-rollout-probe.sh / train-step.sh: sum tx_bytes+send_bytes+rdma_write_bytes+ rdma_read_resp_bytes for the EFA-TX assert, not tx_bytes alone. Under FI_EFA_USE_DEVICE_RDMA=1 the bytes land on rdma_write_bytes, so a tx_bytes-only check false-FAILs a healthy RDMA run (KeitaW review). - both: wrap torchrun in 'timeout 900' so a wedged rendezvous returns a verdict instead of holding the nodes; pairs with the python init_process_group timeout. - train-step.sh: port the probe's three-part transport gate (rc + >=1MiB EFA-TX delta + efa provider banner) instead of PASS on rc==0 alone; default NCCL_DEBUG to INFO so the provider banner is available as proof. Signed-off-by: Anton Alexander --- .../deepep-v2-efa/recipe/run-rollout-probe.sh | 12 ++++- .../deepep-v2-efa/recipe/train-step.sh | 49 ++++++++++++++++--- 2 files changed, 53 insertions(+), 8 deletions(-) diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/run-rollout-probe.sh b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/run-rollout-probe.sh index fd5c5437e..af68e357a 100644 --- a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/run-rollout-probe.sh +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/run-rollout-probe.sh @@ -41,9 +41,14 @@ echo "===== DeepEP-V2 rollout probe: role=$ROLE node_rank=$NODE_RANK nnodes=$NNO # 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; do + 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" @@ -53,7 +58,10 @@ 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. -torchrun --nnodes="$NNODES" --nproc-per-node="$GPUS_PER_NODE" --node-rank="$NODE_RANK" \ +# 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]} diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train-step.sh b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train-step.sh index 6ff86241f..643c788d1 100644 --- a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train-step.sh +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train-step.sh @@ -41,19 +41,56 @@ 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:-WARN}" +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 -torchrun --nnodes="$NNODES" --nproc-per-node="$GPUS_PER_NODE" --node-rank="$NODE_RANK" \ +# 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]} -if [ "$rc" -eq 0 ]; then - echo "TRAIN-STEP GATE PASS (node_rank=$NODE_RANK, dispatcher=$MOE_DISPATCHER)" +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 FAIL (node_rank=$NODE_RANK, torchrun rc=$rc) — see /tmp/train-step.$NODE_RANK.log" -exit 1 +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 From 55a3f6b3e7e5324d0637e4422d9df4100ea3e513 Mon Sep 17 00:00:00 2001 From: Anton Alexander Date: Wed, 26 Aug 2026 18:02:52 +0000 Subject: [PATCH 08/17] nemo-rl/deepep-v2-efa: per-change patch probes + loader-resolved NCCL check MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses KeitaW review on the draft-PR patch layer and the image gate. patches/apply_nemo_rl_patches.py: - Replace the single per-PR content probe with a per-CHANGE probe LIST (one needle per independently-required commit). A PR is skipped as already-present only when EVERY probe passes; a partially-merged upstream tree (some probes pass, some don't) is NOT skipped — its commits apply and a colliding hunk fails --check loud, instead of silently trusting a single needle hit. Post-assert re-checks the FULL list before the marker is written, so the marker can never record a patch state the tree lacks. - Drop the per-commit `git apply --reverse --check` short-circuit (it false-negatives on multi-commit PRs whose later commits touch the same hunks); the whole-PR probe gate now decides skip-vs-apply. - Needles derived from the real fetched .patch files: DeepEP#612 = EP_EFA_MAX_QPS / EP_EFA_RDMA_GBS / kScaleoutUpdateInterval=16; Megatron#4632 = ElasticBuffer / deep_ep.utils.event / _handle_num_experts (commit 4 is a comment/URL cleanup, no probe by design). recipe/verify-image.sh: - Patch-marker check now probes per-change too (adds EP_EFA_RDMA_GBS in utils/envs.py and _handle_num_experts in fused_a2a.py), covering the two named false-positive cases where a single needle passes but a later required change is absent. - Resolve libnccl via the actual loader (dlopen + /proc/self/maps) instead of `ldconfig -p | head -1` (cache order). dlopen honors LD_LIBRARY_PATH before the cache — the same evidence deep_ep.check_nccl_so() acts on — so this catches a scrubbed LD_LIBRARY_PATH or LD_PRELOAD shadow the cache-order check reads past. Signed-off-by: Anton Alexander --- .../patches/apply_nemo_rl_patches.py | 113 +++++++++++++----- .../deepep-v2-efa/recipe/verify-image.sh | 30 ++++- 2 files changed, 109 insertions(+), 34 deletions(-) diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/patches/apply_nemo_rl_patches.py b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/patches/apply_nemo_rl_patches.py index bbd846d96..e700b8639 100644 --- a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/patches/apply_nemo_rl_patches.py +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/patches/apply_nemo_rl_patches.py @@ -20,15 +20,21 @@ * 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. Before applying anything, each PR's content probe is - checked: if it already passes — because the PR merged upstream and the - tree pin advanced past it — the whole PR is skipped and reported (a - per-commit reverse-check would false-negative on multi-commit PRs whose - later commits touch the same hunks). When a PR reports already-present, - delete its entry here. - * Post-asserted. After each PR a content probe (needle in file) confirms the - intended change is really in the tree, catching apply-succeeded-but-wrong- - tree mistakes. + * 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. @@ -48,8 +54,19 @@ # 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). -# `probe` = (relative path, needle) asserted present after the PR applies; -# needle None = the file's existence is the assertion. +# `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": "deepseek-ai/DeepEP#612 — aws-efa: QP cap, get_rdma_gbs fast path, scaleout interval", @@ -61,7 +78,12 @@ "922a1fa7c0cd3ef047c0919638a87f9a2360346b", # EFA fast path in get_rdma_gbs (SM auto-sizing) "28d1f7fb173f728be51632ce0026fea23243e350", # dispatch kScaleoutUpdateInterval 6 -> 16 ], - "probe": ("deep_ep/buffers/elastic.py", "EP_EFA_MAX_QPS"), + "probes": [ + ("deep_ep/buffers/elastic.py", "EP_EFA_MAX_QPS"), # commit 1 + ("deep_ep/utils/envs.py", "EP_EFA_RDMA_GBS"), # commit 2 + ("deep_ep/include/deep_ep/impls/hybrid_dispatch.cuh", + "kScaleoutUpdateInterval = 16"), # commit 3 + ], }, { "name": "NVIDIA/Megatron-LM#4632 — moe: DeepEP V2 ElasticBuffer support in the flex dispatcher", @@ -72,9 +94,13 @@ "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 + "8056d6d489c73a353d590b8079497fceda4f9aa7", # drop downstream repro URLs + dead conditional (cleanup — no probe) ], - "probe": ("megatron/core/transformer/moe/fused_a2a.py", "ElasticBuffer"), + "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", @@ -93,7 +119,9 @@ ], # The PR also ships the worked 2-node EFA GRPO recipe config this test # case's README points at for the full rollout path. - "probe": ("examples/configs/recipes/llm/aws-efa-grpo-qwen3-30ba3b-2n8g-megatron.yaml", None), + "probes": [ + ("examples/configs/recipes/llm/aws-efa-grpo-qwen3-30ba3b-2n8g-megatron.yaml", None), + ], }, ] @@ -112,15 +140,19 @@ def fetch_patch(repo: str, sha: str) -> bytes: def apply_commit(root: Path, repo: str, sha: str) -> str: - """Apply one pinned commit into the tree at `root`. Returns a status word.""" + """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: - # Already present? (PR merged upstream and the pin moved past it.) - if run(["git", "apply", "--reverse", "--check", str(patch_path)], root).returncode == 0: - return "already-present-upstream" check = run(["git", "apply", "--check", str(patch_path)], root) if check.returncode != 0: raise RuntimeError( @@ -144,6 +176,16 @@ def probe_ok(root: Path, probe: tuple[str, str | None]) -> bool: 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) @@ -165,23 +207,38 @@ def main() -> int: print(f"FATAL: {root} is not a git checkout — cannot apply {pset['name']}", file=sys.stderr) return 1 print(f"== {pset['name']} ==") - if probe_ok(root, pset["probe"]): - # PR-level neutralization: the content probe already passes, so the - # tree pin advanced past this PR's merge — delete its entry here. - print(f" already present in the tree (probe satisfied) — skipping; retire this entry") + 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}") - if not probe_ok(root, pset["probe"]): - rel, needle = pset["probe"] + # 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"{'missing file' if needle is None else f'needle {needle!r} absent in'} {root / rel}", - file=sys.stderr) + f"{'; '.join(describe_probe(p) for p in still_missing)} " + f"absent under {root}", file=sys.stderr) return 1 - print(f" post-assert OK ({pset['url']})") + 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). diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/verify-image.sh b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/verify-image.sh index 5dcb23137..79317ae8a 100644 --- a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/verify-image.sh +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/verify-image.sh @@ -28,9 +28,17 @@ docker run --rm --gpus all "${DEV_ARGS[@]}" -e HAVE_EFA_DEV="${HAVE_EFA_DEV}" "$ echo " (run this script on an EFA host for the live efa-direct fabric check)" fi - echo "== single GIN-capable libnccl wins (path + version string — the NGC base bakes its own NCCL in a DIFFERENT dir, so path matters) ==" - NCCL_SO=$(ldconfig -p | grep "libnccl.so.2 " | head -1 | awk "{print \$NF}") - echo "$NCCL_SO" | grep -q "/opt/nccl/build" || { echo "FAIL: baked libnccl shadows the GIN build ($NCCL_SO)"; exit 1; } + 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 ==" @@ -50,11 +58,21 @@ docker run --rm --gpus all "${DEV_ARGS[@]}" -e HAVE_EFA_DEV="${HAVE_EFA_DEV}" "$ echo "== patch-marker consistency (marker and trees/site-packages must agree) ==" if [ -f /opt/.draft-rollout-patches-applied ]; then + # 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. EP_EFA_MAX_QPS present but the get_rdma_gbs EFA fast path absent; + # ElasticBuffer present but the num_experts backward-dispatch fix absent). + # This mirrors the per-change probe list in patches/apply_nemo_rl_patches.py. DEEP_EP_DIR=$(python3 -c "import deep_ep, pathlib; print(pathlib.Path(deep_ep.__file__).parent)") grep -q EP_EFA_MAX_QPS "$DEEP_EP_DIR/buffers/elastic.py" \ - || { echo "FAIL: marker says patched but DeepEP#612 EFA cap not in installed deep_ep"; exit 1; } - grep -q ElasticBuffer /opt/Megatron-LM/megatron/core/transformer/moe/fused_a2a.py \ - || { echo "FAIL: marker says patched but Megatron-LM#4632 not in the flex dispatcher"; exit 1; } + || { echo "FAIL: marker says patched but DeepEP#612 QP cap (EP_EFA_MAX_QPS) not in installed deep_ep"; exit 1; } + grep -q EP_EFA_RDMA_GBS "$DEEP_EP_DIR/utils/envs.py" \ + || { echo "FAIL: marker says patched but DeepEP#612 get_rdma_gbs EFA fast path (EP_EFA_RDMA_GBS) not in installed deep_ep — partially-applied PR"; exit 1; } + 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 (3 draft PRs baked — full GRPO rollout-over-DeepEP path staged)" From 316c1275a32d6a3cd4fbeaa22be4596beb817550 Mon Sep 17 00:00:00 2001 From: Anton Alexander Date: Wed, 26 Aug 2026 18:03:09 +0000 Subject: [PATCH 09/17] nemo-rl/deepep-v2-efa: scope `|| true` to removal, reset OPAL_PREFIX, drop unused GIN knob MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses KeitaW review on the Dockerfile and env contract. nemo-rl.Dockerfile: - Brace-scope the `|| true` to the `apt-get remove` of the distro verbs packages ONLY. apt-get update + install stay fail-hard — a transient mirror miss must not silently ship an image missing cmake/autoconf/libtool and fail 100s of lines later inside the gdrcopy/aws-ofi-nccl build. - Set OPAL_PREFIX=/opt/amazon/openmpi after `rm -rf /opt/hpcx/ompi`. The NGC base points OPAL_PREFIX at the deleted HPC-X Open MPI; reset it to the EFA installer's Open MPI so a dangling prefix can't trip an mpirun from the image. Drop OFI_NCCL_GIN_MAX_REQUESTS=512 everywhere it appeared (Dockerfile runtime ENV, env_vars.example, recipe/run-rollout-probe.sh, recipe/train-step.sh): it is not a knob the measured NCCL-GIN substrate set, and carrying an unverified value in the shipped transport contract is exactly the kind of guessed pin the review flagged. The remaining GIN vars (NCCL_GIN_TYPE/ENABLE, OFI_NCCL_GIN_GDAKI) are the measured ones. env_vars.example: - Add a "HOW THESE REACH THE GATES" note: sourcing sets vars in the local shell only and does NOT cross `kubectl exec`; `envsubst` substitutes NUM_NODES / the EP_* shape into the pod env: block, and the manifest bridges NUM_NODES->NNODES. - Correct the transport-contract header: a `kubectl exec` shell inherits these from the container env (baked + manifest), NOT from a client-side `source`; the local export line matters only for a non-k8s `docker run` on an EFA host. Signed-off-by: Anton Alexander --- .../nemo-rl/deepep-v2-efa/env_vars.example | 16 +++++++++++++--- .../nemo-rl/deepep-v2-efa/nemo-rl.Dockerfile | 15 ++++++++++++--- .../deepep-v2-efa/recipe/run-rollout-probe.sh | 2 +- .../nemo-rl/deepep-v2-efa/recipe/train-step.sh | 2 +- 4 files changed, 27 insertions(+), 8 deletions(-) diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/env_vars.example b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/env_vars.example index 72fdee08e..d8c34d661 100644 --- a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/env_vars.example +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/env_vars.example @@ -39,6 +39,14 @@ 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 @@ -62,8 +70,11 @@ 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; exported -# here too so ad-hoc shells (kubectl exec) carry the same 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 @@ -71,7 +82,6 @@ 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_GIN_MAX_REQUESTS=512 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 diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/nemo-rl.Dockerfile b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/nemo-rl.Dockerfile index c21d611d7..83c1b2c70 100644 --- a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/nemo-rl.Dockerfile +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/nemo-rl.Dockerfile @@ -81,12 +81,22 @@ SHELL ["/bin/bash", "-c"] # 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 + && { 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 @@ -212,7 +222,6 @@ ENV FI_PROVIDER=efa \ NCCL_GIN_TYPE=2 \ NCCL_GIN_ENABLE=1 \ OFI_NCCL_GIN_GDAKI=0 \ - OFI_NCCL_GIN_MAX_REQUESTS=512 \ OFI_NCCL_PROTOCOL=RDMA \ NCCL_NVLS_ENABLE=0 \ NCCL_DEBUG=WARN \ diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/run-rollout-probe.sh b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/run-rollout-probe.sh index af68e357a..69272c1a8 100644 --- a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/run-rollout-probe.sh +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/run-rollout-probe.sh @@ -23,7 +23,7 @@ PROBE="/opt/probe_rollout.py" # ---- 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 OFI_NCCL_GIN_MAX_REQUESTS=512 +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 diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train-step.sh b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train-step.sh index 643c788d1..b1829f546 100644 --- a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train-step.sh +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train-step.sh @@ -36,7 +36,7 @@ if [ "$MOE_DISPATCHER" = "flex" ] && [ ! -f /opt/.draft-rollout-patches-applied 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 OFI_NCCL_GIN_MAX_REQUESTS=512 +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 From 72d66adc46f77eb34d8bb72c2078e73f3ed31d21 Mon Sep 17 00:00:00 2001 From: Anton Alexander Date: Wed, 26 Aug 2026 18:23:23 +0000 Subject: [PATCH 10/17] docs(nemo-rl/deepep-v2-efa): correct provenance + substrate claims from PR #1242 review MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Truthfulness fixes for KeitaW's inline review — no measurement is invented; the edits make wording match what the build actually does and what was measured. README (folder): - NCCL row: state the NGC base bakes an OLDER NCCL line (2.29.x), so the source v2.30.4-1 build + 00-nccl-gin.conf ld.so.conf entry is LOAD-BEARING (makes the GIN-capable copy win the loader search), not a no-op 'same line' override. - aws-ofi-nccl row: describe 9c44d34 + PR#1351 as the GIN CPU-proxy lineage this folder standardises on; note the SHAs POSTDATE Wave-28 (not that run's pins). - trap 3: verify-image.sh resolves libnccl.so.2 via dlopen + /proc/self/maps, not ldconfig -p cache order. - runtime env table: drop OFI_NCCL_GIN_MAX_REQUESTS=512 row (the code no longer sets it — was removed from Dockerfile/scripts/env_vars; table was stale). - Quick Start: push the -draftprs flavor too (distinct tag); flex train-step must set MOE_DISPATCHER=flex on BOTH nodes (one torchrun job — ranks disagree and hang otherwise). - benchmarks refs: micro-benchmarks EP benchmark runs EFA-GDA GIN (NCCL_GIN_TYPE=5), NOT this folder's CPU-proxy (NCCL_GIN_TYPE=2) — drop the 'same substrate' claim in two places. README (parent nemo-rl/): same EFA-GDA-vs-CPU-proxy correction on the micro-benchmarks cross-link. setup_nemo_rl_deepep_efa.sh header: fix the 'links NO NVSHMEM' contradiction — deep_ep _C.so IS build-linked against NVSHMEM (Dockerfile Layer 7b makes the pip nvshmem win the loader search); NVSHMEM is just not the run-time transport on the V2 NCCL-GIN path. Retarget the no-vendor-sync contrast at the repo's canonical V2 builder setup_deepep_gin.sh (only the V1 setup_deepep_efa.sh copies are synced). Signed-off-by: Anton Alexander --- 3.test_cases/pytorch/nemo-rl/README.md | 4 +-- .../pytorch/nemo-rl/deepep-v2-efa/README.md | 28 +++++++++++++------ .../deepep-v2-efa/setup_nemo_rl_deepep_efa.sh | 18 ++++++++++-- 3 files changed, 36 insertions(+), 14 deletions(-) diff --git a/3.test_cases/pytorch/nemo-rl/README.md b/3.test_cases/pytorch/nemo-rl/README.md index afc030ea8..bf834c522 100644 --- a/3.test_cases/pytorch/nemo-rl/README.md +++ b/3.test_cases/pytorch/nemo-rl/README.md @@ -18,6 +18,6 @@ expert-parallel MoE all-to-all. For RL post-training with a different stack (SLIME + SGLang) on the same HyperPod-EKS Ray-cluster pattern, see [`3.test_cases/pytorch/slime`](../slime). For kernel-level expert-parallelism -dispatch/combine benchmarks over EFA — including a DeepEP V2 benchmark on the same NCCL-GIN -substrate this test case uses — see +dispatch/combine benchmarks over EFA — including a DeepEP V2 benchmark on the EFA-GDA NCCL-GIN +backend (this test case runs the CPU-proxy one) — see [`micro-benchmarks/expert-parallelism`](../../../micro-benchmarks/expert-parallelism). diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/README.md b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/README.md index 53df830f3..86bab0c99 100644 --- a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/README.md +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/README.md @@ -46,8 +46,8 @@ stock `alltoall` dispatcher (`train-step.sh`). Never read a build-gate as an E2E |---|---|---| | Base image | `nvcr.io/nvidia/pytorch:26.02-py3` | 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. | | EFA installer | 1.48.0 | The userspace of the measured NCCL-GIN substrate. Bumping is a re-measure event. | -| NCCL | `v2.30.4-1` (source build) | GIN device API generation (`nccl_device.h` asserted at build). Same NCCL line the NGC base bakes — the ld.so.conf override introduces no drift; source-built so DeepEP has one controlled header/lib root. | -| aws-ofi-nccl | commit `9c44d34` + PR#1351 head `c2e773d` | The GIN CPU-proxy plugin pins of the measured substrate (same as the TensorRT-LLM NcclEP sibling). Immutable SHAs — `refs/pull/N/head` is a moving ref. | +| NCCL | `v2.30.4-1` (source build) | GIN device API generation (`nccl_device.h` asserted at build). 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 | `01dc3aa` (upstream `deepseek-ai/DeepEP` main) | Carries EPv2 (ElasticBuffer + NCCL backend) and is the **exact base of draft PR #612**, so the opt-in layer applies `--check`-clean. Stock upstream — NOT a fork. | | Megatron-LM | `19deef67` (main) | The **exact base of draft PR #4632** (ElasticBuffer in the flex dispatcher). | @@ -85,7 +85,8 @@ PR states below are as of 2026-08-25 — check them before relying on this table 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 path** `libnccl.so.2` resolves to *and* its version string. + `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, @@ -113,7 +114,6 @@ PR states below are as of 2026-08-25 — check them before relying on this table |---|---| | `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 | -| `OFI_NCCL_GIN_MAX_REQUESTS=512` | GIN request-ring depth of the measured substrate | | `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 | @@ -158,6 +158,8 @@ 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 @@ -201,8 +203,13 @@ 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: -# MOE_DISPATCHER=flex /opt/train-step.sh leader ${W0_IP} +# 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) @@ -230,8 +237,10 @@ Treat results as your own measurement — this folder publishes none for this pa 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 on the same NCCL-GIN substrate see - [`micro-benchmarks/expert-parallelism`](../../../../micro-benchmarks/expert-parallelism). + 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 @@ -264,4 +273,5 @@ deepep-v2-efa/ mirrors its shape), [`sglang/dsr1-deepep-efa`](../../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 (NCCL GIN) benchmark + 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/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/setup_nemo_rl_deepep_efa.sh b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/setup_nemo_rl_deepep_efa.sh index a9625e505..d974a5531 100644 --- a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/setup_nemo_rl_deepep_efa.sh +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/setup_nemo_rl_deepep_efa.sh @@ -12,9 +12,21 @@ # compiled against the GIN-capable NCCL under $NCCL_HOME. Run AFTER # the optional draft-PR patch layer so patched kernels compile in. # -# Distinct from the NVSHMEM-path setup_deepep_efa.sh (vendor-synced via -# .github/workflows/deepep-vendor-sync.yml) — this script builds the NCCL-GIN -# V2 backend, links NO NVSHMEM, and is intentionally NOT vendor-synced. +# 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 stock-upstream +# deepseek-ai/DeepEP tree the Dockerfile pins (@01dc3aa, the base of draft PR +# #612) rather than gin's amazon-contributing default. +# 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 From af2bf8f443ea9725e7848c2616e3caf612c8cc41 Mon Sep 17 00:00:00 2001 From: Anton Alexander Date: Wed, 26 Aug 2026 18:26:00 +0000 Subject: [PATCH 11/17] chmod +x nemo-rl/deepep-v2-efa shell entrypoints (PR #1242 review) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit KeitaW: the scripts ship mode 100644, so README §3's `recipe/verify-image.sh ${FULL_IMAGE}` fails with Permission denied straight from a fresh clone. Set the executable bit on the three shell entrypoints the docs invoke directly (verify-image.sh, run-rollout-probe.sh, train-step.sh) plus setup_nemo_rl_deepep_efa.sh, via `git update-index --chmod=+x` so the mode lands in the tree. The two .py files stay 100644 — they are run via torchrun/python, never executed directly. Signed-off-by: Anton Alexander --- .../pytorch/nemo-rl/deepep-v2-efa/recipe/run-rollout-probe.sh | 0 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train-step.sh | 0 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/verify-image.sh | 0 .../pytorch/nemo-rl/deepep-v2-efa/setup_nemo_rl_deepep_efa.sh | 0 4 files changed, 0 insertions(+), 0 deletions(-) mode change 100644 => 100755 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/run-rollout-probe.sh mode change 100644 => 100755 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train-step.sh mode change 100644 => 100755 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/verify-image.sh mode change 100644 => 100755 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/setup_nemo_rl_deepep_efa.sh diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/run-rollout-probe.sh b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/run-rollout-probe.sh old mode 100644 new mode 100755 diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train-step.sh b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train-step.sh old mode 100644 new mode 100755 diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/verify-image.sh b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/verify-image.sh old mode 100644 new mode 100755 diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/setup_nemo_rl_deepep_efa.sh b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/setup_nemo_rl_deepep_efa.sh old mode 100644 new mode 100755 From 9bbea803e861da6f4ba00249168b0e5319d9cd0a Mon Sep 17 00:00:00 2001 From: Anton Alexander Date: Wed, 26 Aug 2026 18:37:18 +0000 Subject: [PATCH 12/17] build(nemo-rl/deepep-v2-efa): digest-pin NGC base, commit-pin NCCL, document plugin/fork trade-offs (PR #1242 review) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address KeitaW review threads on pin provenance and honesty: - NGC base: pin @sha256:bbc2b67e… alongside the readable :26.02-py3 tag. It is the ABI anchor the whole image builds around (torch 2.11/CUDA 13 + baked TE/apex/flash-attn), so a silent tag re-push is the most consequential drift possible here. Override NGC_PYTORCH_BASE to move it (a re-measure event). - NCCL: pin the commit 1933fdd6 that v2.30.4-1 resolves to, not the bare tag — held to the same moving-ref standard as the gdrcopy/DeepEP/Megatron SHA pins. Correct the stale "same NCCL line the NGC base bakes / no drift" comment: the base bakes an OLDER 2.29.x line, so the 00-nccl-gin.conf ld.so.conf entry is LOAD-BEARING (it makes this GIN-capable source build win the loader search). - aws-ofi-nccl #1351: note it is CLOSED-UNMERGED, so unlike the draft PRs in patches/ it does NOT self-neutralize — it is a PERMANENT baseline carry on every build. Clarify the "baseline has zero dependence on unmerged PRs" line refers to the four draft PRs in the opt-in layer, not this plugin pin. - DeepEP fork trade-off: document honestly that amazon-contributing/DeepEP fixes trap-2 structurally (sysfs get_rdma_gbs, restructured QP allocator) while baseline stays stock 01dc3aa on purpose (the #612 --check-clean base + the no-fork convention); the knob clamps are the portable equivalent, and fork adoption is a deliberate future re-pin + re-measure. README pins table updated to match (base digest, NCCL commit). Signed-off-by: Anton Alexander --- .../pytorch/nemo-rl/deepep-v2-efa/README.md | 4 +- .../nemo-rl/deepep-v2-efa/nemo-rl.Dockerfile | 49 +++++++++++++++---- 2 files changed, 42 insertions(+), 11 deletions(-) diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/README.md b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/README.md index 86bab0c99..76a7602d3 100644 --- a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/README.md +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/README.md @@ -44,9 +44,9 @@ stock `alltoall` dispatcher (`train-step.sh`). Never read a build-gate as an E2E | Component | Pin | Why | |---|---|---| -| Base image | `nvcr.io/nvidia/pytorch:26.02-py3` | 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. | +| 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` (source build) | GIN device API generation (`nccl_device.h` asserted at build). 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. | +| 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 | `01dc3aa` (upstream `deepseek-ai/DeepEP` main) | Carries EPv2 (ElasticBuffer + NCCL backend) and is the **exact base of draft PR #612**, so the opt-in layer applies `--check`-clean. Stock upstream — NOT a fork. | diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/nemo-rl.Dockerfile b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/nemo-rl.Dockerfile index 83c1b2c70..3b2f8f59f 100644 --- a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/nemo-rl.Dockerfile +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/nemo-rl.Dockerfile @@ -23,7 +23,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. -ARG NGC_PYTORCH_BASE=nvcr.io/nvidia/pytorch:26.02-py3 +# 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 @@ -32,11 +36,16 @@ 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: the GIN device API generation the measured substrate ran -# (ships include/nccl_device.h — asserted below). Same NCCL line the NGC base -# bakes, so the ld.so.conf override below introduces no version drift; built -# from source so the DeepEP build has one controlled root of headers + libs. -ARG NCCL_VERSION=v2.30.4-1 +# 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 @@ -44,16 +53,38 @@ ARG EFA_INSTALLER_VERSION=1.48.0 # 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 pins -# of the measured NCCL-GIN substrate (same pins as the TensorRT-LLM NcclEP +# 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: upstream deepseek-ai/DeepEP main @01dc3aa — carries EPv2 (the # ElasticBuffer + NCCL backend, merged upstream in PR#605) and is the EXACT # base of draft PR deepseek-ai/DeepEP#612, so the opt-in layer applies -# --check-clean. NOT the amazon-contributing fork: baseline is stock upstream. +# --check-clean. NOT the amazon-contributing fork: baseline is stock upstream +# (the repo's NGC-from-scratch, no-fork-base convention — this folder mirrors +# the slime sibling's stock-upstream stance). +# TRADE-OFF, stated honestly: amazon-contributing/DeepEP main = 01dc3aa + ~8 AWS +# commits that fix the trap-2 failure modes STRUCTURALLY rather than by the knob +# clamps this folder documents — a sysfs-based get_rdma_gbs() (provider-agnostic, +# reads /sys/class/infiniband//ports/*/rate) and a restructured GIN QP +# allocator (more QPs on the unordered path than EP_EFA_MAX_QPS=2). Most of those +# commits are days old (2026-08-21/25). The baseline stays on stock 01dc3aa on +# purpose: it is the #612 --check-clean base AND keeps to the no-fork convention; +# the knob clamps (via #612 in the opt-in layer, or the inert ENV defaults on +# baseline) are the portable equivalent. Adopting the fork's structural fixes is +# a deliberate future re-pin + re-measure, tracked separately — not a silent +# freeze. See README traps 2 + the pins table. ARG DEEPEP_SHA=01dc3aaac82068020353dce2c302e38153c0bfaa # 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. From f30b92913882702667f5f9ca9c0e6d690db41370 Mon Sep 17 00:00:00 2001 From: Anton Alexander Date: Wed, 26 Aug 2026 18:49:24 +0000 Subject: [PATCH 13/17] docs(nemo-rl/deepep-v2-efa): make requirements.txt truthful vs NEMO_RL_SHA=46be4e8 (PR #1242 review) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit KeitaW M2_d: the header claimed the deps are carried "at the same specifiers the pinned SHA declares," but that was not literally true. Diffed live against NeMo-RL 46be4e8's pyproject [project.dependencies] and corrected the drift: - wandb: restore upstream's >=0.25.0 floor (was unpinned). - decord -> decord2: upstream declares decord2 (the maintained fork that keeps the `decord` import name); match the spelling exactly, and drop the false "it's an extra, not in deps" note. - Header reworded: deps are at the SHA's specifiers EXCEPT for documented inline deviations, MINUS base-provided/conflicting/unused-feature-path entries — and the omitted ones are now named with why (nccl4py/cuda-bindings = non-colocated refit, pybase64 = SGLang refit, soundfile = audio multimodal — none on this colocated Megatron-GRPO text recipe). - Explain why `pip check` is deliberately NOT a build gate: --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. The maintainer's diff table was computed vs cc75cad (#2411's base); this branch pins 46be4e8 (build-fix 6c14c2e2), where torch/ray/transformers/pillow/mlflow already match line-for-line — so only the wandb floor and decord2 spelling were real deviations. Signed-off-by: Anton Alexander --- .../nemo-rl/deepep-v2-efa/requirements.txt | 28 +++++++++++++------ 1 file changed, 19 insertions(+), 9 deletions(-) diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/requirements.txt b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/requirements.txt index 868fa1bdc..bcc40ccf9 100644 --- a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/requirements.txt +++ b/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/requirements.txt @@ -7,11 +7,20 @@ # 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 same specifiers that SHA -# declares, MINUS the entries the base image already provides or that conflict -# with this test case's substrate: +# 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 @@ -28,7 +37,7 @@ colored==2.2.3 ray[default]==2.49.2 transformers==4.57.1 -wandb +wandb>=0.25.0 numpy datasets>=4.0.0 rich @@ -60,11 +69,12 @@ nvidia-nvshmem-cu13 # megatron.core (PYTHONPATH tree) import-time dependency not in NeMo-RL's set. einops -# decord: imported 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; not in NeMo-RL's own -# [project.dependencies] at 46be4e8 (it's an extra), so it must be listed here. -decord +# 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), From a2e017c4f83b5dc152014f87668c57c99c8dce6c Mon Sep 17 00:00:00 2001 From: Anton Alexander Date: Thu, 3 Sep 2026 03:37:02 +0000 Subject: [PATCH 14/17] refactor(training): migrate nemo-rl/deepep-v2-efa to examples/training (reorg #1119) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reorg #1119 de-numbered the top-level dirs and moved test cases from 3.test_cases/pytorch// to examples/{training,inference,use-cases}/. This PR was opened against the pre-reorg path, so its files landed under the now-deleted 3.test_cases/ tree. Adding files to a deleted directory does not conflict (git silently recreates it), so the PR read mergeable while being orphaned at a dead path — it needs a migration, not a rebase. Per AGENTS.md ("a training framework example goes to examples/training//" + "extend before create"), the sample nests as a variant subdirectory under the existing examples/training/nemo-rl/ example rather than as a new sibling — mirroring examples/training/megatron-bridge/'s populated-framework-dir-with-named-variants shape. - git mv the deepep-v2-efa subtree to examples/training/nemo-rl/deepep-v2-efa (history preserved; 14 of 15 files R100, README.md carries a one-line link fix). - git rm the stale 3.test_cases/pytorch/nemo-rl/README.md parent index (it only pointed at the pre-reorg ../slime sibling). - Add a "## Variants" pointer row to examples/training/nemo-rl/README.md so the nested example is discoverable from the framework dir (the merge's rename detection does not synthesize this — nemo-rl/README.md is a single-example README, not a per-framework case index). Case-README directory depth is unchanged (3.test_cases/pytorch/nemo-rl/ deepep-v2-efa and examples/training/nemo-rl/deepep-v2-efa are both 4 levels), so root-relative links resolve unchanged: ../../../../micro-benchmarks/expert-parallelism -> repo root (unchanged) Sibling links across the reorg's inference/training split needed one fix: ../../slime -> examples/training/slime (unchanged) ../../sglang/dsr1-deepep-efa -> BROKEN (sglang moved to inference/) ../../../inference/sglang/dsr1-deepep-efa -> examples/inference/sglang/... (fixed) Merged upstream/main rather than rebased to preserve the commit SHAs cited in the resolved review-thread replies. No functional change to the sample. Signed-off-by: Anton Alexander --- 3.test_cases/pytorch/nemo-rl/README.md | 23 ------------------- examples/training/nemo-rl/README.md | 9 ++++++++ .../nemo-rl/deepep-v2-efa/.dockerignore | 0 .../nemo-rl/deepep-v2-efa/.gitignore | 0 .../training}/nemo-rl/deepep-v2-efa/README.md | 4 ++-- .../nemo-rl/deepep-v2-efa/env_vars.example | 0 .../kubernetes/data-prep-pod.yaml | 0 .../deepep-v2-efa/kubernetes/raycluster.yaml | 0 .../nemo-rl/deepep-v2-efa/nemo-rl.Dockerfile | 0 .../patches/apply_nemo_rl_patches.py | 0 .../deepep-v2-efa/recipe/probe_rollout.py | 0 .../deepep-v2-efa/recipe/run-rollout-probe.sh | 0 .../deepep-v2-efa/recipe/train-step.sh | 0 .../deepep-v2-efa/recipe/train_moe_step.py | 0 .../deepep-v2-efa/recipe/verify-image.sh | 0 .../nemo-rl/deepep-v2-efa/requirements.txt | 0 .../deepep-v2-efa/setup_nemo_rl_deepep_efa.sh | 0 17 files changed, 11 insertions(+), 25 deletions(-) delete mode 100644 3.test_cases/pytorch/nemo-rl/README.md rename {3.test_cases/pytorch => examples/training}/nemo-rl/deepep-v2-efa/.dockerignore (100%) rename {3.test_cases/pytorch => examples/training}/nemo-rl/deepep-v2-efa/.gitignore (100%) rename {3.test_cases/pytorch => examples/training}/nemo-rl/deepep-v2-efa/README.md (99%) rename {3.test_cases/pytorch => examples/training}/nemo-rl/deepep-v2-efa/env_vars.example (100%) rename {3.test_cases/pytorch => examples/training}/nemo-rl/deepep-v2-efa/kubernetes/data-prep-pod.yaml (100%) rename {3.test_cases/pytorch => examples/training}/nemo-rl/deepep-v2-efa/kubernetes/raycluster.yaml (100%) rename {3.test_cases/pytorch => examples/training}/nemo-rl/deepep-v2-efa/nemo-rl.Dockerfile (100%) rename {3.test_cases/pytorch => examples/training}/nemo-rl/deepep-v2-efa/patches/apply_nemo_rl_patches.py (100%) rename {3.test_cases/pytorch => examples/training}/nemo-rl/deepep-v2-efa/recipe/probe_rollout.py (100%) rename {3.test_cases/pytorch => examples/training}/nemo-rl/deepep-v2-efa/recipe/run-rollout-probe.sh (100%) rename {3.test_cases/pytorch => examples/training}/nemo-rl/deepep-v2-efa/recipe/train-step.sh (100%) rename {3.test_cases/pytorch => examples/training}/nemo-rl/deepep-v2-efa/recipe/train_moe_step.py (100%) rename {3.test_cases/pytorch => examples/training}/nemo-rl/deepep-v2-efa/recipe/verify-image.sh (100%) rename {3.test_cases/pytorch => examples/training}/nemo-rl/deepep-v2-efa/requirements.txt (100%) rename {3.test_cases/pytorch => examples/training}/nemo-rl/deepep-v2-efa/setup_nemo_rl_deepep_efa.sh (100%) diff --git a/3.test_cases/pytorch/nemo-rl/README.md b/3.test_cases/pytorch/nemo-rl/README.md deleted file mode 100644 index bf834c522..000000000 --- a/3.test_cases/pytorch/nemo-rl/README.md +++ /dev/null @@ -1,23 +0,0 @@ - - -# NeMo-RL test cases - -[NeMo-RL](https://github.com/NVIDIA-NeMo/RL) is NVIDIA's scalable post-training library (GRPO, -DPO, SFT) for models from 1 GPU to thousands, with Megatron-core and DTensor training backends. -The samples in this directory deploy NeMo-RL on AWS with high-performance EFA networking and -expert-parallel MoE all-to-all. - -## Available test cases - -| Test case | Orchestrator | Description | -| --- | --- | --- | -| [`deepep-v2-efa`](./deepep-v2-efa) | Kubernetes (Ray cluster, 2-node) | GRPO post-training with MoE expert-parallel dispatch/combine via **DeepEP V2's NCCL backend** (`ElasticBuffer`) over **AWS EFA**, using aws-ofi-nccl with **GIN** (GPU-Initiated Networking) CPU-proxy. Image built NGC-from-scratch from public sources; the recipe runs image build → static verify → cross-node rollout probe → N-step Megatron MoE training gate. Mechanism chain measured on 2× p5.48xlarge (H100); this folder's image assembly is build-staged with an opt-in draft-PR layer for the full rollout path (honest measured-vs-staged breakdown in its README). | - -For RL post-training with a different stack (SLIME + SGLang) on the same HyperPod-EKS Ray-cluster -pattern, see [`3.test_cases/pytorch/slime`](../slime). For kernel-level expert-parallelism -dispatch/combine benchmarks over EFA — including a DeepEP V2 benchmark on the EFA-GDA NCCL-GIN -backend (this test case runs the CPU-proxy one) — see -[`micro-benchmarks/expert-parallelism`](../../../micro-benchmarks/expert-parallelism). 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/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/.dockerignore b/examples/training/nemo-rl/deepep-v2-efa/.dockerignore similarity index 100% rename from 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/.dockerignore rename to examples/training/nemo-rl/deepep-v2-efa/.dockerignore diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/.gitignore b/examples/training/nemo-rl/deepep-v2-efa/.gitignore similarity index 100% rename from 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/.gitignore rename to examples/training/nemo-rl/deepep-v2-efa/.gitignore diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/README.md b/examples/training/nemo-rl/deepep-v2-efa/README.md similarity index 99% rename from 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/README.md rename to examples/training/nemo-rl/deepep-v2-efa/README.md index 76a7602d3..ded473eae 100644 --- a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/README.md +++ b/examples/training/nemo-rl/deepep-v2-efa/README.md @@ -270,8 +270,8 @@ deepep-v2-efa/ - [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), [`sglang/dsr1-deepep-efa`](../../sglang/dsr1-deepep-efa) (the NVSHMEM-path - DeepEP serving sample) + 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/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/env_vars.example b/examples/training/nemo-rl/deepep-v2-efa/env_vars.example similarity index 100% rename from 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/env_vars.example rename to examples/training/nemo-rl/deepep-v2-efa/env_vars.example diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/kubernetes/data-prep-pod.yaml b/examples/training/nemo-rl/deepep-v2-efa/kubernetes/data-prep-pod.yaml similarity index 100% rename from 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/kubernetes/data-prep-pod.yaml rename to examples/training/nemo-rl/deepep-v2-efa/kubernetes/data-prep-pod.yaml diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/kubernetes/raycluster.yaml b/examples/training/nemo-rl/deepep-v2-efa/kubernetes/raycluster.yaml similarity index 100% rename from 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/kubernetes/raycluster.yaml rename to examples/training/nemo-rl/deepep-v2-efa/kubernetes/raycluster.yaml diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/nemo-rl.Dockerfile b/examples/training/nemo-rl/deepep-v2-efa/nemo-rl.Dockerfile similarity index 100% rename from 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/nemo-rl.Dockerfile rename to examples/training/nemo-rl/deepep-v2-efa/nemo-rl.Dockerfile diff --git a/3.test_cases/pytorch/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 similarity index 100% rename from 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/patches/apply_nemo_rl_patches.py rename to examples/training/nemo-rl/deepep-v2-efa/patches/apply_nemo_rl_patches.py diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/probe_rollout.py b/examples/training/nemo-rl/deepep-v2-efa/recipe/probe_rollout.py similarity index 100% rename from 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/probe_rollout.py rename to examples/training/nemo-rl/deepep-v2-efa/recipe/probe_rollout.py diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/run-rollout-probe.sh b/examples/training/nemo-rl/deepep-v2-efa/recipe/run-rollout-probe.sh similarity index 100% rename from 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/run-rollout-probe.sh rename to examples/training/nemo-rl/deepep-v2-efa/recipe/run-rollout-probe.sh diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train-step.sh b/examples/training/nemo-rl/deepep-v2-efa/recipe/train-step.sh similarity index 100% rename from 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train-step.sh rename to examples/training/nemo-rl/deepep-v2-efa/recipe/train-step.sh diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train_moe_step.py b/examples/training/nemo-rl/deepep-v2-efa/recipe/train_moe_step.py similarity index 100% rename from 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/train_moe_step.py rename to examples/training/nemo-rl/deepep-v2-efa/recipe/train_moe_step.py diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/verify-image.sh b/examples/training/nemo-rl/deepep-v2-efa/recipe/verify-image.sh similarity index 100% rename from 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/recipe/verify-image.sh rename to examples/training/nemo-rl/deepep-v2-efa/recipe/verify-image.sh diff --git a/3.test_cases/pytorch/nemo-rl/deepep-v2-efa/requirements.txt b/examples/training/nemo-rl/deepep-v2-efa/requirements.txt similarity index 100% rename from 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/requirements.txt rename to examples/training/nemo-rl/deepep-v2-efa/requirements.txt diff --git a/3.test_cases/pytorch/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 similarity index 100% rename from 3.test_cases/pytorch/nemo-rl/deepep-v2-efa/setup_nemo_rl_deepep_efa.sh rename to examples/training/nemo-rl/deepep-v2-efa/setup_nemo_rl_deepep_efa.sh From fe0250a920553d3b08379459ed20549169c0c2a5 Mon Sep 17 00:00:00 2001 From: Anton Alexander Date: Thu, 3 Sep 2026 08:13:47 +0000 Subject: [PATCH 15/17] refactor(training/nemo-rl/deepep-v2-efa): pin DeepEP source to amazon-contributing fork; drop superseded #612 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Repoint the DeepEP source from stock deepseek-ai/DeepEP @ 01dc3aa plus a draft deepseek-ai/DeepEP#612 opt-in patch to the amazon-contributing/DeepEP fork, pinned at the immutable HEAD 97d8f9bcc1be31e9036db2ab591ef9b9f4e38619. This mirrors the sibling examples/inference/vllm/deepep-v2-gdaki-efa repoint and answers the review note that pinning the pre-fix upstream fork-point forfeits the AWS EFA fixes: the fork is the AWS EPv2/NCCL-GIN tree and carries the EFA delta IN-CODE, so no #612 patch is applied on any flavor. The fork carries both correctness halves of what was draft #612 structurally: - 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 the allocated QP count to _C.{min,max}_unordered_gin_qps). It does NOT carry #612's third commit (a kScaleoutUpdateInterval 6->16 latency micro-opt); the fork keeps =6. No gate in this example depends on that value, so the two correctness fixes are what "supersedes #612" means here. Because the fork carries those fixes on every flavor, this collapses the baseline/opt-in distinction for DeepEP only: - the #612 entry drops out of the opt-in draft-PR layer entirely (the layer now bakes only Megatron-LM#4632 + NeMo-RL#2410); - the two dead env vars EP_EFA_MAX_QPS / EP_EFA_RDMA_GBS (the old #612 patch knobs, zero readers on the fork) are removed from env_vars.example and kubernetes/raycluster.yaml, replaced with an absence-note; and - the verify-image / setup fail-loud gate is rewritten to fork discriminators — it now asserts _get_sysfs_rdma_gbs and unordered_gin_qps are present in the installed deep_ep tree on EVERY flavor, not gated on the draft-PR marker. EP_NUM_QPS=2 is a different variable (the probe's explicit num_allocated_qps) and stays: it survives the fork clamp unchanged (max(2, min(2, max)) == 2). The Dockerfile pins the full 40-char SHA as DEEPEP_SHA; setup and verify-image fail-loud if the clone lacks either fix-half. This is the same fork the repo's micro-benchmarks/expert-parallelism/deepep-v2-benchmark/setup_deepep_gin.sh already clones. ## Test Results docs-and-packaging change (source-repoint + provenance). No functional change to the example's runtime behavior on either flavor; the fork already carried the fixes the removed patch supplied. markdownlint-cli2 clean on the changed README against the root .markdownlint.jsonc. Signed-off-by: Anton Alexander --- .../training/nemo-rl/deepep-v2-efa/README.md | 44 +++++++----- .../nemo-rl/deepep-v2-efa/env_vars.example | 29 ++++---- .../deepep-v2-efa/kubernetes/raycluster.yaml | 8 ++- .../nemo-rl/deepep-v2-efa/nemo-rl.Dockerfile | 69 +++++++++++-------- .../patches/apply_nemo_rl_patches.py | 27 ++------ .../deepep-v2-efa/recipe/probe_rollout.py | 20 +++--- .../deepep-v2-efa/recipe/run-rollout-probe.sh | 6 +- .../deepep-v2-efa/recipe/train-step.sh | 2 +- .../deepep-v2-efa/recipe/verify-image.sh | 41 ++++++----- .../deepep-v2-efa/setup_nemo_rl_deepep_efa.sh | 16 ++++- 10 files changed, 145 insertions(+), 117 deletions(-) diff --git a/examples/training/nemo-rl/deepep-v2-efa/README.md b/examples/training/nemo-rl/deepep-v2-efa/README.md index ded473eae..72f89c79e 100644 --- a/examples/training/nemo-rl/deepep-v2-efa/README.md +++ b/examples/training/nemo-rl/deepep-v2-efa/README.md @@ -14,7 +14,7 @@ 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) @@ -34,7 +34,7 @@ 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 3 draft upstream PRs** (opt-in image layer, default **OFF** — +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 (upstream-only)** image: 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 @@ -49,21 +49,24 @@ stock `alltoall` dispatcher (`train-step.sh`). Never read a build-gate as an E2E | 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 | `01dc3aa` (upstream `deepseek-ai/DeepEP` main) | Carries EPv2 (ElasticBuffer + NCCL backend) and is the **exact base of draft PR #612**, so the opt-in layer applies `--check`-clean. Stock upstream — NOT a fork. | +| 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 3 draft upstream PRs (opt-in layer, default OFF) +## 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 | |---|---|---| -| [deepseek-ai/DeepEP#612](https://github.com/deepseek-ai/DeepEP/pull/612) | open | EFA awareness: auto-QP capped at `EP_EFA_MAX_QPS` (aws-ofi-nccl's GIN request ring is 128 slots; upstream's auto formula overruns it → `CUDA_ERROR_LAUNCH_FAILED` at first dispatch), `get_rdma_gbs()` EFA fast path (SM auto-sizing), dispatch scaleout-interval tuning. | | [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. | @@ -76,12 +79,15 @@ PR states below are as of 2026-08-25 — check them before relying on this table 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. **Upstream DeepEP's SM/QP auto-sizers are EFA-blind** at the pinned SHA: auto-QP - (`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 probe passes `num_allocated_qps`/`num_sms`/`num_qps` **explicitly** (`EP_NUM_QPS=2` is the - value the #612 evidence validated on p5en) so the baseline image stays probeable; the opt-in - layer (#612) fixes the auto-sizers. +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 @@ -119,7 +125,6 @@ PR states below are as of 2026-08-25 — check them before relying on this table | `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 | -| `EP_EFA_MAX_QPS=2`, `EP_EFA_RDMA_GBS=25.0` | read only by the **patched** deep_ep (#612); inert on baseline | ## Hardware requirements @@ -153,7 +158,7 @@ aws ecr create-repository --repository-name ${IMAGE} --region ${AWS_REGION} || t # baseline (upstream-only) docker build -f nemo-rl.Dockerfile -t ${FULL_IMAGE} . -# opt-in flavor with the 3 draft PRs baked (use a DISTINCT tag — never overwrite the baseline) +# 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 . @@ -226,11 +231,12 @@ Treat results as your own measurement — this folder publishes none for this pa - **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.** Three upstream PRs, closed-unmerged upstream +- **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 - does not certify upstream's auto-sizing on EFA — that certification is exactly PR #612. + 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 @@ -244,7 +250,7 @@ Treat results as your own measurement — this folder publishes none for this pa ## File structure -``` +```text deepep-v2-efa/ ├── README.md <- you are here ├── nemo-rl.Dockerfile <- NGC-from-scratch image (baseline + opt-in draft-PR layer) @@ -252,7 +258,7 @@ deepep-v2-efa/ ├── 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 3 draft PRs, pinned SHAs, fail-loud, self-neutralizing +│ └── 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 @@ -266,7 +272,9 @@ deepep-v2-efa/ ## References -- [DeepEP](https://github.com/deepseek-ai/DeepEP) — EPv2 / ElasticBuffer (PR#605) +- [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 diff --git a/examples/training/nemo-rl/deepep-v2-efa/env_vars.example b/examples/training/nemo-rl/deepep-v2-efa/env_vars.example index d8c34d661..efb59d31f 100644 --- a/examples/training/nemo-rl/deepep-v2-efa/env_vars.example +++ b/examples/training/nemo-rl/deepep-v2-efa/env_vars.example @@ -22,11 +22,12 @@ export TAG="v1-20260825" export FULL_IMAGE="${REGISTRY}${IMAGE}:${TAG}" # ----- Opt-in draft-PR image flavor ----- -# 1 = bake the 3 draft upstream PRs (NeMo-RL#2410, Megatron-LM#4632, -# DeepEP#612) that the full GRPO rollout-over-DeepEP path needs; 0 = the -# upstream-only baseline. 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. +# 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) ----- @@ -61,11 +62,11 @@ export EP_EXPERTS=128 export EP_TOPK=8 export EP_HIDDEN=2048 export EP_TOKENS=128 -# Explicit SM/QP counts for the probe. WHY: upstream DeepEP's auto-sizers are -# EFA-blind at the pinned SHA (auto-QP overruns aws-ofi-nccl's 128-slot GIN -# ring; get_rdma_gbs() returns 0) - exactly what draft PR DeepEP#612 fixes. -# Explicit values keep the BASELINE image probeable; the patched image also -# works with the auto-sizers. +# 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 @@ -85,9 +86,11 @@ export OFI_NCCL_GIN_GDAKI=0 # GPU-initiated GIN is not the shipped path 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 -# Read only by the PATCHED deep_ep (DeepEP#612); inert on the baseline image. -export EP_EFA_MAX_QPS=2 -export EP_EFA_RDMA_GBS=25.0 +# 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 diff --git a/examples/training/nemo-rl/deepep-v2-efa/kubernetes/raycluster.yaml b/examples/training/nemo-rl/deepep-v2-efa/kubernetes/raycluster.yaml index 291b475c5..dc172ccf9 100644 --- a/examples/training/nemo-rl/deepep-v2-efa/kubernetes/raycluster.yaml +++ b/examples/training/nemo-rl/deepep-v2-efa/kubernetes/raycluster.yaml @@ -180,9 +180,11 @@ spec: - { name: NCCL_SOCKET_IFNAME, value: "^lo,docker,veth" } # ---- DeepEP V2 knobs ---- - { name: DEEP_EP_USE_V2_SHIM, value: "0" } - # Read only by the PATCHED deep_ep (DeepEP#612); inert on baseline. - - { name: EP_EFA_MAX_QPS, value: "2" } - - { name: EP_EFA_RDMA_GBS, value: "25.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 diff --git a/examples/training/nemo-rl/deepep-v2-efa/nemo-rl.Dockerfile b/examples/training/nemo-rl/deepep-v2-efa/nemo-rl.Dockerfile index 3b2f8f59f..9c2a36014 100644 --- a/examples/training/nemo-rl/deepep-v2-efa/nemo-rl.Dockerfile +++ b/examples/training/nemo-rl/deepep-v2-efa/nemo-rl.Dockerfile @@ -15,7 +15,7 @@ # -> 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 3 DRAFT upstream PRs (see the patches/ +# -> 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. @@ -68,24 +68,24 @@ ARG GDRCOPY_SHA=c91ad9f178e5fb729fc5b6dc62a77c3bb364d6c9 ARG AWS_OFI_NCCL_SHA=9c44d34476f90ddbf4a12d0ac4fc412d46bd8ab4 ARG AWS_OFI_NCCL_PR=1351 ARG AWS_OFI_NCCL_PR_SHA=c2e773dfb2c75b765b3415f8ffd1b47e7c239a7b -# DeepEP: upstream deepseek-ai/DeepEP main @01dc3aa — carries EPv2 (the -# ElasticBuffer + NCCL backend, merged upstream in PR#605) and is the EXACT -# base of draft PR deepseek-ai/DeepEP#612, so the opt-in layer applies -# --check-clean. NOT the amazon-contributing fork: baseline is stock upstream -# (the repo's NGC-from-scratch, no-fork-base convention — this folder mirrors -# the slime sibling's stock-upstream stance). -# TRADE-OFF, stated honestly: amazon-contributing/DeepEP main = 01dc3aa + ~8 AWS -# commits that fix the trap-2 failure modes STRUCTURALLY rather than by the knob -# clamps this folder documents — a sysfs-based get_rdma_gbs() (provider-agnostic, -# reads /sys/class/infiniband//ports/*/rate) and a restructured GIN QP -# allocator (more QPs on the unordered path than EP_EFA_MAX_QPS=2). Most of those -# commits are days old (2026-08-21/25). The baseline stays on stock 01dc3aa on -# purpose: it is the #612 --check-clean base AND keeps to the no-fork convention; -# the knob clamps (via #612 in the opt-in layer, or the inert ENV defaults on -# baseline) are the portable equivalent. Adopting the fork's structural fixes is -# a deliberate future re-pin + re-measure, tracked separately — not a silent -# freeze. See README traps 2 + the pins table. -ARG DEEPEP_SHA=01dc3aaac82068020353dce2c302e38153c0bfaa +# 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 @@ -178,10 +178,14 @@ RUN chmod +x /opt/setup_nemo_rl_deepep_efa.sh \ 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 upstream trees at their pinned SHAs ---------- -# Clones only here; the DeepEP BUILD is deferred to Layer 9 so the opt-in -# patch layer (Layer 8) can land its .cuh/.py edits BEFORE kernels compile. -RUN git clone https://github.com/deepseek-ai/DeepEP.git /opt/DeepEP \ +# ---- 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} \ @@ -217,10 +221,13 @@ RUN ln -sfn "$(python3 -c 'import nvidia.nvshmem; print(nvidia.nvshmem.__path__[ || { 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 3 draft upstream PRs ---------------- -# NVIDIA-NeMo/RL#2410, NVIDIA/Megatron-LM#4632, deepseek-ai/DeepEP#612 +# ---- 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. Commits are pinned inside patches/apply_nemo_rl_patches.py at +# 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). @@ -260,10 +267,12 @@ ENV FI_PROVIDER=efa \ RAY_memory_monitor_refresh_ms=0 \ TOKENIZERS_PARALLELISM=false \ DEEP_EP_USE_V2_SHIM=0 -# Read only by the PATCHED deep_ep (DeepEP#612 adds the EFA awareness); inert -# on the baseline image. Kept here so both flavors run with one manifest. -ENV EP_EFA_MAX_QPS=2 \ - EP_EFA_RDMA_GBS=25.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 index e700b8639..7597a7ee5 100644 --- 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 @@ -3,12 +3,14 @@ # 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 (deepseek-ai/DeepEP, +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 three upstream +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 fourth, -NVIDIA-NeMo/RL#2411, is intentionally excluded — see the #2410 entry below.) +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 @@ -68,23 +70,6 @@ # 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": "deepseek-ai/DeepEP#612 — aws-efa: QP cap, get_rdma_gbs fast path, scaleout interval", - "url": "https://github.com/deepseek-ai/DeepEP/pull/612", - "repo": "deepseek-ai/DeepEP", - "root_arg": "deepep_root", - "commits": [ - "4eddba396006b8e4aa1a3a9a505396020aba4ef7", # cap auto-QP at 2 on EFA (128-slot GIN ring) - "922a1fa7c0cd3ef047c0919638a87f9a2360346b", # EFA fast path in get_rdma_gbs (SM auto-sizing) - "28d1f7fb173f728be51632ce0026fea23243e350", # dispatch kScaleoutUpdateInterval 6 -> 16 - ], - "probes": [ - ("deep_ep/buffers/elastic.py", "EP_EFA_MAX_QPS"), # commit 1 - ("deep_ep/utils/envs.py", "EP_EFA_RDMA_GBS"), # commit 2 - ("deep_ep/include/deep_ep/impls/hybrid_dispatch.cuh", - "kScaleoutUpdateInterval = 16"), # commit 3 - ], - }, { "name": "NVIDIA/Megatron-LM#4632 — moe: DeepEP V2 ElasticBuffer support in the flex dispatcher", "url": "https://github.com/NVIDIA/Megatron-LM/pull/4632", 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 index 4d30b4779..48c9f7900 100644 --- a/examples/training/nemo-rl/deepep-v2-efa/recipe/probe_rollout.py +++ b/examples/training/nemo-rl/deepep-v2-efa/recipe/probe_rollout.py @@ -13,15 +13,17 @@ 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 draft PR -deepseek-ai/DeepEP#612 fixes) only fires at the first real cross-node dispatch. - -Baseline-image note: upstream DeepEP's SM/QP auto-sizers are EFA-blind at the -pinned SHA (auto-QP overruns aws-ofi-nccl's 128-slot GIN request ring; -`get_rdma_gbs()` reads 0 on EFA) — the exact gaps DeepEP#612 fixes. The probe -therefore passes `num_allocated_qps`/`num_sms`/`num_qps` EXPLICITLY -(EP_NUM_QPS=2 is the value the #612 evidence validated on p5en), which keeps -the unpatched baseline probeable; on a patched image the auto-sizers also work. +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 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 index 69272c1a8..5f643fbe7 100755 --- 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 @@ -31,8 +31,10 @@ 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: upstream DeepEP's auto-sizers are EFA-blind at the pinned SHA -# (DeepEP#612 is the fix; opt-in layer). 2 QPs is the value the #612 evidence validated. +# 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 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 index b1829f546..57e86e268 100755 --- a/examples/training/nemo-rl/deepep-v2-efa/recipe/train-step.sh +++ b/examples/training/nemo-rl/deepep-v2-efa/recipe/train-step.sh @@ -29,7 +29,7 @@ DRIVER="/opt/train_moe_step.py" 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 + DeepEP#612 + NeMo-RL#2410)." + 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 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 index 79317ae8a..564d04eb3 100755 --- a/examples/training/nemo-rl/deepep-v2-efa/recipe/verify-image.sh +++ b/examples/training/nemo-rl/deepep-v2-efa/recipe/verify-image.sh @@ -50,24 +50,35 @@ docker run --rm --gpus all "${DEV_ARGS[@]}" -e HAVE_EFA_DEV="${HAVE_EFA_DEV}" "$ 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/site-packages must agree) ==" + 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 - # 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. EP_EFA_MAX_QPS present but the get_rdma_gbs EFA fast path absent; - # ElasticBuffer present but the num_experts backward-dispatch fix absent). - # This mirrors the per-change probe list in patches/apply_nemo_rl_patches.py. - DEEP_EP_DIR=$(python3 -c "import deep_ep, pathlib; print(pathlib.Path(deep_ep.__file__).parent)") - grep -q EP_EFA_MAX_QPS "$DEEP_EP_DIR/buffers/elastic.py" \ - || { echo "FAIL: marker says patched but DeepEP#612 QP cap (EP_EFA_MAX_QPS) not in installed deep_ep"; exit 1; } - grep -q EP_EFA_RDMA_GBS "$DEEP_EP_DIR/utils/envs.py" \ - || { echo "FAIL: marker says patched but DeepEP#612 get_rdma_gbs EFA fast path (EP_EFA_RDMA_GBS) not in installed deep_ep — partially-applied PR"; exit 1; } 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; } @@ -75,13 +86,9 @@ docker run --rm --gpus all "${DEV_ARGS[@]}" -e HAVE_EFA_DEV="${HAVE_EFA_DEV}" "$ || { 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 (3 draft PRs baked — full GRPO rollout-over-DeepEP path staged)" + echo " patched image (2 draft PRs baked — full GRPO rollout-over-DeepEP path staged)" else - DEEP_EP_DIR=$(python3 -c "import deep_ep, pathlib; print(pathlib.Path(deep_ep.__file__).parent)") - if grep -q EP_EFA_MAX_QPS "$DEEP_EP_DIR/buffers/elastic.py" 2>/dev/null; then - echo "FAIL: no marker but the DeepEP#612 patch IS present — ambiguous patch state"; exit 1 - fi - echo " unpatched baseline (upstream-only trees — probe/train-step run with explicit SM/QP counts)" + 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/setup_nemo_rl_deepep_efa.sh b/examples/training/nemo-rl/deepep-v2-efa/setup_nemo_rl_deepep_efa.sh index d974a5531..747be2e94 100755 --- 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 @@ -16,9 +16,11 @@ # 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 stock-upstream -# deepseek-ai/DeepEP tree the Dockerfile pins (@01dc3aa, the base of draft PR -# #612) rather than gin's amazon-contributing default. +# 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 @@ -82,6 +84,14 @@ build_deepep() { # 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}}" From 4dc63a6de8e7804f5de4789665a71350b9c34578 Mon Sep 17 00:00:00 2001 From: Anton Alexander Date: Sat, 5 Sep 2026 10:07:06 +0000 Subject: [PATCH 16/17] =?UTF-8?q?docs(nemo-rl/deepep-v2-efa):=20scope=20th?= =?UTF-8?q?e=20'baseline'=20framing=20=E2=80=94=20it=20is=20not=20literall?= =?UTF-8?q?y=20upstream-only?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit KeitaW review thread (nemo-rl.Dockerfile aws-ofi-nccl PR#1351 cherry-pick): the README headline called the default image 'baseline (upstream-only)', which reads as 'no unmerged-PR carry'. That is inaccurate. Decision: keep PR#1351 as a PERMANENT baseline carry (it is the load-bearing GIN CPU-proxy plugin lineage the whole substrate standardises on; gating it opt-in would be technically wrong) and fix the CLAIM instead. The default image carries, and always did (both disclosed in the Pins table): the standardized AWS GIN plugin lineage incl. closed-unmerged aws-ofi-nccl PR#1351, and the amazon-contributing/DeepEP fork. What 'baseline' actually means here is 'the 2 draft rollout PRs OFF' (Megatron-LM#4632 + NeMo-RL#2410) — those are the only opt-in toggle, and they DO self-neutralize once merged. #1351 does not; it is permanent. - README intro: 'baseline (upstream-only)' -> 'baseline' + explicit note that it carries the AWS GIN lineage (incl. #1351) and the DeepEP fork. - Build snippet comment: same clarification so the term can't recur. docs-only (claim-accuracy fix; no build/behavior change). Author-verify: gh api .../pulls/1242/commits --jq .[].author.login == dmvevents Signed-off-by: Anton Alexander --- examples/training/nemo-rl/deepep-v2-efa/README.md | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/examples/training/nemo-rl/deepep-v2-efa/README.md b/examples/training/nemo-rl/deepep-v2-efa/README.md index 72f89c79e..1bdc03c1f 100644 --- a/examples/training/nemo-rl/deepep-v2-efa/README.md +++ b/examples/training/nemo-rl/deepep-v2-efa/README.md @@ -35,7 +35,9 @@ 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 (upstream-only)** image: static +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. @@ -156,7 +158,7 @@ source env_vars 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 (upstream-only) +# 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 \ From 5bd1da87e79c3e1b47b32d39b43aec92c1c2edf1 Mon Sep 17 00:00:00 2001 From: Anton Alexander Date: Sat, 5 Sep 2026 10:30:27 +0000 Subject: [PATCH 17/17] =?UTF-8?q?docs(nemo-rl/deepep-v2-efa):=20env=5Fvars?= =?UTF-8?q?=20NVSHMEM=20note=20=E2=80=94=20build-time=20link=20vs=20run-ti?= =?UTF-8?q?me=20transport?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit env_vars.example said 'This image links NO NVSHMEM', which contradicts requirements.txt (deep_ep links NVSHMEM unconditionally at build time via upstream setup.py) and setup_nemo_rl_deepep_efa.sh (deep_ep's _C.so IS build-linked against NVSHMEM). Reword to match: NVSHMEM is a build-time link dependency; the V2 NCCL-GIN backend does not USE it at run time, so the values below stay inert. Closes KeitaW review thread on requirements.txt (cid 3856239817) residual. Signed-off-by: Anton Alexander --- .../training/nemo-rl/deepep-v2-efa/env_vars.example | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/examples/training/nemo-rl/deepep-v2-efa/env_vars.example b/examples/training/nemo-rl/deepep-v2-efa/env_vars.example index efb59d31f..56952a77d 100644 --- a/examples/training/nemo-rl/deepep-v2-efa/env_vars.example +++ b/examples/training/nemo-rl/deepep-v2-efa/env_vars.example @@ -96,11 +96,13 @@ export NCCL_NET_PLUGIN=/opt/aws-ofi-nccl/lib/libnccl-net-ofi.so 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 on this image - kept for the rebuild case) ----- -# This image links NO NVSHMEM (the V2 NCCL-GIN backend replaces it). 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. +# ----- 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