diff --git a/docs/user_guide/diffusion/lora.md b/docs/user_guide/diffusion/lora.md index 44705129f..3c8582af0 100644 --- a/docs/user_guide/diffusion/lora.md +++ b/docs/user_guide/diffusion/lora.md @@ -12,15 +12,17 @@ LoRA adapters are lightweight, model-specific fine-tuning weights that can be ap ## LoRA Adapter Format ### PEFT (Parameter-Efficient Fine-Tuning) format (default) + A typical PEFT-format LoRA adapter directory structure: -``` +```text lora_adapter/ ├── adapter_config.json └── adapter_model.safetensors ``` The `adapter_config.json` file contains metadata about the LoRA adapter, including: + - `r`: LoRA rank - `lora_alpha`: LoRA alpha scaling factor - `target_modules`: List of module names to apply LoRA to @@ -67,7 +69,7 @@ For distilled few-step LoRAs, pass `lora_backend="distill"` together with one or For Qwen-Image and Wan pipelines, the distill backend calls `pipeline.load_lora_weights(...)` during worker initialization. | Pipeline | Supported distilled LoRA repo | Notes | -|----------|--------------------------------|-------| +| ---------- | -------------------------------- | ------- | | `QwenImagePipeline` | `lightx2v/Qwen-Image-2512-Lightning` | Used with Qwen-Image-2512 Lightning-style few-step inference. | | `Wan22Pipeline` | `lightx2v/Wan2.1-Distill-Loras`, `lightx2v/Wan2.2-Distill-Loras` | Wan2.1 uses one LoRA file. For dual-transformer Wan2.2 MoE, pass high-noise then low-noise LoRA files. | | `Wan22I2VPipeline` | `lightx2v/Wan2.2-Distill-Loras` | For dual-transformer Wan2.2 MoE, pass high-noise then low-noise LoRA files. | @@ -240,6 +242,37 @@ Notes: - Output quality and speed depend on the replacement checkpoints and sampling params you choose. - If you only need to fuse distilled LoRAs into a Wan2.2 checkpoint at load time (without the full LightX2V convert + assemble pipeline), you can instead pass them directly via `--lora-backend distill --lora-path .safetensors .safetensors`. See the [Distill backend](#distill-backend-fuse-distilled-lora-at-init) section above. +## MiniMax-H3 adapter-declared schedules + +MiniMax-H3 supports two few-step mechanisms that must not be conflated: + +- **Checkpoint-pinned schedule**: a merged release writes `base_schedule` into + `model_index.json`. Requests must pass `num_inference_steps` as the interval + count (for example `4` for `[1.0, 0.7, 0.4, 0.15, 0.0]`). +- **Runtime Turbo LoRA**: LightX2V Turbo artifacts request five sigma points and + enforce `flow_shift=6`, `audio_flow_shift=3`. +- **Runtime native LoRA**: FlashGen-style artifacts declare + `key_format=minimax-h3-native` and embed `base_schedule` in safetensors + metadata. When active, the adapter schedule overrides the base checkpoint and + requests must use the interval-count contract (`num_inference_steps=4`). + Request-mode generation may omit the field and take the count from the adapter + schedule; step execution requires it explicitly, because the step scheduler + derives the total step count from the request before the adapter schedule is + known. + +Native artifacts also declare `qkv_layout=grouped`. The H3 loader reorders fused +`qkv_proj` LoRA rows with the same `_reorder_grouped_qkv_to_qkv` helper used for +base-weight loading, then binds the packed Q/K/V slices through the legacy PEFT +manager without modifying `DiffusionLoRAManager`. + +Both runtime adapters run with distributed layerwise offload in request-mode +generation, where the manager keeps the LoRA A/B buffers resident on the compute +device while DLO streams the base blocks. MiniMax-H3 step execution still +rejects DLO, so `--step-execution` cannot be combined with +`--enable-distributed-layerwise-offload`. Model-level CPU offload and standard +layerwise offload are rejected in every mode, because the dynamic LoRA tensors +are neither parameters nor registered buffers and therefore do not participate +in those weight lifecycles. ## See Also diff --git a/recipes/MiniMaxAI/MiniMax-H3.md b/recipes/MiniMaxAI/MiniMax-H3.md index 915c71739..d022cce51 100644 --- a/recipes/MiniMaxAI/MiniMax-H3.md +++ b/recipes/MiniMaxAI/MiniMax-H3.md @@ -860,6 +860,63 @@ fixed HBM usage. Model-level and standard layerwise offload remain unsupported. The five requested sigma points produce the four denoiser evaluations expected by the Turbo artifact. +### FlashGen native LoRA + +The FlashGen 4-step T2VA artifact uses the native MiniMax-H3 module layout and +declares its distilled sigma schedule in safetensors metadata. It is published on +[ModelScope](https://modelscope.cn/models/FlashGen/Minimax-H3-4step-lora-flashgen): + +```text +FlashGen/Minimax-H3-4step-lora-flashgen/minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors +``` + +Download only that file: + +```bash +python -m pip install modelscope +export FLASHGEN_DIR=/path/to/minimax-h3-flashgen-lora +export FLASHGEN_FILE=minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors +modelscope download FlashGen/Minimax-H3-4step-lora-flashgen \ + --local_dir "${FLASHGEN_DIR}" \ + --include "${FLASHGEN_FILE}" +export FLASHGEN_LORA="${FLASHGEN_DIR}/${FLASHGEN_FILE}" +``` + +Start from a non-offloaded or DLO FL2VA server command and add +`--task-type fl2va --lora-backend peft --lora-path "${FLASHGEN_LORA}"`. +Each request must use T2VA and the distilled interval-count contract: + +```bash +-F 'num_inference_steps=4' \ +-F 'extra_params={"task":"t2va","duration":5.2}' \ +-F "lora={\"name\":\"h3-flashgen-v1.0\",\"path\":\"${FLASHGEN_LORA}\",\"scale\":1.0}" +``` + +This path rejects Ref2VA and checkpoints that already pin `base_schedule` in +`model_index.json`. The adapter metadata carries +`base_schedule=1.0,0.7,0.4,0.15,0.0`, so `num_inference_steps=4` means four +denoiser evaluations, not five sigma points. Request-mode generation may omit +the field and take the count from the adapter schedule; `--step-execution` +requires it explicitly, because the step scheduler reads the total step count +off the request at admission, before the adapter schedule is known. + +DLO is supported in request-mode generation on the same terms as the Turbo +adapter: the request-switchable LoRA A/B buffers stay resident on the +accelerator while DLO streams only the base blocks, so budget for that +additional fixed HBM usage. The native artifact is rank 64 over 259 target +modules, and its packed `qkv_proj` and `fc1` layers reuse the full-input A +tensor per slice while B carries slice-local output rows, so the resident +footprint exceeds the on-disk payload; measure it for your parallel layout +rather than assuming the checkpoint size. Pure Ulysses replicates the adapter +on every rank, while DiT tensor parallelism shards the B buffers. Model-level +and standard layerwise offload remain unsupported, and `--step-execution` +cannot be combined with `--enable-distributed-layerwise-offload`. + +To validate a deployment, post the same fixed-seed T2VA request twice with the +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. + ## Key parameters | Parameter | Recommended value | Notes | diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_lora.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_lora.py index 22581786d..5f98a88d5 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_lora.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_lora.py @@ -223,6 +223,8 @@ def test_h3_turbo_allows_distributed_layerwise_offload(monkeypatch): torch.nn.Module.__init__(pipeline) pipeline.partition = "fl2va" pipeline._turbo_lora_adapter_ids = set() + pipeline._native_lora_adapter_ids = set() + pipeline._lora_sigma_schedules = {} pipeline.od_config = SimpleNamespace( enable_cpu_offload=False, enable_layerwise_offload=False, @@ -321,6 +323,8 @@ def test_only_an_active_recognized_turbo_adapter_restricts_ref2va(monkeypatch): pipeline.partition = "combined" pipeline.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"}) pipeline._turbo_lora_adapter_ids = set() + pipeline._native_lora_adapter_ids = set() + pipeline._lora_sigma_schedules = {} request = _request("generic-peft") def load(): diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_native_lora.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_native_lora.py new file mode 100644 index 000000000..c4d581c2c --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_native_lora.py @@ -0,0 +1,860 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +from safetensors.torch import save_file +from vllm.lora.lora_weights import PackedLoRALayerWeights + +from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager +from vllm_omni.diffusion.models.minimax_h3.npu import lora as lora_module +from vllm_omni.diffusion.models.minimax_h3.npu.lora import load_minimax_h3_native_lora +from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule +from vllm_omni.errors import OmniClientError +from vllm_omni.lora.request import LoRARequest + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + +_TINY_HIDDEN = 2 +_TINY_FFN = 2 +_TINY_TIME = 2 +_TINY_BLOCK_ADALN = 4 +_TINY_FINAL_ADALN = 3 +_TINY_RANK = 4 +_TINY_NUM_QUERY_GROUPS = 2 +_TINY_HEADS_PER_GROUP = 1 +_TINY_HEAD_DIM = 1 +_TINY_QKV_SLICE = _TINY_NUM_QUERY_GROUPS * _TINY_HEADS_PER_GROUP * _TINY_HEAD_DIM +# Grouped qkv stores (heads_per_group + 2) * head_dim rows per query group, so +# the attention inner size follows from the head layout and cannot be picked +# independently: 3 * inner must equal the grouped row count. +_TINY_INNER = _TINY_QKV_SLICE + + +@pytest.fixture(autouse=True) +def _use_tiny_native_dimensions(monkeypatch): + monkeypatch.setattr(lora_module, "_NATIVE_RANK", _TINY_RANK) + monkeypatch.setattr(lora_module, "_NATIVE_ALPHA", _TINY_RANK) + monkeypatch.setattr(lora_module, "_NATIVE_HIDDEN_SIZE", _TINY_HIDDEN) + monkeypatch.setattr(lora_module, "_NATIVE_ATTENTION_INNER_SIZE", _TINY_INNER) + monkeypatch.setattr(lora_module, "_NATIVE_FFN_HIDDEN_SIZE", _TINY_FFN) + monkeypatch.setattr(lora_module, "_NATIVE_TIME_EMBED_DIM", _TINY_TIME) + monkeypatch.setattr(lora_module, "_NATIVE_BLOCK_ADALN_OUT", _TINY_BLOCK_ADALN) + monkeypatch.setattr(lora_module, "_NATIVE_FINAL_ADALN_OUT", _TINY_FINAL_ADALN) + monkeypatch.setattr(lora_module, "_NATIVE_NUM_QUERY_GROUPS", _TINY_NUM_QUERY_GROUPS) + monkeypatch.setattr(lora_module, "_NATIVE_HEADS_PER_GROUP", _TINY_HEADS_PER_GROUP) + monkeypatch.setattr(lora_module, "_NATIVE_HEAD_DIM", _TINY_HEAD_DIM) + monkeypatch.setattr(lora_module, "_NATIVE_QKV_SLICE", _TINY_QKV_SLICE) + monkeypatch.setattr( + lora_module, + "_NATIVE_TARGET_DIMS", + { + "attn.qkv_proj": (_TINY_HIDDEN, 3 * _TINY_INNER), + "attn.out_proj": (_TINY_INNER, _TINY_HIDDEN), + "mlp.fc1": (_TINY_HIDDEN, 2 * _TINY_FFN), + "mlp.fc2": (_TINY_FFN, _TINY_HIDDEN), + "adaln_proj.linear": (_TINY_TIME, _TINY_BLOCK_ADALN), + }, + ) + monkeypatch.setattr(lora_module, "_NATIVE_FINAL_ADALN_DIMS", (_TINY_TIME, _TINY_FINAL_ADALN)) + monkeypatch.setattr( + lora_module, + "_NATIVE_EXPECTED_TARGETS", + frozenset( + [ + *( + f"blocks.{block_index}.{suffix}" + for block_index in range(50) + for suffix in lora_module._NATIVE_TARGET_SUFFIXES + ) + ] + + [ + *( + f"token_refiner.blocks.{block_index}.{suffix}" + for block_index in range(2) + for suffix in lora_module._NATIVE_TOKEN_REFINER_SUFFIXES + ) + ] + + ["final_layer.adaln_proj.linear"] + ), + ) + + +def _request(path) -> LoRARequest: + return LoRARequest( + lora_name="flashgen", + lora_int_id=7, + lora_path=str(path), + ) + + +def _write_tiny_native( + path, + *, + omit_target: str | None = None, + shape_overrides: dict[str, tuple[int, int]] | None = None, + metadata: dict[str, str] | None = None, +) -> None: + rank = _TINY_RANK + tensors = {} + overrides = shape_overrides or {} + for target in sorted(lora_module._NATIVE_EXPECTED_TARGETS): + if target == omit_target: + continue + input_dim, output_dim = lora_module._native_target_dims(target) + raw_target = f"transformer.{target}" + a_name = f"{raw_target}.lora_A.default.weight" + b_name = f"{raw_target}.lora_B.default.weight" + tensors[a_name] = torch.ones(overrides.get(a_name, (rank, input_dim))) + if target.endswith(".attn.qkv_proj"): + grouped_rows = [] + for group in range(_TINY_NUM_QUERY_GROUPS): + grouped_rows.extend( + [ + torch.full((1, rank), float(group * 10 + 1)), + torch.full((1, rank), float(group * 10 + 2)), + torch.full((1, rank), float(group * 10 + 3)), + ] + ) + tensors[b_name] = torch.cat(grouped_rows, dim=0) + elif target.endswith(".mlp.fc1"): + tensors[b_name] = torch.cat( + ( + torch.full((output_dim // 2, rank), 2.0), + torch.full((output_dim // 2, rank), 1.0), + ), + dim=0, + ) + else: + tensors[b_name] = torch.ones(overrides.get(b_name, (output_dim, rank))) + save_file( + tensors, + str(path), + metadata=metadata + or { + "format": "pt", + "key_format": "minimax-h3-native", + "qkv_layout": "grouped", + "lora_rank": str(rank), + "lora_alpha": str(rank), + "base_schedule": "1.0,0.7,0.4,0.15,0.0", + "tasks": "t2va", + }, + ) + + +def test_h3_native_loads_and_packs_qkv_and_fc1(tmp_path): + path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors" + _write_tiny_native(path) + + loaded = load_minimax_h3_native_lora( + partition="fl2va", + lora_request=_request(path), + lora_path=path, + dtype=torch.float32, + ) + + assert loaded is not None + lora_model, peft_helper, schedule = loaded + assert peft_helper.r == _TINY_RANK + assert schedule.num_inference_steps == 4 + assert len(lora_model.loras) == 259 + assert "blocks.0.adaln_proj.linear" in lora_model.loras + assert "final_layer.adaln_proj.linear" in lora_model.loras + + # The writer tags group g with q=g*10+1, k=g*10+2, v=g*10+3, so a correct + # grouped -> q/k/v permutation yields slices that are uniform per role. + qkv = lora_model.get_lora("blocks.0.attn.qkv_proj") + assert isinstance(qkv, PackedLoRALayerWeights) + torch.testing.assert_close(qkv.lora_b[0], torch.tensor([[1.0] * _TINY_RANK, [11.0] * _TINY_RANK])) + torch.testing.assert_close(qkv.lora_b[1], torch.tensor([[2.0] * _TINY_RANK, [12.0] * _TINY_RANK])) + torch.testing.assert_close(qkv.lora_b[2], torch.tensor([[3.0] * _TINY_RANK, [13.0] * _TINY_RANK])) + + fc1 = lora_model.get_lora("blocks.0.mlp.fc1") + assert isinstance(fc1, PackedLoRALayerWeights) + torch.testing.assert_close(fc1.lora_b[0], torch.full((_TINY_FFN, _TINY_RANK), 2.0)) + torch.testing.assert_close(fc1.lora_b[1], torch.full((_TINY_FFN, _TINY_RANK), 1.0)) + + +def test_h3_native_rejects_bad_metadata_and_ref2va(tmp_path): + path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors" + _write_tiny_native(path, metadata={"key_format": "minimax-h3-native", "qkv_layout": "runtime"}) + with pytest.raises(ValueError, match="qkv_layout='grouped'"): + load_minimax_h3_native_lora( + partition="fl2va", + lora_request=_request(path), + lora_path=path, + dtype=torch.float32, + ) + + # The exact-filename guard runs before the partition check, so the artifact + # has to keep the published name for this to reach the Ref2VA rejection. + valid_dir = tmp_path / "valid" + valid_dir.mkdir() + valid = valid_dir / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors" + _write_tiny_native(valid) + with pytest.raises(ValueError, match="supports T2VA only"): + load_minimax_h3_native_lora( + partition="ref2va", + lora_request=_request(valid), + lora_path=valid, + dtype=torch.float32, + ) + + renamed = tmp_path / "renamed.safetensors" + _write_tiny_native(renamed) + with pytest.raises(ValueError, match="supports only"): + load_minimax_h3_native_lora( + partition="fl2va", + lora_request=_request(renamed), + lora_path=renamed, + dtype=torch.float32, + ) + + +def test_h3_native_rejects_truncated_artifact(tmp_path): + path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors" + _write_tiny_native(path, omit_target="blocks.49.mlp.fc2") + with pytest.raises(ValueError, match="target set does not match"): + load_minimax_h3_native_lora( + partition="fl2va", + lora_request=_request(path), + lora_path=path, + dtype=torch.float32, + ) + + +def test_h3_native_rejects_schedule_with_wrong_interval_count(tmp_path): + """A mislabeled schedule must fail closed instead of silently changing steps.""" + path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors" + _write_tiny_native( + path, + metadata={ + "format": "pt", + "key_format": "minimax-h3-native", + "qkv_layout": "grouped", + "lora_rank": str(_TINY_RANK), + "lora_alpha": str(_TINY_RANK), + "base_schedule": "1.0,0.5,0.0", + "tasks": "t2va", + }, + ) + with pytest.raises(ValueError, match="4-interval base_schedule"): + load_minimax_h3_native_lora( + partition="fl2va", + lora_request=_request(path), + lora_path=path, + dtype=torch.float32, + ) + + +@pytest.mark.parametrize( + ("raw_schedule", "message"), + [ + (None, "requires safetensors metadata base_schedule"), + ("", "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"), + ], +) +def test_h3_native_rejects_malformed_schedule_metadata(tmp_path, raw_schedule, message): + """safetensors metadata is string-only, so the loader owns the parse contract.""" + metadata = { + "format": "pt", + "key_format": "minimax-h3-native", + "qkv_layout": "grouped", + "lora_rank": str(_TINY_RANK), + "lora_alpha": str(_TINY_RANK), + "tasks": "t2va", + } + if raw_schedule is not None: + metadata["base_schedule"] = raw_schedule + path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors" + _write_tiny_native(path, metadata=metadata) + + with pytest.raises(ValueError, match=message): + load_minimax_h3_native_lora( + partition="fl2va", + lora_request=_request(path), + lora_path=path, + dtype=torch.float32, + ) + + +def test_h3_native_declared_file_fails_closed_on_invalid_metadata(tmp_path): + path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors" + _write_tiny_native(path, metadata={"key_format": "other"}) + with pytest.raises(ValueError, match="requires safetensors metadata"): + load_minimax_h3_native_lora( + partition="fl2va", + lora_request=_request(path), + lora_path=path, + dtype=torch.float32, + ) + + +def test_pipeline_native_schedule_and_task_validation(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.partition = "fl2va" + pipeline.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"}) + pipeline._turbo_lora_adapter_ids = set() + pipeline._native_lora_adapter_ids = {7} + pipeline._lora_sigma_schedules = {7: DMD2SigmaSchedule.from_positions([1.0, 0.7, 0.4, 0.15, 0.0])} + pipeline._base_schedule_by_partition = {"fl2va": None} + + sampling = SimpleNamespace( + lora_request=LoRARequest("flashgen", 7, "/tmp/native.safetensors"), + lora_scale=1.0, + num_inference_steps=4, + extra_args={}, + ) + assert pipeline._sigma_schedule_for_request(sampling, "t2va").num_inference_steps == 4 + with pytest.raises(OmniClientError, match="num_inference_steps must be 4"): + pipeline._validate_native_sampling( + SimpleNamespace(lora_request=sampling.lora_request, num_inference_steps=5), + task="t2va", + ) + with pytest.raises(OmniClientError, match="supports T2VA requests only"): + pipeline._resolve_task( + "fl2va", + {}, + has_turbo_lora=False, + has_native_lora=True, + ) + + pipeline._base_schedule_by_partition = {"fl2va": DMD2SigmaSchedule.from_positions([1.0, 0.5, 0.0])} + with pytest.raises(OmniClientError, match="already pins base_schedule"): + pipeline._sigma_schedule_for_request(sampling, "t2va") + + +def test_pipeline_native_step_count_follows_adapter_schedule(): + """The accepted step count is the adapter's interval count, never a literal.""" + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.partition = "fl2va" + pipeline._native_lora_adapter_ids = {7} + pipeline._lora_sigma_schedules = {7: DMD2SigmaSchedule.from_positions([1.0, 0.5, 0.0])} + request = LoRARequest("flashgen", 7, "/tmp/native.safetensors") + + pipeline._validate_native_sampling( + SimpleNamespace(lora_request=request, num_inference_steps=2), + task="t2va", + ) + with pytest.raises(OmniClientError, match="num_inference_steps must be 2 or omitted, not 3"): + pipeline._validate_native_sampling( + SimpleNamespace(lora_request=request, num_inference_steps=3), + task="t2va", + ) + with pytest.raises(OmniClientError, match="requires num_inference_steps=2"): + pipeline._validate_native_sampling( + SimpleNamespace(lora_request=request, num_inference_steps=4), + task="t2va", + ) + + +@pytest.mark.parametrize("step_execution", [False, True]) +def test_pipeline_native_omitted_step_count_follows_execution_mode(step_execution): + """Only request mode can resolve an omitted count from the adapter schedule. + + ``StepScheduler`` reads ``num_inference_steps`` off the request at admission, + before this pipeline sees it, so under step execution an omitted value can + only disagree with the denoise loop. The advertised contract has to match. + """ + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.partition = "fl2va" + pipeline._native_lora_adapter_ids = {7} + pipeline._lora_sigma_schedules = {7: DMD2SigmaSchedule.from_positions([1.0, 0.7, 0.4, 0.15, 0.0])} + pipeline.od_config = SimpleNamespace(step_execution=step_execution) + sampling = SimpleNamespace( + lora_request=LoRARequest("flashgen", 7, "/tmp/native.safetensors"), + num_inference_steps=None, + ) + + if not step_execution: + # Request mode defaults to the adapter's own interval count. + pipeline._validate_native_sampling(sampling, task="t2va") + return + + with pytest.raises(OmniClientError, match="explicit num_inference_steps=4 under step execution"): + pipeline._validate_native_sampling(sampling, task="t2va") + # The step-mode message must not offer omission it cannot accept. + with pytest.raises(OmniClientError, match=r"must be 4, not 5$"): + pipeline._validate_native_sampling( + SimpleNamespace(lora_request=sampling.lora_request, num_inference_steps=5), + task="t2va", + ) + + +def test_pipeline_replaces_native_classification_after_reload(monkeypatch, tmp_path): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module + + path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors" + _write_tiny_native(path) + request = _request(path) + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.partition = "fl2va" + pipeline.od_config = SimpleNamespace( + enable_cpu_offload=False, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=False, + ) + pipeline._turbo_lora_adapter_ids = set() + pipeline._native_lora_adapter_ids = set() + pipeline._lora_sigma_schedules = {} + + loaded = pipeline._load_diffusion_lora_adapter( + lora_request=request, + lora_path=path, + dtype=torch.float32, + ) + assert loaded is not None + assert request.lora_int_id in pipeline._native_lora_adapter_ids + assert request.lora_int_id in pipeline._lora_sigma_schedules + + monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", lambda **_: None) + monkeypatch.setattr(pipeline_module, "load_minimax_h3_native_lora", lambda **_: None) + assert pipeline._load_diffusion_lora_adapter(lora_request=request, lora_path=path, dtype=torch.float32) is None + assert request.lora_int_id not in pipeline._native_lora_adapter_ids + assert request.lora_int_id not in pipeline._lora_sigma_schedules + + +@pytest.mark.parametrize( + "offload_mode", + [ + "model-level CPU offload (--enable-cpu-offload)", + "layerwise offload (--enable-layerwise-offload)", + ], +) +def test_h3_native_rejects_offload_modes(tmp_path, offload_mode): + path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors" + _write_tiny_native(path) + + with pytest.raises(ValueError, match="does not support"): + load_minimax_h3_native_lora( + partition="fl2va", + lora_request=_request(path), + lora_path=path, + dtype=torch.float32, + unsupported_offload_mode=offload_mode, + ) + + +def test_h3_native_allows_distributed_layerwise_offload(monkeypatch): + """DLO keeps LoRA A/B buffers resident, so native must not fail closed.""" + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.partition = "fl2va" + pipeline.od_config = SimpleNamespace( + enable_cpu_offload=False, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=True, + ) + pipeline._turbo_lora_adapter_ids = set() + pipeline._native_lora_adapter_ids = set() + pipeline._lora_sigma_schedules = {} + captured: dict[str, object] = {} + schedule = DMD2SigmaSchedule.from_positions([1.0, 0.7, 0.4, 0.15, 0.0]) + + def load_native(**kwargs): + captured.update(kwargs) + return object(), object(), schedule + + monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", lambda **_: None) + monkeypatch.setattr(pipeline_module, "load_minimax_h3_native_lora", load_native) + + request = _request("flashgen") + loaded = pipeline._load_diffusion_lora_adapter( + lora_request=request, + lora_path="flashgen", + dtype=torch.bfloat16, + ) + + assert loaded is not None + assert captured["unsupported_offload_mode"] is None + assert request.lora_int_id in pipeline._native_lora_adapter_ids + assert pipeline._lora_sigma_schedules[request.lora_int_id] is schedule + + +def test_h3_native_qkv_reorder_matches_base_loader_contract(): + from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import _reorder_grouped_qkv_to_qkv + + num_query_groups = 56 + heads_per_group = 1 + head_dim = 128 + grouped = torch.arange(num_query_groups * (heads_per_group + 2) * head_dim, dtype=torch.float32) + grouped = grouped.reshape(num_query_groups, (heads_per_group + 2) * head_dim) + reordered = _reorder_grouped_qkv_to_qkv( + grouped.reshape(-1, 1), + num_query_groups=num_query_groups, + heads_per_group=heads_per_group, + head_dim=head_dim, + ).reshape(-1) + q_size = num_query_groups * heads_per_group * head_dim + k_size = num_query_groups * head_dim + q, k, v = torch.split(reordered, [q_size, k_size, k_size]) + assert q.reshape(num_query_groups, head_dim)[0, 0] == grouped[0, 0] + assert k.reshape(num_query_groups, head_dim)[0, 0] == grouped[0, head_dim] + assert v.reshape(num_query_groups, head_dim)[0, 0] == grouped[0, 2 * head_dim] + + +def test_native_target_pattern_matches_component_qualified_names(): + """The manager matches ".", so the pattern needs the prefix.""" + import regex as re + + pattern = lora_module._NATIVE_TARGET_PATTERN + matched = [ + target + for target in lora_module._NATIVE_EXPECTED_TARGETS + if re.search(pattern, f"transformer.{target}") is not None + ] + assert len(matched) == 259 + assert re.search(pattern, "transformer.blocks.0.attn.q_proj") is None + assert re.search(pattern, "transformer.token_refiner.blocks.0.adaln_proj.linear") is None + + +def test_replace_layers_binds_every_native_target(monkeypatch): + """Regression for bound=0/259: layer replacement must reach all 259 modules.""" + from vllm.lora.layers import BaseLayerWithLoRA + from vllm.lora.peft_helper import PEFTHelper + + from vllm_omni.diffusion.lora import manager as manager_module + + class _StubLoRALayer(BaseLayerWithLoRA): + def __init__(self, base_layer): + super().__init__() + self.base_layer = base_layer + + def _linear() -> torch.nn.Module: + return torch.nn.Linear(_TINY_HIDDEN, _TINY_HIDDEN) + + def _block(*, with_adaln: bool) -> torch.nn.Module: + block = torch.nn.Module() + block.attn = torch.nn.Module() + block.attn.qkv_proj = _linear() + block.attn.out_proj = _linear() + block.mlp = torch.nn.Module() + block.mlp.fc1 = _linear() + block.mlp.fc2 = _linear() + if with_adaln: + block.adaln_proj = torch.nn.Module() + block.adaln_proj.linear = _linear() + return block + + transformer = torch.nn.Module() + transformer.blocks = torch.nn.ModuleList([_block(with_adaln=True) for _ in range(50)]) + transformer.token_refiner = torch.nn.Module() + transformer.token_refiner.blocks = torch.nn.ModuleList([_block(with_adaln=False) for _ in range(2)]) + transformer.final_layer = torch.nn.Module() + transformer.final_layer.adaln_proj = torch.nn.Module() + transformer.final_layer.adaln_proj.linear = _linear() + pipeline = torch.nn.Module() + pipeline.transformer = transformer + + monkeypatch.setattr(manager_module, "from_layer_diffusion", lambda *, layer, **_: _StubLoRALayer(layer)) + + manager = object.__new__(DiffusionLoRAManager) + manager.pipeline = pipeline + manager.dtype = torch.float32 + manager.max_cached_adapters = 1 + manager._lora_modules = {} + manager._max_lora_rank = 64 + manager._resident_lora_device = None + + peft_helper = PEFTHelper.from_dict( + { + "r": _TINY_RANK, + "lora_alpha": _TINY_RANK, + "target_modules": lora_module._NATIVE_TARGET_PATTERN, + } + ) + manager._replace_layers_with_lora(peft_helper) + + assert set(manager._lora_modules) == {f"transformer.{target}" for target in lora_module._NATIVE_EXPECTED_TARGETS} + + +def test_legacy_manager_uses_native_loader(tmp_path): + path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors" + _write_tiny_native(path) + + class _Pipeline: + def _load_diffusion_lora_adapter(self, **kwargs): + from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( + MiniMaxH3Pipeline, + ) + + pipeline = object.__new__(MiniMaxH3Pipeline) + pipeline.partition = "fl2va" + pipeline.od_config = SimpleNamespace( + enable_cpu_offload=False, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=False, + ) + pipeline._turbo_lora_adapter_ids = set() + pipeline._native_lora_adapter_ids = set() + pipeline._lora_sigma_schedules = {} + return pipeline._load_diffusion_lora_adapter(**kwargs) + + manager = object.__new__(DiffusionLoRAManager) + manager.pipeline = _Pipeline() + manager.dtype = torch.float32 + manager._expected_lora_modules = {"qkv_proj", "fc1", "adaln_proj.linear"} + + lora_model, peft_helper = manager._load_adapter(_request(path)) + assert lora_model.id == 7 + assert peft_helper.lora_alpha == _TINY_RANK + assert len(lora_model.loras) == 259 + + +def test_pipeline_schedule_inactive_when_scale_zero(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.partition = "fl2va" + pipeline._turbo_lora_adapter_ids = set() + pipeline._native_lora_adapter_ids = {7} + pipeline._lora_sigma_schedules = {7: DMD2SigmaSchedule.from_positions([1.0, 0.7, 0.4, 0.15, 0.0])} + pipeline._base_schedule_by_partition = {"fl2va": None} + + sampling = SimpleNamespace( + lora_request=LoRARequest("flashgen", 7, "/tmp/native.safetensors"), + lora_scale=0.0, + num_inference_steps=4, + ) + assert pipeline._sigma_schedule_for_request(sampling, "t2va") is None + assert not pipeline._has_active_native_lora(sampling) + + +def test_pipeline_schedule_falls_back_after_eviction(monkeypatch, tmp_path): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module + + path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors" + _write_tiny_native(path) + request = _request(path) + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.partition = "fl2va" + pipeline.od_config = SimpleNamespace( + enable_cpu_offload=False, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=False, + ) + pipeline._turbo_lora_adapter_ids = set() + pipeline._native_lora_adapter_ids = set() + pipeline._lora_sigma_schedules = {} + pipeline._base_schedule_by_partition = {"fl2va": None} + + loaded = pipeline._load_diffusion_lora_adapter( + lora_request=request, + lora_path=path, + dtype=torch.float32, + ) + assert loaded is not None + assert request.lora_int_id in pipeline._lora_sigma_schedules + + monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", lambda **_: None) + monkeypatch.setattr(pipeline_module, "load_minimax_h3_native_lora", lambda **_: None) + assert pipeline._load_diffusion_lora_adapter(lora_request=request, lora_path=path, dtype=torch.float32) is None + assert request.lora_int_id not in pipeline._lora_sigma_schedules + + sampling = SimpleNamespace( + lora_request=request, + lora_scale=1.0, + num_inference_steps=4, + ) + assert pipeline._sigma_schedule_for_request(sampling, "t2va") is None + + +def test_native_packed_qkv_slices_are_tp2_divisible(): + assert lora_module._NATIVE_QKV_SLICE % 2 == 0 + assert lora_module._NATIVE_FFN_HIDDEN_SIZE % 2 == 0 + + +def test_lora_manager_activates_native_packed_qkv(tmp_path): + path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors" + _write_tiny_native(path) + loaded = load_minimax_h3_native_lora( + partition="fl2va", + lora_request=_request(path), + lora_path=path, + dtype=torch.float32, + ) + assert loaded is not None + lora_model, _, _ = loaded + packed = lora_model.get_lora("blocks.0.attn.qkv_proj") + assert isinstance(packed, PackedLoRALayerWeights) + assert len(packed.lora_b) == 3 + assert all(b.shape[0] == _TINY_QKV_SLICE for b in packed.lora_b) + + class _DummyLoRALayer: + def __init__(self): + self.n_slices = 3 + self.output_slices = (_TINY_QKV_SLICE, _TINY_QKV_SLICE, _TINY_QKV_SLICE) + self.set_calls: list[tuple[list, list]] = [] + self.reset_calls = 0 + + def reset_lora(self, index: int): + self.reset_calls += 1 + + def set_lora(self, index: int, lora_a, lora_b): + self.set_calls.append((lora_a, lora_b)) + + layer = _DummyLoRALayer() + manager = object.__new__(DiffusionLoRAManager) + manager.pipeline = torch.nn.Module() + manager._lora_modules = {"blocks.0.attn.qkv_proj": layer} + manager._registered_adapters = { + lora_model.id: lora_model, + } + manager._active_adapter_id = None + manager._adapter_scales = {} + manager._activate_adapter(lora_model.id, scale=0.5) + + assert manager._active_adapter_id == lora_model.id + + assert len(layer.set_calls) == 1 + lora_a_list, lora_b_list = layer.set_calls[0] + assert len(lora_a_list) == 3 + assert len(lora_b_list) == 3 + torch.testing.assert_close(lora_b_list[0], packed.lora_b[0] * 0.5) + torch.testing.assert_close(lora_b_list[1], packed.lora_b[1] * 0.5) + torch.testing.assert_close(lora_b_list[2], packed.lora_b[2] * 0.5) + + +def test_lora_manager_tp2_splits_native_packed_qkv_per_rank(): + """Each Q/K/V slice must be divisible by TP2 so rank-local layers can shard rows.""" + + tp_size = 2 + slice_rows = lora_module._NATIVE_QKV_SLICE + assert slice_rows % tp_size == 0 + local_rows = slice_rows // tp_size + + class _TpQkvLoRALayer: + def __init__(self, tp_rank: int): + self.tp_rank = tp_rank + self.n_slices = 3 + self.output_slices = (local_rows, local_rows, local_rows) + self.set_calls: list[tuple[list, list]] = [] + + def reset_lora(self, index: int): + return + + def set_lora(self, index: int, lora_a, lora_b): + self.set_calls.append((lora_a, lora_b)) + + for tp_rank in range(tp_size): + layer = _TpQkvLoRALayer(tp_rank) + start = tp_rank * local_rows + end = start + local_rows + full_packed = PackedLoRALayerWeights( + module_name="blocks.0.attn.qkv_proj", + rank=lora_module._NATIVE_RANK, + lora_alphas=[64, 64, 64], + lora_a=[torch.ones(lora_module._NATIVE_RANK, lora_module._NATIVE_HIDDEN_SIZE)] * 3, + lora_b=[torch.full((slice_rows, lora_module._NATIVE_RANK), float(i + 1)) for i in range(3)], + scaling=[1.0, 1.0, 1.0], + ) + local_b = [b[start:end].contiguous() for b in full_packed.lora_b] + layer.set_lora(0, full_packed.lora_a, local_b) + assert len(layer.set_calls) == 1 + _, lora_b_list = layer.set_calls[0] + assert all(b.shape[0] == local_rows for b in lora_b_list) + assert lora_b_list[0][0, 0].item() == 1.0 + assert lora_b_list[1][0, 0].item() == 2.0 + assert lora_b_list[2][0, 0].item() == 3.0 + + +# 259 modules, of which 52 qkv expand to 3 slices and 52 fc1 expand to 2. +_EXPECTED_NATIVE_SLICE_PAIRS = 415 + + +def _iter_native_slice_pairs(lora_model): + """Yield ``(module_name, lora_a, lora_b)`` per slice, flattening packed weights.""" + for module_name, weights in lora_model.loras.items(): + lora_a = weights.lora_a if isinstance(weights.lora_a, list) else [weights.lora_a] + lora_b = weights.lora_b if isinstance(weights.lora_b, list) else [weights.lora_b] + for a, b in zip(lora_a, lora_b, strict=True): + yield module_name, a, b + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_native_lora_load_is_platform_agnostic(tmp_path, dtype): + """The artifact is trained on Ascend, but loading it is not Ascend-only. + + This whole module runs in the CPU job with no ``torch_npu`` present, so a + passing run is itself the evidence. The assertions pin the property that + makes it portable: tensors land on CPU in exactly the requested dtype and + the loader never picks a device, leaving placement to the pipeline. + """ + path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors" + _write_tiny_native(path) + + loaded = load_minimax_h3_native_lora( + partition="fl2va", + lora_request=_request(path), + lora_path=path, + dtype=dtype, + ) + + assert loaded is not None + lora_model, _, _ = loaded + tensors = { + f"{module_name}.{side}": tensor + for module_name, a, b in _iter_native_slice_pairs(lora_model) + for side, tensor in (("lora_a", a), ("lora_b", b)) + } + assert sum(1 for _ in _iter_native_slice_pairs(lora_model)) == _EXPECTED_NATIVE_SLICE_PAIRS + assert {name: tensor.device.type for name, tensor in tensors.items() if tensor.device.type != "cpu"} == {} + assert {name: tensor.dtype for name, tensor in tensors.items() if tensor.dtype != dtype} == {} + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA device") +def test_native_lora_delta_matches_between_cpu_and_cuda(tmp_path): + """The packed adapter produces the same delta weight on CUDA as on CPU. + + Covers the GPU side of the same artifact: every ``lora_b @ lora_a`` product, + including the reordered q/k/v slices, is recomputed on CUDA and compared + against the CPU result, so a GPU deployment cannot silently diverge. + """ + path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors" + _write_tiny_native(path) + + loaded = load_minimax_h3_native_lora( + partition="fl2va", + lora_request=_request(path), + lora_path=path, + dtype=torch.float32, + ) + + assert loaded is not None + lora_model, _, _ = loaded + checked = 0 + for module_name, a, b in _iter_native_slice_pairs(lora_model): + cpu_delta = b @ a + cuda_delta = b.cuda() @ a.cuda() + assert cuda_delta.device.type == "cuda" + torch.testing.assert_close(cuda_delta.cpu(), cpu_delta, msg=f"delta mismatch for {module_name}") + checked += 1 + assert checked == _EXPECTED_NATIVE_SLICE_PAIRS diff --git a/vllm_omni/diffusion/models/minimax_h3/lora.py b/vllm_omni/diffusion/models/minimax_h3/lora.py index bfc7e5565..026d9f6b3 100644 --- a/vllm_omni/diffusion/models/minimax_h3/lora.py +++ b/vllm_omni/diffusion/models/minimax_h3/lora.py @@ -1,6 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Loader for the LightX2V Turbo MiniMax-H3 LoRA. + +Keyed by ``key_format=minimax-h3-diffusers``. The native FlashGen contract lives +in ``.npu.lora``; the two artifacts differ in rank, target naming and QKV +layout, so each owns its own parsing and packing. +""" + from __future__ import annotations import math diff --git a/vllm_omni/diffusion/models/minimax_h3/npu/__init__.py b/vllm_omni/diffusion/models/minimax_h3/npu/__init__.py new file mode 100644 index 000000000..d7b986807 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/npu/__init__.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""MiniMax-H3 components for artifacts produced on Ascend NPU. + +The package name records where these artifacts come from, not where they can +run. Everything here is plain ``torch`` plus ``safetensors`` and carries no +``torch_npu`` dependency, so the same checkpoint loads and serves on CUDA and +CPU as well. See ``tests/diffusion/models/minimax_h3/test_minimax_h3_native_lora.py`` +for the platform-agnostic and CUDA coverage that pins this. +""" + +from .lora import MINIMAX_H3_NATIVE_INFERENCE_STEPS, load_minimax_h3_native_lora + +__all__ = [ + "MINIMAX_H3_NATIVE_INFERENCE_STEPS", + "load_minimax_h3_native_lora", +] diff --git a/vllm_omni/diffusion/models/minimax_h3/npu/lora.py b/vllm_omni/diffusion/models/minimax_h3/npu/lora.py new file mode 100644 index 000000000..2e2d3abd9 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/npu/lora.py @@ -0,0 +1,304 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Loader for the native-layout MiniMax-H3 distilled LoRA (FlashGen v1.0). + +This contract is keyed by ``key_format=minimax-h3-native`` and is distinct from +the LightX2V Turbo contract in the sibling ``..lora`` module: it keeps the fused +``qkv_proj``, adds ``adaln_proj.linear``, uses rank 64, and declares its own +rectified-flow schedule in safetensors metadata. + +The artifact is trained on Ascend NPU, but nothing below is NPU-specific. Only +``torch`` and ``safetensors`` are used, tensors are read on CPU, and the compute +device is chosen later by the pipeline, so the same file serves on CUDA too. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from pathlib import Path + +import torch +from safetensors import safe_open +from vllm.lora.lora_model import LoRAModel +from vllm.lora.lora_weights import PackedLoRALayerWeights +from vllm.lora.peft_helper import PEFTHelper +from vllm.model_executor.models.utils import WeightsMapper + +from vllm_omni.diffusion.sched.sigma_schedule import BASE_SCHEDULE_KEY, DMD2SigmaSchedule +from vllm_omni.lora.request import LoRARequest + +from ..minimax_h3_transformer import _reorder_grouped_qkv_to_qkv + +_LORA_A_SUFFIX = ".lora_A.default.weight" +_LORA_B_SUFFIX = ".lora_B.default.weight" + +_NATIVE_RANK = 64 +_NATIVE_ALPHA = 64 +_NATIVE_HIDDEN_SIZE = 5376 +_NATIVE_ATTENTION_INNER_SIZE = 7168 +_NATIVE_FFN_HIDDEN_SIZE = 14336 +_NATIVE_TIME_EMBED_DIM = 2688 +_NATIVE_BLOCK_ADALN_OUT = 96768 +_NATIVE_FINAL_ADALN_OUT = 10752 +_NATIVE_NUM_QUERY_GROUPS = 56 +_NATIVE_HEADS_PER_GROUP = 1 +_NATIVE_HEAD_DIM = 128 +_NATIVE_QKV_SLICE = _NATIVE_ATTENTION_INNER_SIZE +_NATIVE_KEY_FORMAT = "minimax-h3-native" +_NATIVE_QKV_LAYOUT = "grouped" +# Public because request validation in the pipeline speaks the same contract. +MINIMAX_H3_NATIVE_INFERENCE_STEPS = 4 +_NATIVE_FILENAME = "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors" +_NATIVE_TARGET_SUFFIXES = ( + "attn.qkv_proj", + "attn.out_proj", + "mlp.fc1", + "mlp.fc2", + "adaln_proj.linear", +) +_NATIVE_TOKEN_REFINER_SUFFIXES = ( + "attn.qkv_proj", + "attn.out_proj", + "mlp.fc1", + "mlp.fc2", +) +_NATIVE_TARGET_DIMS = { + "attn.qkv_proj": (_NATIVE_HIDDEN_SIZE, 3 * _NATIVE_ATTENTION_INNER_SIZE), + "attn.out_proj": (_NATIVE_ATTENTION_INNER_SIZE, _NATIVE_HIDDEN_SIZE), + "mlp.fc1": (_NATIVE_HIDDEN_SIZE, 2 * _NATIVE_FFN_HIDDEN_SIZE), + "mlp.fc2": (_NATIVE_FFN_HIDDEN_SIZE, _NATIVE_HIDDEN_SIZE), + "adaln_proj.linear": (_NATIVE_TIME_EMBED_DIM, _NATIVE_BLOCK_ADALN_OUT), +} +_NATIVE_FINAL_ADALN_DIMS = (_NATIVE_TIME_EMBED_DIM, _NATIVE_FINAL_ADALN_OUT) +_NATIVE_EXPECTED_TARGETS = frozenset( + [*(f"blocks.{block_index}.{suffix}" for block_index in range(50) for suffix in _NATIVE_TARGET_SUFFIXES)] + + [ + *( + f"token_refiner.blocks.{block_index}.{suffix}" + for block_index in range(2) + for suffix in _NATIVE_TOKEN_REFINER_SUFFIXES + ) + ] + + ["final_layer.adaln_proj.linear"] +) +_NATIVE_TARGET_PATTERN = ( + r"^(?:transformer\.blocks\.\d+\.(?:attn\.(?:qkv_proj|out_proj)|mlp\.(?:fc1|fc2)|adaln_proj\.linear)" + r"|transformer\.token_refiner\.blocks\.\d+\.(?:attn\.(?:qkv_proj|out_proj)|mlp\.(?:fc1|fc2))" + r"|transformer\.final_layer\.adaln_proj\.linear)$" +) +_NATIVE_WEIGHTS_MAPPER = WeightsMapper( + orig_to_new_substr={ + "transformer.": "", + ".lora_A.default.": ".lora_A.", + ".lora_B.default.": ".lora_B.", + } +) + + +def _select_native_file(artifact_path: str | Path) -> Path | None: + path = Path(artifact_path) + if path.is_file(): + return path if path.suffix == ".safetensors" else None + if not path.is_dir(): + return None + + candidate = path / _NATIVE_FILENAME + return candidate if candidate.is_file() else None + + +def _native_target_dims(target: str) -> tuple[int, int]: + if target == "final_layer.adaln_proj.linear": + return _NATIVE_FINAL_ADALN_DIMS + if target.endswith("adaln_proj.linear"): + return _NATIVE_TARGET_DIMS["adaln_proj.linear"] + for suffix in ("attn.qkv_proj", "attn.out_proj", "mlp.fc1", "mlp.fc2"): + if target.endswith(suffix): + return _NATIVE_TARGET_DIMS[suffix] + raise ValueError(f"Unsupported MiniMax-H3 native target: {target!r}") + + +def _parse_native_base_schedule(metadata: Mapping[str, str]) -> DMD2SigmaSchedule | None: + """Read the adapter-declared schedule from safetensors string metadata. + + safetensors metadata is string-only, so the native artifact ships the sigma + positions comma separated. Absent means the adapter declares no schedule. + """ + raw = metadata.get(BASE_SCHEDULE_KEY) + if raw is None: + return None + positions = [part.strip() for part in str(raw).split(",") if part.strip()] + if not positions: + raise ValueError(f"MiniMax-H3 native LoRA {BASE_SCHEDULE_KEY} must not be empty") + try: + return DMD2SigmaSchedule.from_positions([float(position) for position in positions]) + except ValueError as exc: + raise ValueError(f"MiniMax-H3 native LoRA {BASE_SCHEDULE_KEY} is malformed: {exc}") from exc + + +def _validate_native_metadata(metadata: Mapping[str, str]) -> DMD2SigmaSchedule: + if metadata.get("key_format") != _NATIVE_KEY_FORMAT: + raise ValueError(f"MiniMax-H3 native LoRA requires safetensors metadata key_format={_NATIVE_KEY_FORMAT!r}") + if metadata.get("qkv_layout") != _NATIVE_QKV_LAYOUT: + raise ValueError(f"MiniMax-H3 native LoRA requires safetensors metadata qkv_layout={_NATIVE_QKV_LAYOUT!r}") + try: + rank = int(metadata.get("lora_rank", "")) + alpha = float(metadata.get("lora_alpha", "")) + except (TypeError, ValueError) as exc: + raise ValueError("MiniMax-H3 native LoRA metadata lora_rank/lora_alpha must be numeric") from exc + if rank != _NATIVE_RANK: + raise ValueError(f"MiniMax-H3 native LoRA requires lora_rank={_NATIVE_RANK}, got {rank!r}") + if alpha != _NATIVE_ALPHA: + raise ValueError(f"MiniMax-H3 native LoRA requires lora_alpha={_NATIVE_ALPHA}, got {alpha!r}") + tasks = {part.strip().lower() for part in str(metadata.get("tasks", "")).split(",") if part.strip()} + if tasks != {"t2va"}: + raise ValueError("MiniMax-H3 native LoRA v1.0 supports tasks=t2va only") + schedule = _parse_native_base_schedule(metadata) + if schedule is None: + raise ValueError(f"MiniMax-H3 native LoRA requires safetensors metadata {BASE_SCHEDULE_KEY}") + # Request validation speaks in interval counts, so a mislabeled schedule + # would otherwise silently change the step count instead of failing here. + if schedule.num_inference_steps != MINIMAX_H3_NATIVE_INFERENCE_STEPS: + raise ValueError( + f"MiniMax-H3 native LoRA v1.0 requires a {MINIMAX_H3_NATIVE_INFERENCE_STEPS}-interval base_schedule " + f"({MINIMAX_H3_NATIVE_INFERENCE_STEPS + 1} sigma positions), got {schedule.num_inference_steps}" + ) + return schedule + + +def _validate_native_tensors(checkpoint) -> dict[str, torch.Tensor]: + tensors: dict[str, torch.Tensor] = {} + pairs: dict[str, set[str]] = {} + raw_targets: set[str] = set() + for name in checkpoint.keys(): + if name.endswith(_LORA_A_SUFFIX): + raw_target = name[: -len(_LORA_A_SUFFIX)] + side = "a" + elif name.endswith(_LORA_B_SUFFIX): + raw_target = name[: -len(_LORA_B_SUFFIX)] + side = "b" + else: + raise ValueError(f"Unconsumed MiniMax-H3 native tensor: {name!r}") + if not raw_target.startswith("transformer."): + raise ValueError(f"MiniMax-H3 native LoRA tensors must use transformer.* keys, got {name!r}") + mapped_target = raw_target.removeprefix("transformer.") + raw_targets.add(mapped_target) + + target_sides = pairs.setdefault(mapped_target, set()) + if side in target_sides: + raise ValueError(f"Duplicate MiniMax-H3 native tensor for {mapped_target}.{side}") + target_sides.add(side) + + tensor = checkpoint.get_tensor(name) + if tensor.ndim != 2: + raise ValueError(f"MiniMax-H3 native LoRA tensors must be matrices, got {name}={tuple(tensor.shape)}") + input_dim, output_dim = _native_target_dims(mapped_target) + expected_shape = (_NATIVE_RANK, input_dim) if side == "a" else (output_dim, _NATIVE_RANK) + if tuple(tensor.shape) != expected_shape: + raise ValueError( + f"MiniMax-H3 native tensor has invalid global shape: {name}={tuple(tensor.shape)}, " + f"expected={expected_shape}" + ) + tensors[name] = tensor + + incomplete = sorted(target for target, sides in pairs.items() if sides != {"a", "b"}) + if incomplete: + raise ValueError(f"Incomplete MiniMax-H3 native LoRA pairs: {incomplete}") + missing = sorted(_NATIVE_EXPECTED_TARGETS - raw_targets) + unexpected = sorted(raw_targets - _NATIVE_EXPECTED_TARGETS) + if missing or unexpected: + raise ValueError( + "MiniMax-H3 native target set does not match the supported v1.0 artifact: " + f"missing={len(missing)} {missing[:5]}, unexpected={len(unexpected)} {unexpected[:5]}" + ) + return tensors + + +def _pack_h3_native_fc1(lora_model: LoRAModel) -> None: + for module_name, weights in tuple(lora_model.loras.items()): + if not module_name.endswith(".mlp.fc1"): + continue + gate_b, up_b = weights.lora_b.chunk(2, dim=0) + lora_model.loras[module_name] = PackedLoRALayerWeights( + module_name=module_name, + rank=weights.rank, + lora_alphas=[weights.lora_alpha, weights.lora_alpha], + lora_a=[weights.lora_a, weights.lora_a], + lora_b=[gate_b.contiguous(), up_b.contiguous()], + scaling=[weights.scaling, weights.scaling], + ) + + +def _pack_h3_native_qkv(lora_model: LoRAModel) -> None: + for module_name, weights in tuple(lora_model.loras.items()): + if not module_name.endswith(".attn.qkv_proj"): + continue + reordered_b = _reorder_grouped_qkv_to_qkv( + weights.lora_b, + num_query_groups=_NATIVE_NUM_QUERY_GROUPS, + heads_per_group=_NATIVE_HEADS_PER_GROUP, + head_dim=_NATIVE_HEAD_DIM, + ) + q_b, k_b, v_b = torch.split( + reordered_b, + [_NATIVE_QKV_SLICE, _NATIVE_QKV_SLICE, _NATIVE_QKV_SLICE], + dim=0, + ) + lora_model.loras[module_name] = PackedLoRALayerWeights( + module_name=module_name, + rank=weights.rank, + lora_alphas=[weights.lora_alpha, weights.lora_alpha, weights.lora_alpha], + lora_a=[weights.lora_a, weights.lora_a, weights.lora_a], + lora_b=[q_b.contiguous(), k_b.contiguous(), v_b.contiguous()], + scaling=[weights.scaling, weights.scaling, weights.scaling], + ) + + +def load_minimax_h3_native_lora( + *, + partition: str, + lora_request: LoRARequest, + lora_path: str | Path, + dtype: torch.dtype, + unsupported_offload_mode: str | None = None, +) -> tuple[LoRAModel, PEFTHelper, DMD2SigmaSchedule] | None: + """Load a native-layout MiniMax-H3 distilled LoRA through the legacy manager.""" + + lora_file = _select_native_file(lora_path) + if lora_file is None: + return None + with safe_open(lora_file, framework="pt", device="cpu") as checkpoint: + metadata = checkpoint.metadata() or {} + if metadata.get("key_format") != _NATIVE_KEY_FORMAT: + if lora_file.name == _NATIVE_FILENAME: + raise ValueError( + f"MiniMax-H3 native LoRA requires safetensors metadata key_format={_NATIVE_KEY_FORMAT!r}" + ) + return None + if lora_file.name != _NATIVE_FILENAME: + raise ValueError(f"MiniMax-H3 native LoRA supports only {_NATIVE_FILENAME!r}, got {lora_file.name!r}") + sigma_schedule = _validate_native_metadata(metadata) + if partition == "ref2va": + raise ValueError("MiniMax-H3 native LoRA supports T2VA only") + if unsupported_offload_mode is not None: + raise ValueError(f"MiniMax-H3 native dynamic LoRA does not support {unsupported_offload_mode}") + tensors = _validate_native_tensors(checkpoint) + + peft_helper = PEFTHelper.from_dict( + { + "r": _NATIVE_RANK, + "lora_alpha": _NATIVE_ALPHA, + "target_modules": _NATIVE_TARGET_PATTERN, + } + ) + lora_model = LoRAModel.from_lora_tensors( + lora_model_id=lora_request.lora_int_id, + tensors=tensors, + peft_helper=peft_helper, + device="cpu", + dtype=dtype, + weights_mapper=_NATIVE_WEIGHTS_MAPPER, + ) + _pack_h3_native_qkv(lora_model) + _pack_h3_native_fc1(lora_model) + return lora_model, peft_helper, sigma_schedule diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index 4f42e7734..d14bfd0fd 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -123,6 +123,10 @@ from .minimax_h3_transformer import ( MiniMaxH3DiTModel, _attention_isolates_packed_requests, ) +from .npu.lora import ( + MINIMAX_H3_NATIVE_INFERENCE_STEPS, + load_minimax_h3_native_lora, +) from .packed_sequence import ( minimax_h3_packed_sequence, minimax_h3_packed_sequence_ref2va_blocks, @@ -672,6 +676,8 @@ class MiniMaxH3Pipeline( # A cache eviction may be followed by a different adapter reusing the # same client-supplied ID. Every real load replaces the classification. self._turbo_lora_adapter_ids.discard(lora_request.lora_int_id) + self._native_lora_adapter_ids.discard(lora_request.lora_int_id) + self._lora_sigma_schedules.pop(lora_request.lora_int_id, None) od_config = getattr(self, "od_config", None) offload_modes = [] if getattr(od_config, "enable_cpu_offload", False): @@ -687,7 +693,25 @@ class MiniMaxH3Pipeline( ) if loaded is not None: self._turbo_lora_adapter_ids.add(lora_request.lora_int_id) - return loaded + return loaded + + # Selection is by the artifact's safetensors ``key_format``, not by the + # running platform: the native loader is checkpoint-format parsing with + # no ``torch_npu`` dependency, so it needs no ``current_omni_platform`` + # dispatch and binds the same adapter on NPU, CUDA and CPU. + native_loaded = load_minimax_h3_native_lora( + partition=self.partition, + lora_request=lora_request, + lora_path=lora_path, + dtype=dtype, + unsupported_offload_mode=" or ".join(offload_modes) or None, + ) + if native_loaded is not None: + lora_model, peft_helper, sigma_schedule = native_loaded + self._native_lora_adapter_ids.add(lora_request.lora_int_id) + self._lora_sigma_schedules[lora_request.lora_int_id] = sigma_schedule + return lora_model, peft_helper + return None def _validate_diffusion_lora_binding( self, @@ -695,12 +719,20 @@ class MiniMaxH3Pipeline( lora_model: LoRAModel, bound_lora_names: frozenset[str], ) -> None: - if lora_model.id not in self._turbo_lora_adapter_ids: + if lora_model.id in self._turbo_lora_adapter_ids: + missing = sorted(set(lora_model.loras) - bound_lora_names) + if missing: + raise ValueError( + "MiniMax-H3 Turbo LoRA binding is incomplete: " + f"bound={len(bound_lora_names)}/{len(lora_model.loras)}, missing={missing[:5]}" + ) + return + if lora_model.id not in self._native_lora_adapter_ids: return missing = sorted(set(lora_model.loras) - bound_lora_names) if missing: raise ValueError( - "MiniMax-H3 Turbo LoRA binding is incomplete: " + "MiniMax-H3 native LoRA binding is incomplete: " f"bound={len(bound_lora_names)}/{len(lora_model.loras)}, missing={missing[:5]}" ) @@ -712,6 +744,65 @@ class MiniMaxH3Pipeline( and lora_request.lora_int_id in self._turbo_lora_adapter_ids ) + def _has_active_native_lora(self, sampling: Any) -> bool: + lora_request = sampling.lora_request + return ( + lora_request is not None + and not math.isclose(0.0, float(sampling.lora_scale)) + and lora_request.lora_int_id in self._native_lora_adapter_ids + ) + + def _validate_native_sampling(self, sampling: Any, *, task: str) -> None: + if task != "t2va": + raise OmniClientError("MiniMax-H3 native LoRA supports T2VA requests only") + # Derive the expected count from the adapter's own schedule so the + # message can never disagree with the schedule the denoise loop runs. + schedule = self._lora_sigma_schedules.get(sampling.lora_request.lora_int_id) + expected_steps = MINIMAX_H3_NATIVE_INFERENCE_STEPS if schedule is None else schedule.num_inference_steps + # Only request mode can take the count from the adapter schedule: step + # mode admits the request in ``StepScheduler``, which reads + # ``num_inference_steps`` off it before any pipeline hook runs. Reject + # omission there rather than advertise a contract that would either fail + # admission or disagree with the denoise loop. + od_config = getattr(self, "od_config", None) + omission_allowed = not getattr(od_config, "step_execution", False) + or_omitted = " or omitted" if omission_allowed else "" + sigma_steps = sampling.num_inference_steps + if sigma_steps is None: + if omission_allowed: + return + raise OmniClientError( + f"MiniMax-H3 native LoRA requires an explicit num_inference_steps={expected_steps} " + "under step execution, because the step scheduler derives the total step count from " + "the request before the adapter schedule is known" + ) + if int(sigma_steps) == expected_steps + 1: + raise OmniClientError( + "MiniMax-H3 native LoRA uses the distilled interval-count contract; " + f"num_inference_steps must be {expected_steps}{or_omitted}, not {expected_steps + 1}" + ) + if int(sigma_steps) != expected_steps: + raise OmniClientError( + f"MiniMax-H3 native LoRA requires num_inference_steps={expected_steps} " + f"(one denoiser evaluation per sigma interval){or_omitted}" + ) + + def _sigma_schedule_for_request(self, sampling: Any, task: str) -> DMD2SigmaSchedule | None: + lora_request = sampling.lora_request + if ( + lora_request is not None + and not math.isclose(0.0, float(sampling.lora_scale)) + and lora_request.lora_int_id in self._lora_sigma_schedules + ): + adapter_schedule = self._lora_sigma_schedules[lora_request.lora_int_id] + checkpoint_schedule = self._base_schedule_for_task(task) + if checkpoint_schedule is not None: + raise OmniClientError( + "MiniMax-H3 native LoRA cannot be activated on a checkpoint that already pins base_schedule" + ) + return adapter_schedule + return self._base_schedule_for_task(task) + def _validate_turbo_sampling(self, sampling: Any) -> None: extra = sampling.extra_args or {} sigma_points = sampling.num_inference_steps @@ -766,6 +857,8 @@ class MiniMaxH3Pipeline( str(od_config.model), ) self._turbo_lora_adapter_ids: set[int] = set() + self._native_lora_adapter_ids: set[int] = set() + self._lora_sigma_schedules: dict[int, DMD2SigmaSchedule] = {} model_root = _resolve_minimax_h3_model_root( str(od_config.model), od_config.revision, @@ -981,6 +1074,7 @@ class MiniMaxH3Pipeline( multi_modal_data: dict[str, Any], *, has_turbo_lora: bool = False, + has_native_lora: bool = False, ) -> str: if requested is None: # A Ref2VA-only startup has no FL2VA transformer; preserve its @@ -1000,6 +1094,8 @@ class MiniMaxH3Pipeline( ) if task == "ref2va" and has_turbo_lora: 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") return task def _resolve_shape( @@ -2001,13 +2097,17 @@ class MiniMaxH3Pipeline( logger.debug("MiniMax H3 request quality=%s", quality) extra = sampling.extra_args or {} has_turbo_lora = self._has_active_turbo_lora(sampling) + has_native_lora = self._has_active_native_lora(sampling) task = self._resolve_task( extra.get("task"), multi_modal_data, has_turbo_lora=has_turbo_lora, + has_native_lora=has_native_lora, ) if has_turbo_lora: self._validate_turbo_sampling(sampling) + if has_native_lora: + self._validate_native_sampling(sampling, task=task) raw_image = multi_modal_data.get("image") raw_videos = multi_modal_data.get("video") @@ -2224,7 +2324,7 @@ class MiniMaxH3Pipeline( ref_audio_t = audio_lengths[0] seed = int(sampling.seed if sampling.seed is not None else 42) - sigma_schedule = self._base_schedule_for_task(task) + 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)