From 1e25be4e5e5eedfd7f55839c90079c9225df33a0 Mon Sep 17 00:00:00 2001 From: Greninja44 Date: Sat, 29 Aug 2026 08:53:46 +0000 Subject: [PATCH] Fix MPS float64 crash in MiniMax-H3 layout preparation position_ids is built in float64 for cumsum precision over long rotary time spans, but MPS has no float64 support, so moving it to an MPS device crashed. Downcast to float32 only on device transfer, only for MPS/NPU, using the existing maybe_adjust_dtype_for_device helper already used throughout the codebase for this exact limitation. The rope embedding already casts position_ids to float32 on every device before use, so this changes nothing downstream. Fixes #14639 --- .../minimax_h3/before_denoise.py | 12 +++- .../test_modular_pipeline_minimax_h3.py | 64 +++++++++++++++++++ 2 files changed, 73 insertions(+), 3 deletions(-) diff --git a/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py b/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py index c670467a9307..9014bc0a6ccf 100644 --- a/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py +++ b/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py @@ -17,7 +17,7 @@ from ...schedulers import MiniMaxH3Scheduler from ...utils import logging -from ...utils.torch_utils import randn_tensor +from ...utils.torch_utils import maybe_adjust_dtype_for_device, randn_tensor from ..modular_pipeline import ModularPipelineBlocks, PipelineState from ..modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, OutputParam from .modular_pipeline import ( @@ -441,7 +441,10 @@ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) - components.video_tag, block_state.keyframe_anchors, ) - block_state.position_ids = position_ids.to(device) + # `position_ids` is built in float64 for cumsum precision over long sequences (see `_temporal_position_grid`); + # MPS/NPU have no float64 support, so only those devices get the float32 downcast the rope embedding applies + # anyway (`MiniMaxH3RotaryPosEmbed.forward`). + block_state.position_ids = position_ids.to(device, dtype=maybe_adjust_dtype_for_device(torch.float64, device)) block_state.token_tags = token_tags.to(device) block_state.video_indices = video_indices.to(device) block_state.audio_indices = audio_indices.to(device) @@ -765,7 +768,10 @@ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) - components.audio_tag, components.video_tag, ) - block_state.position_ids = position_ids.to(device) + # `position_ids` is built in float64 for cumsum precision over long sequences (see `_temporal_position_grid`); + # MPS/NPU have no float64 support, so only those devices get the float32 downcast the rope embedding applies + # anyway (`MiniMaxH3RotaryPosEmbed.forward`). + block_state.position_ids = position_ids.to(device, dtype=maybe_adjust_dtype_for_device(torch.float64, device)) block_state.token_tags = token_tags.to(device) block_state.video_indices = video_indices.to(device) block_state.audio_indices = audio_indices.to(device) diff --git a/tests/modular_pipelines/minimax_h3/test_modular_pipeline_minimax_h3.py b/tests/modular_pipelines/minimax_h3/test_modular_pipeline_minimax_h3.py index 2f653b28e296..2d1a5cfbfbbb 100644 --- a/tests/modular_pipelines/minimax_h3/test_modular_pipeline_minimax_h3.py +++ b/tests/modular_pipelines/minimax_h3/test_modular_pipeline_minimax_h3.py @@ -27,6 +27,7 @@ MiniMaxH3ImageReference, MiniMaxH3VideoReference, ) +from diffusers.modular_pipelines.minimax_h3.before_denoise import MiniMaxH3PrepareLayoutStep from diffusers.modular_pipelines.minimax_h3.before_encoder import MiniMaxH3Ref2VASetupStep from diffusers.modular_pipelines.minimax_h3.encoders import ( MiniMaxH3FL2VATextEncoderStep, @@ -35,7 +36,9 @@ MiniMaxH3TextEncoderStep, ) from diffusers.modular_pipelines.minimax_h3.modular_pipeline import MINIMAX_H3_FPS +from diffusers.utils.torch_utils import maybe_adjust_dtype_for_device +from ...testing_utils import torch_device from ..testing_utils import ( BaseModularPipelineTesterConfig, ModularLoadingTesterMixin, @@ -894,3 +897,64 @@ def test_video_reference_resampled_to_the_model_frame_rate(self): frames, fps=float(MINIMAX_H3_FPS), num_frames=124, **canvas ) assert np.shares_memory(untouched, frames) + + +class TestMiniMaxH3PositionIdsDevice: + """ + `MiniMaxH3PrepareLayoutStep.build_packed_sequence` lays `position_ids` out in float64: the temporal axis is a + `cumsum` over a `5/3`-scaled span (see `_temporal_position_grid`), and float32 drifts perceptibly over the + thousands of latent frames a long video packs. MPS (and NPU) have no float64 support, so the layout steps have to + downcast when they move `position_ids` onto the execution device — see + https://github.com/huggingface/diffusers/issues/14639. Neither this nor the checkpoint-free + [`TestMiniMaxH3Reference`] above needs a checkpoint or an accelerator to run. + """ + + def test_position_ids_are_float64_on_the_host(self): + r"""The packed layout is always built in float64, regardless of what device it later moves to.""" + text_token_tags = torch.ones(4, dtype=torch.int64) + position_ids, *_ = MiniMaxH3PrepareLayoutStep.build_packed_sequence( + text_token_tags, + num_latent_frames=2, + latent_height=16, + latent_width=16, + num_audio_latents=2, + patch_size=(1, 8, 8), + audio_channels=1, + audio_tag=2, + video_tag=0, + ) + assert position_ids.dtype == torch.float64 + + @pytest.mark.parametrize( + "device_type, expected_dtype", + [("cpu", torch.float64), ("cuda", torch.float64), ("mps", torch.float32)], + ) + def test_device_transfer_dtype_matches_the_execution_device(self, device_type, expected_dtype): + r""" + The dtype the layout steps pass to `position_ids.to(device, dtype=...)` is float32 exactly on the devices + that cannot hold float64, and float64 (i.e. a no-op cast) everywhere else — including on a machine that has + none of these backends, since building a `torch.device` object needs no hardware. NPU is also downcast by + `maybe_adjust_dtype_for_device`, but constructing a `torch.device("npu")` needs the `torch_npu` plugin + installed, so it is exercised only by that helper's own device-agnostic logic, not by an actual device here. + """ + assert maybe_adjust_dtype_for_device(torch.float64, torch.device(device_type)) == expected_dtype + + @pytest.mark.skipif(torch_device != "mps", reason="exercises the actual MPS float64 tensor-move limitation") + def test_position_ids_move_to_mps_as_float32(self): + r"""On real MPS hardware, the layout step's device transfer no longer tries to materialize a float64 tensor.""" + text_token_tags = torch.ones(4, dtype=torch.int64) + position_ids, *_ = MiniMaxH3PrepareLayoutStep.build_packed_sequence( + text_token_tags, + num_latent_frames=2, + latent_height=16, + latent_width=16, + num_audio_latents=2, + patch_size=(1, 8, 8), + audio_channels=1, + audio_tag=2, + video_tag=0, + ) + device = torch.device(torch_device) + moved = position_ids.to(device, dtype=maybe_adjust_dtype_for_device(torch.float64, device)) + assert moved.dtype == torch.float32 + assert moved.device.type == "mps"