Fix silent weight corruption when loading with device_map on MPS (torch < 2.13)

[RudraMantri123]

What does this PR do?

Fixes #13227 — model weights are silently corrupted when loading with device_map="mps" on torch < 2.13 (reported for zai-org/GLM-Image, but not model-specific).

Root cause

load_model_dict_into_meta unconditionally passes non_blocking=True to accelerate's set_module_tensor_to_device. On MPS with torch < 2.13, a non-blocking CPU→MPS copy can read source storage that was already released before the stream synchronizes — filed upstream as pytorch/pytorch#189690 (module: correctness (silent)). The loader casts each tensor to torch_dtype first, creating a temporary CPU tensor, then enqueues the async copy; across a sharded checkpoint hundreds of copies are queued while their source temporaries are freed, so parameters silently take values from freed memory (garbage up to ~1e30 — the issue's "~1e37, LayerNorm overflow, black images").

Verified with a synthetic sharded bf16 checkpoint (3-layer GLM-width transformer, 11 shards) loaded with torch_dtype=torch.float32, device_map="mps" and compared tensor-by-tensor against a CPU load, on Apple M5 Pro / macOS 26.4.1:

torch	raw async-copy probe	model-level load
2.10.0	17/65 tensors corrupted	39–40 corrupted (reporter's exact modules: image_projector, time_condition_embed)
2.11.0	43/65	—
2.12.0	41/65	current main: 41–51 corrupted
2.13.0	0/65	0
any, non_blocking=False	0	0

So the underlying bug is fixed in torch 2.13.0, but every MPS user on 2.10–2.12 hits silent corruption today, including with current main.

The fix

Determine the parameter's target device first, and fall back to blocking copies only for MPS on torch < 2.13. All other devices — and MPS on torch ≥ 2.13 — keep the non-blocking fast path. Verified: torch 2.12 + this branch loads clean (0 corrupted, repeated trials); torch 2.13 unchanged.

Tests

test_sharded_checkpoint_device_map_matches_cpu_load (in TestModelUtils): saves a tiny bf16 sharded checkpoint, loads it with device_map=torch_device plus a dtype conversion, and asserts exact equality with a CPU load. Device-generic (meaningful on CUDA/XPU too), no Hub access. It fails on torch 2.12 + MPS without this fix and passes with it; passes on torch 2.13 either way.

Notes for reviewers

AI disclosure: Claude Code assisted with debugging and drafting; all experiments were run and verified by the author on real hardware.

• The version boundary (2.13) is established empirically via the matrix above; pytorch#189690 is still open upstream, so there is no fix-commit to cite for an exact boundary. If it turns out a 2.12.x patch release picked it up, the guard is conservative, never wrong.
• Audit of the other non_blocking sites in src/: the parallel shard loader (_load_shard_files_with_threadpool) routes through the same load_model_dict_into_meta, so it inherits the guard — the regression test now covers both the serial and threadpool paths (both fail on torch 2.12 without the fix). Pipeline.to() / ModelMixin.to() default to non_blocking=False, and group offloading moves persistent module tensors (with pinned staging buffers), not soon-to-be-freed temporaries — so neither shares this failure mode.
• str(param_device).startswith("mps") is used instead of constructing a torch.device because param_device can be "disk" or an int ordinal on other paths.

Who can review?

@yiyixuxu @sayakpaul @asomoza

🤖 Generated with Claude Code

[RudraMantri123]

Two quick notes for reviewers:

cc @a-r-r-o-w — this touches the non-blocking loading fast path you introduced in #11904, so you're likely the right reviewer. To be clear about scope: the 4–5x speedup is preserved everywhere except MPS on torch < 2.13, where non-blocking CPU→MPS copies can read already-released source storage (pytorch/pytorch#189690) and silently corrupt loaded weights — the fallback to blocking copies applies only to that combination, per-parameter, based on the resolved target device.

On CI: the red "Secret Leaks" run some may see is the trufflehog scan on my fork's branch push — the known new-branch flake tracked in #14376, unrelated to this change. The PR's own test suites are all action_required, awaiting first-time-contributor workflow approval; the regression test included here fails on torch 2.12 + MPS without the fix and passes with it, verified locally both ways.