Quantize Qwen2 embed_tokens before moving module tensors to CUDA to avoid fp16 embedding residency during load.

Result: {"status":"keep","peak_rss_mb":4261.316,"cuda_peak_allocated_mb":1753.854,"log_mel_mse":0,"snr_db":167.147,"length_delta_ms":0,"duration_s":2.16,"elapsed_s":45.363}

cosyvoice/cli/model.py

@@ -120,7 +120,11 @@ class CosyVoiceModel:
 
     def _move_module_tensors(self, module: torch.nn.Module):
         target_dtype = torch.float16 if self.fp16 else torch.float32
-        for child in module.children():
+        for name, child in list(module.named_children()):
+            if os.environ.get('COSYVOICE_LLM_EMBED_INT8') == '1' and name == 'embed_tokens' and isinstance(child, torch.nn.Embedding):
+                from cosyvoice.llm.llm import Int8WeightOnlyEmbedding
+                child = Int8WeightOnlyEmbedding(child)
+                module._modules[name] = child
             self._move_module_tensors(child)
         for name, param in list(module._parameters.items()):
             if param is None:

cosyvoice/llm/llm.py

@@ -45,6 +45,8 @@ class Int8WeightOnlyEmbedding(torch.nn.Module):
 
 
 def quantize_qwen_embed_int8(qwen_model: torch.nn.Module):
+    if isinstance(qwen_model.embed_tokens, Int8WeightOnlyEmbedding):
+        return
     qwen_model.embed_tokens = Int8WeightOnlyEmbedding(qwen_model.embed_tokens)