From 04bfbbae614ff9cb9ea108f6e6b3971cab2f18d4 Mon Sep 17 00:00:00 2001
From: "Guanzhe (Leo) Huang" <116847811+leohuang257@users.noreply.github.com>
Date: Fri, 1 May 2026 20:13:16 -0400
Subject: [PATCH] [BugFix] voxcpm2: eager-init tts_model in __init__ so
 load_format: dummy works (#3229)

Signed-off-by: leohuang257 <masugchds@gmail.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
---
 .../models/voxcpm2/voxcpm2_talker.py          | 53 ++++++++++---------
 1 file changed, 28 insertions(+), 25 deletions(-)

diff --git a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py
index 0a9246251..15d125b53 100644
--- a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py
+++ b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py
@@ -408,14 +408,32 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
         )
         self.make_empty_intermediate_tensors = self.model.make_empty_intermediate_tensors
 
-        self._tts: nn.Module | None = None
+        # Eager-init tts_model so it registers in self.state_dict() before vLLM's
+        # post-__init__ profiling claims the remaining GPU memory for KV cache.
+        # Required for load_format=dummy: DummyModelLoader only randomizes
+        # already-registered nn.Parameters.
+        # NOTE: from_pretrained() is unconditional, so load_format=dummy still pays
+        # the checkpoint download/read cost at construction time; DummyModelLoader
+        # will then randomize the just-loaded _tts params — this is intended.
+        model_path = vllm_config.model_config.model
+        VoxCPM = import_voxcpm2_core()
+        native = VoxCPM.from_pretrained(model_path, load_denoiser=False, optimize=False)
+        self._tts: nn.Module = native.tts_model.to("cuda")
+        self._side_dtype = self._tts.fusion_concat_proj.weight.dtype
         self._device = "cuda"
-        self._side_dtype = torch.bfloat16
-
-        self._patch_size = getattr(self.config, "patch_size", 4)
-        self._feat_dim = getattr(self.config, "feat_dim", 64)
+        self._patch_size = self._tts.patch_size
+        self._feat_dim = self._tts.feat_dim
         self._sample_rate = getattr(self.config, "sample_rate", 48000)
 
+        # base_lm/residual_lm in native tts_model duplicate self.model and
+        # self.residual_model: copy residual weights over, drop both submodules.
+        self.residual_model.load_weights_from_native(self._tts.residual_lm)
+        del self._tts.base_lm
+        self._tts.base_lm = None
+        del self._tts.residual_lm
+        self._tts.residual_lm = None
+        torch.cuda.empty_cache()
+
         self._inference_timesteps = 10
         self._cfg_value = 2.0
         self._cfg_cutoff_ratio = 1.0
@@ -452,7 +470,6 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
 
     @property
     def tts(self) -> nn.Module:
-        assert self._tts is not None, "Model not loaded yet"
         return self._tts
 
     # -------------------- request state management --------------------
@@ -1264,25 +1281,11 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
         loader = AutoWeightsLoader(self)
         loaded = loader.load_weights(_base_lm_only(weights), mapper=self.hf_to_vllm_mapper)
 
-        model_path = self.vllm_config.model_config.model
-        VoxCPM = import_voxcpm2_core()
-        native = VoxCPM.from_pretrained(model_path, load_denoiser=False, optimize=False)
-        self._tts = native.tts_model.to("cuda")
-        self._side_dtype = self._tts.fusion_concat_proj.weight.dtype
-        self._device = "cuda"
-        self._patch_size = self._tts.patch_size
-        self._feat_dim = self._tts.feat_dim
-
-        n = self.residual_model.load_weights_from_native(self._tts.residual_lm)
-        for name, _ in self.residual_model.named_parameters():
-            loaded.add(f"residual_model.{name}")
-        logger.info("VoxCPM2: loaded %d params into paged residual_model", n)
-
-        del self._tts.base_lm
-        self._tts.base_lm = None
-        del self._tts.residual_lm
-        self._tts.residual_lm = None
-        torch.cuda.empty_cache()
+        # _tts and residual_model are constructed and populated eagerly in
+        # __init__ via VoxCPM.from_pretrained; here we only need to mark their
+        # params as loaded so AutoWeightsLoader's strict check doesn't flag
+        # them as missing from the checkpoint.
+        loaded |= {name for name, _ in self.named_parameters() if name.startswith(("_tts.", "residual_model."))}
 
         logger.info(
             "Loaded VoxCPM2 (patch=%d, feat_dim=%d, dtype=%s)", self._patch_size, self._feat_dim, self._side_dtype