File size: 1,155 Bytes
5be2256 | 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 | import json, os, shutil
def rewrite(path, fix):
# break any hardlink first: write to temp then replace
tmp = path + ".new"
cfg = json.load(open(path))
fix(cfg)
with open(tmp, "w") as f:
json.dump(cfg, f, indent=2)
os.replace(tmp, path)
print("fixed", path)
def fix_bf16head(cfg):
q = cfg["quantization_config"]
t = q["config_groups"]["group_0"]["targets"]
q["config_groups"]["group_0"]["targets"] = [x for x in t if "lm_head" not in x]
ig = q.setdefault("ignore", [])
if "lm_head" not in ig:
ig.append("lm_head")
def fix_original(cfg):
q = cfg["quantization_config"]
q["ignore"] = [x for x in q.get("ignore", []) if x != "lm_head"]
rewrite("/p/unsloth-nvfp4-bf16head/config.json", fix_bf16head)
rewrite("/p/unsloth-nvfp4/config.json", fix_original)
# show proof
for p in ["/p/unsloth-nvfp4-bf16head/config.json", "/p/unsloth-nvfp4/config.json"]:
q = json.load(open(p))["quantization_config"]
print(p.split("/")[2], "| group_0 lm_head target:", any("lm_head" in x for x in q["config_groups"]["group_0"]["targets"]), "| ignore lm_head:", "lm_head" in q.get("ignore", []))
|