Correct the model card (pooling is at the LABELS suffix's last token); add code/classifier from project commit 4ced7554f
3d7c238 verified | """Train the heads classifier on the export. Run on the pod as ``python -m classifier.train``. | |
| Adapted from the read-only reference ~/yt_live_analysis/training/train.py (arm H only), with the spec's | |
| fixes: steps of a fixed number of posts split into token-budget micro-batches with per-step loss | |
| normalisation; nulls masked; a cumulative ``--max-hours`` kept in state.json (a resume cannot extend it); | |
| ``--stop-after`` to test resume without changing the schedule; and a forced step on the longest rows to | |
| measure peak memory. | |
| Writes OUT/config.json, OUT/train_log.jsonl, OUT/last/ (resume point), OUT/best/ (lowest dev mean field | |
| log-loss), OUT/predictions/dev.jsonl and OUT/metrics_dev.json for the best checkpoint. Exit code 3: stopped by | |
| --max-hours or --stop-after at a saved resume point; rerun with --resume. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import math | |
| from pathlib import Path | |
| import platform | |
| import random | |
| import time | |
| from .calibrate import nll, probs as calibrated_probs | |
| from .data import inference_batches, load_split, micro_batches, step_counts, steps, subset, write_jsonl | |
| from .formats import FORMAT, encode | |
| from .taxonomy import Taxonomy, argmax | |
| CORE = ('data', 'model', 'subset', 'epochs', 'max_steps', 'posts_per_step', 'token_budget', 'max_len', 'seed', | |
| 'lr_lora', 'lr_heads', 'warmup', 'weight_decay', 'lora_r', 'lora_alpha', 'lora_dropout', 'head_dropout', | |
| 'dtype', 'kernels') | |
| def parse(argv=None): | |
| p = argparse.ArgumentParser(description=__doc__.splitlines()[0]) | |
| p.add_argument('--data', required=True, help='Export directory (train/dev.jsonl, taxonomy.json, manifest.json)') | |
| p.add_argument('--model', required=True, help='Local checkpoint directory') | |
| p.add_argument('--out', required=True) | |
| p.add_argument('--subset', type=int, default=0, help='Train on this many rows in seeded order (0 = all)') | |
| p.add_argument('--epochs', type=float, default=2.0) | |
| p.add_argument('--max-steps', type=int, default=0, help='Cap the schedule at this many optimiser steps') | |
| p.add_argument('--posts-per-step', type=int, default=32) | |
| p.add_argument('--token-budget', type=int, default=16384, help='Padded tokens per training micro-batch') | |
| p.add_argument('--eval-token-budget', type=int, default=65536) | |
| p.add_argument('--max-len', type=int, default=2048) | |
| p.add_argument('--lr-lora', type=float, default=2e-4) | |
| p.add_argument('--lr-heads', type=float, default=1e-3) | |
| p.add_argument('--warmup', type=float, default=0.03) | |
| p.add_argument('--weight-decay', type=float, default=0.01) | |
| p.add_argument('--lora-r', type=int, default=32) | |
| p.add_argument('--lora-alpha', type=int, default=32) | |
| p.add_argument('--lora-dropout', type=float, default=0.05) | |
| p.add_argument('--head-dropout', type=float, default=0.1) | |
| p.add_argument('--eval-every', type=int, default=200, help='Optimiser steps between dev evaluations and saves') | |
| p.add_argument('--dev-limit', type=int, default=0, help='Evaluate on the first N dev rows in seeded order') | |
| p.add_argument('--log-every', type=int, default=10) | |
| p.add_argument('--max-hours', type=float, default=0, help='Cumulative training hours across resumes') | |
| p.add_argument('--stop-after', type=int, default=0, help='Stop this process after N more steps (resume test)') | |
| p.add_argument('--force-longest', type=int, default=0, help='First run one forward/backward on the N longest rows') | |
| p.add_argument('--seed', type=int, default=20261017) | |
| p.add_argument('--resume', action='store_true', | |
| help='Continue from last/ if it exists (settings must match); otherwise start fresh') | |
| p.add_argument('--kernels', default='fast', choices=['fast', 'torch'], | |
| help='fast: require fla + causal-conv1d (asserted); torch: the PyTorch reference path') | |
| p.add_argument('--device', default='cuda') | |
| p.add_argument('--dtype', default='bfloat16', choices=['bfloat16', 'float32']) | |
| return p.parse_args(argv) | |
| def versions() -> dict: | |
| import peft | |
| import torch | |
| import transformers | |
| out = dict(python=platform.python_version(), torch=torch.__version__, transformers=transformers.__version__, | |
| peft=peft.__version__, cuda=torch.version.cuda) | |
| for pkg in ('triton', 'fla', 'causal_conv1d', 'tilelang'): | |
| try: | |
| out[pkg] = getattr(__import__(pkg), '__version__', 'unknown') | |
| except Exception: | |
| out[pkg] = None | |
| commit = Path(__file__).with_name('.source_commit') | |
| out['source_commit'] = commit.read_text().strip() if commit.exists() else None | |
| return out | |
| class Log: | |
| def __init__(self, path: Path): | |
| self.fh = open(path, 'a', encoding='utf-8') | |
| def __call__(self, **row): | |
| row = dict(time=round(time.time(), 1), **row) | |
| line = json.dumps(row, ensure_ascii=False) | |
| self.fh.write(line + '\n') | |
| self.fh.flush() | |
| print(line, flush=True) | |
| def prediction_rows(rows, logits_by_index, infos, split=None) -> list[dict]: | |
| return [{'id': r['id'], 'split': split or r.get('split'), 'n_tokens': infos[i]['n_tokens'], 'truncated': infos[i]['cut'], | |
| 'logits': logits_by_index[i]} for i, r in enumerate(rows)] | |
| def dev_report(preds: list[dict], rows: list[dict], tax: Taxonomy) -> dict: | |
| pairs = [(p, r['targets']) for p, r in zip(preds, rows)] | |
| per_field = {f: nll(pairs, tax, f, 1.0) for f in tax.names} | |
| scored = [v for v in per_field.values() if v is not None] | |
| acc = {} | |
| for f in tax.names: | |
| gold = [(p, t[f]) for p, t in pairs if t.get(f) is not None] | |
| acc[f] = sum(1 for p, g in gold if argmax(calibrated_probs(p, tax))[f] == g) / len(gold) if gold else None | |
| return {'selection_loss': sum(scored) / len(scored), 'log_loss': per_field, 'accuracy': acc, 'rows': len(rows)} | |
| def main(argv=None): | |
| args = parse(argv) | |
| import torch | |
| from .model import build, pad, predict, save_checkpoint, step_loss, load_tokenizer, tokenizer_fn | |
| out = Path(args.out) | |
| out.mkdir(parents=True, exist_ok=True) | |
| cfg_path = out / 'config.json' | |
| cfg = {k: v for k, v in vars(args).items() if k not in ('resume', 'device', 'stop_after', 'force_longest')} | |
| if cfg_path.exists(): | |
| if not args.resume: | |
| raise SystemExit('Refusing to start a new run over %s; use --resume or a new --out' % out) | |
| old = json.loads(cfg_path.read_text()) | |
| changed = sorted(k for k in CORE if old.get(k) != cfg.get(k)) | |
| if changed: | |
| raise SystemExit('--resume with different settings: %s' % changed) | |
| if old.get('max_hours') != cfg.get('max_hours'): | |
| raise SystemExit('--resume with a different --max-hours; the cap is cumulative') | |
| manifest = json.loads((Path(args.data) / 'manifest.json').read_text()) | |
| cfg.update(versions=versions(), format=FORMAT, data_version=manifest.get('version'), | |
| data_files_sha256=manifest.get('files_sha256') or manifest.get('sha256')) | |
| cfg_path.write_text(json.dumps(cfg, indent=1)) | |
| log = Log(out / 'train_log.jsonl') | |
| random.seed(args.seed) | |
| torch.manual_seed(args.seed) | |
| device = torch.device(args.device) | |
| tax = Taxonomy.load(args.data) | |
| (out / 'taxonomy.json').write_text(json.dumps(tax.to_json(), ensure_ascii=False, indent=1)) | |
| tok = load_tokenizer(args.model) | |
| tokenize = tokenizer_fn(tok) | |
| train_rows = subset(load_split(args.data, 'train'), args.subset, args.seed) | |
| dev_rows = subset(load_split(args.data, 'dev'), args.dev_limit, args.seed) | |
| t0 = time.time() | |
| enc_train = [encode(r, tokenize, args.max_len) for r in train_rows] | |
| enc_dev = [encode(r, tokenize, args.max_len) for r in dev_rows] | |
| ids = [e[0] for e in enc_train] | |
| lengths = [len(x) for x in ids] | |
| labels = [tax.encode(r['targets']) for r in train_rows] | |
| dev_ids = [e[0] for e in enc_dev] | |
| dev_infos = [e[1] for e in enc_dev] | |
| dev_batches = inference_batches([len(x) for x in dev_ids], args.eval_token_budget) | |
| cut = sum(1 for e in enc_train if e[1]['cut']) | |
| log(event='data', train=len(train_rows), dev=len(dev_rows), tokens=sum(lengths), max_len=max(lengths), | |
| mean_len=round(sum(lengths) / len(lengths), 1), truncated_rows=cut, | |
| dev_tokens=sum(len(x) for x in dev_ids), seconds=round(time.time() - t0, 1)) | |
| last_dir, best_dir = out / 'last', out / 'best' | |
| resume = args.resume and (last_dir / 'state.json').exists() | |
| model = build(cfg, tax.sizes(), device, train=True) | |
| from .check_env import fast_path, torch_path | |
| if args.kernels == 'torch': | |
| torch_path().__enter__() # for the whole process: the pre-declared PyTorch-path fallback | |
| cfg['fast_path'] = fast_path(model) | |
| cfg_path.write_text(json.dumps(cfg, indent=1)) | |
| if args.kernels == 'fast' and not cfg['fast_path']['ok']: | |
| raise SystemExit('Fast kernels are not active: %s' % json.dumps(cfg['fast_path'])) | |
| if resume: | |
| from peft import set_peft_model_state_dict | |
| from safetensors.torch import load_file | |
| set_peft_model_state_dict(model.backbone, load_file(str(last_dir / 'adapter' / 'adapter_model.safetensors'))) | |
| model.heads.load_state_dict(torch.load(last_dir / 'heads.pt', map_location='cpu')) | |
| lora = [p for n, p in model.named_parameters() if p.requires_grad and 'lora_' in n] | |
| heads = [p for n, p in model.named_parameters() if p.requires_grad and 'lora_' not in n] | |
| opt = torch.optim.AdamW([dict(params=lora, lr=args.lr_lora), dict(params=heads, lr=args.lr_heads)], | |
| weight_decay=args.weight_decay) | |
| steps_per_epoch = math.ceil(len(train_rows) / args.posts_per_step) | |
| total = max(1, int(steps_per_epoch * args.epochs)) | |
| if args.max_steps: | |
| total = min(total, args.max_steps) | |
| from transformers import get_cosine_schedule_with_warmup | |
| sched = get_cosine_schedule_with_warmup(opt, int(args.warmup * total), total) | |
| state = dict(step=0, epoch=0, step_in_epoch=0, best=None, best_step=None, hours_used=0.0) | |
| if resume: | |
| state = json.loads((last_dir / 'state.json').read_text()) | |
| trainer = torch.load(last_dir / 'trainer.pt', map_location='cpu', weights_only=False) | |
| opt.load_state_dict(trainer['opt']) | |
| sched.load_state_dict(trainer['sched']) | |
| log(event='resume', **state) | |
| log(event='start', trainable=sum(p.numel() for p in lora + heads), steps=total, steps_per_epoch=steps_per_epoch) | |
| amp = dict(device_type=device.type, dtype=torch.bfloat16, enabled=device.type == 'cuda') | |
| process_start = time.time() | |
| hours_before = state['hours_used'] | |
| def hours_now(): | |
| return hours_before + (time.time() - process_start) / 3600 | |
| def save_last(): | |
| state['hours_used'] = round(hours_now(), 4) | |
| save_checkpoint(model, last_dir, dict(state), dict(opt=opt.state_dict(), sched=sched.state_dict())) | |
| def evaluate(): | |
| t = time.time() | |
| logits = predict(model, dev_ids, tok.pad_token_id, dev_batches, device, tax.names) | |
| preds = prediction_rows(dev_rows, logits, dev_infos, 'dev') | |
| rep = dev_report(preds, dev_rows, tax) | |
| seconds = time.time() - t | |
| log(event='eval', step=state['step'], selection_loss=round(rep['selection_loss'], 5), | |
| log_loss={f: (round(v, 4) if v is not None else None) for f, v in rep['log_loss'].items()}, | |
| accuracy={f: (round(v, 4) if v is not None else None) for f, v in rep['accuracy'].items()}, | |
| seconds=round(seconds, 1), posts_per_second=round(len(dev_rows) / seconds, 1)) | |
| if state['best'] is None or rep['selection_loss'] < state['best']: | |
| state['best'], state['best_step'] = rep['selection_loss'], state['step'] | |
| save_checkpoint(model, best_dir, dict(state)) | |
| write_jsonl(out / 'predictions' / 'dev.jsonl', preds) | |
| (out / 'metrics_dev.json').write_text(json.dumps(dict(step=state['step'], **rep), indent=1)) | |
| if args.force_longest and not resume: | |
| # The worst-case micro-batch for peak memory: the N longest rows of the whole train split (not the | |
| # subset), which fill the token budget at max_len. One forward/backward, no optimiser step. | |
| pool = sorted(load_split(args.data, 'train'), key=lambda r: -sum(len(x['text']) for x in r['segments'])) | |
| cand = [(encode(r, tokenize, args.max_len)[0], tax.encode(r['targets'])) for r in pool[:4 * args.force_longest]] | |
| cand = sorted(cand, key=lambda c: -len(c[0]))[:args.force_longest] | |
| f_ids, f_labels = [c[0] for c in cand], [c[1] for c in cand] | |
| f_len = [len(x) for x in f_ids] | |
| parts = micro_batches(list(range(len(cand))), f_len, args.token_budget) | |
| for mb in parts: | |
| x, m, last = pad([f_ids[i] for i in mb], tok.pad_token_id, device) | |
| y = torch.tensor([f_labels[i] for i in mb], device=device) | |
| with torch.autocast(**amp): | |
| loss, _ = step_loss(model(x, m, last), y, step_counts(f_labels)) | |
| loss.backward() | |
| opt.zero_grad(set_to_none=True) | |
| mem = round(torch.cuda.max_memory_allocated() / 2 ** 30, 2) if device.type == 'cuda' else None | |
| log(event='forced_longest', rows=len(cand), lengths=f_len, micro_batches=len(parts), | |
| padded_tokens=[max(f_len[i] for i in mb) * len(mb) for mb in parts], max_memory_gb=mem) | |
| if device.type == 'cuda': | |
| torch.cuda.reset_peak_memory_stats() | |
| model.train() | |
| try: | |
| stop = run_steps(args, state, total, model, opt, sched, lora, heads, ids, lengths, labels, tax, device, amp, | |
| tok, log, hours_now, evaluate, save_last) | |
| except KeyboardInterrupt: # timeout sends SIGINT: keep a resume point at the last completed step | |
| log(event='interrupted', step=state['step']) | |
| save_last() | |
| return 3 | |
| if stop == 4: | |
| return 4 | |
| save_last() | |
| if state['best'] is None: | |
| evaluate() | |
| log(event='done' if stop is None else 'stopped', **state) | |
| return stop or 0 | |
| def run_steps(args, state, total, model, opt, sched, lora, heads, ids, lengths, labels, tax, device, amp, tok, log, | |
| hours_now, evaluate, save_last): | |
| """The optimiser loop. Returns None (finished), 3 (cap or planned stop, resume point saved by the caller) | |
| or 4 (non-finite loss or gradient: nothing is saved, so last/ keeps the previous good resume point).""" | |
| import torch | |
| from .model import pad, step_loss | |
| window = dict(loss=0.0, tokens=0, padded=0, t=time.time(), field_sums=[0.0] * len(tax.names), | |
| field_counts=[0] * len(tax.names), n=0) | |
| stop, done_here = None, 0 | |
| while state['step'] < total and stop is None: | |
| epoch_steps = steps(lengths, args.posts_per_step, args.seed * 1000 + state['epoch']) | |
| while state['step_in_epoch'] < len(epoch_steps) and state['step'] < total: | |
| if args.max_hours and hours_now() >= args.max_hours: | |
| stop = 3 | |
| break | |
| step = epoch_steps[state['step_in_epoch']] | |
| step_labels = [labels[i] for i in step] | |
| counts = step_counts(step_labels) | |
| step_total = 0.0 | |
| for mb in micro_batches(step, lengths, args.token_budget): | |
| x, m, last = pad([ids[i] for i in mb], tok.pad_token_id, device) | |
| y = torch.tensor([labels[i] for i in mb], device=device) | |
| with torch.autocast(**amp): | |
| loss, per_field = step_loss(model(x, m, last), y, counts) | |
| loss.backward() | |
| step_total += float(loss.detach()) | |
| window['field_sums'] = [a + b for a, b in zip(window['field_sums'], per_field)] | |
| window['tokens'] += int(m.sum()) | |
| window['padded'] += int(m.numel()) | |
| window['field_counts'] = [a + sum(1 for r in step_labels if r[f] >= 0) | |
| for f, a in enumerate(window['field_counts'])] | |
| norm = float(torch.nn.utils.clip_grad_norm_(lora + heads, 1.0)) | |
| if not (math.isfinite(step_total) and math.isfinite(norm)): | |
| log(event='non_finite', step=state['step'] + 1, loss=step_total, grad_norm=norm) | |
| opt.zero_grad(set_to_none=True) | |
| return 4 | |
| opt.step() | |
| sched.step() | |
| opt.zero_grad(set_to_none=True) | |
| state['step'] += 1 | |
| state['step_in_epoch'] += 1 | |
| done_here += 1 | |
| window['loss'] += step_total | |
| window['n'] += 1 | |
| if state['step'] % args.log_every == 0 or state['step'] == total: | |
| dt = time.time() - window['t'] | |
| row = dict(event='train', step=state['step'], epoch=state['epoch'], loss=round(window['loss'] / window['n'], 4), | |
| fields={f: round(s / c, 4) if c else None for f, s, c in | |
| zip(tax.names, window['field_sums'], window['field_counts'])}, | |
| lr=sched.get_last_lr()[0], tokens_per_second=round(window['tokens'] / dt), | |
| padded_tokens_per_second=round(window['padded'] / dt), hours=round(hours_now(), 3)) | |
| if device.type == 'cuda': | |
| row['max_memory_gb'] = round(torch.cuda.max_memory_allocated() / 2 ** 30, 1) | |
| log(**row) | |
| window = dict(loss=0.0, tokens=0, padded=0, t=time.time(), field_sums=[0.0] * len(tax.names), | |
| field_counts=[0] * len(tax.names), n=0) | |
| if state['step'] % args.eval_every == 0 or state['step'] == total: | |
| evaluate() | |
| save_last() | |
| if args.stop_after and done_here >= args.stop_after: | |
| stop = 3 | |
| break | |
| if state['step_in_epoch'] >= len(epoch_steps): | |
| state['epoch'] += 1 | |
| state['step_in_epoch'] = 0 | |
| return stop | |
| if __name__ == '__main__': | |
| raise SystemExit(main()) | |