From 59cf3add400be1744f308f497bd31dcceb7be262 Mon Sep 17 00:00:00 2001 From: Mu GuanLin <1203789601@qq.com> Date: Wed, 26 Aug 2026 16:15:47 +0800 Subject: [PATCH] [Feature] Support MiniMax-H3 Turbo LoRA with DLO (#6550) Signed-off-by: mglyn <1203789601@qq.com> Co-authored-by: Hongsheng Liu --- recipes/MiniMaxAI/MiniMax-H3.md | 9 ++--- tests/diffusion/lora/test_base_linear.py | 25 +++++++++++++- tests/diffusion/lora/test_lora_manager.py | 33 +++++++++++++++++++ .../models/minimax_h3/test_minimax_h3_lora.py | 33 ++++++++++++++++++- .../diffusion/lora/layers/base_linear.py | 18 +++++++++- vllm_omni/diffusion/lora/manager.py | 13 ++++++++ .../models/minimax_h3/pipeline_minimax_h3.py | 2 -- 7 files changed, 124 insertions(+), 9 deletions(-) diff --git a/recipes/MiniMaxAI/MiniMax-H3.md b/recipes/MiniMaxAI/MiniMax-H3.md index 722b5f8b0..915c71739 100644 --- a/recipes/MiniMaxAI/MiniMax-H3.md +++ b/recipes/MiniMaxAI/MiniMax-H3.md @@ -839,7 +839,7 @@ hf download lightx2v/Minimax-h3-Turbo "${TURBO_FILE}" --local-dir "${TURBO_DIR}" export TURBO_LORA="${TURBO_DIR}/${TURBO_FILE}" ``` -Start from a non-offloaded FL2VA server command and add +Start from a non-offloaded or DLO FL2VA server command and add `--task-type fl2va --lora-backend peft --lora-path "${TURBO_LORA}"`. `--lora-path` preloads the adapter; each request still activates it and uses the published sampling settings: @@ -853,9 +853,10 @@ the published sampling settings: For FL2VA, change `task` and add `input_reference` as shown above. The 8-step, ComfyUI, Ref2VA, and v1.1 artifacts are not supported. This integration is -dynamic-only and does not support prefusion, DLO, or LoRA composition. -It also rejects model-level, layerwise, and distributed layerwise offload; -the legacy dynamic LoRA tensors do not participate in those weight lifecycles. +dynamic-only and does not support prefusion or LoRA composition. DLO is +supported by keeping the request-switchable LoRA A/B buffers resident on the +accelerator while DLO streams only the base blocks; budget for this additional +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. diff --git a/tests/diffusion/lora/test_base_linear.py b/tests/diffusion/lora/test_base_linear.py index 9386d0909..b2f76be05 100644 --- a/tests/diffusion/lora/test_base_linear.py +++ b/tests/diffusion/lora/test_base_linear.py @@ -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 @@ -89,6 +89,29 @@ def test_diffusion_base_linear_apply_multi_slice(): assert torch.allclose(out, expected) +def test_diffusion_lora_buffer_device_persists_across_reallocation(monkeypatch): + from vllm.lora.layers.base_linear import BaseLinearLayerWithLoRA + + layer = DiffusionBaseLinearLayerWithLoRA.__new__(DiffusionBaseLinearLayerWithLoRA) + torch.nn.Module.__init__(layer) + layer.base_layer = torch.nn.Module() + layer.n_slices = 1 + layer.lora_a_stacked = (torch.zeros(2),) + layer.lora_b_stacked = (torch.zeros(3),) + + layer._set_diffusion_lora_buffer_device(torch.device("meta")) + assert all(tensor.is_meta for tensor in (*layer.lora_a_stacked, *layer.lora_b_stacked)) + + def recreate_on_cpu(self, max_loras, lora_config, model_config): + self.lora_a_stacked = (torch.zeros(4),) + self.lora_b_stacked = (torch.zeros(5),) + + monkeypatch.setattr(BaseLinearLayerWithLoRA, "create_lora_weights", recreate_on_cpu) + layer.create_lora_weights(max_loras=1, lora_config=object()) + + assert all(tensor.is_meta for tensor in (*layer.lora_a_stacked, *layer.lora_b_stacked)) + + def test_diffusion_base_linear_reset_lora_disables_fast_path(monkeypatch): # Verify that after reset_lora(), apply() skips LoRA matmuls even if the # LoRA tensors are still allocated and non-empty. diff --git a/tests/diffusion/lora/test_lora_manager.py b/tests/diffusion/lora/test_lora_manager.py index fa80ca7c8..8edc6ce4b 100644 --- a/tests/diffusion/lora/test_lora_manager.py +++ b/tests/diffusion/lora/test_lora_manager.py @@ -3,6 +3,8 @@ from __future__ import annotations +from types import SimpleNamespace + import pytest import torch from vllm.lora.lora_weights import LoRALayerWeights @@ -143,6 +145,37 @@ def test_lora_manager_replace_layers_does_not_rewrap_base_layer(monkeypatch): assert replace_calls == ["foo"] +def test_lora_manager_keeps_dlo_lora_buffers_on_compute_device(monkeypatch): + import vllm_omni.diffusion.lora.manager as manager_mod + + class _DLOLoRALayer(DummyBaseLayerWithLoRA): + lora_buffer_device = None + + def _set_diffusion_lora_buffer_device(self, device): + self.lora_buffer_device = device + + monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DLOLoRALayer) + monkeypatch.setattr( + manager_mod, + "from_layer_diffusion", + lambda *, layer, **_kwargs: _DLOLoRALayer(layer), + ) + + pipeline = torch.nn.Module() + pipeline.od_config = SimpleNamespace(enable_distributed_layerwise_offload=True) + pipeline.transformer = torch.nn.Module() + pipeline.transformer.foo = _FakeLinearBase() + manager = DiffusionLoRAManager( + pipeline=pipeline, + device=torch.device("cpu"), + dtype=torch.bfloat16, + ) + + manager._replace_layers_with_lora(type("_PH", (), {"r": 1})()) + + assert pipeline.transformer.foo.lora_buffer_device == torch.device("cpu") + + def test_lora_manager_replaces_packed_layer_when_targeting_sublayers(monkeypatch): import vllm_omni.diffusion.lora.manager as manager_mod 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 58480f4a9..22581786d 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_lora.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_lora.py @@ -197,7 +197,6 @@ def test_h3_turbo_rejects_wrong_alpha_and_ref2va(tmp_path): [ "model-level CPU offload (--enable-cpu-offload)", "layerwise offload (--enable-layerwise-offload)", - "distributed layerwise offload (--enable-distributed-layerwise-offload)", ], ) def test_h3_turbo_rejects_offload_modes(tmp_path, offload_mode): @@ -214,6 +213,38 @@ def test_h3_turbo_rejects_offload_modes(tmp_path, offload_mode): ) +def test_h3_turbo_allows_distributed_layerwise_offload(monkeypatch): + 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._turbo_lora_adapter_ids = set() + pipeline.od_config = SimpleNamespace( + enable_cpu_offload=False, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=True, + ) + captured = {} + + def load_turbo(**kwargs): + captured.update(kwargs) + return object(), object() + + monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", load_turbo) + loaded = pipeline._load_diffusion_lora_adapter( + lora_request=_request("turbo"), + lora_path="turbo", + dtype=torch.bfloat16, + ) + + assert loaded is not None + assert captured["unsupported_offload_mode"] is None + + def test_h3_turbo_accepts_only_the_declared_v1_artifact(tmp_path): unsupported = tmp_path / "minimax_h3_fl2v_turbo_8step_v1.0_bf16.safetensors" _write_tiny_turbo(unsupported) diff --git a/vllm_omni/diffusion/lora/layers/base_linear.py b/vllm_omni/diffusion/lora/layers/base_linear.py index fe32868d0..fae385c36 100644 --- a/vllm_omni/diffusion/lora/layers/base_linear.py +++ b/vllm_omni/diffusion/lora/layers/base_linear.py @@ -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 @@ -22,6 +22,19 @@ class DiffusionBaseLinearLayerWithLoRA(BaseLinearLayerWithLoRA): is inherited from vLLM's BaseLinearLayerWithLoRA. """ + lora_a_stacked: tuple[torch.Tensor, ...] + lora_b_stacked: tuple[torch.Tensor, ...] + + def _move_lora_buffers(self, device: torch.device) -> None: + self.lora_a_stacked = tuple(tensor.to(device=device) for tensor in self.lora_a_stacked) + self.lora_b_stacked = tuple(tensor.to(device=device) for tensor in self.lora_b_stacked) + + def _set_diffusion_lora_buffer_device(self, device: torch.device) -> None: + """Keep dynamic LoRA sidecars resident on the compute device.""" + + self._diffusion_lora_buffer_device = device + self._move_lora_buffers(device) + def create_lora_weights( self, max_loras: int, @@ -29,6 +42,9 @@ class DiffusionBaseLinearLayerWithLoRA(BaseLinearLayerWithLoRA): model_config=None, ) -> None: super().create_lora_weights(max_loras, lora_config, model_config) + buffer_device = getattr(self, "_diffusion_lora_buffer_device", None) + if buffer_device is not None: + self._move_lora_buffers(buffer_device) # Keep a direct reference for attribute forwarding: `base_layer` is a # registered submodule (stored under `_modules`), so direct access via # `object.__getattribute__` will not find it. We stash a ref in diff --git a/vllm_omni/diffusion/lora/manager.py b/vllm_omni/diffusion/lora/manager.py index 61f6572f6..1b63692fc 100644 --- a/vllm_omni/diffusion/lora/manager.py +++ b/vllm_omni/diffusion/lora/manager.py @@ -68,6 +68,12 @@ class DiffusionLoRAManager: self.pipeline = pipeline self.device = device self.dtype = dtype + od_config = getattr(pipeline, "od_config", None) + # DLO owns the base-weight lifecycle. Keep request-switchable LoRA + # sidecars resident instead of rebuilding DLO host shards per request. + self._resident_lora_device = ( + device if getattr(od_config, "enable_distributed_layerwise_offload", False) is True else None + ) # Cache supported/expected module suffixes once, before any layer # replacement happens. After LoRA layers are injected, the original @@ -450,6 +456,13 @@ class DiffusionLoRAManager: ) if lora_layer is not module and isinstance(lora_layer, BaseLayerWithLoRA): + if self._resident_lora_device is not None: + set_buffer_device = getattr(lora_layer, "_set_diffusion_lora_buffer_device", None) + if not callable(set_buffer_device): + raise RuntimeError( + f"{type(lora_layer).__name__} cannot keep dynamic LoRA buffers resident for DLO" + ) + set_buffer_device(self._resident_lora_device) replace_submodule(component, module_name, lora_layer) self._lora_modules[full_module_name] = lora_layer logger.debug("Replaced layer: %s -> %s", full_module_name, type(lora_layer).__name__) 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 ff3990b32..33c3cd3b0 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -724,8 +724,6 @@ class MiniMaxH3Pipeline( offload_modes.append("model-level CPU offload (--enable-cpu-offload)") if getattr(od_config, "enable_layerwise_offload", False): offload_modes.append("layerwise offload (--enable-layerwise-offload)") - if getattr(od_config, "enable_distributed_layerwise_offload", False): - offload_modes.append("distributed layerwise offload (--enable-distributed-layerwise-offload)") loaded = load_minimax_h3_turbo_lora( partition=self.partition, lora_request=lora_request,