[Model] Add VoxCPM2 startup LoRA adapter support (#7108)

Signed-off-by: bezdarnost <urumbekovaman@gmail.com>
Co-authored-by: bezdarnost <urumbekovaman@gmail.com>

docs/serving/speech_api.md

@@ -777,6 +777,40 @@ hidden-state dtype before the output projection. Consequently, a float32 stage c
 | ------- | ------------- |
 | `openbmb/VoxCPM2` | TTS + voice cloning with built-in speaker presets and uploaded-voice support. Accepts `voice` (preset or uploaded) or `ref_audio` + optional `ref_text`. |
 
+#### Startup LoRA adapter
+
+To serve a fine-tuned VoxCPM2 voice, set
+`voxcpm2_runtime_config.startup_lora_path` in the stage's `hf_overrides`.
+Copy `vllm_omni/deploy/voxcpm2.yaml` to a local deploy config and add this key
+alongside its existing runtime settings:
+
+```yaml
+# Under stages[0].engine_extras.hf_overrides.voxcpm2_runtime_config:
+startup_lora_path: /absolute/path/to/adapter
+```
+
+Then launch with the modified config:
+
+```bash
+vllm serve openbmb/VoxCPM2 --omni --deploy-config /absolute/path/to/voxcpm2-lora.yaml
+```
+
+The directory must be accessible to the worker and contain the native VoxCPM2
+training export: `lora_config.json` (with a `lora_config` object containing
+`r`, `alpha`, enabled groups, and target module names) and
+`lora_weights.safetensors`. PEFT checkpoints and pickle checkpoints are not
+accepted. Missing, unexpected, non-finite, or incorrectly shaped adapter
+tensors fail model loading rather than silently loading a partial adapter.
+
+The adapter is merged into the base LM, residual LM, LocDiT, and optional
+projection layers selected by its configuration, after base weight loading
+and before compilation or CUDA graph capture. All requests use that adapter;
+`voice` still selects a reference voice, not a LoRA adapter. Changing adapters
+requires restarting the server. Runtime loading/unloading, per-request
+multi-LoRA selection, and `load_format=dummy` are not supported by this path.
+Weight fusion rounds to the base model's dtype, so numerical and speech-quality
+parity should be checked against the upstream adapter on your deployment.
+
 ### MOSS-TTS-Nano
 
 | Model | Description |

tests/e2e/online_serving/test_voxcpm2_tts_lora.py

@@ -0,0 +1,106 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Serve a native VoxCPM2 adapter through the speech API."""
+
+import json
+from pathlib import Path
+
+import pytest
+import torch
+import yaml
+from huggingface_hub import hf_hub_download
+from safetensors import safe_open
+from safetensors.torch import save_file
+
+from tests.helpers.fixtures.runtime import omni_fixture_lock
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams, get_model_prefix, iter_omni_server
+from tests.helpers.stage_config import get_deploy_config_path
+
+pytestmark = [pytest.mark.slow, pytest.mark.tts]
+MODEL = "openbmb/VoxCPM2"
+
+
+@pytest.fixture(scope="module")
+def omni_server(request, run_level, tmp_path_factory):
+    directory = tmp_path_factory.mktemp("voxcpm2-lora")
+    model = get_model_prefix() + MODEL
+    weights = Path(model) / "model.safetensors"
+    if not weights.is_file():
+        weights = Path(hf_hub_download(MODEL, "model.safetensors"))
+    projections = ["enc_to_lm_proj", "lm_to_dit_proj", "res_to_dit_proj", "fusion_concat_proj"]
+    generator = torch.Generator().manual_seed(42)
+    tensors = {}
+    with safe_open(weights, framework="pt", device="cpu") as checkpoint:
+        for key in checkpoint.keys():
+            if not key.endswith(".weight"):
+                continue
+            name = key.removesuffix(".weight")
+            attention = name.startswith(("base_lm.", "residual_lm.", "feat_decoder.estimator.")) and name.rsplit(
+                ".", 1
+            )[-1] in {"q_proj", "k_proj", "v_proj", "o_proj"}
+            if not attention and name not in projections:
+                continue
+            out_features, in_features = checkpoint.get_slice(key).get_shape()
+            # Small nonzero deltas exercise every group without retraining a voice.
+            tensors[f"{name}.lora_A"] = torch.randn(2, in_features, generator=generator) * 0.001
+            tensors[f"{name}.lora_B"] = torch.randn(out_features, 2, generator=generator) * 0.001
+    assert tensors
+    save_file(tensors, str(directory / "lora_weights.safetensors"))
+    (directory / "lora_config.json").write_text(
+        json.dumps({"lora_config": {"r": 2, "alpha": 2, "enable_lm": True, "enable_dit": True, "enable_proj": True}})
+    )
+
+    with open(get_deploy_config_path("voxcpm2.yaml")) as stream:
+        config = yaml.safe_load(stream)
+    stage = config["stages"][0]
+    stage.update(
+        max_num_seqs=4,
+        kv_cache_memory_bytes=1024**3,
+        gpu_memory_utilization=0.65,
+        max_model_len=2048,
+        max_num_batched_tokens=2048,
+    )
+    stage["default_sampling_params"]["max_tokens"] = 256
+    runtime = stage["engine_extras"]["hf_overrides"]["voxcpm2_runtime_config"]
+    runtime.update(startup_lora_path=str(directory), unified_decode_graph_max_batch_size=4)
+    deploy_path = directory / "deploy.yaml"
+    deploy_path.write_text(yaml.safe_dump(config))
+    request.param = OmniServerParams(
+        model=MODEL,
+        stage_config_path=str(deploy_path),
+        server_args=["--trust-remote-code", "--disable-log-stats"],
+    )
+    yield from iter_omni_server(request, run_level, omni_fixture_lock)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_startup_lora_concurrent_speech(omni_server, online_client):
+    online_client.send_audio_speech_request(
+        {
+            "model": omni_server.model,
+            "input": "The weather is nice today, perfect for a walk in the park.",
+            "voice": "default",
+            "response_format": "wav",
+            "stream": False,
+            "timeout": 300,
+            "min_audio_bytes": 40_000,
+        },
+        request_num=4,
+    )
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_startup_lora_streaming_speech(omni_server, online_client):
+    online_client.send_audio_speech_request(
+        {
+            "model": omni_server.model,
+            "input": "The weather is nice today, perfect for a walk in the park.",
+            "voice": "default",
+            "response_format": "wav",
+            "stream": True,
+            "stream_format": "audio",
+            "timeout": 300,
+            "min_audio_bytes": 40_000,
+        }
+    )

tests/model_executor/models/voxcpm2/test_startup_lora.py

@@ -0,0 +1,177 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Native VoxCPM2 adapter validation and merge parity, without model weights."""
+
+import importlib.util
+import json
+from pathlib import Path
+
+import pytest
+import torch
+from safetensors.torch import save_file
+from torch import nn
+from torch.nn import functional as F
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture
+def merge_lora():
+    # Keep these L1 checks independent of vLLM custom-op/platform imports.
+    path = Path(__file__).resolve().parents[4] / "vllm_omni/model_executor/models/voxcpm2/lora.py"
+    spec = importlib.util.spec_from_file_location("voxcpm2_startup_lora", path)
+    assert spec is not None and spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module.merge_voxcpm2_lora
+
+
+@pytest.fixture
+def adapter(tmp_path):
+    def attention():
+        return nn.ModuleDict({name: nn.Linear(3, 4, bias=True) for name in ("q_proj", "k_proj", "v_proj", "o_proj")})
+
+    base_lm = nn.ModuleDict({"layers": nn.ModuleList([attention()])})
+    residual_lm = nn.ModuleDict({"layers": nn.ModuleList([attention()])})
+    tts = nn.ModuleDict(
+        {
+            "feat_decoder": nn.ModuleDict({"estimator": nn.ModuleDict({"decoder": attention()})}),
+            "fusion_concat_proj": nn.Linear(3, 4),
+            "stop_head": nn.Linear(3, 2),
+        }
+    )
+    roots = {"base_lm": base_lm, "residual_lm": residual_lm, "tts": tts}
+    native = nn.ModuleDict({"base_lm": base_lm, "residual_lm": residual_lm, **dict(tts.items())})
+    targets = {name: module for name, module in native.named_modules() if isinstance(module, nn.Linear)}
+    targets.pop("stop_head")
+    generator = torch.Generator().manual_seed(42)
+    tensors = {
+        f"{name}.lora_{suffix}": torch.randn(shape, generator=generator)
+        for name in targets
+        for suffix, shape in (("A", (2, 3)), ("B", (4, 2)))
+    }
+    config = {
+        "r": 2,
+        "alpha": 3,
+        "enable_lm": True,
+        "enable_dit": True,
+        "enable_proj": True,
+        "target_proj_modules": ["fusion_concat_proj"],
+    }
+
+    def save():
+        (tmp_path / "lora_config.json").write_text(json.dumps({"lora_config": config}))
+        save_file(tensors, str(tmp_path / "lora_weights.safetensors"))
+
+    save()
+    return tmp_path, roots, targets, tensors, config, save
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_merge_matches_native_linear_output(merge_lora, adapter, dtype):
+    path, roots, targets, tensors, _, _ = adapter
+    inputs = torch.tensor([[0.5, -0.2, 0.1]], dtype=dtype)
+    expected = {}
+    identities = {}
+    for name, layer in targets.items():
+        layer.to(dtype=dtype)
+        identities[name] = id(layer.weight)
+        # Native LoRALinear: linear(x, W, bias) + linear(linear(x, A), B) * alpha/r.
+        expected[name] = F.linear(inputs.float(), layer.weight.float(), layer.bias.float()) + 1.5 * F.linear(
+            F.linear(inputs.float(), tensors[f"{name}.lora_A"]), tensors[f"{name}.lora_B"]
+        )
+    untouched = roots["tts"].stop_head.weight.detach().clone()
+    assert merge_lora(str(path), **roots) == 13
+    for name, layer in targets.items():
+        tolerance = 0.025 if dtype == torch.bfloat16 else 1e-6
+        torch.testing.assert_close(layer(inputs).float(), expected[name], atol=tolerance, rtol=tolerance)
+        assert id(layer.weight) == identities[name]
+        assert layer.weight.dtype == dtype
+        assert set(dict(layer.named_parameters())) == {"weight", "bias"}
+    torch.testing.assert_close(roots["tts"].stop_head.weight, untouched, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("group", ["lm", "dit", "proj"])
+def test_independently_enabled_groups(merge_lora, adapter, group):
+    path, roots, targets, tensors, config, save = adapter
+    for key in ("lm", "dit", "proj"):
+        config[f"enable_{key}"] = key == group
+    prefixes = {"lm": ("base_lm.", "residual_lm."), "dit": ("feat_decoder.",), "proj": ("fusion_concat_proj.",)}
+    for key in list(tensors):
+        if not key.startswith(prefixes[group]):
+            del tensors[key]
+    original = {name: layer.weight.detach().clone() for name, layer in targets.items()}
+    save()
+    assert merge_lora(str(path), **roots) == {"lm": 8, "dit": 4, "proj": 1}[group]
+    for name, layer in targets.items():
+        if not (name + ".").startswith(prefixes[group]):
+            torch.testing.assert_close(layer.weight, original[name], rtol=0, atol=0)
+
+
+@pytest.mark.parametrize(
+    "corruption", ["missing", "extra", "shape", "rank", "alpha", "nonfinite", "dtype", "target", "flag"]
+)
+def test_rejects_invalid_checkpoint_before_mutating_weights(merge_lora, adapter, corruption):
+    path, roots, targets, tensors, config, save = adapter
+    key = "fusion_concat_proj.lora_B"  # Validate even the last component before merging any layer.
+    if corruption == "missing":
+        del tensors[key]
+    elif corruption == "extra":
+        tensors["audio_vae.lora_A"] = torch.ones(2, 3)
+    elif corruption == "shape":
+        tensors[key] = torch.ones(5, 2)
+    elif corruption == "rank":
+        config["r"] = 0
+    elif corruption == "alpha":
+        config["alpha"] = float("nan")
+    elif corruption == "nonfinite":
+        tensors[key][0, 0] = float("nan")
+    elif corruption == "dtype":
+        tensors[key] = tensors[key].long()
+    elif corruption == "target":
+        config["target_modules_lm"] = ["does_not_exist"]
+    elif corruption == "flag":
+        config["enable_dit"] = "false"
+    save()
+    original = {name: layer.weight.detach().clone() for name, layer in targets.items()}
+    with pytest.raises(ValueError, match="VoxCPM2"):
+        merge_lora(str(path), **roots)
+    for name, layer in targets.items():
+        torch.testing.assert_close(layer.weight, original[name], rtol=0, atol=0)
+
+
+def test_requires_native_metadata_and_safetensors(merge_lora, adapter):
+    path, roots, _, _, _, save = adapter
+    (path / "lora_config.json").write_text(json.dumps({"r": 2, "alpha": 3}))
+    with pytest.raises(ValueError, match="lora_config object"):
+        merge_lora(str(path), **roots)
+    save()
+    (path / "lora_weights.safetensors").unlink()
+    with pytest.raises(FileNotFoundError):
+        merge_lora(str(path), **roots)
+
+
+def test_talker_loads_base_before_merge_and_rejects_reload(adapter):
+    pytest.importorskip("vllm")
+    from vllm_omni.model_executor.models.voxcpm2.runtime_config import _VoxCPM2RuntimeConfig
+    from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import VoxCPM2TalkerForConditionalGeneration
+
+    path, roots, targets, tensors, _, _ = adapter
+    talker = VoxCPM2TalkerForConditionalGeneration.__new__(VoxCPM2TalkerForConditionalGeneration)
+    nn.Module.__init__(talker)
+    talker.model = roots["base_lm"]
+    talker.residual_model = roots["residual_lm"]
+    talker._tts = roots["tts"]
+    talker._runtime_config = _VoxCPM2RuntimeConfig(startup_lora_path=str(path))
+    talker._startup_lora_applied = False
+    talker._patch_size, talker._feat_dim, talker._side_dtype = 4, 64, torch.float32
+    weights = [(f"base_lm.{name}", torch.ones_like(param)) for name, param in talker.model.named_parameters()]
+    loaded = talker.load_weights(iter(weights))
+    name = "base_lm.layers.0.q_proj"
+    expected = 1 + 1.5 * tensors[f"{name}.lora_B"] @ tensors[f"{name}.lora_A"]
+    torch.testing.assert_close(targets[name].weight, expected)
+    assert "model.layers.0.q_proj.weight" in loaded
+    assert talker._startup_lora_applied
+    with pytest.raises(ValueError, match="restart the server"):
+        talker.load_weights(iter(weights))
+    torch.testing.assert_close(targets[name].weight, expected)

vllm_omni/model_executor/models/voxcpm2/lora.py

@@ -0,0 +1,95 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Merge a native VoxCPM2 LoRA checkpoint at startup for all requests."""
+
+from __future__ import annotations
+
+import json
+import math
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file
+from torch import nn
+
+
+def merge_voxcpm2_lora(
+    adapter_path: str,
+    *,
+    base_lm: nn.Module,
+    residual_lm: nn.Module,
+    tts: nn.Module,
+) -> int:
+    """Validate and merge one local native safetensors adapter.
+
+    Call once after base weight loading and before fusion or graph capture.
+    Validate all checkpoint keys and tensors before changing weights.
+    """
+    path = Path(adapter_path)
+    with (path / "lora_config.json").open(encoding="utf-8") as f:
+        metadata = json.load(f)
+    config = metadata.get("lora_config") if isinstance(metadata, dict) else None
+    if not isinstance(config, dict):
+        raise ValueError("VoxCPM2 lora_config.json must contain a lora_config object")
+    rank, alpha = config.get("r"), config.get("alpha")
+    if type(rank) is not int or rank <= 0:
+        raise ValueError("VoxCPM2 LoRA r must be a positive integer")
+    if isinstance(alpha, bool) or not isinstance(alpha, (int, float)) or not math.isfinite(alpha):
+        raise ValueError("VoxCPM2 LoRA alpha must be finite")
+
+    targets: dict[str, nn.Linear] = {}
+    groups = (
+        ("lm", {"base_lm": base_lm, "residual_lm": residual_lm}, "target_modules_lm"),
+        ("dit", {"feat_decoder.estimator": tts.feat_decoder.estimator}, "target_modules_dit"),
+        ("proj", {"": tts}, "target_proj_modules"),
+    )
+    for group, roots, target_key in groups:
+        enabled = config.get(f"enable_{group}", False)
+        if type(enabled) is not bool:
+            raise ValueError(f"VoxCPM2 LoRA enable_{group} must be a boolean")
+        if not enabled:
+            continue
+        defaults = (
+            ["enc_to_lm_proj", "lm_to_dit_proj", "res_to_dit_proj", "fusion_concat_proj"]
+            if group == "proj"
+            else ["q_proj", "v_proj", "k_proj", "o_proj"]
+        )
+        names = config.get(target_key, defaults)
+        if not isinstance(names, list) or not names or any(not isinstance(n, str) or not n for n in names):
+            raise ValueError(f"VoxCPM2 LoRA {target_key} must be a non-empty list of names")
+        for prefix, root in roots.items():
+            matched = set()
+            for name, module in root.named_modules():
+                selector = name if group == "proj" else name.rsplit(".", 1)[-1]
+                if selector in names and isinstance(module, nn.Linear):
+                    targets[f"{prefix}.{name}" if prefix else name] = module
+                    matched.add(selector)
+            if missing := set(names) - matched:
+                raise ValueError(f"VoxCPM2 LoRA targets not found in {prefix or 'tts'}: {sorted(missing)}")
+    if not targets:
+        raise ValueError("VoxCPM2 LoRA config enables no target layers")
+
+    state = load_file(str(path / "lora_weights.safetensors"), device="cpu")
+    expected = {f"{name}.lora_{suffix}" for name in targets for suffix in ("A", "B")}
+    if set(state) != expected:
+        raise ValueError(
+            "VoxCPM2 LoRA checkpoint keys do not match the configured targets: "
+            f"missing={sorted(expected - set(state))}, unexpected={sorted(set(state) - expected)}"
+        )
+    for name, module in targets.items():
+        if not module.weight.is_floating_point() or module.weight.is_meta:
+            raise ValueError(f"VoxCPM2 LoRA requires materialized floating-point weights: {name}")
+        for suffix, shape in (("A", (rank, module.in_features)), ("B", (module.out_features, rank))):
+            key = f"{name}.lora_{suffix}"
+            tensor = state[key]
+            if tensor.shape != shape or not tensor.is_floating_point() or not torch.isfinite(tensor).all():
+                raise ValueError(f"Invalid VoxCPM2 LoRA tensor {key}: expected finite floating-point shape {shape}")
+
+    with torch.no_grad():
+        for name, module in targets.items():
+            # Merge one layer at a time in fp32, preserving the base Parameter.
+            a = state[f"{name}.lora_A"].to(device=module.weight.device, dtype=torch.float32)
+            b = state[f"{name}.lora_B"].to(device=module.weight.device, dtype=torch.float32)
+            merged = torch.addmm(module.weight.float(), b, a, alpha=alpha / rank)
+            module.weight.copy_(merged)
+    return len(targets)

vllm_omni/model_executor/models/voxcpm2/runtime_config.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
 
@@ -13,6 +13,7 @@ logger = init_logger(__name__)
 
 @dataclasses.dataclass(frozen=True)
 class _VoxCPM2RuntimeConfig:
+    startup_lora_path: str | None = None
     enable_profiling: bool = False
     enable_nvtx_profile: bool = False
     enable_loc_dit_layer_nvtx: bool = False
@@ -105,6 +106,10 @@ class _VoxCPM2RuntimeConfig:
 
     @staticmethod
     def _coerce_value(key: str, value: Any, default: Any) -> Any:
+        if key == "startup_lora_path":
+            if value is not None and (not isinstance(value, str) or not value.strip()):
+                raise ValueError("VoxCPM2 startup_lora_path must be a non-empty local directory path or null")
+            return value
         if isinstance(default, bool):
             if isinstance(value, str):
                 return value.strip().lower() in ("1", "true", "yes", "on")

vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py

@@ -46,6 +46,7 @@ from vllm_omni.utils.speaker_cache import (
 )
 from vllm_omni.worker.runner_assisted_metadata import RunnerAssistedFullAttentionMetadataRequest
 
+from .lora import merge_voxcpm2_lora
 from .minicpm4_paged import MiniCPM4PagedForVoxCPM2, MiniCPM4PagedResidualLM
 from .runtime_config import _VoxCPM2RuntimeConfig
 from .voxcpm2_import_utils import import_voxcpm2_core
@@ -834,6 +835,9 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
         self.vllm_config = vllm_config
         self.config = vllm_config.model_config.hf_config
         self._runtime_config = _VoxCPM2RuntimeConfig.from_vllm_config(vllm_config)
+        self._startup_lora_applied = False
+        if self._runtime_config.startup_lora_path and vllm_config.load_config.load_format == "dummy":
+            raise ValueError("VoxCPM2 startup LoRA requires real base weights; load_format=dummy is unsupported")
         global _ENABLE_NVTX_PROFILE
         _ENABLE_NVTX_PROFILE = self._runtime_config.enable_nvtx_profile
 
@@ -3284,6 +3288,9 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
     hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"base_lm.": "model."})
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        if self._startup_lora_applied:
+            raise ValueError("Reloading VoxCPM2 weights with a startup LoRA is unsupported; restart the server")
+
         def _base_lm_only(ws):
             for name, tensor in ws:
                 if name.startswith("base_lm."):
@@ -3298,6 +3305,13 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
         # them as missing from the checkpoint.
         loaded |= {name for name, _ in self.named_parameters() if name.startswith(("_tts.", "residual_model."))}
 
+        if adapter_path := self._runtime_config.startup_lora_path:
+            merged = merge_voxcpm2_lora(
+                adapter_path, base_lm=self.model, residual_lm=self.residual_model, tts=self._tts
+            )
+            self._startup_lora_applied = True
+            logger.info("Merged VoxCPM2 startup LoRA from %s into %d linear layers", adapter_path, merged)
+
         logger.info(
             "Loaded VoxCPM2 (patch=%d, feat_dim=%d, dtype=%s)", self._patch_size, self._feat_dim, self._side_dtype
         )