diff --git a/docker/rl/Dockerfile.nmp-rl-base b/docker/rl/Dockerfile.nmp-rl-base index 7d418c25f2..ff7d61eda2 100644 --- a/docker/rl/Dockerfile.nmp-rl-base +++ b/docker/rl/Dockerfile.nmp-rl-base @@ -5,8 +5,8 @@ # from source (rather than the prebuilt NGC nemo-rl image) lets us pin any RL commit via # NEMO_RL_REF; Gym follows as RL's submodule pin. # -# DPO trains on `fsdp`; GRPO trains on `fsdp` and generates on `vllm`. mcore, automodel, sglang -# and trtllm are deliberately not built. +# DPO trains on `fsdp`. GRPO (a future release) trains on `automodel` (DTensor V2) or `mcore` +# (Megatron) and generates on `vllm`. sglang and trtllm are deliberately not built. # # Two things drive most of what follows, both covered in README.md: # - Workers do not run in the base venv. The `uv sync --extra` calls only warm the uv cache; the @@ -101,9 +101,9 @@ ENV RAY_USAGE_STATS_ENABLED=0 \ NEMO_GYM_VENV_DIR=/opt/gym_venvs # ---- builder: uv sync RL + Gym with the extras we use ---- -# Only fsdp (DPO/GRPO training), vllm (GRPO generation), nemo_gym and modelopt are built. The mcore -# and automodel extras are deliberately EXCLUDED: alternative training backends the customizer does -# not use, and the only source of Transformer-Engine (the longest CUDA compile). +# fsdp (DPO training), automodel (GRPO training on DTensor V2, the only LoRA-capable DTensor +# backend), mcore (Megatron training + Megatron-native generation), vllm (generation), nemo_gym and +# modelopt are built. sglang and trtllm are not: vLLM covers generation for every config we compile. FROM base AS builder WORKDIR /opt/nemo-rl @@ -154,9 +154,9 @@ ENV NVTE_BUILD_MAX_JOBS=${NVTE_BUILD_MAX_JOBS} \ # layer's cache key depends on the dependency graph, not on every RL source file. See README.md, # "Layering for fast CI rebuilds". # -# TensorRT-LLM is excluded (we never build the trtllm extra). Megatron-Bridge must stay, even though -# the mcore extra is not built: `uv run` validates the lock, which needs metadata for every path -# source, so removing it fails the prefetch step below. +# TensorRT-LLM is excluded (we never build the trtllm extra). Megatron-Bridge must stay: it is the +# mcore extra's entry point, and its nested Megatron-LM submodule supplies megatron-core. Both are +# editable path sources, so `uv run` also needs their metadata to validate the lock. COPY --from=nemo-rl pyproject.toml uv.lock ./ COPY --from=nemo-rl nemo_rl/__init__.py nemo_rl/package_info.py ./nemo_rl/ COPY --from=nemo-rl --exclude=TensorRT-LLM-workspace --exclude=TensorRT-LLM-workspace/** \ @@ -168,12 +168,16 @@ COPY --from=nemo-rl research/ ./research/ ARG UV_SYNC_MODE=--frozen # Warm the uv CACHE, one extra at a time (they conflict). This is where the heavy compiles happen -# (mamba-ssm, causal-conv1d, deep_ep, deep_gemm); the wheels +# (mamba-ssm, causal-conv1d, deep_ep, deep_gemm, and Transformer-Engine via automodel); the wheels # land in the cache so the prefetch below materializes each venv without recompiling. # # `uv sync` is EXACT, so these do NOT accumulate - each --extra prunes the previous one and the base # venv ends up with default deps + groups only. The populated cache is the real product. # +# Conflicting extras resolve in separate forks but still share cache entries when a fork lands on the +# same pinned source, so mcore reuses the Transformer-Engine wheel automodel built. See README.md, +# "Why the `mcore` backend is built". +# # --no-install-project installs dependencies but not nemo-rl itself; the root package is added # editable in a cheap step after the full-source COPY. RUN <<"EOF" bash -exu @@ -184,7 +188,9 @@ uv venv --seed uv sync ${UV_SYNC_MODE} --no-install-project uv sync ${UV_SYNC_MODE} --extra vllm --no-install-project # GRPO generation (vllm 0.20 cu130 + flashinfer + deep_gemm/deep_ep) -uv sync ${UV_SYNC_MODE} --extra fsdp --no-install-project # DPO/GRPO policy training (flash-attn, mamba-ssm, causal-conv1d) +uv sync ${UV_SYNC_MODE} --extra fsdp --no-install-project # DPO policy training (flash-attn, mamba-ssm, causal-conv1d) +uv sync ${UV_SYNC_MODE} --extra automodel --no-install-project # GRPO policy training, DTensor V2 (+ LoRA); builds Transformer-Engine +uv sync ${UV_SYNC_MODE} --extra mcore --no-install-project # Megatron training + Megatron-native generation; reuses automodel's TE uv sync ${UV_SYNC_MODE} --extra modelopt --no-install-project # quantization / model-opt uv sync ${UV_SYNC_MODE} --extra nemo_gym --no-install-project # NeMo-Gym (uv workspace member) uv sync ${UV_SYNC_MODE} --all-groups --no-install-project # build/test/etc groups @@ -231,8 +237,13 @@ RUN uv pip install --python /opt/nemo_rl_venv/bin/python --no-deps -e /opt/nemo- # cache. See README.md, "Prefetching the per-worker venvs". # # Filters are SUBSTRING matches on actor FQNs, and one venv is built PER ACTOR, so actors sharing -# an extra still need their own filter. Six filters -> seven venvs: -# dtensor_policy_worker.DTensorPolicyWorker -> fsdp (DPO/GRPO training) +# an extra still need their own filter. Eight filters -> nine venvs: +# dtensor_policy_worker.DTensorPolicyWorker -> fsdp (DPO training) +# dtensor_policy_worker_v2.DTensorPolicyWorkerV2 -> automodel (GRPO training; the only LoRA-capable +# DTensor worker). The V1 filter does not match it - the FQN reads `dtensor_policy_worker_v2.`. +# megatron_policy_worker.MegatronPolicyWorker -> mcore (GRPO training on Megatron; also +# hosts Megatron-native generation, which runs in-process rather than as its own actor). +# Does not match modelopt's `megatron_quant_policy_worker.MegatronQuantPolicyWorker`. # vllm.vllm_worker -> vllm (GRPO generation) - matches both the # sync and async workers, which is wanted: NeMo-Gym forces async rollouts. # sync_rollout_actor.SyncRolloutActor -> vllm (GRPO rollout driver) @@ -241,12 +252,14 @@ RUN uv pip install --python /opt/nemo_rl_venv/bin/python --no-deps -e /opt/nemo- # broker_actor.SandboxEpisodeBrokerActor -> nemo_gym (per-episode sandbox broker) # Kept specific on purpose: a bare "vllm" would also match modelopt's vllm_quant_worker. # -# NOT prefetched: automodel, mcore, sglang, trtllm, modelopt-quant -# and the async-GRPO bookkeeping actors; those build on the node if a config selects them. +# NOT prefetched: sglang, trtllm, modelopt-quant, MegatronValueWorker / DTensorValueWorkerV2 (GRPO +# is critic-free) and the async-GRPO bookkeeping actors; those build on the node if selected. RUN <<"EOF" bash -exu export UV_LINK_MODE=symlink uv run nemo_rl/utils/prefetch_venvs.py \ dtensor_policy_worker.DTensorPolicyWorker \ + dtensor_policy_worker_v2.DTensorPolicyWorkerV2 \ + megatron_policy_worker.MegatronPolicyWorker \ vllm.vllm_worker \ sync_rollout_actor.SyncRolloutActor \ nemo_gym.NemoGym \ @@ -258,6 +271,9 @@ uv run nemo_rl/utils/prefetch_venvs.py \ # The vLLM-tier actors are listed too: the publish stage privatizes `vllm/` in whatever venvs exist, # so a filter that stops producing one would leave that worker sharing ray_executor.py again. for actor in \ + nemo_rl.models.policy.workers.dtensor_policy_worker.DTensorPolicyWorker \ + nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2 \ + nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorker \ nemo_rl.environments.nemo_gym.NemoGym \ nemo_rl.environments.sandbox.nemo_gym_actor.SandboxedGymActor \ nemo_rl.environments.sandbox.broker_actor.SandboxEpisodeBrokerActor \ @@ -344,6 +360,11 @@ FROM builder AS nmp-rl-base # their own dependencies, so RL's `setuptools>=80.10.2` override cannot reach them. The result is an # older setuptools left in the shipped cache that no venv uses. Runs here, after every venv exists, # and only drops archives nothing symlinks into. +# +# Still correct now that the mcore extra IS built: installing Megatron-Bridge does not change what +# its BUILD environment resolves, so the old setuptools remains cached and unreferenced. Deleting +# only unreferenced archives is what makes that safe - the mcore venv links the locked setuptools, +# so that archive is kept. RUN <<"EOF" bash -exu for d in "${UV_CACHE_DIR}"/archive-v0/*/; do ls "${d}"setuptools-*.dist-info >/dev/null 2>&1 || continue diff --git a/docker/rl/README.md b/docker/rl/README.md index 8f4f535998..da21b43ec7 100644 --- a/docker/rl/README.md +++ b/docker/rl/README.md @@ -1,8 +1,14 @@ # NeMo-RL training images (`nmp-rl-base`, `nmp-rl-training`) -GPU images for NeMo Platform's RL customization — **DPO and GRPO** — plus the +GPU images for NeMo Platform's RL customization — **DPO**, and **GRPO** (will be added in +future releases) — plus the **NeMo-Gym** environment runtime. A single training image serves both algorithms. +> **(will be added in future releases)** marks capability the image carries but the platform does +> not expose yet: DPO is +> the shipping algorithm, and the GRPO backend is still in review. The image builds the extras +> and prefetches the venvs GRPO needs so it works the day that lands. + The images are built **from source** on NVIDIA's `cuda-dl-base` (CUDA 13, Python 3.13), rather than layered on a prebuilt PyTorch container. This README records the decisions that shape the build so the non-obvious constraints don't have to be @@ -119,19 +125,52 @@ resolve the same interpreter as their parent. | Algorithm | Actor | Extra / venv | Notable contents | |---|---|---|---| -| DPO + GRPO | `DTensorPolicyWorker` (policy training) | **`fsdp`** | `flash-attn` (prebuilt multi-arch wheel), `mamba-ssm`, `causal-conv1d` | -| GRPO | `VllmGenerationWorker`, `SyncRolloutActor` | **`vllm`** | `vllm`, `deep_ep`, `deep_gemm`, `flashinfer` | -| GRPO (Gym) | `NemoGym` | **`nemo_gym`** | NeMo-Gym workspace member | +| DPO | `DTensorPolicyWorker` (policy training) | **`fsdp`** | `flash-attn` (prebuilt multi-arch wheel), `mamba-ssm`, `causal-conv1d` | +| GRPO (future release) | `DTensorPolicyWorkerV2` (policy training) | **`automodel`** | `nemo-automodel`, Transformer-Engine, `megatron-fsdp` | +| GRPO (future release) | `MegatronPolicyWorker` (policy training) | **`mcore`** | `megatron-bridge`, `megatron-core`, Transformer-Engine | +| GRPO (future release) | `VllmGenerationWorker`, `SyncRolloutActor` | **`vllm`** | `vllm`, `deep_ep`, `deep_gemm`, `flashinfer` | +| GRPO (Gym, future release) | `NemoGym` | **`nemo_gym`** | NeMo-Gym workspace member | + +DPO resolves to the **V1** DTensor worker (`dtensor_cfg._v2` defaults to false and the DPO +compiler does not set it), which maps to the `fsdp` extra. GRPO's compiler sets +`policy.dtensor_cfg._v2: true`, which selects **V2** and therefore the `automodel` extra. -DPO resolves to the **V1** DTensor worker (`dtensor_cfg._v2` defaults to false and the -customizer does not set it), which maps to the `fsdp` extra. +**LoRA requires V2 or Megatron.** DTensor V1 asserts `lora_cfg.enabled is False` +(`nemo_rl/models/policy/lm_policy.py`); the DTensor LoRA implementation lives in +`nemo_automodel.components._peft.lora`, and Megatron carries its own `peft` path. So +`automodel` and `mcore` are the parameter-efficient-training tiers in this image. -### Excluded on purpose: the `automodel` and `mcore` backends +GRPO is **critic-free** — no value model — so neither `DTensorValueWorkerV2` (automodel tier) +nor `MegatronValueWorker` (mcore tier) is prefetched. -We deliberately **do not build the `automodel` or `mcore` extras**. They are *alternative -training backends*, and the customizer uses neither for now. +### Why the `mcore` backend is built -**Why that matters: they are the only things that pull Transformer-Engine.** TE is a +`mcore` is a second GRPO **training** tier alongside `automodel`. It is what backs the +capabilities DTensor does not implement in NeMo-RL today: + +| Capability | DTensor (`fsdp` / `automodel`) | Megatron (`mcore`) | +|---|---|---| +| Pipeline parallelism | ❌ `lm_policy.py` reads `pp_size` only from `megatron_cfg` | ✅ | +| FP8 training / FP8 rollouts / FP8 KV-cache | ❌ | ✅ | +| NVFP4 quantization-aware RL | ❌ | ✅ (with `modelopt`) | +| Draft models / EAGLE3 speculative decoding | ❌ hard-gated in `lm_policy.py` | ✅ | +| Megatron-native generation (no refit conversion) | ❌ | ✅ | +| Largest models (`qwen3.5-397ba17b`, `glm5.1`, DeepSeek-V3) | ❌ tops out ~120B MoE | ✅ | +| MoE with expert parallelism, context parallel, LoRA, VLM | ✅ | ✅ | + +Upstream weights this way too: of NeMo-RL's seven recipes that reference `nemo_gym`, six run +on Megatron. + +Megatron-native generation needs no separate venv — it runs **in-process inside +`MegatronPolicyWorker`**, which is the point of it (training and inference share one weight +layout, so there is no refit conversion). That is why it has no registry entry of its own. + +`mcore` also requires `TORCH_CUDA_ARCH_LIST` to be set at **runtime**, not just at build +(`lm_policy.py` raises without it; mcore's inference unified-memory API calls a torch API +that reads it). The publish stage is `FROM builder`, so it inherits the builder's +`TORCH_CUDA_ARCH_LIST="9.0 10.0"` — the same value NeMo-RL's own `docker/Dockerfile` sets. + +**Transformer-Engine is shared between `automodel` and `mcore`.** TE is a training-time *transformer layer* library (fused attention / LayerNorm / GEMM kernels, fp8). Whether a backend needs it comes down to **who implements the transformer layer**: @@ -142,29 +181,31 @@ fp8). Whether a backend needs it comes down to **who implements the transformer | `mcore` | **Megatron-Core** builds its parallel layers on TE primitives | ✅ | | `vllm` | Its own hand-written **inference** kernels (paged attention, fused MoE) | ❌ inference engine, not a training-layer library | -So excluding these two removes the single longest CUDA compile (TE) from the build. - -What is excluded, precisely: - -- `nemo-automodel` (`automodel` extra) — the HF-native DTensor **V2** backend - (`DTensorPolicyWorkerV2`), reachable only via `policy.dtensor_cfg._v2: true`. -- `megatron-core` / `megatron-bridge` (`mcore` extra) — `MegatronPolicyWorker`, reachable - only via `policy.megatron_cfg.enabled: true`, which the customizer explicitly disables. - -**Effect on build time and image size:** the saving is almost entirely **build time** — the -Transformer-Engine compile is the longest step in the image, and it is now skipped. Image -size barely changes, for three reasons: the warmup syncs never persisted those packages -anyway (`uv sync` is exact, so each `--extra` is pruned by the next — see "Consequences" -above), the uv cache is a build-time mount that never enters the image, and we do not -prefetch `automodel` / `mcore` worker venvs. Prefetched worker venvs are the only place an -extra contributes to image size, so image size would only grow if we *also* prefetched -those two. The Automodel and Megatron-Bridge **source trees still ship** regardless: they -are RL git submodules referenced by `uv.lock`. - -**Re-enabling either** is a config change plus two lines in `Dockerfile.nmp-rl-base` — add -the extra to the warmup sync and add the worker FQN to the prefetch filters. Transformer-Engine -is built from source when that happens: `uv.lock` pins it as a git source, so `uv sync --frozen` -compiles it and caches the wheel, and every later venv reuses that. +TE is the single longest CUDA compile in the image, and `automodel` and `mcore` **share one +build of it**. RL's `[tool.uv] override-dependencies` collapses every TE requirement — +`automodel`'s declared `v2.14.1` and `megatron-bridge[te]`'s own rev — onto +`git+…/TransformerEngine.git@release_v2.15`, so `uv.lock` holds exactly one resolved +`transformer-engine 2.15.0+42b8400`. Conflicting extras resolve in separate forks, but uv's +built-wheel cache is keyed on the resolved source identity (git URL + commit + platform +tags), not on which extra requested it. The `mcore` sync therefore reuses the wheel the +`automodel` sync built. `deep_ep`, `mamba-ssm` and `causal-conv1d` are pinned identically +across the extras and are reused the same way. + +TE is built from source, once per image build: `uv.lock` pins it as a **git source**, so +`uv sync --frozen` compiles it and caches the resulting wheel, and every later venv reuses +that. A prebuilt wheel would not be picked up without changing how the lock sources TE. + +**What `mcore` actually adds:** `megatron-bridge` and `megatron-core` are **editable path +sources** from submodules already on disk, so installing them is near-free. The rest is +wheel-only downloads (`flashinfer-*==0.6.8.post1`, `nvshmem4py-cu13`, `cupy-cuda13x`) plus +`nvidia-modelopt` from git. Build-time cost is small; the real cost is one more prefetched +worker venv, which is the only place an extra contributes to image size (`uv sync` is exact, +so the warmup syncs do not accumulate). + +Megatron-Bridge pins `setuptools<80.0.0` for its own build, so uv builds its metadata in a +PEP 517 environment against an old setuptools and leaves it in the shipped cache. RL's +`setuptools>=80.10.2` override cannot reach a build environment, so the publish stage drops +cached setuptools archives that no venv symlinks into. ### NeMo-Gym environments: a second, separate venv layer @@ -215,11 +256,18 @@ IMAGE (built once) RUNTIME (each `uv sync --extra` is pruned by (math / code / VLM environments) the next; extras never persist here) -/opt/ray_venvs/ per-ACTOR venvs, prefetched at build - ├─ …DTensorPolicyWorker [fsdp] ────────> policy training (DPO + GRPO) - ├─ …VllmGenerationWorker [vllm] ────────> generation (GRPO only) - ├─ …SyncRolloutActor [vllm] ────────> rollout driver (GRPO only) - └─ …NemoGym [nemo_gym] ────────> Gym actor (GRPO only) +/opt/ray_venvs/ per-ACTOR venvs, prefetched at build (nine) + ├─ …DTensorPolicyWorker [fsdp] ────> policy training (DPO) + ├─ …DTensorPolicyWorkerV2 [automodel] ────> policy training (GRPO, DTensor V2) + ├─ …MegatronPolicyWorker [mcore] ────> policy training (GRPO, Megatron; + │ also hosts Megatron-native generation) + ├─ …VllmGenerationWorker [vllm] ────> generation (GRPO, sync) + ├─ …VllmAsyncGenerationWorker [vllm] ────> generation (GRPO, async — Gym + │ forces async rollouts) + ├─ …SyncRolloutActor [vllm] ────> rollout driver (GRPO, sync path) + ├─ …NemoGym [nemo_gym] ────> Gym actor (mode A, colocated) + ├─ …SandboxedGymActor [nemo_gym] ────> Gym proxy actor (mode B, sandboxed) + └─ …SandboxEpisodeBrokerActor [nemo_gym] ────> per-episode sandbox broker (mode B) /opt/gym_venvs/ per-ENVIRONMENT venvs — empty in the shipped image ├─ built-in Gym envs ── created at RUNTIME (prefetch off; opt in via @@ -236,11 +284,13 @@ driver (base venv) (init_reference_model=True — no extra worker, no extra venv) ``` -**GRPO job** — training + generation + (optionally) Gym: +**GRPO job (future release), DTensor V2 path** — training + generation + (optionally) Gym. +This is the path the customizer's GRPO compiler emits (`policy.dtensor_cfg._v2: true`, +`megatron_cfg.enabled: false`); see below for the Megatron alternative. ```text driver (base venv) - ├─ Policy ─ DTensorPolicyWorker ×N → /opt/ray_venvs/…DTensorPolicyWorker [fsdp] + ├─ Policy ─ DTensorPolicyWorkerV2 ×N → /opt/ray_venvs/…DTensorPolicyWorkerV2 [automodel] ├─ Generation ─ VllmGenerationWorker ×N → /opt/ray_venvs/…VllmGenerationWorker [vllm] │ └─ deep_ep / deep_gemm → Hopper+ only ├─ Rollout ─ SyncRolloutActor → /opt/ray_venvs/…SyncRolloutActor [vllm] @@ -255,6 +305,22 @@ driver (base venv) wheels-v1 avoids PyPI, native-v1 needs egress) ``` +**GRPO job, Megatron path** — selected by `policy.megatron_cfg.enabled: true`. Only the policy +tier differs; generation, rollout and Gym are unchanged from the diagram above: + +```text +driver (base venv) + ├─ Policy ─ MegatronPolicyWorker ×N → /opt/ray_venvs/…MegatronPolicyWorker [mcore] + │ └─ Megatron-native generation runs IN-PROCESS here when + │ policy.generation.mcore_generation_config is set — training and + │ inference share one weight layout, so there is no refit conversion + │ and no separate generation venv + └─ (generation / rollout / Gym as above) +``` + +Megatron and DTensor are mutually exclusive — `lm_policy.py` raises if both `megatron_cfg` +and `dtensor_cfg` are enabled. + ## Hardware: which GPUs run what The image bundles both **multi-arch prebuilt wheels** and **source-compiled CUDA @@ -284,13 +350,14 @@ warp specialization, NVSHMEM GPU-initiated RDMA) and only build/run on Upstream NeMo-RL pins `TORCH_CUDA_ARCH_LIST="9.0 10.0"`; upstream Automodel builds DeepEP for `"9.0 10.0 12.0"` — the same Hopper floor. -**Megatron backend: unused today** Transformer-Engine — the heavy -fused-kernel library the Megatron backend depends on — is currently pinned to -`NVTE_CUDA_ARCHS=90;100`, so the Megatron backend would be Hopper-only as built. -The customizer's DPO disables Megatron and trains on DTensor; whether GRPO will need -Megatron is still being determined in testing. TE is **not** inherently Hopper-only — -upstream Automodel builds it for `80;90;100;120` — so if the Megatron backend is later -required on A100, TE can be rebuilt to include `8.0`. +**Megatron backend: Hopper / Blackwell as built.** Transformer-Engine — the heavy +fused-kernel library both the Megatron and Automodel backends depend on — is pinned to +`NVTE_CUDA_ARCHS=90;100`, matching NeMo-RL's own `docker/Dockerfile`. TE is **not** +inherently Hopper-only — upstream Automodel builds it for `80;90;100;120` — so if +Megatron or DTensor V2 is later required on A100, TE can be rebuilt to include `8.0`. +`TORCH_CUDA_ARCH_LIST` is additionally a **runtime** requirement for the Megatron +backend, which raises if it is unset; the publish stage inherits `"9.0 10.0"` from the +builder stage. **Genuinely Hopper-only pieces:** only `deep_ep` and `deep_gemm` (their source builds fail for `8.0`). Everything else — the `torch`/`vllm`/`flash-attn`/`flashinfer` @@ -347,8 +414,8 @@ the uv cache + venv prefetch rather than via wheel images: commits than RL (a cp312/cu128 wheel cannot import on cp313/cu130, and the version deltas would fail `uv sync --frozen`). Reusing the pattern would mean new cp313/cu130 stages pinned to RL's exact commits, kept in lockstep with `uv.lock`. -- **Transformer-Engine** is the longest compile, but it is not built at all now — it - only exists in the unused `automodel` / `mcore` extras (see the note above). +- **Transformer-Engine** is the longest compile. It comes in with the `automodel` extra, + which the GRPO (future release) policy worker needs, so it is built from source here. `.python-version` pinning an exact patch release, which uv honours over whatever `uv python install` provisioned. Bumping `PYTHON_VERSION` alone therefore fixed nothing: every venv came up on RL's version while ours sat unused on disk, so the image shipped @@ -371,9 +438,8 @@ whenever the *dependency graph* hasn't changed: - To iterate on `services/rl` without re-entering the base build at all: build the base once and point training at it with `USE_PREBUILT_BASES=1 BASE_TAG_RL=` or `RL_BASE_CONTEXT=docker-image://`. -- Only the extras that are actually used are synced (`vllm`, `fsdp`, `modelopt`, - `nemo_gym`); dropping `mcore` / `automodel` removes the Transformer-Engine compile - entirely. +- Only the extras that are actually used are synced (`fsdp`, `automodel`, `mcore`, `vllm`, + `modelopt`, `nemo_gym`); `sglang` and `trtllm` are dropped. ### Prefetching the per-worker venvs (build once, not per job) @@ -389,16 +455,24 @@ Prefetched (the filters match **actor FQNs**, not extra names): | Filter | Extra | Needed by | |---|---|---| -| `dtensor_policy_worker.DTensorPolicyWorker` | `fsdp` | DPO + GRPO policy training | -| `vllm.vllm_worker` | `vllm` | GRPO generation — matches **both** `VllmGenerationWorker` and `VllmAsyncGenerationWorker` (NeMo-Gym forces async rollouts, so both are on the path) | -| `sync_rollout_actor.SyncRolloutActor` | `vllm` | GRPO rollout driver (sync path) | +| `dtensor_policy_worker.DTensorPolicyWorker` | `fsdp` | DPO policy training | +| `dtensor_policy_worker_v2.DTensorPolicyWorkerV2` | `automodel` | GRPO policy training (future release) — selected by `policy.dtensor_cfg._v2: true`; the only LoRA-capable DTensor worker. The V1 filter does not match it (`dtensor_policy_worker_v2.`) | +| `megatron_policy_worker.MegatronPolicyWorker` | `mcore` | GRPO policy training on Megatron (future release) — selected by `policy.megatron_cfg.enabled: true`; also hosts Megatron-native generation in-process. Does not match modelopt's `megatron_quant_policy_worker.MegatronQuantPolicyWorker` | +| `vllm.vllm_worker` | `vllm` | GRPO generation (future release) — matches **both** `VllmGenerationWorker` and `VllmAsyncGenerationWorker` (NeMo-Gym forces async rollouts, so both are on the path) | +| `sync_rollout_actor.SyncRolloutActor` | `vllm` | GRPO rollout driver (future, sync path) | | `nemo_gym.NemoGym` | `nemo_gym` | Gym environment actor (mode A, colocated) | | `nemo_gym_actor.SandboxedGymActor` | `nemo_gym` | Sandboxed Gym (mode B) — the trusted proxy actor in the training pod | | `broker_actor.SandboxEpisodeBrokerActor` | `nemo_gym` | Trusted episode broker — creates per-episode sandboxes so the job sandbox never holds the OpenSandbox credential | -Filters are **substring matches on actor FQNs**, so six filters yield seven venvs. They are -deliberately specific — a bare `vllm` would also match `nemo_rl.modelopt`'s -`vllm_quant_worker` and pull in the modelopt+vllm combination. +Filters are **substring matches on actor FQNs**, so eight filters yield nine venvs +(`vllm.vllm_worker` matches the sync and async workers alike). They are deliberately specific — +a bare `vllm` would also match `nemo_rl.modelopt`'s `vllm_quant_worker` and pull in the +modelopt+vllm combination, and a bare `megatron` would pull in `MegatronValueWorker` and +modelopt's `MegatronQuantPolicyWorker`. + +`tests/smoke_gpu/test_rl_training.py` asserts this set exactly +(`test_prefetched_venvs_match_expected_set`), so a filter that drifts fails the build rather +than silently shipping an image whose workers rebuild their venv on the node at job start. One venv is built **per actor, not per extra** — `prefetch_venvs.py` passes the actor FQN as the venv name — so the three `nemo_gym`-extra actors above each get their own directory and @@ -407,9 +481,10 @@ RL declares `nemo_gym = ["nemo_gym[sandbox]"]`. Without this, each venv is built **on the node at first run**, re-resolving and recompiling `deep_ep` / `mamba-ssm` / `causal-conv1d` against a cold uv cache on every -job. Not prefetched (they build on the node if a config selects them): `automodel`, -`mcore`, `sglang`, `trtllm`, modelopt-quant workers, and the async-GRPO actors -(`AsyncTrajectoryCollector`, `ReplayBuffer`). +job. Not prefetched (they build on the node if a config selects them): `sglang`, `trtllm`, +modelopt-quant workers, `DTensorValueWorkerV2` / `MegatronValueWorker` (GRPO is critic-free, +so there is no value model), and the async-GRPO actors (`AsyncTrajectoryCollector`, +`ReplayBuffer`). ### Link mode: why the uv cache ships inside the image @@ -444,7 +519,8 @@ the node at runtime (Gym environment venvs, non-prefetched actors) starts from a cache and re-downloads. That is the main argument for prefetching: a venv that is not baked in is not just -un-materialized, it is rebuilt against a cold cache. +un-materialized, it is rebuilt against a cold cache — which for the `automodel` tier would +mean compiling Transformer-Engine on every node at job time. #### `vllm/` is a private copy per vLLM venv @@ -529,7 +605,7 @@ every import. |---|---|---| | `NEMO_RL_REPO` / `NEMO_RL_REF` | see `docker-bake.hcl` | NeMo-RL source + pinned commit. | | `TORCH_CUDA_ARCH_LIST` | `"9.0 10.0"` | Archs for the torch-based source extensions (deep_ep, deep_gemm, mamba, causal-conv1d). | -| `NVTE_CUDA_ARCHS` | `90;100` | Archs for Transformer-Engine. Inert today (TE is only in the unused `automodel`/`mcore` extras); kept for when either is enabled. | +| `NVTE_CUDA_ARCHS` | `90;100` | Archs for Transformer-Engine, which the `automodel` extra compiles from source. Narrow to one arch for a faster dev build. | | `UV_SYNC_MODE` | `--frozen` | Reproducible sync. Set to empty to relock if a bumped RL commit's lock has drifted. | | `NEMO_GYM_PREFETCH_CONFIGS` | *(empty — prefetch off)* | Space-separated Gym config paths whose environment venvs are baked into `/opt/gym_venvs`. Empty means every environment installs at runtime on first use. Set to `examples/nemo_gym/prefetch_super_all_envs.yaml` to bake NeMo-RL's curated set back in. | diff --git a/tests/smoke_gpu/test_rl_training.py b/tests/smoke_gpu/test_rl_training.py index 3d60775110..446741c1a4 100644 --- a/tests/smoke_gpu/test_rl_training.py +++ b/tests/smoke_gpu/test_rl_training.py @@ -34,18 +34,29 @@ # Actor FQN -> packages that must import inside that actor's venv. # -# MUST list every venv the build prefetches (the six filters in docker/rl/Dockerfile.nmp-rl-base -# resolve to SEVEN actors, because `vllm.vllm_worker` matches the sync and async workers alike). -# Listing all seven is what makes a broken prefetch filter fail here instead of silently shipping an +# MUST list every venv the build prefetches (the eight filters in docker/rl/Dockerfile.nmp-rl-base +# resolve to NINE actors, because `vllm.vllm_worker` matches the sync and async workers alike). +# Listing all nine is what makes a broken prefetch filter fail here instead of silently shipping an # image whose workers rebuild their venv on the node at job start. WORKER_VENV_IMPORTS = { - # DPO + GRPO policy training (--extra fsdp) + # DPO policy training, DTensor V1 (--extra fsdp) "nemo_rl.models.policy.workers.dtensor_policy_worker.DTensorPolicyWorker": [ "torch", "flash_attn", "mamba_ssm", "causal_conv1d", ], + # GRPO policy training, DTensor V2 (--extra automodel). nemo_automodel and transformer_engine + # are checked through distribution metadata instead — see WORKER_VENV_DRIVER_LINKED below. + "nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2": [ + "torch", + ], + # GRPO policy training, Megatron (--extra mcore). Also hosts Megatron-native generation in + # process. megatron_core / megatron_bridge / transformer_engine are checked through + # distribution metadata — see WORKER_VENV_DRIVER_LINKED below. + "nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorker": [ + "torch", + ], # GRPO generation (--extra vllm). deep_ep is checked separately — see # WORKER_VENV_DRIVER_LINKED below. This is the full check for the vllm tier; the two actors # after it share the same extra. @@ -95,8 +106,23 @@ # `ImportError: libcuda.so.1: cannot open shared object file` for purely environmental reasons, so # presence is verified through distribution metadata instead — which still catches the case that # matters here: the extra failing to install the package at all. +# +# Transformer-Engine belongs here for the same reason: importing it loads +# libtransformer_engine.so, which resolves against the driver and cuDNN. It is listed under BOTH +# training tiers on purpose — RL's [tool.uv] override-dependencies collapses every TE requirement +# onto one git rev, so a single build is shared, and asserting it in each venv is what proves the +# sharing actually happened rather than one tier silently missing it. WORKER_VENV_DRIVER_LINKED = { "nemo_rl.models.generation.vllm.vllm_worker.VllmGenerationWorker": ["deep_ep"], + "nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2": [ + "nemo-automodel", + "transformer-engine", + ], + "nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorker": [ + "megatron-core", + "megatron-bridge", + "transformer-engine", + ], } WORKER_IMPORT_CASES = [(fqn, mod) for fqn, mods in sorted(WORKER_VENV_IMPORTS.items()) for mod in mods]