| |
| """Merge a LoRA checkpoint into the Gemma 4 E4B base model. |
| |
| Merges LoRA deltas directly into the base model's state dict (rather than |
| reconstructing a PeftModel via peft.PeftModel.from_pretrained), because the |
| adapter's auto-generated `target_modules` regex also matches the vision |
| tower's Gemma4ClippableLinear wrapper modules, which peft's standard tuner |
| injection does not support and raises on. Manual state-dict merging sidesteps |
| that entirely and is also simpler to verify. |
| |
| Verified (see checkpoint audit): every vision_tower.* lora_B tensor in the |
| adapter is exactly all-zero -- expected, since training never passed images |
| through the vision tower (finetune_vision_layers=False, text-only corpus), so |
| those slots never received gradient. This script only merges |
| language_model.* deltas; vision and the multi-modal projector are left as |
| pure base-model weights, byte-identical to upstream. |
| |
| Runs on GPU by default (there is normally ample headroom alongside a training |
| run, and merging is a handful of quick matmuls, not sustained compute -- it |
| does not meaningfully compete with a concurrent training job). Pass |
| --device cpu to run without touching the GPU at all. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import re |
| from pathlib import Path |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--adapter-dir", type=Path, required=True, help="Path to a LoRA checkpoint directory.") |
| parser.add_argument("--base-model", default=None, help="Override base model id (default: read from adapter_config.json).") |
| parser.add_argument("--output-dir", type=Path, required=True) |
| parser.add_argument("--device", default="cuda", choices=["cuda", "cpu"]) |
| args = parser.parse_args() |
|
|
| import torch |
| from safetensors import safe_open |
| from safetensors.torch import load_file as load_safetensors |
| from transformers import AutoModelForImageTextToText, AutoProcessor |
|
|
| adapter_config = json.loads((args.adapter_dir / "adapter_config.json").read_text()) |
| base_model_id = args.base_model or adapter_config["base_model_name_or_path"] |
| lora_r = adapter_config["r"] |
| lora_alpha = adapter_config["lora_alpha"] |
| scale = lora_alpha / lora_r |
| print(f"Base model: {base_model_id}") |
| print(f"Adapter: {args.adapter_dir}") |
| print(f"LoRA r={lora_r} alpha={lora_alpha} scale={scale} dropout={adapter_config['lora_dropout']}") |
| print(f"Device: {args.device}") |
|
|
| args.output_dir.mkdir(parents=True, exist_ok=True) |
|
|
| adapter_path = args.adapter_dir / "adapter_model.safetensors" |
| with safe_open(str(adapter_path), framework="pt") as f: |
| adapter_keys = list(f.keys()) |
| lora_a_keys = [k for k in adapter_keys if k.endswith(".lora_A.weight") and ".language_model." in k] |
| vision_keys = [k for k in adapter_keys if ".vision_tower." in k or ".multi_modal_projector." in k] |
| print(f"Adapter tensors total: {len(adapter_keys)}") |
| print(f" language_model lora_A/B pairs to merge: {len(lora_a_keys)}") |
| print(f" vision/projector tensors present but SKIPPED (verified all-zero, untrained): {len(vision_keys)}") |
|
|
| print(f"Loading base model in bf16 on {args.device} (a few minutes)...") |
| model = AutoModelForImageTextToText.from_pretrained( |
| base_model_id, |
| dtype=torch.bfloat16, |
| device_map=args.device if args.device == "cuda" else "cpu", |
| low_cpu_mem_usage=True, |
| ) |
| base_state = model.state_dict() |
|
|
| print("Loading adapter weight tensors...") |
| adapter_state = load_safetensors(str(adapter_path)) |
|
|
| |
| |
| prefix_re = re.compile(r"^base_model\.model\.") |
|
|
| merged = 0 |
| skipped_zero = 0 |
| for a_key in lora_a_keys: |
| b_key = a_key.replace(".lora_A.weight", ".lora_B.weight") |
| base_key = prefix_re.sub("", a_key).replace(".lora_A.weight", ".weight") |
| if base_key not in base_state: |
| raise KeyError(f"Merged base key {base_key!r} (from adapter key {a_key!r}) not found in base model state dict") |
|
|
| A = adapter_state[a_key].to(dtype=torch.float32, device=args.device) |
| B = adapter_state[b_key].to(dtype=torch.float32, device=args.device) |
| if B.abs().max().item() == 0.0: |
| skipped_zero += 1 |
| continue |
|
|
| delta = (scale * (B @ A)).to(dtype=base_state[base_key].dtype) |
| with torch.no_grad(): |
| base_state[base_key].add_(delta.to(base_state[base_key].device)) |
| merged += 1 |
|
|
| print(f"Merged {merged} language_model layers, skipped {skipped_zero} all-zero (never trained) pairs.") |
| if merged == 0: |
| raise RuntimeError("No LoRA deltas were merged -- refusing to save an unmodified base model.") |
|
|
| print(f"Saving merged model to {args.output_dir} ...") |
| model.save_pretrained(str(args.output_dir), safe_serialization=True) |
|
|
| print("Saving processor/tokenizer + chat template...") |
| processor = AutoProcessor.from_pretrained(base_model_id) |
| processor.save_pretrained(str(args.output_dir)) |
|
|
| metadata = { |
| "schema_version": "tenaos_lora_merge_v1", |
| "base_model": base_model_id, |
| "adapter_dir": str(args.adapter_dir), |
| "adapter_config_step": adapter_config.get("_source_step"), |
| "lora_r": lora_r, |
| "lora_alpha": lora_alpha, |
| "merged_dtype": "bfloat16", |
| "language_layers_merged": merged, |
| "vision_layers_skipped_all_zero": skipped_zero, |
| "note": "Vision tower and multi-modal projector are unmodified base-model weights; only the text decoder was fine-tuned.", |
| } |
| (args.output_dir / "merge_metadata.json").write_text(json.dumps(metadata, indent=2) + "\n") |
| print("Done.") |
| print(json.dumps(metadata, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
