[Diffusion][Model] Add direct FastH3 8-Step V2 checkpoint loading (#7610)
Signed-off-by: david6666666 <530634352@qq.com>
This commit is contained in:
1 parent
8c3d3407d1
commit
cb439f3e84
11 files changed
+1053
-14
No files matched your search
@@ -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")
|
||||
|
||||
|
||||
Reference in new issue
Block a user