diff --git a/docker-bake.hcl b/docker-bake.hcl index f88abb3d93..c4dbd5843b 100644 --- a/docker-bake.hcl +++ b/docker-bake.hcl @@ -52,10 +52,6 @@ variable "FFMPEG_VLM_WHEEL_CONTEXT" { default = "" } -variable "TRANSFORMER_ENGINE_WHEEL_CONTEXT" { - default = "" -} - variable "DISTROLESS_BASE" { default = "nvcr.io/nvidia/distroless/python:3.11-v4.0.8" } @@ -171,11 +167,6 @@ function "get_ffmpeg_vlm_wheel_image" { result = "${WHEELS_REGISTRY}/ffmpeg-vlm-wheel:${WHEELS_TAG}" } -function "get_transformer_engine_wheel_image" { - params = [] - result = "${WHEELS_REGISTRY}/transformer-engine-wheel:${WHEELS_TAG}" -} - function "get_arch_tag" { params = [] result = BUILD_ARCH == "linux/arm64" ? "linux-arm64" : "linux-amd64" @@ -213,11 +204,6 @@ function "ffmpeg_vlm_wheel_context" { result = notequal(FFMPEG_VLM_WHEEL_CONTEXT, "") ? FFMPEG_VLM_WHEEL_CONTEXT : notequal(USE_LOCAL_WHEELS, "") ? "target:ffmpeg-vlm-wheel" : "docker-image://${get_ffmpeg_vlm_wheel_image()}" } -function "transformer_engine_wheel_context" { - params = [] - result = notequal(TRANSFORMER_ENGINE_WHEEL_CONTEXT, "") ? TRANSFORMER_ENGINE_WHEEL_CONTEXT : notequal(USE_LOCAL_WHEELS, "") ? "target:transformer-engine-wheel" : "docker-image://${get_transformer_engine_wheel_image()}" -} - function "wheel_tags" { params = [name] result = [ @@ -704,30 +690,6 @@ target "ffmpeg-vlm-wheel" { platforms = get_platforms() } -# transformer-engine wheel (cp313 / cu130), pinned to NeMo-RL's TE ref (release_v2.15, from RL's -# override-dependencies) and arch (90;100, matching RL). -# -# NOT consumed by any image today, and deliberately not in a build group: Transformer-Engine only -# ships in the `automodel` and `mcore` extras, and the customizer's DPO + GRPO path uses neither -# (DPO -> fsdp, GRPO generation -> vllm). Kept because it builds cleanly and is the drop-in for -# nmp-rl-base if a Megatron/automodel backend is ever adopted. Build on demand: -# docker buildx bake transformer-engine-wheel -target "transformer-engine-wheel" { - target = "transformer-engine-wheel" - context = "." - dockerfile = "docker/base/Dockerfile.python-wheels" - cache-to = maybe_registry_cache_to("transformer-engine-wheel") - cache-from = maybe_registry_cache_from("transformer-engine-wheel") - tags = wheel_tags("transformer-engine-wheel") - output = image_output() - args = { - TE_REF = "release_v2.15" - NVTE_CUDA_ARCHS = "90;100" - } - platforms = get_platforms() -} - - target "safe-synthesizer-tasks-docker" { target = "runtime" context = "." diff --git a/docker/base/Dockerfile.python-wheels b/docker/base/Dockerfile.python-wheels index 45f6024da2..a38d605c58 100644 --- a/docker/base/Dockerfile.python-wheels +++ b/docker/base/Dockerfile.python-wheels @@ -5,7 +5,6 @@ # Builds Python wheels for: # - causal-conv1d (CUDA extension required by mamba-ssm) # - mamba-ssm (selective state space model) -# - transformer-engine (cp313 / cu130, for nmp-rl-base — the longest source compile) # - av + opencv-python-headless + decord2 (FFmpeg 8.1.2 bundled wheels for nspect / VLM) # # causal-conv1d / mamba-ssm only ship source on PyPI and require nvcc. @@ -278,102 +277,6 @@ COPY --from=mamba-ssm-23-wheel-builder-py312 /wheels /wheels COPY --from=mamba-ssm-23-wheel-builder-py312-cu13.1.1 /wheels /wheels/cu13.1.1 COPY --from=mamba-ssm-23-wheel-builder-py312-cu13.2 /wheels /wheels/cu13.2 -# ============================================================================= -# transformer-engine wheel — Python 3.13 / CUDA 13.2 -# -# NOT CONSUMED BY ANY IMAGE TODAY, and intentionally so. Transformer-Engine ships only in -# NeMo-RL's `automodel` and `mcore` extras, and the customizer's DPO + GRPO path uses -# neither (DPO/GRPO train on `fsdp`; GRPO generates on `vllm`). nmp-rl-base therefore does -# not build TE at all. This stage is kept ready — TE is the longest CUDA compile in that -# image, so if the Megatron/automodel backend is ever enabled, wiring this wheel in -# (`uv pip install --no-deps`, plus --no-install-package on the syncs) avoids recompiling -# it on every base rebuild. -# -# ABI must match the consuming image EXACTLY or the wheel won't import: same base image -# (cuda-dl-base 26.03-cuda13.2), cp313, and the torch + TE ref NeMo-RL pins (torch 2.11.0 -# cu130; TE release_v2.15 from override-dependencies). Bumping RL's TE ref means rebuilding -# this wheel. -# -# Verified to build and produce a single self-contained wheel -# (`transformer_engine--cp313-cp313-linux_x86_64.whl`) — the pytorch extension -# (`transformer_engine_torch*.so`) is bundled inside it, there is no separate distribution. -# Build on demand with: -# docker buildx bake transformer-engine-wheel -# ============================================================================= -FROM nvcr.io/nvidia/cuda-dl-base:26.03-cuda13.2-devel-ubuntu24.04 AS te-wheel-base-py313 - -ARG PYTHON_VERSION=3.13.13 -ARG CMAKE_VERSION=4.0.3 -# Mirror NeMo-RL's pins (do not drift): torch 2.11.0 on the cu130 index; cuDNN wheel -# versions from RL's override-dependencies. -ARG TORCH_VERSION=2.11.0 - -ENV DEBIAN_FRONTEND=noninteractive -RUN apt-get update && apt-get install -y --no-install-recommends \ - build-essential ca-certificates git curl ninja-build libopenmpi-dev && \ - rm -rf /var/lib/apt/lists/* - -RUN ARCH="$(uname -m)" && F="cmake-${CMAKE_VERSION}-linux-${ARCH}" && \ - curl --retry 3 --retry-delay 2 -fsSL -o "${F}.tar.gz" \ - "https://github.com/Kitware/CMake/releases/download/v${CMAKE_VERSION}/${F}.tar.gz" && \ - tar -xzf "${F}.tar.gz" && cp -r "${F}/bin/"* /usr/local/bin/ && \ - cp -r "${F}/share/"* /usr/local/share/ && rm -rf "${F}" "${F}.tar.gz" - -# uv 0.11.18 (matches the nmp-rl base). Older uv (0.9.14, used by the mamba stages for -# their lock's required-version) has no python-build-standalone entry for 3.13.13, so -# `uv python install 3.13.13` fails with "No download found". -COPY --from=ghcr.io/astral-sh/uv:0.11.18 /uv /usr/local/bin/uv -ENV UV_PROJECT_ENVIRONMENT=/opt/venv UV_LINK_MODE=copy -# cuda-dl-base / Ubuntu 24.04 ships no Python 3.13 — have uv provision the managed -# interpreter before creating the venv (same as the nmp-rl base does). -RUN uv python install ${PYTHON_VERSION} && uv venv --python ${PYTHON_VERSION} /opt/venv -ENV PATH="/opt/venv/bin:$PATH" VIRTUAL_ENV=/opt/venv - -# Build environment: torch is the ABI anchor (match nmp-rl-base); TE also links cuDNN -# and (for some kernels) mathdx. Versions mirror NeMo-RL's pyproject. -RUN uv pip install --torch-backend=cu130 "torch==${TORCH_VERSION}" && \ - uv pip install setuptools wheel packaging pybind11 ninja \ - "nvidia-cudnn-cu13==9.20.0.48" \ - "nvidia-cudnn-frontend==1.23.0" \ - "nvidia-mathdx==25.1.1" - -ENV CUDA_HOME=/usr/local/cuda \ - CPLUS_INCLUDE_PATH=/usr/local/cuda/include/cccl \ - CUDNN_HOME=/opt/venv/lib/python3.13/site-packages/nvidia/cudnn \ - CUDNN_PATH=/opt/venv/lib/python3.13/site-packages/nvidia/cudnn \ - LD_LIBRARY_PATH="/opt/venv/lib/python3.13/site-packages/nvidia/cudnn/lib:${LD_LIBRARY_PATH}" - -FROM te-wheel-base-py313 AS transformer-engine-wheel-builder - -# NeMo-RL pins TE via override-dependencies (@release_v2.15 wins across all extras). -ARG TE_REF=release_v2.15 -# Includes 80 (A100) on purpose. Transformer-Engine is NOT Hopper-only - upstream Automodel -# builds it for "80;90;100;120". Only deep_ep / deep_gemm have a hard SM90 floor, and neither is -# part of TE. Building 80;90;100 here means the wheel works on A100 as well as Hopper/Blackwell, -# so the Megatron (mcore) / automodel backends could run on A100 if they are ever enabled - -# without having to rediscover that TE was the blocker. Cost is a longer compile (one extra arch). -# NeMo-RL itself pins 90;100; docker-bake.hcl currently overrides this arg back to 90;100 to keep -# today's behavior byte-identical. Drop that override to get an A100-capable wheel. -ARG NVTE_CUDA_ARCHS=80;90;100 -ARG NVTE_BUILD_THREADS_PER_JOB=8 - -RUN mkdir -p /wheels && \ - git clone https://github.com/NVIDIA/TransformerEngine.git /src/TransformerEngine && \ - cd /src/TransformerEngine && \ - git fetch origin ${TE_REF} && git checkout FETCH_HEAD && \ - git submodule update --init --recursive && \ - NVTE_FRAMEWORK=pytorch NVTE_CUDA_ARCHS="${NVTE_CUDA_ARCHS}" \ - NVTE_BUILD_THREADS_PER_JOB=${NVTE_BUILD_THREADS_PER_JOB} \ - uv build --wheel --no-build-isolation --out-dir /wheels . && \ - rm -rf /src/TransformerEngine - -# The final transformer-engine-wheel image contains (cp313 / cu130): -# - transformer_engine-*-cp313-*.whl -# - transformer_engine_torch-*-cp313-*.whl (if TE emits it separately — verify) -# Consumers install with `uv pip install --no-deps /wheels/*.whl`. -FROM scratch AS transformer-engine-wheel -COPY --from=transformer-engine-wheel-builder /wheels /wheels - # ============================================================================= # FFmpeg / VLM wheels (manylinux_2_28) — av + opencv-python-headless + decord2 # ============================================================================= diff --git a/docker/rl/Dockerfile.nmp-rl-base b/docker/rl/Dockerfile.nmp-rl-base index 9f738afd74..7d418c25f2 100644 --- a/docker/rl/Dockerfile.nmp-rl-base +++ b/docker/rl/Dockerfile.nmp-rl-base @@ -1,26 +1,19 @@ # syntax=docker/dockerfile:1 # nmp-rl base - NeMo-RL + NeMo-Gym built FROM SOURCE on cuda-dl-base (Python 3.13, CUDA 13). # -# Mirrors the build in NeMo-RL's own docker/Dockerfile (base + hermetic + release stages), -# stripped to the extras nemo-platform needs (fsdp, vllm, modelopt, nemo_gym). Building from source -# - rather than basing on the prebuilt NGC nemo-rl image - lets us pin ANY RL commit via build-arg -# (NEMO_RL_REF); Gym follows as RL's submodule pin. +# Mirrors NeMo-RL's own docker/Dockerfile, stripped to the extras nemo-platform needs. Building +# 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. # -# Deliberately NOT built: mcore and automodel (alternative training backends the customizer does not -# use, and the only source of Transformer-Engine - the longest CUDA compile), plus sglang and trtllm. -# DPO trains on `fsdp`; GRPO trains on `fsdp` and generates on `vllm`. See README.md. +# DPO trains on `fsdp`; GRPO trains on `fsdp` and generates on `vllm`. mcore, automodel, sglang +# and trtllm are deliberately not built. # -# NOTE (venvs): NeMo-RL does not run its workers in the base venv - each Ray actor launches with its -# own per-node venv under /opt/ray_venvs. The `uv sync --extra` calls below only warm the uv CACHE; -# the venvs training actually uses are materialized by the prefetch steps further down. See README.md -# ("How NeMo-RL runs: the base venv is not where training happens"). -# -# NOTE (publish stage): it inherits the builder filesystem (like RL's release stage) rather than -# COPY-ing /opt/nemo_rl_venv onto a fresh base. cuda-dl-base has NO system Python, so `uv venv` -# symlinks the venv's python to uv's managed interpreter (which lives OUTSIDE the venv); a bare -# COPY-venv-to-fresh-base would leave a broken python symlink. The image is slimmed by pruning -# build-only tooling in the publish stage; the CUDA devel toolkit (nvcc) and gcc/build-essential -# stay because deep_ep JIT and native-v1 env installs need them at runtime. +# 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 +# venvs training uses are materialized by the prefetch steps further down. +# - The publish stage inherits the builder filesystem rather than COPY-ing the venv onto a fresh +# base - cuda-dl-base has no system Python, so the venv's python symlinks to uv's managed +# interpreter outside the venv and a bare COPY would break it. # Kept in lockstep with the pinned NEMO_RL_REF: RL's own Dockerfile builds on this tag, and its lock # overrides nvidia-nccl-cu13 to the version that ships here. Bump both together. @@ -94,22 +87,13 @@ RUN curl -LsSf "https://astral.sh/uv/${UV_VERSION}/install.sh" | sh && \ uv python install "${PYTHON_VERSION}" # Ray / NeMo-RL runtime behavior (affects the running container, not the build): -# RAY_USAGE_STATS_ENABLED=0 - disable Ray's anonymous usage telemetry (no phone-home). -# RAY_ENABLE_UV_RUN_RUNTIME_ENV=0 - Ray >=2.47 auto-builds a uv venv PER TASK for `uv run` -# executables; with RL's huge, sometimes-compiled deps that causes severe contention, so we -# disable it and let NeMo-RL build worker venvs ONCE PER NODE (nemo_rl/utils/venv.py). -# NEMO_RL_VENV_DIR=/opt/ray_venvs - dir where RL materializes those per-node Ray worker venvs. -# NEMO_GYM_VENV_DIR=/opt/gym_venvs - dir where Gym materializes per-ENVIRONMENT venvs at runtime; -# an environment's own deps (rdkit, spacy, numpy 2.x, ...) install here, isolated from the -# training venv so they can't clobber RL's pinned stack (e.g. numpy 1.X vs 2.x). -# NRL_CONTAINER=1 - marks this as a NeMo-RL container. REQUIRED, not cosmetic: NeMo-RL gates -# `get_nemo_gym_uv_cache_dir()` on it (nemo_rl/environments/nemo_gym.py), and without it the -# function returns None, so NeMo-Gym silently falls back to its own default cache inside the -# source tree (3rdparty/Gym-workspace/Gym/cache/uv) - a multi-GB duplicate of the shared uv -# cache. Note the venv-dir helper is NOT gated the same way, which is why venvs land correctly -# while the cache would not. It also enables the frozen-environment wrapper scripts in -# prefetch_venvs.py and the container fingerprint check (a warning-only no-op when the -# fingerprint file is absent). RL's own Dockerfile sets this in its base stage. +# RAY_USAGE_STATS_ENABLED=0 - no telemetry phone-home. +# RAY_ENABLE_UV_RUN_RUNTIME_ENV=0 - Ray >=2.47 builds a uv venv PER TASK for `uv run`; with RL's +# deps that means severe contention, so let NeMo-RL build worker venvs once per node instead. +# NEMO_RL_VENV_DIR / NEMO_GYM_VENV_DIR - where the per-actor and per-Gym-environment venvs land. +# Gym environments stay isolated so their deps cannot clobber RL's pinned stack. +# NRL_CONTAINER=1 - REQUIRED, not cosmetic: without it NeMo-Gym falls back to a cache inside its +# own source tree, a multi-GB duplicate. See README.md, "Why NRL_CONTAINER=1 is required". ENV RAY_USAGE_STATS_ENABLED=0 \ RAY_ENABLE_UV_RUN_RUNTIME_ENV=0 \ NRL_CONTAINER=1 \ @@ -117,28 +101,19 @@ 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 policy training), vllm (GRPO generation), nemo_gym and modelopt are built. -# The mcore and automodel extras are deliberately EXCLUDED: they are alternative training backends -# the customizer does not use, and they are the only source of Transformer-Engine (the longest CUDA -# compile in this image). See README.md, "Excluded on purpose: the automodel and mcore backends". +# 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). FROM base AS builder WORKDIR /opt/nemo-rl -# Arch knobs + cuDNN wiring (mirrors RL's own hermetic stage). NVTE_CUDA_ARCHS / -# TORCH_CUDA_ARCH_LIST control which GPU archs the SOURCE-compiled kernels (mamba-ssm, -# causal-conv1d, deep_ep, deep_gemm; TE too when built) target — these are the ONLY packages the -# arch flags touch. The big prebuilt wheels (torch, vllm, flash-attn, flashinfer) already ship "fat" -# binaries covering many archs; nvcc is never re-run for them, so the arch flags neither narrow nor -# widen the GPUs they support. -# -# BUILD ARCH MUST BE >= 9.0. deep_ep / deep_gemm (pulled by the `vllm` extra, i.e. GRPO generation) -# use Hopper-class features and will NOT compile for SM80/86/89 - building with 8.0 fails the nvcc -# step. This is a BUILD-time constraint only: nvcc cross-compiles, RL upstream pins the same "9.0 10.0". +# Arch knobs + cuDNN wiring (mirrors RL's own hermetic stage). The arch flags only affect the +# SOURCE-compiled kernels (mamba-ssm, causal-conv1d, deep_ep, deep_gemm, TE); prebuilt wheels +# (torch, vllm, flash-attn, flashinfer) ship fat binaries and are unaffected. # -# It does NOT make the whole image Hopper-only at RUNTIME. What a workload needs depends on which -# kernels it launches: DPO trains on the `fsdp` extra (flash-attn / mamba / causal-conv1d - no -# deep_ep) and is expected to run on A100 and up, while GRPO generation uses deep_ep/deep_gemm and -# needs Hopper/Blackwell. See README.md for more details. +# BUILD ARCH MUST BE >= 9.0: deep_ep / deep_gemm use Hopper-class features and fail nvcc at 8.0. +# This is a BUILD-time constraint only and does NOT make the image Hopper-only at runtime - see +# README.md, "Hardware: which GPUs run what". # # Keep at 90 (H100/H200) and/or 100 (B200/GB200); narrow to one for a faster dev build, e.g. # --build-arg NVTE_CUDA_ARCHS=90 --build-arg TORCH_CUDA_ARCH_LIST="9.0" @@ -147,36 +122,21 @@ ARG NVTE_BUILD_THREADS_PER_JOB=2 ARG NVTE_CUDA_ARCHS=90;100 ARG NVTE_FRAMEWORK=pytorch ARG TORCH_CUDA_ARCH_LIST="9.0 10.0" -# Path / toolchain wiring for the source compiles above and for runtime JIT: -# CUDA_HOME - where the build (and DeepEP's runtime JIT) find nvcc and CUDA libs. -# CPLUS_INCLUDE_PATH - CUDA 13 relocated its C++ standard-library headers () -# under $CUDA_HOME/include/cccl. deep_gemm #includes them, so this dir -# must be on the C++ include path or its compile fails with a -# missing-header error. -# CUDNN_HOME/CUDNN_PATH - point Transformer-Engine's Python at the pip-installed cuDNN wheel -# instead of the base image's system cuDNN. The two can differ, and TE -# prioritizes the system copy by default -> version-mismatch crash at -# runtime. -# LD_LIBRARY_PATH - prepend that same pip cuDNN lib dir so the dynamic linker (ld.so) -# resolves libcudnn*.so from pip when loading libtransformer_engine.so. -# This mirrors the cuDNN/arch wiring in NeMo-RL's own docker/Dockerfile (hermetic stage). +# Path / toolchain wiring for the source compiles and for runtime JIT (mirrors RL's own +# docker/Dockerfile): +# CUDA_HOME - where the build and DeepEP's runtime JIT find nvcc and CUDA libs. +# CPLUS_INCLUDE_PATH - CUDA 13 moved under $CUDA_HOME/include/cccl; deep_gemm +# includes them, so this dir must be on the C++ include path. +# CUDNN_HOME/CUDNN_PATH - point Transformer-Engine at the pip cuDNN wheel, not the base image's +# system cuDNN, which TE prefers by default -> version-mismatch crash. +# LD_LIBRARY_PATH - same pip cuDNN lib dir, so ld.so resolves libcudnn*.so when loading +# libtransformer_engine.so. # -# UV_CACHE_DIR + UV_LINK_MODE work as a PAIR and must not be changed independently: -# UV_LINK_MODE=symlink - set PER-COMMAND on the venv-creating steps below, NOT as a persistent -# ENV: uv rejects symlink installs combined with `--no-cache` whenever it has to build a PEP -# 517 build environment, and downstream images (the training image's editable glue install) -# do exactly that. Scoping it per-command keeps it out of anything we do not control. -# Symlinking means venv files point INTO the uv cache instead of being copies. This image -# materializes ~30 venvs (5 per-worker Ray venvs + the bundled Gym environment venvs), and the -# two vLLM-tier worker venvs alone would otherwise be full duplicates of the same multi-GB -# stack. Symlinking means the image carries ONE copy of each package (in the cache) plus -# near-free symlinks. NeMo-RL's own release stage does the same. -# UV_CACHE_DIR=/opt/uv_cache - because those symlinks resolve at RUNTIME, the cache must ship -# INSIDE the image and stay readable by the non-root user (the training image runs as UID -# 1000). uv's default ~/.cache/uv is under /root, which is mode 700 and unreadable to that -# user - the same reason the managed Python lives under /opt. It also cannot be a BuildKit -# `--mount=type=cache`, since those never enter the image and every symlink would dangle. -# Consequence: NEVER delete /opt/uv_cache after the venvs are built - it would break every venv. +# UV_CACHE_DIR + UV_LINK_MODE are a PAIR - see README.md, "Link mode". In short: venvs symlink into +# the cache instead of copying, so the cache must ship inside the image and stay world-readable, and +# deleting /opt/uv_cache breaks every venv. UV_LINK_MODE is set PER-COMMAND, not as an ENV: uv +# rejects symlink installs combined with `--no-cache` when it builds a PEP 517 environment, which +# downstream images do. ENV NVTE_BUILD_MAX_JOBS=${NVTE_BUILD_MAX_JOBS} \ NVTE_BUILD_THREADS_PER_JOB=${NVTE_BUILD_THREADS_PER_JOB} \ NVTE_CUDA_ARCHS=${NVTE_CUDA_ARCHS} \ @@ -190,13 +150,13 @@ ENV NVTE_BUILD_MAX_JOBS=${NVTE_BUILD_MAX_JOBS} \ CUDNN_PATH=/opt/nemo_rl_venv/lib/python3.13/site-packages/nvidia/cudnn \ LD_LIBRARY_PATH="/opt/nemo_rl_venv/lib/python3.13/site-packages/nvidia/cudnn/lib:${LD_LIBRARY_PATH}" -# CI-cache layering: copy ONLY the resolver inputs before the heavy uv sync, so the expensive -# compile layer's cache key depends on the DEPENDENCY GRAPH, not on every RL source file. A -# source-only RL bump (nemo_rl/*.py changed, deps unchanged) then reuses the compiled layer instead -# of recompiling TE/deep_ep/... These inputs are: the lock + project metadata; the workspace-member -# trees under 3rdparty (Gym, Automodel, Megatron-Bridge + nested Megatron-LM) and research/; and the -# top-level nemo_rl package stub (its __version__ is read when building workspace metadata). -# TensorRT-LLM is excluded (we never build the trtllm extra). Full source is copied after the sync. +# CI-cache layering: copy ONLY the resolver inputs before the heavy uv sync, so the compile +# 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. 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/** \ @@ -207,20 +167,15 @@ COPY --from=nemo-rl research/ ./research/ # RL's uv.lock. Set to "" to relock if you bump RL to a commit whose lock has drifted. ARG UV_SYNC_MODE=--frozen -# Warm the uv CACHE with every backend's dependencies, one extra at a time (they can conflict), -# then a final default sync. This is where the one-time heavy compiles happen (mamba-ssm, -# causal-conv1d, deep_ep, deep_gemm) - the resulting wheels land in the uv cache so the per-worker -# venv prefetch below can materialize each venv without recompiling anything. +# 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 +# land in the cache so the prefetch below materializes each venv without recompiling. # -# NOTE: `uv sync` is EXACT - each --extra run prunes the previous one's packages from the venv. So -# these calls do NOT accumulate; the base venv ends up with default deps + groups only. Their real -# product is the populated uv cache. (Transformer-Engine is not compiled at all: it only ships in -# the mcore/automodel extras, which we do not build.) +# `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. # -# --no-install-project = resolve and install all DEPENDENCIES, but do NOT build/install the root -# project package (nemo-rl) itself. This lets the expensive dependency compiles happen exactly once -# here; the root package is then installed editable in a cheap step afterwards (the `uv pip install -# -e /opt/nemo-rl` below), so `import nemo_rl` works without recompiling anything. +# --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 # Symlink the base venv into the uv cache (see UV_CACHE_DIR/UV_LINK_MODE note above). Scoped to # this command rather than the image ENV so downstream `--no-cache` installs are unaffected. @@ -258,15 +213,11 @@ find "${UV_CACHE_DIR}" -type d \ \( -name PyNvVideoCodec -o -name 'pynvvideocodec-*.data' \) -exec rm -rf {} + || true EOF -# Full RL source including submodules (Automodel, Megatron-Bridge + nested Megatron-LM, Gym). Copied -# AFTER the heavy uv sync so a source-only RL bump (nemo_rl/*.py changed but uv.lock / pyproject / -# 3rdparty unchanged) is a CACHE HIT on the compile layer above, and only the cheap steps below -# re-run. This is the main lever for fast CI rebuilds. +# Full RL source (Automodel + Gym). Copied AFTER the heavy uv sync so a source-only RL bump is a +# cache hit on the compile layer above. # -# --exclude=**/.git: omitting --keep-git-dir on the ADD drops only the TOP-LEVEL repository's .git; -# BuildKit still checks out each submodule (Gym, Automodel, Megatron-Bridge, Megatron-LM) with its -# own .git. Excluding here keeps them out of this layer entirely - deleting them in a later layer -# would leave them intact in this one and still extractable from the published image. +# --exclude=**/.git: the ADD drops only the top-level .git, so each submodule still carries its own. +# Excluding here keeps them out of this layer - removing them later would leave them extractable. COPY --from=nemo-rl --exclude=**/.git . /opt/nemo-rl # Install the workspace ROOT (nemo-rl) editable. --no-install-project above installed every @@ -275,43 +226,23 @@ COPY --from=nemo-rl --exclude=**/.git . /opt/nemo-rl RUN uv pip install --python /opt/nemo_rl_venv/bin/python --no-deps -e /opt/nemo-rl # ---- Per-worker Ray venvs (the venvs training actually runs in) ---- -# NeMo-RL does NOT run its workers in the base venv. Each Ray actor is mapped to a py_executable in -# nemo_rl/distributed/ray_actor_environment_registry.py; the uv-based ones ("uv run --extra ") -# are materialized as SEPARATE per-node venvs under NEMO_RL_VENV_DIR (/opt/ray_venvs) by -# nemo_rl/utils/venvs.py. Only "system Python" actors use the base venv. -# -# Without this step those venvs are built on the NODE at first run - re-resolving and recompiling -# deep_ep / mamba-ssm / causal-conv1d on every job, with a cold uv cache. Prefetching bakes them in -# (same approach as NeMo-RL's own release stage). -# -# Prefetched (filters are SUBSTRING matches against ACTOR FQNs, not extra names). One venv is built -# PER ACTOR, not per extra - prefetch_venvs.py passes the actor FQN as the venv name - so actors -# sharing an extra still need their own filter. The six filters below resolve to SEVEN venvs, -# exactly what DPO + GRPO + sandboxed Gym need: -# dtensor_policy_worker.DTensorPolicyWorker -> --extra fsdp (DPO/GRPO policy training) -# vllm.vllm_worker -> --extra vllm (GRPO generation) - matches BOTH -# VllmGenerationWorker and VllmAsyncGenerationWorker, which is wanted: NeMo-Gym forces async -# rollouts, so the async engine is on the GRPO+Gym path. -# sync_rollout_actor.SyncRolloutActor -> --extra vllm (GRPO rollout driver, sync path) -# nemo_gym.NemoGym -> --extra nemo_gym (Gym environment actor, mode A) -# nemo_gym_actor.SandboxedGymActor -> --extra nemo_gym (trusted proxy actor, mode B - the -# Gym stack itself runs in the job sandbox; this actor never imports user environment code) -# broker_actor.SandboxEpisodeBrokerActor -> --extra nemo_gym (trusted episode broker: creates -# per-episode sandboxes on the agent's behalf so the untrusted job sandbox never holds the -# OpenSandbox credential) -# The filters are intentionally specific: a bare "vllm" would also match nemo_rl.modelopt's -# vllm_quant_worker and drag in the modelopt+vllm combination, and a bare "nemo_gym" would not -# distinguish the three nemo_gym-extra actors (it would match all three, which is fine, but naming -# them keeps this list readable against ray_actor_environment_registry.py). +# Workers do not run in the base venv - each Ray actor gets its own venv under /opt/ray_venvs. +# Prefetching bakes them in; otherwise they are built on the node at first run against a cold +# cache. See README.md, "Prefetching the per-worker venvs". # -# The opensandbox / tenacity SDKs the sandbox path needs arrive through the extra itself: RL -# declares `nemo_gym = ["nemo_gym[sandbox]"]`, and Gym's `sandbox` extra pins opensandbox>=0.1.9 and -# tenacity>=9.1.4. +# 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) +# 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) +# nemo_gym.NemoGym -> nemo_gym (Gym environment actor, mode A) +# nemo_gym_actor.SandboxedGymActor -> nemo_gym (mode B proxy actor) +# 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 (unused by our DPO+GRPO path - see README), sglang, trtllm, -# modelopt-quant, and the async-GRPO bookkeeping actors AsyncTrajectoryCollector / ReplayBuffer -# (nemo_rl.algorithms.async_utils - they use the vllm extra but match none of the filters above); -# those build on the node if a config selects them. +# NOT prefetched: automodel, mcore, sglang, trtllm, modelopt-quant +# and the async-GRPO bookkeeping actors; those build on the node if a config selects them. RUN <<"EOF" bash -exu export UV_LINK_MODE=symlink uv run nemo_rl/utils/prefetch_venvs.py \ @@ -343,14 +274,12 @@ done EOF # ---- Bundled NeMo-Gym environment venvs (GRPO only; OFF by default, opt in below) ---- -# SECOND, separate venv layer from /opt/ray_venvs above: NeMo-Gym gives every environment server its -# own venv under NEMO_GYM_VENV_DIR (/opt/gym_venvs), built from that server's own pyproject.toml / -# requirements.txt. The isolation is required - env deps (rdkit, spacy, numpy 2.x) would otherwise -# clobber RL's pinned stack. +# A SECOND venv layer, separate from /opt/ray_venvs: one venv per Gym environment server, built +# from that server's own pyproject/requirements, so env deps cannot clobber RL's pinned stack. +# See README.md, "NeMo-Gym environments: a second, separate venv layer". # -# Nothing is prefetched: every environment installs into /opt/gym_venvs at first use on the node. -# To bake a set in, pass a space-separated list of Gym config paths - e.g. NeMo-RL's curated -# examples/nemo_gym/prefetch_super_all_envs.yaml. +# Nothing is prefetched by default - environments install on the node at first use. Pass a +# space-separated list of Gym config paths to bake a set in. ARG NEMO_GYM_PREFETCH_CONFIGS="" # Gym's model servers pin a plain `vllm==0.20.0`, which resolves to the CUDA 12 wheel on PyPI - # it would fail to load CUDA libraries inside this CUDA 13 image. Override it to vLLM's prebuilt @@ -410,20 +339,29 @@ RUN /opt/nemo_rl_venv/bin/python tools/generate_fingerprint.py \ # ---- publish base (inherits builder FS; keeps the uv-managed python valid - see header NOTE) ---- FROM builder AS nmp-rl-base +# setuptools CVEs (GHSA-58pv-8j8x-9vj2, GHSA-8rrh-rw8j-w5fx, in its vendored jaraco.context/wheel): +# Megatron-Bridge pins `setuptools<80.0.0` for its build, and PEP 517 build environments resolve +# 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. +RUN <<"EOF" bash -exu +for d in "${UV_CACHE_DIR}"/archive-v0/*/; do + ls "${d}"setuptools-*.dist-info >/dev/null 2>&1 || continue + find /opt/nemo_rl_venv /opt/ray_venvs /opt/gym_venvs -lname "*$(basename "${d}")*" -print -quit 2>/dev/null | grep -q . && continue + rm -rf "${d}" +done +EOF + # Prune build-only tooling to shrink the image. KEPT on purpose (needed at runtime): CUDA devel # toolkit / nvcc (deep_ep JIT-compiles kernels at runtime), gcc/build-essential + libibverbs-dev # (JIT + native-v1 env-dep compiles), git + curl (git+https env deps, downloads). # -# /opt/uv_cache is NOT pruned and must never be: the per-worker venvs under /opt/ray_venvs are -# symlinked into it (UV_LINK_MODE=symlink), so deleting it breaks every one of them - as it would -# any /opt/gym_venvs environment baked in via NEMO_GYM_PREFETCH_CONFIGS. It also doubles as a warm -# cache for everything installed at runtime (Gym environments, user FileSets, non-prefetched -# actors), making those installs faster. +# /opt/uv_cache is NOT pruned and must never be: every venv symlinks into it, so deleting it breaks +# all of them. # -# Nsight Systems / Nsight Compute come in with cuda-dl-base and carry Go stdlib CVEs (3 Critical, 25 -# High). They are profilers - nothing on the DPO/GRPO path runs nsys/ncu. This clears them from the -# flattened rootfs the scanner walks; it does NOT shrink the pulled image, since they live in an -# ancestor layer and this stage inherits the builder filesystem (see header NOTE). +# Nsight Systems / Nsight Compute arrive with cuda-dl-base and carry Go stdlib CVEs. They are +# profilers, unused here. Removing them clears the flattened rootfs the scanner walks; it does not +# shrink the pulled image, since they live in an ancestor layer. RUN apt-get purge -y ccache vim vim-common less >/dev/null 2>&1 || true; \ apt-get autoremove -y >/dev/null 2>&1 || true; \ rm -rf \ @@ -455,14 +393,9 @@ RUN chmod -R a+rX /opt/uv /opt/uv_cache /opt/nemo_rl_venv /opt/nemo-rl /opt/gym_ # directories must be writable by the non-root runtime user. Only the top-level directories are # chowned - NOT -R - so the layer stays tiny and the prefetched venvs underneath stay root-owned. # -# /opt/uv_cache is deliberately left read-only. Every prefetched venv symlinks its package files -# into it (UV_LINK_MODE=symlink), so it holds the actual code the trainer executes; making it -# writable would let any code running in this container rewrite the training stack's own Python -# packages. That matters because this container also runs user-authored environment code: NeMo-Gym -# environments execute here, and a `native-v1` environment FileSet is user-supplied source. The -# sandboxed-Gym design does not remove this: its `sandboxed` mode is opt-in (colocated is the -# default), and its per-episode sandbox targets model-generated code, not the environment itself. -# Downstream images point uv at a writable per-user cache instead (see Dockerfile.nmp-rl-training). +# /opt/uv_cache is deliberately left read-only: it holds the actual code the trainer executes, and +# this container also runs user-authored Gym environment code, which must not be able to rewrite +# it. Downstream images point uv at a writable per-user cache (see Dockerfile.nmp-rl-training). ARG RUNTIME_UID=1000 ARG RUNTIME_GID=1000 RUN chown ${RUNTIME_UID}:${RUNTIME_GID} /opt/ray_venvs /opt/gym_venvs diff --git a/docker/rl/README.md b/docker/rl/README.md index a1617c638c..8f4f535998 100644 --- a/docker/rl/README.md +++ b/docker/rl/README.md @@ -13,7 +13,7 @@ rediscovered. | Dockerfile | Image | Role | |---|---|---| | `Dockerfile.nmp-rl-base` | `nmp-rl-base` | Heavy base. Clones NeMo-RL (with submodules), warms the uv cache for the extras we need (`vllm`, `fsdp`, `modelopt`, `nemo_gym`), then **prefetches the per-worker Ray venvs**. All one-time CUDA compiles (mamba-ssm, causal-conv1d, deep_ep, deep_gemm) happen here. | -| `Dockerfile.nmp-rl-training` | `nmp-rl-training` | Thin layer on the base. Adds the pure-Python platform glue editably; entrypoint runs `python -m nmp.rl.tasks.training` (Ray bootstrap → DPO/GRPO). Also the Gym environment runtime. A `smoke-test` stage runs CPU-only import checks during the build. | +| `Dockerfile.nmp-rl-training` | `nmp-rl-training` | Thin layer on the base. Adds the pure-Python platform glue editably; entrypoint runs `python -m nmp.rl.tasks.training` (Ray bootstrap → DPO). Also the Gym environment runtime. A `smoke-test` stage runs CPU-only import checks during the build. | `docker-bake.hcl` wires them: `nmp-rl-training`'s base context defaults to building `nmp-rl-base` as a dependency, unless `USE_PREBUILT_BASES` / `RL_BASE_CONTEXT` point @@ -110,12 +110,12 @@ resolve the same interpreter as their parent. - **Per node, not per worker.** Eight GPU workers on one node share one venv on that node's disk; `venvs.py` uses a `STARTED_ENV_BUILDER` lock so it is built once, not eight times. - **Different workers can have different GPU requirements.** DPO's worker launches with the - `fsdp` venv (no `deep_ep`); GRPO's generation worker with the `vllm` venv (`deep_ep`, - Hopper-only) — in the same image. + `fsdp` venv (no `deep_ep`); GRPO's (will be added in future) generation worker with the + `vllm` venv (`deep_ep`, Hopper-only) — in the same image. - `NEMO_RL_PY_EXECUTABLES_SYSTEM=1` collapses every actor into the base venv. We do **not** set it. -### Which extras DPO and GRPO actually use +### Which extras DPO and GRPO (will be added in future releases) actually use | Algorithm | Actor | Extra / venv | Notable contents | |---|---|---|---| @@ -162,10 +162,9 @@ those two. The Automodel and Megatron-Bridge **source trees still ship** regardl 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. A ready-made -`transformer-engine-wheel` target (cp313 / cu130, pinned to RL's TE ref) already exists in -`docker-bake.hcl` / `docker/base/Dockerfile.python-wheels` for that day; nothing builds it -today. +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. ### NeMo-Gym environments: a second, separate venv layer @@ -277,12 +276,13 @@ Mamba or hybrid model on A100 would fail at first kernel launch with the current `9.0 10.0` build. Add `8.0` to the arch list to cover that case — unlike `deep_ep`, these two do compile for `sm_80`. -**GRPO targets Hopper / Blackwell (for now).** GRPO generation uses vLLM, whose MoE -path pulls `deep_ep` (and `deep_gemm` for fp8) — source-compiled CUDA extensions that -use Hopper-class features (TMA async copy, warp specialization, NVSHMEM GPU-initiated -RDMA) and only build/run on **SM 9.0 (Hopper) / 10.0 (Blackwell)**. They do not -compile for SM 8.0/8.6/8.9. (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.) +**GRPO targets Hopper / Blackwell (for now).** GRPO (will be added in future releases) +generation uses vLLM, whose MoE path pulls `deep_ep` (and `deep_gemm` for fp8), +source-compiled CUDA extensions that use Hopper-class features (TMA async copy, +warp specialization, NVSHMEM GPU-initiated RDMA) and only build/run on +**SM 9.0 (Hopper) / 10.0 (Blackwell)**. They do not compile for SM 8.0/8.6/8.9. +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 @@ -434,8 +434,17 @@ readable by the non-root user. Hence `UV_CACHE_DIR=/opt/uv_cache`: - **`/opt/uv_cache` must never be deleted** — pruning it breaks every venv that points into it. The prune step in the publish stage deliberately leaves it alone. -A useful side effect: environments installed at runtime (user Gym FileSets, actors that -were not prefetched) resolve against a warm cache instead of downloading from scratch. +**It does not warm runtime installs in the training image.** `Dockerfile.nmp-rl-training` +repoints `UV_CACHE_DIR` at a writable per-user path, because `/opt/uv_cache` holds the code +the trainer executes and must stay read-only (user-authored Gym environment code runs in +this container). uv reads only `UV_CACHE_DIR` — there is no read-through to a second cache — +and NeMo-Gym resolves its own cache by shelling out to `uv cache dir` +(`nemo_rl/environments/nemo_gym.py`), which honours that same variable. So anything built on +the node at runtime (Gym environment venvs, non-prefetched actors) starts from an **empty** +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. #### `vllm/` is a private copy per vLLM venv