build(automodel): bump Automodel to r0.6.0 by jQizhang · Pull Request #3498 · NVIDIA-NeMo/RL

jQizhang · 2026-08-05T12:44:59Z

What does this PR do?

Upgrade the NeMo-RL DTensor v2 Automodel backend to Automodel r0.6.0, align compatible dependency forks on Transformers 5.12.1, and migrate context parallelism (CP) to Automodel's model-owned ContextParallelSharder / ShardLayout protocol.

For detailed API changes, support boundaries, before/after GRPO and x-token workflows, and experiment results, see docs/design-docs/automodel-context-parallel.md.

Full before/after results for all 10 workloads (pre-upgrade CP1 and post-upgrade CP1/CP2/CP4) are available in the W&B nemorl-cp project.

Summary

Automodel r0.6.0 upgrade

Bump the Automodel submodule from 24b47e85 to 7e9493d7 and track the r0.6.0 branch.
Migrate model construction to MeshContext.build(...), ParallelismSizes, and one DistributedSetup passed to from_pretrained(); follow relocated imports and updated checkpoint APIs.
Wait for asynchronous DCP writes before promoting tmp_step_N, preventing incomplete model or optimizer checkpoints.
Mirror r0.6.0's megatron-fsdp Git constraint, pin TileLang / apache-tvm-ffi 0.1.11, and expose Z3 libraries to symlinked worker environments.
Remove obsolete token-classification, checkpoint-dtype, and Gemma 4 KV-sharing workarounds.

Transformers 5.12.1 upgrade

Raise the shared Transformers resolution to 5.12.1 and add exact pins to compatible backend extras. TRT-LLM retains its backend-constrained version.

Module / extra	`origin/main` declaration	PR declaration	`origin/main` lock	PR lock	Upgrade status
Base / main environment	`>=5.5.0,<5.9.0`	`>=5.5.0,<=5.12.1`	`5.5.4`	`5.12.1`	Upgraded
FSDP	Inherits Base	`==5.12.1`	`5.5.4`	`5.12.1`	Added an exact pin and upgraded
Automodel	`>=5.5.0,<5.6.0`	`==5.12.1`	`5.5.0`	`5.12.1`	Upgraded for Automodel r0.6.0
MCore / Megatron	Depends on Base and transitive constraints from Megatron-Bridge	`==5.12.1`	`5.8.1`	`5.12.1`	Added an exact pin and upgraded
vLLM	No standalone pin; inherits Base / transitive dependencies	`==5.12.1`	`5.8.1`	`5.12.1`	Added an exact pin and upgraded
SGLang	`==5.12.1`	`==5.12.1`	`5.12.1`	`5.12.1`	Unchanged from current `origin/main`
TRT-LLM	No root-level Transformers pin; constrained by TRT-LLM	No root-level Transformers pin; constrained by TRT-LLM	`5.5.4`	`5.5.4`	Unchanged
ModelOpt	Inherits Base	Inherits Base	`5.5.4`	`5.12.1`	Upgraded with Base
NVRx	Inherits Base	Inherits Base	`5.5.4`	`5.12.1`	Upgraded with Base
NeMo Gym	Inherits Base	Inherits Base	`5.5.4`	`5.12.1`	Upgraded with Base

uv conflicts isolate TRT-LLM from extras using Transformers 5.12.1, including SGLang. The SGLang fork also preserves non-CUDA tensor reductions.

Context-parallel integration

Keep canonical RL/loss tensors separate from the private model batch, preserve the direct CP=1 path, and construct a ContextParallelSharder only for CP>1.
Use the actual ShardLayout to shard next-token targets and restore logprobs, top-k/full logits, and loss inputs instead of reconstructing the legacy round-robin layout.
Make gradient fanout explicit: replicated GRPO-style losses use loss × DP, while partitioned x-token KD/CE retains loss × DP × CP; preserve x-token CUDA IPC, heterogeneous TP/CP, projection, alignment, and multi-teacher aggregation.

