From 54b6ed5421d44dec2bc420f6ad8b89c2120e19c3 Mon Sep 17 00:00:00 2001 From: Pujitha Paladugu <10557236+pujitha24@users.noreply.github.com> Date: Sun, 23 Aug 2026 15:45:06 -0700 Subject: [PATCH] Initialize default process group when loading torch_dist checkpoints 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: https://github.com/NVIDIA/Megatron-LM/issues/1818 Signed-off-by: Pujitha Paladugu <10557236+pujitha24@users.noreply.github.com> Assisted-by: claude-sonnet-5 (via Claude Code) --- .../tools/checkpoint/test_loader_base.py | 90 +++++++++++++++++++ tools/checkpoint/loader_base.py | 9 ++ 2 files changed, 99 insertions(+) create mode 100644 tests/unit_tests/tools/checkpoint/test_loader_base.py diff --git a/tests/unit_tests/tools/checkpoint/test_loader_base.py b/tests/unit_tests/tools/checkpoint/test_loader_base.py new file mode 100644 index 00000000000..7e0a397ed91 --- /dev/null +++ b/tests/unit_tests/tools/checkpoint/test_loader_base.py @@ -0,0 +1,90 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +""" +Unit tests for MegatronCheckpointLoaderBase.initialize_megatron_env. + +Regression test for a bug where the checkpoint conversion loader never +initialized a default torch.distributed process group, causing +dist_checkpointing.load (via torch.distributed.get_world_size()) to fail +with "Default process group has not been initialized" when loading +torch_dist checkpoints through tools/checkpoint/convert.py. +""" + +import os +import sys +from types import SimpleNamespace +from unittest import mock + +import pytest +import torch +import torch.distributed as dist + +sys.path.insert( + 0, os.path.join(os.path.dirname(__file__), '..', '..', '..', '..', 'tools', 'checkpoint') +) + +from loader_base import MegatronCheckpointLoaderBase + + +# These scenarios assume a single-rank (or not-yet-initialized) default +# torch.distributed process group, matching how convert.py actually runs the +# loader (plain `python`, never torchrun). When pytest is launched under this +# repo's multi-rank CI harness (torch.distributed.run --nproc-per-node 8), the +# default PG is already multi-rank before collection, so initialize_megatron_env's +# `if not is_initialized()` guard correctly leaves it untouched -- skip here +# rather than asserting a single-rank world size against it. See the identical +# guard/rationale in test_gpt_hybrid_conversion_parallelism.py. +@pytest.fixture(autouse=True) +def _skip_when_multi_rank_pg(): + if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1: + pytest.skip( + "Single-rank process-group init test skipped under a multi-rank " + "default process group." + ) + + +class TestInitializeMegatronEnv: + def _make_loader(self): + loader = MegatronCheckpointLoaderBase.__new__(MegatronCheckpointLoaderBase) + loader.build_tokenizer = False + loader.margs = SimpleNamespace( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + virtual_pipeline_model_parallel_size=None, + expert_model_parallel_size=1, + ) + return loader + + def test_default_process_group_available_after_init(self): + was_initialized = torch.distributed.is_initialized() + loader = self._make_loader() + + try: + with mock.patch('megatron.training.global_vars.set_global_variables'): + loader.initialize_megatron_env() + + assert torch.distributed.is_initialized() + # This is the exact call (megatron/core/dist_checkpointing/validation.py, + # determine_global_metadata) that raised ValueError before the fix. + assert torch.distributed.get_world_size() == 1 + finally: + if not was_initialized and torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + def test_does_not_reinitialize_existing_process_group(self): + if not torch.distributed.is_initialized(): + os.environ.setdefault('MASTER_ADDR', 'localhost') + os.environ.setdefault('MASTER_PORT', '12356') + torch.distributed.init_process_group(backend='gloo', rank=0, world_size=1) + initialized_here = True + else: + initialized_here = False + + try: + loader = self._make_loader() + with mock.patch('megatron.training.global_vars.set_global_variables'): + loader.initialize_megatron_env() # must not raise re-init errors + assert torch.distributed.get_world_size() == 1 + finally: + if initialized_here: + torch.distributed.destroy_process_group() diff --git a/tools/checkpoint/loader_base.py b/tools/checkpoint/loader_base.py index 3cf717fd486..f015c4f7b78 100644 --- a/tools/checkpoint/loader_base.py +++ b/tools/checkpoint/loader_base.py @@ -140,6 +140,15 @@ def initialize_megatron_env(self): sys.exit(1) set_global_variables(self.margs, build_tokenizer=self.build_tokenizer) + + # Initialize torch.distributed with a minimal single-process backend so that + # dist_checkpointing (which calls torch.distributed.get_world_size()) has a + # default process group to query when loading torch_dist checkpoints. + if not torch.distributed.is_initialized(): + os.environ.setdefault('MASTER_ADDR', 'localhost') + os.environ.setdefault('MASTER_PORT', '12356') + torch.distributed.init_process_group(backend='gloo', rank=0, world_size=1) + mpu.set_tensor_model_parallel_world_size(self.margs.tensor_model_parallel_size) mpu.set_pipeline_model_parallel_world_size(self.margs.pipeline_model_parallel_size) mpu.set_virtual_pipeline_model_parallel_world_size(self.margs.virtual_pipeline_model_parallel_size)