[BugFix][Diffusion] Fix MiniMax-H3 VAE hang when decoder tiles are fewer than ranks (#6345)

Signed-off-by: linzhenpl07 <linzhenpl07@gmail.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

tests/diffusion/models/minimax_h3/test_minimax_h3_vae_tiling.py

@@ -0,0 +1,214 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""CPU tests for the decoder tile-shortage guard in the MiniMax-H3 video VAE."""
+
+import sys
+import types
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.minimax_h3.vae import MiniMaxH3VideoVAE
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+RATIO = 16
+# The checkpoint keeps two independent grids -- ``split_tiles`` selects
+# ``decoder_tile_size``/``decoder_tile_overlap_min`` when ``is_decoder`` is set
+# and ``tile_size``/``tile_overlap_min`` otherwise. The shipped H3 config leaves
+# both at 256/64, so the two happen to coincide there, but that is a config
+# coincidence. The double deliberately gives them different values so that
+# passing the wrong flag changes the tile count and fails a test.
+TILE_SIZE, OVERLAP = 256, 64  # decoder, matching the shipped config
+ENCODER_TILE_SIZE, ENCODER_OVERLAP = 128, 32
+
+
+class _FakeCheckpointModel:
+    """The parts of the checkpoint's AutoencoderKL the guard touches."""
+
+    def __init__(self):
+        self.vae_ratio = RATIO
+        self.parallel_tiling = True
+        self.decoded_with = []
+
+    def split_tiles(self, input_len, is_decoder=False):
+        """Same grid arithmetic as the checkpoint, on pixel dimensions."""
+        tile = TILE_SIZE if is_decoder else ENCODER_TILE_SIZE
+        overlap = OVERLAP if is_decoder else ENCODER_OVERLAP
+        if tile >= input_len:
+            return [0], [input_len], []
+        n = -(-input_len // tile)
+        while tile * n - overlap * (n - 1) - input_len < 0:
+            n += 1
+        return list(range(n)), [tile] * n, [overlap] * (n - 1)
+
+
+def _vae(parallel_size, state):
+    """A stand-in instance: the guard only reads .model/.remote/.parallel_size."""
+    module = types.ModuleType("fake_ckpt.parallel")
+    module.get_parallel_state = lambda: state
+    sys.modules.setdefault("fake_ckpt", types.ModuleType("fake_ckpt"))
+    sys.modules["fake_ckpt.parallel"] = module
+
+    remote = type("Remote", (), {"__module__": "fake_ckpt.klvae"})()
+
+    vae = object.__new__(MiniMaxH3VideoVAE)
+    vae.model = _FakeCheckpointModel()
+    vae.remote = remote
+    vae.parallel_size = parallel_size
+    return vae
+
+
+def _latent(h, w):
+    return torch.zeros(1, 24, 9, h, w)
+
+
+@pytest.mark.parametrize(
+    ("h", "w", "expected"),
+    [
+        (16, 16, 1),  # 256x256 px -> single tile
+        (24, 16, 2),  # 384x256 px
+        (24, 24, 4),  # 384x384 px
+        (48, 84, 28),  # 768x1344 px, the shipped H3 request
+        (96, 168, 112),  # 1536x2688 px, a larger canvas the request path also allows
+    ],
+)
+def test_tile_count_matches_the_checkpoint_grid(h, w, expected):
+    assert MiniMaxH3VideoVAE._decoder_tile_count(_vae(1, {}), _latent(h, w)) == expected
+
+
+@pytest.mark.parametrize(("h", "w", "decoder_tiles", "encoder_tiles"), [(16, 16, 1, 9), (24, 24, 4, 16)])
+def test_tile_count_uses_the_decoder_grid_not_the_encoder_one(h, w, decoder_tiles, encoder_tiles):
+    """Passing ``is_decoder=False`` would silently change the count, so pin it."""
+    vae = _vae(1, {})
+    assert MiniMaxH3VideoVAE._decoder_tile_count(vae, _latent(h, w)) == decoder_tiles
+
+    px_h, px_w = h * RATIO, w * RATIO
+    as_encoder = len(vae.model.split_tiles(px_h, False)[0]) * len(vae.model.split_tiles(px_w, False)[0])
+    assert as_encoder == encoder_tiles != decoder_tiles
+
+
+def test_rank_local_tiling_restores_the_group_state():
+    state = {"sp_size": 4, "sp_rank": 2, "sp_enabled": True, "sp_process_group": "pg"}
+    vae = _vae(4, state)
+
+    with vae._rank_local_tiling():
+        assert state["sp_size"] == 1
+        assert state["sp_rank"] == 0
+        assert state["sp_enabled"] is False
+        assert state["sp_process_group"] is None
+        assert vae.model.parallel_tiling is False
+
+    assert state == {"sp_size": 4, "sp_rank": 2, "sp_enabled": True, "sp_process_group": "pg"}
+    assert vae.model.parallel_tiling is True
+
+
+def test_rank_local_tiling_restores_after_an_exception():
+    state = {"sp_size": 4, "sp_rank": 1, "sp_enabled": True, "sp_process_group": "pg"}
+    vae = _vae(4, state)
+
+    with pytest.raises(RuntimeError, match="decode failed"):
+        with vae._rank_local_tiling():
+            raise RuntimeError("decode failed")
+
+    assert state["sp_size"] == 4
+    assert vae.model.parallel_tiling is True
+
+
+# ---------------------------------------------------------------------------
+# Dispatch: every one of these calls ``decode_latent``, so deleting the guard
+# turns the whole matrix red rather than leaving it green.
+# ---------------------------------------------------------------------------
+
+
+class _RecordingProcessor:
+    @staticmethod
+    def revert_tensor(decoded):
+        return decoded
+
+
+def _dispatch_vae(parallel_size):
+    """A VAE whose ``decode_base`` records the parallel state it ran under."""
+    state = {
+        "group_size": parallel_size,
+        "group_rank": 0,
+        "local_process_group": object(),
+        "sp_size": parallel_size,
+        "sp_rank": 0,
+        "sp_enabled": True,
+        "sp_process_group": object(),
+        "tp_size": 1,
+        "tp_rank": 0,
+    }
+    vae = _vae(parallel_size, state)
+    vae.config_dict = {
+        "latent_channels": 24,
+        "latents_mean": [0.0] * 24,
+        "latents_std": [1.0] * 24,
+    }
+    vae.model.processor = _RecordingProcessor()
+
+    seen = {}
+
+    def decode_base(sample):
+        seen["sp_size"] = state["sp_size"]
+        seen["sp_enabled"] = state["sp_enabled"]
+        seen["sp_process_group"] = state["sp_process_group"]
+        seen["parallel_tiling"] = vae.model.parallel_tiling
+        # the guard never inspects the decoded shape, so keep it small
+        return torch.zeros(1, 3, 2, 4, 4)
+
+    vae.model.decode_base = decode_base
+    return vae, state, seen
+
+
+@pytest.mark.parametrize(
+    ("parallel_size", "h", "w", "falls_back"),
+    [
+        (4, 24, 16, True),  # 2 tiles, ranks 2 and 3 would get none
+        (4, 16, 16, True),  # 1 tile
+        (4, 24, 24, False),  # 4 tiles, exactly enough
+        (4, 48, 84, False),  # 28 tiles, the shipped 768x1344 request
+        (4, 96, 168, False),  # 112 tiles, 1536x2688
+        (2, 24, 16, False),  # 2 tiles across 2 ranks
+        (1, 16, 16, False),  # never parallel, nothing to guard
+    ],
+)
+def test_decode_latent_falls_back_exactly_when_tiles_are_short(parallel_size, h, w, falls_back):
+    """Fewer tiles than ranks is the hang condition; equal or more is fine."""
+    vae, _, seen = _dispatch_vae(parallel_size)
+
+    MiniMaxH3VideoVAE.decode_latent(vae, _latent(h, w))
+
+    # only the guard clears ``parallel_tiling``, so it is the signal that it fired
+    assert (seen["parallel_tiling"] is False) is falls_back
+    assert seen["sp_size"] == (1 if falls_back else parallel_size)
+
+
+def test_decode_latent_isolates_and_restores_the_whole_group_state():
+    """1 tile, 4 ranks: ``decode_base`` must observe rank-local tiling."""
+    vae, state, seen = _dispatch_vae(4)
+
+    frames = MiniMaxH3VideoVAE.decode_latent(vae, _latent(16, 16))
+
+    assert seen["sp_size"] == 1
+    assert seen["sp_enabled"] is False
+    assert seen["sp_process_group"] is None
+    assert seen["parallel_tiling"] is False
+    assert frames.ndim == 5 and frames.dtype is torch.float32
+    # the group state is restored once decode returns
+    assert state["sp_size"] == 4 and state["sp_enabled"] is True
+    assert vae.model.parallel_tiling is True
+
+
+def test_decode_latent_leaves_the_group_state_alone_when_the_grid_is_large_enough():
+    """The shipped 768x1344 request is 28 tiles: the guard must not fire."""
+    vae, state, seen = _dispatch_vae(4)
+
+    MiniMaxH3VideoVAE.decode_latent(vae, _latent(48, 84))
+
+    assert seen["sp_size"] == 4
+    assert seen["sp_enabled"] is True
+    assert seen["sp_process_group"] is state["sp_process_group"]
+    assert seen["parallel_tiling"] is True

vllm_omni/diffusion/models/minimax_h3/vae.py

@@ -5,7 +5,8 @@ from __future__ import annotations
 
 import importlib
 import json
-from contextlib import AbstractContextManager
+from collections.abc import Iterator
+from contextlib import AbstractContextManager, contextmanager, nullcontext
 from pathlib import Path
 from typing import Any
 
@@ -14,6 +15,7 @@ import torch.distributed as dist
 import torch.nn as nn
 from PIL import Image
 from transformers.dynamic_module_utils import get_class_from_dynamic_module
+from vllm.logger import init_logger
 
 from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
     DistributedVaeMixin,
@@ -28,6 +30,9 @@ MINIMAX_H3_AUDIO_SAMPLE_RATE = 32000
 MINIMAX_H3_AUDIO_CHANNELS = 2
 
 
+logger = init_logger(__name__)
+
+
 def _load_component_config(component_path: str) -> dict[str, Any]:
     config_path = Path(component_path) / "config.json"
     config = json.loads(config_path.read_text(encoding="utf-8"))
@@ -170,9 +175,7 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
         self.parallel_size = parallel_size
         enabled = parallel_size > 1
 
-        package = self.remote.__class__.__module__.rsplit(".", 1)[0]
-        parallel_module = importlib.import_module(f"{package}.parallel")
-        state = parallel_module.get_parallel_state()
+        state = self._native_parallel_state()
         state.clear()
         state.update(
             group_size=parallel_size,
@@ -187,6 +190,57 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
         )
         self.model.parallel_tiling = enabled
 
+    def _native_parallel_state(self) -> dict[str, Any]:
+        """Return the checkpoint's own mutable parallel-state dict."""
+
+        package = self.remote.__class__.__module__.rsplit(".", 1)[0]
+        parallel_module = importlib.import_module(f"{package}.parallel")
+        return parallel_module.get_parallel_state()
+
+    def _decoder_tile_count(self, latent: torch.Tensor) -> int:
+        """Number of decoder tiles the checkpoint will split ``latent`` into.
+
+        Mirrors the checkpoint's ``decode_tiled``: the grid is computed from the
+        pixel-space dimensions, so it is a pure function of the latent shape and
+        resolves identically on every rank.
+        """
+
+        ratio = int(self.model.vae_ratio)
+        rows, _, _ = self.model.split_tiles(int(latent.shape[-2]) * ratio, True)
+        cols, _, _ = self.model.split_tiles(int(latent.shape[-1]) * ratio, True)
+        return len(rows) * len(cols)
+
+    @contextmanager
+    def _rank_local_tiling(self) -> Iterator[None]:
+        """Run one decode with tiling kept on this rank, then restore the group.
+
+        Used only when there are fewer tiles than ranks. Every rank then decodes
+        every tile, which is slower than sharing the work but is correct and
+        involves no collective.
+        """
+
+        state = self._native_parallel_state()
+        saved_state = dict(state)
+        saved_tiling = self.model.parallel_tiling
+        state.update(
+            group_size=1,
+            group_rank=0,
+            local_process_group=None,
+            sp_size=1,
+            sp_rank=0,
+            sp_enabled=False,
+            sp_process_group=None,
+            tp_size=1,
+            tp_rank=0,
+        )
+        self.model.parallel_tiling = False
+        try:
+            yield
+        finally:
+            state.clear()
+            state.update(saved_state)
+            self.model.parallel_tiling = saved_tiling
+
     def is_distributed_enabled(self) -> bool:
         return self.parallel_size > 1 and dist.is_initialized()
 
@@ -293,7 +347,27 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
             device=latent.device,
             dtype=latent.dtype,
         ).view(1, channels, 1, 1, 1)
-        decoded = self.model.decode_base(latent * std + mean)
+        # The checkpoint hands rank r the tiles ``range(r, num_tiles, sp_size)``
+        # and then rejects an empty share inside the gather. A rank with no
+        # tiles raises and leaves the collective while the others block in it
+        # forever, so too few tiles hangs the whole stage rather than failing
+        # it. Tile count depends only on the latent shape, so every rank takes
+        # this branch together.
+        num_tiles = self._decoder_tile_count(latent)
+        if self.parallel_size > 1 and num_tiles < self.parallel_size:
+            logger.warning_once(
+                "MiniMax-H3 VAE decode splits into %d tile(s) but the tile group has "
+                "%d ranks; decoding rank-locally for this shape instead, which is "
+                "slower but avoids ranks without tiles hanging the collective.",
+                num_tiles,
+                self.parallel_size,
+            )
+            tiling_context: AbstractContextManager = self._rank_local_tiling()
+        else:
+            tiling_context = nullcontext()
+
+        with tiling_context:
+            decoded = self.model.decode_base(latent * std + mean)
         frames = self.model.processor.revert_tensor(decoded)
         if frames.ndim == 4:
             frames = frames.unsqueeze(0).transpose(1, 2)