File size: 4,917 Bytes
fb68eb7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 | diff --git a/mlx_audio/tts/models/fish_drama3/convert.py b/mlx_audio/tts/models/fish_drama3/convert.py
index 79fa5df..ca91be9 100644
--- a/mlx_audio/tts/models/fish_drama3/convert.py
+++ b/mlx_audio/tts/models/fish_drama3/convert.py
@@ -82,13 +82,19 @@ class ShardWriter:
self.count = 0
self.quantization: dict[str, dict] = {}
- def add(self, name: str, value: mx.array, bits: Optional[int] = None) -> None:
+ def add(
+ self,
+ name: str,
+ value: mx.array,
+ bits: Optional[int] = None,
+ group_size: int = GROUP_SIZE,
+ ) -> None:
items = {name: value}
if bits:
path = name.removesuffix(".weight")
- weight, scales, biases = mx.quantize(value, GROUP_SIZE, bits)
+ weight, scales, biases = mx.quantize(value, group_size, bits)
items = {name: weight, f"{path}.scales": scales, f"{path}.biases": biases}
- self.quantization[path] = {"group_size": GROUP_SIZE, "bits": bits}
+ self.quantization[path] = {"group_size": group_size, "bits": bits}
mx.eval(list(items.values()))
self.pending.update(items)
self.pending_bytes += sum(v.nbytes for v in items.values())
@@ -111,6 +117,9 @@ def convert(
dense_bits: int = 0,
tower_bits: int = 0,
codec: bool = True,
+ expert_group_size: int = GROUP_SIZE,
+ head_bits: Optional[int] = None,
+ expert_bits_map: Optional[dict[int, int]] = None,
) -> None:
"""Writes the MLX checkpoint into dst, which must not exist or be empty:
the loader reads every safetensors file there, so a leftover shard from an
@@ -122,6 +131,8 @@ def convert(
num_experts = config["text_config"]["num_experts"]
band = slice(config["semantic_start_token_id"], config["semantic_end_token_id"] + 1)
eos = config["eos_token_id"]
+ if head_bits is None:
+ head_bits = tower_bits
writer = ShardWriter(dst)
experts: dict[int, dict[tuple, mx.array]] = defaultdict(dict)
@@ -138,7 +149,8 @@ def convert(
]
)
name = f"backbone.model.layers.{layer}.mlp.switch_mlp.{proj}_proj.weight"
- writer.add(name, stacked, expert_bits)
+ bits = (expert_bits_map or {}).get(layer, expert_bits)
+ writer.add(name, stacked, bits, expert_group_size)
index = json.loads((src / "model.safetensors.index.json").read_text())
for shard in sorted(set(index["weight_map"].values())):
@@ -166,7 +178,11 @@ def convert(
out = TOWER_PREFIXES[old] + name.removeprefix(old)
if TOWER_NORMS.search(out):
value = value.astype(mx.float32) + 1.0
- writer.add(out, value, tower_bits if TOWER.search(out) else None)
+ if match := TOWER.search(out):
+ bits = head_bits if match[1].startswith("codebook_heads") else tower_bits
+ else:
+ bits = None
+ writer.add(out, value, bits)
else:
raise ValueError(f"unexpected tensor {name}")
del tensors
@@ -181,7 +197,7 @@ def convert(
# also has its own entry.
config["quantization"] = {
"group_size": GROUP_SIZE,
- "bits": dense_bits or tower_bits or expert_bits,
+ "bits": dense_bits or tower_bits or head_bits or expert_bits,
**writer.quantization,
}
config["model_type"] = "fish_drama3"
@@ -244,8 +260,41 @@ def main() -> None:
# their top-1 agreement with bf16 dropped from about 88% to 40%, so quantize
# it to 8 at most.
parser.add_argument("--tower-bits", type=int, default=0, help="0 keeps bf16")
+ parser.add_argument(
+ "--head-bits",
+ type=int,
+ default=None,
+ help="bits for the tower's codebook heads; defaults to --tower-bits, 0 keeps bf16",
+ )
+ parser.add_argument(
+ "--expert-group-size",
+ type=int,
+ default=GROUP_SIZE,
+ choices=(32, 64, 128),
+ help="quantization group size for the experts",
+ )
+ parser.add_argument(
+ "--expert-bits-map",
+ type=Path,
+ default=None,
+ help='JSON file mapping layer index to expert bits, e.g. {"0": 4}; '
+ "layers not listed use --expert-bits",
+ )
args = parser.parse_args()
- convert(args.src, args.dst, args.expert_bits, args.dense_bits, args.tower_bits)
+ convert(
+ args.src,
+ args.dst,
+ args.expert_bits,
+ args.dense_bits,
+ args.tower_bits,
+ expert_group_size=args.expert_group_size,
+ head_bits=args.head_bits,
+ expert_bits_map={
+ int(k): v for k, v in json.loads(args.expert_bits_map.read_text()).items()
+ }
+ if args.expert_bits_map
+ else None,
+ )
if __name__ == "__main__":
|