Loading HF model into megatron has Checkpoint conversion bug

[snowmanwwg]

[ananthsub]

@snowmanwwg could you share the container version / script used for reproducing?

our functional tests exercising this code path are not showing any errors:

• Megatron-Bridge/tests/functional_tests/training/test_pretrain.py

  Lines 71 to 81 in e1155b4

  	model_cfg = Llama32ModelProvider1B(
  	tensor_model_parallel_size=1,
  	pipeline_model_parallel_size=1,
  	context_parallel_size=1,
  	sequence_parallel=False,
  	attention_softmax_in_fp32=True,
  	pipeline_dtype=torch.bfloat16,
  	bf16=True,
  	seq_length=seq_length,
  	make_vocab_size_divisible_by=128,
  	)

-

Megatron-Bridge/src/megatron/bridge/training/setup.py

Lines 200 to 205 in e1155b4

	model = cfg.model.provide_distributed_model(
	ddp_config=cfg.ddp,
	use_torch_fsdp2=cfg.dist.use_torch_fsdp2,
	overlap_param_gather_with_optimizer_step=cfg.optimizer.overlap_param_gather_with_optimizer_step,
	data_parallel_random_init=cfg.rng.data_parallel_random_init,
	)

• Megatron-Bridge/tests/functional_tests/models/test_qwen2_conversion.py

  Lines 201 to 223 in e1155b4

  	# Run multi_gpu_hf.py with specified parallelism configuration on our toy model
  	cmd = [
  	"python",
  	"-m",
  	"torch.distributed.run",
  	"--nproc_per_node=2",
  	"--nnodes=1",
  	"-m",
  	"coverage",
  	"run",
  	"--data-file=/workspace/.coverage",
  	"--source=/workspace/",
  	"--parallel-mode",
  	"examples/models/multi_gpu_hf.py",
  	"--hf-model-id",
  	qwen2_toy_model_path, # Use our local toy model instead of downloading
  	"--output-dir",
  	str(test_output_dir),
  	"--tp",
  	str(tp),
  	"--pp",
  	str(pp),
  	]

• Megatron-Bridge/examples/models/multi_gpu_hf.py

  Lines 81 to 97 in e1155b4

  	if megatron_load_path:
  	model_provider = bridge.to_megatron_provider(load_weights=False)
  	model_provider.tensor_model_parallel_size = tp
  	model_provider.pipeline_model_parallel_size = pp
  	model_provider.expert_model_parallel_size = ep
  	model_provider.expert_tensor_parallel_size = etp
  	model_provider.initialize_model_parallel(seed=0)
  	megatron_model = bridge.load_megatron_model(megatron_load_path, wrap_with_ddp=False)
  	megatron_model = [m.cuda() for m in megatron_model]
  	else:
  	model_provider = bridge.to_megatron_provider(load_weights=True)
  	model_provider.tensor_model_parallel_size = tp
  	model_provider.pipeline_model_parallel_size = pp
  	model_provider.expert_model_parallel_size = ep
  	model_provider.initialize_model_parallel(seed=0)
  	megatron_model = model_provider.provide_distributed_model(wrap_with_ddp=False)

[jadeCurl]

Hi @ananthsub ,

Thanks for looking into this. Here’s the information you asked for:

Container / Environment

=== GPU/Driver ===
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2024 NVIDIA Corporation
Built on Thu_Sep_12_02:18:05_PDT_2024
Cuda compilation tools, release 12.6, V12.6.77
Build cuda_12.6.r12.6/compiler.34841621_0
=== Python/PyTorch ===
python: 3.10.12
platform: Linux-5.15.120.bsk.6-amd64-x86_64-with-glibc2.35
torch: 2.5.0a0+e000cf0ad9.nv24.10
cuda_available: True
cuda: 12.6
cudnn: 90500
transformers: 4.52.4
megatron.core: 0.15.0rc0
flash_attn: 2.7.4.post1
triton: 3.0.0
accelerate: 1.10.0
onnx: 1.16.2
nvidia.dali: 1.42.0
=== Relevant Pip Freeze ===
megatron-bridge==0.1.0a0
torchprofile==0.0.4
transformer_engine_torch==2.5.0

