| import math |
| import os |
| import time |
| import argparse |
|
|
| import torch |
| import torch.nn.functional as F |
|
|
| from ar.model.config import ARConfig, small_config, tiny_config |
| from ar.model.model import ARModel |
| from ar.data.dataset import BinDataset |
| from ar.train.train_config import TrainConfig |
|
|
|
|
| def lr_at(step: int, cfg: TrainConfig) -> float: |
| """Warmup then cosine decay: ramp up linearly, then ease down to min_lr.""" |
| if step < cfg.warmup_steps: |
| return cfg.peak_lr * (step + 1) / cfg.warmup_steps |
| |
| progress = (step - cfg.warmup_steps) / max(1, cfg.max_steps - cfg.warmup_steps) |
| progress = min(1.0, progress) |
| coeff = 0.5 * (1.0 + math.cos(math.pi * progress)) |
| return cfg.min_lr + coeff * (cfg.peak_lr - cfg.min_lr) |
|
|
|
|
| def train(bin_paths, weights, tcfg: TrainConfig, mcfg: ARConfig, device="cuda"): |
| torch.manual_seed(0) |
| model = ARModel(mcfg).to(device) |
| model.train() |
|
|
| data = BinDataset(bin_paths, tcfg.seq_len, weights) |
|
|
| |
| opt = torch.optim.AdamW( |
| model.parameters(), lr=tcfg.peak_lr, |
| betas=tcfg.betas, weight_decay=tcfg.weight_decay, fused=True, |
| ) |
|
|
| os.makedirs(tcfg.ckpt_dir, exist_ok=True) |
| t0 = time.time() |
|
|
| for step in range(tcfg.max_steps): |
| |
| lr = lr_at(step, tcfg) |
| for g in opt.param_groups: |
| g["lr"] = lr |
|
|
| x, y = data.get_batch(tcfg.batch_size, device) |
|
|
| |
| logits = model(x) |
| loss = F.cross_entropy( |
| logits.view(-1, logits.size(-1)), y.view(-1) |
| ) |
|
|
| |
| opt.zero_grad(set_to_none=True) |
| loss.backward() |
| torch.nn.utils.clip_grad_norm_(model.parameters(), tcfg.grad_clip) |
| opt.step() |
|
|
| if step % tcfg.log_every == 0: |
| dt = time.time() - t0 |
| toks = (step + 1) * tcfg.batch_size * tcfg.seq_len |
| print(f"step {step:6d} | loss {loss.item():.4f} | lr {lr:.2e} " |
| f"| {toks/dt/1e3:.1f}k tok/s") |
|
|
| if step > 0 and step % tcfg.ckpt_every == 0: |
| path = os.path.join(tcfg.ckpt_dir, f"ar_step{step}.pt") |
| torch.save({"model": model.state_dict(), "step": step}, path) |
| print(f" saved {path}") |
|
|
|
|
| if __name__ == "__main__": |
| p = argparse.ArgumentParser() |
| p.add_argument("--bin", nargs="+", required=True, help="one or more .bin paths") |
| p.add_argument("--weights", nargs="+", type=float, default=None) |
| p.add_argument("--max_steps", type=int, default=100) |
| p.add_argument("--batch_size", type=int, default=4) |
| p.add_argument("--seq_len", type=int, default=4096) |
| p.add_argument("--warmup_steps", type=int, default=2000) |
| p.add_argument("--model_size", choices=["small", "full", "tiny"], default="small") |
| args = p.parse_args() |
|
|
| tcfg = TrainConfig(max_steps=args.max_steps, batch_size=args.batch_size, |
| seq_len=args.seq_len, warmup_steps=args.warmup_steps) |
| mcfg = {"small": small_config, "full": ARConfig, "tiny": tiny_config}[args.model_size]() |
| train(args.bin, args.weights, tcfg, mcfg) |
