File size: 1,333 Bytes
67ebcef | 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 | """Extrapolate along the update an RL step made: theta = theta_new + beta * (theta_new - theta_old).
The delta is computed in fp32 per tensor; non-weight files are copied from <new>.
usage: extrap.py <new_dir> <old_dir> <beta> <out_dir>"""
import glob, os, shutil, sys, torch
from safetensors.torch import load_file, save_file
new_dir, old_dir, beta, out = sys.argv[1], sys.argv[2], float(sys.argv[3]), sys.argv[4]
def load_all(d):
sd = {}
for f in sorted(glob.glob(os.path.join(d, "*.safetensors"))):
sd.update(load_file(f))
return sd
new, old = load_all(new_dir), load_all(old_dir)
res, rel = {}, []
for k, v in new.items():
o = old.get(k)
if o is None or o.shape != v.shape:
res[k] = v
continue
d = v.float() - o.float()
res[k] = (v.float() + beta * d).to(v.dtype)
if d.abs().sum() > 0:
rel.append(d.norm().item() / o.float().norm().item())
os.makedirs(out, exist_ok=True)
save_file(res, os.path.join(out, "model.safetensors"), metadata={"format": "pt"})
for f in os.listdir(new_dir):
p = os.path.join(new_dir, f)
if os.path.isfile(p) and not f.endswith(".safetensors") and f != "model.safetensors.index.json":
shutil.copy(p, out)
print(f"{len(rel)} tensors changed, mean relative update {sum(rel) / max(1, len(rel)):.2e}, beta {beta} -> {out}")
|