From 8d12f46f40ef9b89a86334c6206e7ad93a3c58d6 Mon Sep 17 00:00:00 2001 From: Chakshu Dhannawat Date: Thu, 27 Aug 2026 14:53:37 +0900 Subject: [PATCH] Align cross-device text encoder inputs in MiniMax-H3 transformer forward --- .../transformers/transformer_minimax_h3.py | 5 +++- .../test_models_transformer_minimax_h3.py | 27 ++++++++++++++++++- 2 files changed, 30 insertions(+), 2 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_minimax_h3.py b/src/diffusers/models/transformers/transformer_minimax_h3.py index f49cdaca2eb6..71d553f88a17 100644 --- a/src/diffusers/models/transformers/transformer_minimax_h3.py +++ b/src/diffusers/models/transformers/transformer_minimax_h3.py @@ -92,7 +92,7 @@ def __init__(self, rope_freq_dim: int = 16, rope_theta: float = 10000.0): def forward(self, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: # position_ids: (seq_len, 3) -> cos/sin: (seq_len, 2 * 3 * rope_freq_dim) - position_ids = position_ids.to(torch.float32) + position_ids = position_ids.to(self.inv_freq.device).to(torch.float32) freqs = position_ids.unsqueeze(-1) * self.inv_freq.view(1, 1, -1) # (seq_len, 3, rope_freq_dim) freqs_t, freqs_h, freqs_w = freqs.unbind(dim=1) freqs = torch.cat((freqs_t, freqs_h, freqs_w), dim=-1) @@ -619,6 +619,9 @@ def forward( f"{list(token_tags.shape)} and {list(timestep_indices.shape)} for seq_len={sequence_length}." ) + if encoder_hidden_states is not None and encoder_hidden_states.device != self.device: + encoder_hidden_states = encoder_hidden_states.to(device=self.device) + rotary_emb = self.rope(position_ids) # 1. Project each modality and scatter the rows into the packed sequence buffer. The checkpoint is diff --git a/tests/models/transformers/test_models_transformer_minimax_h3.py b/tests/models/transformers/test_models_transformer_minimax_h3.py index 00baa37c84a0..bdb7dfef7f31 100644 --- a/tests/models/transformers/test_models_transformer_minimax_h3.py +++ b/tests/models/transformers/test_models_transformer_minimax_h3.py @@ -19,7 +19,7 @@ from diffusers.models.transformers.transformer_minimax_h3 import MiniMaxH3TransformerOutput from diffusers.utils.torch_utils import randn_tensor -from ...testing_utils import enable_full_determinism, torch_device +from ...testing_utils import enable_full_determinism, require_torch_multi_gpu, torch_device from ..testing_utils import ( AttentionTesterMixin, BaseModelTesterConfig, @@ -189,3 +189,28 @@ class TestMiniMaxH3TransformerContextParallel(MiniMaxH3TransformerTesterConfig, class TestMiniMaxH3TransformerLoRA(MiniMaxH3TransformerTesterConfig, LoraTesterMixin): """LoRA tests for the MiniMax-H3 transformer.""" + + +class TestMiniMaxH3TransformerMultiGPU(MiniMaxH3TransformerTesterConfig): + """Multi-GPU tests for the MiniMax-H3 transformer.""" + + @require_torch_multi_gpu + def test_cross_device_text_encoder_forward(self): + """The documented multi-GPU setup places the text encoder on a different device than the transformer; make + sure `encoder_hidden_states` and `position_ids` are moved to the transformer's device before use.""" + model = self.model_class(**self.get_init_dict()).to("cuda:0").eval() + inputs = self.get_dummy_inputs() + + # Simulate text-encoder-on-cuda:1, transformer-on-cuda:0. + cross_device_inputs = dict(inputs) + cross_device_inputs["encoder_hidden_states"] = cross_device_inputs["encoder_hidden_states"].to("cuda:1") + cross_device_inputs["position_ids"] = cross_device_inputs["position_ids"].to("cuda:1") + + with torch.no_grad(): + baseline_output = model(**inputs) + cross_device_output = model(**cross_device_inputs) + + assert cross_device_output.sample.device == torch.device("cuda:0") + assert cross_device_output.audio_sample.device == torch.device("cuda:0") + torch.testing.assert_close(cross_device_output.sample, baseline_output.sample) + torch.testing.assert_close(cross_device_output.audio_sample, baseline_output.audio_sample)