| """Text LM training from a world-model (iVideoGPT) init vs. controls. |
| |
| Every arm uses the same LlamaConfig as iVideoGPT's transformer, re-sized to the |
| GPT-2 vocab. Only the body (everything but embed_tokens / lm_head) differs: |
| |
| --init scratch HF default random init |
| --init ivgpt iVideoGPT body weights; embeddings + head freshly initialized |
| --init shuffled iVideoGPT body with each tensor's entries randomly permuted |
| (same per-tensor weight statistics, no learned structure) |
| --init rescaled iVideoGPT body with every matrix scaled to std 0.02 (the scratch |
| init scale); keeps learned structure, removes the ~2.5x larger scale |
| |
| --init sp shrink-and-perturb (Ash & Adams 2020): lambda * iVideoGPT + fresh init |
| (--sp_lambda); RMSNorm gains become lambda * g + (1 - lambda) |
| --init nobody no pretrained body at all: only the --extra_layers (baseline for them) |
| |
| --freeze body train only embed_tokens, lm_head and RMSNorm gains (+ LoRA / extra layers) |
| --lora_rank R LoRA on every body linear (frozen base); merged into the weights at save |
| --extra_layers K append K freshly initialized trainable layers on top of the body |
| |
| Launch: torchrun --nproc_per_node=N wp/train.py --init ivgpt --lr 1e-3 ... |
| """ |
|
|
| import argparse |
| import json |
| import math |
| import os |
| import time |
| from contextlib import nullcontext |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import torch.distributed as dist |
| from huggingface_hub import hf_hub_download |
| from safetensors.torch import load_file |
| from transformers import AutoTokenizer, LlamaConfig, LlamaForCausalLM |
|
|
| GPT2_EOT = 50256 |
| VOCAB = 50304 |
| IO_PREFIXES = ("model.embed_tokens.", "lm_head.") |
|
|
|
|
| def parse_args(): |
| p = argparse.ArgumentParser() |
| p.add_argument("--init", choices=["scratch", "ivgpt", "shuffled", "rescaled", "sp", "nobody"], required=True) |
| p.add_argument("--freeze", choices=["none", "body"], default="none") |
| p.add_argument("--repo", default="thuml/ivideogpt-oxe-64-act-free-medium") |
| p.add_argument("--data_dir", default="/data/fineweb-edu-gpt2") |
| p.add_argument("--out_dir", required=True) |
| p.add_argument("--ckpt_dir", default=None, help="rolling checkpoint location (default: out_dir)") |
| p.add_argument("--final_dir", default=None, help="save final bf16 HF model here") |
| p.add_argument("--tokens", type=float, default=2e9, help="training token budget") |
| p.add_argument("--batch_tokens", type=int, default=524288, help="global tokens per step") |
| p.add_argument("--seq_len", type=int, default=1024) |
| p.add_argument("--micro_bs", type=int, default=32) |
| p.add_argument("--lr", type=float, default=1e-3) |
| p.add_argument("--min_lr_frac", type=float, default=0.1) |
| p.add_argument("--warmup_steps", type=int, default=200, help="-1: 5%% of steps, capped at 200") |
| p.add_argument("--wd", type=float, default=0.1) |
| p.add_argument("--grad_clip", type=float, default=1.0) |
| p.add_argument("--eval_every", type=int, default=100) |
| p.add_argument("--val_tokens", type=int, default=4_000_000) |
| p.add_argument("--ckpt_every", type=int, default=500) |
| p.add_argument("--seed", type=int, default=0) |
| p.add_argument("--no_compile", action="store_true") |
| p.add_argument("--sp_lambda", type=float, default=0.5) |
| p.add_argument("--sp_only", choices=["all", "attn", "mlp", "lower", "upper"], default="all", |
| help="apply S&P only to these weights; the rest keep the scratch init") |
| p.add_argument("--sp_shuffle", action="store_true", help="shrink-and-perturb the shuffled body (statistics-matched control)") |
| p.add_argument("--lora_rank", type=int, default=0) |
| p.add_argument("--extra_layers", type=int, default=0) |
| return p.parse_args() |
|
|
|
|
| def build_model(args): |
| cfg = LlamaConfig.from_pretrained(args.repo, subfolder="transformer") |
| n_body = 0 if args.init == "nobody" else cfg.num_hidden_layers |
| cfg.num_hidden_layers = n_body + args.extra_layers |
| cfg.vocab_size = VOCAB |
| cfg.bos_token_id = cfg.eos_token_id = GPT2_EOT |
| cfg.attention_dropout = 0.0 |
| cfg.max_position_embeddings = args.seq_len |
| cfg.tie_word_embeddings = False |
| torch.manual_seed(args.seed) |
| model = LlamaForCausalLM(cfg) |
| model.n_body = n_body |
| if args.init in ("scratch", "nobody"): |
| return model |
|
|
| sd = load_file(hf_hub_download(args.repo, "transformer/model.safetensors")) |
| |
| sd = {k.removeprefix("llm."): v for k, v in sd.items() if not k.startswith("action_linear.")} |
| body = {k: v for k, v in sd.items() if not k.startswith(IO_PREFIXES) and "rotary_emb" not in k} |
| if args.init == "shuffled" or (args.init == "sp" and args.sp_shuffle): |
| g = torch.Generator().manual_seed(args.seed + 1) |
| body = {k: v.flatten()[torch.randperm(v.numel(), generator=g)].view_as(v) for k, v in body.items()} |
| if args.init == "rescaled": |
| body = {k: v * (cfg.initializer_range / v.std()) if v.dim() == 2 else v for k, v in body.items()} |
| if args.init == "sp": |
| fresh = model.state_dict() |
| lam = args.sp_lambda |
| def selected(k): |
| layer = int(k.split(".")[2]) if k.startswith("model.layers.") else -1 |
| return {"all": True, "attn": "self_attn" in k or "input_layernorm" in k, |
| "mlp": ".mlp." in k or "post_attention_layernorm" in k, |
| "lower": 0 <= layer < cfg.num_hidden_layers // 2, |
| "upper": layer >= cfg.num_hidden_layers // 2}[args.sp_only] |
| body = {k: (lam * v + (fresh[k] if v.dim() == 2 else (1 - lam))) if selected(k) else fresh[k] |
| for k, v in body.items()} |
| missing, unexpected = model.load_state_dict(body, strict=False) |
| assert not unexpected, f"unexpected keys: {unexpected}" |
| is_extra = lambda k: k.startswith("model.layers.") and int(k.split(".")[2]) >= n_body |
| assert all(k.startswith(IO_PREFIXES) or is_extra(k) for k in missing), f"missing body keys: {missing}" |
| return model |
|
|
|
|
| class LoRALinear(torch.nn.Module): |
| """y = W x + (B A x) * alpha/r with W frozen; B starts at zero so the init function is unchanged.""" |
|
|
| def __init__(self, base, r, alpha): |
| super().__init__() |
| self.base = base |
| self.lora_A = torch.nn.Parameter(torch.empty(r, base.in_features)) |
| torch.nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) |
| self.lora_B = torch.nn.Parameter(torch.zeros(base.out_features, r)) |
| self.scale = alpha / r |
|
|
| def forward(self, x): |
| return self.base(x) + (x @ self.lora_A.t()) @ self.lora_B.t() * self.scale |
|
|
| def merged(self): |
| lin = self.base |
| lin.weight.data += (self.lora_B @ self.lora_A).to(lin.weight.dtype) * self.scale |
| return lin |
|
|
|
|
| LORA_TARGETS = {"self_attn": ("q_proj", "k_proj", "v_proj", "o_proj"), "mlp": ("gate_proj", "up_proj", "down_proj")} |
|
|
|
|
| def add_lora(model, r): |
| for layer in model.model.layers[: model.n_body]: |
| for parent, names in LORA_TARGETS.items(): |
| mod = getattr(layer, parent) |
| for n in names: |
| setattr(mod, n, LoRALinear(getattr(mod, n), r, alpha=r)) |
|
|
|
|
| def merge_lora(model): |
| for layer in model.model.layers: |
| for parent, names in LORA_TARGETS.items(): |
| mod = getattr(layer, parent) |
| for n in names: |
| if isinstance(getattr(mod, n), LoRALinear): |
| setattr(mod, n, getattr(mod, n).merged()) |
|
|
|
|
| def apply_freeze(model, freeze): |
| if freeze == "none": |
| return |
| for name, p in model.named_parameters(): |
| extra = name.startswith("model.layers.") and int(name.split(".")[2]) >= model.n_body |
| p.requires_grad = name.startswith(IO_PREFIXES) or "norm" in name or "lora_" in name or extra |
|
|
|
|
| class Batches: |
| """Random seq_len+1 windows from a uint16 token memmap; same order for every arm.""" |
|
|
| def __init__(self, path, seq_len, micro_bs, seed, rank): |
| self.data = np.memmap(path, dtype=np.uint16, mode="r") |
| self.seq_len, self.micro_bs = seq_len, micro_bs |
| self.rng = np.random.default_rng([seed, rank]) |
|
|
| def next(self, device): |
| idx = self.rng.integers(0, len(self.data) - self.seq_len - 1, self.micro_bs) |
| x = np.stack([self.data[i : i + self.seq_len + 1] for i in idx]).astype(np.int64) |
| x = torch.from_numpy(x).pin_memory().to(device, non_blocking=True) |
| return x[:, :-1], x[:, 1:] |
|
|
|
|
| @torch.no_grad() |
| def evaluate(model, val, args, rank, world, device): |
| """Mean loss over a fixed set of non-overlapping windows from the start of val.bin.""" |
| model.eval() |
| n_seq = args.val_tokens // args.seq_len |
| starts = np.arange(n_seq)[rank::world] * args.seq_len |
| tot, cnt = torch.zeros((), device=device), torch.zeros((), device=device) |
| for i in range(0, len(starts), args.micro_bs): |
| x = np.stack([val[s : s + args.seq_len + 1] for s in starts[i : i + args.micro_bs]]).astype(np.int64) |
| x = torch.from_numpy(x).to(device) |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| logits = model(input_ids=x[:, :-1]).logits |
| tot += torch.nn.functional.cross_entropy( |
| logits.float().flatten(0, 1), x[:, 1:].flatten(), reduction="sum" |
| ) |
| cnt += x[:, 1:].numel() |
| if world > 1: |
| dist.all_reduce(tot) |
| dist.all_reduce(cnt) |
| model.train() |
| return (tot / cnt).item() |
|
|
|
|
| def lr_at(step, total, args): |
| if step < args.warmup_steps: |
| return args.lr * (step + 1) / args.warmup_steps |
| t = (step - args.warmup_steps) / max(1, total - args.warmup_steps) |
| return args.lr * (args.min_lr_frac + (1 - args.min_lr_frac) * 0.5 * (1 + math.cos(math.pi * t))) |
|
|
|
|
| def main(): |
| args = parse_args() |
| ddp = "RANK" in os.environ |
| if ddp: |
| dist.init_process_group("nccl") |
| rank, world = dist.get_rank(), dist.get_world_size() |
| local = int(os.environ["LOCAL_RANK"]) |
| else: |
| rank, world, local = 0, 1, 0 |
| device = torch.device("cuda", local) |
| torch.cuda.set_device(device) |
| torch.backends.cuda.matmul.allow_tf32 = True |
| torch.backends.cudnn.allow_tf32 = True |
| master = rank == 0 |
| out = Path(args.out_dir) |
| if master: |
| out.mkdir(parents=True, exist_ok=True) |
|
|
| tok_per_micro = args.micro_bs * args.seq_len * world |
| assert args.batch_tokens % tok_per_micro == 0, "batch_tokens must divide micro_bs*seq_len*world" |
| accum = args.batch_tokens // tok_per_micro |
| total_steps = int(args.tokens // args.batch_tokens) |
| if args.warmup_steps < 0: |
| args.warmup_steps = min(200, max(10, int(0.05 * total_steps))) |
|
|
| model = build_model(args) |
| if args.lora_rank: |
| add_lora(model, args.lora_rank) |
| model = model.to(device) |
| apply_freeze(model, args.freeze) |
| params = [p for p in model.parameters() if p.requires_grad] |
| n_total = sum(p.numel() for p in model.parameters()) |
| n_train = sum(p.numel() for p in params) |
| decay = [p for p in params if p.dim() >= 2] |
| no_decay = [p for p in params if p.dim() < 2] |
| opt = torch.optim.AdamW( |
| [{"params": decay, "weight_decay": args.wd}, {"params": no_decay, "weight_decay": 0.0}], |
| lr=args.lr, betas=(0.9, 0.95), eps=1e-8, fused=True, |
| ) |
|
|
| step = 0 |
| ckpt_path = Path(args.ckpt_dir or out) / "ckpt.pt" |
| if master: |
| ckpt_path.parent.mkdir(parents=True, exist_ok=True) |
| if ckpt_path.exists(): |
| ck = torch.load(ckpt_path, map_location=device) |
| model.load_state_dict(ck["model"]) |
| opt.load_state_dict(ck["opt"]) |
| step = ck["step"] |
| if master: |
| print(f"resumed from step {step}") |
|
|
| fwd = model if args.no_compile else torch.compile(model) |
| if ddp: |
| fwd = torch.nn.parallel.DistributedDataParallel(fwd, device_ids=[local]) |
|
|
| train = Batches(Path(args.data_dir) / "train.bin", args.seq_len, args.micro_bs, args.seed, rank) |
| for _ in range(step * accum): |
| train.rng.integers(0, len(train.data) - args.seq_len - 1, args.micro_bs) |
| val = np.memmap(Path(args.data_dir) / "val.bin", dtype=np.uint16, mode="r") |
|
|
| def log(rec): |
| if master: |
| print(json.dumps(rec), flush=True) |
| with open(out / "metrics.jsonl", "a") as f: |
| f.write(json.dumps(rec) + "\n") |
|
|
| if master and step == 0: |
| (out / "args.json").write_text(json.dumps(vars(args) | {"world": world, "accum": accum}, indent=2)) |
| log({"step": 0, "tokens": 0, "val_loss": evaluate(model, val, args, rank, world, device), |
| "n_params": n_total, "n_trainable": n_train}) |
| elif step == 0: |
| evaluate(model, val, args, rank, world, device) |
|
|
| t0 = time.time() |
| while step < total_steps: |
| lr = lr_at(step, total_steps, args) |
| for g in opt.param_groups: |
| g["lr"] = lr |
| loss_acc = 0.0 |
| for micro in range(accum): |
| x, y = train.next(device) |
| sync = nullcontext() if (not ddp or micro == accum - 1) else fwd.no_sync() |
| with sync, torch.autocast("cuda", dtype=torch.bfloat16): |
| logits = fwd(input_ids=x).logits |
| loss = torch.nn.functional.cross_entropy(logits.float().flatten(0, 1), y.flatten()) |
| (loss / accum).backward() |
| loss_acc += loss.item() / accum |
| gnorm = torch.nn.utils.clip_grad_norm_(params, args.grad_clip).item() |
| opt.step() |
| opt.zero_grad(set_to_none=True) |
| step += 1 |
|
|
| if step % 10 == 0 and master: |
| dt = time.time() - t0 |
| t0 = time.time() |
| log({"step": step, "tokens": step * args.batch_tokens, "train_loss": round(loss_acc, 4), |
| "lr": lr, "gnorm": round(gnorm, 3), "tok_per_s": int(10 * args.batch_tokens / dt)}) |
| if step % args.eval_every == 0 or step == total_steps: |
| vl = evaluate(model, val, args, rank, world, device) |
| log({"step": step, "tokens": step * args.batch_tokens, "val_loss": vl}) |
| if master and (step % args.ckpt_every == 0 or step == total_steps): |
| torch.save({"model": model.state_dict(), "opt": opt.state_dict(), "step": step}, ckpt_path.with_suffix(".tmp")) |
| os.replace(ckpt_path.with_suffix(".tmp"), ckpt_path) |
|
|
| if master: |
| if args.final_dir: |
| merge_lora(model) |
| model.to(torch.bfloat16).save_pretrained(args.final_dir) |
| AutoTokenizer.from_pretrained("gpt2").save_pretrained(args.final_dir) |
| (out / "DONE").touch() |
| if ddp: |
| dist.destroy_process_group() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
