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

src/diffusers/models/model_loading_utils.py

@@ -253,10 +253,6 @@ def load_model_dict_into_meta(
                 param = param.to(dtype)
                 set_module_kwargs["dtype"] = dtype
 
-        if is_accelerate_version(">", "1.8.1"):
-            set_module_kwargs["non_blocking"] = True
-            set_module_kwargs["clear_cache"] = False
-
         # For compatibility with PyTorch load_state_dict which converts state dict dtype to existing dtype in model, and which
         # uses `param.copy_(input_param)` that preserves the contiguity of the parameter in the model.
         # Reference: https://github.com/pytorch/pytorch/blob/db79ceb110f6646523019a59bbd7b838f43d4a86/torch/nn/modules/module.py#L2040C29-L2040C29
@@ -277,6 +273,15 @@ def load_model_dict_into_meta(
 
         param_device = _determine_param_device(param_name, device_map)
 
+        if is_accelerate_version(">", "1.8.1"):
+            # On MPS with torch < 2.13, a non-blocking CPU->MPS copy can read source storage that
+            # was already released before the stream synchronizes, silently corrupting the loaded
+            # weights (https://github.com/pytorch/pytorch/issues/189690,
+            # https://github.com/huggingface/diffusers/issues/13227). Use blocking copies there.
+            unsafe_mps_non_blocking = str(param_device).startswith("mps") and is_torch_version("<", "2.13")
+            set_module_kwargs["non_blocking"] = not unsafe_mps_non_blocking
+            set_module_kwargs["clear_cache"] = False
+
         # bnb params are flattened.
         # gguf quants have a different shape based on the type of quantization applied
         if empty_state_dict[param_name].shape != param.shape:

tests/models/test_modeling_common.py

@@ -278,6 +278,47 @@ def get_dummy_inputs():
 
         SD3Transformer2DModel._keep_in_fp32_modules = fp32_modules
 
+    @require_torch_accelerator
+    @pytest.mark.parametrize("parallel_loading", [False, True])
+    def test_sharded_checkpoint_device_map_matches_cpu_load(self, parallel_loading, monkeypatch):
+        # Loading a sharded checkpoint directly onto an accelerator with a dtype conversion must
+        # produce exactly the same weights as loading on CPU. Regression test for silent weight
+        # corruption on MPS with torch < 2.13, where the loader's non-blocking copies could read
+        # already-released source memory (https://github.com/huggingface/diffusers/issues/13227,
+        # https://github.com/pytorch/pytorch/issues/189690). Runs against both the serial and the
+        # threadpool shard loaders, which share the same per-parameter device placement.
+        if parallel_loading:
+            import diffusers.models.modeling_utils as modeling_utils
+
+            monkeypatch.setattr(modeling_utils, "HF_ENABLE_PARALLEL_LOADING", True)
+        torch.manual_seed(0)
+        config = {
+            "block_out_channels": (32, 64),
+            "down_block_types": ("CrossAttnDownBlock2D", "DownBlock2D"),
+            "up_block_types": ("UpBlock2D", "CrossAttnUpBlock2D"),
+            "cross_attention_dim": 32,
+            "attention_head_dim": 8,
+            "out_channels": 4,
+            "in_channels": 4,
+            "layers_per_block": 1,
+            "sample_size": 16,
+        }
+        model = UNet2DConditionModel(**config).to(torch.bfloat16)
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            # several shards so the loader frees per-shard state dicts while copies are queued
+            model.save_pretrained(tmpdir, max_shard_size="200KB")
+            del model
+
+            reference = UNet2DConditionModel.from_pretrained(tmpdir, torch_dtype=torch.float32)
+            reference_sd = reference.state_dict()
+
+            loaded = UNet2DConditionModel.from_pretrained(tmpdir, torch_dtype=torch.float32, device_map=torch_device)
+            for name, value in loaded.state_dict().items():
+                assert torch.equal(value.detach().cpu(), reference_sd[name]), (
+                    f"{name} differs between device_map={torch_device} load and CPU load"
+                )
+
 
 class UNetTesterMixin:
     @staticmethod