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

[anubhutivyas]

Summary

• build the automodel and mcore extras in the RL base image so both GRPO training tiers, DTensor V2 and Megatron
• ship prebuilt instead of being compiled on the node at job time.

Changes

automodel (DTensor V2)

• GRPO's compiler sets policy.dtensor_cfg._v2: true, which selects DTensorPolicyWorkerV2 -> PY_EXECUTABLES.AUTOMODEL. That extra wasn't built, so the venv would have been constructed on the node at first run against a cold cache, including a Transformer-Engine source compile per node, at job time.
• Adds --extra automodel to the warmup syncs.
• Prefetches DTensorPolicyWorkerV2. This is also the only LoRA-capable DTensor tier, V1 asserts lora_cfg.enabled is False in nemo_rl/models/policy/lm_policy.py.

mcore (Megatron)

• Adds --extra mcore to the warmup syncs and prefetches MegatronPolicyWorker.
• mcore backs the GRPO capabilities DTensor does not implement in NeMo-RL today: pipeline parallelism, FP8 training / FP8 rollouts / FP8 KV-cache, NVFP4 quantization-aware RL, draft models + EAGLE3 speculative decoding, Megatron-native generation (no refit weight conversion)
• Transformer-Engine is not rebuilt. 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 keys on resolved source identity (git URL + commit + platform tags), not on the requesting extra, so the mcore sync 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.
• what mcore genuinely adds is cheap: megatron-bridge and megatron-core are editable path sources from submodules already on disk, plus wheel-only flashinfer-*==0.6.8.post1, nvshmem4py-cu13, cupy-cuda13x, and
  nvidia-modelopt from git.

Type of Change

• Code change (feature, bug fix, or refactor)
• Code change with documentation updates
• Documentation only
• Contributor tooling or automation
• CI, build, or test infrastructure

Quality Gates

• Tests added or updated for changed behavior
• Existing tests cover changed behavior — justification:
• Tests not applicable — justification:
• Documentation updated for user-visible behavior
• Documentation not applicable — justification:

Verification

• Pull request title follows the repository's Conventional Commit format
• Every commit includes an appropriate Signed-off-by: trailer
• uv run pre-commit run -a passes, or any blocked checks are identified below
• Targeted tests pass, or tests are marked not applicable above
• No secrets, API keys, or credentials are included

Targeted validation:

Summary by CodeRabbit

• Documentation
  • Updated reinforcement learning container documentation to clarify supported DPO and GRPO training workflows.
  • Documented GRPO’s v2 training path, LoRA support, Transformer Engine compilation, and critic-free execution.
  • Updated environment mappings, job diagrams, prefetch behavior, excluded components, caching details, image-size considerations, and build options.
  • Clarified dependency handling for DPO and GRPO policy training and generation.

[coderabbitai]

📝 Walkthrough

Walkthrough

The RL image now builds and documents Automodel, mcore, Transformer-Engine, and separate DTensor V2 and Megatron policy-worker environments. Smoke tests validate the new worker imports and distributions.

Changes

RL worker environments and build configuration

Layer / File(s)	Summary
Backend and dependency build contracts docker/rl/Dockerfile.nmp-rl-base, docker/rl/README.md	DPO uses fsdp. GRPO policy training uses automodel, Megatron training and generation use mcore, and generation uses vllm. Transformer-Engine is built through automodel and reused by mcore.
Worker prefetch and validation docker/rl/Dockerfile.nmp-rl-base, docker/rl/README.md	The image defines, prefetched, and validates separate DTensor V2 and Megatron policy-worker environments. The README updates the GRPO actor mapping and runtime status.
Smoke checks and cache behavior tests/smoke_gpu/test_rl_training.py, docker/rl/Dockerfile.nmp-rl-base, docker/rl/README.md	Smoke tests validate imports and distributions for the new workers. Documentation covers mcore archive retention and cold-runtime Transformer-Engine recompilation.

Sequence Diagram(s)

sequenceDiagram
  participant DockerBuild
  participant WorkerPrefetch
  participant ActorRegistry
  participant SmokeTest
  DockerBuild->>WorkerPrefetch: Build and prefetch worker environments
  WorkerPrefetch->>ActorRegistry: Register DTensor V2 and Megatron workers
  ActorRegistry->>SmokeTest: Expose worker environments
  SmokeTest->>SmokeTest: Check imports and distributions

Loading

Possibly related PRs

• NVIDIA-NeMo/nemo-platform#1187: Updates the same RL Dockerfile and README for Automodel, DTensor worker prefetching, and cache cleanup.

Suggested reviewers: gabwow

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the main change: adding Automodel and MCore backends to the RL image.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch anubhutiv/add-automodel-be

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@docker/rl/README.md`:
- Around line 3-10: Update the Docker RL README image inventory to include the
automodel extra and Transformer-Engine compilation. Also revise the per-worker
Ray venv count from five to eight to reflect the seven filters and V2
policy-worker addition.

🪄 Autofix

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 92bf6bad-027c-48be-b599-df3a56387fde

📥 Commits

Reviewing files that changed from the base of the PR and between 60ad37b and 30046ae.

📒 Files selected for processing (2)

• docker/rl/Dockerfile.nmp-rl-base
• docker/rl/README.md

[github-actions]

Suite	Lines Covered	Line Rate	Branch Rate
Unit Tests	32002/40624	78.8%	63.5%
Integration Tests	18561/38550	48.1%	20.8%

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@docker/rl/Dockerfile.nmp-rl-base`:
- Around line 8-9: Update docker/rl/Dockerfile.nmp-rl-base lines 8-9 to describe
mcore as a built Megatron backend. In docker/rl/README.md lines 260-264, add
MegatronPolicyWorker, VllmAsyncGenerationWorker, SandboxedGymActor, and
SandboxEpisodeBrokerActor; at lines 281-285, identify the diagram as the
Automodel path or include the Megatron policy alternative; at lines 434-442, add
the mcore filter and revise the inventory to eight filters and nine venvs.

🪄 Autofix

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 4e1c24b1-9509-4201-8dab-e4edb8585fb2

📥 Commits

Reviewing files that changed from the base of the PR and between 6d9163f and 64c0df7.

📒 Files selected for processing (3)

• docker/rl/Dockerfile.nmp-rl-base
• docker/rl/README.md
• tests/smoke_gpu/test_rl_training.py