[MiniMax-H3] Distribute VAE tiles across a process group

docs/source/en/api/models/autoencoderkl_minimax_h3.md

@@ -25,11 +25,32 @@ The temporal geometry is fixed by `clip_length` (17 pixel frames per encoder chu
 import torch
 from diffusers import AutoencoderKLMiniMaxH3
 
-vae = AutoencoderKLMiniMaxH3.from_pretrained(
-    "MiniMaxAI/MiniMax-H3", subfolder="vae", dtype=torch.float32
-).to("cuda")  # or "mps", "xpu", "cpu"
+vae = AutoencoderKLMiniMaxH3.from_pretrained("MiniMaxAI/MiniMax-H3", subfolder="vae", dtype=torch.float32).to(
+    "cuda"
+)  # or "mps", "xpu", "cpu"
 ```
 
+## Parallel tiled decoding
+
+Pass an initialized PyTorch distributed process group to `enable_tiling` to distribute the original spatial tiles
+across its ranks. Temporal chunks, tile overlaps, and blending stay unchanged. Each rank holds the full VAE and
+receives the complete decoded output; encoding remains local.
+
+Initialize the process group and select each rank's device before loading the VAE. Every member must use identical
+weights, latents, and tile settings and call `decode` together. For example, in an existing distributed application:
+
+```python
+import torch.distributed as dist
+
+vae.enable_tiling(process_group=dist.group.WORLD)  # or an existing subgroup
+with torch.inference_mode():
+    decoded = vae.decode(latents).sample
+```
+
+This option requires inference without gradients and runs outside a full-graph `torch.compile` region. Call
+`vae.enable_tiling()` without a process group to return to local tiled decoding, or `vae.disable_tiling()` to disable
+spatial tiling altogether.
+
 ## AutoencoderKLMiniMaxH3
 
 [[autodoc]] AutoencoderKLMiniMaxH3

src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py

@@ -15,6 +15,7 @@
 import math
 
 import torch
+import torch.distributed as dist
 import torch.nn as nn
 import torch.nn.functional as F
 
@@ -617,6 +618,7 @@ def __init__(
         # into smaller tiles, running the encoder/decoder per tile and blending the overlaps, the memory requirement
         # can be lowered. MiniMax-H3 ships with tiling enabled.
         self.use_tiling = True
+        self._tiling_process_group = None
 
         # The tile size in pixel space, and the minimum overlap between two neighbouring tiles. The actual overlaps are
         # widened (in multiples of `spatial_compression_ratio`) so that the tiles cover the frame exactly.
@@ -631,6 +633,7 @@ def enable_tiling(
         tile_sample_min_width: int | None = None,
         tile_sample_min_overlap_height: int | None = None,
         tile_sample_min_overlap_width: int | None = None,
+        process_group: dist.ProcessGroup | None = None,
     ) -> None:
         r"""
         Enable tiled VAE encoding/decoding. When this option is enabled, the VAE splits the frames into tiles, encodes
@@ -646,8 +649,20 @@ def enable_tiling(
                 The minimum overlap, in pixels, between two consecutive vertical tiles.
             tile_sample_min_overlap_width (`int`, *optional*):
                 The minimum overlap, in pixels, between two consecutive horizontal tiles.
+            process_group (`torch.distributed.ProcessGroup`, *optional*):
+                Distribute decoding tiles across this initialized process group. Every member must have the same model,
+                latents, and tile settings and call `decode` together under `torch.no_grad()` or
+                `torch.inference_mode()`. The model and latents must be on each rank's local device. Encoding remains
+                local. Defaults to local decoding when unset.
         """
+        if process_group is not None:
+            if not dist.is_available() or not dist.is_initialized():
+                raise ValueError("Initialize torch.distributed before enabling parallel tiled decoding.")
+            if dist.get_rank(process_group) < 0:
+                raise ValueError("The current rank must belong to the tiled decoding process group.")
+
         self.use_tiling = True
+        self._tiling_process_group = process_group
         self.tile_sample_min_height = tile_sample_min_height or self.tile_sample_min_height
         self.tile_sample_min_width = tile_sample_min_width or self.tile_sample_min_width
         self.tile_sample_min_overlap_height = tile_sample_min_overlap_height or self.tile_sample_min_overlap_height
@@ -759,16 +774,46 @@ def _decode_clip(self, z: torch.Tensor) -> torch.Tensor:
         )
 
         ratio = self.spatial_compression_ratio
-        rows = []
-        for i_pos, i_len in zip(y_indices, y_lengths):
-            row = []
-            for j_pos, j_len in zip(x_indices, x_lengths):
+        group = self._tiling_process_group
+        world_size = dist.get_world_size(group) if group is not None else 1
+        rank = dist.get_rank(group) if group is not None else 0
+        if world_size > 1 and torch.is_grad_enabled():
+            raise RuntimeError("Parallel tiled decoding requires torch.no_grad() or torch.inference_mode().")
+
+        tile_layout = [
+            (i_pos, i_len, j_pos, j_len)
+            for i_pos, i_len in zip(y_indices, y_lengths)
+            for j_pos, j_len in zip(x_indices, x_lengths)
+        ]
+        decoded_tiles = {}
+        for tile_index, (i_pos, i_len, j_pos, j_len) in enumerate(tile_layout):
+            if tile_index % world_size == rank:
                 tile = z[
                     ...,
                     i_pos // ratio : i_pos // ratio + i_len // ratio,
                     j_pos // ratio : j_pos // ratio + j_len // ratio,
                 ]
-                row.append(self.decoder(self.post_quant_conv(tile)))
+                decoded = self.decoder(self.post_quant_conv(tile))
+                decoded_tiles[tile_index] = decoded.contiguous() if world_size > 1 else decoded
+
+        if world_size > 1:
+            # The splitter uses equal tile extents; autocast can change the decoder's output dtype.
+            metadata = [(decoded_tiles[0].shape, decoded_tiles[0].dtype) if rank == 0 else None]
+            dist.broadcast_object_list(metadata, src=dist.get_global_rank(group, 0), group=group, device=z.device)
+            tile_shape, tile_dtype = metadata[0]
+
+        rows = []
+        for i in range(len(y_indices)):
+            row = []
+            for j in range(len(x_indices)):
+                tile_index = i * len(x_indices) + j
+                if tile_index in decoded_tiles:
+                    tile = decoded_tiles.pop(tile_index)
+                else:
+                    tile = torch.empty(tile_shape, dtype=tile_dtype, device=z.device)
+                if world_size > 1:
+                    dist.broadcast(tile, src=dist.get_global_rank(group, tile_index % world_size), group=group)
+                row.append(tile)
             rows.append(row)
 
         return self._stitch_tiles(rows, y_overlaps, x_overlaps)

tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py

@@ -13,8 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from datetime import timedelta
+
 import pytest
 import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
 
 from diffusers import AutoencoderKLMiniMaxH3
 from diffusers.utils.torch_utils import randn_tensor
@@ -40,6 +44,52 @@
 HEIGHT = WIDTH = 8
 
 
+def _parallel_tiled_decode_worker(rank, world_size, init_file, model_config):
+    torch.set_num_threads(1)
+    dist.init_process_group(
+        "gloo", init_method=f"file://{init_file}", rank=rank, world_size=world_size, timeout=timedelta(seconds=60)
+    )
+    try:
+        torch.manual_seed(0)
+        model = AutoencoderKLMiniMaxH3(**model_config).eval()
+        subgroup = dist.new_group([1, 2])
+        tile_settings = {
+            "tile_sample_min_height": 8,
+            "tile_sample_min_width": 8,
+            "tile_sample_min_overlap_height": 4,
+            "tile_sample_min_overlap_width": 4,
+        }
+        calls = []
+        handle = model.decoder.register_forward_hook(lambda *_: calls.append(1))
+        for group in (dist.group.WORLD, subgroup):
+            if dist.get_rank(group) < 0:
+                with pytest.raises(ValueError, match="must belong"):
+                    model.enable_tiling(**tile_settings, process_group=group)
+            else:
+                # Multiple tiles, then more ranks than tiles with the released temporal overlap and autocast.
+                for latent_shape, autocast in (((2, 4, 7, 3, 4), False), ((1, 4, 12, 1, 2), True)):
+                    z = torch.randn(latent_shape, generator=torch.Generator().manual_seed(1))
+                    model.enable_tiling(**tile_settings)
+                    calls.clear()
+                    with torch.inference_mode(), torch.autocast("cpu", dtype=torch.bfloat16, enabled=autocast):
+                        expected = model.decode(z).sample
+                        serial_calls = len(calls)
+                        calls.clear()
+                        model.enable_tiling(**tile_settings, process_group=group)
+                        actual = model.decode(z).sample
+                    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+                    assert actual.dtype == (torch.bfloat16 if autocast else torch.float32)
+                    local_calls = torch.tensor(len(calls))
+                    dist.all_reduce(local_calls, group=group)
+                    assert local_calls.item() == serial_calls
+                with pytest.raises(RuntimeError, match="requires torch.no_grad"):
+                    model.decode(z)
+            dist.barrier()
+        handle.remove()
+    finally:
+        dist.destroy_process_group()
+
+
 class AutoencoderKLMiniMaxH3TesterConfig(BaseModelTesterConfig):
     @property
     def model_class(self):
@@ -192,3 +242,14 @@ class TestAutoencoderKLMiniMaxH3Attention(AutoencoderKLMiniMaxH3TesterConfig, At
 
 class TestAutoencoderKLMiniMaxH3SlicingTiling(AutoencoderKLMiniMaxH3TesterConfig, AutoencoderTesterMixin):
     """Slicing and tiling tests for the MiniMax-H3 video autoencoder."""
+
+    def test_parallel_tiled_decoding(self, tmp_path):
+        """Original tiles and temporal blending match across WORLD and a subgroup without global rank zero."""
+        if not dist.is_available() or not dist.is_gloo_available():
+            pytest.skip("Parallel tiled decoding requires torch.distributed with Gloo.")
+        mp.spawn(
+            _parallel_tiled_decode_worker,
+            args=(3, str(tmp_path / "gloo_init"), self.get_init_dict()),
+            nprocs=3,
+            join=True,
+        )