| import argparse |
| import collections |
| import json |
| import logging |
| import math |
| import os |
| import sys |
| import time |
| from pathlib import Path |
|
|
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") |
|
|
| import torch |
| import torch.distributed as dist |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| _root = Path(__file__).parent.parent |
| if str(_root) not in sys.path: |
| sys.path.insert(0, str(_root)) |
|
|
| import model_cpu_gpt2 as _m |
|
|
| _m._fla_available = False |
| from model_cpu_gpt2 import CPUGPT, CPUGPTConfig |
|
|
| logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") |
| log = logging.getLogger(__name__) |
|
|
|
|
| def fela_250m_config() -> CPUGPTConfig: |
| return CPUGPTConfig( |
| n_layer=20, |
| n_embd=1024, |
| n_head=16, |
| ffn_hidden=2752, |
| fno_modes=256, |
| gla_chunk=256, |
| seq_len=8192, |
| layer_pattern="SSSL", |
| vocab_size=50257, |
| bias=False, |
| ) |
|
|
|
|
| polar_express_coeffs = [ |
| (8.156554524902461, -22.48329292557795, 15.878769915207462), |
| (4.042929935166739, -2.808917465908714, 0.5000178451051316), |
| (3.8916678022926607, -2.772484153217685, 0.5060648178503393), |
| (3.285753657755655, -2.3681294933425376, 0.46449024233003106), |
| (2.3465413258596377, -1.7097828382687081, 0.42323551169305323), |
| ] |
|
|
|
|
| def adamw_step(p, grad, m, v, step, lr, b1, b2, eps, wd): |
| p.mul_(1 - lr * wd) |
| m.lerp_(grad, 1 - b1) |
| v.lerp_(grad.square(), 1 - b2) |
| bc1 = 1 - b1**step |
| bc2 = 1 - b2**step |
| p.addcdiv_(m / bc1, (v / bc2).sqrt_().add_(eps), value=-lr) |
|
|
|
|
| def muon_step(grads_stack, params, mom_buf, lr, momentum=0.95, ns_steps=3): |
| mom_buf.lerp_(grads_stack, 1 - momentum) |
| X = mom_buf.float() |
| X = X / (X.norm(dim=(-2, -1), keepdim=True) * 1.02 + 1e-6) |
| for a, b, c in polar_express_coeffs[:ns_steps]: |
| if X.size(-2) >= X.size(-1): |
| A = X.mT @ X |
| X = a * X + X @ (b * A + c * (A @ A)) |
| else: |
| A = X @ X.mT |
| X = a * X + (b * A + c * (A @ A)) @ X |
| torch._foreach_sub_(params, list((X * lr).to(params[0].dtype).unbind(0))) |
|
|
|
|
| class MuonAdamW(torch.optim.Optimizer): |
| def __init__(self, param_groups): |
| super().__init__(param_groups, defaults={}) |
|
|
| @torch.no_grad() |
| def step(self): |
| for g in self.param_groups: |
| if g["kind"] == "adamw": |
| for p in g["params"]: |
| if p.grad is None: |
| continue |
| st = self.state[p] |
| if not st: |
| st["step"] = 0 |
| st["m"] = torch.zeros_like(p) |
| st["v"] = torch.zeros_like(p) |
| st["step"] += 1 |
| adamw_step( |
| p, |
| p.grad, |
| st["m"], |
| st["v"], |
| st["step"], |
| g["lr"], |
| *g["betas"], |
| g["eps"], |
| g.get("wd", 0), |
| ) |
| elif g["kind"] == "muon": |
| params = g["params"] |
| if not params: |
| continue |
| p0 = params[0] |
| st = self.state[p0] |
| stacked = torch.stack( |
| [p.grad for p in params if p.grad is not None] |
| ).float() |
| if not st: |
| st["mom"] = torch.zeros_like(stacked) |
| lr = g["lr"] * max(1.0, p0.shape[-2] / max(p0.shape[-1], 1)) ** 0.5 |
| muon_step(stacked, params, st["mom"], lr, g.get("momentum", 0.95)) |
|
|
|
|
| def build_optimizer(model, cfg, lr_matrix=0.01, lr_emb=0.02, lr_lm=0.002): |
| scale = (cfg.n_embd / 768) ** -0.5 |
| raw = ( |
| model.module |
| if isinstance(model, nn.parallel.DistributedDataParallel) |
| else model |
| ) |
| matrix_params, scalar_params = [], [] |
| for block in raw.blocks: |
| for p in block.parameters(): |
| (matrix_params if p.ndim == 2 else scalar_params).append(p) |
| groups = [ |
| dict( |
| kind="adamw", |
| params=list(raw.wte.parameters()), |
| lr=lr_emb * scale, |
| betas=(0.8, 0.95), |
| eps=1e-8, |
| ), |
| dict( |
| kind="adamw", |
| params=list(raw.lm_head.parameters()) |
| if raw.lm_head.weight is not raw.wte.weight |
| else [], |
| lr=lr_lm * scale, |
| betas=(0.8, 0.95), |
| eps=1e-8, |
| ), |
| dict( |
| kind="adamw", |
| params=scalar_params, |
| lr=lr_matrix * scale, |
| betas=(0.8, 0.95), |
| eps=1e-8, |
| ), |
| ] |
| for shape in sorted({p.shape for p in matrix_params}): |
| ps = [p for p in matrix_params if p.shape == shape] |
| groups.append(dict(kind="muon", params=ps, lr=lr_matrix, momentum=0.95)) |
| opt = MuonAdamW(groups) |
| for g in opt.param_groups: |
| g["initial_lr"] = g["lr"] |
| return opt |
|
|
|
|
| def lr_multiplier(progress, warmup=0.01, min_ratio=0.1): |
| if progress < warmup: |
| return progress / warmup |
| t = (progress - warmup) / (1.0 - warmup) |
| return min_ratio + (1.0 - min_ratio) * 0.5 * (1.0 + math.cos(math.pi * t)) |
|
|
|
|
| def make_loader(data_dir, seq_len, device, batch_size=1, rank=0, world_size=1): |
| import glob |
| import queue |
| import threading |
|
|
| all_parquet = sorted(glob.glob(os.path.join(data_dir, "*.parquet"))) |
| all_bins = sorted(glob.glob(os.path.join(data_dir, "*.bin"))) |
|
|
| if all_parquet: |
| import pyarrow.parquet as pq |
|
|
| shards = all_parquet[rank::world_size] or all_parquet |
|
|
| def load_fn(path): |
| tbl = pq.read_table(path, columns=["tokens"]) |
| t = torch.tensor(tbl["tokens"].to_pylist()[0], dtype=torch.long) |
| n = (len(t) // (seq_len + 1)) * (seq_len + 1) |
| return t[:n].view(-1, seq_len + 1) |
| elif all_bins: |
| import numpy as np |
|
|
| shards = all_bins[rank::world_size] or all_bins |
|
|
| def load_fn(path): |
| data = np.memmap(path, dtype=np.int32, mode="r") |
| n = (len(data) // (seq_len + 1)) * (seq_len + 1) |
| return torch.from_numpy(data[:n].copy()).long().view(-1, seq_len + 1) |
| else: |
| raise FileNotFoundError(f"No parquet or bin shards in {data_dir}") |
|
|
| buf: queue.Queue = queue.Queue(maxsize=4) |
|
|
| def _producer(): |
| ep = 1 |
| while True: |
| for shard in shards: |
| seqs = load_fn(shard) |
| if len(seqs) == 0: |
| continue |
| idx = torch.randperm(len(seqs)) |
| batch = [] |
| for i in idx: |
| batch.append(seqs[i]) |
| if len(batch) == batch_size: |
| chunk = torch.stack(batch) |
| buf.put((chunk[:, :-1], chunk[:, 1:], ep)) |
| batch = [] |
| ep += 1 |
|
|
| threading.Thread(target=_producer, daemon=True).start() |
| while True: |
| x, y, ep = buf.get() |
| yield x.to(device), y.to(device), ep |
|
|
|
|
| def eval_wikitext_bpb(model, cfg, device): |
| import tiktoken |
|
|
| try: |
| from datasets import load_dataset |
|
|
| try: |
| ds = load_dataset( |
| "Salesforce/wikitext", |
| "wikitext-103-raw-v1", |
| split="test", |
| trust_remote_code=True, |
| ) |
| except Exception: |
| ds = load_dataset( |
| "wikitext", "wikitext-103-raw-v1", split="test", trust_remote_code=True |
| ) |
| text = "\n\n".join(ds["text"]) |
| except Exception as e: |
| log.warning(f"WikiText load failed: {e}") |
| return {"bpb": float("nan"), "nats": float("nan")} |
|
|
| enc = tiktoken.get_encoding("gpt2") |
| tokens = torch.tensor(enc.encode(text[:5_000_000]), dtype=torch.long)[:1_000_000] |
|
|
| raw = ( |
| model.module |
| if isinstance(model, nn.parallel.DistributedDataParallel) |
| else model |
| ) |
| raw.eval() |
| seq_len = cfg.seq_len |
| n_chunks = min(100, (len(tokens) - 1) // seq_len) |
| total_nll, total_tok = 0.0, 0 |
|
|
| with torch.no_grad(): |
| for i in range(n_chunks): |
| x = tokens[i * seq_len : (i + 1) * seq_len].unsqueeze(0).to(device) |
| y = tokens[i * seq_len + 1 : (i + 1) * seq_len + 1].unsqueeze(0).to(device) |
| if x.shape[1] < seq_len: |
| break |
| pad = cfg.gla_chunk - (x.shape[1] % cfg.gla_chunk or cfg.gla_chunk) |
| if pad < cfg.gla_chunk: |
| x = F.pad(x, (0, pad)) |
| logits = raw(x) |
| lgt = logits[0, :seq_len].float() |
| loss = F.cross_entropy(lgt, y[0], reduction="sum") |
| total_nll += loss.item() |
| total_tok += seq_len |
|
|
| nats = total_nll / max(total_tok, 1) |
| bpb = nats / math.log(2) |
| ppl = math.exp(nats) |
| log.info(f"WikiText eval: BPB={bpb:.4f} nats={nats:.4f} PPL={ppl:.2f}") |
| return {"bpb": bpb, "nats": nats, "ppl": ppl, "tokens_evaluated": total_tok} |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--data-dir", default="/data/openwebtext") |
| parser.add_argument("--out-dir", default="results/scaling_midpoint_v2") |
| parser.add_argument("--device-batch", type=int, default=4) |
| parser.add_argument("--smoke", action="store_true") |
| parser.add_argument("--resume", type=str, default=None) |
| args = parser.parse_args() |
|
|
| use_ddp = "RANK" in os.environ |
| local_rank, global_rank = 0, 0 |
| world_size = 1 |
| if use_ddp: |
| dist.init_process_group(backend="nccl") |
| local_rank = dist.get_rank() % torch.cuda.device_count() |
| global_rank = dist.get_rank() |
| world_size = dist.get_world_size() |
| torch.cuda.set_device(local_rank) |
|
|
| device = f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu" |
| is_master = global_rank == 0 |
| num_gpus = world_size |
|
|
| cfg = fela_250m_config() |
| model = CPUGPT(cfg).to(device) |
| if use_ddp: |
| model = nn.parallel.DistributedDataParallel(model, device_ids=[local_rank]) |
|
|
| raw = model.module if use_ddp else model |
| actual_params = raw.param_count() |
|
|
| if is_master: |
| log.info(f"FELA ~250M v2: {actual_params / 1e6:.1f}M params") |
| log.info(f" n_layer={cfg.n_layer} n_embd={cfg.n_embd} seq_len={cfg.seq_len}") |
| log.info(f" lr_emb=0.02 lr_matrix=0.01 warmup=1% [FIXED from v1]") |
|
|
| opt = build_optimizer(model, cfg) |
|
|
| chinchilla_tokens = actual_params * 20 |
| if args.smoke: |
| chinchilla_tokens = 10_000_000 |
| log.info("Smoke: 10M tokens") |
| else: |
| if is_master: |
| log.info(f"Chinchilla target: {chinchilla_tokens / 1e9:.2f}B tokens") |
|
|
| total_batch_tokens = args.device_batch * cfg.seq_len * num_gpus |
| total_steps = max(1, int(chinchilla_tokens / total_batch_tokens)) |
|
|
| if is_master: |
| log.info(f" total_steps={total_steps} batch_tokens={total_batch_tokens:,}") |
|
|
| out_dir = Path(args.out_dir) |
| ckpt_dir = out_dir / "ckpt" |
| if is_master: |
| ckpt_dir.mkdir(parents=True, exist_ok=True) |
|
|
| start_step = 0 |
| if args.resume: |
| ck = torch.load(args.resume, map_location="cpu") |
| sd = { |
| k.replace("module.", "").replace("_orig_mod.", ""): v |
| for k, v in ck["model"].items() |
| } |
| raw.load_state_dict(sd) |
| try: |
| opt.load_state_dict(ck["opt"]) |
| except Exception: |
| pass |
| start_step = ck.get("step", 0) |
| if is_master: |
| log.info(f"Resumed from {args.resume} at step {start_step}") |
|
|
| loader = make_loader( |
| args.data_dir, cfg.seq_len, device, args.device_batch, global_rank, world_size |
| ) |
|
|
| loss_window: collections.deque = collections.deque(maxlen=50) |
| t0 = time.perf_counter() |
| step = start_step |
|
|
| for x, y, _ in loader: |
| if step >= total_steps: |
| break |
|
|
| model.train() |
| progress = step / total_steps |
| mult = lr_multiplier(progress) |
| for g in opt.param_groups: |
| g["lr"] = g["initial_lr"] * mult |
|
|
| with torch.autocast(device_type=device.split(":")[0], dtype=torch.bfloat16): |
| loss = model(x, y) |
|
|
| loss.backward() |
| nn.utils.clip_grad_norm_(model.parameters(), 1.0) |
| opt.step() |
| opt.zero_grad(set_to_none=True) |
| step += 1 |
|
|
| raw_loss = loss.item() |
| loss_window.append(raw_loss) |
|
|
| if is_master and step % 100 == 0: |
| elapsed = time.perf_counter() - t0 |
| tps = (step - start_step) * total_batch_tokens / elapsed |
| eta_h = (total_steps - step) * total_batch_tokens / max(tps, 1) / 3600 |
| avg_loss = sum(loss_window) / len(loss_window) |
| log.info( |
| f"step={step}/{total_steps} ({100 * step / total_steps:.0f}%) " |
| f"loss={raw_loss:.4f} avg50={avg_loss:.4f} " |
| f"tok/s={tps:,.0f} ETA={eta_h:.1f}h" |
| ) |
|
|
| if step >= 10000 and avg_loss > 6.0: |
| log.error( |
| f"DIVERGENCE DETECTED: avg50={avg_loss:.4f} > 6.0 at step {step}" |
| ) |
| log.error("Aborting. Check LR config or data.") |
| if use_ddp: |
| dist.destroy_process_group() |
| sys.exit(1) |
|
|
| if is_master and step % 5000 == 0: |
| ck_path = ckpt_dir / f"step_{step:06d}.pt" |
| torch.save( |
| {"step": step, "model": raw.state_dict(), "opt": opt.state_dict()}, |
| ck_path, |
| ) |
| avg_loss = sum(loss_window) / len(loss_window) |
| log.info(f"Checkpoint saved: {ck_path} avg50_loss={avg_loss:.4f}") |
|
|
| if is_master: |
| torch.save( |
| {"step": step, "model": raw.state_dict(), "opt": opt.state_dict()}, |
| ckpt_dir / "final.pt", |
| ) |
| log.info(f"Training complete at step {step}. Saved final.pt") |
|
|
| if use_ddp: |
| dist.barrier() |
|
|
| if is_master: |
| log.info("=== WikiText BPB eval ===") |
| raw.eval() |
| wt = eval_wikitext_bpb(raw, cfg, device) |
|
|
| tokens_trained = step * total_batch_tokens |
| summary = { |
| "model": "FELA_250M_v2", |
| "params_M": actual_params / 1e6, |
| "n_layer": cfg.n_layer, |
| "n_embd": cfg.n_embd, |
| "seq_len": cfg.seq_len, |
| "tokens_trained": tokens_trained, |
| "chinchilla_optimal_tokens": chinchilla_tokens, |
| "pct_chinchilla": round(tokens_trained / chinchilla_tokens * 100, 1), |
| "steps": step, |
| "wikitext_bpb": wt["bpb"], |
| "wikitext_nats": wt.get("nats"), |
| "wikitext_ppl": wt.get("ppl"), |
| "lr_emb": 0.02, |
| "lr_matrix": 0.01, |
| "fix_note": "v1 had lr_emb=0.1 (5x too high), caused divergence at step 25K", |
| } |
|
|
| (out_dir / "wikitext.json").write_text(json.dumps(wt, indent=2)) |
| (out_dir / "summary.json").write_text(json.dumps(summary, indent=2)) |
|
|
| print("\n" + "=" * 60) |
| print("SCALING LAW DATA POINTS") |
| print("=" * 60) |
| print(f" {'Model':<20} {'Params':>8} {'Tokens':>8} {'BPB':>8}") |
| print(f" {'-' * 48}") |
| print(f" {'FELA 124M (v9)':<20} {'124M':>8} {'2B':>8} {'1.297':>8}") |
| print( |
| f" {'FELA ~250M (v2)':<20} {actual_params / 1e6:>7.0f}M " |
| f"{tokens_trained / 1e9:>7.1f}B {wt['bpb']:>8.4f}" |
| ) |
| print(f" {'FELA 1.13B (22B)':<20} {'1130M':>8} {'22B':>8} {'TBD':>8}") |
| print("=" * 60) |
| log.info(f"Outputs: {out_dir}/") |
|
|
| if use_ddp: |
| dist.destroy_process_group() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
