SIGN IN SIGN UP

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