diff --git a/src/diffusers/models/model_loading_utils.py b/src/diffusers/models/model_loading_utils.py index be2e01e4ce18..f0fa2d2257ef 100644 --- a/src/diffusers/models/model_loading_utils.py +++ b/src/diffusers/models/model_loading_utils.py @@ -247,10 +247,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 @@ -271,6 +267,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..c407078c6764 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 @@ -405,3 +446,56 @@ def test_push_to_hub_library_name(self): # Reset repo delete_repo(repo_id, token=TOKEN) + + +class TestLoadModelDictIntoMetaNonBlocking: + # Companion to test_sharded_checkpoint_device_map_matches_cpu_load, which needs a + # real accelerator and is therefore skipped on CPU CI. This pins the same decision + # without touching a device: accelerate's setter is stubbed out, so the only thing + # under test is which copies load_model_dict_into_meta asks for. + + @pytest.mark.parametrize( + ("device", "torch_below_213", "expected_non_blocking"), + [ + ("mps", True, False), # the unsafe combination -- must fall back to blocking copies + ("mps", False, True), # fixed upstream in torch 2.13, so non-blocking is safe again + ("mps:0", True, False), # indexed form of the same device must not slip through + ("cpu", True, True), # unrelated devices keep the fast path + ("cuda", True, True), + ], + ) + def test_non_blocking_disabled_only_for_unsafe_mps( + self, device, torch_below_213, expected_non_blocking, monkeypatch + ): + from diffusers.models import model_loading_utils + + if not model_loading_utils.is_accelerate_version(">", "1.8.1"): + pytest.skip("non_blocking is only passed on accelerate > 1.8.1") + + real_is_torch_version = model_loading_utils.is_torch_version + + def fake_is_torch_version(operation, version): + if (operation, version) == ("<", "2.13"): + return torch_below_213 + return real_is_torch_version(operation, version) + + recorded = {} + + def fake_set_module_tensor_to_device(model, param_name, param_device, value=None, **kwargs): + recorded[param_name] = (param_device, kwargs) + + monkeypatch.setattr(model_loading_utils, "is_torch_version", fake_is_torch_version) + monkeypatch.setattr(model_loading_utils, "set_module_tensor_to_device", fake_set_module_tensor_to_device) + + model = torch.nn.Linear(4, 4) + state_dict = {"weight": torch.randn(4, 4), "bias": torch.randn(4)} + + model_loading_utils.load_model_dict_into_meta(model, state_dict, device_map={"": device}) + + assert set(recorded) == {"weight", "bias"} + for param_name, (param_device, kwargs) in recorded.items(): + assert param_device == device, param_name + assert kwargs["non_blocking"] is expected_non_blocking, ( + f"{param_name} on {device} (torch<2.13={torch_below_213}) asked for " + f"non_blocking={kwargs['non_blocking']}, expected {expected_non_blocking}" + )