How to Merge LoRA Weights into the Base Fish-Speech Model
Merging LoRA weights into the Fish-Speech base model requires instantiating the base architecture with matching LoRA hyperparameters, unioning the base and adapter state dictionaries, triggering weight folding via eval() mode, and saving the result with drop_lora=True to strip adapter parameters and produce a clean inference checkpoint.
When fine-tuning Fish-Speech using Low-Rank Adaptation (LoRA), the training process produces a small adapter checkpoint containing only the low-rank matrices. To deploy this efficiently or distribute a standalone model, you must merge LoRA weights into the base model, folding the adapters back into the original linear and embedding layers. The fishaudio/fish-speech repository implements this workflow in tools/llama/merge_lora.py with support for both CLI and programmatic execution.
The Three-Step Merge Process
The merge operation follows a strict sequence to ensure numerical correctness and weight compatibility.
Step 1: Load the Base Model with LoRA Configuration
First, instantiate the base model while passing the same LoRA hyperparameters (r, lora_alpha, lora_dropout) used during training. In tools/llama/merge_lora.py (lines 33-38), the code calls:
llama_model = BaseTransformer.from_pretrained(
base_weight,
load_weights=True,
lora_config=lora_config,
)
This creates a model whose internal layers contain both the original pretrained weights and uninitialized LoRA adapter slots (added by setup_lora() in fish_speech/models/text2semantic/lora.py). The architecture must match exactly to ensure dimension compatibility between the base weights and the adapter matrices.
Step 2: Merge the Base and LoRA State Dictionaries
Next, load the LoRA checkpoint (typically model.pth from your training run) and combine it with the base model's state. The script handles potential "model." key prefixes (lines 51-63) and then performs the merge using Python's dictionary union operator (line 68):
merged_state_dict = llama_state_dict | lora_state_dict
llama_model.load_state_dict(merged_state_dict, strict=True)
This operation overlays the trained low-rank matrices (A and B) onto the empty adapter slots created in Step 1, while preserving all original base parameters.
Step 3: Fold Adapters and Save the Merged Checkpoint
Finally, trigger the mathematical folding of LoRA weights into the base weights and export the result. Calling llama_model.eval() (line 73) forces every LoRA-augmented layer to compute its effective weight (W + (A · B) * (alpha / r)). The save_pretrained() method with drop_lora=True (line 74 in merge_lora.py, implemented in fish_speech/models/text2semantic/llama.py) writes a new checkpoint containing only the folded weights:
llama_model.eval()
llama_model.save_pretrained(output, drop_lora=True)
The script concludes with a validation loop (lines 77-92) that compares specific keys between the merged model and the original base to verify numerical tolerance (1e-5).
CLI Method: Using merge_lora.py
The repository provides a ready-to-use command-line interface in tools/llama/merge_lora.py. Execute the merge with a single command:
python -m tools.llama.merge_lora \
--lora-config r_8_alpha_16 \
--base-weight checkpoints/fish-speech-1.4 \
--lora-weight checkpoints/fish-speech-lora/model.pth \
--output merged/fish-speech-merged
Required arguments:
--lora-config: The Hydra config name (e.g.,r_8_alpha_16) located infish_speech/configs/lora/that specifiesr,lora_alpha, andlora_dropout--base-weight: Path to the original Fish-Speech checkpoint directory--lora-weight: Path to the LoRA adapter file (model.pth)--output: Destination directory for the merged checkpoint
The utility uses loguru for progress logging and automatically validates the merge before completion.
Python API: Programmatic Merging
For integration into custom pipelines, replicate the CLI logic using the Python API. The get_merged_state_dict() helper in fish_speech/models/text2semantic/lora.py (lines 81-91) demonstrates the folding mechanism, though the manual approach below offers more control:
from pathlib import Path
import torch
from fish_speech.models.text2semantic.llama import BaseTransformer
from hydra import compose, initialize
from hydra.utils import instantiate
def merge_lora(base_path: str, lora_path: str, cfg_name: str, out_path: str):
# Load LoRA configuration via Hydra
with initialize(version_base="1.3", config_path="fish_speech/configs/lora"):
cfg = compose(config_name=cfg_name)
lora_cfg = instantiate(cfg)
# Initialize base model with empty LoRA adapters
model = BaseTransformer.from_pretrained(
base_path, load_weights=True, lora_config=lora_cfg
)
# Load LoRA weights and merge via dictionary union
lora_state = torch.load(lora_path, map_location="cpu")["state_dict"]
base_state = {k: v for k, v in model.state_dict().items() if "lora" not in k}
merged_state = base_state | lora_state
model.load_state_dict(merged_state, strict=True)
# Fold adapters and save clean checkpoint
model.eval()
model.save_pretrained(Path(out_path), drop_lora=True)
# Example execution
merge_lora(
base_path="checkpoints/fish-speech-1.4",
lora_path="checkpoints/fish-speech-lora/model.pth",
cfg_name="r_8_alpha_16",
out_path="merged/fish-speech-merged",
)
This approach mirrors the CLI implementation exactly but allows for custom preprocessing, batch merging, or cloud-based inference pipelines.
Key Source Files and Implementation Details
Understanding these core files helps debug merge issues or extend functionality:
-
tools/llama/merge_lora.py: The command-line entry point containing the full merge logic, argument parsing (lines 16-21), and validation checks. -
fish_speech/models/text2semantic/lora.py: Implementssetup_lora()to inject LoRA wrappers into linear and embedding layers, andget_merged_state_dict()(lines 81-91) as a programmatic alternative for folding and stripping LoRA parameters. -
fish_speech/models/text2semantic/llama.py: DefinesBaseTransformerwithfrom_pretrained()for loading andsave_pretrained()supporting thedrop_loraflag to exclude adapter tensors from the output. -
fish_speech/configs/lora/: Hydra configuration directory containing YAML files (e.g.,r_8_alpha_16.yaml) that must match the hyperparameters used during your original LoRA training.
Summary
- Initialize with config: Load the base model using
BaseTransformer.from_pretrained()with the identical LoRA configuration to create compatible adapter slots. - Union state dicts: Merge base and LoRA checkpoints using the dictionary union operator (
|) after stripping any "model." prefixes, then load viaload_state_dict(strict=True). - Fold and export: Call
eval()to trigger weight folding (W + A·B), thensave_pretrained(drop_lora=True)to write a clean inference checkpoint. - Validate: The merge script automatically verifies key weights match within 1e-5 tolerance to ensure numerical correctness.
- Use provided tools: The
merge_lora.pyCLI handles the entire workflow, while the Python API offers flexibility for custom integration.
Frequently Asked Questions
What does drop_lora=True do when saving the merged model?
The drop_lora=True parameter in save_pretrained() filters the model's state dictionary to exclude all keys containing "lora" (specifically the low-rank matrices A and B) before writing the checkpoint. According to the implementation in fish_speech/models/text2semantic/llama.py, this ensures the output file contains only the folded base weights, reducing file size and eliminating adapter overhead during inference.
Why must the LoRA configuration match when loading the base model?
The LoRA configuration (r, lora_alpha, lora_dropout) determines the dimensions and scaling of the low-rank matrices. When you call BaseTransformer.from_pretrained() with lora_config, setup_lora() in lora.py replaces standard layers with LoRA-wrapped versions expecting specific shapes. If the configuration mismatches the trained adapter checkpoint, the state dictionary merge will fail due to tensor shape incompatibilities or missing keys.
Can I merge multiple LoRA checkpoints into one base model?
The standard merge_lora.py utility handles single LoRA checkpoint merging. To merge multiple LoRA adapters, you would need to manually load each LoRA state dictionary, average or sum the adapter weights (depending on your ensemble strategy), merge the result with the base state dict, and then proceed with the eval-and-save workflow. The repository does not provide a built-in multi-LoRA merge command.
How do I verify the merge completed successfully?
The merge script includes an automatic validation step (lines 77-92 in merge_lora.py) that loads the newly saved checkpoint and compares specific weight keys against the pre-merge model with a tolerance of 1e-5. You can also perform manual verification by loading the merged checkpoint without LoRA configuration and attempting inference—successful generation confirms the weights folded correctly into the base architecture.
Have a question about this repo?
These articles cover the highlights, but your codebase questions are specific. Give your agent direct access to the source. Share this with your agent to get started:
curl -s "https://instagit.com/install.md" Maintain an open-source project? Get it listed too →