feat: Add pinned memory optimizer offload for Megatron policy worker#2248
Open
snivertynv wants to merge 1 commit intomainfrom
Open
feat: Add pinned memory optimizer offload for Megatron policy worker#2248snivertynv wants to merge 1 commit intomainfrom
snivertynv wants to merge 1 commit intomainfrom
Conversation
…olled using the use_pinned_optimizer_offload setting - set to false in a couple of grpo_math* yaml configs as an example. Added test cases for this feature in test_megatron_worker.py Signed-off-by: Sriharsha Niverty <sniverty@nvidia.com>
Author
|
/ok to test 13e7f01 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This significantly improves performance for the optimizer_offload_before_refit pass which is quite expensive in co-located/syncRL cases.
Enabled/disabled using the
use_pinned_optimizer_offloadsetting (default=disabled). It has been set tofalsein a couple of grpo_math* yaml configs as an example. Added test cases for this feature intest_megatron_worker.pyWhat does this PR do ?
Optimizer D2H/H2D transfers used per-tensor pageable allocations, causing expensive cudaHostAlloc calls and synchronous memcpy on every step. This adds an opt-in mode (use_pinned_optimizer_offload) that coalesces all optimizer state into a single cached pinned buffer, eliminating cudaHostAlloc from the hot path and enabling non-blocking DMA transfers.
Issues
List issues that this PR closes (syntax):
Usage
# Add a code snippet demonstrating how to use thisBefore your PR is "Ready for review"
Pre checks:
Additional Information