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