| """Resumable mixed causal / whole-word-span training, without external weights.""" |
| import argparse |
| from dataclasses import asdict |
| import json |
| import math |
| from pathlib import Path |
| import shutil |
| import time |
|
|
| import numpy as np |
| import torch |
| from torch.nn import functional as F |
|
|
| from diffuthink.v2.model import Config, Denoiser, PAD, BOS, MASK |
| from diffuthink.v2.train import amp, atomic_save |
| from diffuthink.v2.hybrid import causal_evaluate |
| from .data import load |
|
|
|
|
| def span_batch(clean, ranges, rng=None, partial=False): |
| positions = torch.arange(clean.shape[1], device=clean.device)[None] |
| selected = (positions >= ranges[:, :1]) & (positions < ranges[:, 1:]) |
| if partial: |
| |
| keep = torch.rand(clean.shape, device=clean.device, generator=rng) < .5 |
| keep.scatter_(1, ranges[:, :1], True) |
| selected &= keep |
| if not selected.any(-1).all() or (selected & (clean.eq(PAD) | clean.eq(BOS))).any(): |
| raise ValueError('Invalid supervised spans') |
| noise = selected.sum(-1).float() / (clean.ne(PAD) & clean.ne(BOS)).sum(-1) |
| return clean.masked_fill(selected, MASK), selected, noise |
|
|
|
|
| @torch.inference_mode() |
| def span_evaluate(model, pair, limit=400, precision='bf16', batch_size=16): |
| model.eval() |
| device = next(model.parameters()).device |
| total = correct = exact = examples = 0 |
| loss = 0. |
| array, spans = pair |
| for start in range(0, min(limit, len(array)), batch_size): |
| stop = min(start+batch_size, limit, len(array)) |
| clean = torch.tensor(np.array(array[start:stop]), dtype=torch.long, device=device) |
| clean = clean[:, :int(clean.ne(PAD).sum(-1).max())] |
| ranges = torch.tensor(np.array(spans[start:stop, 0]), dtype=torch.long, device=device) |
| noisy, selected, noise = span_batch(clean, ranges) |
| with amp(device.type, precision): |
| logits = model(noisy, noise, selected) |
| target = clean[selected] |
| matched = logits.argmax(-1).eq(target) |
| loss += F.cross_entropy(logits.float(), target, reduction='sum').item() |
| correct += int(matched.sum()) |
| offset = 0 |
| for count in selected.sum(-1).tolist(): |
| exact += int(matched[offset:offset+count].all()) |
| offset += count |
| total += target.numel() |
| examples += len(clean) |
| return {'nll': loss/total, 'token_accuracy': correct/total, 'parallel_exact_match': exact/examples, |
| 'targets': total, 'examples': examples, |
| 'meaning': 'Exact reference reconstruction; valid alternative wording can score zero'} |
|
|
|
|
| def train(args): |
| if min(args.steps, args.batch_size, args.eval_every, args.eval_samples, args.warmup) < 1 or args.steps <= args.warmup: |
| raise ValueError('Positive counts and steps > warmup required') |
| if args.span_every < 0: |
| raise ValueError('span-every must be nonnegative') |
| torch.set_num_threads(4) |
| torch.manual_seed(args.seed) |
| arrays, manifest = load(args.data) |
| source_manifest = manifest['source_manifest'] |
| out = Path(args.output) |
| if (out/'last.pt').exists() and not args.resume: |
| raise ValueError('Existing run; use --resume') |
| out.mkdir(parents=True, exist_ok=True) |
| config = Config(vocab_size=source_manifest['vocab_size'], width=args.width, heads=8, |
| layers=args.layers, max_length=source_manifest['max_length']) |
| lineage = None |
| if args.source: |
| lineage = json.loads((Path(args.source)/'training_info.json').read_text()) |
| if lineage.get('pretrained_weights') is not False: |
| raise ValueError('Only our own from-scratch lineage is accepted') |
| if lineage['manifest']['tokenizer_sha256'] != manifest['tokenizer_sha256']: |
| raise ValueError('Cannot change vocabulary while reusing weights') |
| model = Denoiser.from_pretrained(args.source, args.device) |
| config = model.config |
| if config.max_length != source_manifest['max_length']: |
| raise ValueError('Context mismatch') |
| else: |
| model = Denoiser(config).to(args.device) |
| optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=(.9, .95), weight_decay=.1) |
| sampling = torch.Generator().manual_seed(args.seed) |
| noise_rng = torch.Generator(device=args.device).manual_seed(args.seed+1) |
| settings = {k: v for k, v in vars(args).items() if k not in ('resume', 'stop_after')} |
| step = seen = supervised = 0 |
| elapsed, best = 0., float('inf') |
| if args.resume: |
| state = torch.load(out/'last.pt', map_location='cpu', weights_only=True) |
| restored_settings = dict(state['settings']) |
| restored_settings.setdefault('span_every', 2) |
| restored_settings.setdefault('selection_mode', 'mixed') |
| if restored_settings != settings or state['manifest'] != manifest: |
| raise ValueError('Recovery settings/data mismatch') |
| model.load_state_dict(state['model']) |
| optimizer.load_state_dict(state['optimizer']) |
| sampling.set_state(state['sampling_rng']) |
| noise_rng.set_state(state['noise_rng']) |
| step, seen, supervised, elapsed, best = (state[k] for k in ('step','seen_tokens','supervised_tokens','seconds','best')) |
| (out/'settings.json').write_text(json.dumps(settings, indent=2)) |
| (out/'config.json').write_text(json.dumps(asdict(config), indent=2)) |
| |
| train_rows, spans = arrays['train'] |
| lengths = np.count_nonzero(train_rows, axis=1) |
| order = np.argsort(lengths, kind='stable') |
| started = time.perf_counter() |
| end = min(args.steps, args.stop_after) if args.stop_after else args.steps |
| print(json.dumps({'parameters': sum(p.numel() for p in model.parameters()), 'from_step': step, |
| 'to_step': end, 'source': args.source, 'device': args.device}), flush=True) |
| while step < end: |
| step += 1 |
| model.train() |
| anchor = int(torch.randint(len(order), (1,), generator=sampling)) |
| offsets = torch.randint(min(2048, len(order)), (args.batch_size,), generator=sampling).numpy() |
| indices = order[(anchor+offsets) % len(order)] |
| clean = torch.tensor(np.array(train_rows[indices, :int(lengths[indices].max())]), dtype=torch.long, device=args.device) |
| lr = args.lr * min(1., step/args.warmup) * (.1+.9*.5*(1+math.cos(math.pi*max(0, step-args.warmup)/(args.steps-args.warmup)))) |
| for group in optimizer.param_groups: |
| group['lr'] = lr |
| optimizer.zero_grad(set_to_none=True) |
| editing = args.span_every > 0 and step % args.span_every == 0 |
| with amp(args.device, args.precision): |
| if editing: |
| choices = torch.randint(8, (args.batch_size,), generator=sampling).numpy() |
| ranges = torch.tensor(np.array(spans[indices, choices]), dtype=torch.long, device=args.device) |
| noisy, selected, noise = span_batch(clean, ranges, noise_rng, partial=(step//args.span_every) % 2 == 0) |
| logits = model(noisy, noise, selected) |
| targets = clean[selected] |
| else: |
| targets = clean[:, 1:] |
| selected = targets.ne(PAD) |
| logits = model(clean[:, :-1], torch.zeros(len(clean), device=args.device), selected, causal=True, right_padded=True) |
| targets = targets[selected] |
| loss = F.cross_entropy(logits.float(), targets) |
| if not torch.isfinite(loss): |
| raise RuntimeError('Nonfinite loss') |
| loss.backward() |
| norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.) |
| if not torch.isfinite(norm): |
| raise RuntimeError('Nonfinite gradient') |
| optimizer.step() |
| seen += int(clean.ne(PAD).sum()) |
| supervised += targets.numel() |
| if step == 1 or step % 100 in (0, 1): |
| row = {'step': step, 'objective': 'span' if editing else 'causal', 'loss': loss.item(), |
| 'lr': lr, 'seen_tokens': seen, 'supervised_tokens': supervised, |
| 'seconds': elapsed+time.perf_counter()-started} |
| with (out/'metrics.jsonl').open('a') as stream: |
| stream.write(json.dumps(row)+'\n') |
| print(json.dumps(row), flush=True) |
| if step % args.eval_every == 0 or step == end: |
| metric = span_evaluate(model, arrays['validation'], args.eval_samples, args.precision) |
| causal = causal_evaluate(model, arrays['validation'][0], args.eval_samples, precision=args.precision) |
| |
| score = causal['nll'] if args.selection_mode == 'causal' else metric['nll'] + .25*causal['nll'] |
| row = {'step': step, 'validation_span': metric, 'validation_causal': causal, 'selection_score': score} |
| if score < best: |
| best = score |
| model.save_pretrained(out/'best') |
| shutil.copy2(Path(args.data)/'tokenizer.json', out/'best'/'tokenizer.json') |
| info = {'step': step, 'parameters': sum(p.numel() for p in model.parameters()), |
| 'seen_tokens': seen, 'supervised_tokens': supervised, 'pretrained_weights': False, |
| 'manifest': source_manifest, 'editing_manifest': manifest, 'previous_phase': lineage, |
| 'continuation_mode': 'autoregressive', |
| 'objective': f'Span update every {args.span_every} steps (0 = causal only), alternating full and partially revealed word spans; all other updates causal', |
| 'selection_mode': args.selection_mode, |
| 'selection_score': score, 'validation': row} |
| (out/'best'/'training_info.json').write_text(json.dumps(info, indent=2)) |
| atomic_save({'model': model.state_dict(), 'optimizer': optimizer.state_dict(), |
| 'sampling_rng': sampling.get_state(), 'noise_rng': noise_rng.get_state(), |
| 'step': step, 'seen_tokens': seen, 'supervised_tokens': supervised, |
| 'seconds': elapsed+time.perf_counter()-started, 'best': best, |
| 'settings': settings, 'manifest': manifest}, out/'last.pt') |
| with (out/'metrics.jsonl').open('a') as stream: |
| stream.write(json.dumps(row)+'\n') |
| print(json.dumps(row), flush=True) |
| (out/'status.json').write_text(json.dumps({'completed': step == args.steps, 'step': step, |
| 'planned_steps': args.steps, 'seconds': elapsed+time.perf_counter()-started, |
| 'best_selection_score': best}, indent=2)) |
| print(f'Saved {out}; completed={step == args.steps}', flush=True) |
|
|
|
|
| def parser(): |
| p = argparse.ArgumentParser() |
| p.add_argument('--data', default='data/processed/storypatch-v2') |
| p.add_argument('--output', default='runs/storypatch-58m') |
| p.add_argument('--source', help='Optional continuation of our own smaller scratch-trained model') |
| p.add_argument('--device', choices=['cuda', 'cpu'], default='cuda') |
| p.add_argument('--precision', choices=['bf16', 'fp32'], default='bf16') |
| for key, value in [('width',512), ('layers',12), ('steps',12000), ('batch-size',24), |
| ('warmup',400), ('eval-every',500), ('eval-samples',256), ('seed',48)]: |
| p.add_argument('--'+key, type=int, default=value) |
| p.add_argument('--lr', type=float, default=.0003) |
| p.add_argument('--resume', action='store_true') |
| p.add_argument('--span-every', type=int, default=2, help='0 = causal only; otherwise one span update every N steps') |
| p.add_argument('--selection-mode', choices=['mixed','causal'], default='mixed') |
| p.add_argument('--stop-after', type=int, help='Save recoverable state early without changing the LR schedule') |
| return p |
|
|
|
|
| if __name__ == '__main__': |
| train(parser().parse_args()) |
|
|
