| """Interpolate a fine-tuned Dory checkpoint with the base model (WiSE-FT style), no training. |
| |
| .venv/bin/python -I scripts/merge_dory.py --base models/dory-4b --tuned models/rA/final --t 0.5 --out models/rA-t050 |
| |
| w = (1 - t) * base + t * tuned for every tensor (in fp32), then matrices are projected back onto the |
| unit hypersphere along the hidden dimension, as DoryForCausalLM.normalize_matrices_ does after each |
| optimizer step. For unit rows this is a normalized lerp, i.e. close to slerp. Scalars (alphas, scales, |
| logit_scale) are interpolated linearly. |
| """ |
|
|
| import argparse |
| import json |
| import shutil |
| import sys |
| from pathlib import Path |
|
|
| import torch |
| from safetensors import safe_open |
| from safetensors.torch import save_file |
| from torch import nn |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) |
| from dory.modeling import DoryConfig, DoryForCausalLM, justnorm |
|
|
| BASE_FILES = ("config.json", "generation_config.json", "tokenizer.json", "tokenizer_config.json", |
| "special_tokens_map.json", "chat_template.jinja") |
|
|
|
|
| def norm_dims(cfg): |
| """Parameter name -> dim to normalize along, mirroring normalize_matrices_.""" |
| with torch.device("meta"): |
| model = DoryForCausalLM(cfg) |
| dims = {"backbone.embeddings.weight": 1, "lm_head.weight": 1} |
| for name, mod in model.backbone.named_modules(): |
| if isinstance(mod, nn.Linear): |
| dims[f"backbone.{name}.weight"] = 1 if mod.in_features == cfg.hidden_size else 0 |
| return dims |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--base", required=True) |
| ap.add_argument("--tuned", required=True) |
| ap.add_argument("--t", type=float, required=True, help="0 = base, 1 = tuned") |
| ap.add_argument("--out", required=True) |
| args = ap.parse_args() |
| base, tuned, out = Path(args.base), Path(args.tuned), Path(args.out) |
|
|
| cfg = DoryConfig.from_json(tuned / "config.json") |
| dims = norm_dims(cfg) |
| merged, n_norm = {}, 0 |
| with safe_open(base / "model.safetensors", "pt") as fb, safe_open(tuned / "model.safetensors", "pt") as ft: |
| assert set(fb.keys()) == set(ft.keys()), "checkpoints have different tensors" |
| assert set(dims) <= set(fb.keys()), sorted(set(dims) - set(fb.keys()))[:5] |
| for k in fb.keys(): |
| a, b = fb.get_tensor(k).float(), ft.get_tensor(k).float() |
| w = torch.lerp(a, b, args.t) |
| if k in dims: |
| w = justnorm(w, dim=dims[k]) |
| n_norm += 1 |
| merged[k] = w.to(torch.bfloat16).contiguous() |
|
|
| out.mkdir(parents=True, exist_ok=True) |
| save_file(merged, out / "model.safetensors", metadata={"format": "pt"}) |
| for f in BASE_FILES: |
| if (tuned / f).exists(): |
| shutil.copy(tuned / f, out / f) |
| (out / "merge_args.json").write_text(json.dumps({**vars(args), "normalized_matrices": n_norm}, indent=1)) |
| print(f"merged {len(merged)} tensors ({n_norm} renormalized), t={args.t} -> {out}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
