File size: 3,008 Bytes
f0975ca | 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 | """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 # noqa: E402
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()
|