[Diffusion][Model] Add direct FastH3 8-Step V2 checkpoint loading (#7610)

Signed-off-by: david6666666 <530634352@qq.com>
This commit is contained in:
WeiQing Chen authored and GitHub committed 2026-09-17 08:41:52 +00:00
1 parent 8c3d3407d1
commit cb439f3e84
11 files changed
+1053 -14

No files matched your search

+62
View File
@@ -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 |
@@ -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))
@@ -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")]
@@ -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)
@@ -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__()
@@ -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,
@@ -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),
@@ -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")
@@ -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
@@ -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,
@@ -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")