fix(checkpoint): load base weights under DDP and MegatronFSDP

[ralovets]

What does this PR do ?

Loads the pretrained checkpoint for MODEL_ARCH_MAPPING models under strategy: ddp and
strategy: megatron_fsdp. These models previously trained from random weights, with no
error or warning.

The pre-shard load in apply_model_infrastructure was gated on is_meta_device. Its
else branch assumes the weights already came from from_pretrained, which is true only
on the HuggingFace fallback path. AutoModel's own models are built architecture-only, and
auto_model.py:517 excludes DDPManager and MegatronFSDPManager from meta init, so the
checkpoint read was skipped. checkpoint_already_loaded = True then suppressed the
post-shard load. The fix branches on weights_already_loaded, which is already computed
in auto_model.py and already used at the post-shard load site.

Changelog

Fix

• nemo_automodel/_transformers/infrastructure.py
  • In the load_before_shard path, branch on weights_already_loaded instead of
    is_meta_device when deciding whether to read the base checkpoint.
  • initialize_model_weights stays gated on is_meta_device. Only meta models need
    their parameter shells materialized.

Tests

• tests/unit_tests/_transformers/test_infrastructure.py (3 CPU tests, load decision)
  • test_load_before_shard_loads_checkpoint_when_init_left_weights_unloaded: reads the
    checkpoint when init left the weights unloaded. Fails without the fix.
  • test_load_before_shard_skips_checkpoint_when_init_already_loaded_weights: skips the
    read when from_pretrained already populated the weights.
  • test_load_before_shard_populates_unloaded_model_from_checkpoint: end to end against
    a real safetensors checkpoint, asserts every parameter holds its value. Fails without
    the fix.
• tests/unit_tests/_transformers/test_auto_model.py (1 CPU test, flag computation)
  • test_custom_model_under_ddp_still_needs_its_checkpoint: asserts is_meta_device and
    weights_already_loaded are both False for a MODEL_ARCH_MAPPING model under
    DDPManager.

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 add or update any necessary documentation? Not applicable. No public API,
  config key, or documented behavior changes.

Additional Information

Reproducing

# Saves a tiny Llama checkpoint whose every weight is 0.5, loads it back through
# NeMoAutoModelForCausalLM.from_pretrained, and counts how many parameters hold that
# value. Needs a GPU: from_pretrained calls torch.cuda.current_device().
#
#   torchrun --nproc-per-node 1 repro.py ddp             # 0/21 before the fix, 21/21 after
#   torchrun --nproc-per-node 2 repro.py megatron_fsdp   # 0/12 before the fix, 12/12 after
#   torchrun --nproc-per-node 1 repro.py fsdp2           # 21/21 either way

import sys
import tempfile

import torch
from transformers import LlamaConfig, LlamaForCausalLM

from nemo_automodel import NeMoAutoModelForCausalLM
from nemo_automodel.components.distributed import DistributedSetup, initialize_distributed

config = LlamaConfig(hidden_size=64, intermediate_size=128, num_hidden_layers=2, num_attention_heads=4, vocab_size=256)
reference = LlamaForCausalLM(config)
with torch.no_grad():
    for param in reference.parameters():
        param.fill_(0.5)
checkpoint = tempfile.mkdtemp()
reference.save_pretrained(checkpoint)

strategy = sys.argv[1] if len(sys.argv) > 1 else "ddp"
dist_env = initialize_distributed("nccl")
setup = DistributedSetup.build(strategy=strategy, world_size=dist_env.world_size)
model = NeMoAutoModelForCausalLM.from_pretrained(checkpoint, distributed_setup=setup)

# Skip 0-element shards: megatron_fsdp leaves some ranks holding empty slices,
# and `.all()` on an empty tensor is vacuously True.
shards = [p.detach().to_local() if hasattr(p, "to_local") else p.detach() for p in model.parameters()]
shards = [t for t in shards if t.numel() > 0]
loaded = sum(int(t.eq(0.5).all()) for t in shards)
print(f"\n[{strategy}] {loaded}/{len(shards)} params hold the checkpoint value")

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[akoumpa]

/ok to test ee5f50a

[yuhezhang-ai]

Thanks for fixing it! LGTM

[akoumpa]

/ok to test 6078fbc