""" Training script with qv_variant choices extended for Phase 7 normed Q-only variants. """ import os import time import math import pickle import argparse from contextlib import nullcontext import numpy as np import torch from torch.nn.parallel import DistributedDataParallel as DDP from torch.distributed import init_process_group, destroy_process_group from model import GPTConfig, GPT out_dir = 'out-shakespeare-char' eval_interval = 250 log_interval = 10 eval_iters = 200 eval_only = False always_save_checkpoint = False wandb_log = False wandb_project = 'shakespeare-char' wandb_run_name = 'mini-gpt' dataset = 'shakespeare_char' gradient_accumulation_steps = 1 batch_size = 64 block_size = 256 n_layer = 6 n_head = 6 n_embd = 384 dropout = 0.2 bias = False learning_rate = 1e-3 max_iters = 5000 lr_decay_iters = 5000 min_lr = 1e-4 beta1 = 0.9 beta2 = 0.99 warmup_iters = 100 device = 'cuda' dtype = 'bfloat16' if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else 'float16' compile = False config_keys = [k for k, v in globals().items() if not k.startswith('_') and isinstance(v, (int, float, bool, str))] parser = argparse.ArgumentParser(description='Train a mini-GPT') parser.add_argument('--out_dir', type=str, default=out_dir) parser.add_argument('--eval_interval', type=int, default=eval_interval) parser.add_argument('--log_interval', type=int, default=log_interval) parser.add_argument('--eval_iters', type=int, default=eval_iters) parser.add_argument('--eval_only', action='store_true') parser.add_argument('--always_save_checkpoint', action='store_true') parser.add_argument('--wandb_log', action='store_true') parser.add_argument('--wandb_project', type=str, default=wandb_project) parser.add_argument('--wandb_run_name', type=str, default=wandb_run_name) parser.add_argument('--dataset', type=str, default=dataset) parser.add_argument('--gradient_accumulation_steps', type=int, default=gradient_accumulation_steps) parser.add_argument('--batch_size', type=int, default=batch_size) parser.add_argument('--block_size', type=int, default=block_size) parser.add_argument('--n_layer', type=int, default=n_layer) parser.add_argument('--n_head', type=int, default=n_head) parser.add_argument('--n_embd', type=int, default=n_embd) parser.add_argument('--dropout', type=float, default=dropout) parser.add_argument('--bias', action='store_true') parser.add_argument('--learning_rate', type=float, default=learning_rate) parser.add_argument('--max_iters', type=int, default=max_iters) parser.add_argument('--lr_decay_iters', type=int, default=lr_decay_iters) parser.add_argument('--min_lr', type=float, default=min_lr) parser.add_argument('--beta1', type=float, default=beta1) parser.add_argument('--beta2', type=float, default=beta2) parser.add_argument('--warmup_iters', type=int, default=warmup_iters) parser.add_argument('--device', type=str, default=device) parser.add_argument('--dtype', type=str, default=dtype) parser.add_argument('--compile', action='store_true') parser.add_argument('--seed', type=int, default=1337, help='random seed') parser.add_argument('--qv_variant', type=str, default='none', choices=[ # baseline 'none', # normalization variants 'vnorm', 'qvnorm', # static gates 'static_gate', 'static_gate_prehead', # post-attn normalization 'post_rmsnorm_y', # Q-conditioned gates 'dynamic', 'dynamic_swiglu', 'dynamic_qconditioned_mlp128', 'dynamic_qconditioned_normed', 'dynamic_qconditioned_mlp128_normed', 'dynamic_qconditioned_mlp192', 'dynamic_qconditioned_fullwidth', 'dynamic_qconditioned_fullwidth_headspecific', 'dynamic_q_headshared_elementwise', # X-conditioned gates 'dynamic_xconditioned_g1', 'dynamic_xconditioned_fullwidth_headspecific', 'dynamic_xconditioned_bottleneck', 'dynamic_x_g1_headspecific_elementwise', 'dynamic_x_g1_headspecific_headwise', 'dynamic_x_g1_headshared_elementwise', # Random / ablation gates 'dynamic_random_gate', 'dynamic_ones_gate', 'dynamic_random_normal', 'dynamic_bernoulli_gate', # Zero-parameter dot product gates 'dynamic_dot_scalar', 'dynamic_dot_elementwise', # A-only gate 'dynamic_a_conditioned', # Q+A dual-signal gates 'dynamic_qa_conditioned', 'dynamic_qa_conditioned_headspecific', 'dynamic_qa_conditioned_mlp128', 'dynamic_qa_headshared_elementwise', 'dynamic_qa_bilinear_diag', 'dynamic_qa_conditioned_normed', 'dynamic_qa_conditioned_lowrank16', # legacy 'dynamic_postconcat_matched', ], help='QV experiment variant') args = parser.parse_args() for k in config_keys: if hasattr(args, k): globals()[k] = getattr(args, k) config = {k: globals()[k] for k in config_keys} ddp = int(os.environ.get('RANK', -1)) != -1 if ddp: init_process_group(backend='nccl') ddp_rank = int(os.environ['RANK']) ddp_local_rank = int(os.environ['LOCAL_RANK']) ddp_world_size = int(os.environ['WORLD_SIZE']) device = f'cuda:{ddp_local_rank}' torch.cuda.set_device(device) master_process = ddp_rank == 0 seed_offset = ddp_rank gradient_accumulation_steps //= ddp_world_size else: master_process = True seed_offset = 0 ddp_world_size = 1 tokens_per_iter = gradient_accumulation_steps * ddp_world_size * batch_size * block_size print(f"tokens per iteration will be: {tokens_per_iter:,}") if master_process: os.makedirs(out_dir, exist_ok=True) torch.manual_seed(args.seed + seed_offset) torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True device_type = 'cuda' if 'cuda' in device else 'cpu' ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[dtype] ctx = nullcontext() if device_type == 'cpu' else torch.amp.autocast(device_type=device_type, dtype=ptdtype) data_dir = os.path.join('data', dataset) # OPEN MEMMAPS ONCE, REUSE FOREVER train_data = np.memmap(os.path.join(data_dir, 'train.bin'), dtype=np.uint16, mode='r') val_data = np.memmap(os.path.join(data_dir, 'val.bin'), dtype=np.uint16, mode='r') def get_batch(split): data = train_data if split == 'train' else val_data ix = torch.randint(len(data) - block_size, (batch_size,)) x = torch.stack([torch.from_numpy((data[i:i+block_size]).astype(np.int64)) for i in ix]) y = torch.stack([torch.from_numpy((data[i+1:i+1+block_size]).astype(np.int64)) for i in ix]) if device_type == 'cuda': x, y = x.pin_memory().to(device, non_blocking=True), y.pin_memory().to(device, non_blocking=True) else: x, y = x.to(device), y.to(device) return x, y iter_num = 0 best_val_loss = 1e9 meta_path = os.path.join(data_dir, 'meta.pkl') meta_vocab_size = None if os.path.exists(meta_path): with open(meta_path, 'rb') as f: meta = pickle.load(f) meta_vocab_size = meta['vocab_size'] print(f"found vocab_size = {meta_vocab_size} (inside {meta_path})") model_args = dict(n_layer=n_layer, n_head=n_head, n_embd=n_embd, block_size=block_size, bias=bias, vocab_size=None, dropout=dropout) if meta_vocab_size is None: print("defaulting to vocab_size of GPT-2 to 50304 (50257 rounded up for efficiency)") model_args['vocab_size'] = meta_vocab_size if meta_vocab_size is not None else 50304 gptconf = GPTConfig(**model_args) gptconf.qv_variant = args.qv_variant model = GPT(gptconf) model.to(device) scaler = torch.cuda.amp.GradScaler(enabled=(dtype == 'float16')) optimizer = model.configure_optimizers(1e-1, learning_rate, (beta1, beta2), device_type) if compile: print("compiling the model... (takes a ~minute)") model = torch.compile(model) @torch.no_grad() def estimate_loss(): out = {} model.eval() for split in ['train', 'val']: losses = torch.zeros(eval_iters) for k in range(eval_iters): X, Y = get_batch(split) with ctx: logits, loss = model(X, Y) losses[k] = loss.item() out[split] = losses.mean() model.train() return out def get_lr(it): if it < warmup_iters: return learning_rate * (it + 1) / (warmup_iters + 1) if it > lr_decay_iters: return min_lr decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters) coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio)) return min_lr + coeff * (learning_rate - min_lr) if wandb_log and master_process: import wandb wandb.init(project=wandb_project, name=wandb_run_name, config=config) X, Y = get_batch('train') t0 = time.time() local_iter_num = 0 raw_model = model.module if hasattr(model, 'module') else model running_mfu = -1.0 while True: lr = get_lr(iter_num) for param_group in optimizer.param_groups: param_group['lr'] = lr if iter_num % eval_interval == 0 and master_process: losses = estimate_loss() print(f"step {iter_num}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}") if losses['val'] < best_val_loss or always_save_checkpoint: best_val_loss = losses['val'] if iter_num > 0: checkpoint = { 'model': raw_model.state_dict(), 'optimizer': optimizer.state_dict(), 'model_args': model_args, 'iter_num': iter_num, 'best_val_loss': best_val_loss, 'config': config, } print(f"saving checkpoint to {out_dir}") torch.save(checkpoint, os.path.join(out_dir, 'ckpt.pt')) if iter_num == 0 and eval_only: break for micro_step in range(gradient_accumulation_steps): with ctx: logits, loss = model(X, Y) loss = loss / gradient_accumulation_steps X, Y = get_batch('train') scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) scaler.step(optimizer) scaler.update() optimizer.zero_grad(set_to_none=True) t1 = time.time() dt = t1 - t0 t0 = t1 if iter_num % log_interval == 0 and master_process: lossf = loss.item() * gradient_accumulation_steps if local_iter_num >= 5: mfu = raw_model.estimate_mfu(batch_size * gradient_accumulation_steps, dt) running_mfu = mfu if running_mfu == -1.0 else 0.9 * running_mfu + 0.1 * mfu print(f"iter {iter_num}: loss {lossf:.4f}, time {dt*1000:.2f}ms, mfu {running_mfu*100:.2f}%") iter_num += 1 local_iter_num += 1 if iter_num > max_iters: break if ddp: destroy_process_group()