Reproduction Script
Here is the minimal script that reproduces the issue. It just loads the local llama3.1 8B HF model via AutoBridge, sets TP/PP configs, and builds the Megatron model.

from megatron.bridge import AutoBridge

# Load Llama from HuggingFace Hub and convert to Megatron
bridge = AutoBridge.from_hf_pretrained("/.../model_zoo/Llama-3.1-8B")
provider = bridge.to_megatron_provider()

# The provider is lazy - configure parallelism before creating models
provider.tensor_model_parallel_size = 2
provider.pipeline_model_parallel_size = 1

model = provider.provide_distributed_model(wrap_with_ddp=False)

Happy to provide additional logs or try different TP/PP topologies if helpful. Please let me know what else would be useful.

[ananthsub]

@jadeCurl thanks for providing the reproducer script.

regarding the API not being available, the method was renamed last month, so it should be part of the 0.1.0rc1 version.

an additional difference found is that our fallback logic for multi-gpu settings was not correct, and is being addressed here: #454

otherwise, on the current code, you'd need to explicitly call model_provider.initialize_model_parallel(...) as is done here:

Megatron-Bridge/examples/models/multi_gpu_hf.py

Lines 87 to 96 in e1155b4

	model_provider.initialize_model_parallel(seed=0)
	megatron_model = bridge.load_megatron_model(megatron_load_path, wrap_with_ddp=False)
	megatron_model = [m.cuda() for m in megatron_model]
	else:
	model_provider = bridge.to_megatron_provider(load_weights=True)
	model_provider.tensor_model_parallel_size = tp
	model_provider.pipeline_model_parallel_size = pp
	model_provider.expert_model_parallel_size = ep
	model_provider.initialize_model_parallel(seed=0)

from megatron.bridge import AutoBridge

# Load Llama from HuggingFace Hub and convert to Megatron
bridge = AutoBridge.from_hf_pretrained("/.../model_zoo/Llama-3.1-8B")
provider = bridge.to_megatron_provider()

# The provider is lazy - configure parallelism before creating models
provider.tensor_model_parallel_size = 2
provider.pipeline_model_parallel_size = 1

provider.initialize_model_parallel(seed=0)  # <-- Needed on current main without #454 to avoid hangs when the script is run on multiple devices


model = provider.provide_distributed_model(wrap_with_ddp=False)

after #454 this wouldn't be necessary

[jadeCurl]

Hi @ananthsub — thanks for the quick follow-up and pointers.

Two notes on my side:

1. Repro on single-GPU
   I can reproduce the same failure on a single GPU:

2. About initialize_model_parallel(...) in the snippet
   In your example, model_provider isn’t defined; I assumed model_provider = provider. When I add:

provider.initialize_model_parallel(seed=0)

and run, I get:

[ananthsub]

thanks @jadeCurl - this indeed looks like an older version leading to the error. are you able to upgrade to 0.1.0rc1 ? https://pypi.org/project/megatron-bridge/0.1.0rc1/

[ananthsub]

@jadeCurl the script should be easily extensible for multi-gpu runs. for instance, single-node multi-gpu would look like this:

from megatron.bridge import AutoBridge
import torch

def main():
    bridge = AutoBridge.from_hf_pretrained("meta-llama/Llama-3.1-8B")

    provider = bridge.to_megatron_provider()

    provider.tensor_model_parallel_size = 2
    provider.pipeline_model_parallel_size = 1

    model = provider.provide_distributed_model(wrap_with_ddp=False)
    if torch.distributed.is_initialized():
        torch.distributed.destroy_process_group()


if __name__ == "__main__":
    main()

Launched with

 python -m torch.distributed.run --nproc-per-node 2 <script_name>.py

for a full example, please check out https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/models/multi_gpu_hf.py