diff options
Diffstat (limited to 'src/loss_reweight.py')
| -rw-r--r-- | src/loss_reweight.py | 197 |
1 files changed, 197 insertions, 0 deletions
diff --git a/src/loss_reweight.py b/src/loss_reweight.py new file mode 100644 index 0000000..3274807 --- /dev/null +++ b/src/loss_reweight.py @@ -0,0 +1,197 @@ +""" +Loss-reweighting ablation (GPT-5.6-Terra's design, adapted to faithful J-lens). + +Tests whether increasing a token's EFFECTIVE frequency/importance reduces its +faithful J-lens norm — without the confounds of the old random-insertion +ablation (which shifted positions, destroyed n-grams, and used an unmatched +control run). + +Design per seed (identical init + identical minibatch order for all three): + q-upweight: cross-entropy terms whose target is 'q' are weighted x2. + control: ordinary loss. + ctrl_random: same-total-loss control: weight x2 on the SAME NUMBER of + randomly chosen non-'q' target positions (deterministic per + batch index, so all models share the same control positions). + +If increased effective frequency causally reduces the faithful J-lens norm of +'q', the q-upweight model must show a lower norm than BOTH controls. + +Steps: + python3 src/loss_reweight.py --step train --mode q --seed 0 [--max_iters 3000] + python3 src/loss_reweight.py --step jlens --mode q --seed 0 [--layers 2,3,4] + python3 src/loss_reweight.py --step summary +""" +import sys, os, argparse, subprocess, pickle +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from typing import Any +import numpy as np +import torch +import torch.nn.functional as F + +DATA_DIR = 'data/shakespeare_char' +OUT_ROOT = 'out-loss-reweight' +JLENS_OUT = 'outputs/loss_reweight' +TARGET = 'q' +WEIGHT = 2.0 +MODEL_ARGS: dict[str, Any] = dict(n_layer=6, n_head=6, n_embd=384, block_size=128, + bias=False, dropout=0.2) +LOSS_MODES = ('q', 'control', 'ctrl_random') + + +def _batch(data, blk, bs, g, device): + ix = torch.randint(len(data) - blk, (bs,), generator=g) + x = torch.stack([torch.from_numpy(data[i:i+blk].astype(np.int64)) for i in ix]) + y = torch.stack([torch.from_numpy(data[i+1:i+1+blk].astype(np.int64)) for i in ix]) + return x.to(device), y.to(device) + + +def _weighted_loss(logits, y, mode, q_id, batch_k, V, device): + """Per-token weighted CE. Returns scalar loss.""" + logp = F.log_softmax(logits.view(-1, V), dim=-1) + nll = -logp.gather(1, y.view(-1, 1)).squeeze(1) # (B*T,) + w = torch.ones_like(nll) + if mode == 'q': + w[y.view(-1) == q_id] = WEIGHT + elif mode == 'ctrl_random': + g = torch.Generator(device=device).manual_seed(1000 + batch_k) + n_q = int((y == q_id).sum().item()) + flat = torch.arange(y.numel(), device=device) + non_q = flat[y.view(-1) != q_id] + if len(non_q) > 0 and n_q > 0: + pick = non_q[torch.randperm(len(non_q), generator=g)[:min(n_q, len(non_q))]] + w[pick] = WEIGHT + return (nll * w).mean() + + +def train(mode, seed, max_iters, batch_size): + sys.path.insert(0, '.') + from model import GPT, GPTConfig + torch.manual_seed(seed) + np.random.seed(seed) + device = 'cuda' + train_data = np.memmap(f'{DATA_DIR}/train.bin', dtype=np.uint16, mode='r') + val_data = np.memmap(f'{DATA_DIR}/val.bin', dtype=np.uint16, mode='r') + with open(f'{DATA_DIR}/meta.pkl', 'rb') as f: + meta = pickle.load(f) + q_id = meta['stoi'][TARGET] + args = dict( + n_layer=int(MODEL_ARGS['n_layer']), n_head=int(MODEL_ARGS['n_head']), + n_embd=int(MODEL_ARGS['n_embd']), block_size=int(MODEL_ARGS['block_size']), + bias=bool(MODEL_ARGS['bias']), dropout=float(MODEL_ARGS['dropout']), + vocab_size=int(meta['vocab_size']), + ) + model = GPT(GPTConfig(**args)).to(device) + print(f"[{mode}] seed {seed}: params={sum(p.numel() for p in model.parameters())/1e6:.2f}M") + + opt = model.configure_optimizers(weight_decay=0.1, learning_rate=1e-3, + betas=(0.9, 0.99), device_type='cuda') + bs = batch_size + blk = args['block_size'] + V = args['vocab_size'] + # identical minibatch order for every model: per-seed generator, fixed start + g = torch.Generator(device=device).manual_seed(20260731 + seed) + gval = torch.Generator(device=device).manual_seed(777 + seed) + + def get_batch(split): + d = train_data if split == 'train' else val_data + gg = g if split == 'train' else gval + return _batch(d, blk, bs, gg, device) + + best_val = 1e9 + out_dir = f'{OUT_ROOT}/{mode}/seed{seed}' + os.makedirs(out_dir, exist_ok=True) + for it in range(max_iters): + if it % 500 == 0: + model.eval() + lv = [] + for _ in range(50): + X, Y = get_batch('val') + with torch.no_grad(): + logits = model(X)[0] + lv.append(F.cross_entropy(logits.view(-1, V), Y.view(-1)).item()) + v = np.mean(lv) + model.train() + if v < best_val: + best_val = v + torch.save({'model': model.state_dict(), 'model_args': args, + 'best_val_loss': best_val}, f'{out_dir}/ckpt.pt') + if it % 1000 == 0: + print(f" iter {it}: val={v:.4f}") + X, Y = get_batch('train') + logits = model(X)[0] + loss = _weighted_loss(logits, Y, mode, q_id, it, V, device) + loss.backward() + opt.step() + opt.zero_grad(set_to_none=True) + print(f"[{mode}] seed {seed} done. best_val={best_val:.4f}") + + +def jlens(mode, seed, layers, n_prompts): + out_dir = f'{JLENS_OUT}/{mode}/seed{seed}' + cmd = ["python3", "-u", "src/jlens_v3.py", + "--checkpoint", f'{OUT_ROOT}/{mode}/seed{seed}/ckpt.pt', + "--data_dir", DATA_DIR, + "--n_prompts", str(n_prompts), + "--layers", layers, + "--chunk", "16", + "--output_dir", out_dir] + print("running:", " ".join(cmd)) + r = subprocess.run(cmd, cwd='/workspace/code') + assert r.returncode == 0, "jlens_v3 failed" + + +def summary(layers): + with open(f'{DATA_DIR}/meta.pkl', 'rb') as f: + meta = pickle.load(f) + q_id = meta['stoi'][TARGET] + seeds = sorted(set( + d.split('seed')[1] for mode in LOSS_MODES + for d in os.listdir(f'{JLENS_OUT}/{mode}') + if d.startswith('seed'))) + print(f"\n{'='*78}") + print(f"LOSS-REWEIGHTING: faithful J-lens norm of '{TARGET}' " + f"({WEIGHT}x CE) vs controls") + print(f"{'='*78}") + print(f"{'seed':<5}{'layer':<6}" + "".join(f"{m:>14}" for m in LOSS_MODES)) + for s in seeds: + for l in map(int, layers.split(',')): + row = [s, str(l)] + for m in LOSS_MODES: + d = torch.load(f'{JLENS_OUT}/{m}/seed{s}/layer{l}.pt', + map_location='cpu') + row.append(f"{d['faithful_norms'][q_id]:.4f}") + print(f"{row[0]:<5}{row[1]:<6}" + "".join(f"{v:>14}" for v in row[2:])) + # mean over middle layers per mode + mids = [l for l in map(int, layers.split(','))] + means = {} + for m in LOSS_MODES: + vals = [] + for l in mids: + d = torch.load(f'{JLENS_OUT}/{m}/seed{s}/layer{l}.pt', + map_location='cpu') + vals.append(d['faithful_norms'][q_id]) + means[m] = np.mean(vals) + print(f" -> mean over layers: q={means['q']:.4f} " + f"control={means['control']:.4f} ctrl_random={means['ctrl_random']:.4f}") + print(f" -> q/control = {means['q']/max(means['control'],1e-9):.3f} " + f"q/ctrl_random = {means['q']/max(means['ctrl_random'],1e-9):.3f}") + + +if __name__ == '__main__': + ap = argparse.ArgumentParser() + ap.add_argument('--step', required=True, choices=['train', 'jlens', 'summary']) + ap.add_argument('--mode', choices=LOSS_MODES) + ap.add_argument('--seed', type=int, default=0) + ap.add_argument('--max_iters', type=int, default=3000) + ap.add_argument('--batch_size', type=int, default=16) + ap.add_argument('--layers', default='2,3,4') + ap.add_argument('--n_prompts', type=int, default=10) + a = ap.parse_args() + if a.step == 'train': + assert a.mode, "need --mode" + train(a.mode, a.seed, a.max_iters, a.batch_size) + elif a.step == 'jlens': + assert a.mode, "need --mode" + jlens(a.mode, a.seed, a.layers, a.n_prompts) + else: + summary(a.layers) |