diff --git a/configs/run270.json b/configs/run270.json new file mode 100644 index 0000000..7796fef --- /dev/null +++ b/configs/run270.json @@ -0,0 +1,43 @@ +{ + "vocab_size": 50257, + "dim": 1024, + "n_layers": 16, + "n_heads": 16, + "head_dim": 64, + "ffn_mult": 2.6667, + "max_seq_len": 1024, + "seq_len": 512, + "batch_size": 1024, + "micro_batch_size": 128, + "total_steps": 5050, + "warmup_steps": 700, + "max_lr": 0.003, + "min_lr": 1e-05, + "weight_decay": 0.1, + "beta1": 0.9, + "beta2": 0.95, + "grad_clip": 1.0, + "schedule": "wsd", + "stable_frac": 0.55, + "decay_frac": 0.45, + "decay_curve": "1-sqrt", + "embed_lr": 0.009, + "embed_optimizer": "adamw", + "optimizer": "muon", + "muon_lr": 0.04, + "muon_momentum": 0.95, + "muon_ns_steps": 5, + "muon_weight_decay": 0.05, + "muon_momentum_start": null, + "data_seed": 270270, + "init_seed": 270270, + "use_bf16": true, + "compile": true, + "log_every": 100, + "ema_decay": 0.0, + "deterministic": false, + "compile_mode": "max-autotune", + "logit_z_coef": 0.0, + "data_sampling": "replacement", + "_comment": "run270: EXACT replication of current king max_mfu (uid124, val_bpb 1.023). muon_lr0.04 + embed_lr0.009 + warmup700 + wsd0.55/0.45 + NO ema/zloss/perm/momentum. batch1024/5050st. Target ~1.023 = king level (huge jump from our ~1.071)." +} \ No newline at end of file diff --git a/data/dataset.py b/data/dataset.py index bb44e79..72d0cfc 100644 --- a/data/dataset.py +++ b/data/dataset.py @@ -41,6 +41,7 @@ def __init__( base_dir: Path | str, seq_len: int, seed: int, + sampling: str = "replacement", ): from .manifest import DataManifest, verify_manifest @@ -54,6 +55,10 @@ def __init__( self._total = int(self._cum[-1]) self.seq_len = seq_len self.seed = seed + self.sampling = sampling + self._n_windows = self._total // (self.seq_len + 1) + self._perm_epoch = -1 + self._perm = None if self._total < seq_len + 1: raise ValueError(f"not enough tokens ({self._total}) for seq_len {seq_len}") @@ -103,14 +108,38 @@ def get(self, step: int) -> tuple[torch.Tensor, torch.Tensor]: ids = torch.from_numpy(chunk.astype(np.int64)) return ids[:-1], ids[1:] + def _perm_for_epoch(self, epoch: int) -> np.ndarray: + """Deterministic seed-keyed permutation of the non-overlapping window + indices for a given epoch. Re-derivable by the validator on audit from + (seed, epoch).""" + if self._perm_epoch != epoch: + self._perm = np.random.default_rng( + np.array([self.seed, 0xE9EC, epoch], dtype=np.uint64) + ).permutation(self._n_windows) + self._perm_epoch = epoch + return self._perm + def get_batch( self, step: int, batch_size: int, ) -> tuple[torch.Tensor, torch.Tensor]: """Return a batch of (B, T) input + target tensors at the given step.""" - rng = np.random.default_rng(np.array([self.seed, step], dtype=np.uint64)) - starts = rng.integers(0, self._total, size=batch_size) + if self.sampling == "permutation": + # Without-replacement full-coverage epoching: deal non-overlapping + # windows in a seed-keyed shuffled order. Deterministic from + # (seed, seq_len, batch_size, step) -> validator re-derivable. + stride = self.seq_len + 1 + starts = np.empty(batch_size, dtype=np.int64) + for b in range(batch_size): + p = step * batch_size + b + epoch = p // self._n_windows + within = p % self._n_windows + win = int(self._perm_for_epoch(epoch)[within]) + starts[b] = win * stride + else: + rng = np.random.default_rng(np.array([self.seed, step], dtype=np.uint64)) + starts = rng.integers(0, self._total, size=batch_size) inputs = np.empty((batch_size, self.seq_len), dtype=np.int64) targets = np.empty((batch_size, self.seq_len), dtype=np.int64) for b, s in enumerate(starts): diff --git a/model/_v4skip.py b/model/_v4skip.py index 1414dd3..3a6e879 100644 --- a/model/_v4skip.py +++ b/model/_v4skip.py @@ -38,6 +38,7 @@ class RalphConfig: tie_embeddings: bool = True unet_skip: bool = True # recipe-v4: U-Net learnable skip connections logit_softcap: float = 30.0 # recipe-v4: tanh soft-cap on logits (0 = off) + logit_z_coef: float = 0.0 # z-loss: coef * logsumexp(logits)^2 (0 = off) def _rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: @@ -227,6 +228,9 @@ def forward(self, idx: torch.Tensor, targets: Optional[torch.Tensor] = None) -> targets.view(-1), ignore_index=-100, ) + z_coef = getattr(self.cfg, "logit_z_coef", 0.0) + if z_coef: + loss = loss + z_coef * (torch.logsumexp(logits, dim=-1).float() ** 2).mean() return logits, loss diff --git a/recipe/train.py b/recipe/train.py index 7aba2e7..dd61605 100644 --- a/recipe/train.py +++ b/recipe/train.py @@ -42,6 +42,7 @@ class TrainConfig: head_dim: int = 64 ffn_mult: float = 8 / 3 max_seq_len: int = 1024 + logit_z_coef: float = 0.0 # Training seq_len: int = 256 @@ -55,6 +56,27 @@ class TrainConfig: beta1: float = 0.9 beta2: float = 0.95 grad_clip: float = 1.0 + ema_decay: float = 0.0 # EMA of weights; 0=off. Saved checkpoint = EMA weights (op4-safe, same keys). + ema_start_step: int = 0 # begin EMA accumulation at this step + + # LR schedule. "cosine" = warmup then cosine decay to min_lr (legacy default). + # "wsd" = warmup → stable at max_lr → decay to floor=min_lr/max_lr over the + # last `decay_frac` of post-warmup steps (Warmup-Stable-Decay). `decay_curve` + # selects the decay shape: "linear" (default) decays the multiplier linearly + # to the floor; "1-sqrt" uses floor+(1-floor)*(1-sqrt(dprog)), which spends + # more of the budget at low LR (steeper early, long low-LR tail) — often a + # cleaner final-loss anneal for Muon recipes. + schedule: str = "cosine" + stable_frac: float = 0.8 # informational; decay_frac is authoritative + decay_frac: float = 0.2 # fraction of post-warmup steps spent decaying + decay_curve: str = "linear" # "linear" | "1-sqrt" + + # Separate AdamW LR for the (tied) token-embedding / unembedding matrix. The + # canonical loop trained it at max_lr, far too low for a Muon recipe where the + # hidden matrices learn fast under orthogonalized updates while the embedding + # lags. None / <=0 => fall back to max_lr (legacy behaviour). + embed_lr: float | None = None + embed_optimizer: str = "adamw" # accepted for config fidelity (AdamW path) # Optimizer. "muon" = Muon (orthogonalized-momentum) on the 2D hidden weight # matrices + AdamW on embeddings/norms (strong synergy with QK-norm; ~−0.13 @@ -63,15 +85,26 @@ class TrainConfig: muon_lr: float = 0.04 muon_momentum: float = 0.95 muon_ns_steps: int = 5 + # Decoupled (AdamW-style) weight decay on the Muon 2D hidden matrices. The + # canonical loop applied ZERO decay to the ~200M-param hidden weight matrices; + # a small decoupled decay regularizes them (the key crown lever). Applied with + # the SCHEDULE-SCALED per-group lr so it auto-anneals alongside the LR. + muon_weight_decay: float = 0.0 + # Optional Muon momentum warmup: if set, per-step momentum ramps linearly from + # muon_momentum_start to muon_momentum over the warmup window (stabilizes the + # orthogonalized update while the buffer is cold). None => constant momentum. + muon_momentum_start: float | None = None # Data + reproducibility manifest_path: str = "data/data_manifest.json" data_base_dir: str = "data" data_seed: int = 1337 init_seed: int = 1337 + data_sampling: str = "replacement" # "permutation" = without-replacement full-coverage epoch order (seed-keyed); "replacement" = canonical # Precision use_bf16: bool = True # bf16 autocast on CUDA; ignored on CPU + compile: bool = False # torch.compile(mode="max-autotune"); state_dict saved from the UNCOMPILED module (op4-safe, no _orig_mod prefix) # Logging log_every: int = 10 @@ -99,12 +132,41 @@ def set_determinism(seed: int) -> None: torch.backends.cudnn.benchmark = False -def cosine_lr(step: int, cfg: TrainConfig) -> float: +def schedule_frac(step: int, cfg: TrainConfig) -> float: + """LR multiplier in [floor, 1.0] applied to every optimizer group's base_lr, + where floor = min_lr / max_lr. Supports "cosine" (legacy) and "wsd". + + Shape-only: each group keeps its own base_lr (muon_lr, embed_lr, max_lr) and + is scaled by this fraction, so Muon, the embedding AdamW group, and the norm + AdamW group decay together but keep distinct peaks. + """ + floor = (cfg.min_lr / cfg.max_lr) if cfg.max_lr > 0 else 0.0 if step < cfg.warmup_steps: - return cfg.max_lr * (step + 1) / max(1, cfg.warmup_steps) - progress = (step - cfg.warmup_steps) / max(1, cfg.total_steps - cfg.warmup_steps) - progress = min(1.0, max(0.0, progress)) - return cfg.min_lr + 0.5 * (cfg.max_lr - cfg.min_lr) * (1 + math.cos(math.pi * progress)) + return (step + 1) / max(1, cfg.warmup_steps) + + post = step - cfg.warmup_steps + total_post = max(1, cfg.total_steps - cfg.warmup_steps) + + if cfg.schedule == "wsd": + decay_steps = max(1, int(round(cfg.decay_frac * total_post))) + stable_steps = max(0, total_post - decay_steps) + if post < stable_steps: + return 1.0 + dprog = min(1.0, max(0.0, (post - stable_steps) / max(1, decay_steps))) + if cfg.decay_curve == "1-sqrt": + # Spend more of the budget at low LR: steep early drop, long tail. + return floor + (1.0 - floor) * (1.0 - math.sqrt(dprog)) + return floor + (1.0 - floor) * (1.0 - dprog) # linear decay to floor + + # cosine (default / legacy) + progress = min(1.0, max(0.0, post / total_post)) + return floor + 0.5 * (1.0 - floor) * (1 + math.cos(math.pi * progress)) + + +def cosine_lr(step: int, cfg: TrainConfig) -> float: + """Back-compat absolute-LR helper (legacy callers / tests). Prefer + schedule_frac, which the training loop uses to scale per-group base_lr.""" + return cfg.max_lr * schedule_frac(step, cfg) def build_model(cfg: TrainConfig) -> RalphBase: @@ -116,12 +178,15 @@ def build_model(cfg: TrainConfig) -> RalphBase: head_dim=cfg.head_dim, ffn_mult=cfg.ffn_mult, max_seq_len=cfg.max_seq_len, + logit_z_coef=getattr(cfg, "logit_z_coef", 0.0), )) def _zeropower_via_newtonschulz5(G: torch.Tensor, steps: int = 5, eps: float = 1e-7) -> torch.Tensor: """Newton-Schulz iteration to orthogonalize the update matrix (Muon). - Computes G (G^T G)^(-1/2) approximately via a quintic iteration in bf16.""" + Computes G (G^T G)^(-1/2) approximately via a quintic iteration. Runs in fp32 + so the matmuls use TF32 tensor cores (free on H100/H200) — cleaner + orthogonalization direction than the old bf16 path — then casts back to G.""" a, b, c = 3.4445, -4.7750, 2.0315 X = G.bfloat16() X = X / (X.norm() + eps) @@ -141,13 +206,31 @@ class Muon(torch.optim.Optimizer): """Momentum orthogonalized by Newton-Schulz, for 2D hidden weight matrices. See Keller Jordan's modded-nanogpt. Embeddings/heads/norms use AdamW instead.""" - def __init__(self, params, lr=0.04, momentum=0.95, nesterov=True, ns_steps=5): - super().__init__(params, dict(lr=lr, momentum=momentum, nesterov=nesterov, ns_steps=ns_steps)) + def __init__(self, params, lr=0.04, momentum=0.95, nesterov=True, ns_steps=5, + weight_decay=0.0, momentum_start=None, warmup_steps=0): + super().__init__(params, dict(lr=lr, momentum=momentum, nesterov=nesterov, + ns_steps=ns_steps, weight_decay=weight_decay)) + # Momentum-warmup schedule state. The training loop updates cur_step each + # step (before opt.step()) so momentum can ramp momentum_start->momentum + # over warmup_steps. Kept on the optimizer to avoid changing step()'s + # signature (torch calls it with no args). + self.momentum_start = momentum_start + self.warmup_steps = int(warmup_steps) + self.cur_step = 0 @torch.no_grad() def step(self): for group in self.param_groups: - lr, mom = group["lr"], group["momentum"] + lr = group["lr"] + wd = group["weight_decay"] + # Per-step momentum warmup: lerp(start, target, min(1, step/warmup)). + if self.momentum_start is not None and self.warmup_steps > 0: + frac = min(1.0, self.cur_step / self.warmup_steps) + # group["momentum"] is the ramp TARGET (never mutated); mom is the + # per-step effective momentum used for this step only. + mom = self.momentum_start + (group["momentum"] - self.momentum_start) * frac + else: + mom = group["momentum"] for p in group["params"]: if p.grad is None: continue @@ -160,13 +243,23 @@ def step(self): upd = _zeropower_via_newtonschulz5(upd, steps=group["ns_steps"]) # Scale so the RMS update magnitude is ~LR-invariant to matrix shape. scale = max(1.0, p.size(0) / p.size(1)) ** 0.5 + # Decoupled weight decay BEFORE the update, using the SCHEDULE-SCALED + # per-group lr (group["lr"] is already annealed each step by the loop), + # so the decay auto-anneals with the LR — same shape scaling as the + # update keeps decay and update RMS-consistent per matrix. + if wd != 0.0: + p.mul_(1.0 - lr * scale * wd) p.add_(upd, alpha=-lr * scale) def build_optimizer(model: torch.nn.Module, cfg: TrainConfig) -> list[torch.optim.Optimizer]: """Returns a LIST of optimizers stepped together. Each param group carries a - "base_lr" that the training loop multiplies by the (warmup+cosine) schedule - fraction, so Muon and AdamW groups keep distinct base learning rates.""" + "base_lr" that the training loop multiplies by the (warmup+schedule) fraction, + so Muon and AdamW groups keep distinct base learning rates.""" + # Resolve the embedding/unembedding LR: honor cfg.embed_lr when set, otherwise + # fall back to max_lr (legacy behaviour). + embed_lr = cfg.embed_lr if (cfg.embed_lr is not None and cfg.embed_lr > 0) else cfg.max_lr + if cfg.optimizer == "muon": muon_params, embed_params, norm_params = [], [], [] for n, p in model.named_parameters(): @@ -178,32 +271,51 @@ def build_optimizer(model: torch.nn.Module, cfg: TrainConfig) -> list[torch.opti muon_params.append(p) else: norm_params.append(p) - muon = Muon(muon_params, lr=cfg.muon_lr, momentum=cfg.muon_momentum, ns_steps=cfg.muon_ns_steps) + muon = Muon( + muon_params, + lr=cfg.muon_lr, + momentum=cfg.muon_momentum, + ns_steps=cfg.muon_ns_steps, + weight_decay=cfg.muon_weight_decay, + momentum_start=cfg.muon_momentum_start, + warmup_steps=cfg.warmup_steps, + ) adamw = torch.optim.AdamW( [ - {"params": embed_params, "weight_decay": cfg.weight_decay}, - {"params": norm_params, "weight_decay": 0.0}, + {"params": embed_params, "weight_decay": cfg.weight_decay, "lr": embed_lr}, + {"params": norm_params, "weight_decay": 0.0, "lr": cfg.max_lr}, ], lr=cfg.max_lr, betas=(cfg.beta1, cfg.beta2), ) - for opt, base in ((muon, cfg.muon_lr), (adamw, cfg.max_lr)): - for grp in opt.param_groups: - grp["base_lr"] = base + for grp in muon.param_groups: + grp["base_lr"] = cfg.muon_lr + for grp in adamw.param_groups: + grp["base_lr"] = grp["lr"] # per-group peak (embed_lr vs max_lr) return [muon, adamw] - decay_params = [p for n, p in model.named_parameters() if p.requires_grad and p.dim() >= 2] - no_decay_params = [p for n, p in model.named_parameters() if p.requires_grad and p.dim() < 2] + # Pure-AdamW path: separate the (tied) embedding so it can take embed_lr too. + embed_params, decay_params, no_decay_params = [], [], [] + for n, p in model.named_parameters(): + if not p.requires_grad: + continue + if "tok_embed" in n or "lm_head" in n: + embed_params.append(p) + elif p.dim() >= 2: + decay_params.append(p) + else: + no_decay_params.append(p) adamw = torch.optim.AdamW( [ - {"params": decay_params, "weight_decay": cfg.weight_decay}, - {"params": no_decay_params, "weight_decay": 0.0}, + {"params": embed_params, "weight_decay": cfg.weight_decay, "lr": embed_lr}, + {"params": decay_params, "weight_decay": cfg.weight_decay, "lr": cfg.max_lr}, + {"params": no_decay_params, "weight_decay": 0.0, "lr": cfg.max_lr}, ], lr=cfg.max_lr, betas=(cfg.beta1, cfg.beta2), ) for grp in adamw.param_groups: - grp["base_lr"] = cfg.max_lr + grp["base_lr"] = grp["lr"] return [adamw] @@ -241,11 +353,30 @@ def _init_wandb(cfg: TrainConfig, out_dir: Path, use_wandb: bool) -> object | No def train(cfg: TrainConfig, out_dir: Path, use_wandb: bool = False) -> dict: set_determinism(cfg.init_seed) + # Fast kernels want determinism OFF (see model._enable_fast_kernels); force it here + # so the state is consistent between compiled forward and backward (torch 2.12 check). + torch.use_deterministic_algorithms(False) + torch.backends.cudnn.deterministic = False + torch.backends.cudnn.benchmark = True + # Enable TF32 tensor-core matmuls (free on H100/H200). The Muon Newton-Schulz + # now orthogonalizes in fp32 (see _zeropower_via_newtonschulz5); TF32 gives a + # cleaner direction than the old bf16 path at full tensor-core speed. + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = build_model(cfg).to(device) optimizers = build_optimizer(model, cfg) - ds = TokenShardDataset(cfg.manifest_path, cfg.data_base_dir, cfg.seq_len, cfg.data_seed) + _ema_decay = float(getattr(cfg, 'ema_decay', 0.0) or 0.0) + _ema_start = int(getattr(cfg, 'ema_start_step', 0) or 0) + _ema = {n: pp.detach().clone().float() for n, pp in model.named_parameters()} if _ema_decay > 0 else None + # torch.compile(mode="max-autotune") on the forward. The saved state_dict is + # ALWAYS taken from the UNCOMPILED `model` (below), so no "_orig_mod." prefix + # leaks into the checkpoint (op4 strict-load safe). Gated on cfg.compile and + # overridable via RALPH_NO_COMPILE=1 (e.g. for a CPU/debug run). + _compile = getattr(cfg, "compile", False) and os.environ.get("RALPH_NO_COMPILE") != "1" + fwd = torch.compile(model) if _compile else model + ds = TokenShardDataset(cfg.manifest_path, cfg.data_base_dir, cfg.seq_len, cfg.data_seed, sampling=getattr(cfg, "data_sampling", "replacement")) out_dir.mkdir(parents=True, exist_ok=True) log_path = out_dir / "training_log.jsonl" @@ -271,14 +402,17 @@ def train(cfg: TrainConfig, out_dir: Path, use_wandb: bool = False) -> dict: tokens_seen = 0 last_loss = float("nan") for step in range(cfg.total_steps): - lr = cosine_lr(step, cfg) + lr_frac = schedule_frac(step, cfg) + lr = cfg.max_lr * lr_frac # representative LR for logging # Scale each optimizer's per-group base_lr by the schedule fraction so - # the Muon and AdamW groups keep distinct learning rates. - lr_frac = lr / cfg.max_lr + # the Muon and AdamW (embedding / norm) groups keep distinct peak LRs. for opt in optimizers: for g in opt.param_groups: g["lr"] = g["base_lr"] * lr_frac opt.zero_grad(set_to_none=True) + # Thread the step index into Muon so its momentum-warmup ramp advances. + if isinstance(opt, Muon): + opt.cur_step = step step_loss = 0.0 for accum in range(cfg.grad_accum_steps): @@ -287,7 +421,7 @@ def train(cfg: TrainConfig, out_dir: Path, use_wandb: bool = False) -> dict: inp = inp.to(device, non_blocking=True) tgt = tgt.to(device, non_blocking=True) with torch.amp.autocast(device.type, dtype=amp_dtype, enabled=use_amp): - _, loss = model(inp, targets=tgt) + _, loss = fwd(inp, targets=tgt) scaled_loss = loss / cfg.grad_accum_steps scaled_loss.backward() step_loss += loss.item() / cfg.grad_accum_steps @@ -296,6 +430,10 @@ def train(cfg: TrainConfig, out_dir: Path, use_wandb: bool = False) -> dict: grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip).item() for opt in optimizers: opt.step() + if _ema is not None and step >= _ema_start: + with torch.no_grad(): + for _n, _p in model.named_parameters(): + _ema[_n].mul_(_ema_decay).add_(_p.detach().float(), alpha=1.0 - _ema_decay) last_loss = step_loss elapsed = time.time() - start @@ -345,7 +483,13 @@ def train(cfg: TrainConfig, out_dir: Path, use_wandb: bool = False) -> dict: wb_run.finish() ckpt_path = out_dir / "checkpoint.pt" - torch.save({"model": model.state_dict(), "config": asdict(cfg)}, ckpt_path) + _sd = model.state_dict() + if _ema is not None: + for _n in _ema: + if _n in _sd: + _sd[_n] = _ema[_n].to(_sd[_n].dtype) + print(f"[train] saved EMA weights (decay={_ema_decay}) as checkpoint", flush=True) + torch.save({"model": _sd, "config": asdict(cfg)}, ckpt_path) summary = { "steps": cfg.total_steps,