From ae59fb6db57add84b65bcee5bdb99cfcb7dc7ac6 Mon Sep 17 00:00:00 2001 From: Rudra Mantri <189433012+RudraMantri123@users.noreply.github.com> Date: Fri, 21 Aug 2026 10:13:28 +0530 Subject: [PATCH 1/2] Fix silent weight corruption when loading with device_map on MPS (torch < 2.13) load_model_dict_into_meta unconditionally passes non_blocking=True to 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 (pytorch/pytorch#189690): the loader's temporary dtype-cast tensors are freed while hundreds of copies are still queued, so loaded weights are silently corrupted (values from freed memory, up to ~1e30 garbage). Empirically: torch 2.10/2.11/2.12 corrupt 17-43 of 65 tensors when loading a sharded bf16 checkpoint with device_map="mps" and a dtype conversion; torch 2.13 is clean; blocking copies are clean on every version. 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 fixed torch versions, keep the non-blocking fast path. Adds a device-generic regression test: loading a sharded checkpoint with device_map plus dtype conversion must match a CPU load exactly. The test fails on torch 2.12 + MPS without this fix and passes with it. Fixes #13227 --- src/diffusers/models/model_loading_utils.py | 13 +++++--- tests/models/test_modeling_common.py | 35 +++++++++++++++++++++ 2 files changed, 44 insertions(+), 4 deletions(-) diff --git a/src/diffusers/models/model_loading_utils.py b/src/diffusers/models/model_loading_utils.py index abbde8082bb5..b9baf59da54b 100644 --- a/src/diffusers/models/model_loading_utils.py +++ b/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: diff --git a/tests/models/test_modeling_common.py b/tests/models/test_modeling_common.py index 9968add19dd9..924cb6841071 100644 --- a/tests/models/test_modeling_common.py +++ b/tests/models/test_modeling_common.py @@ -278,6 +278,41 @@ def get_dummy_inputs(): SD3Transformer2DModel._keep_in_fp32_modules = fp32_modules + @require_torch_accelerator + def test_sharded_checkpoint_device_map_matches_cpu_load(self): + # 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). + 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 From 42497fe87adc67559aa9417589c34f29762e7ee9 Mon Sep 17 00:00:00 2001 From: Rudra Mantri <189433012+RudraMantri123@users.noreply.github.com> Date: Fri, 21 Aug 2026 14:37:04 +0530 Subject: [PATCH 2/2] Cover the parallel shard loader in the device_map regression test _load_shard_files_with_threadpool routes through the same load_model_dict_into_meta, so it inherits the MPS guard; parametrize the regression test over both the serial and threadpool loaders. Verified on torch 2.12 + MPS: both variants fail without the fix and pass with it. --- tests/models/test_modeling_common.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/models/test_modeling_common.py b/tests/models/test_modeling_common.py index 924cb6841071..9a68bb7bea64 100644 --- a/tests/models/test_modeling_common.py +++ b/tests/models/test_modeling_common.py @@ -279,12 +279,18 @@ def get_dummy_inputs(): SD3Transformer2DModel._keep_in_fp32_modules = fp32_modules @require_torch_accelerator - def test_sharded_checkpoint_device_map_matches_cpu_load(self): + @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). + # 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),