| import argparse |
| import gc |
| import json |
| import logging |
| import math |
| import os |
| import random |
| 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 |
|
|
| _script_dir = os.path.dirname(os.path.abspath(__file__)) |
| _root = os.path.dirname(_script_dir) |
| if _root not in sys.path: |
| sys.path.insert(0, _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__) |
|
|
| CHUNK_VARIANTS = [64, 256, 512, 1024] |
|
|
|
|
| def make_124m_config(gla_chunk: int) -> CPUGPTConfig: |
| return CPUGPTConfig( |
| n_layer=12, |
| n_embd=768, |
| n_head=12, |
| ffn_hidden=2048, |
| fno_modes=256, |
| gla_chunk=gla_chunk, |
| seq_len=8192, |
| layer_pattern="SSSL", |
| vocab_size=50257, |
| ) |
|
|
|
|
| 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.02, 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): |
| import glob |
| import queue |
| import threading |
|
|
| shards = sorted(glob.glob(os.path.join(data_dir, "*.parquet"))) |
| if not shards: |
| shards = sorted(glob.glob(os.path.join(data_dir, "*.bin"))) |
| if not shards: |
| raise FileNotFoundError(f"No parquet or bin shards in {data_dir}") |
| import numpy as np |
|
|
| 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: |
| import pyarrow.parquet as pq |
|
|
| def load_fn(path): |
| tbl = pq.read_table(path, columns=["tokens"]) |
| tokens = torch.tensor(tbl["tokens"].to_pylist()[0], dtype=torch.long) |
| n = (len(tokens) // (seq_len + 1)) * (seq_len + 1) |
| return tokens[:n].view(-1, seq_len + 1) |
|
|
| 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 train_variant( |
| chunk: int, |
| args, |
| out_dir: Path, |
| device: str, |
| is_master: bool, |
| use_ddp: bool, |
| local_rank: int, |
| ): |
| log.info(f"=== Training gla_chunk={chunk} ===") |
| cfg = make_124m_config(chunk) |
| model = CPUGPT(cfg).to(device) |
| if use_ddp: |
| model = nn.parallel.DistributedDataParallel(model, device_ids=[local_rank]) |
| opt = build_optimizer(model, cfg) |
| if is_master: |
| raw = model.module if use_ddp else model |
| log.info( |
| f" params: {raw.param_count() / 1e6:.1f}M seq_len={cfg.seq_len} chunk={chunk}" |
| ) |
|
|
| total_batch = args.device_batch * max(args.num_gpus, 1) |
| total_steps = max(1, int(args.tokens / (total_batch * cfg.seq_len))) |
| log.info( |
| f" steps={total_steps} total_batch={total_batch} tokens≈{total_steps * total_batch * cfg.seq_len / 1e6:.0f}M" |
| ) |
|
|
| loader = make_loader(args.data_dir, cfg.seq_len, device, args.device_batch) |
| ctx = ( |
| torch.autocast(device_type="cuda", dtype=torch.bfloat16) |
| if "cuda" in device |
| else torch.no_grad().__class__() |
| ) |
|
|
| step = 0 |
| t0 = time.perf_counter() |
| 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 |
|
|
| if is_master and step % 50 == 0: |
| elapsed = time.perf_counter() - t0 |
| tps = step * total_batch * cfg.seq_len / elapsed |
| log.info( |
| f" chunk={chunk} step={step}/{total_steps} ({100 * step / total_steps:.0f}%) " |
| f"loss={loss.item():.4f} tok/s={tps:,.0f}" |
| ) |
|
|
| ckpt_dir = out_dir / f"variant_{chunk}" / "ckpt" |
| ckpt_dir.mkdir(parents=True, exist_ok=True) |
| if is_master: |
| raw = model.module if use_ddp else model |
| torch.save({"step": step, "model": raw.state_dict()}, ckpt_dir / "final.pt") |
| log.info(f" saved {ckpt_dir}/final.pt") |
|
|
| return model, cfg |
|
|
|
|
| _HAY = ( |
| "The forest was quiet except for the occasional rustle of leaves in the breeze. " |
| "Sunlight filtered through the canopy, casting dappled shadows on the mossy ground. " |
| "A small stream wound its way between ancient oak trees, its water clear and cold. " |
| "Birds called to one another across the branches, their songs filling the still air. " |
| "Somewhere in the distance, a woodpecker drummed steadily against hollow bark. " |
| "The smell of earth and pine needles rose from the path with each careful step. " |
| "Nothing moved except the shadows and the water and the swaying tops of the tallest trees. " |
| ) * 300 |
|
|
|
|
| def eval_needle(model, cfg, device, n_trials=20, seed=42): |
| import tiktoken |
|
|
| enc = tiktoken.get_encoding("gpt2") |
| hay_ids = enc.encode_ordinary(_HAY) |
| rng = random.Random(seed) |
|
|
| ctx_lens = [512, 1024, 2048, 4096, cfg.seq_len] |
| depths = [0.1, 0.3, 0.5, 0.7, 0.9] |
|
|
| raw = ( |
| model.module |
| if isinstance(model, nn.parallel.DistributedDataParallel) |
| else model |
| ) |
| raw.eval() |
| results = [] |
|
|
| for ctx_len in ctx_lens: |
| for depth in depths: |
| n_correct, n_valid = 0, 0 |
| for _ in range(n_trials): |
| code = str(rng.randint(10000, 99999)) |
| dist_ = str(rng.randint(10000, 99999)) |
| while dist_ == code: |
| dist_ = str(rng.randint(10000, 99999)) |
|
|
| needle_ids = enc.encode_ordinary( |
| f" The passkey is {code}. Remember this passkey." |
| ) |
| query_ids = enc.encode_ordinary(" The passkey is") |
| code_ids = enc.encode_ordinary(" " + code) |
| dist_ids = enc.encode_ordinary(" " + dist_) |
|
|
| overhead = ( |
| len(needle_ids) + len(query_ids) + max(len(code_ids), len(dist_ids)) |
| ) |
| hay_budget = ctx_len - overhead |
| if hay_budget < 64: |
| continue |
|
|
| prefix_len = int(hay_budget * depth) |
| suffix_len = hay_budget - prefix_len |
| start = rng.randint(0, max(0, len(hay_ids) - hay_budget - 1)) |
| hay_slice = hay_ids[start : start + hay_budget] |
| context = ( |
| hay_slice[:prefix_len] |
| + needle_ids |
| + hay_slice[prefix_len : prefix_len + suffix_len] |
| + query_ids |
| ) |
|
|
| def avg_nll(answer_ids): |
| combined = context + answer_ids |
| remainder = len(combined) % cfg.gla_chunk |
| extra = (cfg.gla_chunk - remainder) % cfg.gla_chunk |
| full = combined + [0] * extra |
| inp = torch.tensor([full], dtype=torch.long, device=device) |
| with torch.no_grad(): |
| logits = raw(inp) |
| s = len(context) - 1 |
| lgt = logits[0, s : s + len(answer_ids)].float() |
| tgt = torch.tensor(answer_ids, dtype=torch.long, device=device) |
| return F.cross_entropy(lgt, tgt, reduction="mean").item() |
|
|
| try: |
| nll_c = avg_nll(code_ids) |
| nll_d = avg_nll(dist_ids) |
| n_valid += 1 |
| if nll_c < nll_d: |
| n_correct += 1 |
| except Exception: |
| pass |
|
|
| acc = n_correct / n_valid if n_valid else float("nan") |
| results.append( |
| { |
| "ctx_len": ctx_len, |
| "depth": depth, |
| "accuracy": acc, |
| "n_correct": n_correct, |
| "n_valid": n_valid, |
| } |
| ) |
| log.info( |
| f" [needle] chunk={cfg.gla_chunk} ctx={ctx_len} depth={depth:.1f} " |
| f"acc={acc:.0%} ({n_correct}/{n_valid})" |
| ) |
| return results |
|
|
|
|
| def eval_throughput(model, cfg, device): |
| raw = ( |
| model.module |
| if isinstance(model, nn.parallel.DistributedDataParallel) |
| else model |
| ) |
| raw.eval() |
| results = [] |
| ctx_lens = [512, 1024, 2048, 4096, cfg.seq_len] |
| for ctx_len in ctx_lens: |
| padded = math.ceil(ctx_len / cfg.gla_chunk) * cfg.gla_chunk |
| x = torch.zeros(1, padded, dtype=torch.long, device=device) |
| for _ in range(2): |
| with torch.no_grad(): |
| raw(x) |
| torch.cuda.synchronize() if "cuda" in device else None |
| t0 = time.perf_counter() |
| reps = 5 |
| for _ in range(reps): |
| with torch.no_grad(): |
| raw(x) |
| torch.cuda.synchronize() if "cuda" in device else None |
| dt = (time.perf_counter() - t0) / reps |
| tps = padded / dt |
| results.append( |
| {"ctx_len": ctx_len, "tok_per_sec": tps, "ms_per_seq": dt * 1000} |
| ) |
| log.info( |
| f" [throughput] chunk={cfg.gla_chunk} ctx={ctx_len} → {tps:,.0f} tok/s" |
| ) |
| return results |
|
|
|
|
| def eval_wikitext(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: |
| log.warning(" Could not load WikiText-103; skipping BPB eval") |
| return {"bpb": float("nan")} |
|
|
| enc = tiktoken.get_encoding("gpt2") |
| tokens = torch.tensor(enc.encode(text[:4_000_000]), dtype=torch.long) |
| tokens = tokens[:1_000_000] |
|
|
| raw = ( |
| model.module |
| if isinstance(model, nn.parallel.DistributedDataParallel) |
| else model |
| ) |
| raw.eval() |
|
|
| seq_len = cfg.seq_len |
| n_chunks = min(50, (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 or y.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) |
| log.info(f" [wikitext] chunk={cfg.gla_chunk} BPB={bpb:.4f}") |
| return {"bpb": bpb, "nats": nats, "tokens_evaluated": total_tok} |
|
|
|
|
| def make_figures(all_results, out_dir: Path): |
| try: |
| import matplotlib |
|
|
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
| import numpy as np |
| except ImportError: |
| log.warning("matplotlib not available; skipping figures") |
| return |
|
|
| chunks = sorted(all_results.keys()) |
| colors = ["#E84855", "#F4A261", "#2E86AB", "#264653"] |
|
|
| fig, axes = plt.subplots(2, 2, figsize=(14, 9), sharex=True, sharey=True) |
| axes = axes.flatten() |
| for idx, chunk in enumerate(chunks): |
| rows = all_results[chunk]["needle"] |
| ctx_u = sorted(set(r["ctx_len"] for r in rows)) |
| dep_u = sorted(set(r["depth"] for r in rows)) |
| grid = np.full((len(dep_u), len(ctx_u)), float("nan")) |
| for r in rows: |
| ci = ctx_u.index(r["ctx_len"]) |
| di = dep_u.index(r["depth"]) |
| if not (isinstance(r["accuracy"], float) and math.isnan(r["accuracy"])): |
| grid[di, ci] = r["accuracy"] |
| ax = axes[idx] |
| im = ax.imshow(grid, vmin=0, vmax=1, cmap="RdYlGn", aspect="auto") |
| ax.set_title(f"gla_chunk={chunk}", fontsize=12, fontweight="bold") |
| ax.set_xticks(range(len(ctx_u))) |
| ax.set_xticklabels( |
| [f"{c // 1024}K" if c >= 1024 else str(c) for c in ctx_u], fontsize=9 |
| ) |
| ax.set_yticks(range(len(dep_u))) |
| ax.set_yticklabels([f"{int(d * 100)}%" for d in dep_u], fontsize=9) |
| ax.set_xlabel("Context Length", fontsize=10) |
| ax.set_ylabel("Needle Depth", fontsize=10) |
| for di in range(len(dep_u)): |
| for ci in range(len(ctx_u)): |
| v = grid[di, ci] |
| if not math.isnan(v): |
| col = "black" if 0.25 < v < 0.75 else "white" |
| ax.text( |
| ci, |
| di, |
| f"{v:.0%}", |
| ha="center", |
| va="center", |
| fontsize=9, |
| color=col, |
| fontweight="bold", |
| ) |
| plt.colorbar(im, ax=ax) |
| fig.suptitle( |
| "Passkey Retrieval Accuracy by GLA Chunk Size\n(124M FELA, 1B tokens)", |
| fontsize=14, |
| ) |
| fig.tight_layout() |
| fig.savefig(out_dir / "fig_chunk_retrieval.png", dpi=150, bbox_inches="tight") |
| plt.close(fig) |
| log.info(" saved fig_chunk_retrieval.png") |
|
|
| fig, ax = plt.subplots(figsize=(8, 5)) |
| for chunk, color in zip(chunks, colors): |
| rows = all_results[chunk].get("throughput", []) |
| if not rows: |
| continue |
| xs = [r["ctx_len"] for r in rows] |
| ys = [r["tok_per_sec"] for r in rows] |
| ax.plot( |
| xs, |
| ys, |
| "o-", |
| color=color, |
| linewidth=2.5, |
| markersize=6, |
| label=f"chunk={chunk}", |
| ) |
| ax.set_xscale("log", base=2) |
| ax.set_xlabel("Context Length (tokens)", fontsize=12) |
| ax.set_ylabel("Throughput (tok/s)", fontsize=12) |
| ax.set_title( |
| "Throughput vs. Context Length\n(124M FELA — chunk size effect)", fontsize=13 |
| ) |
| ax.legend(fontsize=11) |
| ax.grid(True, alpha=0.3) |
| fig.tight_layout() |
| fig.savefig(out_dir / "fig_chunk_throughput.png", dpi=150, bbox_inches="tight") |
| plt.close(fig) |
| log.info(" saved fig_chunk_throughput.png") |
|
|
| fig, axes = plt.subplots(1, 2, figsize=(11, 4)) |
| acc_hard, bpb_vals = [], [] |
| for chunk in chunks: |
| needle_rows = all_results[chunk]["needle"] |
| match = [ |
| r |
| for r in needle_rows |
| if r["ctx_len"] == 4096 and abs(r["depth"] - 0.1) < 0.01 |
| ] |
| acc_hard.append(match[0]["accuracy"] if match else float("nan")) |
| bpb_vals.append(all_results[chunk].get("wikitext", {}).get("bpb", float("nan"))) |
|
|
| c_labels = [f"chunk={c}" for c in chunks] |
| ax1 = axes[0] |
| bars = ax1.bar(c_labels, acc_hard, color=colors) |
| ax1.axhline(0.5, color="gray", linestyle="--", alpha=0.7, label="random") |
| ax1.set_ylim(0, 1) |
| ax1.set_ylabel("Retrieval Accuracy") |
| ax1.set_title("Hard case: ctx=4K, depth=10%\n(needle at very start)") |
| ax1.legend(fontsize=9) |
| for bar, v in zip(bars, acc_hard): |
| if not math.isnan(v): |
| ax1.text( |
| bar.get_x() + bar.get_width() / 2, |
| v + 0.02, |
| f"{v:.0%}", |
| ha="center", |
| fontsize=10, |
| fontweight="bold", |
| ) |
|
|
| ax2 = axes[1] |
| bars2 = ax2.bar(c_labels, bpb_vals, color=colors) |
| ax2.set_ylabel("WikiText BPB (lower = better)") |
| ax2.set_title("Language quality\n(WikiText-103 BPB)") |
| for bar, v in zip(bars2, bpb_vals): |
| if not math.isnan(v): |
| ax2.text( |
| bar.get_x() + bar.get_width() / 2, |
| v + 0.005, |
| f"{v:.3f}", |
| ha="center", |
| fontsize=10, |
| fontweight="bold", |
| ) |
|
|
| fig.suptitle("Chunk Size Tradeoff: Retrieval vs. Quality", fontsize=13) |
| fig.tight_layout() |
| fig.savefig(out_dir / "fig_chunk_summary.png", dpi=150, bbox_inches="tight") |
| plt.close(fig) |
| log.info(" saved fig_chunk_summary.png") |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description="GLA chunk-size ablation") |
| parser.add_argument("--data-dir", default="/data/openwebtext") |
| parser.add_argument( |
| "--tokens", |
| type=float, |
| default=1e9, |
| help="training tokens per variant (default 1B)", |
| ) |
| parser.add_argument( |
| "--device-batch", type=int, default=4, help="per-device batch size (sequences)" |
| ) |
| parser.add_argument( |
| "--num-gpus", type=int, default=0, help="override GPU count (0 = auto-detect)" |
| ) |
| parser.add_argument("--out-dir", default="results/chunk_ablation") |
| parser.add_argument("--needle-trials", type=int, default=20) |
| parser.add_argument( |
| "--smoke", |
| action="store_true", |
| help="fast smoke test: 5M tokens, 3 needle trials", |
| ) |
| parser.add_argument( |
| "--eval-only", |
| type=str, |
| default=None, |
| help="path to existing results dir — skip training, just eval+figures", |
| ) |
| parser.add_argument( |
| "--chunks", |
| type=int, |
| nargs="+", |
| default=CHUNK_VARIANTS, |
| help="which chunk sizes to run (default: 64 256 512 1024)", |
| ) |
| args = parser.parse_args() |
|
|
| if args.smoke: |
| args.tokens = 5e6 |
| args.needle_trials = 3 |
|
|
| use_ddp = False |
| local_rank, global_rank = 0, 0 |
| if "RANK" in os.environ: |
| dist.init_process_group(backend="nccl") |
| local_rank = dist.get_rank() % torch.cuda.device_count() |
| global_rank = dist.get_rank() |
| use_ddp = True |
| torch.cuda.set_device(local_rank) |
|
|
| num_gpus = args.num_gpus or ( |
| torch.cuda.device_count() if torch.cuda.is_available() else 0 |
| ) |
| args.num_gpus = max(1, num_gpus) |
| device = f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu" |
| is_master = global_rank == 0 |
|
|
| out_dir = Path(args.out_dir) |
| if is_master: |
| out_dir.mkdir(parents=True, exist_ok=True) |
|
|
| all_results = {} |
|
|
| for chunk in args.chunks: |
| variant_dir = out_dir / f"variant_{chunk}" |
| if is_master: |
| variant_dir.mkdir(parents=True, exist_ok=True) |
|
|
| model, cfg = train_variant( |
| chunk, args, out_dir, device, is_master, use_ddp, local_rank |
| ) |
|
|
| if use_ddp: |
| dist.barrier() |
|
|
| if is_master: |
| raw = model.module if use_ddp else model |
| raw.eval() |
|
|
| log.info(f"--- Evaluating chunk={chunk} ---") |
| needle_res = eval_needle(raw, cfg, device, n_trials=args.needle_trials) |
| (variant_dir / "needle.json").write_text(json.dumps(needle_res, indent=2)) |
|
|
| tput_res = eval_throughput(raw, cfg, device) |
| (variant_dir / "throughput.json").write_text(json.dumps(tput_res, indent=2)) |
|
|
| wiki_res = eval_wikitext(raw, cfg, device) |
| (variant_dir / "wikitext.json").write_text(json.dumps(wiki_res, indent=2)) |
|
|
| all_results[chunk] = { |
| "needle": needle_res, |
| "throughput": tput_res, |
| "wikitext": wiki_res, |
| "params_M": raw.param_count() / 1e6, |
| "gla_chunk": chunk, |
| } |
|
|
| del model |
| gc.collect() |
| if torch.cuda.is_available(): |
| torch.cuda.empty_cache() |
|
|
| if use_ddp: |
| dist.barrier() |
|
|
| if is_master: |
| (out_dir / "summary.json").write_text(json.dumps(all_results, indent=2)) |
| log.info(f"Summary written to {out_dir}/summary.json") |
|
|
| make_figures(all_results, out_dir) |
|
|
| print("\n" + "=" * 65) |
| print("CHUNK ABLATION SUMMARY") |
| print("=" * 65) |
| print( |
| f"{'chunk':>8} {'4K/10% acc':>10} {'4K/50% acc':>10} {'4K/90% acc':>10} {'BPB':>7}" |
| ) |
| print("-" * 65) |
| for chunk in sorted(all_results.keys()): |
| nr = all_results[chunk]["needle"] |
|
|
| def get_acc(ctx, dep): |
| m = [ |
| r |
| for r in nr |
| if r["ctx_len"] == ctx and abs(r["depth"] - dep) < 0.01 |
| ] |
| return m[0]["accuracy"] if m else float("nan") |
|
|
| bpb = all_results[chunk].get("wikitext", {}).get("bpb", float("nan")) |
| print( |
| f"{chunk:>8} {get_acc(4096, 0.1):>10.0%} {get_acc(4096, 0.5):>10.0%} " |
| f"{get_acc(4096, 0.9):>10.0%} {bpb:>7.4f}" |
| ) |
| print("=" * 65) |
|
|
| if use_ddp: |
| dist.destroy_process_group() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
