From e4cbf989f306a9593173ac7ddb66791387739b7b Mon Sep 17 00:00:00 2001 From: Lele Date: Fri, 16 Oct 2026 20:14:01 +0800 Subject: [PATCH] [MiniMax-H3] Add opt-in process-group tiled VAE decoding Decode original tiles across an explicit process group and restore the original tile order before temporal and spatial blending. Co-authored-by: Codex Signed-off-by: LeslieWylie <112309453+LeslieWylie@users.noreply.github.com> --- .../en/api/models/autoencoderkl_minimax_h3.md | 27 +++++++- .../autoencoders/autoencoder_kl_minimax_h3.py | 55 +++++++++++++++-- .../test_models_autoencoder_kl_minimax_h3.py | 61 +++++++++++++++++++ 3 files changed, 135 insertions(+), 8 deletions(-) diff --git a/docs/source/en/api/models/autoencoderkl_minimax_h3.md b/docs/source/en/api/models/autoencoderkl_minimax_h3.md index 1df2aa27241e..bb79ec60e988 100644 --- a/docs/source/en/api/models/autoencoderkl_minimax_h3.md +++ b/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 diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py b/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py index 3f5f44e79769..0c038fd27fa1 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py +++ b/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) diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py b/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py index c7871b0620e2..fef39d8cad6a 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py +++ b/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, + )