[core] support cp in h3. (#14407)

* support cp in h3.

* empty

Author: sayakpaul

src/diffusers/models/transformers/transformer_minimax_h3.py

@@ -21,6 +21,7 @@
 from ...configuration_utils import ConfigMixin, register_to_config
 from ...loaders import PeftAdapterMixin
 from ...utils import BaseOutput, apply_lora_scale, logging
+from .._modeling_parallel import ContextParallelInput, ContextParallelOutput
 from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward
 from ..attention_dispatch import dispatch_attention_fn
 from ..cache_utils import CacheMixin
@@ -446,6 +447,36 @@ class MiniMaxH3Transformer3DModel(ModelMixin, ConfigMixin, AttentionMixin, PeftA
         "audio_proj_out",
         "rope",
     ]
+    # Context parallelism shards the packed sequence, so the split cannot happen on the inputs of `forward`: the rows
+    # of the three modalities are scattered into the packed buffer with sequence-wide indices, which only address the
+    # full sequence. The split therefore happens once the buffer is built, at the first block, and everything that is
+    # indexed per row of the packed sequence is split alongside it: `adaln_indices` for the block stack,
+    # `timestep_indices` for `norm_out`, and the two `rope` outputs, which carry one row of angles each. The two output
+    # heads gather the sequence back, because `forward` then selects the video and audio rows with sequence-wide
+    # indices as well.
+    #
+    # The packed sequence carries no padding, so its length is whatever the caller packed. Splitting it across the
+    # context parallel region requires that length to be divisible by the region size — use `ulysses_anything` or
+    # `ring_anything` for a layout that is not.
+    _cp_plan = {
+        "rope": {
+            0: ContextParallelInput(split_dim=0, expected_dims=2, split_output=True),
+            1: ContextParallelInput(split_dim=0, expected_dims=2, split_output=True),
+        },
+        "transformer_blocks.0": {
+            "hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False),
+        },
+        # `hidden_states` is split once and then flows from block to block already sharded, but every block reads
+        # `adaln_indices` from `forward`'s unsharded tensor, so each of them splits its own copy.
+        "transformer_blocks.*": {
+            "adaln_indices": ContextParallelInput(split_dim=0, expected_dims=1, split_output=False),
+        },
+        "norm_out": {
+            "timestep_indices": ContextParallelInput(split_dim=0, expected_dims=1, split_output=False),
+        },
+        "proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3),
+        "audio_proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3),
+    }
 
     @register_to_config
     def __init__(

tests/models/transformers/test_models_transformer_minimax_h3.py

@@ -23,6 +23,7 @@
 from ..testing_utils import (
     AttentionTesterMixin,
     BaseModelTesterConfig,
+    ContextParallelTesterMixin,
     MemoryTesterMixin,
     ModelTesterMixin,
     TorchCompileTesterMixin,
@@ -121,12 +122,11 @@ def get_packed_layout(self, num_video_tokens: int = NUM_VIDEO_TOKENS) -> dict:
             "text_indices": text_indices,
         }
 
-    def get_dummy_inputs(self, num_video_tokens: int = NUM_VIDEO_TOKENS) -> dict:
+    def get_dummy_inputs(self, num_video_tokens: int = NUM_VIDEO_TOKENS, batch_size: int = 2) -> dict:
         generator = self.generator
         init_dict = self.get_init_dict()
         patch_size = init_dict["patch_size"]
         video_patch_dim = init_dict["in_channels"] * patch_size[0] * patch_size[1] * patch_size[2]
-        batch_size = 2
 
         return {
             "hidden_states": randn_tensor(
@@ -180,3 +180,7 @@ class TestMiniMaxH3TransformerAttention(MiniMaxH3TransformerTesterConfig, Attent
 
 class TestMiniMaxH3TransformerTorchCompile(MiniMaxH3TransformerTesterConfig, TorchCompileTesterMixin):
     """Torch compile tests for the MiniMax-H3 transformer."""
+
+
+class TestMiniMaxH3TransformerContextParallel(MiniMaxH3TransformerTesterConfig, ContextParallelTesterMixin):
+    """Context parallel inference tests for the MiniMax-H3 transformer."""