From ee5f50abfea192396de555bca43e7e2191eec5c4 Mon Sep 17 00:00:00 2001 From: Roman Ralovets Date: Thu, 20 Aug 2026 04:57:30 -0700 Subject: [PATCH] fix(checkpoint): load base weights under DDP and MegatronFSDP AutoModel's own model implementations (anything in MODEL_ARCH_MAPPING) trained from random weights under `strategy: ddp` and `strategy: megatron_fsdp`, with no error or warning. `apply_model_infrastructure` gated the base-checkpoint read on `is_meta_device`, whose `else` branch assumed "not on meta" implies "weights already loaded". That holds for the HuggingFace fallback path, where `from_pretrained` populates the weights during init. It does not hold for custom implementations, which `_init_model` builds with `model_cls(hf_config)` -- architecture only. `auto_model.py` excludes DDPManager and MegatronFSDPManager from meta-device init, so every custom model under those two wrappers took the `else` branch and never read its checkpoint. Setting `checkpoint_already_loaded = True` then suppressed the post-shard load that would otherwise have caught it. Branch on `weights_already_loaded` instead, which is the flag that answers whether the model still needs its checkpoint. It was already computed in `auto_model.py` and already consulted at the post-shard load site. fsdp2 is unaffected above world size 1: its mesh has a real `dp_shard` axis, so `dp_shard_size > 1` routes it to the post-shard load. megatron_fsdp is affected at every world size, its mesh axes being (dp, cp, tp) with no `dp_shard`. Verified on 2 GPUs with a tiny Llama checkpoint of all-0.5 weights and with Qwen3-1.7B, comparing the language-modeling loss on a known sentence: strategy GPUs tensors holding 0.5 Qwen3-1.7B loss ddp 1 0/21 -> 21/21 11.5824 -> 2.9304 ddp 2 0/21 -> 21/21 12.6028 -> 2.9304 megatron_fsdp 1 0/21 -> 21/21 12.4241 -> 2.9304 megatron_fsdp 2 0/12 local -> 12/12 12.2534 -> 2.9304 fsdp2 1 21/21 -> 21/21 2.9304 -> 2.9304 fsdp2 2 21/21 -> 21/21 2.9375 -> 2.9375 ln(151936) = 11.93, so every broken row sat at or above a uniform distribution over the vocabulary, and the values were not reproducible run to run. After the fix all five unsharded configurations agree exactly. The fsdp2 2-GPU value is unchanged and differs only because that is the one configuration applying FSDP2's bf16 MixedPrecisionPolicy. Adds four CPU tests. Three cover the load decision: two pinning it for either value of `weights_already_loaded`, and one end-to-end check that a model built without weights ends up holding the checkpoint's tensors. Two of those three fail without this change. The fourth pins the flag computation that feeds the decision, so a regression in the manager exclusion or in `weights_already_loaded` cannot silently reintroduce random-weight training. Co-Authored-By: Claude Opus 5 Signed-off-by: Roman Ralovets --- .../_transformers/infrastructure.py | 18 +++-- .../_transformers/test_auto_model.py | 37 +++++++++ .../_transformers/test_infrastructure.py | 77 ++++++++++++++++++- 3 files changed, 122 insertions(+), 10 deletions(-) diff --git a/nemo_automodel/_transformers/infrastructure.py b/nemo_automodel/_transformers/infrastructure.py index b2559891e8..87e1e21eb8 100644 --- a/nemo_automodel/_transformers/infrastructure.py +++ b/nemo_automodel/_transformers/infrastructure.py @@ -593,9 +593,16 @@ def apply_model_infrastructure( checkpoint_already_loaded = False if load_before_shard: - if is_meta_device: - lora_a_init = getattr(peft_config, "lora_A_init", None) - checkpointer.initialize_model_weights(model, device, peft_init_method=lora_a_init) + if weights_already_loaded: + # HF's from_pretrained already populated the weights during model init. + # Still call load_base_model with load_base_model=False to + # handle weight tying + checkpointer.load_base_model(model, device, cache_dir, pretrained_model_name_or_path, load_base_model=False) + else: + # Only meta-device models need their parameter shells materialized first. + if is_meta_device: + lora_a_init = getattr(peft_config, "lora_A_init", None) + checkpointer.initialize_model_weights(model, device, peft_init_method=lora_a_init) checkpointer.load_base_model( model, device, @@ -603,11 +610,6 @@ def apply_model_infrastructure( pretrained_model_name_or_path, load_base_model=load_base_model, ) - else: - # Non-meta models already have weights from from_pretrained. - # Still call load_base_model with load_base_model=False to - # handle weight tying - checkpointer.load_base_model(model, device, cache_dir, pretrained_model_name_or_path, load_base_model=False) checkpoint_already_loaded = True # hold a list copy of the model state dict keys before any parallelization. To be used during checkpoint saving in safetensors format. diff --git a/tests/unit_tests/_transformers/test_auto_model.py b/tests/unit_tests/_transformers/test_auto_model.py index f75934dc45..4c2abbf6ce 100644 --- a/tests/unit_tests/_transformers/test_auto_model.py +++ b/tests/unit_tests/_transformers/test_auto_model.py @@ -1727,6 +1727,43 @@ def test_meta_tensor_not_implemented_error_retries_without_meta_device_on_hf_pat assert result is sentinel_model assert mock_init.call_count == 2 + def test_custom_model_under_ddp_still_needs_its_checkpoint(self): + """A MODEL_ARCH_MAPPING model under DDP reaches infrastructure unloaded and off meta. + + ``DDPManager`` is excluded from meta-device init, and custom model constructors + only build the architecture, so ``apply_model_infrastructure`` has to be told the + weights are still missing. If either flag is wrong the model enters training + randomly initialized and nothing is raised. + """ + build_kwargs, mock_config = self._make_build_kwargs() + build_kwargs["is_hf_model"] = False + dummy_manager_cls = type("DummyManager", (), {}) + build_kwargs["model_wrapper"] = dummy_manager_cls() + sentinel_model = MagicMock() + captured = {} + + def capture(**kwargs): + captured.update(kwargs) + return sentinel_model + + with ( + patch("nemo_automodel._transformers.auto_model.DDPManager", dummy_manager_cls), + patch("nemo_automodel._transformers.auto_model._init_model", return_value=(True, sentinel_model)), + patch("nemo_automodel._transformers.auto_model.get_world_size_safe", return_value=1), + patch( + "nemo_automodel._transformers.capabilities.attach_capabilities_and_validate", + return_value=sentinel_model, + ), + patch("nemo_automodel._transformers.auto_model.apply_model_infrastructure", side_effect=capture), + patch("nemo_automodel._transformers.auto_model.get_hf_config", return_value=mock_config), + patch("nemo_automodel._transformers.auto_model._maybe_dequantize_fp8_for_peft", return_value=False), + patch("torch.cuda.current_device", return_value=0), + ): + _BaseNeMoAutoModelClass._build_model(mock_config, **build_kwargs) + + assert captured["is_meta_device"] is False + assert captured["weights_already_loaded"] is False + class TestNeMoAutoModelForMultimodalLM: """Tests for the NeMoAutoModelForMultimodalLM class and its exports.""" diff --git a/tests/unit_tests/_transformers/test_infrastructure.py b/tests/unit_tests/_transformers/test_infrastructure.py index 9468c88376..8aec00068e 100644 --- a/tests/unit_tests/_transformers/test_infrastructure.py +++ b/tests/unit_tests/_transformers/test_infrastructure.py @@ -569,7 +569,9 @@ def test_skips_rotary_fix_when_not_needed(self): # ============================================================================= -def _run_apply_model_infrastructure_load_before_shard(*, peft_config=None): +def _run_apply_model_infrastructure_load_before_shard( + *, peft_config=None, is_meta_device=True, weights_already_loaded=False +): """Helper that invokes apply_model_infrastructure with load_before_shard=True.""" from nemo_automodel._transformers.infrastructure import apply_model_infrastructure @@ -588,12 +590,13 @@ def _run_apply_model_infrastructure_load_before_shard(*, peft_config=None): result = apply_model_infrastructure( model=model, - is_meta_device=True, + is_meta_device=is_meta_device, device=torch.device("cpu"), load_base_model=True, peft_config=peft_config, pretrained_model_name_or_path="test/model", cache_dir="/tmp/cache", + weights_already_loaded=weights_already_loaded, ) return result, mock_ckpt, model @@ -632,6 +635,76 @@ def test_load_before_shard_does_not_call_load_base_model_with_peft_init_method(s _, kwargs = mock_ckpt.load_base_model.call_args assert "peft_init_method" not in kwargs + def test_load_before_shard_loads_checkpoint_when_init_left_weights_unloaded(self): + """A non-meta model whose init did not load weights must still read the checkpoint. + + ``is_meta_device`` is False for every model built under DDP or MegatronFSDP, + including AutoModel's own implementations, whose constructor only creates the + architecture. Gating the read on it left those models randomly initialized. + """ + _, mock_ckpt, model = _run_apply_model_infrastructure_load_before_shard( + is_meta_device=False, weights_already_loaded=False + ) + + mock_ckpt.load_base_model.assert_called_once_with( + model, torch.device("cpu"), "/tmp/cache", "test/model", load_base_model=True + ) + # Nothing is on meta, so there are no parameter shells to materialize. + mock_ckpt.initialize_model_weights.assert_not_called() + + def test_load_before_shard_skips_checkpoint_when_init_already_loaded_weights(self): + """HF's from_pretrained already populated the weights; only re-tie, do not re-read.""" + _, mock_ckpt, model = _run_apply_model_infrastructure_load_before_shard( + is_meta_device=False, weights_already_loaded=True + ) + + mock_ckpt.load_base_model.assert_called_once_with( + model, torch.device("cpu"), "/tmp/cache", "test/model", load_base_model=False + ) + mock_ckpt.initialize_model_weights.assert_not_called() + + +def test_load_before_shard_populates_unloaded_model_from_checkpoint(tmp_path): + """End-to-end guard on CPU: the model must end up holding the checkpoint's values. + + The mocked tests above pin the call; this one pins the outcome, using the real + Checkpointer and a real safetensors checkpoint whose every tensor is 0.5. + """ + from transformers import LlamaConfig + from transformers import LlamaForCausalLM as HFLlamaForCausalLM + + from nemo_automodel._transformers.infrastructure import apply_model_infrastructure + from nemo_automodel.components.models.llama.model import LlamaForCausalLM + + config = LlamaConfig( + hidden_size=64, + intermediate_size=128, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + vocab_size=256, + max_position_embeddings=128, + tie_word_embeddings=False, + ) + reference = HFLlamaForCausalLM(config) + with torch.no_grad(): + for param in reference.parameters(): + param.fill_(0.5) + reference.save_pretrained(tmp_path, safe_serialization=True) + + config.torch_dtype = torch.float32 + model = apply_model_infrastructure( + model=LlamaForCausalLM(config), + is_meta_device=False, + device=torch.device("cpu"), + load_base_model=True, + pretrained_model_name_or_path=str(tmp_path), + weights_already_loaded=False, + ) + + for name, param in model.named_parameters(): + assert torch.equal(param.detach(), torch.full_like(param, 0.5)), f"{name} was not loaded" + # ============================================================================= # Tests for from_config load_base_model kwarg forwarding