The upstream CP API is described in Automodel #2937; the distributed setup migration follows Automodel #2266.

Tests and documentation

Add focused CP layout/sharder, loss, checkpoint, Automodel setup/train/data, policy-worker, and SGLang regression tests.
Add an upgrade design document and CP1/CP2/CP4 comparison plots for Gemma 4 E2B, Gemma 4 26B, Qwen3.5 35B-A3B, and Qwen3.5 9B.

Issues

Closes #3319: Automodel bump up + Refactor to use the generic CP interface to be provided by Automodel.

Follow-up: #3842 — Track nightly OOMs after the Automodel r0.6.0 upgrade tracks fixes for the four affected nightly cases after this PR merges.

Usage

There is no new user-facing training API. Existing Automodel recipes continue to set context_parallel_size; NeMo-RL now delegates model-side token layout to Automodel.

Validation

Targeted CUDA/H100 and Ray validation completed during development (before the final rebase) covered CP migration, Automodel data/setup/train, loss, x-token, worker autocast, and distributed x-token TP2/CP1 and TP1/CP2 tests. Ruff and format checks also passed.

After the final rebase, uv lock --check with uv 0.11.28, git diff --check origin/main...HEAD, and python -m py_compile nemo_rl/models/automodel/train.py passed. Current GitHub checks pass for submodule fast-forward, lockfile freshness, copyright, secret scanning, DCO, and PR-title validation.

Short CP1/CP2/CP4 GRPO/DAPO runs show similar reward, generated-token, and generation-KL trends; these are functional and short-run numerical checks, not long-horizon convergence results. Full before/after results for all 10 workloads (pre-upgrade CP1 and post-upgrade CP1/CP2/CP4) are available in the W&B nemorl-cp project.

Additional information

The CP refactor is limited to Automodel DTensor v2; Megatron CP behavior is unchanged.
VLM training with CP>1 remains rejected. X-token distillation keeps contiguous CP loss/IPC windows and does not support sequence packing.
Full-sequence gather remains a correctness-first migration boundary; a fully CP-native loss implementation can be pursued separately.

copy-pr-bot · 2026-08-05T12:45:03Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

