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..9a68bb7bea64 100644 --- a/tests/models/test_modeling_common.py +++ b/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