Initialize default process group when loading torch_dist checkpoints - #6787
Draft
pujitha24 wants to merge 1 commit into
Draft
Initialize default process group when loading torch_dist checkpoints#6787pujitha24 wants to merge 1 commit into
pujitha24 wants to merge 1 commit into
Conversation
Motivation: tools/checkpoint/convert.py fails with "ValueError: Default process group has not been initialized" when loading torch_dist checkpoints saved with --use-dist-ckpt. dist_checkpointing.load calls torch.distributed.get_world_size() (via determine_global_metadata), but MegatronCheckpointLoaderBase.initialize_megatron_env never calls torch.distributed.init_process_group, so no default process group exists. The equivalent checkpoint saver path (saver_base.py) already guards against this with a single-process gloo process group, added for the save direction in a prior fix, but the loader path never received the same fix. Approach: Mirror saver_base.py's existing pattern in MegatronCheckpointLoaderBase.initialize_megatron_env: if no default process group exists, initialize a minimal single-process gloo backend (rank=0, world_size=1) before the mpu fake-parallelism setup. convert.py always runs the loader as a single process (via multiprocessing, never torchrun), so this cannot conflict with a real multi-rank launch. Validation: Added tests/unit_tests/tools/checkpoint/test_loader_base.py, which calls initialize_megatron_env directly and asserts a default process group is available afterward (skipped under this repo's multi-rank CI harness, matching the existing guard in test_gpt_hybrid_conversion_parallelism.py, since that harness already provides its own multi-rank default group). I could not execute this new pytest file end-to-end in my local sandbox (macOS, no GPU/Docker): importing megatron.core/megatron.training transitively requires triton, which has no macOS wheel. I did directly validate the underlying mechanism: installed a CPU-only torch build and reproduced the exact "torch.distributed.get_world_size()" ValueError prior to any init_process_group call, then confirmed torch.distributed.init_process_group(backend='gloo', rank=0, world_size=1) resolves it -- the same call this change adds. I also confirmed via git history that this identical pattern is already merged and working for the saver side. CI is the first environment that can run the new test end-to-end. Report: NVIDIA#1818 Signed-off-by: Pujitha Paladugu <10557236+pujitha24@users.noreply.github.com> Assisted-by: claude-sonnet-5 (via Claude Code)
svcnvidia-nemo-ci
marked this pull request as draft
August 23, 2026 22:45
Contributor
|
This PR has been automatically converted to draft because all PRs must start as drafts. When you are ready for review, click Ready for Review to begin the review process. This will:
See the contribution guide for more details. |
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
1 participant
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.
What does this PR do?
Initializes a single-process default
torch.distributedprocess group in the checkpoint conversion loader sotools/checkpoint/convert.pycan loadtorch_distcheckpoints without aValueError: Default process group has not been initializedcrash.CI on this repo requires a maintainer to add
ok-to-test— happy to address anything it surfaces.Note: I could not execute the new unit test end-to-end in my local sandbox (macOS, no GPU/Docker), because importing
megatron.core/megatron.trainingtransitively requirestriton, which has no macOS wheel. I validated the underlying mechanism directly (installing a CPU-only torch build and reproducing the exacttorch.distributed.get_world_size()ValueErrorbefore the fix, then confirminginit_process_group(backend='gloo', rank=0, world_size=1)resolves it), and confirmed viagit logthat this exact fix pattern is already merged and working for the checkpoint saver path (tools/checkpoint/saver_base.py, added in commit a3a7a0c "fix checkpointing conversion", #4058) — the loader path never received the equivalent fix. CI is the first environment that can actually execute the new pytest file.Issue tracking
Linked issue:
Report: #1818
Contribution process
Pre-checks
Code review
Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!
All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.
Step 1: Mark PR as "Ready for Review"
.github/CODEOWNERS.Final Review might get declined if these requirements are not fulfilled.
Step 2: Final Review
For PRs that change
megatron/core, once all expert reviewers have approved, theFinal Reviewlabel is applied automatically and final reviewers are assigned.For PRs outside
megatron/core, this step is skipped.Step 3: Approved
Once all required reviewers have approved, the
Approvedlabel is applied automatically.Merge
Any member of mcore-engineers will be able to merge your PR.
Fixes #1818