fix: disable non-blocking tensor copies to MPS during model loading

[agarwalprakhar2511]

Summary

Fixes silent weight corruption when loading models with device_map="mps".

Root cause

load_model_dict_into_meta unconditionally sets non_blocking=True for set_module_tensor_to_device (added for accelerate > 1.8.1). With mmap-backed safetensors, the source CPU tensor can be released or recycled before the asynchronous MPS copy completes, leaving the MPS parameter filled with garbage data.

The corruption is non-deterministic and dtype-dependent:

• float32 + MPS: weights corrupted, biases OK
• float16 + MPS: biases corrupted, weights OK

This produces extreme values (~1e37), LayerNorm overflow, and NaN/zero outputs (all-black images). The user's workaround — loading to CPU first, then calling .to("mps") — works because .to() is synchronous and the state dict is still alive.

Fix

In load_model_dict_into_meta, detect when the target device is MPS and force non_blocking=False for those transfers. All other devices (CUDA, CPU) continue to use non_blocking=True.

is_mps_target = str(param_device) == "mps" or (
    isinstance(param_device, torch.device) and param_device.type == "mps"
)
set_module_kwargs["non_blocking"] = not is_mps_target

Affected pipelines

Any pipeline loaded with device_map="mps", including but not limited to GlmImagePipeline (reported), Flux, StableDiffusion, etc. The fix is in shared loading infrastructure, not pipeline-specific code.

Test plan

• Verified MPS device detection produces correct non_blocking values for "mps", torch.device("mps"), "cpu", "cuda:0", and integer device indices.
• The reporter's reproduction (GlmImagePipeline.from_pretrained(..., device_map="mps")) should no longer produce NaN outputs. I don't have access to the 10GB+ GLM-Image model weights to run end-to-end, but the code path is clear and the fix is minimal.

Fixes #13227