Add training/merge_4b_v43.py

training/merge_4b_v43.py

@@ -0,0 +1,202 @@
+#!/usr/bin/env python3
+"""
+merge_4b_v43.py — Merge mlx_lm LoRA adapter into Qwen3-4B base weights.
+
+Does NOT use mlx_lm.fuse (broken — silently loses LoRA during GGUF conversion).
+Instead: loads safetensors directly, computes delta = (alpha/rank) * B @ A per layer,
+writes merged BF16 safetensors compatible with llama.cpp convert_hf_to_gguf.py.
+
+Usage:
+    python3 merge_4b_v43.py \
+        --base   Qwen/Qwen3-4B \
+        --adapter /tmp/4b_v43_adapter \
+        --out    /tmp/4b_v43_merged
+
+Requires: transformers, safetensors, torch (or mlx)
+"""
+import argparse
+import json
+import shutil
+import sys
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file, save_file
+from transformers import AutoTokenizer, AutoConfig
+
+
+def load_adapter_config(adapter_dir: Path) -> float:
+    """Returns the final scale factor from mlx_lm adapter_config.json.
+
+    mlx_lm stores lora_parameters.scale = alpha/rank pre-computed.
+    Falls back to 20.0 (default for r=8, alpha=160) if not found.
+    """
+    cfg_path = adapter_dir / "adapter_config.json"
+    if cfg_path.exists():
+        cfg = json.loads(cfg_path.read_text())
+        lora_params = cfg.get("lora_parameters", {})
+        if "scale" in lora_params:
+            return float(lora_params["scale"])
+    print("WARN: lora_parameters.scale not found — defaulting to 20.0")
+    return 20.0
+
+
+def find_safetensors(directory: Path) -> list[Path]:
+    files = sorted(directory.glob("*.safetensors"))
+    if not files:
+        print(f"ERROR: no .safetensors files in {directory}", file=sys.stderr)
+        sys.exit(1)
+    return files
+
+
+def load_all_safetensors(directory: Path) -> dict[str, torch.Tensor]:
+    tensors = {}
+    for f in find_safetensors(directory):
+        tensors.update(load_file(str(f), device="cpu"))
+    return tensors
+
+
+def merge(base_dir: Path, adapter_dir: Path, out_dir: Path) -> None:
+    scale = load_adapter_config(adapter_dir)
+    print(f"  LoRA scale: {scale:.4f}  (from adapter_config.json lora_parameters.scale)")
+
+    print("\nLoading base model weights...")
+    base = load_all_safetensors(base_dir)
+    print(f"  {len(base)} tensors loaded from base")
+
+    print("Loading adapter weights...")
+    adapter = load_all_safetensors(adapter_dir)
+    # mlx_lm adapter keys look like: model.layers.0.self_attn.q_proj.lora_a
+    lora_keys = [k for k in adapter if k.endswith(".lora_a")]
+    print(f"  {len(lora_keys)} LoRA A matrices found")
+
+    if not lora_keys:
+        print("ERROR: no lora_a keys found in adapter — wrong adapter format?", file=sys.stderr)
+        sys.exit(1)
+
+    merged = {k: v.clone() for k, v in base.items()}
+    applied = 0
+
+    for a_key in lora_keys:
+        b_key = a_key.replace(".lora_a", ".lora_b")
+        # Derive base weight key: strip .lora_a suffix, map to base weight name
+        # mlx_lm uses e.g. "model.layers.0.self_attn.q_proj.lora_a"
+        # base weight is "model.layers.0.self_attn.q_proj.weight"
+        base_key = a_key.replace(".lora_a", ".weight")
+
+        if b_key not in adapter:
+            print(f"  WARN: missing lora_b for {a_key} — skipping")
+            continue
+        if base_key not in merged:
+            print(f"  WARN: base key {base_key} not found — skipping")
+            continue
+
+        A = adapter[a_key].float()  # mlx_lm: (in_features, rank)
+        B = adapter[b_key].float()  # mlx_lm: (rank, out_features)
+        W = merged[base_key].float()
+
+        # mlx_lm stores weights as (in, out) — delta = scale * A @ B → (in, out)
+        delta = scale * (A @ B)
+
+        if delta.shape != W.shape:
+            # Fallback: try transposed orientation
+            if delta.T.shape == W.shape:
+                delta = delta.T
+            else:
+                print(f"  WARN: shape mismatch {delta.shape} vs {W.shape} for {base_key} — skipping")
+                continue
+
+        merged[base_key] = (W + delta).to(torch.bfloat16)
+        applied += 1
+
+    print(f"\n  Applied {applied}/{len(lora_keys)} LoRA deltas")
+    if applied == 0:
+        print("ERROR: zero deltas applied — check adapter key format", file=sys.stderr)
+        sys.exit(1)
+    if applied < len(lora_keys) * 0.9:
+        print(f"ERROR: only {applied}/{len(lora_keys)} deltas applied (<90%) — likely key mismatch", file=sys.stderr)
+        sys.exit(1)
+
+    # Cast all to bfloat16 for GGUF conversion
+    merged = {k: v.to(torch.bfloat16) if v.is_floating_point() else v for k, v in merged.items()}
+
+    print(f"\nSaving merged model to {out_dir}...")
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    # Split into shards of ~4GB each (llama.cpp prefers <5GB shards)
+    SHARD_BYTES = 4 * 1024 ** 3
+    shard, shard_bytes, shard_idx = {}, 0, 0
+    for k, v in merged.items():
+        size = v.numel() * v.element_size()
+        if shard and shard_bytes + size > SHARD_BYTES:
+            fname = out_dir / f"model-{shard_idx:05d}-of-XXXXX.safetensors"
+            save_file(shard, str(fname))
+            print(f"  Shard {shard_idx}: {fname.name} ({shard_bytes / 1e9:.2f} GB)")
+            shard, shard_bytes, shard_idx = {}, 0, shard_idx + 1
+        shard[k] = v
+        shard_bytes += size
+    if shard:
+        # Single-file model: use standard name
+        fname = out_dir / ("model.safetensors" if shard_idx == 0 else f"model-{shard_idx:05d}-of-XXXXX.safetensors")
+        save_file(shard, str(fname))
+        print(f"  Shard {shard_idx}: {fname.name} ({shard_bytes / 1e9:.2f} GB)")
+
+    # Rename shards with correct total count
+    shards = sorted(out_dir.glob("model-*-of-XXXXX.safetensors"))
+    n = len(shards)
+    if n > 0:
+        index = {"metadata": {"total_size": sum(v.numel() * v.element_size() for v in merged.values())}, "weight_map": {}}
+        for i, p in enumerate(shards):
+            new_name = f"model-{i:05d}-of-{n:05d}.safetensors"
+            p.rename(out_dir / new_name)
+            tensors = load_file(str(out_dir / new_name), device="cpu")
+            for k in tensors:
+                index["weight_map"][k] = new_name
+        (out_dir / "model.safetensors.index.json").write_text(json.dumps(index, indent=2))
+
+    # Copy tokenizer + config from base
+    for fname in ["config.json", "tokenizer.json", "tokenizer_config.json",
+                  "special_tokens_map.json", "generation_config.json", "chat_template.jinja"]:
+        src = base_dir / fname
+        if src.exists():
+            shutil.copy(src, out_dir / fname)
+
+    print(f"\n✅ Merge complete → {out_dir}")
+    print(f"   Applied {applied} LoRA deltas at scale {scale:.4f}")
+    print(f"\nNext: bash export_4b_v43_gguf.sh")
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--base",    type=Path, default=None,
+                   help="Path to HF base model dir (or HF hub id — will be downloaded)")
+    p.add_argument("--adapter", type=Path, default=Path("/tmp/4b_v43_adapter"))
+    p.add_argument("--out",     type=Path, default=Path("/tmp/4b_v43_merged"))
+    args = p.parse_args()
+
+    # Resolve base: try local cache first
+    if args.base is None:
+        from transformers.utils import cached_file
+        try:
+            # Trigger download/cache of config to locate cache dir
+            cfg_path = cached_file("Qwen/Qwen3-4B", "config.json")
+            args.base = Path(cfg_path).parent
+            print(f"Using cached base: {args.base}")
+        except Exception:
+            print("ERROR: --base not specified and Qwen/Qwen3-4B not in cache.", file=sys.stderr)
+            print("Run:   python3 -c \"from transformers import AutoModelForCausalLM; AutoModelForCausalLM.from_pretrained('Qwen/Qwen3-4B')\"", file=sys.stderr)
+            sys.exit(1)
+    elif not args.base.exists():
+        print(f"ERROR: --base path not found: {args.base}", file=sys.stderr)
+        sys.exit(1)
+
+    if not args.adapter.exists():
+        print(f"ERROR: --adapter path not found: {args.adapter}", file=sys.stderr)
+        print("Run training first: bash train_4b_v43_local.sh", file=sys.stderr)
+        sys.exit(1)
+
+    merge(args.base, args.adapter, args.out)
+
+
+if __name__ == "__main__":
+    main()