[Feature][MiniMax-H3] Fuse the FastH3 four-step adapter at load time (#6714)
Signed-off-by: princepride <wangzhipeng628@gmail.com> Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
汪志鹏andHongsheng Liu authored and GitHub committed
1 parent
5e06d2e73c
commit
86b85c078b
13 files changed
+1420
-54
No files matched your search
@@ -918,6 +918,72 @@ adapter and twice without it, then compare the four output digests. The adapter
|
||||
is bound and deterministic when each pair matches internally and the two pairs
|
||||
differ from each other.
|
||||
|
||||
### FastH3 adapter
|
||||
|
||||
[FastH3](https://haoailab.com/blogs/fasth3-preview/) is FastVideo's four-step
|
||||
DMD2 student of H3-Base. It reuses H3's text encoder, VAEs, tokenizers, and
|
||||
schedulers unchanged, replacing 49 denoiser evaluations with four. It is fused
|
||||
into the checkpoint at load time rather than switched per request, because it
|
||||
carries full-rank deltas that no LoRA layer can express.
|
||||
|
||||
The bundle publishes four variants, so download one and point `--lora-path` at
|
||||
it; the repository root is refused rather than guessed at:
|
||||
|
||||
```bash
|
||||
export FASTH3_DIR=/path/to/fasth3
|
||||
hf download FastVideo/FastVideo-FastH3-4-step-Preview-v1-LoRA \
|
||||
dense-datafree/adapter_model.safetensors --local-dir "${FASTH3_DIR}"
|
||||
export FASTH3_LORA="${FASTH3_DIR}/dense-datafree/adapter_model.safetensors"
|
||||
```
|
||||
|
||||
Add `--task-type fl2va --lora-path "${FASTH3_LORA}"` to a non-offloaded server
|
||||
command. T2VA is served by the FL2VA partition, so `--task-type fl2va` is
|
||||
correct even though FastH3 preview v1 distills T2VA only. Because the adapter is
|
||||
fused, `--lora-backend` does not apply and a request carrying a `lora=` field is
|
||||
rejected rather than served without the adapter it asked for.
|
||||
|
||||
```bash
|
||||
-F 'num_inference_steps=4' \
|
||||
-F 'extra_params={"task":"t2va","duration":4.4}'
|
||||
```
|
||||
|
||||
Requests must ask for `num_inference_steps=4` and `task=t2va`: the release's five
|
||||
sigma points bound four denoiser evaluations, and that count is what the step
|
||||
scheduler admits a request on. The server denoises on the release's own ladder,
|
||||
keeping H3's per-modality shifts at the checkpoint values, so a request that
|
||||
overrides `flow_shift` or `audio_flow_shift` is rejected - it would sample the
|
||||
student at noise levels it was never distilled at.
|
||||
|
||||
Only a release that identifies itself as FastH3 is fused; any other
|
||||
`fastvideo-lora-v2` adapter stays on the dynamic LoRA route. A claimed artifact
|
||||
is then held to its own metadata: one that misdeclares its tensor counts or
|
||||
leaves a transformer block unedited is refused at startup instead of serving
|
||||
mostly base H3 weights on a four-step schedule. Offload is refused for the same
|
||||
reason - `--enable-cpu-offload`, `--enable-layerwise-offload` and
|
||||
`--enable-distributed-layerwise-offload` all bypass the fusion, so they fail fast.
|
||||
|
||||
> [!NOTE]
|
||||
> Only the **Dense / Data-Free** variant is supported today. The three VSA
|
||||
> variants carry `.set_weight` compression gates for modules that do not exist
|
||||
> in vLLM-Omni's dense H3 attention, and the release states dense attention is
|
||||
> not a supported substitute for them, so they are refused at startup rather
|
||||
> than silently run wrong. Video Sparse Attention for H3's packed
|
||||
> `[text | cond | audio | video]` sequence is follow-up work.
|
||||
|
||||
Measured on 8x NVIDIA B300 with USP8, VAE patch-parallel 8, `TRTLLM_ATTN`, at
|
||||
1344x768, 4.4 s, seed 1101, one warmup excluded and two runs recorded:
|
||||
|
||||
| Adapter | Steps | End-to-end | Diffusion engine |
|
||||
| --- | ---: | ---: | ---: |
|
||||
| none (base H3) | 50 | 25.8 / 26.4 s | 16.22 / 16.28 s |
|
||||
| FastH3 Dense | 4 (5 sigma points) | 11.7 / 11.8 s | 2.37 / 2.36 s |
|
||||
|
||||
The denoising speedup is 6.9x. End-to-end is 2.2x because text encoding, VAE
|
||||
decoding and muxing are a fixed cost that dominates a clip this short; longer
|
||||
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.
|
||||
|
||||
## Key parameters
|
||||
|
||||
| Parameter | Recommended value | Notes |
|
||||
|
||||
@@ -43,21 +43,16 @@ def test_jit_build_includes_cuda_headers_from_nvidia_wheels(tmp_path, monkeypatc
|
||||
|
||||
@dataclass
|
||||
class _FakeAllocation:
|
||||
"""Byte workspace stand-in that slices and reinterprets like ``symm_mem.empty()``."""
|
||||
|
||||
size: int
|
||||
device: torch.device
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self.buffer = torch.empty(self.size, dtype=torch.uint8)
|
||||
|
||||
class _FakeHandle:
|
||||
rank = 0
|
||||
|
||||
def __init__(self, allocation: _FakeAllocation) -> None:
|
||||
self.allocation = allocation
|
||||
self.views: list[tuple[tuple[int, ...], torch.dtype]] = []
|
||||
|
||||
def get_buffer(self, rank, shape, dtype):
|
||||
assert rank == self.rank
|
||||
self.views.append((shape, dtype))
|
||||
return self.allocation, shape, dtype
|
||||
def __getitem__(self, item: slice) -> torch.Tensor:
|
||||
return self.buffer[item]
|
||||
|
||||
|
||||
def test_workspace_reuses_peak_capacity_across_shapes(monkeypatch) -> None:
|
||||
@@ -93,16 +88,20 @@ def test_workspace_reuses_peak_capacity_across_shapes(monkeypatch) -> None:
|
||||
return allocation
|
||||
|
||||
monkeypatch.setattr(a2a_permute.symm_mem, "empty", empty)
|
||||
monkeypatch.setattr(a2a_permute.symm_mem, "rendezvous", lambda allocation, _group: _FakeHandle(allocation))
|
||||
# The handle is only retained to keep the rendezvous registration alive.
|
||||
monkeypatch.setattr(a2a_permute.symm_mem, "rendezvous", lambda _allocation, _group: object())
|
||||
|
||||
device = torch.device("cuda:0")
|
||||
first = a2a_permute._get_symm_buffer((2, 3), torch.float16, device, "group")
|
||||
smaller = a2a_permute._get_symm_buffer((1, 4), torch.float16, device, "group")
|
||||
larger = a2a_permute._get_symm_buffer((4, 3), torch.float16, device, "group")
|
||||
|
||||
assert first[0].size == 12
|
||||
assert smaller[0] is first[0]
|
||||
assert larger[0].size == 24
|
||||
assert (first.shape, first.dtype) == ((2, 3), torch.float16)
|
||||
assert (smaller.shape, smaller.dtype) == ((1, 4), torch.float16)
|
||||
assert (larger.shape, larger.dtype) == ((4, 3), torch.float16)
|
||||
# A sub-capacity request reuses the workspace; only growth allocates again.
|
||||
assert smaller.untyped_storage().data_ptr() == first.untyped_storage().data_ptr()
|
||||
assert larger.untyped_storage().data_ptr() != first.untyped_storage().data_ptr()
|
||||
assert [allocation.size for allocation in allocations] == [12, 24]
|
||||
assert len(a2a_permute._SYMM_WORKSPACES) == 1
|
||||
assert len(all_reduce_calls) == 1
|
||||
|
||||
@@ -496,13 +496,16 @@ def test_joint_postprocess_is_multiprocessing_picklable():
|
||||
|
||||
postprocess = get_minimax_h3_post_process_func(SimpleNamespace())
|
||||
postprocess = ForkingPickler.loads(ForkingPickler.dumps(postprocess))
|
||||
video = torch.linspace(0, 1, 2 * 3 * 2 * 4 * 5).reshape(2, 3, 2, 4, 5)
|
||||
# decode() hands over quantised frames already laid out as (B, T, H, W, C).
|
||||
video = torch.arange(2 * 2 * 4 * 5 * 3, dtype=torch.uint8).reshape(2, 2, 4, 5, 3)
|
||||
audio = torch.arange(12, dtype=torch.float32).reshape(1, 2, 6)
|
||||
|
||||
result = postprocess((video, audio), output_type="np")
|
||||
|
||||
assert isinstance(result["video"], list)
|
||||
assert result["video"][0].dtype == np.uint8
|
||||
assert result["video"][0].shape == (2, 4, 5, 3)
|
||||
np.testing.assert_array_equal(result["video"][0], video[0].numpy())
|
||||
np.testing.assert_array_equal(result["audio"], audio.numpy())
|
||||
assert result["audio_sample_rate"] == 32000
|
||||
assert result["fps"] == 24
|
||||
@@ -595,7 +598,7 @@ def test_base_schedule_overrides_the_uniform_sigma_positions():
|
||||
@pytest.mark.parametrize(
|
||||
"base_schedule",
|
||||
[
|
||||
[0.9, 0.5, 0.0],
|
||||
[1.5, 0.5, 0.0],
|
||||
[1.0, 0.5, 0.1],
|
||||
[1.0, 0.5, 0.5, 0.0],
|
||||
[1.0],
|
||||
|
||||
@@ -0,0 +1,498 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from safetensors.torch import save_file
|
||||
|
||||
from vllm_omni.diffusion.models.minimax_h3.fasth3 import (
|
||||
FASTH3_BASE_MODEL,
|
||||
FASTH3_BASE_SCHEDULE,
|
||||
FASTH3_DENOISE_STEPS,
|
||||
FASTH3_FORMAT,
|
||||
FastH3AdapterError,
|
||||
FastH3WeightFusion,
|
||||
resolve_fasth3_fusion,
|
||||
)
|
||||
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import _reorder_grouped_qkv_to_qkv
|
||||
from vllm_omni.diffusion.models.minimax_h3.time_request import minimax_h3_time_shift_sigmas
|
||||
from vllm_omni.errors import OmniClientError
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
_RANK = 2
|
||||
_HIDDEN = 4
|
||||
_HEAD_DIM = 2
|
||||
_HEADS = 3
|
||||
_INNER = _HEAD_DIM * _HEADS # attention inner size
|
||||
_FFN = 5
|
||||
|
||||
|
||||
def _factors(out_dim: int, in_dim: int) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""A rank-2 pair whose product is reproducible and non-symmetric."""
|
||||
a = torch.arange(_RANK * in_dim, dtype=torch.float32).reshape(_RANK, in_dim) / (in_dim * _RANK)
|
||||
b = torch.arange(out_dim * _RANK, dtype=torch.float32).reshape(out_dim, _RANK) / (out_dim * _RANK)
|
||||
return a, b
|
||||
|
||||
|
||||
def _write_adapter(
|
||||
path,
|
||||
*,
|
||||
tensors=None,
|
||||
drop: str | None = None,
|
||||
blocks: int = 1,
|
||||
metadata: dict[str, str | None] | None = None,
|
||||
) -> None:
|
||||
"""Write an artifact in the published ``fastvideo-lora-v2`` shape."""
|
||||
payload: dict[str, torch.Tensor] = {}
|
||||
for block in range(blocks):
|
||||
for suffix, (out_dim, in_dim) in {
|
||||
"attn.to_q": (_INNER, _HIDDEN),
|
||||
"attn.to_k": (_INNER, _HIDDEN),
|
||||
"attn.to_v": (_INNER, _HIDDEN),
|
||||
"attn.to_out.0": (_HIDDEN, _INNER),
|
||||
"ff.net.0.proj": (2 * _FFN, _HIDDEN),
|
||||
"ff.net.2": (_HIDDEN, _FFN),
|
||||
}.items():
|
||||
a, b = _factors(out_dim, in_dim)
|
||||
payload[f"transformer_blocks.{block}.{suffix}.lora_A.weight"] = a
|
||||
payload[f"transformer_blocks.{block}.{suffix}.lora_B.weight"] = b
|
||||
payload.update(tensors or {})
|
||||
if drop is not None:
|
||||
payload.pop(drop, None)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
# The published writer records what it emitted; mirror it so the fixtures
|
||||
# are as self-describing as the real artifact.
|
||||
declared = {
|
||||
"format": FASTH3_FORMAT,
|
||||
"finetuned_model": "FastVideo/FastVideo-FastH3-Dense-4-step-v1",
|
||||
"base_model": FASTH3_BASE_MODEL,
|
||||
"rank": str(_RANK),
|
||||
"low_rank_tensors": str(sum(1 for key in payload if key.endswith((".lora_A.weight", ".lora_B.weight")))),
|
||||
"diff_tensors": str(sum(1 for key in payload if key.endswith((".diff", ".diff_b")))),
|
||||
"set_weight_tensors": str(sum(1 for key in payload if key.endswith(".set_weight"))),
|
||||
}
|
||||
# safetensors metadata is string-only, so a ``None`` override drops the key.
|
||||
declared.update(metadata or {})
|
||||
save_file(payload, str(path), metadata={key: value for key, value in declared.items() if value is not None})
|
||||
|
||||
|
||||
def _claim(path, **kwargs) -> FastH3WeightFusion | None:
|
||||
"""Whatever from_path decides, including declining the artifact."""
|
||||
kwargs.setdefault("head_dim", _HEAD_DIM)
|
||||
kwargs.setdefault("num_blocks", 1)
|
||||
kwargs.setdefault("num_refiner_blocks", 0)
|
||||
return FastH3WeightFusion.from_path(path, **kwargs)
|
||||
|
||||
|
||||
def _load(path, **kwargs) -> FastH3WeightFusion:
|
||||
"""The same, for the tests that only make sense on a claimed adapter."""
|
||||
fusion = _claim(path, **kwargs)
|
||||
assert fusion is not None, f"{path} was not claimed as a FastH3 adapter"
|
||||
return fusion
|
||||
|
||||
|
||||
def test_only_an_artifact_carrying_the_release_identity_is_claimed(tmp_path):
|
||||
plain = tmp_path / "peft" / "adapter_model.safetensors"
|
||||
plain.parent.mkdir(parents=True)
|
||||
save_file({"transformer_blocks.0.attn.to_q.lora_A.weight": torch.ones((_RANK, _HIDDEN))}, str(plain))
|
||||
# No fastvideo-lora-v2 metadata: this is somebody else's LoRA and has to
|
||||
# stay on the dynamic route rather than being fused.
|
||||
assert _claim(plain.parent) is None
|
||||
assert _claim(tmp_path / "missing") is None
|
||||
|
||||
claimed = tmp_path / "fasth3" / "adapter_model.safetensors"
|
||||
_write_adapter(claimed)
|
||||
assert _claim(claimed.parent) is not None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"identity",
|
||||
[
|
||||
# An ordinary H3 adapter out of FastVideo's own extraction tools.
|
||||
{"finetuned_model": "someone/minimax-h3-style-lora"},
|
||||
# Somebody else's adapter that merely names the student it imitates.
|
||||
{"finetuned_model": "someone/fasth3-style-lora-for-h3"},
|
||||
{"finetuned_model": ""},
|
||||
{"finetuned_model": None},
|
||||
# A FastH3 name over a different base model.
|
||||
{"base_model": "Wan-AI/Wan2.2-TI2V-5B"},
|
||||
],
|
||||
)
|
||||
def test_a_generic_fastvideo_h3_adapter_stays_on_the_dynamic_route(identity, tmp_path):
|
||||
"""The fixture edits every block, so only the identity separates it from FastH3."""
|
||||
path = tmp_path / "generic" / "adapter_model.safetensors"
|
||||
_write_adapter(path, metadata=identity)
|
||||
|
||||
assert _claim(path.parent) is None
|
||||
|
||||
|
||||
def test_a_claimed_artifact_must_declare_every_tensor_count(tmp_path):
|
||||
"""Declining is for other people's adapters; a FastH3 file is held to its word."""
|
||||
path = tmp_path / "fasth3" / "adapter_model.safetensors"
|
||||
_write_adapter(path, metadata={"diff_tensors": None})
|
||||
|
||||
with pytest.raises(FastH3AdapterError, match="omits diff_tensors"):
|
||||
_claim(path.parent)
|
||||
|
||||
|
||||
def test_the_published_bundle_root_is_refused_rather_than_guessed(tmp_path):
|
||||
root = tmp_path / "FastVideo-FastH3-4-step-Preview-v1-LoRA"
|
||||
for slug in ("dense-datafree", "vsa-datafree"):
|
||||
_write_adapter(root / slug / "adapter_model.safetensors")
|
||||
(root / "adapter_manifest.json").write_text(json.dumps({"schema_version": "fasth3-lora-bundle-v1"}))
|
||||
|
||||
with pytest.raises(FastH3AdapterError, match="point --lora-path at one variant"):
|
||||
_load(root)
|
||||
# One variant inside it is unambiguous.
|
||||
assert _claim(root / "dense-datafree") is not None
|
||||
|
||||
|
||||
def test_a_multi_shard_non_fasth3_adapter_stays_on_the_dynamic_route(tmp_path):
|
||||
# A plain PEFT LoRA saved as several shards must not hard-fail startup with a
|
||||
# FastH3-specific message; it has to fall through to the dynamic LoRA route.
|
||||
directory = tmp_path / "peft-sharded"
|
||||
directory.mkdir()
|
||||
for shard in ("adapter_model-00001-of-00002.safetensors", "adapter_model-00002-of-00002.safetensors"):
|
||||
save_file(
|
||||
{"transformer_blocks.0.attn.to_q.lora_A.weight": torch.ones((_RANK, _HIDDEN))}, str(directory / shard)
|
||||
)
|
||||
|
||||
assert _claim(directory) is None
|
||||
|
||||
|
||||
def test_low_rank_factors_reach_the_fused_projections(tmp_path):
|
||||
path = tmp_path / "fasth3" / "adapter_model.safetensors"
|
||||
_write_adapter(path)
|
||||
fusion = _load(path.parent)
|
||||
|
||||
qkv = torch.zeros((3 * _INNER, _HIDDEN), dtype=torch.float32)
|
||||
fused = fusion.fuse("blocks.0.attn.qkv_proj.weight", qkv).cpu()
|
||||
# The checkpoint stores one head group at a time as [q, k, v], so the delta
|
||||
# has to survive the loader's own unpacking as three separate projections.
|
||||
q, k, v = torch.split(
|
||||
_reorder_grouped_qkv_to_qkv(fused, num_query_groups=_HEADS, heads_per_group=1, head_dim=_HEAD_DIM),
|
||||
[_INNER, _INNER, _INNER],
|
||||
)
|
||||
a, b = _factors(_INNER, _HIDDEN)
|
||||
expected = b @ a
|
||||
for got in (q, k, v):
|
||||
assert torch.allclose(got, expected, atol=1e-5)
|
||||
|
||||
|
||||
def test_the_fused_mlp_delta_is_swapped_into_gate_first_order(tmp_path):
|
||||
path = tmp_path / "fasth3" / "adapter_model.safetensors"
|
||||
_write_adapter(path)
|
||||
fusion = _load(path.parent)
|
||||
|
||||
fused = fusion.fuse("blocks.0.mlp.fc1.weight", torch.zeros((2 * _FFN, _HIDDEN))).cpu()
|
||||
a, b = _factors(2 * _FFN, _HIDDEN)
|
||||
value, gate = (b @ a).chunk(2, dim=0)
|
||||
# The diffusers export is value-first; H3's fc1 is gate-first.
|
||||
assert torch.allclose(fused, torch.cat((gate, value), dim=0), atol=1e-5)
|
||||
|
||||
|
||||
def test_diff_and_diff_b_edit_weights_and_biases(tmp_path):
|
||||
path = tmp_path / "fasth3" / "adapter_model.safetensors"
|
||||
_write_adapter(
|
||||
path,
|
||||
tensors={
|
||||
"transformer_blocks.0.norm1.diff": torch.full((_HIDDEN,), 0.25),
|
||||
"transformer_blocks.0.adaln_proj.linear.diff_b": torch.full((_HIDDEN,), -0.5),
|
||||
"context_embedder.diff": torch.full((_HIDDEN, _HIDDEN), 0.125),
|
||||
"context_embedder.diff_b": torch.full((_HIDDEN,), 2.0),
|
||||
},
|
||||
)
|
||||
fusion = _load(path.parent)
|
||||
|
||||
# A full-rank delta on an RMSNorm vector is exactly what a LoRA layer
|
||||
# cannot express, and why this release is fused instead of switched.
|
||||
assert torch.allclose(fusion.fuse("blocks.0.norm1.weight", torch.ones(_HIDDEN)).cpu(), torch.full((_HIDDEN,), 1.25))
|
||||
assert torch.allclose(
|
||||
fusion.fuse("blocks.0.adaln_proj.linear.bias", torch.ones(_HIDDEN)).cpu(), torch.full((_HIDDEN,), 0.5)
|
||||
)
|
||||
assert torch.allclose(
|
||||
fusion.fuse("condition_proj.weight", torch.zeros((_HIDDEN, _HIDDEN))).cpu(),
|
||||
torch.full((_HIDDEN, _HIDDEN), 0.125),
|
||||
)
|
||||
assert torch.allclose(fusion.fuse("condition_proj.bias", torch.zeros(_HIDDEN)).cpu(), torch.full((_HIDDEN,), 2.0))
|
||||
|
||||
|
||||
def test_a_streamed_checkpoint_keeps_the_parameters_the_adapter_never_edits(tmp_path):
|
||||
path = tmp_path / "fasth3" / "adapter_model.safetensors"
|
||||
_write_adapter(path, tensors={"transformer_blocks.0.norm1.diff": torch.full((_HIDDEN,), 3.0)})
|
||||
fusion = _load(path.parent)
|
||||
|
||||
untouched = torch.arange(6, dtype=torch.bfloat16).reshape(2, 3)
|
||||
streamed = dict(
|
||||
fusion.apply([("blocks.0.norm1.weight", torch.zeros(_HIDDEN)), ("blocks.0.attn.q_norm.weight", untouched)])
|
||||
)
|
||||
assert torch.allclose(streamed["blocks.0.norm1.weight"].cpu(), torch.full((_HIDDEN,), 3.0))
|
||||
assert streamed["blocks.0.attn.q_norm.weight"] is untouched
|
||||
|
||||
# validate_fully_applied() releases the deltas, so a second stream would
|
||||
# fuse nothing at all; it has to fail rather than serve base H3 weights.
|
||||
with pytest.raises(FastH3AdapterError, match="already been fused"):
|
||||
next(fusion.apply([("blocks.0.norm1.weight", torch.zeros(_HIDDEN))]))
|
||||
|
||||
|
||||
def test_the_fused_result_keeps_the_parameter_dtype(tmp_path):
|
||||
path = tmp_path / "fasth3" / "adapter_model.safetensors"
|
||||
_write_adapter(path, tensors={"transformer_blocks.0.norm1.diff": torch.full((_HIDDEN,), 0.5)})
|
||||
fusion = _load(path.parent)
|
||||
|
||||
fused = fusion.fuse("blocks.0.norm1.weight", torch.ones(_HIDDEN, dtype=torch.bfloat16))
|
||||
assert fused.dtype == torch.bfloat16
|
||||
|
||||
|
||||
def test_a_vsa_variant_is_recognised_by_its_compression_gates(tmp_path):
|
||||
dense = tmp_path / "dense" / "adapter_model.safetensors"
|
||||
_write_adapter(dense)
|
||||
assert _load(dense.parent).requires_vsa is False
|
||||
|
||||
sparse = tmp_path / "vsa" / "adapter_model.safetensors"
|
||||
_write_adapter(sparse, tensors={"transformer_blocks.0.attn.to_gate_compress.set_weight": torch.ones((2, 2))})
|
||||
assert _load(sparse.parent).requires_vsa is True
|
||||
|
||||
|
||||
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))})
|
||||
# Dropping it silently would load a model that is not the distilled student.
|
||||
with pytest.raises(FastH3AdapterError, match="name no known"):
|
||||
_load(path.parent)
|
||||
|
||||
|
||||
def test_an_unpaired_low_rank_factor_is_an_error(tmp_path):
|
||||
path = tmp_path / "fasth3" / "adapter_model.safetensors"
|
||||
_write_adapter(path, drop="transformer_blocks.0.attn.to_q.lora_B.weight")
|
||||
|
||||
with pytest.raises(FastH3AdapterError, match="unpaired factor"):
|
||||
_load(path.parent)
|
||||
|
||||
|
||||
def test_a_delta_the_checkpoint_never_offered_is_reported(tmp_path):
|
||||
path = tmp_path / "fasth3" / "adapter_model.safetensors"
|
||||
_write_adapter(path)
|
||||
fusion = _load(path.parent)
|
||||
|
||||
fusion.fuse("blocks.0.attn.qkv_proj.weight", torch.zeros((3 * _INNER, _HIDDEN)))
|
||||
# An unapplied delta is the failure that matters: the model would load and
|
||||
# generate, just not as the distilled student.
|
||||
with pytest.raises(FastH3AdapterError, match="never provided"):
|
||||
fusion.validate_fully_applied()
|
||||
|
||||
for name, shape in (
|
||||
("blocks.0.attn.out_proj.weight", (_HIDDEN, _INNER)),
|
||||
("blocks.0.mlp.fc1.weight", (2 * _FFN, _HIDDEN)),
|
||||
("blocks.0.mlp.fc2.weight", (_HIDDEN, _FFN)),
|
||||
):
|
||||
fusion.fuse(name, torch.zeros(shape))
|
||||
fusion.validate_fully_applied()
|
||||
|
||||
|
||||
def _pipeline_stub(fusion=None):
|
||||
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||
|
||||
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||
torch.nn.Module.__init__(pipeline)
|
||||
pipeline.partition = "combined"
|
||||
pipeline.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"})
|
||||
pipeline._turbo_lora_adapter_ids = set()
|
||||
pipeline._fasth3 = fusion
|
||||
pipeline.od_config = SimpleNamespace()
|
||||
pipeline.default_video_shift, pipeline.default_audio_shift = 12.0, 3.0
|
||||
return pipeline
|
||||
|
||||
|
||||
def _sampling(steps=FASTH3_DENOISE_STEPS, **kwargs) -> SimpleNamespace:
|
||||
fields = {"num_inference_steps": steps, "extra_args": {}, "lora_request": None}
|
||||
return SimpleNamespace(**{**fields, **kwargs})
|
||||
|
||||
|
||||
def _check_request(fusion, sampling) -> None:
|
||||
fusion.check_request(sampling, video_shift=12.0, audio_shift=3.0)
|
||||
|
||||
|
||||
def _fusion(tmp_path, **kwargs) -> FastH3WeightFusion:
|
||||
path = tmp_path / "fasth3" / "adapter_model.safetensors"
|
||||
_write_adapter(path, **kwargs)
|
||||
return _load(path.parent)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("task", ["fl2va", "ref2va"])
|
||||
def test_an_active_fasth3_fusion_restricts_requests_to_t2va(task, tmp_path):
|
||||
pipeline = _pipeline_stub(_fusion(tmp_path))
|
||||
|
||||
with pytest.raises(OmniClientError, match="distills \\['t2va'\\] only"):
|
||||
pipeline._resolve_task(task, {})
|
||||
assert pipeline._resolve_task("t2va", {}) == "t2va"
|
||||
assert _pipeline_stub()._resolve_task(task, {}) == task
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_inference_steps", [None, 5, 50])
|
||||
def test_fasth3_requires_one_forward_per_sigma_interval(num_inference_steps, tmp_path):
|
||||
fusion = _fusion(tmp_path)
|
||||
|
||||
# The five sigma points bound four forwards, and a request states forwards:
|
||||
# the interval contract the pinned schedules use, and the count the step
|
||||
# scheduler admits a request on before any pipeline hook runs.
|
||||
with pytest.raises(OmniClientError, match=f"num_inference_steps={FASTH3_DENOISE_STEPS}"):
|
||||
_check_request(fusion, _sampling(num_inference_steps))
|
||||
_check_request(fusion, _sampling())
|
||||
# The pinned branch coerces the string form; this one cannot disagree.
|
||||
_check_request(fusion, _sampling(str(FASTH3_DENOISE_STEPS)))
|
||||
|
||||
|
||||
def test_the_request_denoises_on_the_release_ladder(tmp_path):
|
||||
pipeline = _pipeline_stub(_fusion(tmp_path))
|
||||
|
||||
positions, num_steps = pipeline._resolve_sigma_positions("t2va", _sampling())
|
||||
|
||||
# dmd_denoising_steps [999, 749, 500, 250] out of 1000, closed with the
|
||||
# terminal 0.0: the pre-shift positions the student was distilled at, not
|
||||
# the uniform ladder num_inference_steps would otherwise derive.
|
||||
assert positions == (0.999, 0.749, 0.5, 0.25, 0.0)
|
||||
# Five sigma points bound four transformer forwards, and forwards is the
|
||||
# unit num_steps carries downstream (Cache-DiT, quality hints).
|
||||
assert num_steps == FASTH3_DENOISE_STEPS == len(positions) - 1
|
||||
# H3's own per-modality shifts still apply on top, and are what turn those
|
||||
# positions into the noise levels the four forwards run at.
|
||||
assert minimax_h3_time_shift_sigmas(shift_scale=12.0, base_schedule=positions) == pytest.approx(
|
||||
[0.999917, 0.972833, 0.923077, 0.8, 0.0], abs=1e-6
|
||||
)
|
||||
assert minimax_h3_time_shift_sigmas(shift_scale=3.0, base_schedule=positions) == pytest.approx(
|
||||
[0.999666, 0.899520, 0.75, 0.5, 0.0], abs=1e-6
|
||||
)
|
||||
# Without a fused adapter the same request keeps the undistilled ladder.
|
||||
assert _pipeline_stub()._resolve_sigma_positions("t2va", _sampling()) == (None, FASTH3_DENOISE_STEPS)
|
||||
|
||||
|
||||
def test_the_contract_leaves_the_checkpoint_shifts_alone(tmp_path):
|
||||
pipeline = _pipeline_stub(_fusion(tmp_path))
|
||||
pipeline.partition = "fl2va"
|
||||
|
||||
pipeline._fasth3.check_serving_contract(
|
||||
partition=pipeline.partition,
|
||||
od_config=pipeline.od_config,
|
||||
video_shift=pipeline.default_video_shift,
|
||||
audio_shift=pipeline.default_audio_shift,
|
||||
)
|
||||
|
||||
# The ladder is stated pre-shift, so H3's 12/3 shifts are what place the
|
||||
# student on the levels it was distilled at; the contract must not touch them.
|
||||
assert pipeline.default_video_shift == 12.0
|
||||
assert pipeline.default_audio_shift == 3.0
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"extra",
|
||||
[{"flow_shift": 1.0}, {"audio_flow_shift": 1.0}, {"flow_shift": "bad"}],
|
||||
)
|
||||
def test_a_request_may_not_move_the_student_off_its_rungs(extra, tmp_path):
|
||||
fusion = _fusion(tmp_path)
|
||||
|
||||
with pytest.raises(OmniClientError, match="FastH3 requires"):
|
||||
_check_request(fusion, _sampling(extra_args=extra))
|
||||
|
||||
_check_request(fusion, _sampling(extra_args={"flow_shift": 12.0, "audio_flow_shift": 3.0}))
|
||||
|
||||
|
||||
def test_a_request_may_not_carry_a_lora_on_a_fused_server(tmp_path):
|
||||
fusion = _fusion(tmp_path)
|
||||
|
||||
# The dynamic LoRA manager is skipped for a fused adapter, so an adapter
|
||||
# asked for per request would be neither applied nor reported.
|
||||
with pytest.raises(OmniClientError, match="per-request lora is unavailable"):
|
||||
_check_request(fusion, _sampling(lora_request=SimpleNamespace(lora_int_id=1)))
|
||||
|
||||
|
||||
def test_offload_is_refused_because_it_bypasses_the_fusion(tmp_path):
|
||||
fusion = _fusion(tmp_path)
|
||||
for flag in ("enable_cpu_offload", "enable_layerwise_offload", "enable_distributed_layerwise_offload"):
|
||||
# A host-weight plan installs the transformer without load_weights(),
|
||||
# so the fusion and its completeness check would both be skipped.
|
||||
with pytest.raises(ValueError, match="cannot be combined with"):
|
||||
fusion.check_serving_contract(
|
||||
partition="fl2va",
|
||||
od_config=SimpleNamespace(**{flag: True}),
|
||||
video_shift=12.0,
|
||||
audio_shift=3.0,
|
||||
)
|
||||
|
||||
|
||||
def test_adopting_the_contract_refuses_a_vsa_variant_and_ref2va(tmp_path):
|
||||
sparse = tmp_path / "vsa" / "adapter_model.safetensors"
|
||||
_write_adapter(sparse, tensors={"transformer_blocks.0.attn.to_gate_compress.set_weight": torch.ones((2, 2))})
|
||||
contract = {"od_config": SimpleNamespace(), "video_shift": 12.0, "audio_shift": 3.0}
|
||||
with pytest.raises(ValueError, match="Video Sparse Attention variant"):
|
||||
_load(sparse.parent).check_serving_contract(partition="fl2va", **contract)
|
||||
|
||||
with pytest.raises(ValueError, match="cannot serve a Ref2VA partition"):
|
||||
_fusion(tmp_path).check_serving_contract(partition="ref2va", **contract)
|
||||
|
||||
|
||||
def test_a_full_rank_delta_on_a_fused_parameter_is_refused(tmp_path):
|
||||
# Only low-rank factors are placed into H3's fused QKV and gate/up layouts,
|
||||
# so a .diff aimed at one would otherwise be added transposed.
|
||||
path = tmp_path / "fasth3" / "adapter_model.safetensors"
|
||||
_write_adapter(path, tensors={"transformer_blocks.0.attn.to_q.diff": torch.ones((_INNER, _HIDDEN))})
|
||||
with pytest.raises(FastH3AdapterError, match="fused layout"):
|
||||
_claim(path.parent)
|
||||
|
||||
|
||||
def test_an_adapter_that_leaves_blocks_untouched_is_refused(tmp_path):
|
||||
path = tmp_path / "fasth3" / "adapter_model.safetensors"
|
||||
_write_adapter(path, blocks=2)
|
||||
|
||||
# Every tensor in this file is well formed and pairs up, so the fusion's
|
||||
# own completeness check passes; only the model's depth reveals that a
|
||||
# third of it would have been served as base H3 weights.
|
||||
with pytest.raises(FastH3AdapterError, match="missing=\\[2\\]"):
|
||||
_claim(path.parent, num_blocks=3)
|
||||
with pytest.raises(FastH3AdapterError, match="unknown=\\[1\\]"):
|
||||
_claim(path.parent, num_blocks=1)
|
||||
assert _claim(path.parent, num_blocks=2) is not None
|
||||
|
||||
|
||||
def test_an_adapter_shorter_than_it_declares_is_refused(tmp_path):
|
||||
path = tmp_path / "fasth3" / "adapter_model.safetensors"
|
||||
_write_adapter(path, metadata={"low_rank_tensors": "724", "diff_tensors": "85"})
|
||||
|
||||
# The writer records what it emitted, which is the file's one statement
|
||||
# about its own completeness.
|
||||
with pytest.raises(FastH3AdapterError, match="declares low_rank_tensors=724 but carries 12"):
|
||||
_claim(path.parent)
|
||||
|
||||
|
||||
def test_only_a_configured_adapter_path_is_read(tmp_path):
|
||||
path = tmp_path / "fasth3" / "adapter_model.safetensors"
|
||||
_write_adapter(path)
|
||||
transformer = SimpleNamespace(
|
||||
arch=SimpleNamespace(attention_head_dim=_HEAD_DIM, num_layers=1, token_refiner_num_layers=0)
|
||||
)
|
||||
unbuilt = object() # a transformer whose arch must not be touched
|
||||
|
||||
assert resolve_fasth3_fusion(SimpleNamespace(lora_path=str(path)), transformer) is not None
|
||||
assert resolve_fasth3_fusion(SimpleNamespace(lora_path=[str(path)]), transformer) is not None
|
||||
# Nothing configured, or more than one artifact, stays on the dynamic route
|
||||
# without reading anything off the model.
|
||||
assert resolve_fasth3_fusion(SimpleNamespace(lora_path=None), unbuilt) is None
|
||||
assert resolve_fasth3_fusion(SimpleNamespace(), unbuilt) is None
|
||||
assert resolve_fasth3_fusion(SimpleNamespace(lora_path=[str(path), str(path)]), unbuilt) is None
|
||||
|
||||
|
||||
def test_a_pipeline_that_fused_its_adapter_needs_no_lora_manager(tmp_path):
|
||||
assert _pipeline_stub(_fusion(tmp_path)).lora_is_fused is True
|
||||
assert _pipeline_stub().lora_is_fused is False
|
||||
|
||||
|
||||
def test_the_ladder_is_the_one_the_release_publishes():
|
||||
# FASTH3_DENOISE_STEPS is derived from it, so the two cannot drift apart.
|
||||
assert FASTH3_BASE_SCHEDULE.base_schedule == (0.999, 0.749, 0.5, 0.25, 0.0)
|
||||
assert FASTH3_BASE_SCHEDULE.num_inference_steps == FASTH3_DENOISE_STEPS == 4
|
||||
@@ -259,7 +259,11 @@ def test_h3_native_rejects_schedule_with_wrong_interval_count(tmp_path):
|
||||
("", "must not be empty"),
|
||||
(",,", "must not be empty"),
|
||||
("1.0,oops,0.15,0.0", "is malformed"),
|
||||
("0.9,0.7,0.4,0.15,0.0", "is malformed"),
|
||||
# A capped opening rung below 1.0 is a legal DMD2 schedule, so the
|
||||
# out-of-range cases are a first position above 1.0 and a tail that
|
||||
# never reaches clean latents.
|
||||
("1.5,0.7,0.4,0.15,0.0", "is malformed"),
|
||||
("1.0,0.7,0.4,0.15", "is malformed"),
|
||||
],
|
||||
)
|
||||
def test_h3_native_rejects_malformed_schedule_metadata(tmp_path, raw_schedule, message):
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
import pytest
|
||||
|
||||
from vllm_omni.diffusion.sched import DMD2SigmaSchedule
|
||||
@@ -18,7 +18,8 @@ def test_step_count_is_intervals_not_boundaries():
|
||||
[
|
||||
[],
|
||||
[1.0],
|
||||
[0.9, 0.5, 0.0],
|
||||
[1.5, 0.5, 0.0],
|
||||
[0.0, 0.5, 0.0],
|
||||
[1.0, 0.5, 0.1],
|
||||
[1.0, 0.4, 0.6, 0.0],
|
||||
[1.0, 0.5, 0.5, 0.0],
|
||||
@@ -31,6 +32,15 @@ def test_malformed_positions_are_rejected(positions):
|
||||
DMD2SigmaSchedule.from_positions(positions)
|
||||
|
||||
|
||||
def test_a_capped_first_position_is_a_schedule_not_a_malformed_one():
|
||||
# A DMD2 student trained with max_timestep_ratio < 1 never sees pure noise,
|
||||
# so its opening rung sits just below 1.0 (FastH3: 999/1000).
|
||||
schedule = DMD2SigmaSchedule.from_positions([0.999, 0.749, 0.5, 0.25, 0.0])
|
||||
|
||||
assert schedule.base_schedule[0] == 0.999
|
||||
assert schedule.num_inference_steps == 4
|
||||
|
||||
|
||||
def test_one_base_schedule_drives_several_shift_scales():
|
||||
schedule = DMD2SigmaSchedule.from_positions([1.0, 0.7, 0.4, 0.15, 0.0])
|
||||
|
||||
|
||||
@@ -520,6 +520,32 @@ def test_interleaved_video_uses_legacy_fallback_automatically(monkeypatch):
|
||||
assert len(coerce_calls) == 1
|
||||
|
||||
|
||||
def test_uint8_frames_reach_the_muxer_without_a_float_round_trip():
|
||||
"""uint8 is the muxer's own dtype, so normalisation must leave it alone.
|
||||
|
||||
MiniMax-H3 quantises on the accelerator to keep the response small across
|
||||
its two process hops. Re-normalising to float here would undo that at the
|
||||
last step and cost a full-size conversion in each direction.
|
||||
"""
|
||||
video = np.arange(2 * 4 * 5 * 3, dtype=np.uint8).reshape(2, 4, 5, 3)
|
||||
|
||||
frames, frame_shape, common_dtype = video_api_utils._prepare_video_frames(video)
|
||||
|
||||
assert common_dtype == np.dtype(np.uint8)
|
||||
assert frame_shape == (4, 5, 3)
|
||||
np.testing.assert_array_equal(np.stack(frames), video)
|
||||
|
||||
|
||||
def test_non_uint8_integer_frames_are_still_normalised():
|
||||
"""Only uint8 is the muxer's dtype; other integer payloads still scale."""
|
||||
video = np.full((2, 4, 5, 3), 255, dtype=np.int32)
|
||||
|
||||
frames, _, common_dtype = video_api_utils._prepare_video_frames(video)
|
||||
|
||||
assert np.issubdtype(common_dtype, np.floating)
|
||||
np.testing.assert_allclose(frames[0], 1.0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"video",
|
||||
[
|
||||
@@ -552,9 +578,10 @@ def test_prepared_automatic_fallback_preserves_legacy_quantization(monkeypatch,
|
||||
monkeypatch.setattr(media_utils, "mux_video_audio_bytes", fake_compat_mux)
|
||||
|
||||
prepared_frames, frame_shape, _ = video_api_utils._prepare_video_frames(video)
|
||||
reference = np.rint(np.clip(np.stack([frame[..., :3] for frame in prepared_frames]), 0.0, 1.0) * 255.0).astype(
|
||||
np.uint8
|
||||
)
|
||||
stacked = np.stack([frame[..., :3] for frame in prepared_frames])
|
||||
# uint8 frames are the muxer's own dtype and reach it unchanged; everything
|
||||
# else is still quantised out of the normalised [0, 1] range.
|
||||
reference = stacked if stacked.dtype == np.uint8 else np.rint(np.clip(stacked, 0.0, 1.0) * 255.0).astype(np.uint8)
|
||||
|
||||
assert video_api_utils._encode_video_bytes(video, fps=12) == b"legacy-video"
|
||||
assert len(mux_inputs) == 1
|
||||
|
||||
@@ -142,7 +142,7 @@ def _get_symm_buffer(
|
||||
stream_id=stream_id,
|
||||
)
|
||||
_SYMM_WORKSPACES[key] = workspace
|
||||
return workspace.handle.get_buffer(workspace.handle.rank, symm_shape, dtype)
|
||||
return workspace.allocation[:required_bytes].view(dtype).view(symm_shape)
|
||||
|
||||
|
||||
def clear_a2a_permute_workspaces() -> None:
|
||||
|
||||
@@ -0,0 +1,679 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""FastVideo FastH3: a four-step DMD2 student of MiniMax-H3.
|
||||
|
||||
FastH3 replaces H3's 49 denoiser evaluations with four. It ships as an adapter
|
||||
over the base checkpoint rather than as a full release, so it reuses H3's text
|
||||
encoder, video VAE, audio VAE, tokenizers and schedulers unchanged.
|
||||
|
||||
The artifact is *not* a PEFT LoRA, and it is not request-switchable. Its own
|
||||
metadata states the reconstruction as::
|
||||
|
||||
W = W_base + lora_B @ lora_A; then .diff/.diff_b added and .set_weight assigned
|
||||
|
||||
so besides rank-64 factors it carries full-rank ``.diff``/``.diff_b`` deltas for
|
||||
RMSNorm weights, biases, patch projections and the final layer - none of which a
|
||||
LoRA layer can express - and the VSA variants add ``.set_weight`` tensors for
|
||||
compression gates that do not exist in the base transformer at all. The adapter
|
||||
is therefore fused into the checkpoint stream at load time, before the weights
|
||||
are sharded, which is also what the release's model card requires.
|
||||
|
||||
The low-rank factors carry no alpha: the reconstruction adds ``lora_B @ lora_A``
|
||||
directly, i.e. a scale of exactly 1.
|
||||
|
||||
Two checkpoint spellings meet here. The adapter is written in the diffusers
|
||||
namespace (``transformer_blocks.0.attn.to_q``) while vLLM-Omni loads H3's native
|
||||
one (``blocks.0.attn.qkv_proj``), whose attention and MLP projections are fused.
|
||||
Every mapping and layout convention below was verified tensor by tensor against
|
||||
the released full checkpoint (``FastVideo-FastH3-4-step-Preview-v1-Dense-DataFree``):
|
||||
``W_base + delta`` reproduces it to bf16 rounding.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from collections.abc import Iterable, Iterator, Mapping
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import torch
|
||||
from safetensors import safe_open
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule
|
||||
from vllm_omni.errors import OmniClientError
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .minimax_h3_transformer import MiniMaxH3DiTModel
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
# FastVideo's generic adapter container, not a FastH3 marker: their tools emit
|
||||
# it for ordinary H3 adapters too.
|
||||
FASTH3_FORMAT = "fastvideo-lora-v2"
|
||||
FASTH3_MANIFEST = "adapter_manifest.json"
|
||||
# The release identity: the distilled student the adapter came from, over the
|
||||
# base it edits. The name is held to FastVideo's own namespace as well as the
|
||||
# student's name, so an unrelated adapter that merely mentions FastH3 is not
|
||||
# claimed.
|
||||
FASTH3_BASE_MODEL = "MiniMaxAI/MiniMax-H3"
|
||||
_FASTH3_IDENTITY_KEY = "finetuned_model"
|
||||
_FASTH3_IDENTITY_NAMESPACE = "fastvideo/"
|
||||
_FASTH3_IDENTITY_MARKER = "fasth3"
|
||||
|
||||
# The rectified-flow positions the student was distilled at, and the ladder the
|
||||
# server samples on. The release states them as `dmd_denoising_steps`
|
||||
# [999, 749, 500, 250] (`sampling.base_timesteps` in the bundle manifest):
|
||||
# timestep indices out of 1000, i.e. pre-shift positions, closed here with the
|
||||
# terminal 0.0 every rectified-flow schedule ends on. The opening rung is 0.999
|
||||
# rather than 1.0 because training capped the noise level there
|
||||
# (`max_timestep_ratio`).
|
||||
#
|
||||
# These are positions, not final sigmas: H3's schedulers apply their own
|
||||
# per-modality shift on top (12 for video, 3 for audio), which is what
|
||||
# reproduces the levels the student saw. Nothing here overrides those shifts.
|
||||
FASTH3_BASE_SCHEDULE = DMD2SigmaSchedule.from_positions((0.999, 0.749, 0.5, 0.25, 0.0))
|
||||
# Its five points bound four transformer forwards, one per sigma interval. That
|
||||
# count is the one a request states, the one Cache-DiT is refreshed with and the
|
||||
# one step execution admits a request on - the same interval contract H3's
|
||||
# pinned checkpoint schedules and native LoRAs already use.
|
||||
FASTH3_DENOISE_STEPS = FASTH3_BASE_SCHEDULE.num_inference_steps
|
||||
# Preview v1 distills the text-to-video-and-audio path only.
|
||||
FASTH3_SUPPORTED_TASKS = frozenset({"t2va"})
|
||||
|
||||
_LORA_A = ".lora_A.weight"
|
||||
_LORA_B = ".lora_B.weight"
|
||||
_DIFF = ".diff"
|
||||
_DIFF_B = ".diff_b"
|
||||
_SET_WEIGHT = ".set_weight"
|
||||
|
||||
# Adapter module prefix -> the native parameter it edits, minus the
|
||||
# ``.weight``/``.bias`` suffix.
|
||||
_MODEL_LEVEL_TARGETS = {
|
||||
"proj_in": "video_patch_proj",
|
||||
"proj_out": "final_layer.video_out",
|
||||
"audio_proj_in": "audio_patch_proj",
|
||||
"audio_proj_out": "final_layer.audio_out",
|
||||
"context_embedder": "condition_proj",
|
||||
"time_embedder.linear_1": "time_embedder.proj_in",
|
||||
"time_embedder.linear_2": "time_embedder.proj_out",
|
||||
"norm_out.linear": "final_layer.adaln_proj.linear",
|
||||
"norm_out.norm": "final_layer.norm",
|
||||
}
|
||||
|
||||
# Per-block adapter suffix -> (native suffix, how a delta enters the native
|
||||
# parameter). H3 stores attention as one grouped QKV matrix and the MLP as one
|
||||
# fused gate/up matrix, so those deltas need placing rather than adding.
|
||||
_PLAIN, _QKV_Q, _QKV_K, _QKV_V, _SWAP_HALVES = "plain", "q", "k", "v", "swap_halves"
|
||||
_QKV_SLOTS = (_QKV_Q, _QKV_K, _QKV_V)
|
||||
_BLOCK_TARGETS = {
|
||||
"attn.to_q": ("attn.qkv_proj", _QKV_Q),
|
||||
"attn.to_k": ("attn.qkv_proj", _QKV_K),
|
||||
"attn.to_v": ("attn.qkv_proj", _QKV_V),
|
||||
"attn.to_out.0": ("attn.out_proj", _PLAIN),
|
||||
"ff.net.0.proj": ("mlp.fc1", _SWAP_HALVES),
|
||||
"ff.net.2": ("mlp.fc2", _PLAIN),
|
||||
"adaln_proj.linear": ("adaln_proj.linear", _PLAIN),
|
||||
"norm1": ("norm1", _PLAIN),
|
||||
"norm2": ("norm2", _PLAIN),
|
||||
}
|
||||
|
||||
# Adapter block prefix -> native block prefix.
|
||||
_BLOCK_PREFIXES = (
|
||||
("token_refiner.refiner_blocks.", "token_refiner.blocks."),
|
||||
("transformer_blocks.", "blocks."),
|
||||
)
|
||||
|
||||
|
||||
class FastH3AdapterError(ValueError):
|
||||
"""The artifact is a FastH3 adapter, but it cannot be applied as one."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class _ParamPatch:
|
||||
"""Everything the adapter contributes to one native parameter."""
|
||||
|
||||
# layout -> (lora_A, lora_B). A grouped QKV parameter collects three.
|
||||
low_rank: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = field(default_factory=dict)
|
||||
diff: torch.Tensor | None = None
|
||||
layout: str = _PLAIN
|
||||
|
||||
|
||||
def _swap_halves(tensor: torch.Tensor) -> torch.Tensor:
|
||||
"""Exchange the two halves of a fused gate/up matrix.
|
||||
|
||||
The diffusers export stores the feed-forward projection value-first while
|
||||
H3's native ``mlp.fc1`` is gate-first, so a delta computed in the diffusers
|
||||
layout has to be swapped before it can be added to the native parameter.
|
||||
"""
|
||||
if tensor.shape[0] % 2:
|
||||
raise FastH3AdapterError(f"fused gate/up delta must split evenly, got {tuple(tensor.shape)}")
|
||||
first, second = tensor.chunk(2, dim=0)
|
||||
return torch.cat((second, first), dim=0)
|
||||
|
||||
|
||||
def _place_in_grouped_qkv(deltas: Mapping[str, torch.Tensor], *, head_dim: int) -> torch.Tensor:
|
||||
"""Interleave per-projection deltas into H3's grouped QKV layout.
|
||||
|
||||
The checkpoint stores one head group at a time as ``[q, k, v]``, which is
|
||||
what :func:`_reorder_grouped_qkv_to_qkv` unpacks on the way in. A delta
|
||||
built from the separate diffusers projections has to be folded back into
|
||||
that order.
|
||||
"""
|
||||
missing = sorted(set(_QKV_SLOTS) - set(deltas))
|
||||
if missing:
|
||||
raise FastH3AdapterError(f"grouped QKV delta is missing its {missing} projections")
|
||||
parts = []
|
||||
for slot in _QKV_SLOTS:
|
||||
delta = deltas[slot]
|
||||
if delta.shape[0] % head_dim:
|
||||
raise FastH3AdapterError(
|
||||
f"QKV {slot} delta rows {delta.shape[0]} are not a multiple of head_dim {head_dim}"
|
||||
)
|
||||
parts.append(delta.reshape(delta.shape[0] // head_dim, head_dim, *delta.shape[1:]))
|
||||
groups = parts[0].shape[0]
|
||||
if any(part.shape[0] != groups for part in parts):
|
||||
raise FastH3AdapterError("QKV projections disagree on the number of head groups")
|
||||
return torch.cat(parts, dim=1).reshape(groups * 3 * head_dim, *parts[0].shape[2:])
|
||||
|
||||
|
||||
def _resolve_native_target(module: str) -> tuple[str, str, tuple[str, int] | None] | None:
|
||||
"""Map an adapter module path to ``(native path, layout, block)``.
|
||||
|
||||
``block`` is the ``(native block prefix, index)`` the module sits in, or
|
||||
``None`` for a model-level one. Coverage is checked per block, so the index
|
||||
has to survive the mapping rather than being folded into the path.
|
||||
"""
|
||||
native = _MODEL_LEVEL_TARGETS.get(module)
|
||||
if native is not None:
|
||||
return native, _PLAIN, None
|
||||
for adapter_prefix, native_prefix in _BLOCK_PREFIXES:
|
||||
if not module.startswith(adapter_prefix):
|
||||
continue
|
||||
remainder = module[len(adapter_prefix) :]
|
||||
index, _, suffix = remainder.partition(".")
|
||||
if not index.isdigit():
|
||||
return None
|
||||
target = _BLOCK_TARGETS.get(suffix)
|
||||
if target is None:
|
||||
return None
|
||||
native_suffix, layout = target
|
||||
return f"{native_prefix}{index}.{native_suffix}", layout, (native_prefix, int(index))
|
||||
return None
|
||||
|
||||
|
||||
def _split_adapter_key(name: str) -> tuple[str, str] | None:
|
||||
"""Split an adapter tensor name into ``(module path, role)``."""
|
||||
for marker, role in (
|
||||
(_LORA_A, "lora_a"),
|
||||
(_LORA_B, "lora_b"),
|
||||
(_DIFF_B, "diff_b"),
|
||||
(_DIFF, "diff"),
|
||||
(_SET_WEIGHT, "set_weight"),
|
||||
):
|
||||
if name.endswith(marker):
|
||||
return name[: -len(marker)], role
|
||||
return None
|
||||
|
||||
|
||||
def _is_fasth3_release(metadata: Mapping[str, str]) -> bool:
|
||||
"""Whether this is a FastH3 release rather than merely FastVideo's format.
|
||||
|
||||
Claiming on the container alone would fuse an ordinary H3 adapter into the
|
||||
checkpoint and put every request on a four-step schedule it was never
|
||||
distilled for, and so would a bare ``fasth3`` substring: an adapter someone
|
||||
else named after the student is not the student.
|
||||
"""
|
||||
if metadata.get("format") != FASTH3_FORMAT:
|
||||
return False
|
||||
identity = metadata.get(_FASTH3_IDENTITY_KEY, "").lower()
|
||||
if not identity.startswith(_FASTH3_IDENTITY_NAMESPACE) or _FASTH3_IDENTITY_MARKER not in identity:
|
||||
return False
|
||||
# Only held against a file that states it.
|
||||
base_model = metadata.get("base_model")
|
||||
return base_model is None or base_model.casefold() == FASTH3_BASE_MODEL.casefold()
|
||||
|
||||
|
||||
def _check_declared_counts(
|
||||
metadata: Mapping[str, str],
|
||||
counted: Mapping[str, int],
|
||||
*,
|
||||
weights_path: Path,
|
||||
) -> None:
|
||||
"""Hold the artifact to the tensor counts its own metadata declares.
|
||||
|
||||
The writer of the ``fastvideo-lora-v2`` format records how many tensors of
|
||||
each kind it emitted, which is the one statement in the file about its own
|
||||
completeness. A truncated or partially re-exported artifact is otherwise
|
||||
indistinguishable from a small one, so a claimed file has to carry all of
|
||||
them rather than opt out by omission.
|
||||
"""
|
||||
for key, seen in counted.items():
|
||||
declared = metadata.get(key)
|
||||
if declared is None:
|
||||
raise FastH3AdapterError(
|
||||
f"{weights_path} declares the FastH3 identity but omits {key}; "
|
||||
"a claimed adapter has to state what it emitted"
|
||||
)
|
||||
try:
|
||||
expected = int(declared)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise FastH3AdapterError(f"{weights_path} declares a non-numeric {key}={declared!r}") from exc
|
||||
if expected != seen:
|
||||
raise FastH3AdapterError(
|
||||
f"{weights_path} declares {key}={expected} but carries {seen}; the adapter is incomplete "
|
||||
"and would leave most of the transformer on base H3 weights"
|
||||
)
|
||||
|
||||
|
||||
def _check_block_coverage(
|
||||
seen: Mapping[str, set[int]],
|
||||
*,
|
||||
expected: Mapping[str, int],
|
||||
weights_path: Path,
|
||||
) -> None:
|
||||
"""Every block of the model this adapter is loaded against must be edited.
|
||||
|
||||
The release drops tensors training left unchanged, so per-parameter
|
||||
coverage is legitimately sparse - but a distilled student touches every
|
||||
block, so a block with no edits at all means the artifact does not match
|
||||
this model.
|
||||
"""
|
||||
for prefix, count in expected.items():
|
||||
indices = seen.get(prefix, set())
|
||||
wanted = set(range(count))
|
||||
if indices == wanted:
|
||||
continue
|
||||
missing = sorted(wanted - indices)
|
||||
extra = sorted(indices - wanted)
|
||||
raise FastH3AdapterError(
|
||||
f"{weights_path} edits {len(indices)} of the model's {count} {prefix}* blocks "
|
||||
f"(missing={missing[:5]}, unknown={extra[:5]}); it is not an adapter for this checkpoint"
|
||||
)
|
||||
|
||||
|
||||
class FastH3WeightFusion:
|
||||
"""Fuse a FastH3 adapter into the H3 checkpoint stream as it is loaded."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
source: Path,
|
||||
patches: Mapping[str, _ParamPatch],
|
||||
head_dim: int,
|
||||
requires_vsa: bool,
|
||||
) -> None:
|
||||
self._source = source
|
||||
self._patches = dict(patches)
|
||||
self._head_dim = head_dim
|
||||
self.requires_vsa = requires_vsa
|
||||
self._applied: set[str] = set()
|
||||
self._device: torch.device | None = None
|
||||
|
||||
@property
|
||||
def source(self) -> Path:
|
||||
return self._source
|
||||
|
||||
@property
|
||||
def base_schedule(self) -> tuple[float, ...]:
|
||||
"""The rectified-flow positions this student samples on.
|
||||
|
||||
The fused checkpoint is a four-step student, so the ladder comes from
|
||||
the release rather than from the many-step teacher's metadata or from
|
||||
the uniform one ``num_inference_steps`` would otherwise derive.
|
||||
"""
|
||||
return FASTH3_BASE_SCHEDULE.base_schedule
|
||||
|
||||
@classmethod
|
||||
def from_path(
|
||||
cls,
|
||||
path: str | Path,
|
||||
*,
|
||||
head_dim: int,
|
||||
num_blocks: int,
|
||||
num_refiner_blocks: int,
|
||||
) -> FastH3WeightFusion | None:
|
||||
"""Build a fusion from an adapter file or directory, else ``None``.
|
||||
|
||||
Returning ``None`` keeps every other ``--lora-path`` artifact on the
|
||||
dynamic LoRA route; only a file carrying the FastH3 release identity is
|
||||
claimed here.
|
||||
|
||||
The block counts are the model's, and the artifact has to cover them:
|
||||
claiming a partial adapter would switch the server onto the four-step
|
||||
contract while most of the transformer still held base H3 weights.
|
||||
"""
|
||||
weights_path = _resolve_adapter_file(path)
|
||||
if weights_path is None:
|
||||
return None
|
||||
|
||||
patches: dict[str, _ParamPatch] = {}
|
||||
gate_tensors: list[str] = []
|
||||
unmapped: list[str] = []
|
||||
blocks_seen: dict[str, set[int]] = {}
|
||||
counted = {"low_rank_tensors": 0, "diff_tensors": 0}
|
||||
with safe_open(weights_path, framework="pt", device="cpu") as checkpoint:
|
||||
metadata = checkpoint.metadata() or {}
|
||||
if not _is_fasth3_release(metadata):
|
||||
return None
|
||||
for name in checkpoint.keys():
|
||||
split = _split_adapter_key(name)
|
||||
if split is None:
|
||||
unmapped.append(name)
|
||||
continue
|
||||
module, role = split
|
||||
if role == "set_weight":
|
||||
# A VSA compression gate: a module the base transformer does
|
||||
# not have, so there is nothing to fuse it into.
|
||||
gate_tensors.append(name)
|
||||
continue
|
||||
target = _resolve_native_target(module)
|
||||
if target is None:
|
||||
unmapped.append(name)
|
||||
continue
|
||||
native_module, layout, block = target
|
||||
if block is not None:
|
||||
blocks_seen.setdefault(block[0], set()).add(block[1])
|
||||
counted["diff_tensors" if role in ("diff", "diff_b") else "low_rank_tensors"] += 1
|
||||
native_param = f"{native_module}.{'bias' if role == 'diff_b' else 'weight'}"
|
||||
patch = patches.setdefault(native_param, _ParamPatch(layout=layout))
|
||||
tensor = checkpoint.get_tensor(name)
|
||||
if role in ("diff", "diff_b"):
|
||||
if patch.diff is not None:
|
||||
raise FastH3AdapterError(f"duplicate {role} for {native_param}")
|
||||
patch.diff = tensor
|
||||
else:
|
||||
a, b = patch.low_rank.get(layout, (None, None))
|
||||
patch.low_rank[layout] = (tensor, b) if role == "lora_a" else (a, tensor)
|
||||
|
||||
if unmapped:
|
||||
raise FastH3AdapterError(
|
||||
f"FastH3 adapter at {weights_path} has {len(unmapped)} tensors that name no known "
|
||||
f"H3 parameter: {sorted(unmapped)[:5]}"
|
||||
)
|
||||
_check_declared_counts(
|
||||
metadata,
|
||||
{**counted, "set_weight_tensors": len(gate_tensors)},
|
||||
weights_path=weights_path,
|
||||
)
|
||||
_check_block_coverage(
|
||||
blocks_seen,
|
||||
expected={"blocks.": num_blocks, "token_refiner.blocks.": num_refiner_blocks},
|
||||
weights_path=weights_path,
|
||||
)
|
||||
for native_param, patch in patches.items():
|
||||
for slot, (a, b) in patch.low_rank.items():
|
||||
if a is None or b is None:
|
||||
raise FastH3AdapterError(f"FastH3 adapter has an unpaired factor for {native_param} slot {slot!r}")
|
||||
# Only the low-rank factors are placed into H3's fused QKV and
|
||||
# gate/up layouts; a full-rank delta is added as it comes, so one
|
||||
# aimed at a fused parameter would silently land transposed.
|
||||
if patch.diff is not None and patch.layout != _PLAIN:
|
||||
raise FastH3AdapterError(
|
||||
f"FastH3 adapter carries a full-rank delta for {native_param}, which H3 stores in the "
|
||||
f"{patch.layout!r} fused layout; this loader can only place low-rank factors there"
|
||||
)
|
||||
|
||||
fusion = cls(
|
||||
source=weights_path,
|
||||
patches=patches,
|
||||
head_dim=head_dim,
|
||||
requires_vsa=bool(gate_tensors),
|
||||
)
|
||||
logger.info(
|
||||
"FastH3 adapter %s: rank=%s, parameters patched=%d, low-rank=%s, diff=%s, set_weight=%d",
|
||||
weights_path,
|
||||
metadata.get("rank", "?"),
|
||||
len(patches),
|
||||
metadata.get("low_rank_tensors", "?"),
|
||||
metadata.get("diff_tensors", "?"),
|
||||
len(gate_tensors),
|
||||
)
|
||||
return fusion
|
||||
|
||||
def _compute_device(self, weight: torch.Tensor) -> torch.device:
|
||||
"""Where to reconstruct a delta.
|
||||
|
||||
H3's per-block modulation projection is 96768x2688, so rebuilding all
|
||||
343 patched parameters is a few TFLOP of rank-64 products. On CPU that
|
||||
adds minutes to a load that already has a startup deadline, so the
|
||||
accelerator does the arithmetic whenever there is one.
|
||||
"""
|
||||
if weight.device.type != "cpu":
|
||||
return weight.device
|
||||
if self._device is None:
|
||||
# Ask the platform rather than PyTorch's global accelerator
|
||||
# registry, so an out-of-tree backend controls its own placement.
|
||||
self._device = current_omni_platform.get_torch_device()
|
||||
return self._device
|
||||
|
||||
@staticmethod
|
||||
def _widen(tensor: torch.Tensor, device: torch.device) -> torch.Tensor:
|
||||
"""Move to ``device``, then widen to float32.
|
||||
|
||||
Asking ``Tensor.to`` for a device and a dtype at once converts on the
|
||||
host and ships twice the bytes; splitting it moves bfloat16 and widens
|
||||
on the accelerator.
|
||||
"""
|
||||
return tensor.to(device, non_blocking=True).to(torch.float32)
|
||||
|
||||
def fuse(self, name: str, weight: torch.Tensor) -> torch.Tensor:
|
||||
"""Return ``weight`` with this adapter's contribution added."""
|
||||
patch = self._patches.get(name)
|
||||
if patch is None:
|
||||
return weight
|
||||
self._applied.add(name)
|
||||
|
||||
device = self._compute_device(weight)
|
||||
|
||||
delta: torch.Tensor | None = None
|
||||
if patch.low_rank:
|
||||
if patch.layout in _QKV_SLOTS:
|
||||
per_slot = {
|
||||
slot: self._widen(b, device) @ self._widen(a, device) for slot, (a, b) in patch.low_rank.items()
|
||||
}
|
||||
delta = _place_in_grouped_qkv(per_slot, head_dim=self._head_dim)
|
||||
else:
|
||||
a, b = patch.low_rank[patch.layout]
|
||||
if patch.layout == _SWAP_HALVES:
|
||||
# Permuting the rows of B permutes the rows of the product,
|
||||
# so swap the rank-64 factor instead of the full delta.
|
||||
b = _swap_halves(b)
|
||||
delta = self._widen(b, device) @ self._widen(a, device)
|
||||
if patch.diff is not None:
|
||||
diff = self._widen(patch.diff, device)
|
||||
delta = diff if delta is None else delta + diff
|
||||
if delta is None:
|
||||
return weight
|
||||
if delta.shape != weight.shape:
|
||||
raise FastH3AdapterError(
|
||||
f"FastH3 delta for {name} has shape {tuple(delta.shape)}, parameter is {tuple(weight.shape)}"
|
||||
)
|
||||
# Leave the result on the compute device. These weights are bound for
|
||||
# the accelerator anyway, so returning them to host memory would pay a
|
||||
# device-to-host copy of the whole checkpoint only for the loader to
|
||||
# send it straight back: measured at 152s against 15s for 60 GiB of
|
||||
# patched projections, against 17s for the unavoidable upload alone.
|
||||
# Fold the base weight into the freshly built delta in place. Promoting
|
||||
# the weight to float32 on its own would allocate two more buffers the
|
||||
# size of the parameter, and H3's largest patched projection is 0.5 GiB.
|
||||
return delta.add_(weight.to(device, non_blocking=True)).to(weight.dtype)
|
||||
|
||||
def apply(self, weights: Iterable[tuple[str, torch.Tensor]]) -> Iterator[tuple[str, torch.Tensor]]:
|
||||
"""Fuse every streamed checkpoint tensor on its way into the model."""
|
||||
if self._applied:
|
||||
# ``validate_fully_applied`` released the deltas, so a second stream
|
||||
# would fuse nothing and then pass its own completeness check: the
|
||||
# server would serve base H3 weights on the student's ladder.
|
||||
raise FastH3AdapterError(f"{self._source} has already been fused into this checkpoint")
|
||||
for name, weight in weights:
|
||||
yield name, self.fuse(name, weight)
|
||||
|
||||
def validate_fully_applied(self) -> None:
|
||||
"""Close the fusion: every edit must have met its parameter.
|
||||
|
||||
A silently unapplied delta is the failure mode that matters here: the
|
||||
model would load and generate, just not as the distilled student. The
|
||||
weights are loaded once, so the mapped payloads are dropped afterwards
|
||||
rather than held for the life of the process.
|
||||
"""
|
||||
missing = sorted(set(self._patches) - self._applied)
|
||||
if missing:
|
||||
raise FastH3AdapterError(
|
||||
f"FastH3 adapter edits {len(missing)} parameters the checkpoint never provided: {missing[:5]}"
|
||||
)
|
||||
for patch in self._patches.values():
|
||||
patch.low_rank.clear()
|
||||
patch.diff = None
|
||||
|
||||
def check_serving_contract(
|
||||
self,
|
||||
*,
|
||||
partition: str,
|
||||
od_config: Any,
|
||||
video_shift: float,
|
||||
audio_shift: float,
|
||||
) -> None:
|
||||
"""Hold a starting server to the ladder this student was trained on."""
|
||||
if partition == "ref2va":
|
||||
raise ValueError("FastH3 preview v1 distills T2VA only, so it cannot serve a Ref2VA partition")
|
||||
offloads = [
|
||||
flag
|
||||
for flag in ("enable_cpu_offload", "enable_layerwise_offload", "enable_distributed_layerwise_offload")
|
||||
if getattr(od_config, flag, False)
|
||||
]
|
||||
if offloads:
|
||||
# A host-weight plan installs the transformer without going through
|
||||
# load_weights(), which is where the fusion and its completeness
|
||||
# check live. Serving base H3 weights under a four-step schedule
|
||||
# would otherwise degrade output with nothing to signal it.
|
||||
raise ValueError(
|
||||
f"FastH3 is fused while the checkpoint streams in, so it cannot be combined with "
|
||||
f"{sorted(offloads)}. Serve it without offload."
|
||||
)
|
||||
if self.requires_vsa:
|
||||
raise ValueError(
|
||||
f"{self.source} is a Video Sparse Attention variant of FastH3, and its "
|
||||
"compression gates have no counterpart in vLLM-Omni's dense H3 attention. Use the "
|
||||
"release's Dense variant until VSA lands for H3's packed sequence."
|
||||
)
|
||||
logger.info(
|
||||
"FastH3 adapter active: sigma points %s for %d transformer forwards, "
|
||||
"flow_shift=%g, audio_flow_shift=%g, tasks=%s",
|
||||
list(self.base_schedule),
|
||||
FASTH3_DENOISE_STEPS,
|
||||
video_shift,
|
||||
audio_shift,
|
||||
sorted(FASTH3_SUPPORTED_TASKS),
|
||||
)
|
||||
|
||||
def check_task(self, task: str) -> None:
|
||||
"""Refuse a task this preview never distilled."""
|
||||
if task not in FASTH3_SUPPORTED_TASKS:
|
||||
raise OmniClientError(
|
||||
f"FastH3 preview v1 distills {sorted(FASTH3_SUPPORTED_TASKS)} only, got task={task!r}"
|
||||
)
|
||||
|
||||
def check_request(self, sampling: Any, *, video_shift: float, audio_shift: float) -> None:
|
||||
"""Refuse a request that would sample the student off its rungs."""
|
||||
if sampling.lora_request is not None:
|
||||
# The adapter is already in the weights and the dynamic LoRA manager
|
||||
# is skipped, so nothing would apply the requested one. Serving the
|
||||
# request anyway would quietly ignore it.
|
||||
raise OmniClientError(
|
||||
f"this server fused {self.source} into the checkpoint at startup, so per-request "
|
||||
"lora is unavailable; drop the lora field"
|
||||
)
|
||||
if int(sampling.num_inference_steps or 0) != FASTH3_DENOISE_STEPS:
|
||||
raise OmniClientError(
|
||||
f"FastH3 is a four-step student and requires num_inference_steps={FASTH3_DENOISE_STEPS} "
|
||||
"(one transformer forward per sigma interval)"
|
||||
)
|
||||
# The checkpoint's per-modality shifts turn the release's positions into
|
||||
# the noise levels the student was distilled at, so a request that moves
|
||||
# them samples where it was never trained.
|
||||
extra = sampling.extra_args or {}
|
||||
for key, expected in (("flow_shift", video_shift), ("audio_flow_shift", audio_shift)):
|
||||
try:
|
||||
requested = float(extra.get(key, expected))
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise OmniClientError(f"FastH3 requires {key}={expected:g}") from exc
|
||||
if not math.isclose(requested, expected):
|
||||
raise OmniClientError(f"FastH3 requires {key}={expected:g}, got {requested:g}")
|
||||
|
||||
|
||||
def resolve_fasth3_fusion(od_config: Any, transformer: MiniMaxH3DiTModel) -> FastH3WeightFusion | None:
|
||||
"""Claim ``--lora-path`` when it points at a FastH3 adapter.
|
||||
|
||||
FastH3 rewrites RMSNorm weights and biases, so it cannot be expressed as a
|
||||
request-switchable LoRA and is fused into the checkpoint instead. Any other
|
||||
artifact returns None here and stays on the dynamic LoRA route.
|
||||
"""
|
||||
lora_path = getattr(od_config, "lora_path", None)
|
||||
if isinstance(lora_path, (list, tuple)):
|
||||
if len(lora_path) != 1:
|
||||
return None
|
||||
lora_path = lora_path[0]
|
||||
if not lora_path:
|
||||
return None
|
||||
# Placing a delta into the fused QKV parameter needs the head size, and
|
||||
# completeness is judged against the model's depth, so the architecture is
|
||||
# only read once an adapter is actually configured.
|
||||
arch = transformer.arch
|
||||
return FastH3WeightFusion.from_path(
|
||||
lora_path,
|
||||
head_dim=arch.attention_head_dim,
|
||||
num_blocks=arch.num_layers,
|
||||
num_refiner_blocks=arch.token_refiner_num_layers,
|
||||
)
|
||||
|
||||
|
||||
def _safetensors_metadata(path: Path) -> Mapping[str, str]:
|
||||
"""The header metadata of a safetensors file, or ``{}`` if unreadable."""
|
||||
try:
|
||||
with safe_open(path, framework="pt", device="cpu") as checkpoint:
|
||||
return checkpoint.metadata() or {}
|
||||
except Exception: # noqa: BLE001 - a path that will not open is simply not ours
|
||||
return {}
|
||||
|
||||
|
||||
def _resolve_adapter_file(path: str | Path) -> Path | None:
|
||||
"""Find the single adapter file at ``path``, or ``None``."""
|
||||
candidate = Path(path)
|
||||
if candidate.is_file():
|
||||
return candidate if candidate.suffix == ".safetensors" else None
|
||||
if not candidate.is_dir():
|
||||
return None
|
||||
named = candidate / "adapter_model.safetensors"
|
||||
if named.is_file():
|
||||
return named
|
||||
files = sorted(candidate.glob("*.safetensors"))
|
||||
if len(files) == 1:
|
||||
return files[0]
|
||||
# The published repository bundles four variants under one root, marked by an
|
||||
# adapter_manifest.json. Such a bundle is ambiguous rather than loadable - but
|
||||
# only hold that against a directory that actually carries FastH3 artifacts;
|
||||
# an unrelated multi-shard LoRA stays on the dynamic route via ``None``.
|
||||
if (candidate / FASTH3_MANIFEST).is_file() or any(
|
||||
_is_fasth3_release(_safetensors_metadata(file)) for file in files
|
||||
):
|
||||
raise FastH3AdapterError(
|
||||
f"{candidate} holds several FastH3 adapters; point --lora-path at one variant "
|
||||
"(for example dense-datafree/adapter_model.safetensors)"
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
__all__ = [
|
||||
"FASTH3_BASE_MODEL",
|
||||
"FASTH3_BASE_SCHEDULE",
|
||||
"FASTH3_DENOISE_STEPS",
|
||||
"FASTH3_FORMAT",
|
||||
"FASTH3_SUPPORTED_TASKS",
|
||||
"FastH3AdapterError",
|
||||
"FastH3WeightFusion",
|
||||
"resolve_fasth3_fusion",
|
||||
]
|
||||
@@ -117,6 +117,7 @@ from .denoise_loop import (
|
||||
minimax_h3_publish_denoise_progress,
|
||||
)
|
||||
from .encoder import MiniMaxH3Qwen3VLEncoder
|
||||
from .fasth3 import FastH3WeightFusion, resolve_fasth3_fusion
|
||||
from .lora import load_minimax_h3_turbo_lora
|
||||
from .minimax_h3_transformer import (
|
||||
MiniMaxH3Attention,
|
||||
@@ -354,17 +355,24 @@ def _minimax_h3_post_process(output, output_type: str = "np"):
|
||||
|
||||
The callable crosses the multiprocessing result queue, so it must remain a
|
||||
module-level function that the standard pickle module can resolve.
|
||||
|
||||
``_prepare_minimax_h3_video_output`` already quantises the video to uint8
|
||||
frames on the accelerator, so there is nothing left to scale or transpose
|
||||
here.
|
||||
"""
|
||||
if not isinstance(output, tuple) or len(output) != 2:
|
||||
return output
|
||||
video, audio = output
|
||||
if video.dtype != torch.uint8 or video.ndim != 5 or video.shape[-1] not in (3, 4):
|
||||
# Float or channel-first frames would reach the muxer as a black or
|
||||
# banded video rather than as an error.
|
||||
raise ValueError(
|
||||
f"MiniMax-H3 post-processing expects (B, T, H, W, C) uint8, got {tuple(video.shape)} {video.dtype}"
|
||||
)
|
||||
if output_type == "latent":
|
||||
return output
|
||||
if output_type == "np":
|
||||
if video.dtype == torch.uint8 and video.ndim == 5 and video.shape[-1] == 3:
|
||||
video = video.detach().cpu().numpy()
|
||||
else:
|
||||
video = video.detach().float().cpu().permute(0, 2, 3, 4, 1).clamp(0, 1).numpy()
|
||||
video = video.detach().cpu().numpy()
|
||||
audio = audio.detach().float().cpu().numpy()
|
||||
video = [sample for sample in video]
|
||||
return {
|
||||
@@ -678,6 +686,8 @@ class MiniMaxH3Pipeline(
|
||||
# Only distilled releases pin a schedule, so the default keeps the legacy
|
||||
# uniform path available to partially constructed pipelines.
|
||||
_base_schedule_by_partition: ClassVar[Mapping[str, DMD2SigmaSchedule | None]] = {}
|
||||
# Set from --lora-path during construction; absent means no FastH3 adapter.
|
||||
_fasth3: FastH3WeightFusion | None = None
|
||||
|
||||
def _load_diffusion_lora_adapter(
|
||||
self,
|
||||
@@ -948,6 +958,15 @@ class MiniMaxH3Pipeline(
|
||||
quant_config=transformer_quant_config,
|
||||
)
|
||||
|
||||
self._fasth3 = resolve_fasth3_fusion(od_config, self.transformer)
|
||||
if self._fasth3 is not None:
|
||||
self._fasth3.check_serving_contract(
|
||||
partition=self.partition,
|
||||
od_config=od_config,
|
||||
video_shift=self.default_video_shift,
|
||||
audio_shift=self.default_audio_shift,
|
||||
)
|
||||
|
||||
if self.load_text_encoder:
|
||||
self.tokenizer = Qwen2TokenizerFast.from_pretrained(
|
||||
str(model_path),
|
||||
@@ -1056,7 +1075,12 @@ class MiniMaxH3Pipeline(
|
||||
raise ValueError(f"MiniMax-H3 weight source {prefix!r} is not contiguous")
|
||||
loaded_prefixes.add(prefix)
|
||||
component = getattr(self, prefix.removesuffix("."))
|
||||
loaded = component.load_weights((name[len(prefix) :], tensor) for name, tensor in grouped_weights)
|
||||
stream = ((name[len(prefix) :], tensor) for name, tensor in grouped_weights)
|
||||
if prefix == "transformer." and self._fasth3 is not None:
|
||||
# Fuse before the model shards anything, which is also the only
|
||||
# point where the checkpoint's fused QKV/MLP layouts are intact.
|
||||
stream = self._fasth3.apply(stream)
|
||||
loaded = component.load_weights(stream)
|
||||
if prefix != "text_encoder.":
|
||||
component.post_load_weights()
|
||||
loaded_with_prefix.update(prefix + name for name in loaded)
|
||||
@@ -1069,13 +1093,49 @@ class MiniMaxH3Pipeline(
|
||||
if component is None:
|
||||
continue
|
||||
loaded_with_prefix.update(f"{component_name}.{name}" for name, _ in component.named_parameters())
|
||||
if self._fasth3 is not None:
|
||||
self._fasth3.validate_fully_applied()
|
||||
return loaded_with_prefix
|
||||
|
||||
@property
|
||||
def lora_is_fused(self) -> bool:
|
||||
"""True when --lora-path was consumed as a load-time weight fusion."""
|
||||
return self._fasth3 is not None
|
||||
|
||||
def _transformer_for_task(self, task: str) -> MiniMaxH3DiTModel:
|
||||
if task == "ref2va" and hasattr(self, "transformers_ref"):
|
||||
return self.transformers_ref
|
||||
return self.transformer
|
||||
|
||||
def _resolve_sigma_positions(self, task: str, sampling: Any) -> tuple[tuple[float, ...] | None, int]:
|
||||
"""Pick the rectified-flow positions this request denoises on.
|
||||
|
||||
Returns them explicitly, or ``None`` to leave the uniform ladder to be
|
||||
derived from the step count, together with the count the rest of the
|
||||
request speaks in.
|
||||
"""
|
||||
if self._fasth3 is not None:
|
||||
# A fused student carries its own positions; the checkpoint
|
||||
# underneath it is the many-step teacher, whose schedule does not
|
||||
# apply. Its five points bound four transformer forwards, and
|
||||
# forwards is the unit ``check_request``, the pinned-schedule branch
|
||||
# below and Cache-DiT all speak in.
|
||||
positions = self._fasth3.base_schedule
|
||||
return positions, len(positions) - 1
|
||||
sigma_schedule = self._sigma_schedule_for_request(sampling, task)
|
||||
if sigma_schedule is None:
|
||||
return None, int(sampling.num_inference_steps or 50)
|
||||
# The schedule lists sigma boundaries; the denoise loop runs one step
|
||||
# per interval, and that count is what requests and Cache-DiT speak in.
|
||||
num_steps = sigma_schedule.num_inference_steps
|
||||
requested_steps = sampling.num_inference_steps
|
||||
if requested_steps is not None and int(requested_steps) != num_steps:
|
||||
raise OmniClientError(
|
||||
"this MiniMax H3 checkpoint pins a distilled sigma schedule; num_inference_steps "
|
||||
f"must be {num_steps} or omitted, got {int(requested_steps)}"
|
||||
)
|
||||
return sigma_schedule.base_schedule, num_steps
|
||||
|
||||
def _base_schedule_for_task(self, task: str) -> DMD2SigmaSchedule | None:
|
||||
"""Return the distilled schedule of the partition that serves ``task``."""
|
||||
partition = "ref2va" if task == "ref2va" else "fl2va"
|
||||
@@ -1109,6 +1169,8 @@ class MiniMaxH3Pipeline(
|
||||
raise OmniClientError("MiniMax-H3 Turbo LoRA supports T2VA/FL2VA requests only")
|
||||
if has_native_lora and task != "t2va":
|
||||
raise OmniClientError("MiniMax-H3 native LoRA supports T2VA requests only")
|
||||
if self._fasth3 is not None:
|
||||
self._fasth3.check_task(task)
|
||||
return task
|
||||
|
||||
def _resolve_shape(
|
||||
@@ -2121,6 +2183,12 @@ class MiniMaxH3Pipeline(
|
||||
self._validate_turbo_sampling(sampling)
|
||||
if has_native_lora:
|
||||
self._validate_native_sampling(sampling, task=task)
|
||||
if self._fasth3 is not None:
|
||||
self._fasth3.check_request(
|
||||
sampling,
|
||||
video_shift=self.default_video_shift,
|
||||
audio_shift=self.default_audio_shift,
|
||||
)
|
||||
|
||||
raw_image = multi_modal_data.get("image")
|
||||
raw_videos = multi_modal_data.get("video")
|
||||
@@ -2337,22 +2405,7 @@ class MiniMaxH3Pipeline(
|
||||
ref_audio_t = audio_lengths[0]
|
||||
|
||||
seed = int(sampling.seed if sampling.seed is not None else 42)
|
||||
sigma_schedule = self._sigma_schedule_for_request(sampling, task)
|
||||
if sigma_schedule is None:
|
||||
base_schedule = None
|
||||
num_steps = int(sampling.num_inference_steps or 50)
|
||||
else:
|
||||
# The schedule lists sigma boundaries; the denoise loop runs one
|
||||
# step per interval, and that count is what requests and Cache-DiT
|
||||
# speak in.
|
||||
base_schedule = sigma_schedule.base_schedule
|
||||
num_steps = sigma_schedule.num_inference_steps
|
||||
requested_steps = sampling.num_inference_steps
|
||||
if requested_steps is not None and int(requested_steps) != num_steps:
|
||||
raise OmniClientError(
|
||||
"this MiniMax H3 checkpoint pins a distilled sigma schedule; num_inference_steps "
|
||||
f"must be {num_steps} or omitted, got {int(requested_steps)}"
|
||||
)
|
||||
base_schedule, num_steps = self._resolve_sigma_positions(task, sampling)
|
||||
video_shift = float(extra.get("flow_shift", self.default_video_shift))
|
||||
audio_shift = float(extra.get("audio_flow_shift", self.default_audio_shift))
|
||||
quality_plan = self._quality_policy.resolve(
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -37,8 +37,11 @@ class DMD2SigmaSchedule:
|
||||
# have to be rejected before the monotonicity check rather than after it.
|
||||
if any(not math.isfinite(value) for value in values):
|
||||
raise ValueError("DMD2 base_schedule entries must be finite")
|
||||
if values[0] != 1.0 or values[-1] != 0.0:
|
||||
raise ValueError("DMD2 base_schedule must start at 1.0 and end at 0.0")
|
||||
# A student whose training capped the noise level (``max_timestep_ratio``)
|
||||
# starts just below 1.0 rather than at it, so the opening position is
|
||||
# bounded rather than pinned; the terminal one is always clean latents.
|
||||
if not 0.0 < values[0] <= 1.0 or values[-1] != 0.0:
|
||||
raise ValueError("DMD2 base_schedule must start in (0.0, 1.0] and end at 0.0")
|
||||
if any(curr <= nxt for curr, nxt in zip(values, values[1:], strict=False)):
|
||||
raise ValueError("DMD2 base_schedule must be strictly decreasing")
|
||||
|
||||
|
||||
@@ -538,7 +538,16 @@ class DiffusionWorker:
|
||||
|
||||
def init_lora_manager(self) -> None:
|
||||
"""Initialize the LoRA manager for this worker."""
|
||||
if self.model_runner.pipeline is None:
|
||||
pipeline = self.model_runner.pipeline
|
||||
if pipeline is None:
|
||||
return
|
||||
|
||||
# A release whose weights cannot be expressed as switchable LoRA layers
|
||||
# is fused into the checkpoint while the pipeline loads. There is then
|
||||
# no adapter left to register, and handing the same path to the manager
|
||||
# would only fail on a format it does not accept.
|
||||
if getattr(pipeline, "lora_is_fused", False):
|
||||
logger.info("LoRA was fused into the checkpoint at load time; skipping the dynamic LoRA manager.")
|
||||
return
|
||||
|
||||
lora_path = self.od_config.lora_path
|
||||
@@ -739,11 +748,15 @@ class DiffusionWorker:
|
||||
logger.warning("LoRA activation skipped: %s", exc)
|
||||
|
||||
def remove_lora(self, adapter_id: int) -> bool:
|
||||
if self.lora_manager is None:
|
||||
return False
|
||||
return self.lora_manager.remove_adapter(adapter_id)
|
||||
|
||||
def add_lora(self, lora_request: LoRARequest) -> bool:
|
||||
# NOTE (Alex): We have not implemented the API routing
|
||||
# for the frontend server yet.
|
||||
if self.lora_manager is None:
|
||||
return False
|
||||
return self.lora_manager.add_adapter(lora_request)
|
||||
|
||||
def submit_interaction(
|
||||
@@ -756,9 +769,13 @@ class DiffusionWorker:
|
||||
self.model_runner.submit_interaction(request_id, interaction)
|
||||
|
||||
def list_loras(self) -> list[int]:
|
||||
if self.lora_manager is None:
|
||||
return []
|
||||
return self.lora_manager.list_adapters()
|
||||
|
||||
def pin_lora(self, adapter_id: int) -> bool:
|
||||
if self.lora_manager is None:
|
||||
return False
|
||||
return self.lora_manager.pin_adapter(adapter_id)
|
||||
|
||||
def sleep(self, level: int = 1) -> int:
|
||||
|
||||
@@ -396,7 +396,7 @@ def _normalize_video_tensor(video_tensor: torch.Tensor) -> np.ndarray:
|
||||
# Cast to float32 first: bf16 (e.g. SANA-WM's refiner output) has no
|
||||
# numpy dtype, so ``.numpy()`` below raises on it.
|
||||
video_tensor = video_tensor.float().clamp(-1, 1) * 0.5 + 0.5
|
||||
else:
|
||||
elif video_tensor.dtype != torch.uint8:
|
||||
video_tensor = video_tensor.to(torch.float32) / 255.0
|
||||
video_array = video_tensor.numpy()
|
||||
return _normalize_single_video_array(video_array)
|
||||
@@ -419,7 +419,7 @@ def _normalize_single_video_array(video_array: np.ndarray) -> np.ndarray:
|
||||
if np.issubdtype(video_array.dtype, np.floating):
|
||||
if video_array.size and (video_array.min() < 0.0 or video_array.max() > 1.0):
|
||||
video_array = np.clip(video_array, -1.0, 1.0) * 0.5 + 0.5
|
||||
elif np.issubdtype(video_array.dtype, np.integer):
|
||||
elif video_array.dtype != np.uint8 and np.issubdtype(video_array.dtype, np.integer):
|
||||
video_array = video_array.astype(np.float32) / 255.0
|
||||
return video_array
|
||||
|
||||
@@ -435,7 +435,7 @@ def _normalize_video_array(video_array: np.ndarray) -> list[np.ndarray] | np.nda
|
||||
|
||||
|
||||
def _normalize_frames(frames: list[Any]) -> list[np.ndarray]:
|
||||
"""Normalize a list of frames into numpy arrays with values in [0,1]."""
|
||||
"""Normalize a list of frames into numpy arrays, uint8 ones unchanged."""
|
||||
normalized: list[np.ndarray] = []
|
||||
for frame in frames:
|
||||
if isinstance(frame, torch.Tensor):
|
||||
@@ -453,7 +453,7 @@ def _normalize_frames(frames: list[Any]) -> list[np.ndarray]:
|
||||
if np.issubdtype(frame_array.dtype, np.floating):
|
||||
if frame_array.size and (frame_array.min() < 0.0 or frame_array.max() > 1.0):
|
||||
frame_array = np.clip(frame_array, -1.0, 1.0) * 0.5 + 0.5
|
||||
elif np.issubdtype(frame_array.dtype, np.integer):
|
||||
elif frame_array.dtype != np.uint8 and np.issubdtype(frame_array.dtype, np.integer):
|
||||
frame_array = frame_array.astype(np.float32) / 255.0
|
||||
|
||||
normalized.append(frame_array)
|
||||
@@ -461,7 +461,14 @@ def _normalize_frames(frames: list[Any]) -> list[np.ndarray]:
|
||||
|
||||
|
||||
def _coerce_video_to_frames(video: Any) -> list[np.ndarray]:
|
||||
"""Convert a video payload into a list of normalized float32 frames."""
|
||||
"""Convert a video payload into a list of normalized frames.
|
||||
|
||||
Frames are float32 in [0, 1], except uint8 ones, which pass through: the
|
||||
muxer's own dtype is uint8, so normalising here would only pay for a
|
||||
full-size conversion each way. The direct planar path additionally needs
|
||||
contiguous channel planes and falls back for interleaved RGB whatever the
|
||||
dtype; the standard muxer takes the uint8 frames as they are.
|
||||
"""
|
||||
if isinstance(video, torch.Tensor):
|
||||
video_array = _normalize_video_tensor(video)
|
||||
return list(video_array)
|
||||
|
||||
Reference in new issue
Block a user