Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
326 lines
19 KiB
Python
326 lines
19 KiB
Python
import json
|
|
from pathlib import Path
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
import torch
|
|
import gguf
|
|
import numpy as np
|
|
from safetensors.torch import save_file
|
|
from safetensors import SafetensorError
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "tools/lora"))
|
|
from bindings import BaseGGUF, BaseSources, Bindings
|
|
from merge_lora import main, normalize, read_adapter, sidecars
|
|
from checkpoint import EXTENSIONS, find_checkpoint, load_checkpoint
|
|
|
|
|
|
class MergeLoraTests(unittest.TestCase):
|
|
def test_cosy_mixed_package_selection_and_sidecars(self):
|
|
name = "model.layers.0.q_proj"
|
|
bindings = Bindings("cosyvoice3", {
|
|
"llm/llm.model." + name + ".weight": (4, 4),
|
|
"llm/speech_embedding.weight": (4, 4),
|
|
"llm/llm_decoder.weight": (4, 4),
|
|
})
|
|
native = {name + ".lora_A": torch.ones(2, 4), name + ".lora_B": torch.ones(4, 2)}
|
|
peft = {name + ".lora_A.weight": torch.full((1, 4), 3.),
|
|
name + ".lora_B.weight": torch.full((4, 1), 7.)}
|
|
for extension in EXTENSIONS:
|
|
with self.subTest(extension=extension), tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
(root / "config.json").write_text(json.dumps({"config": {"lora_r": 2, "lora_alpha": 4}}))
|
|
(root / "adapter_config.json").write_text(json.dumps({"r": 1, "lora_alpha": 1}))
|
|
for stem, state in (("lora_weights", native), ("adapter_model", peft),
|
|
("speech_embedding", {"weight": torch.full((4, 4), 5.)}),
|
|
("llm_decoder", {"weight": torch.full((4, 4), 6.)})):
|
|
file = root / (stem + extension)
|
|
if extension == ".safetensors":
|
|
save_file(state, file)
|
|
else:
|
|
torch.save(state, file)
|
|
for path, is_native in ((root, True), (root / ("lora_weights" + extension), True),
|
|
(root / ("adapter_model" + extension), False)):
|
|
loaded, config, _, native_config = read_adapter(path, "cosyvoice3", root, None)
|
|
self.assertEqual(native_config is not None, is_native)
|
|
output = {}
|
|
normalize(loaded, config, bindings, output, 1, {}, native_scaling=is_native)
|
|
target = "llm/llm.model." + name + ".weight"
|
|
self.assertEqual(output[target + ".alpha"].item(), 4 if is_native else 1)
|
|
expected_a = native[name + ".lora_A"] if is_native else peft[name + ".lora_A.weight"]
|
|
self.assertTrue(torch.equal(output[target + ".lora_A"], expected_a))
|
|
sidecars("cosyvoice3", root, native_config, None, bindings, output, root)
|
|
self.assertEqual("llm/speech_embedding.weight" in output, is_native)
|
|
self.assertEqual("llm/llm_decoder.weight" in output, is_native)
|
|
(root / ("speech_embedding" + extension)).unlink()
|
|
for path in (root, root / ("lora_weights" + extension)):
|
|
_, _, _, native_config = read_adapter(path, "cosyvoice3", root, None)
|
|
with self.assertRaises(FileNotFoundError):
|
|
sidecars("cosyvoice3", root, native_config, None, bindings, {}, root)
|
|
|
|
def test_cosy_native_rank_mismatch_is_rejected(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
(root / "config.json").write_text(json.dumps({"config": {"lora_r": 3, "lora_alpha": 6}}))
|
|
name = "model.layers.0.q_proj"
|
|
save_file({name + ".lora_A": torch.ones(2, 4), name + ".lora_B": torch.ones(4, 2)}, root / "lora_weights.safetensors")
|
|
loaded, config, _, native_config = read_adapter(root, "cosyvoice3", root, None)
|
|
bindings = Bindings("cosyvoice3", {"llm/llm.model." + name + ".weight": (4, 4)})
|
|
with self.assertRaisesRegex(ValueError, "Configured rank differs"):
|
|
normalize(loaded, config, bindings, {}, 1, {}, native_scaling=native_config is not None)
|
|
|
|
def test_adapter_layout_is_independent_of_container(self):
|
|
for native in (False, True):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
if native:
|
|
(root / "config.json").write_text(json.dumps({"config": {"lora_r": 2, "lora_alpha": 4}}))
|
|
else:
|
|
(root / "adapter_config.json").write_text(json.dumps({"r": 2, "lora_alpha": 4}))
|
|
stem = "lora_weights" if native else "adapter_model"
|
|
suffix = "" if native else ".weight"
|
|
name = "model.layers.0.q_proj"
|
|
state = {name + ".lora_A" + suffix: torch.arange(8.).reshape(2, 4),
|
|
name + ".lora_B" + suffix: torch.arange(8.).reshape(4, 2)}
|
|
bindings = Bindings("cosyvoice3", {"llm/llm.model." + name + ".weight": (4, 4)})
|
|
expected = None
|
|
for extension in EXTENSIONS:
|
|
with self.subTest(native=native, extension=extension):
|
|
path = root / (stem + extension)
|
|
if extension == ".safetensors":
|
|
save_file(state, path)
|
|
else:
|
|
torch.save(state, path)
|
|
self.assertEqual(find_checkpoint(root, stem), path)
|
|
for supplied_path in (root, path):
|
|
loaded, config, _, native_config = read_adapter(supplied_path, "cosyvoice3", root, None)
|
|
self.assertEqual(native_config is not None, native)
|
|
output = {}
|
|
normalize(loaded, config, bindings, output, 1, {}, native_scaling=native)
|
|
if expected is None:
|
|
expected = output
|
|
self.assertEqual(output.keys(), expected.keys())
|
|
for key in output:
|
|
self.assertTrue(torch.equal(output[key], expected[key]), key)
|
|
path.unlink()
|
|
|
|
def test_checkpoint_discovery_does_not_hide_corruption(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
torch.save({"weight": torch.ones(2)}, root / "adapter_model.pt")
|
|
invalid = root / "adapter_model.safetensors"
|
|
invalid.write_bytes(b"invalid")
|
|
self.assertEqual(find_checkpoint(root, "adapter_model"), invalid)
|
|
with self.assertRaises(SafetensorError):
|
|
load_checkpoint(invalid)
|
|
self.assertIsNone(find_checkpoint(root, "missing", required=False))
|
|
with self.assertRaises(FileNotFoundError):
|
|
find_checkpoint(root, "missing")
|
|
|
|
def test_fish_training_layout_is_independent_of_container(self):
|
|
state = {"model.layers.0.attention.wqkv.lora_A": torch.ones(2, 4),
|
|
"model.layers.0.attention.wqkv.lora_B": torch.ones(8, 2)}
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
for extension in EXTENSIONS:
|
|
with self.subTest(extension=extension):
|
|
path = root / ("training" + extension)
|
|
if extension == ".safetensors":
|
|
save_file(state, path)
|
|
else:
|
|
torch.save({"state_dict": state, "epoch": 1}, path)
|
|
loaded, _, _, _ = read_adapter(path, "fish_audio", root, 4)
|
|
self.assertTrue(torch.equal(loaded["layers.0.attention.wqkv.lora_A"], state["model.layers.0.attention.wqkv.lora_A"]))
|
|
self.assertEqual(loaded["layers.0.attention.wqkv.alpha"].item(), 4)
|
|
|
|
def test_prepared_fish_scaling_is_not_overwritten(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
path = root / "prepared.safetensors"
|
|
state = {"layers.0.attention.wqkv.lora_A": torch.ones(2, 4),
|
|
"layers.0.attention.wqkv.lora_B": torch.ones(8, 2),
|
|
"layers.0.attention.wqkv.alpha": torch.tensor(4.)}
|
|
save_file(state, path)
|
|
loaded, _, _, _ = read_adapter(path, "fish_audio", root, 16)
|
|
self.assertEqual(loaded.keys(), state.keys())
|
|
for name in state:
|
|
self.assertTrue(torch.equal(loaded[name], state[name]))
|
|
|
|
def test_embedded_byte_array_is_read_in_bulk(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
file = root / "base.gguf"
|
|
writer = gguf.GGUFWriter(file, "test")
|
|
writer.add_array("audiocpp.embedded_files.names", ["vocab.txt"])
|
|
writer.add_key_value("audiocpp.embedded_files.offsets", [0, 3],
|
|
gguf.GGUFValueType.ARRAY, gguf.GGUFValueType.UINT64)
|
|
writer.add_array("audiocpp.embedded_files.data", b"ABC")
|
|
writer.add_tensor("projection.weight", np.ones((2, 32), dtype=np.float32))
|
|
writer.write_header_to_file()
|
|
writer.write_kv_data_to_file()
|
|
writer.write_tensors_to_file()
|
|
writer.close()
|
|
base = BaseGGUF(file)
|
|
self.assertEqual(base.reader.fields["audiocpp.embedded_files.data"].contents(), [65, 66, 67])
|
|
self.assertEqual(base.shapes["projection.weight"], (2, 32))
|
|
self.assertTrue(torch.equal(base.read_f32("projection.weight"), torch.ones(2, 32)))
|
|
destination = root / "sidecars"
|
|
base.extract_sidecars(destination)
|
|
self.assertEqual((destination / "vocab.txt").read_bytes(), b"ABC")
|
|
|
|
def test_all_family_bindings(self):
|
|
cases = [
|
|
("qwen3_tts", "model_weights/talker.model.layers.0.self_attn.q_proj.weight", "model.layers.0.self_attn.q_proj"),
|
|
("qwen3_asr", "model.language_model.layers.0.self_attn.q_proj.weight", "model.layers.0.self_attn.q_proj"),
|
|
("voxcpm2", "weights/base_lm.layers.0.self_attn.q_proj.weight", "base_lm.layers.0.self_attn.q_proj"),
|
|
("voxcpm1", "blk.0.attn_q.weight", "base_lm.layers.0.self_attn.q_proj"),
|
|
("fish_audio", "model_weights/fast_output.weight", "model.fast_output"),
|
|
("cosyvoice3", "llm/llm.model.model.layers.0.self_attn.q_proj.weight", "model.layers.0.self_attn.q_proj"),
|
|
("f5_tts", "transformer/ema_model.transformer.transformer_blocks.0.attn.to_q.weight", "transformer.transformer_blocks.0.attn.to_q"),
|
|
("outetts", "model_weights/model.layers.0.self_attn.q_proj.weight", "model.layers.0.self_attn.q_proj"),
|
|
("yue2", "model.layers.0.self_attn.q_proj.weight", "layers.0.self_attn.q_proj"),
|
|
("ace_step", "dit_turbo_weights/decoder.layers.0.q_proj.weight", "layers.0.q_proj"),
|
|
("vibevoice", "model.language_model.layers.0.self_attn.q_proj.weight", "layers.0.self_attn.q_proj"),
|
|
]
|
|
for family, name, external in cases:
|
|
with self.subTest(family=family):
|
|
shapes = {name: (4, 4)}
|
|
if family == "yue2":
|
|
shapes.update({name.replace("q_proj", part): (4, 4) for part in ("k_proj", "v_proj")})
|
|
bindings = Bindings(family, shapes)
|
|
self.assertEqual(bindings.resolve("base_model.model." + external)[0].name, name)
|
|
|
|
def test_fused_fish_slices_keep_all_rows(self):
|
|
shapes = {f"model_weights/layers.0.attention.{part}_proj.weight": (rows, 4)
|
|
for part, rows in (("q", 4), ("k", 2), ("v", 2))}
|
|
bindings = Bindings("fish_audio", shapes)
|
|
a = torch.arange(8, dtype=torch.float32).reshape(2, 4)
|
|
b = torch.arange(16, dtype=torch.float32).reshape(8, 2)
|
|
state = {"model.layers.0.attention.wqkv.lora_A": a,
|
|
"model.layers.0.attention.wqkv.lora_B": b,
|
|
"model.layers.0.attention.wqkv.alpha": torch.tensor(4.)}
|
|
output = {}
|
|
normalize(state, {}, bindings, output, 1, {})
|
|
offset = 0
|
|
for part, rows in (("q", 4), ("k", 2), ("v", 2)):
|
|
name = f"model_weights/layers.0.attention.{part}_proj.weight"
|
|
self.assertTrue(torch.equal(output[name + ".lora_A"], a))
|
|
self.assertTrue(torch.equal(output[name + ".lora_B"], b[offset:offset + rows]))
|
|
self.assertEqual(output[name + ".alpha"].item(), 4)
|
|
offset += rows
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
save_file(output, Path(directory) / "normalized.safetensors")
|
|
|
|
def test_fused_slices_reject_missing_and_excess_rows(self):
|
|
for family, prefix, fused in (
|
|
("fish_audio", "model_weights/layers.0.attention", "model.layers.0.attention.wqkv"),
|
|
("yue2", "model.layers.0.self_attn", "layers.0.self_attn.qkv_proj"),
|
|
):
|
|
shapes = {f"{prefix}.{part}_proj.weight": (rows, 4)
|
|
for part, rows in (("q", 4), ("k", 2), ("v", 2))}
|
|
bindings = Bindings(family, shapes)
|
|
for rows in (7, 9):
|
|
with self.subTest(family=family, rows=rows):
|
|
state = {fused + ".lora_A": torch.ones(2, 4),
|
|
fused + ".lora_B": torch.ones(rows, 2),
|
|
fused + ".alpha": torch.tensor(4.)}
|
|
output = {}
|
|
with self.assertRaisesRegex(ValueError, "output row"):
|
|
normalize(state, {}, bindings, output, 1, {})
|
|
self.assertEqual(output, {})
|
|
|
|
def test_mixed_lora_and_lokr_factors_are_rejected(self):
|
|
name = "model.layers.0.q_proj"
|
|
bindings = Bindings("qwen3_asr", {name + ".weight": (4, 4)})
|
|
for factor in ("lokr_w1", "lokr_w2", "lokr_w1_a", "lokr_w1_b", "lokr_w2_a", "lokr_w2_b"):
|
|
with self.subTest(factor=factor):
|
|
state = {name + ".lora_A": torch.ones(2, 4),
|
|
name + ".lora_B": torch.ones(4, 2),
|
|
name + ".alpha": torch.tensor(2.),
|
|
name + "." + factor: torch.ones(2, 2)}
|
|
output = {}
|
|
with self.assertRaisesRegex(ValueError, "Mixed adapter algorithms"):
|
|
normalize(state, {}, bindings, output, 1, {})
|
|
self.assertEqual(output, {})
|
|
|
|
def test_peft_patterns_and_rslora_are_preserved(self):
|
|
name = "model.layers.0.q_proj.weight"
|
|
bindings = Bindings("qwen3_asr", {name: (4, 4)})
|
|
source = "base_model.model.model.layers.0.q_proj"
|
|
state = {source + ".lora_A.weight": torch.ones(2, 4), source + ".lora_B.weight": torch.ones(4, 2)}
|
|
config = {"r": 8, "lora_alpha": 16, "rank_pattern": {"q_proj": 2}, "alpha_pattern": {"q_proj": 6}, "use_rslora": True}
|
|
rs = {"rank_pattern": {}, "alpha_pattern": {}}
|
|
output = {}
|
|
normalize(state, config, bindings, output, 1, rs)
|
|
self.assertEqual(list(rs["rank_pattern"].values()), [2])
|
|
self.assertEqual(list(rs["alpha_pattern"].values()), [6])
|
|
self.assertIn(name + ".lora_A.weight", output)
|
|
|
|
def test_duplicate_aliases_cannot_overwrite_rslora(self):
|
|
name = "model.layers.0.q_proj.weight"
|
|
bindings = Bindings("qwen3_asr", {name: (4, 4)})
|
|
source = "base_model.model.model.layers.0.q_proj"
|
|
state = {source + ".lora_A.weight": torch.ones(2, 4),
|
|
source + ".lora_B.weight": torch.ones(4, 2)}
|
|
config = {"r": 2, "lora_alpha": 2, "use_rslora": True}
|
|
output = {}
|
|
normalize(state, config, bindings, output, 1, {"rank_pattern": {}, "alpha_pattern": {}})
|
|
with self.assertRaisesRegex(ValueError, "Multiple updates"):
|
|
normalize(state, config, bindings, output, 1, {"rank_pattern": {}, "alpha_pattern": {}})
|
|
|
|
def test_native_does_not_inherit_neighbor_peft_config(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
file = root / "native.safetensors"
|
|
save_file({"layers.0.self_attn.q_proj.lora_A": torch.ones(2, 4),
|
|
"layers.0.self_attn.q_proj.lora_B": torch.ones(4, 2)}, file)
|
|
(root / "adapter_config.json").write_text(json.dumps({"r": 128, "lora_alpha": 256, "use_dora": True}))
|
|
state, config, _, _ = read_adapter(file, "yue2", root, None)
|
|
self.assertEqual(config, {})
|
|
output = {}
|
|
shapes = {f"model.layers.0.self_attn.{part}_proj.weight": (4, 4) for part in ("q", "k", "v")}
|
|
normalize(state, config, Bindings("yue2", shapes),
|
|
output, 1, {}, require_scaling=False)
|
|
self.assertEqual(output["model.layers.0.self_attn.q_proj.weight.alpha"].item(), 2)
|
|
|
|
def test_missing_peft_metadata_and_empty_checkpoint_are_rejected(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
file = root / "adapter.safetensors"
|
|
save_file({"base_model.model.proj.lora_A.weight": torch.ones(2, 4),
|
|
"base_model.model.proj.lora_B.weight": torch.ones(4, 2)}, file)
|
|
with self.assertRaisesRegex(ValueError, "adapter_config.json"):
|
|
read_adapter(file, "qwen3_asr", root, None)
|
|
save_file({}, file)
|
|
with self.assertRaisesRegex(ValueError, "empty"):
|
|
read_adapter(file, "qwen3_asr", root, None)
|
|
|
|
def test_component_ambiguity_is_explicit(self):
|
|
shapes = {f"{component}/decoder.proj.weight": (4, 4) for component in ("dit_base_weights", "dit_turbo_weights")}
|
|
with self.assertRaisesRegex(ValueError, "Ambiguous"):
|
|
Bindings("ace_step", shapes)
|
|
bindings = Bindings("ace_step", shapes, "dit_turbo_weights")
|
|
self.assertEqual(bindings.resolve("proj")[0].name, "dit_turbo_weights/decoder.proj.weight")
|
|
|
|
def test_original_sharded_inputs(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
save_file({"model.proj.weight": torch.ones(4, 4)}, root / "part.safetensors")
|
|
index = root / "model.safetensors.index.json"
|
|
index.write_text(json.dumps({"weight_map": {"model.proj.weight": "part.safetensors"}}))
|
|
base = BaseSources([f"model_weights={index}"], root)
|
|
self.assertEqual(base.shapes, {"model_weights/model.proj.weight": (4, 4)})
|
|
self.assertTrue(torch.equal(base.read_f32("model_weights/model.proj.weight"), torch.ones(4, 4)))
|
|
original = (root / "part.safetensors").read_bytes()
|
|
arguments = ["merge_lora.py", "--family", "qwen3_asr", "--input", str(index),
|
|
"--output", str(root / "part.safetensors"), "--overwrite", "--adapter", "unused"]
|
|
with patch.object(sys, "argv", arguments), self.assertRaises(SystemExit) as error:
|
|
main()
|
|
self.assertEqual(error.exception.code, 2)
|
|
self.assertEqual((root / "part.safetensors").read_bytes(), original)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|