From 520e058b953e8a8e0d0bbad0786bc29ab451aef8 Mon Sep 17 00:00:00 2001 From: hsliuustc0106 Date: Tue, 1 Sep 2026 07:59:02 +0800 Subject: [PATCH] [Model][MiniMax-H3] Delegate temporal VAE chunks Signed-off-by: hsliuustc0106 --- .../minimax_h3/test_minimax_h3_vae_chunks.py | 380 ++++++++++++++++++ vllm_omni/diffusion/models/minimax_h3/vae.py | 299 ++++++++++++-- 2 files changed, 640 insertions(+), 39 deletions(-) create mode 100644 tests/diffusion/models/minimax_h3/test_minimax_h3_vae_chunks.py diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_chunks.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_chunks.py new file mode 100644 index 000000000..ba2142261 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_chunks.py @@ -0,0 +1,380 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import multiprocessing as mp +import time +from pathlib import Path +from types import MethodType +from typing import Any + +import pytest +import torch +import torch.distributed as dist +import torch.nn as nn + +from vllm_omni.diffusion.models.minimax_h3.vae import ( + MiniMaxH3ChunkCallbackPeerError, + MiniMaxH3ChunkedDecodeUnsupportedError, + MiniMaxH3VideoVAE, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +class _IdentityProcessor: + @staticmethod + def revert_tensor(tensor: torch.Tensor) -> torch.Tensor: + return tensor + + +class _FakeVendorTemporalModel(nn.Module): + temporal_chunk_callback_api_version = 1 + vae_ratio = 1 + + def __init__( + self, + *, + collective: bool = False, + metadata_mode: str = "valid", + ) -> None: + super().__init__() + values = torch.linspace(0, 1, 124).view(1, 1, 124, 1, 1) + self.expected = values.expand(1, 3, -1, 1, 1).contiguous() + self.processor = _IdentityProcessor() + self.parallel_tiling = collective + self.collective = collective + self.metadata_mode = metadata_mode + self.decode_base_calls = 0 + self.vendor_callback_calls = 0 + + @staticmethod + def split_tiles(size: int, tiled: bool): + del size, tiled + return [0, 1], None, None + + def decode_base( + self, + latent: torch.Tensor, + *, + temporal_chunk_callback=None, + ) -> torch.Tensor: + del latent + self.decode_base_calls += 1 + output = self.expected.clone() + if temporal_chunk_callback is None: + return output + + frame_start = 0 + for chunk_index, chunk_size in enumerate([17] * 7 + [5]): + if self.collective: + collective_probe = torch.tensor([chunk_index], dtype=torch.int32) + dist.all_reduce(collective_probe) + self.vendor_callback_calls += 1 + metadata_index = chunk_index + metadata_total = 8 + metadata_start = frame_start + is_final = chunk_index == 7 + if self.metadata_mode == "gap" and chunk_index == 1: + metadata_start += 1 + elif self.metadata_mode == "bad_total" and chunk_index == 1: + metadata_total += 1 + elif self.metadata_mode == "no_final": + is_final = False + temporal_chunk_callback( + self.expected[:, :, frame_start : frame_start + chunk_size].clone(), + chunk_index=metadata_index, + total_chunks=metadata_total, + frame_start=metadata_start, + is_final=is_final, + ) + frame_start += chunk_size + return output + + +def _vae( + *, + parallel_size: int = 1, + collective: bool = False, + metadata_mode: str = "valid", +) -> tuple[MiniMaxH3VideoVAE, _FakeVendorTemporalModel]: + vae = object.__new__(MiniMaxH3VideoVAE) + nn.Module.__init__(vae) + model = _FakeVendorTemporalModel( + collective=collective, + metadata_mode=metadata_mode, + ) + vae.model = model + vae.config_dict = { + "latent_channels": 3, + "latents_mean": [0.0, 0.0, 0.0], + "latents_std": [1.0, 1.0, 1.0], + } + vae.parallel_size = parallel_size + return vae, model + + +def _latent() -> torch.Tensor: + return torch.zeros(1, 3, 37, 1, 1) + + +def test_default_decode_path_does_not_request_vendor_chunks() -> None: + vae, model = _vae() + + output = vae.decode_latent(_latent()) + + torch.testing.assert_close(output, model.expected, rtol=0, atol=0) + assert model.decode_base_calls == 1 + assert model.vendor_callback_calls == 0 + + +def test_callback_chunks_are_owned_and_reconstruct_the_full_decode() -> None: + vae, model = _vae() + chunks: list[torch.Tensor] = [] + snapshots: list[torch.Tensor] = [] + metadata: list[tuple[int, int, int, bool]] = [] + + def consume( + frames: torch.Tensor, + *, + chunk_index: int, + total_chunks: int, + frame_start: int, + is_final: bool, + ) -> None: + assert frames.dtype is torch.float32 + assert frames.is_contiguous() + assert torch.all((0 <= frames) & (frames <= 1)) + chunks.append(frames) + snapshots.append(frames.clone()) + metadata.append((chunk_index, total_chunks, frame_start, is_final)) + + output = vae.decode_latent_with_chunks(_latent(), consume) + + torch.testing.assert_close(output, model.expected, rtol=0, atol=0) + torch.testing.assert_close(torch.cat(chunks, dim=2), output, rtol=0, atol=0) + for chunk, snapshot in zip(chunks, snapshots, strict=True): + torch.testing.assert_close(chunk, snapshot, rtol=0, atol=0) + assert len({chunk.untyped_storage().data_ptr() for chunk in chunks}) == len(chunks) + assert metadata == [(index, 8, index * 17, index == 7) for index in range(8)] + + +def test_callback_mutation_does_not_change_returned_full_output() -> None: + vae, model = _vae() + + def mutate(frames: torch.Tensor, **metadata: Any) -> None: + del metadata + frames.fill_(1234) + + output = vae.decode_latent_with_chunks(_latent(), mutate) + + torch.testing.assert_close(output, model.expected, rtol=0, atol=0) + + +def test_callback_failure_is_deferred_until_vendor_decode_finishes() -> None: + vae, model = _vae() + callback_calls = 0 + + def fail(*args, **kwargs) -> None: + nonlocal callback_calls + del args, kwargs + callback_calls += 1 + raise LookupError("sink failed") + + with pytest.raises(LookupError, match="sink failed"): + vae.decode_latent_with_chunks(_latent(), fail) + + assert callback_calls == 1 + assert model.vendor_callback_calls == 8 + torch.testing.assert_close(vae.decode_latent(_latent()), model.expected, rtol=0, atol=0) + + +def test_terminal_callback_failure_preserves_the_original_exception() -> None: + vae, model = _vae() + seen: list[tuple[int, int, int, bool]] = [] + + def fail_final( + frames: torch.Tensor, + *, + chunk_index: int, + total_chunks: int, + frame_start: int, + is_final: bool, + ) -> None: + del frames + seen.append((chunk_index, total_chunks, frame_start, is_final)) + if is_final: + raise LookupError("final sink failed") + + with pytest.raises(LookupError, match="final sink failed"): + vae.decode_latent_with_chunks(_latent(), fail_final) + + assert seen == [(index, 8, index * 17, index == 7) for index in range(8)] + assert model.vendor_callback_calls == 8 + + +@pytest.mark.parametrize( + ("metadata_mode", "message"), + [ + ("gap", "discontinuous chunk metadata"), + ("bad_total", "discontinuous chunk metadata"), + ("no_final", "discontinuous chunk metadata"), + ], +) +def test_adapter_rejects_a_broken_vendor_callback_contract( + metadata_mode: str, + message: str, +) -> None: + vae, model = _vae(metadata_mode=metadata_mode) + + with pytest.raises(RuntimeError, match=message): + vae.decode_latent_with_chunks(_latent(), lambda *args, **kwargs: None) + + assert model.vendor_callback_calls == 8 + + +def test_callback_request_fails_closed_without_the_vendor_api() -> None: + vae, model = _vae() + model.temporal_chunk_callback_api_version = -1 + + with pytest.raises(MiniMaxH3ChunkedDecodeUnsupportedError, match="official MiniMax"): + vae.decode_latent_with_chunks(_latent(), lambda *args, **kwargs: None) + + torch.testing.assert_close(vae.decode_latent(_latent()), model.expected, rtol=0, atol=0) + + +def _distributed_worker( + rank: int, + init_file: str, + result_queue: Any, +) -> None: + dist.init_process_group( + "gloo", + init_method=f"file://{init_file}", + rank=rank, + world_size=2, + ) + try: + vae, model = _vae(parallel_size=2, collective=True) + vae._native_parallel_state = MethodType( + lambda self: {"sp_process_group": dist.group.WORLD}, + vae, + ) + + try: + vae.decode_latent_with_chunks( + _latent(), + lambda *args, **kwargs: None, + ) + except BaseException as exc: # noqa: BLE001 + invalid_owner_type = type(exc).__name__ + else: + invalid_owner_type = "none" + + def fail(*args, **kwargs) -> None: + del args, kwargs + raise LookupError("distributed sink failed") + + try: + vae.decode_latent_with_chunks( + _latent(), + fail if rank == 0 else None, + ) + except BaseException as exc: # noqa: BLE001 + failure_type = type(exc).__name__ + else: + failure_type = "none" + + success_calls = 0 + + def succeed(*args, **kwargs) -> None: + nonlocal success_calls + del args, kwargs + success_calls += 1 + + output = vae.decode_latent_with_chunks( + _latent(), + succeed if rank == 0 else None, + ) + + local_vae, local_model = _vae(parallel_size=1) + local_calls = 0 + + def consume_local(*args, **kwargs) -> None: + nonlocal local_calls + del args, kwargs + local_calls += 1 + + local_output = local_vae.decode_latent_with_chunks( + _latent(), + consume_local if rank == 0 else None, + ) + + recovery = torch.tensor([rank + 1], dtype=torch.int32) + dist.all_reduce(recovery) + result_queue.put( + { + "rank": rank, + "invalid_owner_type": invalid_owner_type, + "failure_type": failure_type, + "vendor_calls": model.vendor_callback_calls, + "success_calls": success_calls, + "output_matches": torch.equal(output, model.expected), + "local_calls": local_calls, + "local_output_matches": torch.equal(local_output, local_model.expected), + "recovery": int(recovery.item()), + } + ) + finally: + dist.destroy_process_group() + + +@pytest.mark.parallel +def test_two_rank_callback_failure_is_symmetric_and_group_recovers( + tmp_path: Path, +) -> None: + context = mp.get_context("spawn") + result_queue = context.Queue() + init_file = str(tmp_path / "gloo-init") + processes = [ + context.Process( + target=_distributed_worker, + args=(rank, init_file, result_queue), + ) + for rank in range(2) + ] + for process in processes: + process.start() + deadline = time.monotonic() + 90 + for process in processes: + process.join(timeout=max(0, deadline - time.monotonic())) + for process in processes: + if process.is_alive(): + process.join(timeout=1) + hung = [process for process in processes if process.is_alive()] + for process in hung: + process.terminate() + process.join(timeout=5) + assert not hung, "two-rank chunk callback test deadlocked" + assert [process.exitcode for process in processes] == [0, 0] + + results = sorted( + [result_queue.get(timeout=5) for _ in range(2)], + key=lambda result: result["rank"], + ) + assert [result["failure_type"] for result in results] == [ + "LookupError", + MiniMaxH3ChunkCallbackPeerError.__name__, + ] + assert [result["invalid_owner_type"] for result in results] == [ + "ValueError", + "ValueError", + ] + assert [result["vendor_calls"] for result in results] == [16, 16] + assert [result["success_calls"] for result in results] == [8, 0] + assert all(result["output_matches"] for result in results) + # With VAE parallelism disabled, the caller still designates one request owner. + assert [result["local_calls"] for result in results] == [8, 0] + assert all(result["local_output_matches"] for result in results) + assert [result["recovery"] for result in results] == [3, 3] diff --git a/vllm_omni/diffusion/models/minimax_h3/vae.py b/vllm_omni/diffusion/models/minimax_h3/vae.py index 2e830fa04..7f658d8ae 100644 --- a/vllm_omni/diffusion/models/minimax_h3/vae.py +++ b/vllm_omni/diffusion/models/minimax_h3/vae.py @@ -9,7 +9,7 @@ import json from collections.abc import Iterator from contextlib import AbstractContextManager, contextmanager, nullcontext from pathlib import Path -from typing import Any +from typing import Any, Protocol import torch import torch.distributed as dist @@ -38,6 +38,35 @@ MINIMAX_H3_AUDIO_CHANNELS = 2 logger = init_logger(__name__) +class MiniMaxH3VideoChunkCallback(Protocol): + """Consume one ordered, temporally committed H3 video chunk. + + ``frames`` is an owned, contiguous float32 RGB tensor in ``BCTHW`` layout + and the ``[0, 1]`` value range. The callback runs synchronously on the + producer's current device stream. Before reading it on another stream, the + consumer must establish ordering from the producer stream (for example via + an event) and retain the allocation for that stream's lifetime. + """ + + def __call__( + self, + frames: torch.Tensor, + *, + chunk_index: int, + total_chunks: int, + frame_start: int, + is_final: bool, + ) -> None: ... + + +class MiniMaxH3ChunkedDecodeUnsupportedError(RuntimeError): + """Raised when the loaded MiniMax H3 VAE lacks the vendor chunk API.""" + + +class MiniMaxH3ChunkCallbackPeerError(RuntimeError): + """Raised on peer VAE ranks when the output-rank callback fails.""" + + 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")) @@ -265,6 +294,116 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): def is_distributed_enabled(self) -> bool: return self.parallel_size > 1 and dist.is_initialized() + def _denormalize_latent(self, latent: torch.Tensor) -> torch.Tensor: + channels = int(self.config_dict["latent_channels"]) + mean = torch.tensor( + self.config_dict["latents_mean"], + device=latent.device, + dtype=latent.dtype, + ).view(1, channels, 1, 1, 1) + std = torch.tensor( + self.config_dict["latents_std"], + device=latent.device, + dtype=latent.dtype, + ).view(1, channels, 1, 1, 1) + return latent * std + mean + + def _decode_tiling_context(self, latent: torch.Tensor) -> AbstractContextManager: + # 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, + ) + return self._rank_local_tiling() + return nullcontext() + + def _normalize_decoded_frames(self, decoded: torch.Tensor) -> torch.Tensor: + frames = self.model.processor.revert_tensor(decoded) + if frames.ndim == 4: + frames = frames.unsqueeze(0).transpose(1, 2) + if frames.ndim != 5: + raise ValueError(f"unexpected decoded video shape {tuple(frames.shape)}") + return frames.float() + + def _chunk_process_group(self) -> dist.ProcessGroup | None: + if not self.is_distributed_enabled(): + return None + group = self._native_parallel_state().get("sp_process_group") + if group is None or dist.get_world_size(group) != self.parallel_size: + raise RuntimeError("MiniMax-H3 VAE chunk decode has an invalid spatial-parallel group") + return group + + @staticmethod + def _is_chunk_output_rank( + chunk_callback: MiniMaxH3VideoChunkCallback | None, + *, + group: dist.ProcessGroup | None, + device: torch.device, + ) -> bool: + if group is None: + return chunk_callback is not None + + group_rank = dist.get_rank(group) + callback_owners = torch.tensor( + [int(chunk_callback is not None)], + dtype=torch.int32, + device=device, + ) + dist.all_reduce(callback_owners, group=group) + valid_owner = int(callback_owners.item()) == 1 and ((group_rank == 0) == (chunk_callback is not None)) + if not valid_owner: + raise ValueError("MiniMax-H3 chunk decode requires a callback on exactly VAE group rank 0") + return group_rank == 0 + + @staticmethod + def _synchronize_chunk_callback_error( + callback_error: BaseException | None, + *, + group: dist.ProcessGroup | None, + device: torch.device, + ) -> None: + if group is None: + if callback_error is not None: + raise callback_error + return + + is_output_rank = dist.get_rank(group) == 0 + failed = torch.tensor( + [int(is_output_rank and callback_error is not None)], + dtype=torch.int32, + device=device, + ) + dist.broadcast( + failed, + src=dist.get_global_rank(group, 0), + group=group, + ) + if int(failed.item()) == 0: + return + if is_output_rank: + assert callback_error is not None + raise callback_error + raise MiniMaxH3ChunkCallbackPeerError("MiniMax-H3 video chunk callback failed on this VAE group's output rank") + + def _validate_vendor_chunk_api(self) -> None: + version = getattr(self.model, "temporal_chunk_callback_api_version", None) + if version != 1: + raise MiniMaxH3ChunkedDecodeUnsupportedError( + "Loaded MiniMax-H3 VAE does not provide temporal_chunk_callback API v1 " + f"(found {version!r}); install a model snapshot containing the official " + "MiniMax temporal chunk callback" + ) + @torch.inference_mode() def encode_image(self, image: Image.Image) -> torch.Tensor: previous_parallel = self.model.parallel_tiling @@ -357,44 +496,120 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): @torch.inference_mode() def decode_latent(self, latent: torch.Tensor) -> torch.Tensor: - channels = int(self.config_dict["latent_channels"]) - mean = torch.tensor( - self.config_dict["latents_mean"], - device=latent.device, - dtype=latent.dtype, - ).view(1, channels, 1, 1, 1) - std = torch.tensor( - self.config_dict["latents_std"], - device=latent.device, - dtype=latent.dtype, - ).view(1, channels, 1, 1, 1) - # 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 self._decode_tiling_context(latent): + decoded = self.model.decode_base(self._denormalize_latent(latent)) + return self._normalize_decoded_frames(decoded) - 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) - if frames.ndim != 5: - raise ValueError(f"unexpected decoded video shape {tuple(frames.shape)}") - return frames.float() + @torch.inference_mode() + def decode_latent_with_chunks( + self, + latent: torch.Tensor, + chunk_callback: MiniMaxH3VideoChunkCallback | None, + ) -> torch.Tensor: + """Decode a video while publishing vendor-committed temporal chunks. + + Every rank collaborating in the configured H3 VAE tile group must call + this method, with ``chunk_callback`` supplied only on group rank 0. A + callback failure is deferred until the vendor decoder finishes its + remaining collectives, then propagated to every group rank. With + ``parallel_size=1`` the runtime designates each request's output owner + by supplying a callback there and ``None`` on replica peers. + + This is a model-local producer seam before the typed media boundary in + RFC #6541 / PR #6615. Representation conversion, transport, + presentation policy, and encoding remain runtime-owned. Chunks retain + the VAE canvas dimensions; a runtime consumer must apply the same + requested-size crop as the complete H3 pipeline output. + """ + + self._validate_vendor_chunk_api() + group = self._chunk_process_group() + is_output_rank = self._is_chunk_output_rank( + chunk_callback, + group=group, + device=latent.device, + ) + callback_error: BaseException | None = None + next_chunk_index = 0 + expected_total_chunks: int | None = None + next_frame_start = 0 + saw_final = False + + def publish( + decoded_chunk: torch.Tensor, + *, + chunk_index: int, + total_chunks: int, + frame_start: int, + is_final: bool, + ) -> None: + nonlocal callback_error, expected_total_chunks + nonlocal next_chunk_index, next_frame_start, saw_final + if not is_output_rank or callback_error is not None: + return + try: + if expected_total_chunks is None: + expected_total_chunks = total_chunks + if ( + saw_final + or total_chunks <= 0 + or total_chunks != expected_total_chunks + or chunk_index != next_chunk_index + or chunk_index >= total_chunks + or frame_start != next_frame_start + or is_final != (chunk_index + 1 == total_chunks) + ): + raise RuntimeError( + "MiniMax-H3 vendor callback emitted discontinuous chunk metadata: " + f"index={chunk_index}/{next_chunk_index}, " + f"total={total_chunks}/{expected_total_chunks}, " + f"frame_start={frame_start}/{next_frame_start}, " + f"final={is_final}, after_final={saw_final}" + ) + frames = self._normalize_decoded_frames(decoded_chunk) + frames = frames.clone(memory_format=torch.contiguous_format) + frame_count = int(frames.shape[2]) + if frame_count <= 0: + raise RuntimeError("MiniMax-H3 vendor callback emitted an empty chunk") + next_chunk_index += 1 + next_frame_start += frame_count + saw_final = is_final + assert chunk_callback is not None + chunk_callback( + frames, + chunk_index=chunk_index, + total_chunks=total_chunks, + frame_start=frame_start, + is_final=is_final, + ) + except BaseException as exc: # noqa: BLE001 + # Keep a failed GPU chunk and its callback locals out of the + # traceback while peer ranks finish the remaining collectives. + exc.__traceback__ = None + callback_error = exc + + with self._decode_tiling_context(latent): + decoded = self.model.decode_base( + self._denormalize_latent(latent), + temporal_chunk_callback=publish, + ) + frames = self._normalize_decoded_frames(decoded) + if ( + is_output_rank + and callback_error is None + and (not saw_final or next_chunk_index != expected_total_chunks or next_frame_start != int(frames.shape[2])) + ): + callback_error = RuntimeError( + "MiniMax-H3 vendor callback did not reconstruct the full decode: " + f"final={saw_final}, chunks={next_chunk_index}/{expected_total_chunks}, " + f"frames={next_frame_start}/{frames.shape[2]}" + ) + self._synchronize_chunk_callback_error( + callback_error, + group=group, + device=frames.device, + ) + return frames class MiniMaxH3AudioVAE(nn.Module): @@ -516,4 +731,10 @@ class MiniMaxH3AudioVAE(nn.Module): return waveform.permute(1, 0, 2).contiguous().float() -__all__ = ["MiniMaxH3AudioVAE", "MiniMaxH3VideoVAE"] +__all__ = [ + "MiniMaxH3AudioVAE", + "MiniMaxH3ChunkCallbackPeerError", + "MiniMaxH3ChunkedDecodeUnsupportedError", + "MiniMaxH3VideoChunkCallback", + "MiniMaxH3VideoVAE", +]