feat: Create DTensorPolicyWorkerV2 to integrate nemo-automodel apis by ffrujeri · Pull Request #794 · NVIDIA-NeMo/RL

ffrujeri · 2025-07-29T22:46:46Z

What does this PR do ?

This PR adds automodel repo as a submodule and creates DTensorPolicyWorkerV2 to consume parallelize functionality from nemo_automodel apis recently upstreamed.

Issues

List issues that this PR closes (syntax):
#578

Usage

You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

Make sure you read and followed Contributor guidelines
Did you write any new necessary tests?
Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

...

.github/workflows/_automodel_integration_check.yml

terrykong

thanks @ffrujeri for beginning this migration!

nemo_rl/models/policy/dtensor_policy_worker_v2.py

.github/workflows/_automodel_integration_check.yml

nemo_rl/models/policy/lm_policy.py

pyproject.toml

uv.lock

tests/unit/models/policy/test_dtensor_worker_v2.py

nemo_rl/models/policy/dtensor_policy_worker_v2.py

github-actions · 2025-08-15T23:33:35Z

ℹ️ File Synchronization Check

Check based on commit: 388ad40 (PR #794 from ffrujeri/automodel_integration_dtensor_policy_worker)

✅ DTensor Policy Worker Synchronization Check

Both DTensor policy worker files were modified in this PR:

nemo_rl/models/policy/dtensor_policy_worker.py
nemo_rl/models/policy/dtensor_policy_worker_v2.py

Please ensure that the changes are consistent between both files where applicable.

_{This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.}

github-actions · 2025-08-15T23:34:07Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 388ad40 (PR #794 from ffrujeri/automodel_integration_dtensor_policy_worker)

✅ Submodules that are properly updated:

Automodel: ✅ New submodule being added

All submodule changes look good! ✨

Signed-off-by: Felipe Vieira Frujeri <ffrujeri@nvidia.com>

github-actions · 2025-08-23T05:27:40Z

⚠️ File Synchronization Check

Check based on commit: 5eacd3f (PR #794 from ffrujeri/automodel_integration_dtensor_policy_worker)

⚠️ Parallel Plans Synchronization Warning

The file nemo_rl/models/dtensor/parallelize.py was modified in this PR, but 3rdparty/Automodel-workspace/Automodel/nemo_automodel/components/distributed/optimized_tp_plans.py was not updated.

Why this matters:
These files contain similar parallel plan implementations that should be kept synchronized to ensure consistency across the codebase.

Action required:

Please review if the changes in nemo_rl/models/dtensor/parallelize.py should also be applied to 3rdparty/Automodel-workspace/Automodel/nemo_automodel/components/distributed/optimized_tp_plans.py
Update 3rdparty/Automodel-workspace/Automodel/nemo_automodel/components/distributed/optimized_tp_plans.py if necessary to maintain synchronization
If the files are intentionally different, please add a comment in the PR explaining why

Files to check:

Modified: nemo_rl/models/dtensor/parallelize.py
Not modified: 3rdparty/Automodel-workspace/Automodel/nemo_automodel/components/distributed/optimized_tp_plans.py

✅ DTensor Policy Worker Synchronization Check

Both DTensor policy worker files were modified in this PR:

nemo_rl/models/policy/dtensor_policy_worker.py
nemo_rl/models/policy/dtensor_policy_worker_v2.py

Please ensure that the changes are consistent between both files where applicable.

_{This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.}

github-actions · 2025-08-23T05:28:15Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 5eacd3f (PR #794 from ffrujeri/automodel_integration_dtensor_policy_worker)

✅ Submodules that are properly updated:

Automodel: ✅ New submodule being added

All submodule changes look good! ✨

github-actions · 2025-08-23T05:38:46Z

⚠️ File Synchronization Check

Check based on commit: 5eacd3f (PR #794 from ffrujeri/automodel_integration_dtensor_policy_worker)

⚠️ Parallel Plans Synchronization Warning

The file nemo_rl/models/dtensor/parallelize.py was modified in this PR, but 3rdparty/Automodel-workspace/Automodel/nemo_automodel/components/distributed/optimized_tp_plans.py was not updated.

Why this matters:
These files contain similar parallel plan implementations that should be kept synchronized to ensure consistency across the codebase.

Action required:

Please review if the changes in nemo_rl/models/dtensor/parallelize.py should also be applied to 3rdparty/Automodel-workspace/Automodel/nemo_automodel/components/distributed/optimized_tp_plans.py
Update 3rdparty/Automodel-workspace/Automodel/nemo_automodel/components/distributed/optimized_tp_plans.py if necessary to maintain synchronization
If the files are intentionally different, please add a comment in the PR explaining why

Files to check:

Modified: nemo_rl/models/dtensor/parallelize.py
Not modified: 3rdparty/Automodel-workspace/Automodel/nemo_automodel/components/distributed/optimized_tp_plans.py

✅ DTensor Policy Worker Synchronization Check

Both DTensor policy worker files were modified in this PR:

nemo_rl/models/policy/dtensor_policy_worker.py
nemo_rl/models/policy/dtensor_policy_worker_v2.py

Please ensure that the changes are consistent between both files where applicable.

_{This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.}

github-actions · 2025-08-23T05:39:11Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 5eacd3f (PR #794 from ffrujeri/automodel_integration_dtensor_policy_worker)

✅ Submodules that are properly updated:

Automodel: ✅ New submodule being added

All submodule changes look good! ✨

github-actions · 2025-08-24T21:42:19Z

⚠️ File Synchronization Check

Check based on commit: 5eacd3f (PR #794 from ffrujeri/automodel_integration_dtensor_policy_worker)

⚠️ Parallel Plans Synchronization Warning

The file nemo_rl/models/dtensor/parallelize.py was modified in this PR, but 3rdparty/Automodel-workspace/Automodel/nemo_automodel/components/distributed/optimized_tp_plans.py was not updated.

Why this matters:
These files contain similar parallel plan implementations that should be kept synchronized to ensure consistency across the codebase.

Action required:

Please review if the changes in nemo_rl/models/dtensor/parallelize.py should also be applied to 3rdparty/Automodel-workspace/Automodel/nemo_automodel/components/distributed/optimized_tp_plans.py
Update 3rdparty/Automodel-workspace/Automodel/nemo_automodel/components/distributed/optimized_tp_plans.py if necessary to maintain synchronization
If the files are intentionally different, please add a comment in the PR explaining why

Files to check:

Modified: nemo_rl/models/dtensor/parallelize.py
Not modified: 3rdparty/Automodel-workspace/Automodel/nemo_automodel/components/distributed/optimized_tp_plans.py

✅ DTensor Policy Worker Synchronization Check

Both DTensor policy worker files were modified in this PR:

nemo_rl/models/policy/dtensor_policy_worker.py
nemo_rl/models/policy/dtensor_policy_worker_v2.py

Please ensure that the changes are consistent between both files where applicable.

_{This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.}

github-actions · 2025-08-24T21:42:45Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 5eacd3f (PR #794 from ffrujeri/automodel_integration_dtensor_policy_worker)

✅ Submodules that are properly updated:

Automodel: ✅ New submodule being added

All submodule changes look good! ✨

…794) Signed-off-by: Felipe Vieira Frujeri <ffrujeri@nvidia.com> Signed-off-by: Julien Veron Vialard <jveronvialar@nvidia.com>

…VIDIA-NeMo#794) Signed-off-by: Felipe Vieira Frujeri <ffrujeri@nvidia.com> Signed-off-by: Qidong Su <qidongs@nvidia.com>

…VIDIA-NeMo#794) Signed-off-by: Felipe Vieira Frujeri <ffrujeri@nvidia.com> Signed-off-by: Stanislav Kirdey <stan@inflection.ai>

…VIDIA-NeMo#794) Signed-off-by: Felipe Vieira Frujeri <ffrujeri@nvidia.com> Signed-off-by: Qidong Su <qidongs@nvidia.com>

…VIDIA-NeMo#794) Signed-off-by: Felipe Vieira Frujeri <ffrujeri@nvidia.com>

ffrujeri marked this pull request as ready for review July 30, 2025 00:57

ffrujeri force-pushed the ffrujeri/automodel_integration_dtensor_policy_worker branch 3 times, most recently from a040061 to e565216 Compare July 31, 2025 19:11

github-actions bot added the CI Relating to CI label Aug 1, 2025

ffrujeri force-pushed the ffrujeri/automodel_integration_dtensor_policy_worker branch 2 times, most recently from cdf4cad to ddffb89 Compare August 1, 2025 20:20

terrykong reviewed Aug 3, 2025

View reviewed changes

.github/workflows/_automodel_integration_check.yml Show resolved Hide resolved

terrykong requested changes Aug 3, 2025

View reviewed changes

ffrujeri requested a review from terrykong August 4, 2025 18:15

ffrujeri force-pushed the ffrujeri/automodel_integration_dtensor_policy_worker branch from 327a2b9 to f47ab8d Compare August 5, 2025 04:54

ffrujeri changed the title ~~feat: Create DTensorPolicyWorkerV2 to integrate nemo_automodel apis~~ feat: Create DTensorPolicyWorkerV2 to integrate nemo-automodel apis Aug 5, 2025

ffrujeri force-pushed the ffrujeri/automodel_integration_dtensor_policy_worker branch from f47ab8d to 293d3c9 Compare August 5, 2025 16:33

ffrujeri force-pushed the ffrujeri/automodel_integration_dtensor_policy_worker branch 2 times, most recently from 193d7e2 to 039adbd Compare August 14, 2025 18:11

NVIDIA-NeMo deleted a comment from github-actions bot Aug 15, 2025

ffrujeri commented Aug 15, 2025

View reviewed changes

nemo_rl/models/policy/dtensor_policy_worker_v2.py Show resolved Hide resolved

ffrujeri force-pushed the ffrujeri/automodel_integration_dtensor_policy_worker branch from 7b75e42 to 388ad40 Compare August 15, 2025 23:33

NVIDIA-NeMo deleted a comment from github-actions bot Aug 15, 2025

ffrujeri added 5 commits August 22, 2025 22:27

Update test_dtensor_worker_v2.py import.

c37baa7

Signed-off-by: Felipe Vieira Frujeri <ffrujeri@nvidia.com>

Add a comment on pyproject.toml.

fbeee73

Signed-off-by: Felipe Vieira Frujeri <ffrujeri@nvidia.com>

Resolve review comments, fix test import error.

29a46b5

Signed-off-by: Felipe Vieira Frujeri <ffrujeri@nvidia.com>

Rename sequence_parallel to sp.

4ed7eea

Signed-off-by: Felipe Vieira Frujeri <ffrujeri@nvidia.com>

Remove double definition of AUTOMODEL as PY_EXECUTABLE.

5eacd3f

Signed-off-by: Felipe Vieira Frujeri <ffrujeri@nvidia.com>

ffrujeri force-pushed the ffrujeri/automodel_integration_dtensor_policy_worker branch from 181b9bb to 5eacd3f Compare August 23, 2025 05:27

ffrujeri added the CI:L1 Run doctests, unit tests, and functional tests label Aug 23, 2025

ffrujeri temporarily deployed to nemo-ci August 23, 2025 05:38 — with GitHub Actions Inactive

ffrujeri temporarily deployed to nemo-ci August 23, 2025 06:35 — with GitHub Actions Inactive

terrykong enabled auto-merge August 24, 2025 21:41

terrykong added CI:docs Run doctest and removed CI:L0 Run doctests and unit tests CI:L1 Run doctests, unit tests, and functional tests labels Aug 24, 2025

terrykong temporarily deployed to nemo-ci August 24, 2025 21:42 — with GitHub Actions Inactive

terrykong added this pull request to the merge queue Aug 24, 2025

Merged via the queue into main with commit e861b94 Aug 25, 2025
51 of 56 checks passed

terrykong deleted the ffrujeri/automodel_integration_dtensor_policy_worker branch August 25, 2025 01:43

jveronvialard pushed a commit that referenced this pull request Aug 27, 2025

feat: Create DTensorPolicyWorkerV2 to integrate nemo-automodel apis (#…

8063c09

…794) Signed-off-by: Felipe Vieira Frujeri <ffrujeri@nvidia.com> Signed-off-by: Julien Veron Vialard <jveronvialar@nvidia.com>

soodoshll pushed a commit to soodoshll/RL that referenced this pull request Sep 4, 2025

feat: Create DTensorPolicyWorkerV2 to integrate nemo-automodel apis (N…

9364849

…VIDIA-NeMo#794) Signed-off-by: Felipe Vieira Frujeri <ffrujeri@nvidia.com> Signed-off-by: Qidong Su <qidongs@nvidia.com>

PrinsYin pushed a commit to PrinsYin/RL that referenced this pull request Nov 30, 2025

feat: Create DTensorPolicyWorkerV2 to integrate nemo-automodel apis (N…

14da672

…VIDIA-NeMo#794) Signed-off-by: Felipe Vieira Frujeri <ffrujeri@nvidia.com>

Conversation

ffrujeri commented Jul 29, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Issues

Usage

Before your PR is "Ready for review"

Additional Information

Uh oh!

Uh oh!

terrykong left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions bot commented Aug 15, 2025

ℹ️ File Synchronization Check

✅ DTensor Policy Worker Synchronization Check

Uh oh!

github-actions bot commented Aug 15, 2025

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions bot commented Aug 23, 2025

⚠️ File Synchronization Check

⚠️ Parallel Plans Synchronization Warning

✅ DTensor Policy Worker Synchronization Check

Uh oh!

github-actions bot commented Aug 23, 2025

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions bot commented Aug 23, 2025

⚠️ File Synchronization Check

⚠️ Parallel Plans Synchronization Warning

✅ DTensor Policy Worker Synchronization Check

Uh oh!

github-actions bot commented Aug 23, 2025

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions bot commented Aug 24, 2025

⚠️ File Synchronization Check

⚠️ Parallel Plans Synchronization Warning

✅ DTensor Policy Worker Synchronization Check

Uh oh!

github-actions bot commented Aug 24, 2025

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

ffrujeri commented Jul 29, 2025 •

edited

Loading