From cb439f3e84394c0cdfc252eaff3303fb099cc67d Mon Sep 17 00:00:00 2001 From: WeiQing Chen <40507679+david6666666@users.noreply.github.com> Date: Thu, 17 Sep 2026 16:41:52 +0800 Subject: [PATCH] [Diffusion][Model] Add direct FastH3 8-Step V2 checkpoint loading (#7610) Signed-off-by: david6666666 <530634352@qq.com> --- recipes/MiniMaxAI/MiniMax-H3.md | 62 ++ .../minimax_h3/test_fasth3_checkpoint.py | 679 ++++++++++++++++++ .../minimax_h3/test_minimax_h3_contract.py | 1 + .../minimax_h3/test_minimax_h3_fasth3.py | 5 +- .../test_minimax_h3_quantization.py | 23 +- .../test_minimax_h3_quantization_quality.py | 6 +- .../diffusion/models/minimax_h3/fasth3.py | 3 + .../models/minimax_h3/fasth3_checkpoint.py | 131 ++++ .../minimax_h3/minimax_h3_transformer.py | 76 +- .../models/minimax_h3/pipeline_minimax_h3.py | 52 +- .../models/minimax_h3/checkpoint.py | 29 +- 11 files changed, 1053 insertions(+), 14 deletions(-) create mode 100644 tests/diffusion/models/minimax_h3/test_fasth3_checkpoint.py create mode 100644 vllm_omni/diffusion/models/minimax_h3/fasth3_checkpoint.py diff --git a/recipes/MiniMaxAI/MiniMax-H3.md b/recipes/MiniMaxAI/MiniMax-H3.md index b48a1f472..934e882f7 100644 --- a/recipes/MiniMaxAI/MiniMax-H3.md +++ b/recipes/MiniMaxAI/MiniMax-H3.md @@ -1106,6 +1106,68 @@ generations move the end-to-end figure toward the denoising one. Fusing the adapter does not measurably change startup: weight loading took 77.3 s with it against 85.8 s without. +### FastH3 8-Step V2 full checkpoint + +[FastH3 8-Step V2](https://huggingface.co/FastVideo/FastVideo-FastH3-8-Step-V2) +ships a full Diffusers-format transformer, including learned VSA compression +gates. Serve the Hugging Face model ID or a downloaded snapshot directly: + +```bash +export FASTH3_V2_MODEL=FastVideo/FastVideo-FastH3-8-Step-V2 +``` + +Use a Hugging Face account with access to the release, and pin `--revision` +when reproducing a result. The existing component loader reads its safetensors +shards; the H3 weight loader maps names and loads Q/K/V and MLP projections into +native parameters in memory. No conversion command or second transformer +checkpoint is needed. + +The transformer and text components come from the V2 release. To retain Omni's +native tiled and parallel VAE execution, the loader fetches only the video/audio +VAEs from the `MiniMaxAI/MiniMax-H3` revision pinned in V2's `provenance.json`. +Those components use the normal Hugging Face cache. No separate base-model +argument or full base-transformer download is required. Download time is a +preparation cost, separate from warmup and request latency. + +Install the optional `fastvideo-kernel` build required by the +`FASTVIDEO_VSA` backend, including its H3 block-map entry point. A four-GPU +configuration is: + +```bash +CUDA_VISIBLE_DEVICES=0,1,2,3 vllm serve "${FASTH3_V2_MODEL}" --omni \ + --host 127.0.0.1 --port 8095 --trust-remote-code --task-type fl2va \ + --served-model-name FastH3-V2 \ + --num-gpus 4 --usp 4 --ring 1 \ + --vae-patch-parallel-size 4 --vae-parallel-mode tile --vae-use-tiling \ + --diffusion-attention-backend FASTVIDEO_VSA + +curl --fail-with-body -sS http://127.0.0.1:8095/v1/videos/sync \ + -F 'model=FastH3-V2' \ + -F 'prompt=A red fox runs through fresh snow at dawn, with fast pawsteps, winter wind and distant birds.' \ + -F 'width=1344' -F 'height=768' -F 'aspect_ratio=16:9' -F 'fps=24' \ + -F 'num_inference_steps=8' -F 'guidance_scale=1' -F 'flow_shift=10' \ + -F 'seed=1101' \ + -F 'extra_params={"task":"t2va","duration":4.4,"audio_flow_shift":3.0}' \ + -o fasth3_v2.mp4 +``` + +V2 pins video/audio shifts **10/3**, guidance **1**, and the unshifted ladder +`[0.999, 0.874, 0.749, 0.624, 0.5, 0.375, 0.25, 0.125, 0.0]`. +Use **`num_inference_steps=8`** in Omni: these nine nodes bound eight model +evaluations. FastVideo's `--steps 9` counts nodes instead. The existing H3 ODE +update is reused without stochastic re-noising. + +The trained attention policy is **80% sparsity with 64-token video tiles**. +Omni derives the number of selected video tiles from each request's geometry; +do not supply a fixed `--fastvideo-vsa-topk`. Text/condition/audio prefix keys +remain available to every query, and prefix queries remain dense. Missing VSA +geometry or a failed H3 kernel raises an error instead of serving the student +with dense attention. + +This release supports T2VA only, with local or pure Ulysses attention. Additional +LoRA adapters, Ref2VA and FL2VA conditioning are unsupported. The legacy +four-step adapter keeps its original sampling and fixed-top-k behavior. + ## Key parameters | Parameter | Recommended value | Notes | diff --git a/tests/diffusion/models/minimax_h3/test_fasth3_checkpoint.py b/tests/diffusion/models/minimax_h3/test_fasth3_checkpoint.py new file mode 100644 index 000000000..9a530bf90 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_fasth3_checkpoint.py @@ -0,0 +1,679 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import json +from collections.abc import Callable, Iterable + +import pytest +import torch +import torch.nn as nn + +from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, DiffusionParallelConfig, OmniDiffusionConfig +from vllm_omni.diffusion.models.minimax_h3.fasth3_checkpoint import ( + FASTH3_BASE_MODEL, + FASTH3_V2_BASE_SCHEDULE, + FASTH3_V2_MODEL_ID, + FastH3CheckpointSpec, +) +from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( + MiniMaxH3DiTArchConfig, + MiniMaxH3DiTModel, +) +from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( + resolve_minimax_h3_diffusion_model_path, +) +from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule +from vllm_omni.errors import OmniClientError +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.lora.request import LoRARequest +from vllm_omni.model_executor.models.minimax_h3.checkpoint import ( + is_minimax_h3_modular, + resolve_minimax_h3_model_root, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def _fastvideo_metadata() -> dict[str, object]: + """Return the release-owned metadata read from fastvideo_inference.json.""" + return { + "model_id": FASTH3_V2_MODEL_ID, + "schema_version": "fasth3-inference-contract-v1", + "dmd_denoising_steps": [999, 874, 749, 624, 500, 375, 250, 125], + "transformer_forwards": 8, + "video_scheduler_shift": 10.0, + "audio_scheduler_shift": 3.0, + "guidance_scale": 1.0, + "attention_backend": "VIDEO_SPARSE_ATTN_H3", + "vsa_sparsity": 0.8, + "vsa_tile_size": 64, + "task": "t2av", + } + + +def _sampling( + *, + num_inference_steps: int | str | None = 8, + guidance_scale: float | None = 1.0, + extra_args: dict[str, object] | None = None, + lora_request: LoRARequest | None = None, + timesteps: torch.Tensor | None = None, + sigmas: list[float] | None = None, +) -> OmniDiffusionSamplingParams: + return OmniDiffusionSamplingParams( + # Keep the string form to exercise malformed client input validation. + num_inference_steps=num_inference_steps, # type: ignore[arg-type] + guidance_scale=guidance_scale, + extra_args=extra_args or {}, + lora_request=lora_request, + timesteps=timesteps, + sigmas=sigmas, + ) + + +def _od_config( + *, + backend: str = "FASTVIDEO_VSA", + topk: int | None = None, + per_role: dict[str, AttentionSpec] | None = None, + lora_path: str | None = None, + ring_degree: int = 1, + allgather_degree: int = 1, +) -> OmniDiffusionConfig: + attention_config = AttentionConfig( + default=AttentionSpec(backend=backend, fastvideo_vsa_topk=topk), + per_role=per_role or {}, + ) + return OmniDiffusionConfig( + diffusion_attention_config=attention_config, + lora_path=lora_path, + parallel_config=DiffusionParallelConfig(ring_degree=ring_degree, allgather_degree=allgather_degree), + ) + + +def test_v2_release_has_nine_sigma_nodes_and_eight_intervals(): + assert FASTH3_V2_BASE_SCHEDULE.base_schedule == ( + 0.999, + 0.874, + 0.749, + 0.624, + 0.5, + 0.375, + 0.25, + 0.125, + 0.0, + ) + assert len(FASTH3_V2_BASE_SCHEDULE.base_schedule) == 9 + assert FASTH3_V2_BASE_SCHEDULE.num_inference_steps == 8 + assert FASTH3_V2_BASE_SCHEDULE.shifted_sigmas(10.0) == pytest.approx( + [ + 0.9998999099189271, + 0.985788405143244, + 0.9675752486758817, + 0.94316807738815, + 0.9090909090909091, + 0.8571428571428571, + 0.7692307692307693, + 0.5882352941176471, + 0.0, + ], + abs=1e-12, + ) + assert FASTH3_V2_BASE_SCHEDULE.shifted_sigmas(3.0) == pytest.approx( + [ + 0.9996664442961973, + 0.9541484716157204, + 0.8995196156925539, + 0.8327402135231315, + 0.75, + 0.6428571428571429, + 0.5, + 0.3, + 0.0, + ], + abs=1e-12, + ) + + +def test_fastvideo_metadata_builds_release_sampling_metadata(): + spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()) + + assert spec.vsa_sparsity == 0.8 + assert spec.release_metadata() == { + "partition": "fl2va", + "tasks": ["t2va"], + "sigma_shift_scales": {"video": 10.0, "audio": 3.0}, + "base_schedule": list(FASTH3_V2_BASE_SCHEDULE.base_schedule), + } + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("model_id", "FastVideo/FastH3-other"), + ("schema_version", "fasth3-inference-contract-v0"), + ("dmd_denoising_steps", [999, 749, 500, 250]), + ("transformer_forwards", 4), + ("video_scheduler_shift", 12.0), + ("audio_scheduler_shift", 4.0), + ("guidance_scale", 2.0), + ("attention_backend", "FLASH_ATTN"), + ("vsa_sparsity", 0.9), + ("vsa_tile_size", 32), + ("task", "t2va"), + ], +) +def test_fastvideo_metadata_rejects_contract_drift(field, value): + metadata = _fastvideo_metadata() + metadata[field] = value + + with pytest.raises(ValueError, match=field): + FastH3CheckpointSpec.from_metadata(metadata) + + +def test_fastvideo_metadata_rejects_missing_contract_field(): + metadata = _fastvideo_metadata() + metadata.pop("task") + + with pytest.raises(ValueError, match="task"): + FastH3CheckpointSpec.from_metadata(metadata) + + +def test_legacy_h3_metadata_is_not_treated_as_fastvideo_release(): + with pytest.raises(ValueError, match="model_id"): + FastH3CheckpointSpec.from_metadata( + { + "schema_version": 1, + "partition": "fl2va", + "tasks": ["t2va", "fl2va"], + "sigma_shift_scales": {"video": 12.0, "audio": 3.0}, + } + ) + + +def test_v2_serving_contract_accepts_backend_selected_for_h3_self_role(): + config = _od_config( + backend="TORCH_SDPA", + per_role={"self": AttentionSpec(backend="FASTVIDEO_VSA")}, + ) + FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()).check_serving_contract( + partition="fl2va", od_config=config + ) + + +@pytest.mark.parametrize( + ("kwargs", "partition", "match"), + [ + ({}, "ref2va", "task-type fl2va"), + ({"lora_path": "/tmp/adapter"}, "fl2va", "LoRA"), + ({"backend": "TORCH_SDPA"}, "fl2va", "FASTVIDEO_VSA"), + ({"topk": 2}, "fl2va", "fixed fastvideo_vsa_topk"), + ({"ring_degree": 2}, "fl2va", "local attention"), + ({"allgather_degree": 2}, "fl2va", "local attention"), + ], +) +def test_v2_serving_contract_rejects_incompatible_runtime(kwargs, partition, match): + spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()) + with pytest.raises(ValueError, match=match): + spec.check_serving_contract(partition=partition, od_config=_od_config(**kwargs)) + + +def test_v2_serving_contract_reads_a_per_role_fixed_topk_override(): + spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()) + config = _od_config( + per_role={"self": AttentionSpec(backend="FASTVIDEO_VSA", fastvideo_vsa_topk=1)}, + ) + with pytest.raises(ValueError, match="fixed fastvideo_vsa_topk"): + spec.check_serving_contract(partition="fl2va", od_config=config) + + +def test_v2_request_contract_accepts_omitted_steps_and_default_shifts(): + spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()) + spec.check_request(_sampling(num_inference_steps=None)) + spec.check_request(_sampling(extra_args={"flow_shift": 10, "audio_flow_shift": 3}, guidance_scale=1)) + + +@pytest.mark.parametrize( + "sampling", + [ + _sampling(num_inference_steps=None), + _sampling(timesteps=torch.tensor([0.9, 0.5, 0.0])), + _sampling(sigmas=[1.0, 0.5, 0.0]), + ], +) +def test_v2_step_execution_requires_pinned_count_and_schedule(sampling): + spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()) + with pytest.raises(OmniClientError, match="step execution"): + spec.check_request(sampling, step_execution=True) + + +def test_v2_step_execution_accepts_only_explicit_pinned_count(): + spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()) + spec.check_request(_sampling(num_inference_steps=8), step_execution=True) + + +@pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"num_inference_steps": 7}, "num_inference_steps=8"), + ({"num_inference_steps": 9}, "num_inference_steps=8"), + ({"num_inference_steps": "8"}, "num_inference_steps=8"), + ({"extra_args": {"flow_shift": 9}}, "flow_shift=10"), + ({"extra_args": {"audio_flow_shift": 4}}, "audio_flow_shift=3"), + ({"extra_args": {"flow_shift": "bad"}}, "flow_shift=10"), + ({"guidance_scale": 2.0}, "guidance_scale=1"), + ( + {"lora_request": LoRARequest(lora_name="test", lora_int_id=1, lora_path="/tmp/test.safetensors")}, + "per-request LoRA", + ), + ], +) +def test_v2_request_contract_rejects_sampling_drift(kwargs, match): + spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()) + with pytest.raises(OmniClientError, match=match): + spec.check_request(_sampling(**kwargs)) + + +def test_resolve_native_vaes_uses_pinned_revision_and_only_vae_patterns(tmp_path, monkeypatch): + model_root = tmp_path / "FastH3" + model_root.mkdir() + revision = "base-revision-123" + (model_root / "provenance.json").write_text( + json.dumps({"base_model": f"hf://{FASTH3_BASE_MODEL}@{revision}"}), + encoding="utf-8", + ) + snapshot = tmp_path / "hf-snapshot" + calls: dict[str, object] = {} + + def fake_download(**kwargs: object) -> str: + calls.update(kwargs) + snapshot.mkdir() + return str(snapshot) + + monkeypatch.setattr( + "vllm_omni.diffusion.models.minimax_h3.fasth3_checkpoint.download_weights_from_hf_specific", + fake_download, + ) + + result = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()).resolve_native_vaes(model_root) + + assert result == snapshot / "FL2VA" + assert calls == { + "model_name_or_path": FASTH3_BASE_MODEL, + "cache_dir": None, + "allow_patterns": ["FL2VA/video_vae/**", "FL2VA/audio_vae/**"], + "revision": revision, + "require_all": True, + } + assert sorted(path.name for path in model_root.iterdir()) == ["provenance.json"] + assert not (model_root / "model_index.json").exists() + assert not (model_root / "conversion.json").exists() + + +@pytest.mark.parametrize( + "base_model", + [ + "MiniMaxAI/MiniMax-H3@main", + "hf://MiniMaxAI/MiniMax-H3@", + "hf://Other/MiniMax-H3@revision", + "", + None, + ], +) +def test_resolve_native_vaes_rejects_unpinned_provenance(tmp_path, base_model): + model_root = tmp_path / "FastH3" + model_root.mkdir() + (model_root / "provenance.json").write_text(json.dumps({"base_model": base_model}), encoding="utf-8") + + with pytest.raises(ValueError, match="provenance"): + FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()).resolve_native_vaes(model_root) + + +def test_modular_local_model_root_uses_text_encoder_without_partition_files(tmp_path): + root = tmp_path / "FastH3" + (root / "text_encoder").mkdir(parents=True) + (root / "text_encoder" / "config.json").write_text("{}", encoding="utf-8") + (root / "modular_model_index.json").write_text( + json.dumps({"_class_name": "MiniMaxH3ModularPipeline"}), + encoding="utf-8", + ) + + assert is_minimax_h3_modular(str(root)) is True + assert resolve_minimax_h3_model_root(str(root), None, "fl2va") == str(root / "text_encoder") + assert not (root / "FL2VA").exists() + + +def test_modular_hub_model_root_downloads_only_text_encoder(monkeypatch, tmp_path): + from vllm_omni.model_executor.models.minimax_h3 import checkpoint as checkpoint_module + + snapshot = tmp_path / "modular-snapshot" + calls: dict[str, object] = {} + + def fake_index(model: str, *, revision: str | None = None) -> dict[str, str]: + assert model == "FastVideo/FastVideo-FastH3-8-Step-V2" + assert revision == "release-revision" + return {"_class_name": "MiniMaxH3ModularPipeline"} + + def fake_download(**kwargs: object) -> str: + calls.update(kwargs) + snapshot.mkdir() + return str(snapshot) + + monkeypatch.setattr(checkpoint_module, "get_diffusion_model_index", fake_index) + monkeypatch.setattr(checkpoint_module, "download_weights_from_hf_specific", fake_download) + + result = resolve_minimax_h3_model_root( + "FastVideo/FastVideo-FastH3-8-Step-V2", + "release-revision", + "fl2va", + ) + + assert result == str(snapshot / "text_encoder") + assert calls == { + "model_name_or_path": "FastVideo/FastVideo-FastH3-8-Step-V2", + "cache_dir": None, + "allow_patterns": ["text_encoder/**"], + "revision": "release-revision", + "require_all": True, + } + + +def test_native_hub_model_with_modular_class_keeps_partition_layout(monkeypatch): + from vllm_omni.model_executor.models.minimax_h3 import checkpoint as checkpoint_module + + monkeypatch.setattr( + checkpoint_module, + "get_diffusion_model_index", + lambda _model, *, revision=None: {"_class_name": "MiniMaxH3ModularPipeline"}, + ) + + assert checkpoint_module.is_minimax_h3_modular("MiniMaxAI/MiniMax-H3") is False + + +def test_native_local_model_root_keeps_partition_layout(tmp_path): + root = tmp_path / "MiniMax-H3" + (root / "FL2VA" / "text_encoder").mkdir(parents=True) + (root / "FL2VA" / "text_encoder" / "config.json").write_text("{}", encoding="utf-8") + + assert is_minimax_h3_modular(str(root)) is False + assert resolve_minimax_h3_model_root(str(root), None, "fl2va") == str(root / "FL2VA" / "text_encoder") + + +def test_diffusion_resolver_returns_local_modular_root_without_partition_files(tmp_path): + root = tmp_path / "FastH3" + root.mkdir() + (root / "modular_model_index.json").write_text( + json.dumps({"_class_name": "MiniMaxH3ModularPipeline"}), + encoding="utf-8", + ) + + assert resolve_minimax_h3_diffusion_model_path(str(root), None, "fl2va") == str(root) + assert not (root / "FL2VA").exists() + assert not (root / "Ref2VA").exists() + + +def test_diffusion_resolver_hub_modular_downloads_transformer_contract_only(monkeypatch, tmp_path): + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module + + snapshot = tmp_path / "modular-snapshot" + calls: dict[str, object] = {} + + def fake_modular(_model: str, _revision: str | None = None) -> bool: + return True + + def fake_download(**kwargs: object) -> str: + calls.update(kwargs) + snapshot.mkdir() + return str(snapshot) + + monkeypatch.setattr(pipeline_module, "is_minimax_h3_modular", fake_modular) + monkeypatch.setattr(pipeline_module, "download_weights_from_hf_specific", fake_download) + + result = resolve_minimax_h3_diffusion_model_path( + "FastVideo/FastVideo-FastH3-8-Step-V2", "release-revision", "fl2va" + ) + + assert result == str(snapshot) + assert calls == { + "model_name_or_path": "FastVideo/FastVideo-FastH3-8-Step-V2", + "cache_dir": None, + "allow_patterns": [ + "modular_model_index.json", + "fastvideo_inference.json", + "provenance.json", + "transformer/**", + ], + "revision": "release-revision", + "require_all": True, + } + patterns = calls["allow_patterns"] + assert isinstance(patterns, list) + assert all("video_vae" not in pattern and "audio_vae" not in pattern for pattern in patterns) + + +def test_diffusers_arch_config_normalizes_all_aliases(): + arch = MiniMaxH3DiTArchConfig.from_mapping( + { + "num_refiner_layers": 3, + "ffn_dim": 12, + "in_channels": 4, + "audio_in_channels": 6, + "freq_dim": 8, + "time_embed_hidden_dim": 10, + "rope_freq_dim": 5, + "patch_size": [1, 2, 2], + } + ) + + assert arch.token_refiner_num_layers == 3 + assert arch.ffn_hidden_size == 12 + assert arch.latents_dim == 4 + assert arch.audio_latents_dim == 6 + assert arch.timestep_input_dim == 8 + assert arch.time_embed_hidden_size == 10 + assert arch.rope_inv_freq_len == 5 + assert arch.patch_size == (1, 2, 2) + + +class _WeightTarget(nn.Module): + def __init__(self, shape: tuple[int, ...], loader: Callable[..., None]) -> None: + super().__init__() + self.weight = nn.Parameter(torch.zeros(shape)) + self.weight.weight_loader = loader + + +class _PipelineLoadTransformer(nn.Module): + def __init__(self, loaded: set[str]) -> None: + super().__init__() + self.arch = MiniMaxH3DiTArchConfig(num_layers=1) + self._loaded = loaded + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + list(weights) + return self._loaded + + def post_load_weights(self) -> None: + pass + + +def test_pipeline_load_weights_rejects_missing_fast_h3_compression_gate(): + from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + nn.Module.__init__(pipeline) + pipeline.transformer = _PipelineLoadTransformer({"blocks.0.attn.qkv_proj.weight"}) + pipeline.video_vae = None + pipeline.audio_vae = None + pipeline._fasth3 = None + pipeline._fasth3_checkpoint = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()) + + with pytest.raises(ValueError, match="missing compression gates"): + pipeline.load_weights([("transformer.blocks.0.attn.qkv_proj.weight", torch.zeros(1))]) + + +def _load_fixture( + *, diffusers_weights: bool +) -> tuple[MiniMaxH3DiTModel, list[tuple[str | None, torch.Tensor]], list[tuple[int, torch.Tensor]]]: + model = object.__new__(MiniMaxH3DiTModel) + nn.Module.__init__(model) + model.arch = MiniMaxH3DiTArchConfig( + num_layers=1, + hidden_size=3, + num_attention_heads=2, + attention_head_dim=2, + ffn_hidden_size=6, + rope_inv_freq_len=2, + ) + model._diffusers_weights = diffusers_weights + model._rope_theta = 10000.0 + qkv_calls: list[tuple[str | None, torch.Tensor]] = [] + fc1_calls: list[tuple[int, torch.Tensor]] = [] + + def qkv_loader(param: nn.Parameter, loaded_weight: torch.Tensor, shard_id: str | None = None) -> None: + if shard_id is None: + assert loaded_weight.shape == (12, 3) + param.data.copy_(loaded_weight) + qkv_calls.append((None, loaded_weight.clone())) + return + assert loaded_weight.shape == (4, 3) + assert shard_id in {"q", "k", "v"} + offset = {"q": 0, "k": 4, "v": 8}[shard_id] + param.data[offset : offset + 4].copy_(loaded_weight) + qkv_calls.append((shard_id, loaded_weight.clone())) + + def fc1_loader(param: nn.Parameter, loaded_weight: torch.Tensor, shard_id: int) -> None: + assert loaded_weight.shape == (3, 3) + param.data[shard_id * 3 : (shard_id + 1) * 3].copy_(loaded_weight) + fc1_calls.append((shard_id, loaded_weight.clone())) + + model.blocks = nn.ModuleList([nn.Module()]) + model.blocks[0].attn = nn.Module() + model.blocks[0].attn.qkv_proj = _WeightTarget((12, 3), qkv_loader) + model.blocks[0].mlp = nn.Module() + model.blocks[0].mlp.fc1 = _WeightTarget((6, 3), fc1_loader) + model.rope = nn.Module() + model.rope.register_buffer("inv_freq", torch.zeros(2)) + return model, qkv_calls, fc1_calls + + +def test_diffusers_load_weights_uses_shape_aware_qkv_and_mlp_slices(): + model, qkv_calls, fc1_calls = _load_fixture(diffusers_weights=True) + q = torch.arange(12, dtype=torch.float32).reshape(4, 3) + 10 + k = torch.arange(12, dtype=torch.float32).reshape(4, 3) + 100 + v = torch.arange(12, dtype=torch.float32).reshape(4, 3) + 200 + mlp = torch.arange(18, dtype=torch.float32).reshape(6, 3) + 300 + + loaded = model.load_weights( + [ + ("transformer_blocks.0.attn.to_q.weight", q), + ("transformer_blocks.0.attn.to_k.weight", k), + ("transformer_blocks.0.attn.to_v.weight", v), + ("transformer_blocks.0.ff.net.0.proj.weight", mlp), + ] + ) + + assert loaded == { + "blocks.0.attn.qkv_proj.weight", + "blocks.0.mlp.fc1.weight", + "rope.inv_freq", + } + assert [shard_id for shard_id, _ in qkv_calls] == ["q", "k", "v"] + assert [shard_id for shard_id, _ in fc1_calls] == [0, 1] + torch.testing.assert_close(model.blocks[0].attn.qkv_proj.weight, torch.cat((q, k, v))) + torch.testing.assert_close(model.blocks[0].mlp.fc1.weight, torch.cat((mlp[3:], mlp[:3]))) + torch.testing.assert_close( + model.rope.inv_freq, + 1.0 / (10000.0 ** (torch.arange(0, 4, 2, dtype=torch.float32) / 4)), + rtol=0, + atol=0, + ) + + +def test_native_load_weights_keeps_grouped_qkv_and_gate_first_layouts(): + model, qkv_calls, fc1_calls = _load_fixture(diffusers_weights=False) + grouped_qkv = torch.arange(36, dtype=torch.float32).reshape(12, 3) + fc1 = torch.arange(18, dtype=torch.float32).reshape(6, 3) + + loaded = model.load_weights( + [ + ("blocks.0.attn.qkv_proj.weight", grouped_qkv), + ("blocks.0.mlp.fc1.weight", fc1), + ] + ) + + expected_qkv = torch.cat( + ( + grouped_qkv[:2], + grouped_qkv[6:8], + grouped_qkv[2:4], + grouped_qkv[8:10], + grouped_qkv[4:6], + grouped_qkv[10:], + ) + ) + assert loaded == {"blocks.0.attn.qkv_proj.weight", "blocks.0.mlp.fc1.weight"} + assert [shard_id for shard_id, _ in qkv_calls] == [None] + assert [shard_id for shard_id, _ in fc1_calls] == [0, 1] + torch.testing.assert_close(model.blocks[0].attn.qkv_proj.weight, expected_qkv) + torch.testing.assert_close(model.blocks[0].mlp.fc1.weight, fc1) + assert torch.equal(model.rope.inv_freq, torch.zeros(2)) + + +def test_diffusers_load_weights_rejects_unknown_name(): + model, _, _ = _load_fixture(diffusers_weights=True) + + with pytest.raises(ValueError, match="unsupported Diffusers H3 weight"): + model.load_weights([("transformer_blocks.0.attn.not_a_projection.weight", torch.zeros(4, 3))]) + + +def test_diffusers_load_weights_rejects_duplicate_source_name(): + model, _, _ = _load_fixture(diffusers_weights=True) + weight = torch.zeros(4, 3) + + with pytest.raises(ValueError, match="duplicate Diffusers H3 weight"): + model.load_weights( + [ + ("transformer_blocks.0.attn.to_q.weight", weight), + ("transformer_blocks.0.attn.to_q.weight", weight), + ] + ) + + +def test_diffusers_load_weights_rejects_incomplete_qkv_group(): + model, _, _ = _load_fixture(diffusers_weights=True) + + with pytest.raises(ValueError, match="incomplete Diffusers H3 QKV"): + model.load_weights( + [ + ("transformer_blocks.0.attn.to_q.weight", torch.zeros(4, 3)), + ("transformer_blocks.0.attn.to_k.weight", torch.zeros(4, 3)), + ] + ) + + +def test_diffusers_load_weights_rejects_missing_compression_gate_parameter(): + model, _, _ = _load_fixture(diffusers_weights=True) + + with pytest.raises(ValueError, match="has no model parameter"): + model.load_weights([("transformer_blocks.0.attn.to_gate_compress.weight", torch.zeros(4, 3))]) + + +def test_resolve_sigma_positions_uses_release_metadata_schedule(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + nn.Module.__init__(pipeline) + pipeline._fasth3 = None + pipeline._lora_sigma_schedules = {} + spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()) + pipeline._base_schedule_by_partition = { + "fl2va": DMD2SigmaSchedule.from_metadata(spec.release_metadata()), + } + + positions, steps = pipeline._resolve_sigma_positions("t2va", _sampling()) + assert positions == FASTH3_V2_BASE_SCHEDULE.base_schedule + assert steps == 8 == len(positions) - 1 + with pytest.raises(OmniClientError, match="must be 8"): + pipeline._resolve_sigma_positions("t2va", _sampling(num_inference_steps=9)) 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 4b6fc0d88..cea4862d1 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -423,6 +423,7 @@ def test_pipeline_loads_task_selected_components_with_encoder_ownership( assert pipeline.partition == expected_partition assert pipeline.supported_tasks == expected_tasks assert len(created["dit"]) == expected_dits + assert all(kwargs["diffusers_weights"] is False for _, kwargs in created["dit"]) component_path = tmp_path / component_partition assert created["video_vae"] == [str(component_path / "video_vae")] assert created["audio_vae"] == [str(component_path / "audio_vae")] diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_fasth3.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_fasth3.py index 5c874ff02..c89c37db9 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_fasth3.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_fasth3.py @@ -262,7 +262,10 @@ def test_a_vsa_variant_is_recognised_by_its_compression_gates(tmp_path): def test_a_tensor_naming_no_h3_parameter_is_an_error(tmp_path): path = tmp_path / "fasth3" / "adapter_model.safetensors" - _write_adapter(path, tensors={"transformer_blocks.0.attn.norm_q.lora_A.weight": torch.ones((_RANK, _HIDDEN))}) + _write_adapter( + path, + tensors={"transformer_blocks.0.attn.not_a_projection.lora_A.weight": torch.ones((_RANK, _HIDDEN))}, + ) # Dropping it silently would load a model that is not the distilled student. with pytest.raises(FastH3AdapterError, match="name no known"): _load(path.parent) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py index 92bd09ac2..e1f83fd01 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from types import SimpleNamespace from unittest.mock import Mock @@ -41,7 +42,7 @@ class _FakeAttention(nn.Module): super().__init__() -def _small_od_config(): +def _small_od_config(*, class_name=None): arch = { "num_layers": 1, "token_refiner_num_layers": 1, @@ -60,6 +61,8 @@ def _small_od_config(): "final_adaln_out_features": 2 * 8, "rope_inv_freq_len": 2, } + if class_name is not None: + arch["_class_name"] = class_name return SimpleNamespace( tf_model_config=arch, parallel_config=SimpleNamespace(ulysses_degree=1), @@ -133,6 +136,24 @@ def test_fp8_scope_and_prefix_propagation(monkeypatch): ) +def test_explicit_native_weight_format_overrides_diffusers_class_name(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3 + + monkeypatch.setattr(h3, "ColumnParallelLinear", _FakeLinear) + monkeypatch.setattr(h3, "MergedColumnParallelLinear", _FakeLinear) + monkeypatch.setattr(h3, "QKVParallelLinear", _FakeLinear) + monkeypatch.setattr(h3, "RowParallelLinear", _FakeLinear) + monkeypatch.setattr(h3, "Attention", _FakeAttention) + monkeypatch.setattr(h3, "get_tensor_model_parallel_world_size", lambda: 1) + + model = h3.MiniMaxH3DiTModel( + _small_od_config(class_name="MiniMaxH3Transformer3DModel"), + diffusers_weights=False, + ) + + assert model._diffusers_weights is False + + class _WeightTarget(nn.Module): def __init__(self, loader): super().__init__() diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py index 87b04613a..3beda5acc 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py @@ -187,7 +187,11 @@ def test_minimax_h3_quantization_quality(config: QualityTestConfig): "model": model_ref, "enforce_eager": True, # The fused BF16 baseline only fits on H100-80GB with encoder TP. - "stage_0_tensor_parallel_size": 2, + # This test loads the direct FL2VA diffusion pipeline (not the + # disaggregated two-stage MiniMax-H3 pipeline), so configure both the + # DiT world and the Qwen3-VL encoder TP explicitly. + "tensor_parallel_size": 2, + "text_encoder_tp_size": 2, "vae_use_tiling": True, # MiniMax H3's VAE is implemented by code shipped with the checkpoint. "trust_remote_code": True, diff --git a/vllm_omni/diffusion/models/minimax_h3/fasth3.py b/vllm_omni/diffusion/models/minimax_h3/fasth3.py index c4860c57e..dcc513028 100644 --- a/vllm_omni/diffusion/models/minimax_h3/fasth3.py +++ b/vllm_omni/diffusion/models/minimax_h3/fasth3.py @@ -102,6 +102,7 @@ _MODEL_LEVEL_TARGETS = { "time_embedder.linear_2": "time_embedder.proj_out", "norm_out.linear": "final_layer.adaln_proj.linear", "norm_out.norm": "final_layer.norm", + "token_refiner.final_norm": "token_refiner.final_norm", } # Per-block adapter suffix -> (native suffix, how a delta enters the native @@ -115,6 +116,8 @@ _BLOCK_TARGETS = { "attn.to_v": ("attn.qkv_proj", _QKV_V), "attn.to_out.0": ("attn.out_proj", _PLAIN), "attn.to_gate_compress": ("attn.to_gate_compress", _PLAIN), + "attn.norm_q": ("attn.q_norm", _PLAIN), + "attn.norm_k": ("attn.k_norm", _PLAIN), "ff.net.0.proj": ("mlp.fc1", _SWAP_HALVES), "ff.net.2": ("mlp.fc2", _PLAIN), "adaln_proj.linear": ("adaln_proj.linear", _PLAIN), diff --git a/vllm_omni/diffusion/models/minimax_h3/fasth3_checkpoint.py b/vllm_omni/diffusion/models/minimax_h3/fasth3_checkpoint.py new file mode 100644 index 000000000..0382e9220 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/fasth3_checkpoint.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Sampling contract for FastVideo FastH3 full checkpoints.""" + +from __future__ import annotations + +import json +import math +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule +from vllm_omni.errors import OmniClientError +from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific + +from .fasth3 import FASTH3_BASE_MODEL, _resolve_dit_attention_backend + +if TYPE_CHECKING: + from vllm_omni.diffusion.data import OmniDiffusionConfig + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +FASTH3_V2_MODEL_ID = "FastVideo/FastVideo-FastH3-8-Step-V2" +FASTH3_V2_BASE_SCHEDULE = DMD2SigmaSchedule.from_positions((0.999, 0.874, 0.749, 0.624, 0.5, 0.375, 0.25, 0.125, 0.0)) + + +@dataclass(frozen=True) +class FastH3CheckpointSpec: + """The full V2 release has its own schedule and trained attention policy. + + This is independent of ``FastH3WeightFusion``: the released weights already + contain the entire student, including its learned compression gates. + """ + + vsa_sparsity: float = 0.8 + + @classmethod + def from_metadata(cls, metadata: Mapping[str, object]) -> FastH3CheckpointSpec: + """Validate the release's own fastvideo_inference.json.""" + required = { + "model_id": FASTH3_V2_MODEL_ID, + "schema_version": "fasth3-inference-contract-v1", + "dmd_denoising_steps": [999, 874, 749, 624, 500, 375, 250, 125], + "transformer_forwards": 8, + "video_scheduler_shift": 10.0, + "audio_scheduler_shift": 3.0, + "guidance_scale": 1.0, + "attention_backend": "VIDEO_SPARSE_ATTN_H3", + "vsa_sparsity": 0.8, + "vsa_tile_size": 64, + "task": "t2av", + } + for key, expected in required.items(): + if metadata.get(key) != expected: + raise ValueError(f"unsupported FastH3 contract: {key}={metadata.get(key)!r}, expected {expected!r}") + return cls() + + def release_metadata(self) -> dict[str, Any]: + """Express the sampling policy in the existing H3 pipeline schema.""" + return { + "partition": "fl2va", + "tasks": ["t2va"], + "sigma_shift_scales": {"video": 10.0, "audio": 3.0}, + "base_schedule": list(FASTH3_V2_BASE_SCHEDULE.base_schedule), + } + + def resolve_native_vaes(self, model_root: Path) -> Path: + """Reuse the frozen base VAEs with Omni's native tiled/parallel runtime. + + Only VAE components are fetched; the student and text encoder come + directly from the FastVideo release. The release pins the base revision. + """ + provenance = json.loads((model_root / "provenance.json").read_text(encoding="utf-8")) + base = provenance.get("base_model", "") + prefix = f"hf://{FASTH3_BASE_MODEL}@" + if not isinstance(base, str) or not base.startswith(prefix) or not base[len(prefix) :]: + raise ValueError("FastH3 V2 provenance must pin a MiniMaxAI/MiniMax-H3 base revision") + return ( + Path( + download_weights_from_hf_specific( + model_name_or_path=FASTH3_BASE_MODEL, + cache_dir=None, + allow_patterns=["FL2VA/video_vae/**", "FL2VA/audio_vae/**"], + revision=base[len(prefix) :], + require_all=True, + ) + ) + / "FL2VA" + ) + + def check_serving_contract(self, *, partition: str, od_config: OmniDiffusionConfig) -> None: + if partition != "fl2va": + raise ValueError("FastH3 V2 requires --task-type fl2va (T2VA requests only)") + if getattr(od_config, "lora_path", None): + raise ValueError("FastH3 V2 is a full checkpoint; additional LoRA adapters are unsupported") + if _resolve_dit_attention_backend(od_config) != "FASTVIDEO_VSA": + raise ValueError("FastH3 V2 requires --diffusion-attention-backend FASTVIDEO_VSA") + attention_config = getattr(od_config, "diffusion_attention_config", None) + per_role = getattr(attention_config, "per_role", None) or {} + spec = per_role.get("self") or getattr(attention_config, "default", None) + if getattr(spec, "fastvideo_vsa_topk", None) is not None: + raise ValueError("FastH3 V2 pins VSA sparsity=0.8; remove the fixed fastvideo_vsa_topk override") + parallel = getattr(od_config, "parallel_config", None) + if any(int(getattr(parallel, key, 1) or 1) != 1 for key in ("ring_degree", "allgather_degree")): + raise ValueError("FastH3 V2 supports local attention or pure Ulysses sequence parallelism") + + def check_request(self, sampling: OmniDiffusionSamplingParams, *, step_execution: bool = False) -> None: + if step_execution: + # StepScheduler admits requests before the pipeline hook runs and + # derives its lifetime from these fields. The pinned schedule must + # therefore be explicit and cannot be overridden by custom arrays. + if sampling.num_inference_steps != FASTH3_V2_BASE_SCHEDULE.num_inference_steps: + raise OmniClientError("FastH3 V2 step execution requires num_inference_steps=8") + if sampling.timesteps is not None or sampling.sigmas is not None: + raise OmniClientError("FastH3 V2 step execution does not support custom timesteps or sigmas") + if sampling.lora_request is not None: + raise OmniClientError("FastH3 V2 does not support per-request LoRA adapters") + steps = sampling.num_inference_steps + if steps is not None and steps != FASTH3_V2_BASE_SCHEDULE.num_inference_steps: + raise OmniClientError("FastH3 V2 requires num_inference_steps=8 (nine sigma points), or omitted") + extra = sampling.extra_args or {} + for key, expected in (("flow_shift", 10.0), ("audio_flow_shift", 3.0)): + try: + value = float(extra.get(key, expected)) + except (TypeError, ValueError) as exc: + raise OmniClientError(f"FastH3 V2 requires {key}={expected:g}") from exc + if not math.isclose(value, expected): + raise OmniClientError(f"FastH3 V2 requires {key}={expected:g}, got {value:g}") + if sampling.guidance_scale is not None and sampling.guidance_scale != 1.0: + raise OmniClientError("FastH3 V2 requires guidance_scale=1") 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 e26493931..5662a9641 100644 --- a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py +++ b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py @@ -55,6 +55,8 @@ from vllm_omni.diffusion.layers.rope import RotaryEmbedding from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract from vllm_omni.platforms import current_omni_platform +from .fasth3 import _resolve_native_target + if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, @@ -118,6 +120,18 @@ class MiniMaxH3DiTArchConfig: @classmethod def from_mapping(cls, config: Mapping[str, Any]) -> MiniMaxH3DiTArchConfig: + # The modular Diffusers checkpoint uses these spellings for the same + # architecture. Normalize before constructing the existing native DiT. + aliases = { + "num_refiner_layers": "token_refiner_num_layers", + "ffn_dim": "ffn_hidden_size", + "in_channels": "latents_dim", + "audio_in_channels": "audio_latents_dim", + "freq_dim": "timestep_input_dim", + "time_embed_hidden_dim": "time_embed_hidden_size", + "rope_freq_dim": "rope_inv_freq_len", + } + config = {aliases.get(name, name): value for name, value in config.items()} fields = cls.__dataclass_fields__ values = {name: config[name] for name in fields if name in config} if "patch_size" in values: @@ -394,6 +408,9 @@ def _sdpa_varlen_attention( class MiniMaxH3Attention(nn.Module): + # Full sparse checkpoints pin a ratio; legacy adapters use backend top-k. + vsa_sparsity: float | None = None + def __init__( self, arch: MiniMaxH3DiTArchConfig, @@ -588,6 +605,7 @@ class MiniMaxH3Attention(nn.Module): # (see MINIMAX_H3_LASER_INPUT_SCALE). Ignored by every other # backend/path. "laser_input_scale": MINIMAX_H3_LASER_INPUT_SCALE, + **({"vsa_h3_sparsity": self.vsa_sparsity} if self.vsa_sparsity is not None else {}), # Present only for a VSA artifact; the VSA backend reads it as # the learned compression gate and every other backend ignores it. **({"gate_compress": gate_compress.unsqueeze(0)} if gate_compress is not None else {}), @@ -1127,10 +1145,23 @@ class MiniMaxH3DiTModel(nn.Module): self, od_config: OmniDiffusionConfig, quant_config: QuantizationConfig | None = None, + *, + diffusers_weights: bool | None = None, ) -> None: super().__init__() tf_config = od_config.tf_model_config config_mapping = tf_config.to_dict() if hasattr(tf_config, "to_dict") else dict(tf_config) + # The native MiniMax-H3 Hub snapshot advertises the Diffusers + # transformer class in its root config, while its FL2VA/Ref2VA + # components still contain native weights. The pipeline has already + # resolved the actual source format, so let it override the + # class-name heuristic. Keep the heuristic for standalone callers. + self._diffusers_weights = ( + config_mapping.get("_class_name") == "MiniMaxH3Transformer3DModel" + if diffusers_weights is None + else diffusers_weights + ) + self._rope_theta = float(config_mapping.get("rope_theta", 10000.0)) arch = MiniMaxH3DiTArchConfig.from_mapping(config_mapping) self.arch = arch self.od_config = od_config @@ -1211,7 +1242,7 @@ class MiniMaxH3DiTModel(nn.Module): ) self._mark_missing_params_required() - def enable_vsa_gates(self) -> None: + def enable_vsa_gates(self, *, sparsity: float | None = None) -> None: """Give every DiT block's attention a VSA compression gate. A FastH3 VSA artifact assigns these projections rather than adding to @@ -1223,6 +1254,9 @@ class MiniMaxH3DiTModel(nn.Module): return for block in self.blocks: block.attn.enable_vsa_gate() + block.attn.vsa_sparsity = sparsity + if sparsity is not None: + block.attn.to_gate_compress.weight.missing_param_init = "error" self.vsa_gates_enabled = True def _mark_missing_params_required(self) -> None: @@ -1293,17 +1327,37 @@ class MiniMaxH3DiTModel(nn.Module): self, weights: Iterable[tuple[str, torch.Tensor]], ) -> set[str]: - """Load exact H3 checkpoint names with logical TP-aware loaders.""" + """Load native or Diffusers H3 weights with the existing TP-aware loaders.""" params = dict(self.named_parameters()) params.update(dict(self.named_buffers())) loaded: set[str] = set() + diffusers_weights = getattr(self, "_diffusers_weights", False) + qkv_parts: dict[str, set[str]] = {} + source_names: set[str] = set() for name, loaded_weight in weights: + layout = "plain" + if diffusers_weights: + if name in source_names: + raise ValueError(f"duplicate Diffusers H3 weight: {name}") + source_names.add(name) + module, _, kind = name.rpartition(".") + target = _resolve_native_target(module) + if target is None or kind not in {"weight", "bias"}: + raise ValueError(f"unsupported Diffusers H3 weight: {name}") + name, layout = f"{target[0]}.{kind}", target[1] param = params.get(name) if param is None: + if diffusers_weights: + raise ValueError(f"Diffusers H3 weight has no model parameter: {name}") logger.warning("Skipping MiniMax H3 weight not present in model: %s", name) continue weight_loader = getattr(param, "weight_loader", default_weight_loader) - if name.endswith(".attn.qkv_proj.weight"): + if layout in {"q", "k", "v"}: + # vLLM can load each projection directly into its packed QKV + # parameter, including TP slicing and online quantization. + weight_loader(param, loaded_weight, layout) + qkv_parts.setdefault(name, set()).add(layout) + elif name.endswith(".attn.qkv_proj.weight"): # Transform checkpoint layout before entering vLLM's loader so # online FP8 can keep ``online_process_loader`` outermost. loaded_weight = _reorder_grouped_qkv_to_qkv( @@ -1319,12 +1373,26 @@ class MiniMaxH3DiTModel(nn.Module): "MiniMax H3 fc1 checkpoint rows must split evenly into " f"gate/up matrices, got {tuple(loaded_weight.shape)}" ) - gate, up = loaded_weight.chunk(2, dim=0) + first, second = loaded_weight.chunk(2, dim=0) + gate, up = (second, first) if layout == "swap_halves" else (first, second) weight_loader(param, gate, 0) weight_loader(param, up, 1) else: weight_loader(param, loaded_weight) loaded.add(name) + if diffusers_weights: + for name, parts in qkv_parts.items(): + if parts != {"q", "k", "v"}: + raise ValueError(f"incomplete Diffusers H3 QKV group {name}: {sorted(parts)}") + # Diffusers reconstructs RoPE from config instead of storing this + # native checkpoint buffer. Compute on CPU for identical values. + freq_dim = self.arch.rope_inv_freq_len + rope = 1.0 / ( + self._rope_theta + ** (torch.arange(0, 2 * freq_dim, 2, dtype=torch.float32, device="cpu") / (2 * freq_dim)) + ) + default_weight_loader(params["rope.inv_freq"], rope) + loaded.add("rope.inv_freq") return loaded @staticmethod 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 1b8bd7fea..60d6f4125 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -67,6 +67,7 @@ from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, ) from vllm_omni.model_executor.models.minimax_h3.checkpoint import ( + is_minimax_h3_modular, resolve_minimax_h3_partition, ) from vllm_omni.model_executor.models.minimax_h3.conditioning import ( @@ -106,6 +107,7 @@ from .denoise_loop import ( ) from .encoder import MiniMaxH3Qwen3VLEncoder from .fasth3 import FastH3WeightFusion, resolve_fasth3_fusion +from .fasth3_checkpoint import FastH3CheckpointSpec from .lora import TurboSpec, load_minimax_h3_turbo_lora from .minimax_h3_transformer import ( MiniMaxH3Attention, @@ -213,7 +215,11 @@ def _resolve_minimax_h3_model_root( if path.name in {"FL2VA", "Ref2VA"}: return path.parent return path - if load_text_encoder: + if is_minimax_h3_modular(model, revision): + allow_patterns = ["modular_model_index.json", "fastvideo_inference.json", "provenance.json", "transformer/**"] + if load_text_encoder: + allow_patterns += ["text_encoder/**", "tokenizer/**", "processor/**"] + elif load_text_encoder: allow_patterns = ( MINIMAX_H3_DOWNLOAD_PATTERNS if partition == "combined" else MINIMAX_H3_TASK_DOWNLOAD_PATTERNS[partition] ) @@ -314,6 +320,8 @@ def resolve_minimax_h3_diffusion_model_path( partition, load_text_encoder=False, ) + if is_minimax_h3_modular(str(model_root), revision): + return str(model_root) if partition == "combined": return str(model_root) subdir = "Ref2VA" if partition == "ref2va" else "FL2VA" @@ -575,6 +583,7 @@ class MiniMaxH3Pipeline( _base_schedule_by_partition: ClassVar[Mapping[str, DMD2SigmaSchedule | None]] = {} # Set from --lora-path during construction; absent means no FastH3 adapter. _fasth3: FastH3WeightFusion | None = None + _fasth3_checkpoint: FastH3CheckpointSpec | None = None def _load_diffusion_lora_adapter( self, @@ -800,10 +809,13 @@ class MiniMaxH3Pipeline( self._PROFILER_TARGETS.remove("encode_prompt") if not self.load_vae_encoder: self._PROFILER_TARGETS.remove("_encode_local_media") + modular = is_minimax_h3_modular(str(od_config.model), od_config.revision) self.partition = _minimax_h3_partition_for_task( getattr(od_config, "task_type", None), str(od_config.model), ) + if modular and str(od_config.task_type or "auto").lower() == "auto": + self.partition = "fl2va" self._turbo_lora_specs: dict[int, TurboSpec] = {} self._native_lora_adapter_ids: set[int] = set() self._lora_sigma_schedules: dict[int, DMD2SigmaSchedule] = {} @@ -813,9 +825,19 @@ class MiniMaxH3Pipeline( self.partition, load_text_encoder=self.load_text_encoder, ) - model_path = model_root / ("Ref2VA" if self.partition == "ref2va" else "FL2VA") - model_index = json.loads((model_path / "model_index.json").read_text(encoding="utf-8")) - release = model_index.get("_minimax_h3") or {} + if modular: + model_path = model_root + self._fasth3_checkpoint = FastH3CheckpointSpec.from_metadata( + json.loads((model_root / "fastvideo_inference.json").read_text(encoding="utf-8")) + ) + self._fasth3_checkpoint.check_serving_contract(partition=self.partition, od_config=od_config) + release = self._fasth3_checkpoint.release_metadata() + vae_model_path = self._fasth3_checkpoint.resolve_native_vaes(model_root) + else: + model_path = model_root / ("Ref2VA" if self.partition == "ref2va" else "FL2VA") + model_index = json.loads((model_path / "model_index.json").read_text(encoding="utf-8")) + release = model_index.get("_minimax_h3") or {} + vae_model_path = model_path partition = str(release.get("partition", "")).lower() expected_partition = "ref2va" if self.partition == "ref2va" else "fl2va" if partition != expected_partition: @@ -876,11 +898,18 @@ class MiniMaxH3Pipeline( self.transformer = MiniMaxH3DiTModel( od_config, quant_config=transformer_quant_config, + diffusers_weights=modular, ) + if self._fasth3_checkpoint is not None: + self.transformer.enable_vsa_gates(sparsity=self._fasth3_checkpoint.vsa_sparsity) + logger.info( + "FastH3 V2 full checkpoint: 8 transformer forwards, video/audio shifts 10/3, VSA sparsity=0.8 tile=64" + ) if ref2va_model_path is not None: self.transformers_ref = MiniMaxH3DiTModel( od_config, quant_config=transformer_quant_config, + diffusers_weights=modular, ) self._fasth3 = resolve_fasth3_fusion(od_config, self.transformer) @@ -963,14 +992,14 @@ class MiniMaxH3Pipeline( # so VAEs stay resident for new configurations. component_load_device = torch.device("cpu") if legacy_manual_components else self.device self.video_vae = MiniMaxH3VideoVAE( - os.path.join(model_path, "video_vae"), + os.path.join(vae_model_path, "video_vae"), device=self.device, load_device=component_load_device, decode_only=not self.load_vae_encoder, trust_remote_code=od_config.trust_remote_code, ) self.audio_vae = MiniMaxH3AudioVAE( - os.path.join(model_path, "audio_vae"), + os.path.join(vae_model_path, "audio_vae"), device=self.device, load_device=component_load_device, decode_only=not self.load_vae_encoder, @@ -1046,6 +1075,12 @@ class MiniMaxH3Pipeline( # load_weights only warns on a parameter the model does not have, so # close the adapter against what the DiT actually consumed. self._fasth3.validate_fully_applied(transformer_loaded) + if self._fasth3_checkpoint is not None: + required_gates = { + f"blocks.{i}.attn.to_gate_compress.weight" for i in range(self.transformer.arch.num_layers) + } + if missing_gates := required_gates - transformer_loaded: + raise ValueError(f"FastH3 V2 checkpoint is missing compression gates: {sorted(missing_gates)}") return loaded_with_prefix @property @@ -2056,6 +2091,11 @@ class MiniMaxH3Pipeline( self._validate_turbo_sampling(sampling, turbo_spec) if has_native_lora: self._validate_native_sampling(sampling, task=task) + if self._fasth3_checkpoint is not None: + self._fasth3_checkpoint.check_request( + sampling, + step_execution=bool(getattr(self.od_config, "step_execution", False)), + ) if self._fasth3 is not None: self._fasth3.check_request( sampling, diff --git a/vllm_omni/model_executor/models/minimax_h3/checkpoint.py b/vllm_omni/model_executor/models/minimax_h3/checkpoint.py index 0143e9675..b39af5cf8 100644 --- a/vllm_omni/model_executor/models/minimax_h3/checkpoint.py +++ b/vllm_omni/model_executor/models/minimax_h3/checkpoint.py @@ -1,10 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""MiniMax H3 checkpoint resolution for the AR text-encoder stage.""" + from pathlib import Path +from vllm_omni.diffusion.utils.hf_utils import get_diffusion_model_index from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, ) +_FASTH3_MODULAR_MODEL_PREFIX = "FastVideo/" + MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS = { partition: [ f"{subdir}/text_encoder/**", @@ -16,6 +22,22 @@ MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS = { MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS["combined"] = MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS["fl2va"] +def is_minimax_h3_modular(model: str, revision: str | None = None) -> bool: + path = Path(model) + if path.is_dir(): + # Native MiniMax-H3 snapshots can advertise the Diffusers modular + # class in ``model_index.json`` while retaining the legacy FL2VA / + # Ref2VA layout. FastH3 modular releases carry this marker file. + return (path / "modular_model_index.json").is_file() or (path / "fastvideo_inference.json").is_file() + # The FastVideo release is the modular checkpoint handled by this + # resolver. Do not classify the native MiniMax-H3 Hub repository from + # its class name alone; its partition files are still required here. + if not model.startswith(_FASTH3_MODULAR_MODEL_PREFIX): + return False + index = get_diffusion_model_index(model, revision=revision) or {} + return index.get("_class_name") == "MiniMaxH3ModularPipeline" + + def resolve_minimax_h3_partition( model: str, task_type: str | None, @@ -42,10 +64,13 @@ def resolve_minimax_h3_encoder_model_root( ) -> str: path = Path(model) partition = resolve_minimax_h3_partition(model, task_type, auto_partition="fl2va") + modular = is_minimax_h3_modular(model, revision) if path.is_dir(): if path.name == "text_encoder" and (path / "config.json").is_file(): return str(path) + if modular: + return str(path / "text_encoder") if path.name in {"FL2VA", "Ref2VA"}: path = path.parent subdir = "Ref2VA" if partition == "ref2va" else "FL2VA" @@ -53,11 +78,13 @@ def resolve_minimax_h3_encoder_model_root( snapshot = download_weights_from_hf_specific( model_name_or_path=model, cache_dir=None, - allow_patterns=MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS[partition], + allow_patterns=["text_encoder/**"] if modular else MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS[partition], revision=revision, require_all=True, ) + if modular: + return str(Path(snapshot) / "text_encoder") subdir = "Ref2VA" if partition == "ref2va" else "FL2VA" return str(Path(snapshot) / subdir / "text_encoder")