From 3d035bfa190e303f53d72e3baa10885f60abe682 Mon Sep 17 00:00:00 2001 From: Mu GuanLin <1203789601@qq.com> Date: Mon, 24 Aug 2026 06:21:42 +0800 Subject: [PATCH] [Feature] Support MiniMax-H3 Turbo LoRA with the legacy manager (#6476) Signed-off-by: mglyn <1203789601@qq.com> Co-authored-by: Hongsheng Liu --- recipes/MiniMaxAI/MiniMax-H3.md | 67 ++- tests/diffusion/lora/test_lora_manager.py | 77 +++- .../minimax_h3/test_minimax_h3_contract.py | 88 ++-- .../models/minimax_h3/test_minimax_h3_lora.py | 381 ++++++++++++++++++ vllm_omni/diffusion/lora/manager.py | 100 +++-- vllm_omni/diffusion/models/minimax_h3/lora.py | 205 ++++++++++ .../minimax_h3/minimax_h3_transformer.py | 8 + .../models/minimax_h3/pipeline_minimax_h3.py | 101 ++++- 8 files changed, 956 insertions(+), 71 deletions(-) create mode 100644 tests/diffusion/models/minimax_h3/test_minimax_h3_lora.py create mode 100644 vllm_omni/diffusion/models/minimax_h3/lora.py diff --git a/recipes/MiniMaxAI/MiniMax-H3.md b/recipes/MiniMaxAI/MiniMax-H3.md index fb2edfc3d..56dd04901 100644 --- a/recipes/MiniMaxAI/MiniMax-H3.md +++ b/recipes/MiniMaxAI/MiniMax-H3.md @@ -80,7 +80,7 @@ between the legacy and direct planar paths ([full result](https://github.com/vllm-project/vllm-omni/pull/6288#issuecomment-5337546499)): | Metric | Legacy | Direct planar | Change | -|---|---:|---:|---:| +| --- | ---: | ---: | ---: | | Median wall time | 1.805 s | 1.394 s | -22.8% | | Median process CPU time | 3.613 s | 3.207 s | -11.2% | | Peak RSS | 3182 MiB | 2794 MiB | -387 MiB (-12.2%) | @@ -183,7 +183,7 @@ vllm serve "${MODEL}" \ Use the profile that matches the per-GPU memory capacity: | Profile | GPUs | Starting shape | Resident DiT blocks | Attention | Execution | Status | -|---|---:|---:|---:|---|---|---| +| --- | ---: | ---: | ---: | --- | --- | --- | | `rtx5090` | 2 x 32 GB | 1344x768 | 20 | cuDNN attention | eager | Target-hardware validated | | `rtx4090` | 2 x 24 GB | 1024x576 | 12 | cuDNN attention | eager | Capacity-proxy starting point | @@ -202,8 +202,8 @@ increasing it on a different request shape. At vLLM-Omni commit `ae6577ea`, one full 50-step T2VA request completed on 2 x RTX 5090 without OOM: -| Shape | Frames | Client E2E | Sampled peak/GPU | Output validation | -|---:|---:|---:|---:|---| +| Shape | Frames | Client E2E | Sampled peak/GPU | Output validation | +| -------: | ------------: | ---------: | ---------------------: | -----------------------------------------------------------: | | 1344x768 | 124 at 24 FPS | 8 min 38 s | approximately 22.6 GiB | H.264 video + 32 kHz stereo AAC; full `ffmpeg` decode passed | This is a single end-to-end validation run, not a warmed multi-run latency @@ -272,7 +272,7 @@ vllm serve "${MODEL}" \ Do not add `--enforce-eager` to this performance configuration. The first request includes regional compilation; warm the server once before measuring steady-state latency. H3 is CFG-distilled, so `--cfg-parallel-size` must remain -1. The H3 VAE supports its native `tile` mode, not +`1`. The H3 VAE supports its native `tile` mode, not `spatial_shard_height` or `spatial_shard_width`. ### Attention Backends @@ -530,7 +530,7 @@ vLLM-Omni with MiniMax H3 support, BF16. gfx942 rows measured with the `0.26.0+rocm723` wheel (HIP 7.2). | Workload | Configuration | Observed result | -|----------|---------------|-----------------| +| ---------- | --------------- | ----------------- | | T2VA, 1344x768, 209 frames, 50 steps | 4x gfx942 (MI300X), FLASH_ATTN, USP4, text-enc TP4, VAE PP4 tile | encode 0.09 s, denoise 244.04 s, decode 4.15 s, 267.42 s client E2E; H.264 24 FPS + 32 kHz stereo AAC | | FL2VA, 1344x768, 209 frames, 50 steps | 4x gfx942 (MI300X), FLASH_ATTN, USP4, text-enc TP4, VAE PP4 tile | encode 13.98 s, denoise 257.58 s, decode 4.11 s, 287.07 s client E2E; H.264 24 FPS + 32 kHz stereo AAC | | T2VA, 832x480, ~4 s, 40 steps | 1x gfx950 (MI350), FLASH_ATTN, CPU offload | valid MP4 (H.264 + synced audio); ~0.73 s/denoise-step (~1.37 it/s), ~55 s client E2E incl. warmup | @@ -712,7 +712,7 @@ at most 15 seconds combined. ## Official input matrix and limits | Task | Supported references | Limits | -|------|----------------------|--------| +| ------ | ---------------------- | -------- | | T2VA | text only | prompt must be non-empty | | FL2VA | first image, last image, or ordered first+last images | at most 2 images; `frame_indices` is `[0]`, `[-1]`, or `[0,-1]` | | Ref2VA | image-only, image+image, image+video, video+audio, and mixed image/video/audio | images ≤9, videos ≤3, audios ≤3, total references ≤12; audio requires a visual reference | @@ -743,13 +743,17 @@ Omitting `quality` preserves the startup default: it uses the reference path normally, or the server-configured profile when the server was started with `--cache-backend cache_dit`. +Turbo is independent of this quality switch: `quality` selects a Cache-DiT +policy, while Turbo changes the active LoRA weights and sampling schedule. See +[LoRA](#lora) below. + The following result was measured on 4× NVIDIA H200 with SP4, text-encoder TP4, 1344×768, 124 frames, 24 FPS, and 50 inference steps. One full `lossless` warmup was excluded, followed by three fixed prompt/seed pairs in balanced switch order. | `quality` | Median inference latency | Speedup | SSIM vs `lossless` | PSNR vs `lossless` | Expected trade-off | -|---|---:|---:|---:|---:|---| +| --- | ---: | ---: | ---: | ---: | --- | | `lossless` | 85.49 s | 1.00× | 1.0000 | exact | Native reference path | | `high` | 63.36 s | 1.35× | 0.9709 | 34.98 dB | Faster with measured same-seed deviation | @@ -759,10 +763,49 @@ balanced switch order. > workload. The values above apply to this deployment and are not universal > guarantees. `lossless` remains the exact reference path. +## LoRA + +### Turbo LoRA + +Only the native Diffusers 4-step FL2VA/T2VA v1.0 artifact is supported: + +```text +lightx2v/Minimax-h3-Turbo/minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors +``` + +Download only that file: + +```bash +export TURBO_DIR=/path/to/minimax-h3-turbo +export TURBO_FILE=minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors +hf download lightx2v/Minimax-h3-Turbo "${TURBO_FILE}" --local-dir "${TURBO_DIR}" +export TURBO_LORA="${TURBO_DIR}/${TURBO_FILE}" +``` + +Start from a non-offloaded FL2VA server command and add +`--task-type fl2va --lora-backend peft --lora-path "${TURBO_LORA}"`. +`--lora-path` preloads the adapter; each request still activates it and uses +the published sampling settings: + +```bash +-F 'num_inference_steps=5' \ +-F 'flow_shift=6' \ +-F 'extra_params={"task":"t2va","duration":4.4,"audio_flow_shift":3.0}' \ +-F "lora={\"name\":\"h3-turbo-v1.0\",\"path\":\"${TURBO_LORA}\",\"scale\":1.0}" +``` + +For FL2VA, change `task` and add `input_reference` as shown above. The 8-step, +ComfyUI, Ref2VA, and v1.1 artifacts are not supported. This integration is +dynamic-only and does not support prefusion, DLO, or LoRA composition. +It also rejects model-level, layerwise, and distributed layerwise offload; +the legacy dynamic LoRA tensors do not participate in those weight lifecycles. +The five requested sigma points produce the four denoiser evaluations expected +by the Turbo artifact. + ## Key parameters | Parameter | Recommended value | Notes | -|-----------|-------------------|-------| +| ----------- | ------------------- | ------- | | `quality` | omitted or `lossless` | Request-level quality intent; `high` dynamically installs H3's conservative Cache-DiT profile | | `extra_params.force_refresh_step_hint` | omitted | Optional positive 1-based denoising-step hint for an active Cache-DiT request; pair with `extra_params.force_refresh_step_policy`=`once` or `repeat` | | `task` | `t2va`, `fl2va`, or `ref2va` | Passed in `extra_params`; selects the task-specific DiT | @@ -787,9 +830,9 @@ Users can also use a ComfyUI frontend to interact with a hosted MiniMax-H3 servi The four-GPU recommendation was measured on four NVIDIA B300 GPUs with one excluded warmup followed by three requests. -| Workload | Configuration | Observed result | -|----------|---------------|-----------------| -| FL2VA, 209 frames, 1248x768 | no offload, U4, VPP4 tile, regional compile | 86.964 s mean HTTP client latency | +| Workload | Configuration | Observed result | +| -------------------------------------- | ------------------------------------------- | ------------------------------------ | +| FL2VA, 209 frames, 1248x768 | no offload, U4, VPP4 tile, regional compile | 86.964 s mean HTTP client latency | | Two-video Ref2VA, 362 frames, 1344x768 | no offload, U4, VPP4 tile, regional compile | 784.394 s accounted model-stage mean | These measurements describe the validated shapes rather than a general diff --git a/tests/diffusion/lora/test_lora_manager.py b/tests/diffusion/lora/test_lora_manager.py index 723e4b13a..fa80ca7c8 100644 --- a/tests/diffusion/lora/test_lora_manager.py +++ b/tests/diffusion/lora/test_lora_manager.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -20,9 +20,10 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] class _DummyLoRALayer: - def __init__(self, n_slices: int, output_slices: tuple[int, ...]): + def __init__(self, n_slices: int, output_slices: tuple[int, ...], tp_size: int = 1): self.n_slices = n_slices self.output_slices = output_slices + self.tp_size = tp_size self.set_calls: list[ tuple[list[torch.Tensor | None] | torch.Tensor, list[torch.Tensor | None] | torch.Tensor] ] = [] @@ -236,6 +237,14 @@ def test_lora_manager_activates_fused_lora_on_packed_layer(): def test_lora_manager_activates_packed_lora_from_sublayers(): pipeline = torch.nn.Module() + bound_names = None + + def validate_binding(*, lora_model, bound_lora_names): + nonlocal bound_names + assert lora_model.id == 1 + bound_names = bound_lora_names + + pipeline._validate_diffusion_lora_binding = validate_binding pipeline.stacked_params_mapping = [ (".to_qkv", ".to_q", "q"), (".to_qkv", ".to_k", "k"), @@ -279,6 +288,70 @@ def test_lora_manager_activates_packed_lora_from_sublayers(): assert torch.allclose(lora_b_list[0], torch.ones((2, rank)) * 3 * 2.0) assert torch.allclose(lora_b_list[1], torch.ones((1, rank)) * 4 * 2.0) assert torch.allclose(lora_b_list[2], torch.ones((1, rank)) * 4 * 2.0) + assert bound_names == frozenset(loras) + + +def test_lora_manager_rolls_back_all_layers_when_activation_fails(): + class _FailingLoRALayer(_DummyLoRALayer): + fail = False + + def set_lora(self, index: int, lora_a, lora_b): + if self.fail: + raise RuntimeError("bind failed") + super().set_lora(index, lora_a, lora_b) + + manager = DiffusionLoRAManager( + pipeline=torch.nn.Module(), + device=torch.device("cpu"), + dtype=torch.bfloat16, + max_cached_adapters=2, + ) + first = _DummyLoRALayer(n_slices=1, output_slices=(2,)) + second = _FailingLoRALayer(n_slices=1, output_slices=(2,)) + manager._lora_modules = { + "transformer.first": first, + "transformer.second": second, + } + + def adapter(adapter_id: int): + loras = { + name: LoRALayerWeights( + module_name=name, + rank=2, + lora_alpha=2, + lora_a=torch.full((2, 2), float(adapter_id)), + lora_b=torch.full((2, 2), float(adapter_id)), + ) + for name in manager._lora_modules + } + return type( + "LM", + (), + { + "id": adapter_id, + "loras": loras, + "get_lora": lambda self, key: self.loras.get(key), + }, + )() + + manager._registered_adapters = {1: adapter(1), 2: adapter(2)} + manager._activate_adapter(1, scale=1.0) + first_calls_after_success = len(first.set_calls) + + second.fail = True + with pytest.raises(RuntimeError, match="bind failed"): + manager._activate_adapter(2, scale=1.0) + + assert manager._active_adapter_id is None + assert first.reset_calls == 1 + assert second.reset_calls == 1 + + # The failed activation must not leave the old adapter eligible for the + # fast path: reactivating it binds every layer again. + second.fail = False + manager._activate_adapter(1, scale=1.0) + assert manager._active_adapter_id == 1 + assert len(first.set_calls) == first_calls_after_success + 2 def _dummy_lora_request(adapter_id: int) -> LoRARequest: diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index 0f5a11279..7bede350b 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -1,10 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import json import sys from multiprocessing.reduction import ForkingPickler from types import SimpleNamespace +from typing import Any, TypeVar from unittest.mock import Mock, patch import numpy as np @@ -14,6 +15,14 @@ import torch.nn as nn pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] +_ItemT = TypeVar("_ItemT") +_ResultT = TypeVar("_ResultT") + + +def _append_and_return(items: list[_ItemT], item: _ItemT, result: _ResultT) -> _ResultT: + items.append(item) + return result + def test_h3_prepares_resolved_cache_state_immediately_before_denoise(): from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline @@ -185,6 +194,7 @@ def test_startup_task_rejects_unsupported_value(): def test_combined_task_inference_and_transformer_routing(): from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.errors import OmniClientError pipeline = object.__new__(MiniMaxH3Pipeline) torch.nn.Module.__init__(pipeline) @@ -196,6 +206,9 @@ def test_combined_task_inference_and_transformer_routing(): assert pipeline._resolve_task(None, {"image": object()}) == "fl2va" assert pipeline._resolve_task(None, {"audio": object()}) == "ref2va" assert pipeline._resolve_task(None, {"video": object()}) == "ref2va" + assert pipeline._resolve_task("fl2va", {"image": object()}, has_turbo_lora=True) == "fl2va" + with pytest.raises(OmniClientError, match="supports T2VA/FL2VA requests only"): + pipeline._resolve_task("ref2va", {}, has_turbo_lora=True) pipeline.partition = "ref2va" pipeline.supported_tasks = frozenset({"ref2va"}) assert pipeline._resolve_task(None, {"image": object()}) == "ref2va" @@ -260,7 +273,7 @@ def test_pipeline_loads_task_selected_dits_and_shared_components_once( ): (tmp_path / partition_name / component).mkdir() - created = {"dit": [], "text_encoder": [], "video_vae": [], "audio_vae": []} + created: dict[str, list[Any]] = {"dit": [], "text_encoder": [], "video_vae": [], "audio_vae": []} class FakeModule(torch.nn.Module): def __init__(self, *args, **kwargs): @@ -278,15 +291,15 @@ def test_pipeline_loads_task_selected_dits_and_shared_components_once( return create - tokenizer_calls = [] - processor_calls = [] + tokenizer_calls: list[Any] = [] + processor_calls: list[Any] = [] tokenizer_cls = Mock(spec=pipeline_module.Qwen2TokenizerFast) - tokenizer_cls.from_pretrained.side_effect = lambda *args, **kwargs: ( - tokenizer_calls.append((args, kwargs)) or object() + tokenizer_cls.from_pretrained.side_effect = lambda *args, **kwargs: _append_and_return( + tokenizer_calls, (args, kwargs), object() ) processor_cls = Mock(spec=pipeline_module.Qwen3VLProcessor) - processor_cls.from_pretrained.side_effect = lambda *args, **kwargs: ( - processor_calls.append((args, kwargs)) or object() + processor_cls.from_pretrained.side_effect = lambda *args, **kwargs: _append_and_return( + processor_calls, (args, kwargs), object() ) monkeypatch.setattr(pipeline_module, "MiniMaxH3DiTModel", FakeDiT) monkeypatch.setattr( @@ -626,7 +639,7 @@ def test_partially_constructed_pipeline_falls_back_to_the_uniform_schedule(): def test_distilled_forward_reports_denoising_steps_not_sigma_boundaries(): - diffuse_calls = [] + diffuse_calls: list[dict[str, Any]] = [] base_schedule = [1.0, 0.7, 0.4, 0.15, 0.0] pipeline = _distilled_pipeline(diffuse_calls, {"fl2va": base_schedule, "ref2va": None}) @@ -643,7 +656,7 @@ def test_distilled_forward_reports_denoising_steps_not_sigma_boundaries(): def test_distilled_forward_accepts_the_matching_explicit_step_count(): - diffuse_calls = [] + diffuse_calls: list[dict[str, Any]] = [] pipeline = _distilled_pipeline(diffuse_calls, {"fl2va": [1.0, 0.7, 0.4, 0.15, 0.0], "ref2va": None}) pipeline.forward(_t2va_batch(num_inference_steps=4)) @@ -713,8 +726,10 @@ def test_cudnn_packed_attention_uses_python_length_without_padding_mask(): ) assert output.shape == q.shape - assert attention.attention.metadata.attn_mask is None - assert attention.attention.metadata.extra["valid_kv_length"] == 5 + metadata = attention.attention.metadata + assert metadata is not None + assert metadata.attn_mask is None + assert metadata.extra["valid_kv_length"] == 5 def test_packed_attention_skips_mask_for_packed_mask_free_backend(): @@ -754,9 +769,11 @@ def test_packed_attention_skips_mask_for_packed_mask_free_backend(): packed_total=8, ) - assert attention.attention.metadata.attn_mask is None - assert attention.attention.metadata.extra["valid_kv_length"] == 5 - assert attention.attention.metadata.extra["npu_attn_varlen"] is True + metadata = attention.attention.metadata + assert metadata is not None + assert metadata.attn_mask is None + assert metadata.extra["valid_kv_length"] == 5 + assert metadata.extra["npu_attn_varlen"] is True def test_packed_attention_keeps_padding_mask_for_other_backends(): @@ -796,8 +813,10 @@ def test_packed_attention_keeps_padding_mask_for_other_backends(): packed_total=8, ) + metadata = attention.attention.metadata + assert metadata is not None assert torch.equal( - attention.attention.metadata.attn_mask, + metadata.attn_mask, torch.tensor([[True, True, True, True, True, False, False, False]]), ) @@ -1484,12 +1503,12 @@ def test_r7_r8_ref2va_video_segment_matrix(monkeypatch, tmp_path, case, start_ti "audio_codecs": (), "file_size": 1024, } - transcode_calls = [] + transcode_calls: list[Any] = [] monkeypatch.setattr(reference_video_module, "_probe_video", lambda _path: metadata) monkeypatch.setattr( reference_video_module, "_transcode_reference_video", - lambda source, **kwargs: transcode_calls.append((source, kwargs)) or "prepared.mp4", + lambda source, **kwargs: _append_and_return(transcode_calls, (source, kwargs), "prepared.mp4"), ) prepared = reference_video_module.prepare_reference_videos( @@ -1510,7 +1529,7 @@ def test_ref2va_qwen_sampling_uses_one_selective_decode(monkeypatch): from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module decoded = np.arange(3 * 2 * 2 * 3, dtype=np.uint8).reshape(3, 2, 2, 3) - decode_calls = [] + decode_calls: list[Any] = [] monkeypatch.setattr( reference_video_module, "_probe_video", @@ -1523,7 +1542,7 @@ def test_ref2va_qwen_sampling_uses_one_selective_decode(monkeypatch): monkeypatch.setattr( reference_video_module, "_decode_video_frames_ffmpeg", - lambda path, **kwargs: decode_calls.append((path, kwargs)) or decoded, + lambda path, **kwargs: _append_and_return(decode_calls, (path, kwargs), decoded), ) sampled = reference_video_module.sample_reference_video_frames("prepared.mp4") @@ -1584,7 +1603,7 @@ def test_ref2va_load_video_frames_uses_ffmpeg_without_decord(monkeypatch): from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module expected = np.zeros((2, 2, 3, 3), dtype=np.uint8) - decode_calls = [] + decode_calls: list[Any] = [] monkeypatch.setitem(sys.modules, "decord", None) monkeypatch.setattr( reference_video_module, @@ -1598,7 +1617,7 @@ def test_ref2va_load_video_frames_uses_ffmpeg_without_decord(monkeypatch): monkeypatch.setattr( reference_video_module, "_decode_video_frames_ffmpeg", - lambda path, **kwargs: decode_calls.append((path, kwargs)) or expected, + lambda path, **kwargs: _append_and_return(decode_calls, (path, kwargs), expected), ) actual = reference_video_module.load_video_frames("prepared.mp4") @@ -1638,7 +1657,7 @@ def test_ref2va_ffmpeg_decode_rejects_invalid_metadata(kwargs, message): def test_ref2va_probe_uses_container_frame_count_without_scan(monkeypatch): from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module - calls = [] + calls: list[Any] = [] probe = { "streams": [ { @@ -1656,7 +1675,11 @@ def test_ref2va_probe_uses_container_frame_count_without_scan(monkeypatch): monkeypatch.setattr( reference_video_module.subprocess, "run", - lambda command, **kwargs: calls.append((command, kwargs)) or SimpleNamespace(stdout=json.dumps(probe)), + lambda command, **kwargs: _append_and_return( + calls, + (command, kwargs), + SimpleNamespace(stdout=json.dumps(probe)), + ), ) metadata = reference_video_module._probe_video("prepared.mp4") @@ -1669,7 +1692,7 @@ def test_ref2va_probe_uses_container_frame_count_without_scan(monkeypatch): def test_ref2va_probe_counts_frames_only_when_metadata_is_missing(monkeypatch): from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module - calls = [] + calls: list[Any] = [] stream = { "codec_type": "video", "width": 4, @@ -1689,8 +1712,11 @@ def test_ref2va_probe_counts_frames_only_when_metadata_is_missing(monkeypatch): monkeypatch.setattr( reference_video_module.subprocess, "run", - lambda command, **kwargs: calls.append((command, kwargs)) - or SimpleNamespace(stdout=json.dumps(first_probe if len(calls) == 1 else second_probe)), + lambda command, **kwargs: _append_and_return( + calls, + (command, kwargs), + SimpleNamespace(stdout=json.dumps(first_probe if not calls else second_probe)), + ), ) metadata = reference_video_module._probe_video("prepared.mp4") @@ -1750,12 +1776,16 @@ def test_ref2va_two_video_recipe_tolerates_container_rounding(monkeypatch, tmp_p "file_size": 1024, }, } - transcode_calls = [] + transcode_calls: list[Any] = [] monkeypatch.setattr(reference_video_module, "_probe_video", lambda path: metadata[str(path)]) monkeypatch.setattr( reference_video_module, "_transcode_reference_video", - lambda source, **kwargs: transcode_calls.append((source, kwargs)) or f"{source}.prepared.mp4", + lambda source, **kwargs: _append_and_return( + transcode_calls, + (source, kwargs), + f"{source}.prepared.mp4", + ), ) prepared = reference_video_module.prepare_reference_videos( diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_lora.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_lora.py new file mode 100644 index 000000000..58480f4a9 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_lora.py @@ -0,0 +1,381 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +from safetensors.torch import save_file +from vllm.lora.lora_weights import PackedLoRALayerWeights + +from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager +from vllm_omni.diffusion.models.minimax_h3 import lora as lora_module +from vllm_omni.diffusion.models.minimax_h3.lora import ( + load_minimax_h3_turbo_lora, +) +from vllm_omni.errors import OmniClientError +from vllm_omni.lora.request import LoRARequest + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + +_TINY_TARGET_DIMS = { + "attn.to_q": (2, 3), + "attn.to_k": (2, 3), + "attn.to_v": (2, 3), + "attn.to_out.0": (3, 2), + "ff.net.0.proj": (2, 4), + "ff.net.2": (2, 2), +} + + +@pytest.fixture(autouse=True) +def _use_tiny_h3_dimensions(monkeypatch): + monkeypatch.setattr(lora_module, "_TURBO_TARGET_DIMS", _TINY_TARGET_DIMS) + + +def _request(path) -> LoRARequest: + return LoRARequest( + lora_name="turbo", + lora_int_id=1, + lora_path=str(path), + ) + + +def _write_tiny_turbo( + path, + *, + alpha: str = "128", + key_format: str = "minimax-h3-diffusers", + omit_target: str | None = None, + shape_overrides: dict[str, tuple[int, int]] | None = None, +) -> None: + rank = 128 + tensors = {} + overrides = shape_overrides or {} + target_suffixes = ( + "attn.to_q", + "attn.to_k", + "attn.to_v", + "attn.to_out.0", + "ff.net.0.proj", + "ff.net.2", + ) + for prefix, block_count in ( + ("transformer_blocks", 50), + ("token_refiner.refiner_blocks", 2), + ): + for block_index in range(block_count): + for suffix in target_suffixes: + target = f"{prefix}.{block_index}.{suffix}" + if target == omit_target: + continue + input_dim, output_dim = _TINY_TARGET_DIMS[suffix] + a_name = f"{target}.lora_A.default.weight" + b_name = f"{target}.lora_B.default.weight" + tensors[a_name] = torch.ones(overrides.get(a_name, (rank, input_dim))) + if suffix == "ff.net.0.proj": + if b_name in overrides: + tensors[b_name] = torch.ones(overrides[b_name]) + else: + tensors[b_name] = torch.cat( + ( + torch.ones(output_dim // 2, rank), + torch.full((output_dim // 2, rank), 2.0), + ), + dim=0, + ) + else: + tensors[b_name] = torch.ones(overrides.get(b_name, (output_dim, rank))) + save_file( + tensors, + str(path), + metadata={"alpha": alpha, "key_format": key_format}, + ) + + +def test_h3_turbo_loads_through_legacy_lora_model_and_swaps_ffn(tmp_path): + path = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors" + _write_tiny_turbo(path) + + loaded = load_minimax_h3_turbo_lora( + partition="fl2va", + lora_request=_request(path), + lora_path=path, + dtype=torch.float32, + ) + + assert loaded is not None + lora_model, peft_helper = loaded + assert peft_helper.r == 128 + assert peft_helper.lora_alpha == 128 + assert len(lora_model.loras) == 312 + assert "blocks.0.attn.to_q" in lora_model.loras + assert "blocks.0.mlp.fc1" in lora_model.loras + fc1 = lora_model.get_lora("blocks.0.mlp.fc1") + assert isinstance(fc1, PackedLoRALayerWeights) + assert len(fc1.lora_a) == 2 + assert len(fc1.lora_b) == 2 + torch.testing.assert_close(fc1.lora_a[0], torch.ones(128, 2)) + torch.testing.assert_close(fc1.lora_a[1], torch.ones(128, 2)) + torch.testing.assert_close(fc1.lora_b[0], torch.full((2, 128), 2.0)) + torch.testing.assert_close(fc1.lora_b[1], torch.ones(2, 128)) + + +@pytest.mark.parametrize( + ("tensor_suffix", "bad_shape"), + [ + ("attn.to_q.lora_A.default.weight", (128, 1)), + ("attn.to_q.lora_A.default.weight", (128, 3)), + ("attn.to_q.lora_B.default.weight", (2, 128)), + ("attn.to_q.lora_B.default.weight", (4, 128)), + ], +) +def test_h3_turbo_rejects_under_and_oversized_global_shapes(tmp_path, tensor_suffix, bad_shape): + path = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors" + tensor_name = f"transformer_blocks.0.{tensor_suffix}" + _write_tiny_turbo(path, shape_overrides={tensor_name: bad_shape}) + + with pytest.raises(ValueError, match="invalid global shape"): + load_minimax_h3_turbo_lora( + partition="fl2va", + lora_request=_request(path), + lora_path=path, + dtype=torch.float32, + ) + + +def test_legacy_manager_uses_the_h3_model_loader_without_changing_its_interface(tmp_path): + path = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors" + _write_tiny_turbo(path) + + class _Pipeline: + def _load_diffusion_lora_adapter(self, **kwargs): + return load_minimax_h3_turbo_lora(partition="fl2va", **kwargs) + + manager = object.__new__(DiffusionLoRAManager) + manager.pipeline = _Pipeline() + manager.dtype = torch.float32 + manager._expected_lora_modules = {"to_q", "fc1"} + + lora_model, peft_helper = manager._load_adapter(_request(path)) + + assert lora_model.id == 1 + assert peft_helper.lora_alpha == 128 + assert len(lora_model.loras) == 312 + assert "blocks.0.attn.to_q" in lora_model.loras + assert "blocks.0.mlp.fc1" in lora_model.loras + + +def test_h3_turbo_rejects_wrong_alpha_and_ref2va(tmp_path): + wrong_alpha = tmp_path / "wrong_alpha" / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors" + wrong_alpha.parent.mkdir() + _write_tiny_turbo(wrong_alpha, alpha="8") + with pytest.raises(ValueError, match="requires alpha=128"): + load_minimax_h3_turbo_lora( + partition="fl2va", + lora_request=_request(wrong_alpha), + lora_path=wrong_alpha, + dtype=torch.float32, + ) + + valid = tmp_path / "valid" / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors" + valid.parent.mkdir() + _write_tiny_turbo(valid) + with pytest.raises(ValueError, match="supports FL2VA/T2VA only"): + load_minimax_h3_turbo_lora( + partition="ref2va", + lora_request=_request(valid), + lora_path=valid, + dtype=torch.float32, + ) + + +@pytest.mark.parametrize( + "offload_mode", + [ + "model-level CPU offload (--enable-cpu-offload)", + "layerwise offload (--enable-layerwise-offload)", + "distributed layerwise offload (--enable-distributed-layerwise-offload)", + ], +) +def test_h3_turbo_rejects_offload_modes(tmp_path, offload_mode): + path = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors" + _write_tiny_turbo(path) + + with pytest.raises(ValueError, match="does not support"): + load_minimax_h3_turbo_lora( + partition="fl2va", + lora_request=_request(path), + lora_path=path, + dtype=torch.float32, + unsupported_offload_mode=offload_mode, + ) + + +def test_h3_turbo_accepts_only_the_declared_v1_artifact(tmp_path): + unsupported = tmp_path / "minimax_h3_fl2v_turbo_8step_v1.0_bf16.safetensors" + _write_tiny_turbo(unsupported) + + with pytest.raises(ValueError, match="supports only"): + load_minimax_h3_turbo_lora( + partition="fl2va", + lora_request=_request(unsupported), + lora_path=unsupported, + dtype=torch.float32, + ) + + supported = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors" + _write_tiny_turbo(supported) + # Directory resolution selects the declared artifact even when another + # v1.0 checkpoint is present beside it. + assert ( + load_minimax_h3_turbo_lora( + partition="fl2va", + lora_request=_request(tmp_path), + lora_path=tmp_path, + dtype=torch.float32, + ) + is not None + ) + + +def test_h3_turbo_rejects_a_truncated_declared_artifact(tmp_path): + path = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors" + _write_tiny_turbo(path, omit_target="transformer_blocks.49.ff.net.2") + + with pytest.raises(ValueError, match="target set does not match"): + load_minimax_h3_turbo_lora( + partition="fl2va", + lora_request=_request(path), + lora_path=path, + dtype=torch.float32, + ) + + +def test_non_h3_checkpoint_falls_back_to_the_generic_peft_loader(tmp_path): + path = tmp_path / "other.safetensors" + _write_tiny_turbo(path, key_format="other") + + assert ( + load_minimax_h3_turbo_lora( + partition="fl2va", + lora_request=_request(path), + lora_path=path, + dtype=torch.float32, + ) + is None + ) + + declared = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors" + _write_tiny_turbo(declared, key_format="other") + with pytest.raises(ValueError, match="requires safetensors metadata"): + load_minimax_h3_turbo_lora( + partition="fl2va", + lora_request=_request(declared), + lora_path=declared, + dtype=torch.float32, + ) + + +def test_only_an_active_recognized_turbo_adapter_restricts_ref2va(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.models.minimax_h3 import ( + pipeline_minimax_h3 as pipeline_module, + ) + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.partition = "combined" + pipeline.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"}) + pipeline._turbo_lora_adapter_ids = set() + request = _request("generic-peft") + + def load(): + return pipeline._load_diffusion_lora_adapter( + lora_request=request, + lora_path=request.lora_path, + dtype=torch.float32, + ) + + def resolve(scale): + sampling = SimpleNamespace(lora_request=request, lora_scale=scale) + return pipeline._resolve_task( + "ref2va", + {}, + has_turbo_lora=pipeline._has_active_turbo_lora(sampling), + ) + + recognized = (object(), object()) + monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", lambda **_: recognized) + assert load() is recognized + assert resolve(0.0) == "ref2va" + with pytest.raises(OmniClientError, match="supports T2VA/FL2VA requests only"): + resolve(1.0) + + # Simulate manager eviction followed by a generic PEFT adapter reusing the + # same client-supplied ID. A real reload must replace the stale kind. + monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", lambda **_: None) + assert load() is None + assert resolve(1.0) == "ref2va" + + +def test_h3_turbo_requires_all_loaded_targets_to_bind(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline._turbo_lora_adapter_ids = {1} + lora_model = SimpleNamespace(id=1, loras={"to_q": object(), "to_k": object()}) + + pipeline._validate_diffusion_lora_binding( + lora_model=lora_model, + bound_lora_names=frozenset(lora_model.loras), + ) + with pytest.raises(ValueError, match="bound=1/2"): + pipeline._validate_diffusion_lora_binding( + lora_model=lora_model, + bound_lora_names=frozenset({"to_q"}), + ) + + +@pytest.mark.parametrize( + ("num_inference_steps", "extra_args", "error"), + [ + (None, {"flow_shift": 6.0, "audio_flow_shift": 3.0}, "num_inference_steps=5"), + (4, {"flow_shift": 6.0, "audio_flow_shift": 3.0}, "num_inference_steps=5"), + ("5", {"flow_shift": 6.0, "audio_flow_shift": 3.0}, "num_inference_steps=5"), + (5, {"flow_shift": 7.0, "audio_flow_shift": 3.0}, "flow_shift=6"), + (5, {"flow_shift": "bad", "audio_flow_shift": 3.0}, "flow_shift=6"), + (5, {"flow_shift": 6.0, "audio_flow_shift": 4.0}, "audio_flow_shift=3"), + (5, {"flow_shift": 6.0, "audio_flow_shift": []}, "audio_flow_shift=3"), + ], +) +def test_h3_turbo_rejects_unsupported_sampling(num_inference_steps, extra_args, error): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + pipeline.default_video_shift = 12.0 + pipeline.default_audio_shift = 3.0 + sampling = SimpleNamespace( + num_inference_steps=num_inference_steps, + extra_args=extra_args, + ) + + with pytest.raises(OmniClientError, match=error): + pipeline._validate_turbo_sampling(sampling) + + +def test_h3_turbo_accepts_five_sigma_points_for_four_nfe(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + pipeline.default_video_shift = 12.0 + pipeline.default_audio_shift = 3.0 + pipeline._validate_turbo_sampling( + SimpleNamespace( + num_inference_steps=5, + extra_args={"flow_shift": 6.0, "audio_flow_shift": 3.0}, + ) + ) diff --git a/vllm_omni/diffusion/lora/manager.py b/vllm_omni/diffusion/lora/manager.py index 33094bc49..61f6572f6 100644 --- a/vllm_omni/diffusion/lora/manager.py +++ b/vllm_omni/diffusion/lora/manager.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import math import time @@ -289,11 +289,35 @@ class DiffusionLoRAManager: lora_path = get_adapter_absolute_path(lora_request.lora_path) logger.debug("Resolved LoRA path: %s", lora_path) - peft_helper = PEFTHelper.from_local_dir( - lora_path, - max_position_embeddings=None, # no need in diffusion - tensorizer_config_dict=lora_request.tensorizer_config_dict, - ) + model_loader = getattr(self.pipeline, "_load_diffusion_lora_adapter", None) + loaded = None + if callable(model_loader): + loaded = model_loader( + lora_request=lora_request, + lora_path=lora_path, + dtype=self.dtype, + ) + + if loaded is None: + peft_helper = PEFTHelper.from_local_dir( + lora_path, + max_position_embeddings=None, # no need in diffusion + tensorizer_config_dict=lora_request.tensorizer_config_dict, + ) + + lora_model = LoRAModel.from_local_checkpoint( + lora_path, + expected_lora_modules=self._expected_lora_modules, + peft_helper=peft_helper, + lora_model_id=lora_request.lora_int_id, + device="cpu", # consistent w/ vllm's behavior + dtype=self.dtype, + model_vocab_size=None, + tensorizer_config_dict=lora_request.tensorizer_config_dict, + weights_mapper=None, + ) + else: + lora_model, peft_helper = loaded logger.info( "Loaded PEFT config: r=%d, lora_alpha=%d, target_modules=%s", @@ -302,18 +326,6 @@ class DiffusionLoRAManager: peft_helper.target_modules, ) - lora_model = LoRAModel.from_local_checkpoint( - lora_path, - expected_lora_modules=self._expected_lora_modules, - peft_helper=peft_helper, - lora_model_id=lora_request.lora_int_id, - device="cpu", # consistent w/ vllm's behavior - dtype=self.dtype, - model_vocab_size=None, - tensorizer_config_dict=lora_request.tensorizer_config_dict, - weights_mapper=None, - ) - logger.info( "Loaded LoRA model: id=%d, num_modules=%d, modules=%s", lora_model.id, @@ -522,13 +534,17 @@ class DiffusionLoRAManager: matches_scale = self._adapter_scales.get(adapter_id) == rounded_scale return is_active and matches_scale - def _activate_adapter(self, adapter_id: int, scale: float) -> None: - if self._is_active_at_scale(adapter_id, scale): - logger.debug("Adapter %d already active at scale %.3f skipping", adapter_id, scale) - return + def _bind_adapter_weights(self, lora_model: LoRAModel, scale: float) -> None: + binding_validator = getattr(self.pipeline, "_validate_diffusion_lora_binding", None) + lora_names_by_id = ( + {id(weights): name for name, weights in lora_model.loras.items()} if callable(binding_validator) else {} + ) + bound_lora_names: set[str] = set() - logger.info("Activating adapter: id=%d", adapter_id) - lora_model = self._registered_adapters[adapter_id] + def _record_bound(weights: LoRALayerWeights | PackedLoRALayerWeights) -> None: + name = lora_names_by_id.get(id(weights)) + if name is not None: + bound_lora_names.add(name) # activate weights in each LoRA layer for full_module_name, lora_layer in self._lora_modules.items(): @@ -570,6 +586,9 @@ class DiffusionLoRAManager: lora_b_list.append(sub_lora.lora_b * scale) lora_layer.set_lora(index=0, lora_a=lora_a_list, lora_b=lora_b_list) + for sub_lora in sub_loras: + if sub_lora is not None: + _record_bound(sub_lora) logger.debug( "Activated packed LoRA for %s via submodules=%s (scale=%.2f)", full_module_name, @@ -588,6 +607,7 @@ class DiffusionLoRAManager: for b in lora_weights.lora_b ] lora_layer.set_lora(index=0, lora_a=lora_a_list, lora_b=lora_b_list) + _record_bound(lora_weights) logger.debug( "Activated packed LoRA for %s (scale=%.2f)", full_module_name, @@ -618,6 +638,7 @@ class DiffusionLoRAManager: lora_a_list = [lora_weights.lora_a] * n_slices lora_b_list = [b * scale for b in b_splits] lora_layer.set_lora(index=0, lora_a=lora_a_list, lora_b=lora_b_list) + _record_bound(lora_weights) logger.debug( "Activated fused LoRA for packed layer %s (scale=%.2f)", full_module_name, @@ -627,6 +648,7 @@ class DiffusionLoRAManager: scaled_lora_b = lora_weights.lora_b * scale lora_layer.set_lora(index=0, lora_a=lora_weights.lora_a, lora_b=scaled_lora_b) + _record_bound(lora_weights) logger.debug( "Activated LoRA for %s: lora_a shape=%s, lora_b shape=%s, scale=%.2f", full_module_name, @@ -635,6 +657,33 @@ class DiffusionLoRAManager: scale, ) + if callable(binding_validator): + binding_validator( + lora_model=lora_model, + bound_lora_names=frozenset(bound_lora_names), + ) + + def _reset_lora_layers(self) -> None: + for lora_layer in self._lora_modules.values(): + lora_layer.reset_lora(0) + + def _activate_adapter(self, adapter_id: int, scale: float) -> None: + if self._is_active_at_scale(adapter_id, scale): + logger.debug("Adapter %d already active at scale %.3f skipping", adapter_id, scale) + return + + logger.info("Activating adapter: id=%d", adapter_id) + lora_model = self._registered_adapters[adapter_id] + # Binding overwrites slot 0 incrementally. Invalidate the fast-path + # state before the first mutation and leave every wrapper inactive if + # any set_lora() call or model validator fails. + self._active_adapter_id = None + try: + self._bind_adapter_weights(lora_model, scale) + except Exception: + self._reset_lora_layers() + raise + self._active_adapter_id = adapter_id self._update_adapter_scale(adapter_id, scale) @@ -643,8 +692,7 @@ class DiffusionLoRAManager: logger.debug("All adapters already inactive") return logger.info("Deactivating all adapters: %d layers", len(self._lora_modules)) - for lora_layer in self._lora_modules.values(): - lora_layer.reset_lora(0) + self._reset_lora_layers() self._active_adapter_id = None logger.debug("All adapters deactivated") diff --git a/vllm_omni/diffusion/models/minimax_h3/lora.py b/vllm_omni/diffusion/models/minimax_h3/lora.py new file mode 100644 index 000000000..bfc7e5565 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/lora.py @@ -0,0 +1,205 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import math +from pathlib import Path + +import torch +from safetensors import safe_open +from vllm.lora.lora_model import LoRAModel +from vllm.lora.lora_weights import PackedLoRALayerWeights +from vllm.lora.peft_helper import PEFTHelper +from vllm.model_executor.models.utils import WeightsMapper + +from vllm_omni.lora.request import LoRARequest + +_TURBO_RANK = 128 +_TURBO_ALPHA = 128 +_TURBO_HIDDEN_SIZE = 5376 +_TURBO_ATTENTION_INNER_SIZE = 7168 +_TURBO_FFN_HIDDEN_SIZE = 14336 +_TURBO_FILENAME = "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors" +_LORA_A_SUFFIX = ".lora_A.default.weight" +_LORA_B_SUFFIX = ".lora_B.default.weight" +_TURBO_TARGETS = frozenset({"to_q", "to_k", "to_v", "out_proj", "fc1", "fc2"}) +_TURBO_RAW_TARGET_SUFFIXES = ( + "attn.to_q", + "attn.to_k", + "attn.to_v", + "attn.to_out.0", + "ff.net.0.proj", + "ff.net.2", +) +_TURBO_TARGET_DIMS = { + "attn.to_q": (_TURBO_HIDDEN_SIZE, _TURBO_ATTENTION_INNER_SIZE), + "attn.to_k": (_TURBO_HIDDEN_SIZE, _TURBO_ATTENTION_INNER_SIZE), + "attn.to_v": (_TURBO_HIDDEN_SIZE, _TURBO_ATTENTION_INNER_SIZE), + "attn.to_out.0": (_TURBO_ATTENTION_INNER_SIZE, _TURBO_HIDDEN_SIZE), + "ff.net.0.proj": (_TURBO_HIDDEN_SIZE, 2 * _TURBO_FFN_HIDDEN_SIZE), + "ff.net.2": (_TURBO_FFN_HIDDEN_SIZE, _TURBO_HIDDEN_SIZE), +} +_TURBO_EXPECTED_RAW_TARGETS = frozenset( + f"{prefix}.{block_index}.{suffix}" + for prefix, block_count in ( + ("transformer_blocks", 50), + ("token_refiner.refiner_blocks", 2), + ) + for block_index in range(block_count) + for suffix in _TURBO_RAW_TARGET_SUFFIXES +) +_TURBO_TARGET_PATTERN = ( + r"^transformer\.(?:token_refiner\.blocks|blocks)\.\d+\." + r"(?:attn\.(?:to_q|to_k|to_v|out_proj)|mlp\.(?:fc1|fc2))$" +) + +_TURBO_WEIGHTS_MAPPER = WeightsMapper( + orig_to_new_substr={ + "token_refiner.refiner_blocks.": "token_refiner.blocks.", + "transformer_blocks.": "blocks.", + ".attn.to_out.0.": ".attn.out_proj.", + ".ff.net.0.proj.": ".mlp.fc1.", + ".ff.net.2.": ".mlp.fc2.", + ".lora_A.default.": ".lora_A.", + ".lora_B.default.": ".lora_B.", + } +) + + +def _select_turbo_file(artifact_path: str | Path) -> Path | None: + path = Path(artifact_path) + if path.is_file(): + return path if path.suffix == ".safetensors" else None + if not path.is_dir(): + return None + + candidate = path / _TURBO_FILENAME + return candidate if candidate.is_file() else None + + +def _validate_and_convert_tensors(checkpoint) -> dict[str, torch.Tensor]: + tensors: dict[str, torch.Tensor] = {} + pairs: dict[str, set[str]] = {} + raw_targets: set[str] = set() + for name in checkpoint.keys(): + if name.endswith(_LORA_A_SUFFIX): + raw_target = name[: -len(_LORA_A_SUFFIX)] + side = "a" + elif name.endswith(_LORA_B_SUFFIX): + raw_target = name[: -len(_LORA_B_SUFFIX)] + side = "b" + else: + raise ValueError(f"Unconsumed MiniMax-H3 Turbo tensor: {name!r}") + raw_targets.add(raw_target) + + mapped_name = _TURBO_WEIGHTS_MAPPER.apply_list([name])[0] + mapped_target = mapped_name.rsplit(".lora_", 1)[0] + if mapped_target.rsplit(".", 1)[-1] not in _TURBO_TARGETS: + raise ValueError(f"Unsupported MiniMax-H3 Turbo target: {raw_target!r}") + target_sides = pairs.setdefault(mapped_target, set()) + if side in target_sides: + raise ValueError(f"Duplicate MiniMax-H3 Turbo tensor for {mapped_target}.{side}") + target_sides.add(side) + + tensor = checkpoint.get_tensor(name) + if tensor.ndim != 2: + raise ValueError(f"MiniMax-H3 Turbo LoRA tensors must be matrices, got {name}={tuple(tensor.shape)}") + suffix = next((suffix for suffix in _TURBO_RAW_TARGET_SUFFIXES if raw_target.endswith(suffix)), None) + if suffix is None: + raise ValueError(f"MiniMax-H3 Turbo LoRA contains unsupported target: {raw_target}") + input_dim, output_dim = _TURBO_TARGET_DIMS[suffix] + expected_shape = (_TURBO_RANK, input_dim) if side == "a" else (output_dim, _TURBO_RANK) + if tuple(tensor.shape) != expected_shape: + raise ValueError( + f"MiniMax-H3 Turbo tensor has invalid global shape: {name}={tuple(tensor.shape)}, " + f"expected={expected_shape}" + ) + if side == "b" and ".ff.net.0.proj." in name: + value, gate = tensor.chunk(2, dim=0) + tensor = torch.cat((gate, value), dim=0).contiguous() + tensors[name] = tensor + + incomplete = sorted(target for target, sides in pairs.items() if sides != {"a", "b"}) + if incomplete: + raise ValueError(f"Incomplete MiniMax-H3 Turbo LoRA pairs: {incomplete}") + missing = sorted(_TURBO_EXPECTED_RAW_TARGETS - raw_targets) + unexpected = sorted(raw_targets - _TURBO_EXPECTED_RAW_TARGETS) + if missing or unexpected: + raise ValueError( + "MiniMax-H3 Turbo target set does not match the supported v1.0 artifact: " + f"missing={len(missing)} {missing[:5]}, unexpected={len(unexpected)} {unexpected[:5]}" + ) + return tensors + + +def _pack_h3_turbo_fc1(lora_model: LoRAModel) -> None: + """Represent H3's fused gate/up projection without generic layout guesses.""" + + for module_name, weights in tuple(lora_model.loras.items()): + if not module_name.endswith(".mlp.fc1"): + continue + gate_b, up_b = weights.lora_b.chunk(2, dim=0) + lora_model.loras[module_name] = PackedLoRALayerWeights( + module_name=module_name, + rank=weights.rank, + lora_alphas=[weights.lora_alpha, weights.lora_alpha], + lora_a=[weights.lora_a, weights.lora_a], + lora_b=[gate_b.contiguous(), up_b.contiguous()], + scaling=[weights.scaling, weights.scaling], + ) + + +def load_minimax_h3_turbo_lora( + *, + partition: str, + lora_request: LoRARequest, + lora_path: str | Path, + dtype: torch.dtype, + unsupported_offload_mode: str | None = None, +) -> tuple[LoRAModel, PEFTHelper] | None: + """Load the published LightX2V Turbo v1.0 through the legacy manager.""" + + lora_file = _select_turbo_file(lora_path) + if lora_file is None: + return None + with safe_open(lora_file, framework="pt", device="cpu") as checkpoint: + metadata = checkpoint.metadata() or {} + if metadata.get("key_format") != "minimax-h3-diffusers": + if lora_file.name == _TURBO_FILENAME: + raise ValueError( + "MiniMax-H3 Turbo v1.0 requires safetensors metadata key_format='minimax-h3-diffusers'" + ) + return None + if lora_file.name != _TURBO_FILENAME: + raise ValueError(f"MiniMax-H3 Turbo supports only {_TURBO_FILENAME!r}, got {lora_file.name!r}") + raw_alpha = metadata.get("alpha") + try: + alpha = float(raw_alpha) if raw_alpha is not None else math.nan + except ValueError as exc: + raise ValueError(f"MiniMax-H3 Turbo alpha must be numeric, got {raw_alpha!r}") from exc + if alpha != _TURBO_ALPHA: + raise ValueError(f"MiniMax-H3 Turbo v1.0 requires alpha={_TURBO_ALPHA}, got {raw_alpha!r}") + if partition == "ref2va": + raise ValueError("MiniMax-H3 Turbo LoRA supports FL2VA/T2VA only") + if unsupported_offload_mode is not None: + raise ValueError(f"MiniMax-H3 Turbo dynamic LoRA does not support {unsupported_offload_mode}") + tensors = _validate_and_convert_tensors(checkpoint) + + peft_helper = PEFTHelper.from_dict( + { + "r": _TURBO_RANK, + "lora_alpha": _TURBO_ALPHA, + "target_modules": _TURBO_TARGET_PATTERN, + } + ) + lora_model = LoRAModel.from_lora_tensors( + lora_model_id=lora_request.lora_int_id, + tensors=tensors, + peft_helper=peft_helper, + device="cpu", + dtype=dtype, + weights_mapper=_TURBO_WEIGHTS_MAPPER, + ) + _pack_h3_turbo_fc1(lora_model) + return lora_model, peft_helper 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 91e03c865..1b852969a 100644 --- a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py +++ b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py @@ -954,6 +954,14 @@ class MiniMaxH3DiTModel(nn.Module): # quantization or LoRA to map onto. Address the fused layers directly, e.g. # ignored_layers=["blocks.0.attn.qkv_proj"]. packed_modules_mapping = {} + # Turbo LoRA checkpoints publish separate Q/K/V adapters. This declaration + # lets the legacy diffusion LoRA manager bind them to the packed QKV layer; + # it does not change the fused base-checkpoint loading path above. + stacked_params_mapping = ( + (".attn.qkv_proj", ".attn.to_q", "q"), + (".attn.qkv_proj", ".attn.to_k", "k"), + (".attn.qkv_proj", ".attn.to_v", "v"), + ) def _validate_tp_config(self, *, arch: MiniMaxH3DiTArchConfig, tp_size: int) -> None: if tp_size < 1: diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index 62d10488c..e6e1a03fb 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """vLLM-Omni pipeline for MiniMax H3 FL2VA and Ref2VA partitions.""" from __future__ import annotations @@ -11,7 +12,7 @@ from collections.abc import Iterable, Mapping, Sequence from contextlib import contextmanager, nullcontext from itertools import groupby from pathlib import Path -from typing import Any, ClassVar +from typing import TYPE_CHECKING, Any, ClassVar import numpy as np import torch @@ -74,6 +75,7 @@ from .condition_noise import ( ) from .denoise_loop import MiniMaxH3DenoiseBranch, minimax_h3_denoise_loop from .encoder import MiniMaxH3Qwen3VLEncoder +from .lora import load_minimax_h3_turbo_lora from .minimax_h3_transformer import MiniMaxH3DiTModel from .packed_sequence import ( minimax_h3_packed_sequence, @@ -110,6 +112,12 @@ from .vae import MiniMaxH3AudioVAE, MiniMaxH3VideoVAE logger = init_logger(__name__) +if TYPE_CHECKING: + from vllm.lora.lora_model import LoRAModel + from vllm.lora.peft_helper import PEFTHelper + + from vllm_omni.lora.request import LoRARequest + MINIMAX_H3_FPS = 24 MINIMAX_H3_AUDIO_SAMPLE_RATE = 32000 MINIMAX_H3_IMGVID_COND_TIMESTEP = 0.999 @@ -130,6 +138,9 @@ MINIMAX_H3_MAX_REFERENCE_IMAGE_BYTES = 30 * 1024 * 1024 MINIMAX_H3_REFERENCE_IMAGE_FORMATS = frozenset({"jpeg", "png", "webp", "heic", "heif"}) MINIMAX_H3_MIN_OUTPUT_SECONDS = 4.0 MINIMAX_H3_MAX_OUTPUT_SECONDS = 15.0 +MINIMAX_H3_TURBO_SIGMA_POINTS = 5 +MINIMAX_H3_TURBO_VIDEO_SHIFT = 6.0 +MINIMAX_H3_TURBO_AUDIO_SHIFT = 3.0 MINIMAX_H3_DOWNLOAD_PATTERNS = [ "FL2VA/**", "Ref2VA/model_index.json", @@ -559,6 +570,80 @@ class MiniMaxH3Pipeline( # uniform path available to partially constructed pipelines. _base_schedule_by_partition: ClassVar[Mapping[str, DMD2SigmaSchedule | None]] = {} + def _load_diffusion_lora_adapter( + self, + *, + lora_request: LoRARequest, + lora_path: str | Path, + dtype: torch.dtype, + ) -> tuple[LoRAModel, PEFTHelper] | None: + # A cache eviction may be followed by a different adapter reusing the + # same client-supplied ID. Every real load replaces the classification. + self._turbo_lora_adapter_ids.discard(lora_request.lora_int_id) + od_config = getattr(self, "od_config", None) + offload_modes = [] + if getattr(od_config, "enable_cpu_offload", False): + offload_modes.append("model-level CPU offload (--enable-cpu-offload)") + if getattr(od_config, "enable_layerwise_offload", False): + offload_modes.append("layerwise offload (--enable-layerwise-offload)") + if getattr(od_config, "enable_distributed_layerwise_offload", False): + offload_modes.append("distributed layerwise offload (--enable-distributed-layerwise-offload)") + loaded = load_minimax_h3_turbo_lora( + partition=self.partition, + lora_request=lora_request, + lora_path=lora_path, + dtype=dtype, + unsupported_offload_mode=" or ".join(offload_modes) or None, + ) + if loaded is not None: + self._turbo_lora_adapter_ids.add(lora_request.lora_int_id) + return loaded + + def _validate_diffusion_lora_binding( + self, + *, + lora_model: LoRAModel, + bound_lora_names: frozenset[str], + ) -> None: + if lora_model.id not in self._turbo_lora_adapter_ids: + return + missing = sorted(set(lora_model.loras) - bound_lora_names) + if missing: + raise ValueError( + "MiniMax-H3 Turbo LoRA binding is incomplete: " + f"bound={len(bound_lora_names)}/{len(lora_model.loras)}, missing={missing[:5]}" + ) + + def _has_active_turbo_lora(self, sampling: Any) -> bool: + lora_request = sampling.lora_request + return ( + lora_request is not None + and not math.isclose(0.0, float(sampling.lora_scale)) + and lora_request.lora_int_id in self._turbo_lora_adapter_ids + ) + + def _validate_turbo_sampling(self, sampling: Any) -> None: + extra = sampling.extra_args or {} + sigma_points = sampling.num_inference_steps + if sigma_points != MINIMAX_H3_TURBO_SIGMA_POINTS: + raise OmniClientError( + "MiniMax-H3 Turbo requires num_inference_steps=5 (five sigma points produce four denoiser evaluations)" + ) + try: + video_shift = float(extra.get("flow_shift", self.default_video_shift)) + except (TypeError, ValueError) as exc: + raise OmniClientError(f"MiniMax-H3 Turbo requires flow_shift={MINIMAX_H3_TURBO_VIDEO_SHIFT:g}") from exc + if not math.isclose(video_shift, MINIMAX_H3_TURBO_VIDEO_SHIFT): + raise OmniClientError(f"MiniMax-H3 Turbo requires flow_shift={MINIMAX_H3_TURBO_VIDEO_SHIFT:g}") + try: + audio_shift = float(extra.get("audio_flow_shift", self.default_audio_shift)) + except (TypeError, ValueError) as exc: + raise OmniClientError( + f"MiniMax-H3 Turbo requires audio_flow_shift={MINIMAX_H3_TURBO_AUDIO_SHIFT:g}" + ) from exc + if not math.isclose(audio_shift, MINIMAX_H3_TURBO_AUDIO_SHIFT): + raise OmniClientError(f"MiniMax-H3 Turbo requires audio_flow_shift={MINIMAX_H3_TURBO_AUDIO_SHIFT:g}") + def adopt_cache_dit_backend(self, backend: CacheDiTBackend) -> None: """Adopt runner-installed generic Cache-DiT for request transitions.""" @@ -589,6 +674,7 @@ class MiniMaxH3Pipeline( getattr(od_config, "task_type", None), str(od_config.model), ) + self._turbo_lora_adapter_ids: set[int] = set() model_root = _resolve_minimax_h3_model_root( str(od_config.model), od_config.revision, @@ -781,6 +867,8 @@ class MiniMaxH3Pipeline( self, requested: str | None, multi_modal_data: dict[str, Any], + *, + has_turbo_lora: bool = False, ) -> str: if requested is None: # A Ref2VA-only startup has no FL2VA transformer; preserve its @@ -798,6 +886,8 @@ class MiniMaxH3Pipeline( raise OmniClientError( f"checkpoint partition {self.partition!r} supports {sorted(self.supported_tasks)}, got task={task!r}" ) + if task == "ref2va" and has_turbo_lora: + raise OmniClientError("MiniMax-H3 Turbo LoRA supports T2VA/FL2VA requests only") return task def _resolve_shape( @@ -1655,7 +1745,14 @@ class MiniMaxH3Pipeline( quality = sampling.quality logger.debug("MiniMax H3 request quality=%s", quality) extra = sampling.extra_args or {} - task = self._resolve_task(extra.get("task"), multi_modal_data) + has_turbo_lora = self._has_active_turbo_lora(sampling) + task = self._resolve_task( + extra.get("task"), + multi_modal_data, + has_turbo_lora=has_turbo_lora, + ) + if has_turbo_lora: + self._validate_turbo_sampling(sampling) raw_image = multi_modal_data.get("image") raw_videos = multi_modal_data.get("video")