Skip to content

Commit 730d3cd

Browse files
soluwalanaalbcuianubhutivyas
authored
feat(customizer): add grpo support with gym environments (#1156)
* feat(customizer): add opensandbox installation and verification Signed-off-by: Sam Oluwalana <soluwalana@nvidia.com> * feat(customizer): add grpo support with gym environments Signed-off-by: Sam Oluwalana <soluwalana@nvidia.com> * Cleanup doc strings Signed-off-by: Sam Oluwalana <soluwalana@nvidia.com> * fix(rl): clear pre-existing lint and type-check debt `ruff check` and the `ty` pre-commit hook both fail on this branch today. The ty hook only inspects staged files, so the type errors stay invisible until a commit happens to touch one of these files — at which point they block it. - ruff: 8 auto-fixable errors (import sorting, unused imports) plus formatting across services/rl. - test_compiler: the compiler returns TypedDicts typed as PlatformJobStepSpecParam, which is not assignable to `dict[str, Any]`, and `spec["steps"]` is typed as an Iterable so it cannot be indexed. Loosen the helper signatures and go through a concrete list. - test_compiler: pass `OutputNameType.MODEL` rather than the bare string, and drop a dead attribute branch — step env entries are TypedDicts, so they are plain dicts at runtime. - test_schemas: annotate the parametrized kwargs dict so the heterogeneous bad-value fixtures stop failing per-field argument checks. - config: suppress `unsupported-base` for the dynamic service-config base, matching the existing precedent in services/studio/src/nmp/studio/config.py. No behavior change. Signed-off-by: Albert Cui <albcui@nvidia.com> * refactor(jobs): extract job_storage_subpath into a shared helper The Kubernetes jobs backend spells out `jobs/{workspace}/{job}` in two places when mounting the shared job-storage PVC. Anything that hands the same data to a second consumer mounting that PVC itself has to reproduce the layout exactly, so put it behind one helper next to DEFAULT_JOB_STORAGE_PATH rather than have each call site restate it. No behavior change; the follow-up sandboxed-GRPO work is the first outside consumer. Signed-off-by: Albert Cui <albcui@nvidia.com> * fix(rl): make sandboxed GRPO runnable end to end Sandboxed GRPO is the platform default and the only mode the compiler allows once a cluster is sandbox-capable, but nothing on that path reached a training step. Verified against the NeMo-RL commit the image pins (NEMO_RL_REF in docker-bake.hcl). Sandbox provisioning - SandboxConfig omitted environment_pvc_claim and workspace_pvc_claim, which upstream declares required with no default under `extra="forbid"`, so the Gym host failed validation after the model download and vLLM startup. Add the claim/sub-path triples, sourced from a new NMP_RL_JOB_STORAGE_PVC_CLAIM setting the compiler injects — the training container is handed the storage path but never the claim behind it. Fail at compile time when it is unset. - Emit job_id. Upstream defaults it to a shared constant and stamps it onto every sandbox pod as a label, so concurrent jobs were indistinguishable. Config compilation - data.train.data_path pointed at the sandbox-internal /job/dataset mount, but NeMo-RL opens it in the training master, where that path does not exist. Keep the job-storage path; the sandbox sees the same file at its own mount. - config_paths was written only on the colocated branch, so sandboxed Gym started with no server instances. Read it from the manifest in both modes. - A Gym dataset only has to ship training.jsonl, but validation was left enabled, and upstream asserts a validation dataset exists whenever val_period/val_at_start/val_at_end is set. Disable validation and best-checkpoint ranking when there is no split, and say so in the log. - Reject num_prompts_per_step * num_generations_per_prompt that is not a multiple of batch_size at submit time; it previously asserted at the first optimizer step, after a full rollout. Environment packaging - upload passed api_key= to FilesClient, which takes auth=, so --upload raised TypeError before issuing a request. - A validation_fraction that consumed every row silently made the training and validation files identical, and GRPO then ranks checkpoints on its own training data. Reject it, and bound the CLI flag to [0, 1). - Pick the hub wheel by parsed version rather than filename order, so 0.10.0 wins over 0.9.0 and host-side generation matches what the cluster installs. - Default the offline-install venv outside the package root: site-packages ships .jsonl files, which the layout validator rejects, so job resume and rerun failed on the same PVC. - Guard the optional example_id column, and compare config_paths containment by path components — "environment-attacker" has "environment" as a prefix. Schema and API - training.type has no default. It is the union discriminator, so it is already required on the wire; defaulting it made the generated OpenAPI advertise the field as optional while the server rejected it. - Restore the field descriptions dropped when schemas.py was split into a package, and document the GRPO fields. They drive `rl explain`, the OpenAPI spec, and SDK docstrings. - Expose GRPO max_grad_norm, which the compiler read but never populated. - Add the environment fileset purpose and its metadata content type. OpenSandbox operator assets - install.sh ran under `set -x` and passed the generated API key via --from-literal, so it reached the trace, the terminal, any CI log, and kubectl's argv. Disable xtrace around key handling and pass a 0600 temp file. - The crun template paired a blanket `tolerations: [{operator: Exists}]` with only preferred anti-affinity, so shared-kernel sandboxes running untrusted environment code could land on control-plane nodes under capacity pressure. Add a required rule excluding them. - verify/lib.sh took a bare command substitution for its readiness poll, so under `set -e` one transient curl failure aborted the script and defeated READY_TIMEOUT_S. Signed-off-by: Albert Cui <albcui@nvidia.com> * chore(sdk): regenerate OpenAPI spec and Python SDK Picks up the environment fileset purpose, EnvironmentMetadataContent, and the now-required training.type discriminator. EnvironmentMetadataContent is pinned in $shared rather than left under files/filesets. Every FilesetMetadata member has to live alongside it: homed under `files`, FilesetMetadata's import of it becomes shared -> files -> shared and the generated package fails to import. This mirrors the existing pin on FilesetMetadata itself. Generated by `make refresh-openapi` and `make stainless`. Signed-off-by: Albert Cui <albcui@nvidia.com> * fix(rl): stop selecting the V2 DTensor worker for GRPO `policy.dtensor_cfg._v2` is not a perf toggle — it decides which Ray actor runs policy training, and therefore which venv that actor launches under: _v2: True -> DTensorPolicyWorkerV2 -> PY_EXECUTABLES.AUTOMODEL -> --extra automodel _v2 unset -> DTensorPolicyWorker -> PY_EXECUTABLES.FSDP -> --extra fsdp docker/rl/Dockerfile.nmp-rl-base deliberately does not build the automodel extra: with mcore it is the only source of Transformer-Engine, the longest CUDA compile in that image. The venv prefetch cannot cover for it either — the filter is the substring "dtensor_policy_worker.DTensorPolicyWorker", which does not match the V2 FQN (it has "_v2." where the filter expects "."). So every GRPO job would have compiled nemo-automodel + TE on the node at startup, and failed outright on a cluster without egress to PyPI. This also contradicted our own docs: docker/rl/README.md states DPO and GRPO both train on DTensorPolicyWorker / `fsdp`. DPO was already correct because it never sets the flag. Drop it rather than make it configurable. The documented reason to want V2 is dp_replicate_size > 1, which we never set, and dp_replicate_size is a parallelism concern — so if we adopt Automodel the natural interface is to derive _v2 from a user-facing parallelism field, not an operator boolean. Designing that knob now, before Automodel support exists, would likely pick the wrong shape and leave a deprecated flag in the config schema and Helm values. Add a regression test pinning the config -> actor -> venv coupling for all three backend selections (dtensor, megatron, generation), since re-adding the flag looks harmless in isolation. Confirmed it fails when _v2 is restored. When Automodel is added to the image, that test is the one place to flip — and doing so forces a look at the prefetch filters. Signed-off-by: Albert Cui <albcui@nvidia.com> * make lint-fix Signed-off-by: Albert Cui <albcui@nvidia.com> * chore(licenses): add overrides for antlr4-python3-runtime and wget Both arrive transitively with the GRPO Gym dependencies and osv-scanner reports each as NON-STANDARD, so `make check-licenses` failed with two packages needing overrides. antlr4-python3-runtime is BSD-3-Clause (already an approved license), so it goes with the other parser runtimes. wget declares "License: Public Domain" with no corresponding SPDX identifier, so it goes in the "Licenses not automatically approved" block alongside the existing Unlicense entries. Dependency chains: antlr4-python3-runtime <- latex2sympy2-extended <- math-verify wget <- verifiers (NeMo Gym) Regenerating also picks up nmp-jobs in the nemo_platform_plugin reverse dependency list, which the job_storage_subpath refactor introduced. Signed-off-by: Albert Cui <albcui@nvidia.com> * Accomodate GRPO + LoRA Signed-off-by: Sam Oluwalana <soluwalana@nvidia.com> * lint fix Signed-off-by: Sam Oluwalana <soluwalana@nvidia.com> * fix(rl): select the V2 DTensor worker when GRPO uses LoRA `Accomodate GRPO + LoRA` added lora_cfg to policy.dtensor_cfg but left `_v2` unset, and its own test asserted `_v2 is True` — so the branch had a failing test pointing straight at the gap. LoRA is implemented only in DTensorPolicyWorkerV2. At the pinned NeMo-RL SHA (ace40313), dtensor_policy_worker.py has zero references to lora, while dtensor_policy_worker_v2.py has 23, and lm_policy.py picks between them on `config["dtensor_cfg"].get("_v2", False)`. Handing lora_cfg to V1 therefore does not fail — it silently trains full weights and reports success. `_v2` is set only when LoRA is enabled, not unconditionally, because the two workers resolve to different venvs in ray_actor_environment_registry: V1 -> PY_EXECUTABLES.FSDP, which the image prefetches, and V2 -> PY_EXECUTABLES.AUTOMODEL, which it does not. Full-weight GRPO keeps the prefetched path and is unchanged; only LoRA runs take on the automodel dependency. Until that venv is prefetched, a LoRA actor builds it on the node at job start — slow at best, a hard failure on a deny-egress cluster. The Dockerfile comment claiming automodel is "unused by our DPO+GRPO path" is now false, so it is corrected to describe what LoRA needs. Also fixes five LoRA schema tests that construct GRPOTraining without `type`. Four failed outright, since the discriminator is required and validation rejected them before reaching the LoRA validators under test. The fifth, test_grpo_lora_rejects_lora_merged, passed for the wrong reason: it caught the missing-`type` error rather than the invalid finetuning_type, and would have passed even if lora_merged were accepted. Adds test_lora_and_v2_stay_coupled to pin the relationship in both directions; the existing tests each cover one case but neither fails if the two are decoupled. Signed-off-by: Albert Cui <albcui@nvidia.com> * refactor(rl): omit lora_cfg entirely for full-weight GRPO Both `lora_cfg` and `_v2` are NotRequired in NeMo-RL's DTensorConfig, and both read sites use `.get("lora_cfg", {}).get("enabled", False)` / `.get("_v2", False)`, so a full-weight run can leave them out rather than emit `lora_cfg.enabled: False`. The compiled YAML now says only what applies. Emitting the pair together also makes the coupling structural: `_v2` selects DTensorPolicyWorkerV2, the only worker implementing LoRA, so the two keys can no longer drift apart through an edit that touches one line and not the other. Corrects the reasoning in the previous commit: V1 does not silently train full weights when handed an enabled lora_cfg. lm_policy.py asserts "LoRA is not supported for DTensorPolicyWorker V1" before building the worker, so the failure is loud. The fix stands -- only the described failure mode was wrong. Tests now read the config the way NeMo-RL does, so "key omitted" and "enabled: False" are equivalent to them and they do not pin which we emit. Signed-off-by: Albert Cui <albcui@nvidia.com> * style(rl): drop the LoRA/_v2 commentary from the config builder The dtensor_cfg block reads as data again. Why `_v2` tracks LoRA is covered by test_lora_and_v2_stay_coupled, which fails in both directions, and the image dependency it creates is documented at the prefetch filters in docker/rl/Dockerfile.nmp-rl-base, where acting on it belongs. Signed-off-by: Albert Cui <albcui@nvidia.com> * fix(rl): emit the MasterConfig fields the GRPO driver requires A real GRPO run on dev-blue died at driver startup with four pydantic ValidationErrors against NeMo-RL's MasterConfig: policy.generation.top_k policy.generation.stop_token_ids policy.generation.stop_strings policy.make_sequence_length_divisible_by All four are required with no default. Values mirror upstream's reference config examples/configs/grpo_math_1B.yaml, which sets the three generation fields to null and make_sequence_length_divisible_by to the tensor-parallel size. None is meaningful rather than absent: generation/__init__.py fills stop_token_ids with [tokenizer.eos_token_id] when it sees None. Why this survived until a live run: dpo_config already sets make_sequence_length_divisible_by, and the three generation fields are GRPO-only because DPO runs no vLLM, so no existing path touched them. The compiler's tests assert individual fields rather than validating the emitted config against MasterConfig, and MasterConfig is not importable here (it needs torch), so nothing local could have caught it. The cost of that gap is the reason for the added test: the failure surfaces in the driver, after Ray is up and four pods in, minutes of GPU time after submit. Signed-off-by: Albert Cui <albcui@nvidia.com> * fix(rl): emit the required vllm_cfg keys for GRPO generation Second live failure on dev-blue, one layer past the MasterConfig fix: nmp.rl.tasks.training.errors.parser.KeyError: 'kv_cache_dtype' raised from grpo.py where it reads generation_config["vllm_cfg"]["kv_cache_dtype"] directly. vllm_cfg is a TypedDict (VllmSpecificArgs) whose keys MasterConfig does not validate, so a missing member is not a config error at load — it is a KeyError wherever vLLM first reads it, after Ray is already up. Rather than fix the one key that failed, this adds every non-NotRequired member we were missing. Of the eight required by the TypedDict we emitted five; kv_cache_dtype and expert_parallel_size were absent. expert_parallel_size is read unconditionally at vllm_generation.py:109 and vllm_worker.py:304, so it was the next failure in line. skip_tokenizer_init is required by the TypedDict but deliberately still omitted: generation/__init__.py fills it in when absent, choosing based on stop_strings and expose_http_server, and hardcoding it would override logic that exists for VLMs. The test asserts its absence so a later "complete the required fields" change does not silently break that. Values match upstream's examples/configs/grpo_math_1B.yaml. Signed-off-by: Albert Cui <albcui@nvidia.com> * chore: rebase off of main + resolve conflicts + regenrate Signed-off-by: anubhutiv <anubhutiv@nvidia.com> * fix(rl): omit chat_template when the model ships none NeMo-RL's TokenizerConfig declares chat_template: NotRequired[str] -- absent is valid, None is not. resolve_chat_template returns None whenever the model has no built-in template and the user supplied no override (chat_templates.py:173, right after logging "No chat template found"), and both config builders emitted the key unconditionally. Any such model therefore failed MasterConfig validation in the driver with "Input should be a valid string", after Ray was already up. Extracted as a shared helper because DPO and GRPO had the identical block and the identical bug. Not caught by the dev-blue run because Qwen3 ships a template; it needs a model that does not, which is the kind of gap a single-model GPU run cannot close. Found instead by validating the compiled config against the real MasterConfig locally. Signed-off-by: Albert Cui <albcui@nvidia.com> * test(rl): gate the image build on the compiled config matching NeMo-RL compile_grpo_config and NeMo-RL's schemas only ever share an interpreter inside the training image: services/rl/tests cannot import MasterConfig (it needs torch) and NeMo-RL cannot see the compiler. So the two halves of the contract were never checked against each other, and two config-shape bugs reached a live GPU run -- each costing an image build and a redeploy to find, and each surfacing only after Ray was already up. Runs in the existing smoke-test stage under the smoke_nmp_rl_training marker, so no Dockerfile change: line 82 already runs everything carrying it. Adds ~12s to a stage that already installs pytest and runs on a CPU runner, and failing it fails the nmp-rl-training-smoke-test bake target, so a broken image is not published. Checks MasterConfig construction and VllmSpecificArgs.__required_keys__ (a TypedDict pydantic does not validate the members of), for both the full-weight and LoRA paths, plus the lora_cfg/_v2 coupling. Running inside the image also removes the skew that makes an external validator untrustworthy: an image with pydantic <2.13 accepts configs the deployed one rejects, because 2.13 enforces required keys on nested TypedDicts and 2.12 does not. Here the validator and the validated code are the same artifact. Verified by overlaying current source onto an existing image: 5 passed, and 5 failed when kv_cache_dtype, make_sequence_length_divisible_by and _v2 were removed. The fixture is duplicated rather than shared with services/rl/tests, which is not copied into the image. It only needs one representative config, not the parametrizable one the unit tests build. Signed-off-by: Albert Cui <albcui@nvidia.com> * chore(customizer): add automodel and mcore backends in RL image Signed-off-by: anubhutiv <anubhutiv@nvidia.com> * fix: no module named nemo_gym during training Signed-off-by: anubhutiv <anubhutiv@nvidia.com> * fix: sandbox create timeout + sandbox runtime image + fix relative config paths Signed-off-by: anubhutiv <anubhutiv@nvidia.com> * fix: add policy_model.yaml Signed-off-by: anubhutiv <anubhutiv@nvidia.com> * fix: hook up allow_internet from platform to rl Signed-off-by: anubhutiv <anubhutiv@nvidia.com> * fix: allow configuring the allowlisted domains Signed-off-by: anubhutiv <anubhutiv@nvidia.com> * fix: propagate the venv correctly Signed-off-by: anubhutiv <anubhutiv@nvidia.com> * fix: minor bug Signed-off-by: anubhutiv <anubhutiv@nvidia.com> * fix: make opensandbox server scheme configurable + conversion script fixes Signed-off-by: anubhutiv <anubhutiv@nvidia.com> * fix: find the correct lora adapter Signed-off-by: anubhutiv <anubhutiv@nvidia.com> * chore: expose val_at_start to validate baseline model Signed-off-by: anubhutiv <anubhutiv@nvidia.com> * expose more grpo configs + fix schema Signed-off-by: anubhutiv <anubhutiv@nvidia.com> * expose the sandbox resource Signed-off-by: anubhutiv <anubhutiv@nvidia.com> * chore: make sandbox pod ttl configurable Signed-off-by: anubhutiv <anubhutiv@nvidia.com> * rebase onto main + regenerate licenses and sdk Signed-off-by: anubhutiv <anubhutiv@nvidia.com> * chore: update RL repo ref + lint Signed-off-by: anubhutiv <anubhutiv@nvidia.com> * chore: expose more configs for moe models Signed-off-by: anubhutiv <anubhutiv@nvidia.com> * chore: address comments Signed-off-by: anubhutiv <anubhutiv@nvidia.com> * fix: add automodel patch Signed-off-by: anubhutiv <anubhutiv@nvidia.com> * fix: update how we patch Signed-off-by: anubhutiv <anubhutiv@nvidia.com> * chore: address comments + patch for nemotron Signed-off-by: anubhutiv <anubhutiv@nvidia.com> * fix: fix log_interval + address comments Signed-off-by: anubhutiv <anubhutiv@nvidia.com> * update licenses Signed-off-by: anubhutiv <anubhutiv@nvidia.com> * fix lint Signed-off-by: anubhutiv <anubhutiv@nvidia.com> --------- Signed-off-by: Sam Oluwalana <soluwalana@nvidia.com> Signed-off-by: Albert Cui <albcui@nvidia.com> Signed-off-by: anubhutiv <anubhutiv@nvidia.com> Co-authored-by: Albert Cui <albcui@nvidia.com> Co-authored-by: anubhutiv <anubhutiv@nvidia.com>
1 parent ca93cab commit 730d3cd

97 files changed

Lines changed: 8751 additions & 2466 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

docker/rl/Dockerfile.nmp-rl-base

Lines changed: 63 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,8 @@
88
# from source (rather than the prebuilt NGC nemo-rl image) lets us pin any RL commit via
99
# NEMO_RL_REF; Gym follows as RL's submodule pin.
1010
#
11-
# DPO trains on `fsdp`; GRPO trains on `fsdp` and generates on `vllm`. mcore, automodel, sglang
12-
# and trtllm are deliberately not built.
11+
# DPO trains on `fsdp`. GRPO trains on `automodel` (DTensor V2) or `mcore`
12+
# (Megatron) and generates on `vllm`. sglang and trtllm are deliberately not built.
1313
#
1414
# Two things drive most of what follows, both covered in README.md:
1515
# - Workers do not run in the base venv. The `uv sync --extra` calls only warm the uv cache; the
@@ -129,9 +129,9 @@ ENV RAY_USAGE_STATS_ENABLED=0 \
129129
NEMO_GYM_VENV_DIR=/opt/gym_venvs
130130

131131
# ---- builder: uv sync RL + Gym with the extras we use ----
132-
# Only fsdp (DPO/GRPO training), vllm (GRPO generation), nemo_gym and modelopt are built. The mcore
133-
# and automodel extras are deliberately EXCLUDED: alternative training backends the customizer does
134-
# not use, and the only source of Transformer-Engine (the longest CUDA compile).
132+
# fsdp (DPO training), automodel (GRPO training on DTensor V2, the only LoRA-capable DTensor
133+
# backend), mcore (Megatron training + Megatron-native generation), vllm (generation), nemo_gym and
134+
# modelopt are built. sglang and trtllm are not: vLLM covers generation for every config we compile.
135135
FROM base AS builder
136136
ARG BUILD_UID=2000
137137
ARG BUILD_GID=2000
@@ -185,9 +185,9 @@ ENV NVTE_BUILD_MAX_JOBS=${NVTE_BUILD_MAX_JOBS} \
185185
# layer's cache key depends on the dependency graph, not on every RL source file. See README.md,
186186
# "Layering for fast CI rebuilds".
187187
#
188-
# TensorRT-LLM is excluded (we never build the trtllm extra). Megatron-Bridge must stay, even though
189-
# the mcore extra is not built: `uv run` validates the lock, which needs metadata for every path
190-
# source, so removing it fails the prefetch step below.
188+
# TensorRT-LLM is excluded (we never build the trtllm extra). Megatron-Bridge must stay: it is the
189+
# mcore extra's entry point, and its nested Megatron-LM submodule supplies megatron-core. Both are
190+
# editable path sources, so `uv run` also needs their metadata to validate the lock.
191191
COPY --from=nemo-rl --chown=${BUILD_UID}:${BUILD_GID} pyproject.toml uv.lock ./
192192
COPY --from=nemo-rl --chown=${BUILD_UID}:${BUILD_GID} nemo_rl/__init__.py nemo_rl/package_info.py ./nemo_rl/
193193
COPY --from=nemo-rl --chown=${BUILD_UID}:${BUILD_GID} --exclude=TensorRT-LLM-workspace --exclude=TensorRT-LLM-workspace/** \
@@ -202,12 +202,16 @@ USER ${BUILD_UID}:${BUILD_GID}
202202
ARG UV_SYNC_MODE=--frozen
203203

204204
# Warm the uv CACHE, one extra at a time (they conflict). This is where the heavy compiles happen
205-
# (mamba-ssm, causal-conv1d, deep_ep, deep_gemm); the wheels
205+
# (mamba-ssm, causal-conv1d, deep_ep, deep_gemm, and Transformer-Engine via automodel); the wheels
206206
# land in the cache so the prefetch below materializes each venv without recompiling.
207207
#
208208
# `uv sync` is EXACT, so these do NOT accumulate - each --extra prunes the previous one and the base
209209
# venv ends up with default deps + groups only. The populated cache is the real product.
210210
#
211+
# Conflicting extras resolve in separate forks but still share cache entries when a fork lands on the
212+
# same pinned source, so mcore reuses the Transformer-Engine wheel automodel built. See README.md,
213+
# "Why the `mcore` backend is built".
214+
#
211215
# --no-install-project installs dependencies but not nemo-rl itself; the root package is added
212216
# editable in a cheap step after the full-source COPY.
213217
RUN <<"EOF" bash -exu
@@ -242,7 +246,9 @@ grep -Fq \
242246
"${list_inl}"
243247

244248
uv sync ${UV_SYNC_MODE} --extra vllm --no-install-project # GRPO generation (vllm 0.20 cu130 + flashinfer + deep_gemm/deep_ep)
245-
uv sync ${UV_SYNC_MODE} --extra fsdp --no-install-project # DPO/GRPO policy training (flash-attn, mamba-ssm, causal-conv1d)
249+
uv sync ${UV_SYNC_MODE} --extra fsdp --no-install-project # DPO policy training (flash-attn, mamba-ssm, causal-conv1d)
250+
uv sync ${UV_SYNC_MODE} --extra automodel --no-install-project # GRPO policy training, DTensor V2 (+ LoRA); builds Transformer-Engine
251+
uv sync ${UV_SYNC_MODE} --extra mcore --no-install-project # Megatron training + Megatron-native generation; reuses automodel's TE
246252
uv sync ${UV_SYNC_MODE} --extra modelopt --no-install-project # quantization / model-opt
247253
uv sync ${UV_SYNC_MODE} --extra nemo_gym --no-install-project # NeMo-Gym (uv workspace member)
248254
uv sync ${UV_SYNC_MODE} --all-groups --no-install-project # build/test/etc groups
@@ -256,6 +262,19 @@ uv pip install --python /opt/nemo_rl_venv/bin/python "opensandbox>=0.1.9" "tenac
256262
# and Go stdlib floors even when NeMo-RL's lock has an older wheel.
257263
uv pip install --python /opt/nemo_rl_venv/bin/python "wandb==0.28.2"
258264

265+
# NeMo-Gym in the BASE venv, for the DRIVER - not a duplicate of the nemo_gym extra above.
266+
# Sandboxed (mode B) GRPO imports nemo_rl.environments.sandbox in the driver process:
267+
# host/models.py needs BROKER_TOKEN_ENV / BROKER_URL_ENV and backends/base.py needs
268+
# EpisodeResources, both from nemo_gym.sandbox.broker. The `--all-groups` sync above is EXACT, so
269+
# it prunes the nemo_gym extra back out of the base venv and the driver dies with
270+
# `ModuleNotFoundError: No module named 'nemo_gym'` at Gym spin-up. Worker venvs are unaffected -
271+
# they get nemo_gym from the extra.
272+
# --no-deps is safe and deliberate: the driver's import closure is
273+
# nemo_gym.sandbox.broker -> broker/wire.py -> providers/base.py, whose only third-party import is
274+
# pydantic (already present). Resolving Gym's full dependency set here could move RL's pinned stack.
275+
uv pip install --python /opt/nemo_rl_venv/bin/python --no-deps \
276+
-e /opt/nemo-rl/3rdparty/Gym-workspace/Gym
277+
259278
# CVE GHSA-mqqc-3gqh-h2x8: drop ray's bundled old aiohttp from the uv cache.
260279
find "${UV_CACHE_DIR}" -type d -path "*ray/_private/runtime_env/agent/thirdparty_files/aiohttp*" -exec rm -rf {} + || true
261280

@@ -282,6 +301,21 @@ EOF
282301
# Excluding here keeps them out of this layer - removing them later would leave them extractable.
283302
COPY --from=nemo-rl --chown=${BUILD_UID}:${BUILD_GID} --exclude=**/.git . /opt/nemo-rl
284303

304+
# TEMPORARY: NemotronH on DTensor v2 dies in Automodel's parallelizer, which reaches for
305+
# model.backbone.layers while both the stock-HF and Automodel NemotronH name the trunk
306+
# `model` (only the HF checkpoint KEYS use a backbone. prefix). Fixed on Automodel main
307+
# (parallelizer.py: `inner = model.backbone if hasattr(model, "backbone") else model.model`)
308+
# but NVIDIA-NeMo/RL main still pins 24b47e8, which predates it. Drop this once the
309+
# submodule is bumped past that fix.
310+
#
311+
# Bind-mounted rather than COPY-ed: this stage already runs as ${BUILD_UID}, so a COPY would
312+
# land root-owned and could not be cleaned up afterwards, and anything left behind ships in
313+
# the layer. The mount keeps the patch out of the image entirely.
314+
RUN --mount=type=bind,source=docker/rl/patches/automodel-nemotronh-trunk.patch,target=/patches/nemotronh.patch \
315+
cd /opt/nemo-rl && \
316+
git apply --check -p1 /patches/nemotronh.patch && \
317+
git apply -p1 --verbose /patches/nemotronh.patch
318+
285319
# Install the workspace ROOT (nemo-rl) editable. --no-install-project above installed every
286320
# DEPENDENCY and workspace MEMBER (incl. nemo-gym) but skipped the root; this adds it with --no-deps
287321
# (deps already present). .git is on disk now, so nemo_rl.__version__ picks up its +<sha> suffix.
@@ -293,8 +327,13 @@ RUN umask 022 && uv pip install --python /opt/nemo_rl_venv/bin/python --no-deps
293327
# cache. See README.md, "Prefetching the per-worker venvs".
294328
#
295329
# Filters are SUBSTRING matches on actor FQNs, and one venv is built PER ACTOR, so actors sharing
296-
# an extra still need their own filter. Six filters -> seven venvs:
297-
# dtensor_policy_worker.DTensorPolicyWorker -> fsdp (DPO/GRPO training)
330+
# an extra still need their own filter. Eight filters -> nine venvs:
331+
# dtensor_policy_worker.DTensorPolicyWorker -> fsdp (DPO training)
332+
# dtensor_policy_worker_v2.DTensorPolicyWorkerV2 -> automodel (GRPO training; the only LoRA-capable
333+
# DTensor worker). The V1 filter does not match it - the FQN reads `dtensor_policy_worker_v2.`.
334+
# megatron_policy_worker.MegatronPolicyWorker -> mcore (GRPO training on Megatron; also
335+
# hosts Megatron-native generation, which runs in-process rather than as its own actor).
336+
# Does not match modelopt's `megatron_quant_policy_worker.MegatronQuantPolicyWorker`.
298337
# vllm.vllm_worker -> vllm (GRPO generation) - matches both the
299338
# sync and async workers, which is wanted: NeMo-Gym forces async rollouts.
300339
# sync_rollout_actor.SyncRolloutActor -> vllm (GRPO rollout driver)
@@ -303,13 +342,15 @@ RUN umask 022 && uv pip install --python /opt/nemo_rl_venv/bin/python --no-deps
303342
# broker_actor.SandboxEpisodeBrokerActor -> nemo_gym (per-episode sandbox broker)
304343
# Kept specific on purpose: a bare "vllm" would also match modelopt's vllm_quant_worker.
305344
#
306-
# NOT prefetched: automodel, mcore, sglang, trtllm, modelopt-quant
307-
# and the async-GRPO bookkeeping actors; those build on the node if a config selects them.
345+
# NOT prefetched: sglang, trtllm, modelopt-quant, MegatronValueWorker / DTensorValueWorkerV2 (GRPO
346+
# is critic-free) and the async-GRPO bookkeeping actors; those build on the node if selected.
308347
RUN <<"EOF" bash -exu
309348
umask 022
310349
export UV_LINK_MODE=symlink
311350
uv run nemo_rl/utils/prefetch_venvs.py \
312351
dtensor_policy_worker.DTensorPolicyWorker \
352+
dtensor_policy_worker_v2.DTensorPolicyWorkerV2 \
353+
megatron_policy_worker.MegatronPolicyWorker \
313354
vllm.vllm_worker \
314355
sync_rollout_actor.SyncRolloutActor \
315356
nemo_gym.NemoGym \
@@ -321,6 +362,9 @@ uv run nemo_rl/utils/prefetch_venvs.py \
321362
# The vLLM-tier actors are listed too: the publish stage privatizes `vllm/` in whatever venvs exist,
322363
# so a filter that stops producing one would leave that worker sharing ray_executor.py again.
323364
for actor in \
365+
nemo_rl.models.policy.workers.dtensor_policy_worker.DTensorPolicyWorker \
366+
nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2 \
367+
nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorker \
324368
nemo_rl.environments.nemo_gym.NemoGym \
325369
nemo_rl.environments.sandbox.nemo_gym_actor.SandboxedGymActor \
326370
nemo_rl.environments.sandbox.broker_actor.SandboxEpisodeBrokerActor \
@@ -427,6 +471,11 @@ ARG RUNTIME_GID=1000
427471
# their own dependencies, so RL's `setuptools>=80.10.2` override cannot reach them. The result is an
428472
# older setuptools left in the shipped cache that no venv uses. Runs here, after every venv exists,
429473
# and only drops archives nothing symlinks into.
474+
#
475+
# Still correct now that the mcore extra IS built: installing Megatron-Bridge does not change what
476+
# its BUILD environment resolves, so the old setuptools remains cached and unreferenced. Deleting
477+
# only unreferenced archives is what makes that safe - the mcore venv links the locked setuptools,
478+
# so that archive is kept.
430479
RUN <<"EOF" bash -exu
431480
for d in "${UV_CACHE_DIR}"/archive-v0/*/; do
432481
ls "${d}"setuptools-*.dist-info >/dev/null 2>&1 || continue

0 commit comments

Comments
 (0)