github-actions · 2026-08-05T12:46:28Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 902481f (PR #3498 from bump-automodel-r0.6.0)

✅ Submodules that are properly updated:

Automodel: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

github-actions · 2026-08-10T15:33:13Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 4743058 (PR #3498 from bump-automodel-r0.6.0)

✅ Submodules that are properly updated:

Automodel: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

github-actions · 2026-08-19T10:05:26Z

✅ Submodule Fast-Forward Check Results

Check based on commit: b5ed69c (PR #3498 from bump-automodel-r0.6.0)

✅ Submodules that are properly updated:

Automodel: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

github-actions · 2026-08-24T10:30:04Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 98c3f04 (PR #3498 from bump-automodel-r0.6.0)

✅ Submodules that are properly updated:

Automodel: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

github-actions · 2026-08-25T09:45:43Z

✅ Submodule Fast-Forward Check Results

Check based on commit: f68880b (PR #3498 from bump-automodel-r0.6.0)

✅ Submodules that are properly updated:

Automodel: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

github-actions · 2026-08-26T12:33:41Z

✅ Submodule Fast-Forward Check Results

Check based on commit: bf22989 (PR #3498 from bump-automodel-r0.6.0)

✅ Submodules that are properly updated:

Automodel: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

github-actions · 2026-08-26T13:29:07Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 186f945 (PR #3498 from bump-automodel-r0.6.0)

✅ Submodules that are properly updated:

Automodel: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

Signed-off-by: larkzhang-nv <larkz@nvidia.com>

model_forward now takes the Automodel-prepared model_batch dict, but five TestModelForward multimodal tests still passed a ProcessedInputs dataclass, so model(**model_batch) raised TypeError and aborted the file under -x. Route them through the existing _prepare_cp1 helper and assert on prepared.model_batch. filter_multimodal_kwargs_for_model now runs inside _build_model_batch, so the mixed-resolution guard is asserted around prepare_model_forward rather than model_forward; verified by mutation that disabling the guard fails the test. Also apply ruff import-sort and format fixes to the touched test files. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: larkzhang-nv <larkz@nvidia.com>

Signed-off-by: larkzhang-nv <larkz@nvidia.com>

SGLang now uses Transformers 5.12.1, matching the ModelOpt, NVRx, and NeMo Gym environments. Remove the uv conflict pairs that were only needed while SGLang remained on Transformers 5.6.0. Signed-off-by: larkzhang-nv <larkz@nvidia.com>

Signed-off-by: larkzhang-nv <larkz@nvidia.com>

jQizhang · 2026-08-28T06:23:54Z

/ok to test 254ec1a

github-actions · 2026-08-28T06:24:18Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 254ec1a (PR #3498 from bump-automodel-r0.6.0)

✅ Submodules that are properly updated:

Automodel: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

Signed-off-by: larkzhang-nv <larkz@nvidia.com>

github-actions · 2026-08-28T09:41:03Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 30c26e4 (PR #3498 from bump-automodel-r0.6.0)

✅ Submodules that are properly updated:

Automodel: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

jQizhang · 2026-08-28T09:43:58Z

/ok to test 30c26e4

sharonyu-115

The test case change looks good. Approve.

Signed-off-by: larkzhang-nv <larkz@nvidia.com>

github-actions · 2026-08-28T15:25:56Z

✅ Submodule Fast-Forward Check Results

Check based on commit: f56b445 (PR #3498 from bump-automodel-r0.6.0)

✅ Submodules that are properly updated:

Automodel: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

jQizhang · 2026-08-28T15:28:03Z

/ok to test f56b445

github-actions Bot added the Documentation Improvements or additions to documentation label Aug 5, 2026

jQizhang changed the title ~~feat(automodel): upgrade to r0.6.0 and migrate context parallelism~~ build(automodel): bump Automodel to r0.6.0 Aug 5, 2026

jQizhang force-pushed the bump-automodel-r0.6.0 branch from f3fefe7 to 4743058 Compare August 10, 2026 15:31

zpqiu mentioned this pull request Aug 17, 2026

feat(models): add Qwen3.8-27B Megatron support #3678

Merged

4 tasks

jQizhang force-pushed the bump-automodel-r0.6.0 branch from 28705e2 to b5ed69c Compare August 19, 2026 10:04

sharonyu-115 mentioned this pull request Aug 20, 2026

[Bug] save_consolidated times out and truncates Gemma 4 31B safetensors with FSDP2/DTensor in v0.7.0 #3624

Open

terrykong mentioned this pull request Aug 21, 2026

feat(sglang): megatron backend weight refit for sglang rollouts #3612

Merged

jQizhang force-pushed the bump-automodel-r0.6.0 branch 2 times, most recently from 22ea60a to 98c3f04 Compare August 24, 2026 10:28

jQizhang force-pushed the bump-automodel-r0.6.0 branch from 98c3f04 to f68880b Compare August 25, 2026 09:44

jQizhang force-pushed the bump-automodel-r0.6.0 branch 2 times, most recently from 5ae4753 to bf22989 Compare August 26, 2026 12:32

jQizhang force-pushed the bump-automodel-r0.6.0 branch from bf22989 to 186f945 Compare August 26, 2026 13:26

jQizhang marked this pull request as ready for review August 26, 2026 13:28

jQizhang requested review from a team as code owners August 26, 2026 13:28

jQizhang and others added 12 commits August 27, 2026 23:03

build: use Transformers 5.12.1 for compatible backends

c33f436

Signed-off-by: larkzhang-nv <larkz@nvidia.com>

fix(sglang): preserve non-CUDA tensor reductions

d592efc

Signed-off-by: larkzhang-nv <larkz@nvidia.com>

build: pin transformers for compatible backends

f777c07

Signed-off-by: larkzhang-nv <larkz@nvidia.com>

fix: expose Z3 to symlinked worker venvs

5796a09

Signed-off-by: larkzhang-nv <larkz@nvidia.com>

fix: address review-pr-team feedback

eb5fb10

Signed-off-by: larkzhang-nv <larkz@nvidia.com>

build: bump Automodel r0.6.0

d57c710

Signed-off-by: larkzhang-nv <larkz@nvidia.com>

fix(automodel): check FLA for dense Qwen3.5 CP

4fe7197

Signed-off-by: larkzhang-nv <larkz@nvidia.com>

docs(automodel): update context parallel guide

9121ea2

Signed-off-by: larkzhang-nv <larkz@nvidia.com>

build(automodel): track r0.6.0 branch

e9f2944

Signed-off-by: larkzhang-nv <larkz@nvidia.com>

build: drop obsolete SGLang dependency conflicts

00f52fd

SGLang now uses Transformers 5.12.1, matching the ModelOpt, NVRx, and NeMo Gym environments. Remove the uv conflict pairs that were only needed while SGLang remained on Transformers 5.6.0. Signed-off-by: larkzhang-nv <larkz@nvidia.com>

test(automodel): fix context parallel import ordering

254ec1a

Signed-off-by: larkzhang-nv <larkz@nvidia.com>

jQizhang force-pushed the bump-automodel-r0.6.0 branch from 138c0f1 to 254ec1a Compare August 28, 2026 06:22

terrykong previously approved these changes Aug 28, 2026

View reviewed changes

test(automodel): fix async checkpoint tests

30c26e4

Signed-off-by: larkzhang-nv <larkz@nvidia.com>

jQizhang dismissed terrykong’s stale review via 30c26e4 August 28, 2026 09:39

sharonyu-115 previously approved these changes Aug 28, 2026

View reviewed changes

fix(ci): fix Qwen-VL weight mapping issue

f56b445

Signed-off-by: larkzhang-nv <larkz@nvidia.com>

jQizhang dismissed sharonyu-115’s stale review via f56b445 August 28, 2026 15:24

terrykong approved these changes Aug 29, 2026

View reviewed changes

terrykong merged commit c158373 into NVIDIA-NeMo:main Aug 29, 2026
102 of 103 checks passed

tianyi-zhang-02 mentioned this pull request Aug 29, 2026

perf(loss): defer the fp32 upcast in the distillation top-k path #3496

Open

sharonyu-115 mentioned this pull request Aug 29, 2026

feat(gemma4): add CP recipes and 12B support #3911

Merged

4 tasks

jQizhang mentioned this pull request Aug 31, 2026

fix(automodel): align checkpoint lifecycle with consolidated save modes #3935

Open

Uh oh!

Conversation

jQizhang commented Aug 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Summary

Automodel r0.6.0 upgrade

Transformers 5.12.1 upgrade

Context-parallel integration

Tests and documentation

Issues

Usage

Validation

Additional information

Uh oh!

copy-pr-bot Bot commented Aug 5, 2026

Uh oh!

github-actions Bot commented Aug 5, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions Bot commented Aug 10, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions Bot commented Aug 19, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions Bot commented Aug 24, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions Bot commented Aug 25, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions Bot commented Aug 26, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions Bot commented Aug 26, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

jQizhang commented Aug 28, 2026

Uh oh!

github-actions Bot commented Aug 28, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions Bot commented Aug 28, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

jQizhang commented Aug 28, 2026

Uh oh!

sharonyu-115 left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Aug 28, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

jQizhang commented Aug 28, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

jQizhang commented Aug 5, 2026 •

edited

Loading