| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| """Train the matched byte-Transformer baseline on identical terms to Trident. |
| |
| Same corpus (FineWeb-Edu), same held-out set (wikitext-103), same byte vocab, |
| same optimizer/schedule, same data budget. Pushes a checkpoint + val BPB so the |
| head-to-head is apples-to-apples. |
| """ |
| import json |
| import math |
| import os |
| import sys |
| import time |
| from pathlib import Path |
|
|
| import torch |
|
|
|
|
| def env(k, d=None): |
| v = os.environ.get(k) |
| return v if v not in (None, "") else d |
|
|
|
|
| def env_i(k, d): |
| return int(env(k, d)) |
|
|
|
|
| def env_f(k, d): |
| return float(env(k, d)) |
|
|
|
|
| def log(m): |
| print(f"[baseline] {m}", flush=True) |
|
|
|
|
| def main(): |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| code_repo = env("CODE_REPO", "farguney/trident") |
| ckpt_dir = env("CKPT_DIR", "checkpoints_baseline") |
|
|
| from huggingface_hub import HfApi, snapshot_download |
| work = Path("/tmp/trident_code") |
| snapshot_download(repo_id=code_repo, repo_type="model", |
| allow_patterns=["src/**"], local_dir=str(work)) |
| sys.path.insert(0, str(work / "src")) |
| from trident.baseline import ByteTransformer |
| from trident.data import ByteWindowIterable, collate, load_fixed_byte_windows |
|
|
| if device == "cuda": |
| log(f"gpu: {torch.cuda.get_device_name(0)} torch {torch.__version__}") |
|
|
| seq_len = env_i("SEQ_LEN", 2048) |
| batch = env_i("BATCH", 192) |
| grad_accum = env_i("GRAD_ACCUM", 1) |
| d_model = env_i("D_MODEL", 512) |
| n_layers = env_i("N_LAYERS", 10) |
| n_heads = env_i("N_HEADS", 8) |
| d_ff = env_i("D_FF", 2176) |
|
|
| model = ByteTransformer(d_model=d_model, n_layers=n_layers, n_heads=n_heads, |
| d_ff=d_ff, max_seq=max(seq_len, 8192)).to(device) |
| if env_i("GRAD_CKPT", 0) == 1: |
| pass |
| log(f"baseline params={model.num_params()/1e6:.2f}M d_model={d_model} layers={n_layers} heads={n_heads} d_ff={d_ff}") |
|
|
| dataset = env("DATASET", "HuggingFaceFW/fineweb-edu") |
| ds_name = env("DATASET_NAME", "sample-10BT") |
| split = env("SPLIT", "train") |
| text_field = env("TEXT_FIELD", "text") |
| train_ds = ByteWindowIterable( |
| dataset=dataset, split=split, seq_len=seq_len, text_field=text_field, |
| name=ds_name, shuffle_buffer=env_i("SHUFFLE_BUFFER", 1000), seed=env_i("SEED", 0), |
| ) |
| loader = torch.utils.data.DataLoader(train_ds, batch_size=batch, collate_fn=collate, |
| num_workers=env_i("NUM_WORKERS", 2), drop_last=True) |
|
|
| val_windows = None |
| try: |
| val_windows = load_fixed_byte_windows( |
| dataset=env("VAL_DATASET", "Salesforce/wikitext"), |
| split=env("VAL_SPLIT", "validation"), |
| name=env("VAL_NAME", "wikitext-103-raw-v1"), |
| text_field=env("VAL_TEXT_FIELD", "text"), |
| seq_len=seq_len, num_windows=env_i("VAL_WINDOWS", 64), add_eod=False, |
| ).to(device) |
| log(f"held-out windows: {tuple(val_windows.shape)}") |
| except Exception as e: |
| log(f"held-out load failed ({e})") |
|
|
| max_steps = env_i("MAX_STEPS", 1500) |
| warmup = env_i("WARMUP", 150) |
| lr = env_f("LR", 6e-4) |
| min_lr = env_f("MIN_LR", 6e-5) |
| wd = env_f("WEIGHT_DECAY", 0.1) |
| grad_clip = env_f("GRAD_CLIP", 1.0) |
|
|
| decay = [p for p in model.parameters() if p.ndim >= 2] |
| no_decay = [p for p in model.parameters() if p.ndim < 2] |
| opt = torch.optim.AdamW([{"params": decay, "weight_decay": wd}, |
| {"params": no_decay, "weight_decay": 0.0}], |
| lr=lr, betas=(0.9, 0.95), eps=1e-8) |
|
|
| def lr_at(step): |
| if step < warmup: |
| return lr * (step + 1) / warmup |
| t = min(1.0, (step - warmup) / max(1, max_steps - warmup)) |
| return min_lr + 0.5 * (lr - min_lr) * (1 + math.cos(math.pi * t)) |
|
|
| amp = torch.bfloat16 if device == "cuda" else torch.float32 |
| api = HfApi() |
|
|
| @torch.no_grad() |
| def evaluate(): |
| model.eval() |
| tot_nll, tot_bytes = 0.0, 0 |
| vb = env_i("VAL_BATCH", 16) |
| for i in range(0, val_windows.shape[0], vb): |
| chunk = val_windows[i:i + vb] |
| with torch.autocast(device_type="cuda", dtype=amp, enabled=device == "cuda"): |
| out = model(chunk, return_logits=True) |
| logits = out["logits"].float() |
| nll = torch.nn.functional.cross_entropy( |
| logits.reshape(-1, logits.shape[-1]), |
| chunk.clamp(0, logits.shape[-1] - 1).reshape(-1), reduction="sum") |
| tot_nll += nll.item() |
| tot_bytes += chunk.numel() |
| model.train() |
| return tot_nll / tot_bytes / math.log(2) |
|
|
| def save_checkpoint(tag, step, metrics): |
| from safetensors.torch import save_file |
| outdir = Path("/tmp/ckpt"); outdir.mkdir(exist_ok=True) |
| sd = {k: v.detach().cpu().contiguous() for k, v in model.state_dict().items()} |
| save_file(sd, str(outdir / "model.safetensors")) |
| cfg = dict(vocab=model.vocab, d_model=d_model, n_layers=n_layers, n_heads=n_heads, |
| d_ff=d_ff, max_seq=max(seq_len, 8192)) |
| (outdir / "config.json").write_text(json.dumps(cfg, indent=2)) |
| (outdir / "metrics.json").write_text(json.dumps(metrics, indent=2)) |
| api.upload_folder(folder_path=str(outdir), repo_id=code_repo, |
| path_in_repo=f"{ckpt_dir}/{tag}", repo_type="model", |
| commit_message=f"baseline {ckpt_dir}/{tag} @ step {step}") |
| log(f"pushed checkpoint '{ckpt_dir}/{tag}' (step {step})") |
|
|
| log(f"training: max_steps={max_steps} batch={batch} accum={grad_accum} seq_len={seq_len} " |
| f"eff_bytes/step={batch*grad_accum*seq_len}") |
| model.train() |
| data_iter = iter(loader) |
| t_win = time.time() |
| running = 0.0 |
| log_every = env_i("LOG_EVERY", 20) |
| eval_every = env_i("EVAL_EVERY", 250) |
| save_every = env_i("SAVE_EVERY", 250) |
|
|
| for step in range(1, max_steps + 1): |
| for g in opt.param_groups: |
| g["lr"] = lr_at(step) |
| opt.zero_grad(set_to_none=True) |
| loss_val = 0.0 |
| last = {} |
| for _ in range(grad_accum): |
| try: |
| bd = next(data_iter) |
| except StopIteration: |
| data_iter = iter(loader) |
| bd = next(data_iter) |
| x = bd["bytes_in"].to(device, non_blocking=True) |
| v = bd["valid"].to(device, non_blocking=True) |
| with torch.autocast(device_type="cuda", dtype=amp, enabled=device == "cuda"): |
| out = model(x, valid=v) |
| loss = out["loss"] / grad_accum |
| loss.backward() |
| loss_val += loss.item() |
| last = out |
| gnorm = torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip) |
| if torch.isfinite(gnorm): |
| opt.step() |
| else: |
| opt.zero_grad(set_to_none=True) |
| log(f"step {step}: non-finite grad; skipped") |
| running += loss_val |
|
|
| if step % log_every == 0: |
| dt = time.time() - t_win |
| bps = log_every * batch * grad_accum * seq_len / dt |
| t_win = time.time() |
| log(f"step {step}/{max_steps} loss {running/log_every:.4f} bpb {last['bpb'].item():.4f} " |
| f"lr {lr_at(step):.2e} gnorm {gnorm:.2f} {bps/1e3:.1f} kB/s") |
| running = 0.0 |
| if val_windows is not None and step % eval_every == 0: |
| vb = evaluate() |
| log(f"[eval] step {step} val_bpb {vb:.4f}") |
| if step % save_every == 0 or step == max_steps: |
| vb = evaluate() if val_windows is not None else -1.0 |
| save_checkpoint("latest", step, {"step": step, "val_bpb": vb, |
| "params_M": model.num_params() / 1e6}) |
|
|
| vb = evaluate() if val_windows is not None else -1.0 |
| save_checkpoint("final", max_steps, {"step": max_steps, "val_bpb": vb, |
| "params_M": model.num_params() / 1e6}) |
| log(f"DONE. final val_bpb={vb}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
