From 902d9cf5e1e9e23d312f2e1787343b523dd83ce3 Mon Sep 17 00:00:00 2001 From: Lele Date: Fri, 16 Oct 2026 20:14:21 +0800 Subject: [PATCH] [MiniMax-H3] Reuse Ulysses shard lengths across an inference forward Add an inference-only opt-in that validates the global packed length once per forward and passes an immutable shard plan across H3 transformer blocks. Keep the existing collection path by default and avoid persistent caches. Signed-off-by: LeslieWylie <112309453+LeslieWylie@users.noreply.github.com> Co-authored-by: Codex --- .../en/training/distributed_inference.md | 14 ++++ src/diffusers/models/_modeling_parallel.py | 2 + src/diffusers/models/attention_dispatch.py | 14 +++- .../transformers/transformer_minimax_h3.py | 36 +++++++-- .../test_models_transformer_minimax_h3.py | 79 +++++++++++++++++++ 5 files changed, 136 insertions(+), 9 deletions(-) diff --git a/docs/source/en/training/distributed_inference.md b/docs/source/en/training/distributed_inference.md index 818488b4cfce..d9c8893aeacd 100644 --- a/docs/source/en/training/distributed_inference.md +++ b/docs/source/en/training/distributed_inference.md @@ -354,6 +354,20 @@ The default Ulysses Attention mechanism requires that the sequence length of hid pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2, ulysses_anything=True)) ``` +MiniMax-H3 can reuse its packed-sequence shard lengths across its transformer blocks in one inference forward: + +```python +transformer = pipeline.transformer +with torch.no_grad(): + output = transformer(**model_inputs, attention_kwargs={"reuse_ulysses_metadata": True}) +``` + +Enable Ulysses Anything on the transformer first. Every rank must use the same opt-in setting for a forward and receive +the same complete packed layout. The token refiner retains the default length collection. The model +checks the global sequence length once on every forward, then passes an immutable per-call plan to its self-attention +blocks. It does not cache plans between requests or infer a global layout from a rank-local length. Calls with gradients +enabled are rejected for this option. Other models, normal Ulysses, Ring Attention, and the default H3 path are unchanged. + > [!TIP] To avoid multiple forced CUDA sync caused by H2D and D2H transfers, please add the **gloo** backend in `init_process_group`. This will significantly reduce communication latency. We ran a benchmark for FLUX.1-dev with Ulysses, Ring, Unified Attention and Ulysses Anything Attention with [this script](https://github.com/huggingface/diffusers/pull/12996#issuecomment-3797695999) on a node of 4 L20 GPUs. The results are summarized as follows: diff --git a/src/diffusers/models/_modeling_parallel.py b/src/diffusers/models/_modeling_parallel.py index 58cbc822f8bf..3d7de521c126 100644 --- a/src/diffusers/models/_modeling_parallel.py +++ b/src/diffusers/models/_modeling_parallel.py @@ -210,6 +210,8 @@ class ParallelConfig: context_parallel_config: ContextParallelConfig | None = None tensor_parallel_config: TensorParallelConfig | None = None + _ulysses_sequence_lengths: tuple[int, ...] | None = None + _rank: int = None _world_size: int = None _device: torch.device = None diff --git a/src/diffusers/models/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py index 18e575d8ebbc..9726987780a1 100644 --- a/src/diffusers/models/attention_dispatch.py +++ b/src/diffusers/models/attention_dispatch.py @@ -2326,7 +2326,9 @@ def all_to_all_single_any_qkv_async( # S_LOCAL maybe not equal for all ranks in dynamic shape case, # since we don't know the actual shape before this timing, thus, # we have to use all gather to collect the S_LOCAL first. - output_split_sizes = gather_size_by_comm(S_LOCAL, group) + output_split_sizes = kwargs.get("seq_lengths") + if output_split_sizes is None: + output_split_sizes = gather_size_by_comm(S_LOCAL, group) x = x.flatten(0, 1) # (world_size * S_LOCAL, B, H_LOCAL, D) x = funcol.all_to_all_single(x, output_split_sizes, input_split_sizes, group) @@ -2365,7 +2367,9 @@ def all_to_all_single_any_o_async(x: torch.Tensor, group: dist.ProcessGroup, **k # b.tensor_split(4)[0].shape[1]) S_LOCAL = kwargs.get("Q_S_LOCAL") - input_split_sizes = gather_size_by_comm(S_LOCAL, group) + input_split_sizes = kwargs.get("seq_lengths") + if input_split_sizes is None: + input_split_sizes = gather_size_by_comm(S_LOCAL, group) x = x.permute(1, 0, 2, 3).contiguous() # (S_GLOBAL, B, H_LOCAL, D) output_split_sizes = [S_LOCAL] * world_size x = funcol.all_to_all_single(x, output_split_sizes, input_split_sizes, group) @@ -2780,6 +2784,12 @@ def forward( _, S_KV_LOCAL, _, _ = key.shape metadata = ulysses_anything_metadata(query) + sequence_lengths = _parallel_config._ulysses_sequence_lengths + if sequence_lengths is not None: + local_length = sequence_lengths[dist.get_rank(group)] + if query.shape[1] != local_length or key.shape[1] != local_length: + raise ValueError("The forward sequence plan requires matching self-attention shard lengths.") + metadata = {**metadata, "seq_lengths": list(sequence_lengths)} query_wait = all_to_all_single_any_qkv_async(query, group, **metadata) key_wait = all_to_all_single_any_qkv_async(key, group, **metadata) value_wait = all_to_all_single_any_qkv_async(value, group, **metadata) diff --git a/src/diffusers/models/transformers/transformer_minimax_h3.py b/src/diffusers/models/transformers/transformer_minimax_h3.py index 95ee1f2ac511..b6ca7eb192e2 100644 --- a/src/diffusers/models/transformers/transformer_minimax_h3.py +++ b/src/diffusers/models/transformers/transformer_minimax_h3.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import Any import torch @@ -21,7 +21,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FromOriginalModelMixin, PeftAdapterMixin from ...utils import BaseOutput, apply_lora_scale, logging -from .._modeling_parallel import ContextParallelInput, ContextParallelOutput +from .._modeling_parallel import ContextParallelInput, ContextParallelOutput, ParallelConfig, gather_size_by_comm from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward from ..attention_dispatch import dispatch_attention_fn from ..cache_utils import CacheMixin @@ -171,6 +171,7 @@ def __call__( hidden_states: torch.Tensor, rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, attention_mask: torch.Tensor | None = None, + parallel_config: ParallelConfig | None = None, ) -> torch.Tensor: if attn.fused_projections: query, key, value = attn.to_qkv(hidden_states).chunk(3, dim=-1) @@ -201,7 +202,7 @@ def __call__( dropout_p=0.0, is_causal=False, backend=self._attention_backend, - parallel_config=self._parallel_config, + parallel_config=parallel_config if parallel_config is not None else self._parallel_config, ) hidden_states = hidden_states.flatten(2, 3).type_as(query) hidden_states = attn.to_out[0](hidden_states) @@ -243,8 +244,11 @@ def forward( hidden_states: torch.Tensor, rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, attention_mask: torch.Tensor | None = None, + parallel_config: ParallelConfig | None = None, ) -> torch.Tensor: - return self.processor(self, hidden_states, rotary_emb, attention_mask) + if parallel_config is None: + return self.processor(self, hidden_states, rotary_emb, attention_mask) + return self.processor(self, hidden_states, rotary_emb, attention_mask, parallel_config=parallel_config) class MiniMaxH3TokenRefinerBlock(nn.Module): @@ -351,6 +355,7 @@ def forward( adaln_indices: torch.Tensor, rotary_emb: tuple[torch.Tensor, torch.Tensor], attention_mask: torch.Tensor | None = None, + parallel_config: ParallelConfig | None = None, ) -> torch.Tensor: shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaln_proj(temb) @@ -359,7 +364,7 @@ def forward( norm_hidden_states = norm_hidden_states * ( 1.0 + scale_msa.index_select(0, adaln_indices) ) + shift_msa.index_select(0, adaln_indices) - attn_output = self.attn(norm_hidden_states, rotary_emb, attention_mask) + attn_output = self.attn(norm_hidden_states, rotary_emb, attention_mask, parallel_config=parallel_config) hidden_states = residual + gate_msa.index_select(0, adaln_indices) * attn_output residual = hidden_states @@ -601,7 +606,8 @@ def forward( text_indices (`torch.Tensor` of shape `(num_text_tokens,)`): Positions of the text rows in the packed sequence. attention_kwargs (`dict`, *optional*): - A kwargs dictionary that, if specified, may carry a `scale` entry which is applied to the LoRA layers. + May carry a LoRA `scale` and `reuse_ulysses_metadata=True` to reuse packed-sequence shard lengths + within this inference forward when using Ulysses Anything context parallelism. return_dict (`bool`, defaults to `True`): Whether to return a [`MiniMaxH3TransformerOutput`] instead of a plain tuple. @@ -646,13 +652,29 @@ def forward( # 3. Row -> AdaLN table row. adaln_indices = timestep_indices * MINIMAX_H3_MODALITY_NUM + token_tags + parallel_config = None + if attention_kwargs is not None and attention_kwargs.get("reuse_ulysses_metadata", False): + parallel_config = self._parallel_config + cp_config = parallel_config.context_parallel_config if parallel_config is not None else None + if torch.is_grad_enabled() or cp_config is None or not cp_config.ulysses_anything: + raise ValueError("Sequence metadata reuse requires inference with Ulysses Anything enabled.") + group = cp_config._ulysses_mesh.get_group() + global_lengths = gather_size_by_comm(sequence_length, group) + if any(length != sequence_length for length in global_lengths): + raise ValueError("All Ulysses ranks must agree on the packed sequence length for this forward.") + degree = len(global_lengths) + shard_lengths = tuple( + sequence_length // degree + int(rank < sequence_length % degree) for rank in range(degree) + ) + parallel_config = replace(parallel_config, _ulysses_sequence_lengths=shard_lengths) + for block in self.transformer_blocks: if torch.is_grad_enabled() and self.gradient_checkpointing: hidden_states = self._gradient_checkpointing_func( block, hidden_states, temb, adaln_indices, rotary_emb ) else: - hidden_states = block(hidden_states, temb, adaln_indices, rotary_emb) + hidden_states = block(hidden_states, temb, adaln_indices, rotary_emb, parallel_config=parallel_config) # 5. Both heads run over every row, then the rows of each modality are selected. The heads are listed in # `_keep_in_fp32_modules`, so they stay float32 while the block stack runs in the requested `torch_dtype`; diff --git a/tests/models/transformers/test_models_transformer_minimax_h3.py b/tests/models/transformers/test_models_transformer_minimax_h3.py index 4165f9b5392d..6430f1201f62 100644 --- a/tests/models/transformers/test_models_transformer_minimax_h3.py +++ b/tests/models/transformers/test_models_transformer_minimax_h3.py @@ -13,9 +13,17 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dataclasses import replace +from datetime import timedelta + +import pytest import torch +import torch.distributed as dist +import torch.multiprocessing as mp from diffusers import MiniMaxH3Transformer3DModel +from diffusers.models._modeling_parallel import ContextParallelConfig, gather_size_by_comm +from diffusers.models.attention_dispatch import dispatch_attention_fn from diffusers.models.transformers.transformer_minimax_h3 import MiniMaxH3TransformerOutput from diffusers.utils.torch_utils import randn_tensor @@ -36,6 +44,77 @@ enable_full_determinism() +def _sequence_plan_worker(rank, port): + torch.set_num_threads(1) + dist.init_process_group( + "cpu:gloo", init_method=f"tcp://127.0.0.1:{port}", rank=rank, world_size=2, timeout=timedelta(seconds=60) + ) + try: + fixture = MiniMaxH3TransformerTesterConfig() + torch.manual_seed(7) + model = MiniMaxH3Transformer3DModel(**fixture.get_init_dict()).eval() + model.set_attention_backend("native") + mesh = dist.device_mesh.init_device_mesh("cpu", (1, 2), mesh_dim_names=("ring", "ulysses")) + model.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2, ulysses_anything=True, mesh=mesh)) + original_config = model._parallel_config + with torch.no_grad(): + for video_tokens in (8, 10, 8): + inputs = fixture.get_dummy_inputs(num_video_tokens=video_tokens, batch_size=1) + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU]) as baseline_profile: + expected = model(**inputs, return_dict=False) + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU]) as planned_profile: + actual = model(**inputs, attention_kwargs={"reuse_ulysses_metadata": True}, return_dict=False) + for output, reference in zip(actual, expected): + torch.testing.assert_close(output, reference, rtol=0, atol=0) + counts = [ + sum(event.count for event in profile.key_averages() if "allgather" in event.key.lower()) + for profile in (baseline_profile, planned_profile) + ] + assert counts[1] < counts[0], counts + assert model._parallel_config is original_config + assert original_config._ulysses_sequence_lengths is None + + # Rank zero keeps five rows across 9 -> 10 -> 9, while rank one's length changes. + for sequence_length in (9, 10, 9): + lengths = gather_size_by_comm(sequence_length, mesh["ulysses"].get_group()) + assert lengths == [sequence_length, sequence_length] + plan = replace( + original_config, + _ulysses_sequence_lengths=(sequence_length // 2 + sequence_length % 2, sequence_length // 2), + ) + torch.manual_seed(5) + q, k, v = (torch.randn(1, sequence_length, 2, 16) for _ in range(3)) + q, k, v = (x.tensor_split(2, dim=1)[rank].contiguous() for x in (q, k, v)) + expected = dispatch_attention_fn(q, k, v, backend="native", parallel_config=original_config) + actual = dispatch_attention_fn(q, k, v, backend="native", parallel_config=plan) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + with pytest.raises(ValueError, match="agree on the packed sequence length"): + model( + **fixture.get_dummy_inputs(num_video_tokens=8 + 2 * rank, batch_size=1), + attention_kwargs={"reuse_ulysses_metadata": True}, + ) + assert model._parallel_config is original_config + assert original_config._ulysses_sequence_lengths is None + finally: + dist.destroy_process_group() + + +class TestMiniMaxH3SequencePlan: + def test_forward_plan_reuses_collectives_without_leaking_between_requests(self): + if torch_device != "cpu" or not dist.is_gloo_available(): + pytest.skip("This focused distributed test uses CPU Gloo.") + from ..testing_utils.parallelism import _find_free_port + + mp.spawn(_sequence_plan_worker, args=(_find_free_port(),), nprocs=2, join=True) + + def test_metadata_reuse_requires_inference_context_parallelism(self): + fixture = MiniMaxH3TransformerTesterConfig() + model = MiniMaxH3Transformer3DModel(**fixture.get_init_dict()) + with pytest.raises(ValueError, match="inference with Ulysses Anything"): + model(**fixture.get_dummy_inputs(), attention_kwargs={"reuse_ulysses_metadata": True}) + + # The packed layout the dummy inputs describe: a text block, then the audio rows, then the target video rows, which # is the layout `MiniMaxH3Blocks` builds. NUM_TEXT_TOKENS = 4