diff --git a/recipes/MiniMaxAI/MiniMax-H3.md b/recipes/MiniMaxAI/MiniMax-H3.md index 016c39991..293dcea70 100644 --- a/recipes/MiniMaxAI/MiniMax-H3.md +++ b/recipes/MiniMaxAI/MiniMax-H3.md @@ -115,6 +115,36 @@ steady-state latency. H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1. The H3 VAE supports its native `tile` mode, not `spatial_shard_height` or `spatial_shard_width`. +### Optional exact-schedule AdaLN cache + +H3 can precompute every AdaLN projection for the request's known timestep +schedule and reuse the exact rows in each denoise step: + +```bash +export VLLM_OMNI_H3_ADALN_SCHEDULE_CACHE=1 +``` + +The cache key uses the exact float32 timestep bit patterns. A request with the +same schedule reuses the installed tables; a changed schedule rebuilds them. +This is disabled by default because the tables consume persistent device +memory proportional to the number of denoise steps and layers. + +After the tables are installed, the original projection weights can also move +to CPU and release their device storage: + +```bash +export VLLM_OMNI_H3_ADALN_SCHEDULE_CACHE=1 +export VLLM_OMNI_H3_ADALN_OFFLOAD_WEIGHTS=1 +``` + +Weight offload requires the schedule cache. On a schedule change, the weights +are restored, the replacement tables are built, and the weights are offloaded +again. The optimization falls back to direct projection when the tables would +occupy at least as much memory as the projection weights. DTensor/HSDP +parameters remain under their distributed owner. With layerwise offload, H3's +AdaLN projection parameters stay outside the per-block transfer set while the +exact-schedule cache owns them. + ### Text encoder tensor parallelism The Qwen3-VL text encoder (~51.5 GB in BF16 for the retained 50 layers) is by diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_adaln_cache.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_adaln_cache.py new file mode 100644 index 000000000..1f4cc1bcd --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_adaln_cache.py @@ -0,0 +1,353 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def _make_branch(layout: str): + from vllm_omni.diffusion.models.minimax_h3.denoise_loop import ( + MiniMaxH3DenoiseBranch, + ) + + update_mask = { + "t2va": [True, True], + "fl2va": [False, True], + "ref2va": [True, True], + }[layout] + audio_update_mask = { + "t2va": [True, True], + "fl2va": [True, True], + "ref2va": [False, True], + }[layout] + packed = { + "seq_len": 8, + "text_pos": torch.tensor([0, 1]), + "img_pos": torch.tensor([2, 3]), + "audio_pos": torch.tensor([4, 5]), + "update_mask": torch.tensor(update_mask), + "audio_update_mask": torch.tensor(audio_update_mask), + "img_position_ids": torch.zeros(8, 3, dtype=torch.long), + "cu_seqlens": torch.tensor([0, 6, 8], dtype=torch.int32), + } + return MiniMaxH3DenoiseBranch( + packed=packed, + text_embeddings=torch.zeros(2, 4), + token_tags=torch.zeros(8, dtype=torch.long), + device=torch.device("cpu"), + ) + + +def _reference_full_vector_unique( + branch, + *, + sigma_video: float, + sigma_audio: float, + imgvid_cond_timestep: float, + audio_ref_cond_timestep: float, +) -> tuple[torch.Tensor, torch.Tensor]: + t_video = 1.0 - sigma_video + t_audio = 1.0 - sigma_audio + timesteps = torch.full((branch.seq_len,), t_video, dtype=torch.float32) + timesteps[branch.img_pos[branch.update_mask]] = t_video + timesteps[branch.img_pos[~branch.update_mask]] = max(t_video, imgvid_cond_timestep) + timesteps[branch.audio_pos[branch.audio_update_mask]] = t_audio + timesteps[branch.audio_pos[~branch.audio_update_mask]] = max(t_audio, audio_ref_cond_timestep) + return torch.unique(timesteps, sorted=True, return_inverse=True) + + +@pytest.mark.parametrize("layout", ["t2va", "fl2va", "ref2va"]) +def test_timestep_metadata_matches_full_vector_unique(layout: str): + branch = _make_branch(layout) + sigmas_video = [1.0, 0.5, 0.0] + sigmas_audio = [1.0, 0.25, 0.0] + + cached = branch.prepare_timestep_metadata( + sigmas_video=sigmas_video, + sigmas_audio=sigmas_audio, + imgvid_cond_noise_aug_for_inference=0.999, + audio_cond_noise_aug_for_inference=1.0, + ) + + assert len(cached) == 2 + for step, metadata in enumerate(cached): + expected_unique, expected_inverse = _reference_full_vector_unique( + branch, + sigma_video=sigmas_video[step], + sigma_audio=sigmas_audio[step], + imgvid_cond_timestep=0.999, + audio_ref_cond_timestep=1.0, + ) + torch.testing.assert_close(metadata.unique_timesteps.cpu(), expected_unique, rtol=0, atol=0) + torch.testing.assert_close(metadata.inverse_indices.cpu(), expected_inverse, rtol=0, atol=0) + + +def test_adaln_schedule_key_uses_exact_float32_bits(): + from vllm_omni.diffusion.models.minimax_h3.adaln_schedule_cache import ( + minimax_h3_adaln_schedule_key, + ) + + branch = _make_branch("t2va") + first = branch.prepare_timestep_metadata( + sigmas_video=[1.0, 0.5, 0.0], + sigmas_audio=[1.0, 0.25, 0.0], + imgvid_cond_noise_aug_for_inference=0.999, + audio_cond_noise_aug_for_inference=1.0, + ) + same_after_float32_rounding = branch.prepare_timestep_metadata( + sigmas_video=[1.0, 0.5 + 1e-12, 0.0], + sigmas_audio=[1.0, 0.25 + 1e-12, 0.0], + imgvid_cond_noise_aug_for_inference=0.999, + audio_cond_noise_aug_for_inference=1.0, + ) + changed = branch.prepare_timestep_metadata( + sigmas_video=[1.0, 0.5, 0.0], + sigmas_audio=[1.0, 0.25000003, 0.0], + imgvid_cond_noise_aug_for_inference=0.999, + audio_cond_noise_aug_for_inference=1.0, + ) + + first_key = minimax_h3_adaln_schedule_key(first) + assert first_key == minimax_h3_adaln_schedule_key(same_after_float32_rounding) + assert first_key != minimax_h3_adaln_schedule_key(changed) + + +class _FakeAdalnLinear(torch.nn.Module): + def __init__(self, out_features: int): + super().__init__() + self.weight = torch.nn.Parameter(torch.arange(out_features * 4, dtype=torch.bfloat16).reshape(out_features, 4)) + self.bias = torch.nn.Parameter(torch.arange(out_features, dtype=torch.bfloat16)) + self.out_features = out_features + self.calls = 0 + + def forward(self, x): + assert self.weight.numel() > 0 + assert self.bias.numel() > 0 + self.calls += 1 + rows = torch.arange( + x.shape[0] * self.out_features, + device=x.device, + dtype=torch.float32, + ).reshape(x.shape[0], self.out_features) + rows = rows + x.float().sum(dim=-1, keepdim=True) + return rows.to(torch.bfloat16), None + + +def _make_tiny_adaln_projection(*, expand_ratio: int = 2, modality_num: int = 3): + from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( + MiniMaxH3AdalnProj, + ) + + projection = object.__new__(MiniMaxH3AdalnProj) + torch.nn.Module.__init__(projection) + projection.expand_ratio = expand_ratio + projection.modality_num = modality_num + projection.hidden_size = 2 + projection.linear = _FakeAdalnLinear(out_features=expand_ratio * modality_num * projection.hidden_size) + projection.register_buffer("_precomputed_table", None, persistent=False) + projection._precomputed_step = None + projection._offloaded_projection_state = None + return projection + + +def test_adaln_projection_reuses_exact_precomputed_rows_without_linear_calls(): + projection = _make_tiny_adaln_projection() + first_t_emb = torch.tensor([[0.0, 1.0, 2.0, 3.0]]) + second_t_emb = torch.tensor([[4.0, 5.0, 6.0, 7.0], [8.0, 9.0, 10.0, 11.0]]) + first_direct = projection.compute_flat(first_t_emb) + second_direct = projection.compute_flat(second_t_emb) + first_padded = torch.cat( + [first_direct, first_direct.new_zeros(second_direct.shape[0] - first_direct.shape[0], first_direct.shape[1])] + ) + table = torch.stack([first_padded, second_direct]) + cursor = torch.zeros((), dtype=torch.long) + direct_calls = projection.linear.calls + + projection.install_precomputed(table, cursor) + first_cached = projection(first_t_emb) + cursor.fill_(1) + second_cached = projection(second_t_emb) + + assert projection.linear.calls == direct_calls + for actual, expected in zip(first_cached, first_direct.chunk(2, dim=-1)): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + for actual, expected in zip(second_cached, second_direct.chunk(2, dim=-1)): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_adaln_projection_offload_preserves_cache_and_clear_restores_weights(): + projection = _make_tiny_adaln_projection() + t_emb = torch.tensor([[0.0, 1.0, 2.0, 3.0]]) + direct = projection.compute_flat(t_emb) + expected_state = {name: tensor.clone() for name, tensor in projection.linear.state_dict().items()} + expected_bytes = sum(tensor.numel() * tensor.element_size() for tensor in expected_state.values()) + projection.install_precomputed(direct.unsqueeze(0), torch.zeros((), dtype=torch.long)) + + freed_bytes = projection.offload_projection() + + assert freed_bytes == expected_bytes + assert all(parameter.numel() == 0 for parameter in projection.linear.parameters()) + cached = projection(t_emb) + for actual, expected in zip(cached, direct.chunk(projection.expand_ratio, dim=-1)): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + projection.clear_precomputed() + assert not projection.projection_weights_offloaded + for name, tensor in projection.linear.state_dict().items(): + torch.testing.assert_close(tensor, expected_state[name], rtol=0, atol=0) + torch.testing.assert_close(projection.compute_flat(t_emb), direct, rtol=0, atol=0) + + +class _CountingTimeEmbedder(torch.nn.Module): + def __init__(self): + super().__init__() + self.calls = 0 + + def forward(self, timesteps): + self.calls += 1 + return timesteps[:, None] + + +class _TinyAdalnBlock(torch.nn.Module): + def __init__(self, *, modality_num: int = 3): + super().__init__() + self.adaln_proj = _make_tiny_adaln_projection(modality_num=modality_num) + + +def _make_adaln_schedule_model(): + from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( + MiniMaxH3DiTModel, + ) + + model = object.__new__(MiniMaxH3DiTModel) + torch.nn.Module.__init__(model) + model.time_embedder = _CountingTimeEmbedder() + model.blocks = torch.nn.ModuleList([_TinyAdalnBlock(), _TinyAdalnBlock()]) + model.final_layer = _TinyAdalnBlock(modality_num=1) + model._adaln_projections = tuple(block.adaln_proj for block in model.blocks) + (model.final_layer.adaln_proj,) + model._adaln_schedule_cache = None + return model + + +def _schedule(branch, *, middle_video: float, middle_audio: float): + return branch.prepare_timestep_metadata( + sigmas_video=[1.0, middle_video, 0.0], + sigmas_audio=[1.0, middle_audio, 0.0], + imgvid_cond_noise_aug_for_inference=0.999, + audio_cond_noise_aug_for_inference=1.0, + ) + + +def test_model_cache_hit_and_changed_schedule_rebuild_with_weight_offload(): + from vllm_omni.diffusion.models.minimax_h3.adaln_schedule_cache import ( + minimax_h3_adaln_schedule_key, + ) + + branch = _make_branch("t2va") + first_plan = _schedule(branch, middle_video=0.5, middle_audio=0.25) + changed_plan = _schedule(branch, middle_video=0.4, middle_audio=0.2) + model = _make_adaln_schedule_model() + expected_projection_bytes = sum( + parameter.numel() * parameter.element_size() + for projection in model._adaln_projections + for parameter in projection.linear.parameters() + ) + + first = model.prepare_adaln_schedule_cache( + unique_timestep_plan=tuple(step.unique_timesteps for step in first_plan), + schedule_key=minimax_h3_adaln_schedule_key(first_plan), + offload_weights=True, + ) + calls_after_first = [projection.linear.calls for projection in model._adaln_projections] + hit = model.prepare_adaln_schedule_cache( + unique_timestep_plan=tuple(step.unique_timesteps for step in first_plan), + schedule_key=minimax_h3_adaln_schedule_key(first_plan), + offload_weights=True, + ) + replaced = model.prepare_adaln_schedule_cache( + unique_timestep_plan=tuple(step.unique_timesteps for step in changed_plan), + schedule_key=minimax_h3_adaln_schedule_key(changed_plan), + offload_weights=True, + ) + + assert first is hit + assert replaced is not None and replaced is not first + assert first is not None + assert first.offloaded_projection_bytes == expected_projection_bytes + assert first.net_memory_saved_bytes == expected_projection_bytes - first.table_bytes + assert calls_after_first == [2, 2, 2] + assert [projection.linear.calls for projection in model._adaln_projections] == [4, 4, 4] + assert all(projection.projection_weights_offloaded for projection in model._adaln_projections) + + +def test_model_weight_offload_falls_back_when_tables_are_larger(): + model = _make_adaln_schedule_model() + unique_timestep_plan = tuple(torch.tensor([float(step)]) for step in range(12)) + + cache = model.prepare_adaln_schedule_cache( + unique_timestep_plan=unique_timestep_plan, + schedule_key=tuple((step,) for step in range(12)), + offload_weights=True, + ) + + assert cache is None + assert all(projection._precomputed_table is None for projection in model._adaln_projections) + assert all(not projection.projection_weights_offloaded for projection in model._adaln_projections) + + +def test_denoise_loop_prepares_cache_once_and_passes_step_indices(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import denoise_loop + + monkeypatch.setenv("VLLM_OMNI_H3_ADALN_SCHEDULE_CACHE", "1") + monkeypatch.setenv("VLLM_OMNI_H3_ADALN_OFFLOAD_WEIGHTS", "1") + + class FakeModel: + def __init__(self): + self.prepare_calls = 0 + self.prepare_for_request_calls = 0 + self.step_indices = [] + + def prepare_for_request_caches(self): + assert torch.is_inference_mode_enabled() + self.prepare_for_request_calls += 1 + + def prepare_adaln_schedule_cache(self, *, unique_timestep_plan, schedule_key, offload_weights): + assert torch.is_inference_mode_enabled() + assert len(unique_timestep_plan) == len(schedule_key) == 2 + assert offload_weights + self.prepare_calls += 1 + return object() + + def __call__(self, **kwargs): + self.step_indices.append(kwargs["adaln_step_index"]) + return torch.zeros(2, 96), torch.zeros(2, 32) + + model = FakeModel() + denoise_loop.minimax_h3_denoise_loop( + model=model, + positive=_make_branch("t2va"), + initial_video_rows=torch.zeros(2, 96), + initial_audio_rows=torch.zeros(2, 32), + keyframe_cond_rows=None, + sigmas_video=[1.0, 0.5, 0.0], + sigmas_audio=[1.0, 0.25, 0.0], + device=torch.device("cpu"), + ) + + assert model.prepare_for_request_calls == 1 + assert model.prepare_calls == 1 + assert model.step_indices == [0, 1] + + +def test_weight_offload_env_requires_schedule_cache(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3.adaln_schedule_cache import ( + minimax_h3_adaln_weight_offload_enabled, + ) + + monkeypatch.delenv("VLLM_OMNI_H3_ADALN_SCHEDULE_CACHE", raising=False) + monkeypatch.setenv("VLLM_OMNI_H3_ADALN_OFFLOAD_WEIGHTS", "1") + assert not minimax_h3_adaln_weight_offload_enabled() + monkeypatch.setenv("VLLM_OMNI_H3_ADALN_SCHEDULE_CACHE", "1") + assert minimax_h3_adaln_weight_offload_enabled() diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py index d6f089d0a..7560960a1 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py @@ -25,7 +25,7 @@ def test_grouped_qkv_checkpoint_reorder(): assert reordered[:, 0].tolist() == [0, 3, 1, 4, 2, 5] -def test_transformer_declares_cache_sp_layerwise_offload_and_hsdp(): +def test_transformer_declares_cache_sp_layerwise_offload_and_hsdp(monkeypatch): from cache_dit import ForwardPattern from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( @@ -38,6 +38,11 @@ def test_transformer_declares_cache_sp_layerwise_offload_and_hsdp(): assert not MiniMaxH3DiTModel._cache_dit_adapter_config.has_separate_cfg assert set(MiniMaxH3DiTModel._sp_plan) == {"sp_prepare", "sp_gather"} + monkeypatch.setenv("VLLM_OMNI_H3_ADALN_SCHEDULE_CACHE", "1") + assert MiniMaxH3DiTModel.layerwise_offload_excluded_parameter_prefixes() == ("adaln_proj.linear.",) + monkeypatch.delenv("VLLM_OMNI_H3_ADALN_SCHEDULE_CACHE") + assert MiniMaxH3DiTModel.layerwise_offload_excluded_parameter_prefixes() == () + model = object.__new__(MiniMaxH3DiTModel) nn.Module.__init__(model) model.blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(2)]) @@ -61,6 +66,14 @@ def test_packed_attention_is_a_regional_compile_boundary(): assert getattr(MiniMaxH3Attention._run_packed_attention, "_torchdynamo_disable", False) +def test_adaln_projection_is_a_regional_compile_boundary(): + from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( + MiniMaxH3AdalnProj, + ) + + assert getattr(MiniMaxH3AdalnProj.forward, "_torchdynamo_disable", False) + + @pytest.mark.parametrize( ("tp_size", "message"), [ diff --git a/tests/diffusion/offloader/test_layerwise_backend.py b/tests/diffusion/offloader/test_layerwise_backend.py index 5fd80e75c..62aab0b04 100644 --- a/tests/diffusion/offloader/test_layerwise_backend.py +++ b/tests/diffusion/offloader/test_layerwise_backend.py @@ -135,6 +135,14 @@ class _DummyBlock(nn.Module): self.weight = nn.Parameter(torch.randn(10, 10)) +class _OwnedProjectionBlock(nn.Module): + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.arange(4, dtype=torch.float32)) + self.adaln_proj = nn.Module() + self.adaln_proj.linear = nn.Linear(2, 3) + + class _SingleBlockModel(nn.Module): _layerwise_offload_blocks_attrs = ["blocks"] @@ -224,6 +232,49 @@ class TestGetBlocksFromDit: assert len(blocks) == 2 +def test_layerwise_hook_leaves_excluded_parameters_under_model_ownership(patched_offload_runtime): + current_block = _OwnedProjectionBlock() + next_block = _OwnedProjectionBlock() + current_adaln = {name: tensor.clone() for name, tensor in current_block.adaln_proj.linear.state_dict().items()} + next_adaln = {name: tensor.clone() for name, tensor in next_block.adaln_proj.linear.state_dict().items()} + hook = LayerwiseOffloadHook( + next_block=next_block, + device=torch.device("cpu"), + stream=DummyStream(), + pin_memory=False, + excluded_parameter_prefixes=("adaln_proj.linear.",), + ) + + hook.initialize_hook(current_block) + + assert set(hook.block_parameters) == {"weight"} + assert set(hook.next_block_parameters) == {"weight"} + assert next_block.weight.numel() == 0 + for name, tensor in next_block.adaln_proj.linear.state_dict().items(): + torch.testing.assert_close(tensor, next_adaln[name], rtol=0, atol=0) + + hook.offload_layer() + assert current_block.weight.numel() == 0 + for name, tensor in current_block.adaln_proj.linear.state_dict().items(): + torch.testing.assert_close(tensor, current_adaln[name], rtol=0, atol=0) + + +def test_layerwise_backend_materializes_excluded_parameters_on_target_device(): + block = _OwnedProjectionBlock() + expected_bytes = sum( + parameter.numel() * parameter.element_size() for parameter in block.adaln_proj.linear.parameters() + ) + + moved_bytes = LayerWiseOffloadBackend._move_excluded_parameters_to_device( + block, + torch.device("cpu"), + ("adaln_proj.linear.",), + ) + + assert moved_bytes == expected_bytes + assert all(parameter.device.type == "cpu" for parameter in block.adaln_proj.linear.parameters()) + + class TestGetBlocksAttrNames: def test_get_blocks_attr_names_new_format(self): model = _MultiBlockModel() diff --git a/vllm_omni/diffusion/models/minimax_h3/adaln_schedule_cache.py b/vllm_omni/diffusion/models/minimax_h3/adaln_schedule_cache.py new file mode 100644 index 000000000..7806298e5 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/adaln_schedule_cache.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Exact schedule identity and policy for MiniMax-H3 AdaLN reuse.""" + +from __future__ import annotations + +import os +from collections.abc import Iterable +from dataclasses import dataclass +from typing import Any + +import torch + +MINIMAX_H3_ADALN_SCHEDULE_CACHE_ENV = "VLLM_OMNI_H3_ADALN_SCHEDULE_CACHE" +MINIMAX_H3_ADALN_OFFLOAD_WEIGHTS_ENV = "VLLM_OMNI_H3_ADALN_OFFLOAD_WEIGHTS" +_TRUTHY_ENV_VALUES = frozenset({"1", "true", "yes", "on"}) + +MiniMaxH3AdalnScheduleKey = tuple[tuple[int, ...], ...] + + +@dataclass(frozen=True) +class MiniMaxH3AdalnScheduleCache: + """One model-local exact modulation schedule.""" + + key: MiniMaxH3AdalnScheduleKey + t_embs: tuple[torch.Tensor, ...] + step_cursor: torch.Tensor + steps: int + max_unique_timesteps: int + table_bytes: int + projection_bytes: int = 0 + offloaded_projection_bytes: int = 0 + net_memory_saved_bytes: int = 0 + + +def minimax_h3_adaln_schedule_cache_enabled() -> bool: + """Return whether persistent exact-schedule AdaLN reuse is enabled.""" + return os.getenv(MINIMAX_H3_ADALN_SCHEDULE_CACHE_ENV, "").strip().lower() in _TRUTHY_ENV_VALUES + + +def minimax_h3_adaln_weight_offload_enabled() -> bool: + """Return whether exact tables should take ownership of AdaLN weights.""" + return minimax_h3_adaln_schedule_cache_enabled() and ( + os.getenv(MINIMAX_H3_ADALN_OFFLOAD_WEIGHTS_ENV, "").strip().lower() in _TRUTHY_ENV_VALUES + ) + + +def minimax_h3_float32_bits(values: Any) -> tuple[int, ...]: + """Return the signed int32 bit patterns of a CPU float32 tensor.""" + if not isinstance(values, torch.Tensor) or values.device.type != "cpu" or values.dtype != torch.float32: + raise ValueError("schedule key values must be a CPU float32 tensor") + return tuple(int(value) for value in values.contiguous().view(torch.int32).tolist()) + + +def minimax_h3_adaln_schedule_key( + metadata: Iterable[Any], +) -> MiniMaxH3AdalnScheduleKey: + """Build an immutable schedule key without reading a device tensor.""" + key: list[tuple[int, ...]] = [] + for step in metadata: + bits = getattr(step, "timestep_bits", None) + if bits is None: + raise ValueError("timestep metadata is missing host timestep_bits") + key.append(tuple(int(value) for value in bits)) + if not key: + raise ValueError("AdaLN schedule must contain at least one denoise step") + return tuple(key) + + +def minimax_h3_build_adaln_table( + projection: Any, + t_embs: tuple[torch.Tensor, ...], +) -> torch.Tensor: + """Project each step at its original shape and pad only unused rows.""" + rows_by_step = tuple(projection.compute_flat(t_emb) for t_emb in t_embs) + max_rows = max(int(rows.shape[0]) for rows in rows_by_step) + padded: list[torch.Tensor] = [] + for rows in rows_by_step: + missing = max_rows - int(rows.shape[0]) + if missing: + rows = torch.cat( + [rows, rows.new_zeros((missing, rows.shape[1]))], + dim=0, + ) + padded.append(rows) + return torch.stack(padded, dim=0) + + +__all__ = [ + "MINIMAX_H3_ADALN_OFFLOAD_WEIGHTS_ENV", + "MINIMAX_H3_ADALN_SCHEDULE_CACHE_ENV", + "MiniMaxH3AdalnScheduleCache", + "MiniMaxH3AdalnScheduleKey", + "minimax_h3_build_adaln_table", + "minimax_h3_adaln_schedule_cache_enabled", + "minimax_h3_adaln_weight_offload_enabled", + "minimax_h3_adaln_schedule_key", + "minimax_h3_float32_bits", +] diff --git a/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py b/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py index b051951a7..c41afaa41 100644 --- a/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py +++ b/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py @@ -10,10 +10,17 @@ from __future__ import annotations from collections.abc import Callable from contextlib import AbstractContextManager, nullcontext +from dataclasses import dataclass from typing import Any import torch +from .adaln_schedule_cache import ( + minimax_h3_adaln_schedule_cache_enabled, + minimax_h3_adaln_schedule_key, + minimax_h3_adaln_weight_offload_enabled, + minimax_h3_float32_bits, +) from .scheduling_minimax_h3_euler_ancestral import ( minimax_h3_euler_eta0_step, minimax_h3_rf_v_to_x0, @@ -28,6 +35,15 @@ MINIMAX_H3_VIDEO_ROW_WIDTH = 96 MINIMAX_H3_AUDIO_ROW_WIDTH = 32 +@dataclass(frozen=True) +class MiniMaxH3StepTimestepMetadata: + """Device-ready timestep routing for one denoise forward.""" + + unique_timesteps: torch.Tensor + inverse_indices: torch.Tensor + timestep_bits: tuple[int, ...] + + class MiniMaxH3DenoiseBranch: """Static per-branch state: packed layout + fixed forward kwargs. @@ -45,6 +61,7 @@ class MiniMaxH3DenoiseBranch: device: torch.device, ) -> None: seq_len = int(packed["seq_len"]) + self.device = device self.seq_len = seq_len self.img_pos = packed["img_pos"].view(-1).to(torch.long) self.audio_pos = packed["audio_pos"].view(-1).to(torch.long) @@ -65,6 +82,12 @@ class MiniMaxH3DenoiseBranch: self.audio_pos_dev = self.audio_pos.to(device) self.update_mask_dev = self.update_mask.to(device) self.audio_update_mask_dev = self.audio_update_mask.to(device) + timestep_category_ids = torch.zeros(seq_len, dtype=torch.long) + timestep_category_ids[self.img_pos[~self.update_mask]] = 1 + timestep_category_ids[self.audio_pos[self.audio_update_mask]] = 2 + timestep_category_ids[self.audio_pos[~self.audio_update_mask]] = 3 + self.timestep_category_ids = timestep_category_ids + self.used_timestep_categories = torch.unique(timestep_category_ids, sorted=True) self.x_base = torch.zeros(1, seq_len, MINIMAX_H3_VIDEO_ROW_WIDTH, dtype=torch.float32, device=device) self.audio_x_base = torch.zeros(1, seq_len, MINIMAX_H3_AUDIO_ROW_WIDTH, dtype=torch.float32, device=device) text_pos_dev = packed["text_pos"].view(-1).to(torch.long).to(device) @@ -87,43 +110,103 @@ class MiniMaxH3DenoiseBranch: "max_seqlen_q": text_len, }, } + self._adaln_schedule_cache_active = False + + def prepare_timestep_metadata( + self, + *, + sigmas_video: list[float], + sigmas_audio: list[float], + imgvid_cond_noise_aug_for_inference: float, + audio_cond_noise_aug_for_inference: float, + ) -> list[MiniMaxH3StepTimestepMetadata]: + """Precompute the full schedule's packed-row timestep routing on CPU.""" + if len(sigmas_video) != len(sigmas_audio): + raise ValueError("video/audio sigma schedules must have equal length") + if len(sigmas_video) < 2: + raise ValueError("sigma schedules need at least 2 entries") + + result: list[MiniMaxH3StepTimestepMetadata] = [] + for step in range(len(sigmas_video) - 1): + t_video = 1.0 - float(sigmas_video[step]) + t_audio = 1.0 - float(sigmas_audio[step]) + category_timesteps = torch.tensor( + [ + t_video, + max(t_video, float(imgvid_cond_noise_aug_for_inference)), + t_audio, + max(t_audio, float(audio_cond_noise_aug_for_inference)), + ], + dtype=torch.float32, + ) + active_timesteps = category_timesteps.index_select(0, self.used_timestep_categories) + unique_timesteps, active_inverse = torch.unique( + active_timesteps, + sorted=True, + return_inverse=True, + ) + category_inverse = torch.zeros(4, dtype=torch.long) + category_inverse.index_copy_(0, self.used_timestep_categories, active_inverse) + inverse_indices = category_inverse.index_select(0, self.timestep_category_ids) + result.append( + MiniMaxH3StepTimestepMetadata( + unique_timesteps=unique_timesteps.to(self.device), + inverse_indices=inverse_indices.to(self.device), + timestep_bits=minimax_h3_float32_bits(unique_timesteps), + ) + ) + return result + + def prepare_adaln_schedule_cache( + self, + model: Any, + timestep_metadata: list[MiniMaxH3StepTimestepMetadata], + *, + enabled: bool, + ) -> Any | None: + """Prepare exact model-local modulation tables for this schedule.""" + self._adaln_schedule_cache_active = False + if not enabled: + clear = getattr(model, "clear_adaln_schedule_cache", None) + if clear is not None: + clear() + return None + prepare = getattr(model, "prepare_adaln_schedule_cache", None) + if prepare is None: + return None + with torch.inference_mode(): + cache = prepare( + unique_timestep_plan=tuple(step.unique_timesteps for step in timestep_metadata), + schedule_key=minimax_h3_adaln_schedule_key(timestep_metadata), + offload_weights=minimax_h3_adaln_weight_offload_enabled(), + ) + self._adaln_schedule_cache_active = cache is not None + return cache def forward_kwargs( self, *, video_rows: torch.Tensor, audio_rows: torch.Tensor, - t_video: float, - t_audio: float, - imgvid_cond_timestep: float, - audio_ref_cond_timestep: float, + timestep_metadata: MiniMaxH3StepTimestepMetadata, + step_index: int | None = None, ) -> dict[str, Any]: x = self.x_base.clone() x[0].index_copy_(0, self.img_pos_dev, video_rows) audio_x = self.audio_x_base.clone() audio_x[0].index_copy_(0, self.audio_pos_dev, audio_rows) - # Packed-sequence timestep semantics: non-media rows (text and - # padding) inherit the current video timestep. Later steps must reuse - # the previous step's updated rows; re-initializing from zeros is only - # valid at step 0. - timesteps = torch.full( - (self.seq_len,), - float(t_video), - dtype=torch.float32, - device=x.device, - ) - timesteps[self.img_pos_dev[self.update_mask_dev]] = t_video - timesteps[self.img_pos_dev[~self.update_mask_dev]] = imgvid_cond_timestep - timesteps[self.audio_pos_dev[self.audio_update_mask_dev]] = t_audio - timesteps[self.audio_pos_dev[~self.audio_update_mask_dev]] = audio_ref_cond_timestep - unique_timesteps, inverse_indices = torch.unique(timesteps, sorted=True, return_inverse=True) - return { + dynamic_kwargs = { **self.static_kwargs, "x": x, "audio_x": audio_x, - "unique_timesteps": unique_timesteps, - "inverse_indices": inverse_indices, + "unique_timesteps": timestep_metadata.unique_timesteps, + "inverse_indices": timestep_metadata.inverse_indices, } + if self._adaln_schedule_cache_active: + if step_index is None: + raise ValueError("step_index is required while the AdaLN schedule cache is active") + dynamic_kwargs["adaln_step_index"] = step_index + return dynamic_kwargs def minimax_h3_denoise_loop( @@ -153,6 +236,12 @@ def minimax_h3_denoise_loop( raise ValueError("video/audio sigma schedules must have equal length") if len(sigmas_video) < 2: raise ValueError("sigma schedules need at least 2 entries") + timestep_metadata = positive.prepare_timestep_metadata( + sigmas_video=sigmas_video, + sigmas_audio=sigmas_audio, + imgvid_cond_noise_aug_for_inference=imgvid_cond_noise_aug_for_inference, + audio_cond_noise_aug_for_inference=audio_cond_noise_aug_for_inference, + ) n_cond = int((~positive.update_mask).sum()) if keyframe_cond_rows is None: if n_cond != 0: @@ -187,6 +276,16 @@ def minimax_h3_denoise_loop( if audio_anchor is not None: audio_rows[~audio_update] = audio_anchor + prepare_for_request_caches = getattr(model, "prepare_for_request_caches", None) + if prepare_for_request_caches is not None: + with torch.inference_mode(): + prepare_for_request_caches() + positive.prepare_adaln_schedule_cache( + model, + timestep_metadata, + enabled=minimax_h3_adaln_schedule_cache_enabled(), + ) + num_steps = len(sigmas_video) - 1 for step in range(num_steps): step_cm = step_profiler(step) if step_profiler is not None else nullcontext() @@ -194,16 +293,12 @@ def minimax_h3_denoise_loop( s_v, s_v_next = sigmas_video[step], sigmas_video[step + 1] s_a, s_a_next = sigmas_audio[step], sigmas_audio[step + 1] t_v, t_a = 1.0 - s_v, 1.0 - s_a - imgvid_cond_t = max(t_v, float(imgvid_cond_noise_aug_for_inference)) - audio_ref_cond_t = max(t_a, float(audio_cond_noise_aug_for_inference)) fk = positive.forward_kwargs( video_rows=video_rows, audio_rows=audio_rows, - t_video=t_v, - t_audio=t_a, - imgvid_cond_timestep=imgvid_cond_t, - audio_ref_cond_timestep=audio_ref_cond_t, + timestep_metadata=timestep_metadata[step], + step_index=step, ) with torch.inference_mode(): v_video, v_audio = model(**fk) @@ -245,5 +340,6 @@ __all__ = [ "MINIMAX_H3_IMGVID_COND_TIMESTEP", "MINIMAX_H3_VIDEO_ROW_WIDTH", "MiniMaxH3DenoiseBranch", + "MiniMaxH3StepTimestepMetadata", "minimax_h3_denoise_loop", ] diff --git a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py index fd5ca67e9..ce5ac5064 100644 --- a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py +++ b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py @@ -16,6 +16,7 @@ from typing import TYPE_CHECKING, Any import torch import torch.nn as nn from cache_dit import ForwardPattern +from torch.distributed.tensor import DTensor from vllm.distributed import get_tensor_model_parallel_world_size from vllm.logger import init_logger from vllm.model_executor.layers.linear import ( @@ -33,6 +34,14 @@ from vllm_omni.diffusion.distributed.sp_plan import ( SequenceParallelInput, SequenceParallelOutput, ) +from vllm_omni.platforms import current_omni_platform + +from .adaln_schedule_cache import ( + MiniMaxH3AdalnScheduleCache, + MiniMaxH3AdalnScheduleKey, + minimax_h3_adaln_schedule_cache_enabled, + minimax_h3_build_adaln_table, +) if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import ( @@ -132,6 +141,7 @@ _FORWARD_SUPPORTED_KWARGS = frozenset( "img_pos_for_infer_output_info", "packed_seq_params", "refiner_packed_seq_params", + "adaln_step_index", } ) @@ -528,14 +538,118 @@ class MiniMaxH3AdalnProj(nn.Module): quant_config=quant_config, prefix=f"{prefix}.linear", ) + self.register_buffer("_precomputed_table", None, persistent=False) + self._precomputed_step: torch.Tensor | None = None + self._offloaded_projection_state: dict[str, tuple[torch.Tensor, torch.device, torch.Size]] | None = None + @property + def projection_weights_offloaded(self) -> bool: + return self._offloaded_projection_state is not None + + @property + def projection_nbytes(self) -> int: + state = self._offloaded_projection_state + if state is not None: + return sum(tensor.numel() * tensor.element_size() for tensor, _, _ in state.values()) + total = 0 + for parameter in self.linear.parameters(): + local = parameter.to_local() if isinstance(parameter, DTensor) else parameter + total += local.numel() * local.element_size() + return total + + def can_offload_projection(self) -> bool: + """Whether this projection can safely own ordinary parameter storage.""" + if self.projection_weights_offloaded: + return True + parameters = tuple(self.linear.parameters()) + return bool(parameters) and all( + not isinstance(parameter, DTensor) and not parameter.is_meta and parameter.numel() > 0 + for parameter in parameters + ) + + def offload_projection(self) -> int: + """Move exact projection parameters to CPU and release device storage.""" + if self.projection_weights_offloaded: + return 0 + if self._precomputed_table is None: + raise RuntimeError("AdaLN projection weights require an installed precomputed table before offload") + if not self.can_offload_projection(): + raise RuntimeError("AdaLN projection weight offload does not support DTensor, meta, or empty parameters") + + state: dict[str, tuple[torch.Tensor, torch.device, torch.Size]] = {} + try: + for name, parameter in self.linear.named_parameters(): + cpu_tensor = parameter.detach().to(device="cpu", copy=True) + state[name] = (cpu_tensor, parameter.device, parameter.shape) + parameter.data = torch.empty((0,), device=parameter.device, dtype=parameter.dtype) + except Exception: + for name, (cpu_tensor, device, shape) in state.items(): + self.linear.get_parameter(name).data = cpu_tensor.to(device=device).view(shape) + raise + self._offloaded_projection_state = state + return sum(tensor.numel() * tensor.element_size() for tensor, _, _ in state.values()) + + def restore_projection(self) -> None: + """Restore a previously offloaded projection to its original devices.""" + state = self._offloaded_projection_state + if state is None: + return + restored = { + name: cpu_tensor.to(device=device).view(shape) for name, (cpu_tensor, device, shape) in state.items() + } + for name, tensor in restored.items(): + self.linear.get_parameter(name).data = tensor + self._offloaded_projection_state = None + + def compute_flat(self, t_emb: torch.Tensor) -> torch.Tensor: + """Run the original projection and return `[M*modalities, ratio*H]`.""" + x = nn.functional.silu(t_emb) + x, _ = self.linear(x.to(self.linear.weight.dtype)) + m = x.shape[0] + return x.view(m * self.modality_num, self.expand_ratio * self.hidden_size) + + def install_precomputed( + self, + table: torch.Tensor, + step_cursor: torch.Tensor, + ) -> None: + """Atomically select a non-persistent exact-schedule modulation table.""" + expected_width = self.expand_ratio * self.hidden_size + if table.dim() != 3 or table.shape[-1] != expected_width: + raise ValueError( + f"precomputed AdaLN table must be [steps, rows, {expected_width}], got {list(table.shape)}" + ) + if table.dtype != self.linear.weight.dtype: + raise ValueError(f"precomputed AdaLN table dtype mismatch: {table.dtype} != {self.linear.weight.dtype}") + if step_cursor.dtype != torch.long or step_cursor.dim() != 0: + raise ValueError("precomputed AdaLN step cursor must be a scalar torch.long tensor") + if table.device != step_cursor.device: + raise ValueError(f"precomputed AdaLN table/cursor device mismatch: {table.device} != {step_cursor.device}") + self._buffers["_precomputed_table"] = table + self._precomputed_step = step_cursor + + def clear_precomputed(self) -> None: + """Restore direct projection without changing checkpoint parameters.""" + self.restore_projection() + self._buffers["_precomputed_table"] = None + self._precomputed_step = None + + @torch.compiler.disable def forward(self, t_emb: torch.Tensor) -> tuple[torch.Tensor, ...]: """t_emb: [M, t_dim] -> expand_ratio tensors of [M*modality_num, H].""" - x = nn.functional.silu(t_emb) - x, _ = self.linear(x.to(_BF16_DTYPE)) - m = x.shape[0] - x = x.view(m * self.modality_num, self.expand_ratio * self.hidden_size) - return tuple(x.chunk(self.expand_ratio, dim=-1)) + table = self._precomputed_table + if table is None: + rows = self.compute_flat(t_emb) + return tuple(rows.chunk(self.expand_ratio, dim=-1)) + + step_cursor = self._precomputed_step + if step_cursor is None: + raise RuntimeError("precomputed AdaLN table is installed without a step cursor") + row_count = t_emb.shape[0] * self.modality_num + if row_count > table.shape[1]: + raise ValueError(f"precomputed AdaLN table row capacity is too small: {table.shape[1]} < {row_count}") + rows = table.index_select(0, step_cursor.reshape(1))[0].narrow(0, 0, row_count) + return tuple(rows.chunk(self.expand_ratio, dim=-1)) class MiniMaxH3TokenRefinerBlock(nn.Module): @@ -781,6 +895,13 @@ class MiniMaxH3DiTModel(nn.Module): _repeated_blocks = ["MiniMaxH3DiTBlock"] _layerwise_offload_blocks_attrs = ["blocks"] + @staticmethod + def layerwise_offload_excluded_parameter_prefixes() -> tuple[str, ...]: + """Parameters owned by the exact-schedule cache instead of block hooks.""" + if minimax_h3_adaln_schedule_cache_enabled(): + return ("adaln_proj.linear.",) + return () + @staticmethod def _is_transformer_block(name: str, module: nn.Module) -> bool: del module @@ -921,6 +1042,8 @@ class MiniMaxH3DiTModel(nn.Module): quant_config, prefix="final_layer", ) + self._adaln_projections = tuple(block.adaln_proj for block in self.blocks) + (self.final_layer.adaln_proj,) + self._adaln_schedule_cache: MiniMaxH3AdalnScheduleCache | None = None self._mark_missing_params_required() def _mark_missing_params_required(self) -> None: @@ -935,6 +1058,126 @@ class MiniMaxH3DiTModel(nn.Module): if name in MINIMAX_H3_FP32_BUFFER_NAMES and buffer.dtype != _FP32_DTYPE: raise ValueError(f"{name} must stay fp32 after load, got {buffer.dtype}.") + @torch.inference_mode() + def prepare_for_request_caches(self) -> None: + """Materialize a sequentially offloaded DiT before auxiliary cache work.""" + registry = getattr(self, "_hook_registry", None) + if registry is None: + return + hook = registry.get_hook("sequential_offload") + if hook is not None: + hook.pre_forward(self) + + @torch.inference_mode() + def prepare_adaln_schedule_cache( + self, + *, + unique_timestep_plan: tuple[torch.Tensor, ...], + schedule_key: MiniMaxH3AdalnScheduleKey, + offload_weights: bool = False, + ) -> MiniMaxH3AdalnScheduleCache | None: + """Build or reuse the single exact AdaLN schedule held by this model.""" + if not unique_timestep_plan: + raise ValueError("AdaLN schedule must contain at least one denoise step") + normalized_key = tuple(tuple(int(value) for value in step) for step in schedule_key) + if len(normalized_key) != len(unique_timestep_plan): + raise ValueError( + f"AdaLN schedule key/plan length mismatch: {len(normalized_key)} != {len(unique_timestep_plan)}" + ) + existing = self._adaln_schedule_cache + if existing is not None and existing.key == normalized_key: + return existing + + for projection in self._adaln_projections: + projection.restore_projection() + + t_embs = tuple(self.time_embedder(timestep.view(-1)) for timestep in unique_timestep_plan) + device = t_embs[0].device + if any(t_emb.device != device for t_emb in t_embs): + raise ValueError("AdaLN timestep embeddings must share one device") + step_cursor = torch.zeros((), dtype=torch.long, device=device) + tables = tuple(minimax_h3_build_adaln_table(projection, t_embs) for projection in self._adaln_projections) + table_bytes = sum(table.numel() * table.element_size() for table in tables) + projection_bytes = sum(projection.projection_nbytes for projection in self._adaln_projections) + offload_supported = offload_weights and all( + projection.can_offload_projection() for projection in self._adaln_projections + ) + if offload_supported and table_bytes >= projection_bytes: + for projection in self._adaln_projections: + projection.clear_precomputed() + self._adaln_schedule_cache = None + logger.warning( + "MiniMax H3 AdaLN weight offload skipped because table bytes (%d) " + "would not save memory versus projection bytes (%d)", + table_bytes, + projection_bytes, + ) + return None + + # Build every table before changing the live projection path. + for projection, table in zip(self._adaln_projections, tables): + projection.install_precomputed(table, step_cursor) + + offloaded_projection_bytes = 0 + if offload_supported: + try: + offloaded_projection_bytes = sum( + projection.offload_projection() for projection in self._adaln_projections + ) + except Exception: + for projection in self._adaln_projections: + projection.restore_projection() + logger.exception("MiniMax H3 AdaLN projection offload failed; retaining exact tables and weights") + elif offload_weights: + logger.warning("MiniMax H3 AdaLN projection offload is unsupported for the active parameter layout") + + if offloaded_projection_bytes: + current_omni_platform.empty_cache() + + cache = MiniMaxH3AdalnScheduleCache( + key=normalized_key, + t_embs=t_embs, + step_cursor=step_cursor, + steps=len(t_embs), + max_unique_timesteps=max(int(t_emb.shape[0]) for t_emb in t_embs), + table_bytes=table_bytes, + projection_bytes=projection_bytes, + offloaded_projection_bytes=offloaded_projection_bytes, + net_memory_saved_bytes=offloaded_projection_bytes - table_bytes, + ) + self._adaln_schedule_cache = cache + if offloaded_projection_bytes: + gib = 1024**3 + logger.info( + "MiniMax H3 AdaLN projection weights offloaded: %.3f GiB weights, " + "%.3f GiB tables, %.3f GiB net device-memory reduction", + offloaded_projection_bytes / gib, + table_bytes / gib, + cache.net_memory_saved_bytes / gib, + ) + return cache + + def clear_adaln_schedule_cache(self) -> None: + """Drop the model-local table and restore direct AdaLN projections.""" + if self._adaln_schedule_cache is None: + return + for projection in self._adaln_projections: + projection.clear_precomputed() + self._adaln_schedule_cache = None + + def _select_adaln_schedule_step(self, step_index: Any) -> torch.Tensor | None: + cache = self._adaln_schedule_cache + if cache is None: + if step_index is not None: + raise ValueError("adaln_step_index was supplied without a prepared AdaLN schedule") + return None + if isinstance(step_index, bool) or not isinstance(step_index, int): + raise ValueError("adaln_step_index must be an integer while the AdaLN schedule is active") + if step_index < 0 or step_index >= cache.steps: + raise ValueError(f"adaln_step_index {step_index} is outside [0, {cache.steps})") + cache.step_cursor.fill_(step_index) + return cache.t_embs[step_index] + def load_weights( self, weights: Iterable[tuple[str, torch.Tensor]], @@ -1007,6 +1250,7 @@ class MiniMaxH3DiTModel(nn.Module): refiner_max_seqlen: int, seq_len: int, device: torch.device, + precomputed_t_emb: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Build packed multimodal embeddings for the TP=1/SP=1 inference path. @@ -1032,7 +1276,7 @@ class MiniMaxH3DiTModel(nn.Module): embeddings.index_add_(0, img_pos, video_embed.to(_BF16_DTYPE)[: img_pos.shape[0]]) embeddings.index_add_(0, audio_pos, audio_embed.to(_BF16_DTYPE)[: audio_pos.shape[0]]) - t_emb = self.time_embedder(unique_timesteps) + t_emb = self.time_embedder(unique_timesteps) if precomputed_t_emb is None else precomputed_t_emb return embeddings, t_emb def forward(self, **kwargs: Any) -> tuple[torch.Tensor, torch.Tensor]: @@ -1052,6 +1296,8 @@ class MiniMaxH3DiTModel(nn.Module): f"{sorted(_FORWARD_SUPPORTED_KWARGS)}" ) + precomputed_t_emb = self._select_adaln_schedule_step(kwargs.get("adaln_step_index")) + x = _required_kwarg(kwargs, "x") audio_x = _required_kwarg(kwargs, "audio_x") img_position_ids = _required_kwarg(kwargs, "img_position_ids") @@ -1104,6 +1350,7 @@ class MiniMaxH3DiTModel(nn.Module): refiner_max_seqlen=refiner_max, seq_len=seq_len, device=device, + precomputed_t_emb=precomputed_t_emb, ) combined_indices = (inverse_indices * MINIMAX_H3_ADALN_MODALITY_NUM + token_tags.clamp(min=0)).to(device) diff --git a/vllm_omni/diffusion/offloader/layerwise_backend.py b/vllm_omni/diffusion/offloader/layerwise_backend.py index 9979d01b1..bc0522086 100644 --- a/vllm_omni/diffusion/offloader/layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/layerwise_backend.py @@ -41,6 +41,7 @@ class LayerwiseOffloadHook(ModelHook): device: torch.device, stream: current_omni_platform.Stream | None = None, pin_memory: bool = True, + excluded_parameter_prefixes: tuple[str, ...] = (), ): assert isinstance(next_block, nn.Module), "transformer block must be type `torch.nn.Module`" @@ -48,6 +49,7 @@ class LayerwiseOffloadHook(ModelHook): self.device = device self.copy_stream = stream or current_omni_platform.current_stream() self.pin_memory = pin_memory + self.excluded_parameter_prefixes = tuple(excluded_parameter_prefixes) # Per-block synchronization primitive: set after H2D copy completes. self._prefetch_done: current_omni_platform.Event | None = None @@ -90,10 +92,10 @@ class LayerwiseOffloadHook(ModelHook): # the input module is kept intact module = super().initialize_hook(module) - self.block_parameters: dict[str, nn.Parameter] = dict(module.named_parameters()) + self.block_parameters = self._owned_parameters(module) self.block_buffers: dict[str, torch.Tensor] = dict(module.named_buffers()) - self.next_block_parameters: dict[str, nn.Parameter] = dict(self.next_block.named_parameters()) + self.next_block_parameters = self._owned_parameters(self.next_block) self.next_block_buffers: dict[str, torch.Tensor] = dict(self.next_block.named_buffers()) # Pre-allocate gpu tensors in a flattened way @@ -106,6 +108,13 @@ class LayerwiseOffloadHook(ModelHook): return module + def _owned_parameters(self, module: nn.Module) -> dict[str, nn.Parameter]: + return { + name: parameter + for name, parameter in module.named_parameters() + if not name.startswith(self.excluded_parameter_prefixes) + } + @staticmethod def _to_cpu( params: dict[str, nn.Parameter], @@ -255,9 +264,10 @@ def apply_block_hook( device: torch.device, stream: current_omni_platform.Stream | None = None, pin_memory: bool = True, + excluded_parameter_prefixes: tuple[str, ...] = (), ) -> LayerwiseOffloadHook: registry = HookRegistry.get_or_create(module) - hook = LayerwiseOffloadHook(next_block, device, stream, pin_memory) + hook = LayerwiseOffloadHook(next_block, device, stream, pin_memory, excluded_parameter_prefixes) registry.register_hook(LayerwiseOffloadHook._HOOK_NAME, hook) return hook @@ -284,6 +294,20 @@ class LayerWiseOffloadBackend(OffloadBackend): self.copy_stream = current_omni_platform.Stream() self._blocks: list[list[nn.Module]] = [] + @staticmethod + def _move_excluded_parameters_to_device( + module: nn.Module, + device: torch.device, + excluded_parameter_prefixes: tuple[str, ...], + ) -> int: + moved_bytes = 0 + for name, parameter in module.named_parameters(): + if not name.startswith(excluded_parameter_prefixes): + continue + moved_bytes += parameter.numel() * parameter.element_size() + parameter.data = parameter.data.to(device, non_blocking=True) + return moved_bytes + def enable(self, pipeline: nn.Module) -> None: if self.enabled: logger.warning("LayerWiseOffloadBackend already enabled") @@ -341,6 +365,30 @@ class LayerWiseOffloadBackend(OffloadBackend): dit_module.to(self.device) continue + excluded_parameter_prefixes: tuple[str, ...] = () + get_excluded_prefixes = getattr(dit_module, "layerwise_offload_excluded_parameter_prefixes", None) + if callable(get_excluded_prefixes): + excluded_parameter_prefixes = tuple(get_excluded_prefixes()) + if excluded_parameter_prefixes: + logger.info( + "Leaving layer-wise parameters matching %s under %s ownership", + excluded_parameter_prefixes, + dit_module.__class__.__name__, + ) + moved_bytes = sum( + self._move_excluded_parameters_to_device( + block, + self.device, + excluded_parameter_prefixes, + ) + for block in blocks + ) + logger.info( + "Materialized %.3f GiB of model-owned layer-wise parameters on %s", + moved_bytes / 1024**3, + self.device, + ) + # Move non-block modules to GPU (they stay resident) for name, m in dit_module.named_children(): if name not in blocks_attr_names: @@ -368,6 +416,7 @@ class LayerWiseOffloadBackend(OffloadBackend): self.device, self.copy_stream, self.config.pin_cpu_memory, + excluded_parameter_prefixes, ) last_hook.prefetch_layer(non_blocking=False) @@ -381,6 +430,7 @@ class LayerWiseOffloadBackend(OffloadBackend): self.device, self.copy_stream, self.config.pin_cpu_memory, + excluded_parameter_prefixes, ) block_hooks.append(hook)