feat: refactor dtensor policy v2 into modular functions #1511

hemildesai · 2025-11-12T17:11:43Z

What does this PR do ?

refactors dtensor policy v2 into core modular reusable functions

Issues

List issues that this PR closes (syntax):

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-actions · 2025-11-12T17:12:07Z

⚠️ File Consistency Check

Check based on commit: f4a4d1d (PR #1511 from hemil/dtensor-policy-v2-refactor)

⚠️ DTensor Policy Worker Synchronization Warning

The file nemo_rl/models/policy/dtensor_policy_worker_v2.py was modified in this PR, but nemo_rl/models/policy/dtensor_policy_worker.py was not updated.

Why this matters:
These files contain related DTensor policy worker implementations that should be kept synchronized to ensure consistency across different versions.

Action required:

Please review if the changes in nemo_rl/models/policy/dtensor_policy_worker_v2.py should also be applied to nemo_rl/models/policy/dtensor_policy_worker.py
Update nemo_rl/models/policy/dtensor_policy_worker.py if necessary to maintain consistency
If the files are intentionally different, please add a comment in the PR explaining why

Files to check:

Modified: nemo_rl/models/policy/dtensor_policy_worker_v2.py
Not modified: nemo_rl/models/policy/dtensor_policy_worker.py

_{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-11-12T17:12:40Z

❌ Submodule Fast-Forward Check Failed

Check based on commit: f4a4d1d (PR #1511 from hemil/dtensor-policy-v2-refactor)

❌ Submodules that need attention:

Automodel: ❌ Commits have DIVERGED from a common ancestor
TARGET (main branch): https://github.com/NVIDIA-NeMo/Automodel/commits/a2db048383cd54b3fafc928df4c30bf7bbf7c430/
CURRENT (PR #1511 from hemil/dtensor-policy-v2-refactor): https://github.com/NVIDIA-NeMo/Automodel/commits/8134b0c039802fb3f6161571400ef7085dd1e9cb/

Please ensure all submodule commits are fast-forwards of the main branch before merging.

github-actions · 2025-11-12T17:52:42Z

⚠️ File Consistency Check

Check based on commit: ae04fcc (PR #1511 from hemil/dtensor-policy-v2-refactor)

⚠️ DTensor Policy Worker Synchronization Warning

The file nemo_rl/models/policy/dtensor_policy_worker_v2.py was modified in this PR, but nemo_rl/models/policy/dtensor_policy_worker.py was not updated.

Why this matters:
These files contain related DTensor policy worker implementations that should be kept synchronized to ensure consistency across different versions.

Action required:

Please review if the changes in nemo_rl/models/policy/dtensor_policy_worker_v2.py should also be applied to nemo_rl/models/policy/dtensor_policy_worker.py
Update nemo_rl/models/policy/dtensor_policy_worker.py if necessary to maintain consistency
If the files are intentionally different, please add a comment in the PR explaining why

Files to check:

Modified: nemo_rl/models/policy/dtensor_policy_worker_v2.py
Not modified: nemo_rl/models/policy/dtensor_policy_worker.py

_{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-11-12T17:52:50Z

❌ Submodule Fast-Forward Check Failed

Check based on commit: ae04fcc (PR #1511 from hemil/dtensor-policy-v2-refactor)

❌ Submodules that need attention:

Automodel: ❌ Commits have DIVERGED from a common ancestor
TARGET (main branch): https://github.com/NVIDIA-NeMo/Automodel/commits/a2db048383cd54b3fafc928df4c30bf7bbf7c430/
CURRENT (PR #1511 from hemil/dtensor-policy-v2-refactor): https://github.com/NVIDIA-NeMo/Automodel/commits/8134b0c039802fb3f6161571400ef7085dd1e9cb/

Please ensure all submodule commits are fast-forwards of the main branch before merging.

github-actions · 2025-11-12T18:58:07Z

⚠️ File Consistency Check

Check based on commit: 0ac6167 (PR #1511 from hemil/dtensor-policy-v2-refactor)

⚠️ DTensor Policy Worker Synchronization Warning

The file nemo_rl/models/policy/dtensor_policy_worker_v2.py was modified in this PR, but nemo_rl/models/policy/dtensor_policy_worker.py was not updated.

Why this matters:
These files contain related DTensor policy worker implementations that should be kept synchronized to ensure consistency across different versions.

Action required:

Please review if the changes in nemo_rl/models/policy/dtensor_policy_worker_v2.py should also be applied to nemo_rl/models/policy/dtensor_policy_worker.py
Update nemo_rl/models/policy/dtensor_policy_worker.py if necessary to maintain consistency
If the files are intentionally different, please add a comment in the PR explaining why

Files to check:

Modified: nemo_rl/models/policy/dtensor_policy_worker_v2.py
Not modified: nemo_rl/models/policy/dtensor_policy_worker.py

_{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.}

Signed-off-by: Hemil Desai <[email protected]>

github-actions · 2025-11-12T18:58:38Z

❌ Submodule Fast-Forward Check Failed

Check based on commit: 0ac6167 (PR #1511 from hemil/dtensor-policy-v2-refactor)

❌ Submodules that need attention:

Automodel: ❌ Commits have DIVERGED from a common ancestor
TARGET (main branch): https://github.com/NVIDIA-NeMo/Automodel/commits/a2db048383cd54b3fafc928df4c30bf7bbf7c430/
CURRENT (PR #1511 from hemil/dtensor-policy-v2-refactor): https://github.com/NVIDIA-NeMo/Automodel/commits/8134b0c039802fb3f6161571400ef7085dd1e9cb/

Please ensure all submodule commits are fast-forwards of the main branch before merging.

github-actions · 2025-11-12T23:40:54Z

⚠️ File Consistency Check

Check based on commit: be7fc16 (PR #1511 from hemil/dtensor-policy-v2-refactor)

⚠️ DTensor Policy Worker Synchronization Warning

The file nemo_rl/models/policy/dtensor_policy_worker_v2.py was modified in this PR, but nemo_rl/models/policy/dtensor_policy_worker.py was not updated.

Why this matters:
These files contain related DTensor policy worker implementations that should be kept synchronized to ensure consistency across different versions.

Action required:

Please review if the changes in nemo_rl/models/policy/dtensor_policy_worker_v2.py should also be applied to nemo_rl/models/policy/dtensor_policy_worker.py
Update nemo_rl/models/policy/dtensor_policy_worker.py if necessary to maintain consistency
If the files are intentionally different, please add a comment in the PR explaining why

Files to check:

Modified: nemo_rl/models/policy/dtensor_policy_worker_v2.py
Not modified: nemo_rl/models/policy/dtensor_policy_worker.py

_{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-11-12T23:41:13Z

❌ Submodule Fast-Forward Check Failed

Check based on commit: be7fc16 (PR #1511 from hemil/dtensor-policy-v2-refactor)

❌ Submodules that need attention:

Automodel: ❌ Commits have DIVERGED from a common ancestor
TARGET (main branch): https://github.com/NVIDIA-NeMo/Automodel/commits/a2db048383cd54b3fafc928df4c30bf7bbf7c430/
CURRENT (PR #1511 from hemil/dtensor-policy-v2-refactor): https://github.com/NVIDIA-NeMo/Automodel/commits/8134b0c039802fb3f6161571400ef7085dd1e9cb/

Please ensure all submodule commits are fast-forwards of the main branch before merging.

Signed-off-by: Hemil Desai <[email protected]>

github-actions · 2025-11-13T05:00:47Z

⚠️ File Consistency Check

Check based on commit: a4d2501 (PR #1511 from hemil/dtensor-policy-v2-refactor)

⚠️ DTensor Policy Worker Synchronization Warning

The file nemo_rl/models/policy/dtensor_policy_worker_v2.py was modified in this PR, but nemo_rl/models/policy/dtensor_policy_worker.py was not updated.

Why this matters:
These files contain related DTensor policy worker implementations that should be kept synchronized to ensure consistency across different versions.

Action required:

Please review if the changes in nemo_rl/models/policy/dtensor_policy_worker_v2.py should also be applied to nemo_rl/models/policy/dtensor_policy_worker.py
Update nemo_rl/models/policy/dtensor_policy_worker.py if necessary to maintain consistency
If the files are intentionally different, please add a comment in the PR explaining why

Files to check:

Modified: nemo_rl/models/policy/dtensor_policy_worker_v2.py
Not modified: nemo_rl/models/policy/dtensor_policy_worker.py

_{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-11-13T05:01:06Z

❌ Submodule Fast-Forward Check Failed

Check based on commit: a4d2501 (PR #1511 from hemil/dtensor-policy-v2-refactor)

❌ Submodules that need attention:

Automodel: ❌ Commits have DIVERGED from a common ancestor
TARGET (main branch): https://github.com/NVIDIA-NeMo/Automodel/commits/a2db048383cd54b3fafc928df4c30bf7bbf7c430/
CURRENT (PR #1511 from hemil/dtensor-policy-v2-refactor): https://github.com/NVIDIA-NeMo/Automodel/commits/8134b0c039802fb3f6161571400ef7085dd1e9cb/

Please ensure all submodule commits are fast-forwards of the main branch before merging.

Signed-off-by: Hemil Desai <[email protected]>

github-actions · 2025-11-13T06:09:02Z

⚠️ File Consistency Check

Check based on commit: 3222237 (PR #1511 from hemil/dtensor-policy-v2-refactor)

⚠️ DTensor Policy Worker Synchronization Warning

The file nemo_rl/models/policy/dtensor_policy_worker_v2.py was modified in this PR, but nemo_rl/models/policy/dtensor_policy_worker.py was not updated.

Why this matters:
These files contain related DTensor policy worker implementations that should be kept synchronized to ensure consistency across different versions.

Action required:

Please review if the changes in nemo_rl/models/policy/dtensor_policy_worker_v2.py should also be applied to nemo_rl/models/policy/dtensor_policy_worker.py
Update nemo_rl/models/policy/dtensor_policy_worker.py if necessary to maintain consistency
If the files are intentionally different, please add a comment in the PR explaining why

Files to check:

Modified: nemo_rl/models/policy/dtensor_policy_worker_v2.py
Not modified: nemo_rl/models/policy/dtensor_policy_worker.py

_{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-11-13T06:09:17Z

❌ Submodule Fast-Forward Check Failed

Check based on commit: 3222237 (PR #1511 from hemil/dtensor-policy-v2-refactor)

❌ Submodules that need attention:

Automodel: ❌ Commits have DIVERGED from a common ancestor
TARGET (main branch): https://github.com/NVIDIA-NeMo/Automodel/commits/a2db048383cd54b3fafc928df4c30bf7bbf7c430/
CURRENT (PR #1511 from hemil/dtensor-policy-v2-refactor): https://github.com/NVIDIA-NeMo/Automodel/commits/8134b0c039802fb3f6161571400ef7085dd1e9cb/

Please ensure all submodule commits are fast-forwards of the main branch before merging.

Signed-off-by: Hemil Desai <[email protected]>

github-actions · 2025-11-13T14:30:26Z

⚠️ File Consistency Check

Check based on commit: 14e04d4 (PR #1511 from hemil/dtensor-policy-v2-refactor)

⚠️ DTensor Policy Worker Synchronization Warning

The file nemo_rl/models/policy/dtensor_policy_worker_v2.py was modified in this PR, but nemo_rl/models/policy/dtensor_policy_worker.py was not updated.

Why this matters:
These files contain related DTensor policy worker implementations that should be kept synchronized to ensure consistency across different versions.

Action required:

Please review if the changes in nemo_rl/models/policy/dtensor_policy_worker_v2.py should also be applied to nemo_rl/models/policy/dtensor_policy_worker.py
Update nemo_rl/models/policy/dtensor_policy_worker.py if necessary to maintain consistency
If the files are intentionally different, please add a comment in the PR explaining why

Files to check:

Modified: nemo_rl/models/policy/dtensor_policy_worker_v2.py
Not modified: nemo_rl/models/policy/dtensor_policy_worker.py

_{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-11-13T14:30:51Z

❌ Submodule Fast-Forward Check Failed

Check based on commit: 14e04d4 (PR #1511 from hemil/dtensor-policy-v2-refactor)

❌ Submodules that need attention:

Automodel: ❌ Commits have DIVERGED from a common ancestor
TARGET (main branch): https://github.com/NVIDIA-NeMo/Automodel/commits/a2db048383cd54b3fafc928df4c30bf7bbf7c430/
CURRENT (PR #1511 from hemil/dtensor-policy-v2-refactor): https://github.com/NVIDIA-NeMo/Automodel/commits/8134b0c039802fb3f6161571400ef7085dd1e9cb/

Please ensure all submodule commits are fast-forwards of the main branch before merging.

Signed-off-by: Hemil Desai <[email protected]>

github-actions · 2025-11-13T18:41:25Z

⚠️ File Consistency Check

Check based on commit: dbe1707 (PR #1511 from hemil/dtensor-policy-v2-refactor)

⚠️ DTensor Policy Worker Synchronization Warning

The file nemo_rl/models/policy/dtensor_policy_worker_v2.py was modified in this PR, but nemo_rl/models/policy/dtensor_policy_worker.py was not updated.

Why this matters:
These files contain related DTensor policy worker implementations that should be kept synchronized to ensure consistency across different versions.

Action required:

Please review if the changes in nemo_rl/models/policy/dtensor_policy_worker_v2.py should also be applied to nemo_rl/models/policy/dtensor_policy_worker.py
Update nemo_rl/models/policy/dtensor_policy_worker.py if necessary to maintain consistency
If the files are intentionally different, please add a comment in the PR explaining why

Files to check:

Modified: nemo_rl/models/policy/dtensor_policy_worker_v2.py
Not modified: nemo_rl/models/policy/dtensor_policy_worker.py

_{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-11-13T18:41:37Z

❌ Submodule Fast-Forward Check Failed

Check based on commit: dbe1707 (PR #1511 from hemil/dtensor-policy-v2-refactor)

❌ Submodules that need attention:

Automodel: ❌ Commits have DIVERGED from a common ancestor
TARGET (main branch): https://github.com/NVIDIA-NeMo/Automodel/commits/a2db048383cd54b3fafc928df4c30bf7bbf7c430/
CURRENT (PR #1511 from hemil/dtensor-policy-v2-refactor): https://github.com/NVIDIA-NeMo/Automodel/commits/17055ab4fe820ba0e4c13aed8fd2c810b6551bf7/

Please ensure all submodule commits are fast-forwards of the main branch before merging.

hemildesai added 4 commits November 12, 2025 10:58

feat: refactor dtensor policy v2 into core modular functions

aefe383

Signed-off-by: Hemil Desai <[email protected]>

fixes

5047e18

Signed-off-by: Hemil Desai <[email protected]>

add is_hf_model and is_moe_model to model state

530a216

Signed-off-by: Hemil Desai <[email protected]>

fix

be7fc16

Signed-off-by: Hemil Desai <[email protected]>

hemildesai force-pushed the hemil/dtensor-policy-v2-refactor branch from 0ac6167 to be7fc16 Compare November 12, 2025 23:40

rename

a4d2501

Signed-off-by: Hemil Desai <[email protected]>

hemildesai changed the title ~~feat: refactor dtensor policy v2 into core modular functions~~ feat: refactor dtensor policy v2 into modular functions Nov 13, 2025

add back torch.autocast

3222237

Signed-off-by: Hemil Desai <[email protected]>

simplify

14e04d4

Signed-off-by: Hemil Desai <[email protected]>

hemildesai added the CI:L0 Run doctests and unit tests label Nov 13, 2025

hemildesai temporarily deployed to nemo-ci November 13, 2025 14:34 — with GitHub Actions Inactive

hemildesai temporarily deployed to nemo-ci November 13, 2025 14:53 — with GitHub Actions Inactive

fix

dbe1707

Signed-off-by: Hemil Desai <[email protected]>

hemildesai removed the CI:L0 Run doctests and unit tests label Nov 13, 2025

hemildesai added the CI:L0 Run doctests and unit tests label Nov 13, 2025

hemildesai temporarily deployed to nemo-ci November 13, 2025 18:41 — with GitHub Actions Inactive

hemildesai temporarily deployed to nemo-ci November 13, 2025 19:20 — with GitHub Actions Inactive

hemildesai closed this Nov 19, 2025

feat: refactor dtensor policy v2 into modular functions #1511

feat: refactor dtensor policy v2 into modular functions #1511

Uh oh!

Conversation

hemildesai commented Nov 12, 2025

What does this PR do ?

Issues

Usage

Before your PR is "Ready for review"

Additional Information

Uh oh!

github-actions bot commented Nov 12, 2025

⚠️ File Consistency Check

⚠️ DTensor Policy Worker Synchronization Warning

Uh oh!

github-actions bot commented Nov 12, 2025

❌ Submodule Fast-Forward Check Failed

❌ Submodules that need attention:

Uh oh!

github-actions bot commented Nov 12, 2025

⚠️ File Consistency Check

⚠️ DTensor Policy Worker Synchronization Warning

Uh oh!

github-actions bot commented Nov 12, 2025

❌ Submodule Fast-Forward Check Failed

❌ Submodules that need attention:

Uh oh!

github-actions bot commented Nov 12, 2025

⚠️ File Consistency Check

⚠️ DTensor Policy Worker Synchronization Warning

Uh oh!

github-actions bot commented Nov 12, 2025

❌ Submodule Fast-Forward Check Failed

❌ Submodules that need attention:

Uh oh!

github-actions bot commented Nov 12, 2025

⚠️ File Consistency Check

⚠️ DTensor Policy Worker Synchronization Warning

Uh oh!

github-actions bot commented Nov 12, 2025

❌ Submodule Fast-Forward Check Failed

❌ Submodules that need attention:

Uh oh!

github-actions bot commented Nov 13, 2025

⚠️ File Consistency Check

⚠️ DTensor Policy Worker Synchronization Warning

Uh oh!

github-actions bot commented Nov 13, 2025

❌ Submodule Fast-Forward Check Failed

❌ Submodules that need attention:

Uh oh!

github-actions bot commented Nov 13, 2025

⚠️ File Consistency Check

⚠️ DTensor Policy Worker Synchronization Warning

Uh oh!

github-actions bot commented Nov 13, 2025

❌ Submodule Fast-Forward Check Failed

❌ Submodules that need attention:

Uh oh!

github-actions bot commented Nov 13, 2025

⚠️ File Consistency Check

⚠️ DTensor Policy Worker Synchronization Warning

Uh oh!

github-actions bot commented Nov 13, 2025

❌ Submodule Fast-Forward Check Failed

❌ Submodules that need attention:

Uh oh!

github-actions bot commented Nov 13, 2025

⚠️ File Consistency Check

⚠️ DTensor Policy Worker Synchronization Warning

Uh oh!

github-actions bot commented Nov 13, 2025

❌ Submodule Fast-Forward Check Failed

❌ Submodules that need attention:

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone