fix(checkpoint): load base weights under DDP and MegatronFSDP (#3597)
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. Signed-off-by: Roman Ralovets <roman@ralovets.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
R
Roman Ralovets committed
e79b92c2f11d513fd52dc1b16174dcae6b854a2a
Parent: 36044ad
Committed by GitHub <noreply@github.com>
on 8/22/2026, 10:21:59 PM