| import json, os, shutil, torch |
| from safetensors.torch import load_file, save_file |
|
|
| SRC = os.environ.get("SRC", "/work/unsloth-nvfp4") |
| DST = os.environ.get("DST", "/work/unsloth-nvfp4-bf16head") |
|
|
| if not os.path.isdir(DST): |
| shutil.copytree(SRC, DST, copy_function=os.link) |
| for f in os.listdir(DST): |
| if f.endswith(".safetensors"): |
| os.remove(os.path.join(DST, f)) |
|
|
| st = load_file(os.path.join(SRC, "model.safetensors")) |
| w = st.pop("lm_head.weight").to(torch.float32) |
| scale = st.pop("lm_head.weight_scale").to(torch.float32) |
| st["lm_head.weight"] = (w * scale).to(torch.bfloat16) |
| print("lm_head ->", tuple(st["lm_head.weight"].shape), st["lm_head.weight"].dtype) |
|
|
| save_file(st, os.path.join(DST, "model.safetensors"), metadata={"format": "pt"}) |
|
|
| cfg = json.load(open(os.path.join(DST, "config.json"))) |
| q = cfg.get("quantization_config", {}) |
| ig = q.setdefault("ignore", []) |
| if not any("lm_head" in x for x in ig): |
| ig.append("lm_head") |
| json.dump(cfg, open(os.path.join(DST, "config.json"), "w"), indent=2) |
| print("saved + ignore:", q["ignore"]) |
|
|
