Add generic native LoRA training support

cosyvoice/bin/train.py

@@ -35,6 +35,7 @@
     init_optimizer_and_scheduler,
     init_summarywriter, save_model,
     wrap_cuda_model, check_modify_and_save_config)
+from cosyvoice.utils.lora import inject_lora, load_lora_state_dict
 
 
 def get_args():
@@ -51,6 +52,12 @@ def get_args():
     parser.add_argument('--qwen_pretrain_path', required=False, help='qwen pretrain path')
     parser.add_argument('--onnx_path', required=False, help='onnx path, which is required for online feature extraction')
     parser.add_argument('--checkpoint', help='checkpoint model')
+    parser.add_argument('--lora', action='store_true', default=False,
+                        help='freeze the model and train native LoRA adapters')
+    parser.add_argument('--lora_checkpoint', help='adapter-only checkpoint to resume')
+    parser.add_argument('--lora_rank', type=int, default=16)
+    parser.add_argument('--lora_alpha', type=float, default=32.0)
+    parser.add_argument('--lora_dropout', type=float, default=0.05)
     parser.add_argument('--model_dir', required=True, help='save model dir')
     parser.add_argument('--tensorboard_dir',
                         default='tensorboard',
@@ -143,6 +150,18 @@ def main():
         else:
             logging.warning('checkpoint {} do not exsist!'.format(args.checkpoint))
 
+    if args.lora:
+        target_count = inject_lora(model, args.lora_rank, args.lora_alpha, args.lora_dropout)
+        logging.info('Injected LoRA into %s projection layers; trainable parameters=%s',
+                     target_count, sum(p.numel() for p in model.parameters() if p.requires_grad))
+        if args.lora_checkpoint is not None:
+            adapter_state = torch.load(args.lora_checkpoint, map_location='cpu')
+            load_lora_state_dict(model, adapter_state)
+            if 'step' in adapter_state:
+                start_step = int(adapter_state['step'])
+            if 'epoch' in adapter_state:
+                start_epoch = int(adapter_state['epoch'])
+
     # Dispatch model from cpu to gpu
     model = wrap_cuda_model(args, model)
 

cosyvoice/utils/lora.py

@@ -0,0 +1,113 @@
+"""Small native LoRA implementation for CosyVoice adaptation.
+
+This intentionally has no PEFT dependency. It injects adapters into the Qwen
+projection layers while leaving the original weights available for the
+unmodified base model.
+"""
+
+from __future__ import annotations
+
+from typing import Iterable
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+class LoRALinear(nn.Module):
+    def __init__(self, base: nn.Linear, rank: int, alpha: float, dropout: float):
+        super().__init__()
+        if rank < 1:
+            raise ValueError("LoRA rank must be positive")
+        self.base = base
+        self.rank = rank
+        self.alpha = alpha
+        self.scaling = alpha / rank
+        self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
+        self.lora_A = nn.Parameter(torch.empty(rank, base.in_features))
+        self.lora_B = nn.Parameter(torch.zeros(base.out_features, rank))
+        nn.init.kaiming_uniform_(self.lora_A, a=5**0.5)
+        for parameter in self.base.parameters():
+            parameter.requires_grad = False
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        update = F.linear(F.linear(self.dropout(x), self.lora_A), self.lora_B)
+        return self.base(x) + update * self.scaling
+
+    @torch.no_grad()
+    def merge_(self) -> None:
+        self.base.weight.add_(self.lora_B @ self.lora_A, alpha=self.scaling)
+
+
+TARGET_LINEAR_NAMES = {
+    "q_proj",
+    "k_proj",
+    "v_proj",
+    "o_proj",
+    "gate_proj",
+    "up_proj",
+    "down_proj",
+}
+
+
+def _replace_target_linears(module: nn.Module, rank: int, alpha: float, dropout: float) -> int:
+    count = 0
+    for name, child in list(module.named_children()):
+        if isinstance(child, nn.Linear) and name in TARGET_LINEAR_NAMES:
+            setattr(module, name, LoRALinear(child, rank, alpha, dropout))
+            count += 1
+        else:
+            count += _replace_target_linears(child, rank, alpha, dropout)
+    return count
+
+
+def inject_lora(model: nn.Module, rank: int = 16, alpha: float = 32.0, dropout: float = 0.05) -> int:
+    """Freeze the model and inject LoRA into Qwen projections.
+
+    The CosyVoice speech embedding and decoder remain trainable because they
+    are the language-adaptation head; all other parameters remain frozen.
+    """
+    for parameter in model.parameters():
+        parameter.requires_grad = False
+
+    count = _replace_target_linears(model, rank, alpha, dropout)
+    train_head_names = ("llm_decoder", "speech_embedding")
+    for name, parameter in model.named_parameters():
+        if any(
+            name == head
+            or name.startswith(f"{head}.")
+            or f".{head}." in name
+            for head in train_head_names
+        ):
+            parameter.requires_grad = True
+
+    trainable = [parameter for parameter in model.parameters() if parameter.requires_grad]
+    if not trainable:
+        raise RuntimeError("LoRA injection produced no trainable parameters")
+    model._lora_enabled = True
+    model._lora_target_count = count
+    return count
+
+
+def lora_state_dict(model: nn.Module) -> dict[str, torch.Tensor]:
+    """Return only LoRA and adaptation-head weights for a small checkpoint."""
+    trainable_names = {name for name, p in model.named_parameters() if p.requires_grad}
+    return {
+        name: value.detach().cpu()
+        for name, value in model.state_dict().items()
+        if name in trainable_names
+    }
+
+
+def load_lora_state_dict(model: nn.Module, state: dict[str, torch.Tensor]) -> None:
+    missing, unexpected = model.load_state_dict(state, strict=False)
+    unexpected = [name for name in unexpected if name not in {"step", "epoch"}]
+    if unexpected:
+        raise RuntimeError(f"Unexpected LoRA checkpoint keys: {unexpected[:8]}")
+    missing_trainable = [name for name in missing if name in state]
+    if missing_trainable:
+        raise RuntimeError(f"Could not load LoRA checkpoint keys: {missing_trainable[:8]}")
+
+
+def trainable_parameters(model: nn.Module) -> Iterable[nn.Parameter]:
+    return (parameter for parameter in model.parameters() if parameter.requires_grad)

cosyvoice/utils/train_utils.py

@@ -110,10 +110,13 @@ def wrap_cuda_model(args, model):
 
 def init_optimizer_and_scheduler(args, configs, model, gan):
     if gan is False:
+        trainable_parameters = [p for p in model.parameters() if p.requires_grad]
+        if not trainable_parameters:
+            raise RuntimeError('No trainable parameters found')
         if configs['train_conf']['optim'] == 'adam':
-            optimizer = optim.Adam(model.parameters(), **configs['train_conf']['optim_conf'])
+            optimizer = optim.Adam(trainable_parameters, **configs['train_conf']['optim_conf'])
         elif configs['train_conf']['optim'] == 'adamw':
-            optimizer = optim.AdamW(model.parameters(), **configs['train_conf']['optim_conf'])
+            optimizer = optim.AdamW(trainable_parameters, **configs['train_conf']['optim_conf'])
         else:
             raise ValueError("unknown optimizer: " + configs['train_conf'])
 
@@ -199,7 +202,12 @@ def save_model(model, model_name, info_dict):
 
     if info_dict["train_engine"] == "torch_ddp":
         if rank == 0:
-            torch.save({**model.module.state_dict(), 'epoch': info_dict['epoch'], 'step': info_dict['step']}, save_model_path)
+            if getattr(model.module, '_lora_enabled', False):
+                from cosyvoice.utils.lora import lora_state_dict
+                state = lora_state_dict(model.module)
+            else:
+                state = model.module.state_dict()
+            torch.save({**state, 'epoch': info_dict['epoch'], 'step': info_dict['step']}, save_model_path)
     else:
         with torch.no_grad():
             model.save_checkpoint(save_dir=model_dir,

docs/lora.md

@@ -0,0 +1,36 @@
+# Native LoRA adaptation
+
+CosyVoice can adapt its language-model stage with native LoRA adapters without
+adding a PEFT dependency or changing the base checkpoint.
+
+## Training
+
+Use the existing LLM training entry point and enable LoRA:
+
+```bash
+PYTHONPATH=.:third_party/Matcha-TTS \
+python cosyvoice/bin/train.py \
+  --train_engine torch_ddp \
+  --config path/to/config.yaml \
+  --train_data path/to/train.data.list \
+  --cv_data path/to/dev.data.list \
+  --model llm \
+  --checkpoint path/to/llm.pt \
+  --model_dir experiments/lora \
+  --tensorboard_dir tensorboard/lora \
+  --lora --lora_rank 16 --lora_alpha 32 --lora_dropout 0.05
+```
+
+LoRA freezes the base model, injects adapters into the attention and MLP
+projections, and keeps the speech embedding and decoder heads trainable. The
+optimizer receives only trainable parameters.
+
+To continue from an adapter checkpoint, use `--lora_checkpoint`. The adapter
+checkpoint contains only trainable adapter/head weights plus `epoch` and
+`step`, so the original base checkpoint is still required.
+
+## Inference
+
+Construct the matching base model, inject the same rank and alpha, and load the
+adapter state with `cosyvoice.utils.lora.load_lora_state_dict`. The base model
+remains available for fallback, comparison, or a later merge operation.

tests/test_lora.py

@@ -0,0 +1,56 @@
+import copy
+import unittest
+
+import torch
+from torch import nn
+
+from cosyvoice.utils.lora import (
+    LoRALinear,
+    inject_lora,
+    load_lora_state_dict,
+    lora_state_dict,
+)
+
+
+class TinyLanguageModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.q_proj = nn.Linear(4, 4)
+        self.ff = nn.Linear(4, 4)
+        self.llm_decoder = nn.Linear(4, 4)
+        self.speech_embedding = nn.Embedding(8, 4)
+
+    def forward(self, x):
+        return self.llm_decoder(self.q_proj(x) + self.ff(x))
+
+
+class NativeLoRATest(unittest.TestCase):
+    def test_injection_freezes_base_and_keeps_adaptation_heads_trainable(self):
+        model = TinyLanguageModel()
+        count = inject_lora(model, rank=2, alpha=4, dropout=0.0)
+
+        self.assertEqual(count, 1)
+        self.assertIsInstance(model.q_proj, LoRALinear)
+        self.assertFalse(model.q_proj.base.weight.requires_grad)
+        self.assertTrue(model.q_proj.lora_A.requires_grad)
+        self.assertTrue(model.llm_decoder.weight.requires_grad)
+        self.assertTrue(model.speech_embedding.weight.requires_grad)
+        self.assertFalse(model.ff.weight.requires_grad)
+
+    def test_initial_adapter_is_a_noop_and_state_round_trips(self):
+        model = TinyLanguageModel()
+        baseline = copy.deepcopy(model)
+        inject_lora(model, rank=2, alpha=4, dropout=0.0)
+        sample = torch.randn(3, 4)
+
+        torch.testing.assert_close(model(sample), baseline(sample))
+        state = lora_state_dict(model)
+        restored = TinyLanguageModel()
+        inject_lora(restored, rank=2, alpha=4, dropout=0.0)
+        load_lora_state_dict(restored, state)
+        for name, value in state.items():
+            torch.testing.assert_close(restored.state_dict()[name], value)
+
+
+if __name__ == "__main__":
+    unittest.main()