Skip to content

chore(customizer): add automodel and mcore backends in RL image #1231

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
anubhutivyas merged 1 commit into from
Aug 11, 2026
Merged

chore(customizer): add automodel and mcore backends in RL image #1231

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 35 additions & 14 deletions docker/rl/Dockerfile.nmp-rl-base
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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/** \
Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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)
Expand All @@ -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 \
Expand All @@ -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 \
Expand Down Expand Up @@ -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
Expand Down
Loading