Skip to content

fix: enable expandable segments for hopper+ #594

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 2 commits into from
Jul 3, 2025
Merged

fix: enable expandable segments for hopper+ #594

merged 2 commits into from
Jul 3, 2025

Conversation

parthchadha
Copy link
Contributor

What does this PR do ?

Add a one line overview of what this PR aims to accomplish.

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

  • ...

@parthchadha
fix: enable expandable segments for hopper+
c0ce2bc 
Signed-off-by: Parth Chadha <pchadha@nvidia.com>
@parthchadha parthchadha added the CI:L0 Run doctests and unit tests label Jul 2, 2025
ashors1
ashors1 reviewed Jul 2, 2025
View reviewed changes
Copy link
Contributor

@ashors1 ashors1 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for making this change! I think we actually want to detect compute capability when using megatron: https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/models/policy/megatron_policy_worker.py#L647-L651. @SahilJain314 , correct me if I'm wrong, but it's my understanding that we only hit issues with (A100 + expandable segments) when using Megatron

@parthchadha
add runtime env for megatron worker
77cf050 
Signed-off-by: Parth Chadha <pchadha@nvidia.com>
@parthchadha
Copy link
Contributor Author

parthchadha commented Jul 2, 2025 •
edited
Loading

Thank you for making this change! I think we actually want to detect compute capability when using megatron: https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/models/policy/megatron_policy_worker.py#L647-L651. @SahilJain314 , correct me if I'm wrong, but it's my understanding that we only hit issues with (A100 + expandable segments) when using Megatron

Made the change for both dtensor and megatron worker in the latest commit.

@parthchadha parthchadha added this pull request to the merge queue Jul 3, 2025
Merged via the queue into main with commit f4150bf Jul 3, 2025
13 of 14 checks passed
@parthchadha parthchadha deleted the pchadha/enabled-expandable-segments-hopper branch July 3, 2025 18:36
therealnaveenkamal pushed a commit to therealnaveenkamal/RL that referenced this pull request Jul 7, 2025
@parthchadha @therealnaveenkamal
fix: enable expandable segments for hopper+ (NVIDIA-NeMo#594)
3b33987 
Signed-off-by: Parth Chadha <pchadha@nvidia.com>
YzjiaoNvd pushed a commit to YzjiaoNvd/NeMo-RL that referenced this pull request Jul 14, 2025
@parthchadha
fix: enable expandable segments for hopper+ (NVIDIA-NeMo#594)
ef205e7 
Signed-off-by: Parth Chadha <pchadha@nvidia.com>
jialei777 pushed a commit to jialei777/nemo-rl that referenced this pull request Jul 23, 2025
@parthchadha @jialei777
fix: enable expandable segments for hopper+ (NVIDIA-NeMo#594)
9325915 
Signed-off-by: Parth Chadha <pchadha@nvidia.com>
Signed-off-by: Jialei Chen <jialeic@google.com>
KiddoZhu pushed a commit that referenced this pull request Jul 28, 2025
@parthchadha @KiddoZhu
fix: enable expandable segments for hopper+ (#594)
56fc9c6 
Signed-off-by: Parth Chadha <pchadha@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@ashors1 ashors1 ashors1 approved these changes

@SahilJain314 SahilJain314 Awaiting requested review from SahilJain314

Assignees
No one assigned
Labels
CI:L0 Run doctests and unit tests
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@parthchadha @ashors1