File size: 2,381 Bytes
8a5ffa8 | 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 | """Claim-5 matched-learning-rate control.
The paper's own model-selection rule (pick the learning rate with the best
zero-shot Humanities accuracy) happened to select very different step sizes for
the two arms in our runs -- 1e-5 for all-attention (Q/K/V) but 3e-4 for
value-matrix-only (V). The two fine-tuned models therefore sit at very
different distances from the base model, which confounds any comparison of how
much in-context ability each one retained.
This control removes that confound: it trains BOTH arms at the SAME learning
rate (1e-4, the middle of the paper's grid, where both arms score well on the
zero-shot objective) with the same seeds, and re-measures 0-shot and 7-shot
accuracy on all four categories.
Emits results/claim5_matched.json.
"""
import argparse, json, os, torch
import mmlu_data as MD
import claim5_mmlu_qwen_finetune as C
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--out", default="/work/results")
ap.add_argument("--lr", type=float, default=1e-4)
ap.add_argument("--seeds", type=int, default=2)
ap.add_argument("--bsz", type=int, default=32)
a = ap.parse_args()
os.makedirs(a.out, exist_ok=True)
C.log("building MMLU data ...")
data = MD.build()
tok = C.load_tok()
out = {"model": C.MODEL_ID, "matched_lr": a.lr, "lora_r": C.LORA_R,
"epochs": C.EPOCHS, "n_eval_per_cat": MD.N_EVAL_PER_CAT,
"n_shot": MD.N_SHOT, "arms": {}}
path = os.path.join(a.out, "claim5_matched.json")
def dump():
json.dump(out, open(path, "w"), indent=2)
C.log("=== BASE MODEL (re-measured for a self-contained control) ===")
base = C.load_base()
out["base"], out["base_hits"] = C.eval_all(base, tok, data, batch_size=a.bsz)
del base; torch.cuda.empty_cache()
dump()
for tgt in ("qkv", "v"):
out["arms"][tgt] = {"lr": a.lr, "seeds": []}
for s in range(a.seeds):
C.log(f"=== MATCHED target={tgt} lr={a.lr:g} seed={s} ===")
m, meta = C.train_lora(tok, data, tgt, a.lr, seed=s)
acc, hits = C.eval_all(m, tok, data, batch_size=a.bsz)
out["arms"][tgt]["seeds"].append(
{"seed": s, "acc": acc, "hits": hits, **meta})
del m; torch.cuda.empty_cache()
dump()
C.log("MATCHED CONTROL DONE")
dump()
if __name__ == "__main__":
main()
|