# --- nemo_automodel/recipes/llm/prism_entry.py (trainer) --- """PRISM AutoModel 2.0 entry: self-contained build_model + train. Modern decoder LM (RMSNorm, RoPE, SwiGLU, GQA, QK-norm, value-residual, U-net skips, tied-free head, logit soft-cap) trained from random init with Muon + AdamW (per-group LRs), a WSD schedule, and FlexAttention document masking. The harness resolves this module and calls build_model(ctx) then train(model, ctx). """ import math import os import sys import time os.environ.setdefault("TOKENIZERS_PARALLELISM", "true") import pyarrow.parquet as pq import torch import torch.nn as nn import torch.nn.functional as F from transformers import GPT2TokenizerFast try: from torch.nn.attention.flex_attention import flex_attention as _flex _flex = torch.compile(_flex) except Exception: _flex = None try: from torch.nn.attention.flex_attention import create_block_mask as _create_block_mask except Exception: _create_block_mask = None try: import prism_telemetry except ImportError: class _TelemetryFallback: @staticmethod def report(**_kwargs): return None @staticmethod def finish_evaluation(): return None prism_telemetry = _TelemetryFallback() VOCAB = 8192 # default; real value comes from the shipped tokenizer via ctx["vocab_size"] BLOCK = 1024 D_MODEL = 768 N_LAYER = 15 N_HEAD = 12 N_KV_HEAD = 4 FFN_MULT = 4 DROPOUT = 0.0 LOGIT_SOFTCAP = 15.0 def _ctx_get(ctx, key, default): if hasattr(ctx, "get"): try: v = ctx.get(key, default) return default if v is None else v except Exception: return default return getattr(ctx, key, default) class RMSNorm(nn.Module): def __init__(self, dim: int, eps: float = 1e-6): super().__init__() self.weight = nn.Parameter(torch.ones(dim)) self.eps = eps def forward(self, x): x32 = x.float() x32 = x32 * torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + self.eps) return (x32 * self.weight.float()).type_as(x) def _rope(head_dim: int, seq: int, base: float = 10000.0): inv = 1.0 / (base ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)) t = torch.arange(seq, dtype=torch.float32) f = torch.outer(t, inv) return torch.cos(f), torch.sin(f) def _apply_rope(x, cos, sin): # x: (b, h, t, hd). Interleaved even/odd rotation. xe, xo = x[..., 0::2], x[..., 1::2] cos = cos[None, None, :, :].to(x.dtype) sin = sin[None, None, :, :].to(x.dtype) re = xe * cos - xo * sin ro = xe * sin + xo * cos return torch.stack([re, ro], dim=-1).flatten(-2) class Attention(nn.Module): def __init__(self, d: int, n_head: int, n_kv: int): super().__init__() self.n_head, self.n_kv = n_head, n_kv self.hd = d // n_head self.groups = n_head // n_kv self.wq = nn.Linear(d, n_head * self.hd, bias=False) self.wk = nn.Linear(d, n_kv * self.hd, bias=False) self.wv = nn.Linear(d, n_kv * self.hd, bias=False) self.wo = nn.Linear(n_head * self.hd, d, bias=False) # QK-norm: normalize queries/keys per head for stable attention logits. self.q_norm = RMSNorm(self.hd) self.k_norm = RMSNorm(self.hd) # Value-residual: mix layer-0 values into this layer (sigmoid(lam) in (0,1), # init 0.5). Speeds convergence by giving every block direct access to # first-layer value features (ResFormer / value-residual learning). self.lam = nn.Parameter(torch.zeros(1)) def forward(self, x, cos, sin, v0=None, block_mask=None): b, t, _ = x.shape q = self.q_norm(self.wq(x).view(b, t, self.n_head, self.hd)).transpose(1, 2) k = self.k_norm(self.wk(x).view(b, t, self.n_kv, self.hd)).transpose(1, 2) v = self.wv(x).view(b, t, self.n_kv, self.hd) if v0 is None: v0_out = v else: lam = torch.sigmoid(self.lam) v = lam * v + (1.0 - lam) * v0 v0_out = v0 v = v.transpose(1, 2) q = _apply_rope(q, cos, sin) k = _apply_rope(k, cos, sin) k = k.to(q.dtype) v = v.to(q.dtype) if block_mask is not None and _flex is not None: # Fast document-masked attention (GQA handled natively). o = _flex(q, k, v, block_mask=block_mask, enable_gqa=True) else: if self.groups > 1: k = k.repeat_interleave(self.groups, dim=1) v = v.repeat_interleave(self.groups, dim=1) o = F.scaled_dot_product_attention(q, k, v, is_causal=True) o = o.transpose(1, 2).contiguous().view(b, t, self.n_head * self.hd) return self.wo(o), v0_out class SwiGLU(nn.Module): def __init__(self, d: int, hidden: int): super().__init__() self.w_gate = nn.Linear(d, hidden, bias=False) self.w_up = nn.Linear(d, hidden, bias=False) self.w_down = nn.Linear(hidden, d, bias=False) def forward(self, x): return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x)) class Block(nn.Module): def __init__(self, d, n_head, n_kv, ffn, dropout=0.0): super().__init__() self.n1 = RMSNorm(d) self.attn = Attention(d, n_head, n_kv) self.n2 = RMSNorm(d) self.mlp = SwiGLU(d, ffn) self.drop = nn.Dropout(dropout) def forward(self, x, cos, sin, v0=None, block_mask=None): a, v0_out = self.attn(self.n1(x), cos, sin, v0, block_mask) x = x + self.drop(a) return x + self.drop(self.mlp(self.n2(x))), v0_out class DecoderLM(nn.Module): def __init__(self, vocab, d, n_layer, n_head, n_kv, block, ffn_mult): super().__init__() self.block = block self.tok_emb = nn.Embedding(vocab, d) ffn = ((int(d * ffn_mult * 2 / 3) + 63) // 64) * 64 self.blocks = nn.ModuleList(Block(d, n_head, n_kv, ffn, DROPOUT) for _ in range(n_layer)) self.emb_drop = nn.Dropout(DROPOUT) self.skip_weights = nn.Parameter(torch.zeros(n_layer // 2)) self.ln_f = RMSNorm(d) self.head = nn.Linear(d, vocab, bias=False) # untied head (zero-init), Adam high-LR cos, sin = _rope(d // n_head, block) self.register_buffer("rope_cos", cos, persistent=False) self.register_buffer("rope_sin", sin, persistent=False) self.apply(self._init) # Scaled residual init for depth stability. for name, p in self.named_parameters(): if name.endswith("wo.weight") or name.endswith("w_down.weight"): nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * n_layer)) nn.init.zeros_(self.head.weight) def _init(self, m): if isinstance(m, nn.Linear): nn.init.normal_(m.weight, mean=0.0, std=0.02) elif isinstance(m, nn.Embedding): nn.init.normal_(m.weight, mean=0.0, std=0.02) def forward(self, ids, block_mask=None): t = min(ids.shape[1], self.block) ids = ids[:, -t:] if self.rope_cos.device != ids.device: self.rope_cos = self.rope_cos.to(ids.device) self.rope_sin = self.rope_sin.to(ids.device) cos, sin = self.rope_cos[:t], self.rope_sin[:t] x = self.emb_drop(self.tok_emb(ids)) v0 = None n = len(self.blocks) n_enc = n // 2 skips = [] for i, blk in enumerate(self.blocks): if i >= n - n_enc: x = x + self.skip_weights[i - (n - n_enc)] * skips.pop() x, v0 = blk(x, cos, sin, v0, block_mask) if i < n_enc: skips.append(x) x = self.ln_f(x) logits = self.head(x) return LOGIT_SOFTCAP * torch.tanh(logits / LOGIT_SOFTCAP) def _resolve_tokenizer(ctx): """Prefer the harness-resolved+validated tokenizer (ctx['tokenizer']); else load the shipped in-domain tokenizer from the local source tree. Never hits the network (the miner child runs under unshare --net).""" tok = _ctx_get(ctx, "tokenizer", None) if tok is not None: return tok import os as _os here = _os.path.dirname(_os.path.abspath(__file__)) cands = [ _ctx_get(ctx, "tokenizer_dir", None), _os.path.join(here, "tokenizer"), _os.path.join(here, "tok8k"), _os.path.join(_os.getcwd(), "tokenizer"), _os.path.join(_os.getcwd(), "submission", "tokenizer"), ] for d in cands: if d and _os.path.isfile(_os.path.join(d, "tokenizer.json")): try: from transformers import AutoTokenizer return AutoTokenizer.from_pretrained(d, local_files_only=True) except Exception: from transformers import PreTrainedTokenizerFast return PreTrainedTokenizerFast( tokenizer_file=_os.path.join(d, "tokenizer.json"), eos_token="<|eos|>", pad_token="<|eos|>") # Last-resort fallback: the pinned GPT-2 default (keeps local smoke tests # runnable even without shipped files; never reached in-pod). from transformers import GPT2TokenizerFast return GPT2TokenizerFast.from_pretrained("gpt2") def build_model(ctx): """Recipe contract entrypoint: return an nn.Module given the context.""" torch.manual_seed(int(_ctx_get(ctx, "seed", 0))) d = int(_ctx_get(ctx, "d_model", D_MODEL)) n_layer = int(_ctx_get(ctx, "n_layer", N_LAYER)) n_head = int(_ctx_get(ctx, "n_head", N_HEAD)) n_kv = int(_ctx_get(ctx, "n_kv_head", N_KV_HEAD)) block = int(_ctx_get(ctx, "block", BLOCK)) # Embedding/head are sized from the resolved (shipped) tokenizer, not a # hardcoded constant — the harness passes vocab_size for the tokenizer it # resolved from submission/tokenizer/. Fail loud on mismatch would be ideal, # but the harness guarantees this equals len(shipped tokenizer). vocab = int(_ctx_get(ctx, "vocab_size", VOCAB)) return DecoderLM(vocab, d, n_layer, n_head, n_kv, block, FFN_MULT) HARD_MAX_STEPS = 20000 HARD_MAX_SECONDS = 21600 MUON_LR = 0.036 # peak LR for the transformer weight matrices (Muon) # Embedding LR scales inversely with per-token gradient density. At 50k vocab # (sparse: ~4 hits/token/step) 0.3 was optimal; at 4096 vocab embeddings update # ~12x denser, so hold the effective per-token update constant -> ~0.03. EMB_LR = 0.05 # input embedding LR (tuned for 8k in-domain vocab) HEAD_LR = 3.6e-3 # peak LR for the untied output head (AdamW ~ 0.1/sqrt(d)) SCALAR_LR = 0.015 # peak LR for 1-D params: norms, skip/value-residual scalars MIN_LR_RATIO = 0.01 WARMUP_STEPS = 300 WEIGHT_DECAY = 0.1 GRAD_CLIP = 1.0 REPORT_EVERY = 20 EOT_ID = 50256 # gpt2 <|endoftext|> # Rows of unique data to read (memory-safe as uint16); ~420M tokens, reused a few epochs. MAX_DATA_ROWS = 700000 # Large effective batch via gradient accumulation (the run is step-capped, so # tokens/step is the lever); target a few epochs over the shard, not more. EFFECTIVE_TOKENS_PER_STEP = 262144 MICRO_TOKENS = 16384 TARGET_EPOCHS = 2.0 def _ctx_get(ctx, key, default): if hasattr(ctx, "get"): try: v = ctx.get(key, default) return default if v is None else v except Exception: return default return getattr(ctx, key, default) def _budget(ctx): max_steps, max_seconds = HARD_MAX_STEPS, HARD_MAX_SECONDS b = getattr(ctx, "budget", None) if callable(b): try: bb = b() max_steps = int(getattr(bb, "max_steps", max_steps)) max_seconds = float(getattr(bb, "max_seconds", max_seconds)) except Exception: pass max_steps = int(_ctx_get(ctx, "max_train_steps", max_steps)) max_seconds = float(_ctx_get(ctx, "max_seconds", max_seconds)) return min(max_steps, HARD_MAX_STEPS), min(max_seconds, HARD_MAX_SECONDS) def _packed_tokens(ctx, tok, block, want_tokens): # Concatenate tokenized docs (EOT-separated) into one uint16 stream, then reshape # into (n, block) windows. Tokens are kept as numpy uint16 (gpt2 vocab < 65536) so # ~hundreds of millions of tokens fit in memory instead of a huge Python int list. import numpy as np path = _ctx_get(ctx, "dataset_path", None) tr = int(_ctx_get(ctx, "train_rows", 2048)) vr = int(_ctx_get(ctx, "val_rows", 256)) # Cache the packed token stream so a re-run in the same pod skips re-tokenizing. cache = f"/tmp/prism_pack_{MAX_DATA_ROWS}_{tr}_{vr}.npy" if os.path.exists(cache): stream = np.load(cache) else: # Only materialize the rows we use. The harness scores texts[tr:tr+vr]; the # pinned shard has no <100-char rows, so raw row order matches its filtered # order and slicing raw keeps train disjoint from the scored val cut. col = pq.read_table(path, columns=["text"]).column("text") head = col.slice(0, tr).to_pylist() body = col.slice(tr + vr, MAX_DATA_ROWS).to_pylist() texts = [t for t in head + body if isinstance(t, str) and len(t) >= 100] eot = np.array([EOT_ID], dtype=np.uint16) parts = [] total = 0 for start in range(0, len(texts), 1024): for enc in tok(texts[start : start + 1024], add_special_tokens=False)["input_ids"]: parts.append(np.asarray(enc, dtype=np.uint16)) parts.append(eot) total += len(enc) + 1 if total >= want_tokens: break stream = np.concatenate(parts) if parts else np.zeros(block + 1, dtype=np.uint16) try: np.save(cache, stream) except OSError: pass n = max(1, (len(stream) - 1) // block) stream = stream[: n * block + 1] x = torch.from_numpy(stream[:-1].astype(np.int32)).view(n, block) y = torch.from_numpy(stream[1:].astype(np.int32)).view(n, block) return x, y def _batch_size(device, block): if not str(device).startswith("cuda"): return 4 # ~32k tokens/step; back off if VRAM is tight (handled by the OOM guard). return max(1, min(48, 32768 // max(block, 1))) def _doc_block_mask(xb): # Block-sparse mask: attend only within the same document (EOT-split) + causal. if _create_block_mask is None: return None eot = (xb == EOT_ID) doc_id = torch.cumsum(eot.int(), dim=1) - eot.int() def mm(b, h, q_idx, kv_idx): return (q_idx >= kv_idx) & (doc_id[b, q_idx] == doc_id[b, kv_idx]) return _create_block_mask(mm, xb.shape[0], None, xb.shape[1], xb.shape[1], device=xb.device) def _sched(elapsed, budget, warmup_frac=0.03, decay_frac=0.20): # Warmup -> stable (hold peak LR) -> cosine decay over the last decay_frac. # WSD keeps LR high through most of the run (faster convergence when val is # still descending at the cap) and only anneals near the end. warm = warmup_frac * budget if elapsed < warm: return elapsed / max(warm, 1e-9) decay_start = (1.0 - decay_frac) * budget if elapsed < decay_start: return 1.0 progress = min(1.0, (elapsed - decay_start) / max(1e-9, budget - decay_start)) cosine = 0.5 * (1.0 + math.cos(math.pi * progress)) return MIN_LR_RATIO + (1.0 - MIN_LR_RATIO) * cosine def _is_no_decay(name, p): return p.dim() < 2 or "norm" in name.lower() or "emb" in name.lower() def _newtonschulz5(G, steps=5): # Orthogonalize the gradient matrix via a quintic Newton-Schulz iteration # (Muon; Bernstein/Jordan). Runs in bf16; coefficients tuned for LM training. a, b, c = 3.4445, -4.7750, 2.0315 X = G.bfloat16() transposed = X.shape[0] > X.shape[1] if transposed: X = X.T X = X / (X.norm() + 1e-7) for _ in range(steps): A = X @ X.T B = b * A + c * (A @ A) X = a * X + B @ X return X.T if transposed else X class Muon(torch.optim.Optimizer): """Muon: momentum + Newton-Schulz orthogonalization for 2D weight matrices. Trains transformer matrices markedly faster than AdamW at the same compute.""" def __init__(self, params, lr=0.02, momentum=0.95, nesterov=True, ns_steps=5): super().__init__(params, dict(lr=lr, momentum=momentum, nesterov=nesterov, ns_steps=ns_steps)) @torch.no_grad() def step(self): for grp in self.param_groups: lr, mom, nest, ns = grp["lr"], grp["momentum"], grp["nesterov"], grp["ns_steps"] for p in grp["params"]: if p.grad is None: continue st = self.state[p] if "buf" not in st: st["buf"] = torch.zeros_like(p.grad) buf = st["buf"] buf.mul_(mom).add_(p.grad) g = p.grad.add(buf, alpha=mom) if nest else buf g = _newtonschulz5(g, steps=ns) # Scale so the RMS update size is ~lr regardless of matrix shape. scale = (max(1.0, p.size(0) / p.size(1)) ** 0.5) p.add_(g.to(p.dtype), alpha=-lr * scale) def train(model, ctx): """Recipe contract entrypoint: trains in place, returns a metrics dict.""" device = _ctx_get(ctx, "device", "cpu") torch.manual_seed(int(_ctx_get(ctx, "seed", 0))) guard = _ctx_get(ctx, "guard", lambda: None) if str(device).startswith("cuda"): torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True max_steps, max_seconds = _budget(ctx) warmup = min(WARMUP_STEPS, max(1, max_steps // 20)) tok = _resolve_tokenizer(ctx) # The frozen-val bpb is scored with THIS tokenizer; the doc separator in the # packed training stream + the doc-mask must use its real eos id. global EOT_ID _eos = getattr(tok, "eos_token_id", None) EOT_ID = int(_eos) if _eos is not None else (len(tok) - 1) block = getattr(model, "block", 512) micro = max(1, int(_ctx_get(ctx, "batch_size", MICRO_TOKENS // max(block, 1)))) accum = max(1, EFFECTIVE_TOKENS_PER_STEP // (micro * block)) x, y = _packed_tokens(ctx, tok, block, want_tokens=1_200_000_000) # Exact replica of the harness val cut (texts[train_rows:train_rows+val_rows]) # for a live bpb estimate; these rows are excluded from training data. val_ids = [] try: _vp = _ctx_get(ctx, "dataset_path", None) _tr = int(_ctx_get(ctx, "train_rows", 2048)); _vr = int(_ctx_get(ctx, "val_rows", 256)) _vcol = pq.read_table(_vp, columns=["text"]).column("text").slice(_tr, _vr).to_pylist() _vtexts = [t for t in _vcol if isinstance(t, str) and len(t) >= 100] for _t in _vtexts: _ids = tok(_t, add_special_tokens=False)["input_ids"] if len(_ids) >= 2: val_ids.append(torch.tensor(_ids, dtype=torch.long)) except Exception: val_ids = [] n_windows = x.shape[0] # Train for the whole compute budget (time governs the schedule); reshuffled # windows give ~as many epochs as fit, and the LR decays to its floor by the cap. time_budget = max_seconds * 0.94 # Muon for the 2D transformer weight matrices; AdamW for embeddings + norms. muon_params, adam_emb, adam_head, adam_scalar = [], [], [], [] for n, p in model.named_parameters(): if not p.requires_grad: continue if "tok_emb" in n: adam_emb.append(p) elif "head" in n: adam_head.append(p) elif p.ndim >= 2: muon_params.append(p) else: adam_scalar.append(p) opt_muon = Muon(muon_params, lr=MUON_LR, momentum=0.95) opt_adam = torch.optim.AdamW( [{"params": adam_emb, "lr": EMB_LR}, {"params": adam_head, "lr": HEAD_LR}, {"params": adam_scalar, "lr": SCALAR_LR}], betas=(0.9, 0.95), eps=1e-8, weight_decay=0.0, fused=str(device).startswith("cuda"), ) optimizers = [opt_muon, opt_adam] use_bf16 = str(device).startswith("cuda") and torch.cuda.is_bf16_supported() model.to(device) model.train() EMA_DECAY = 0.999 ema_params = {n: p.detach().clone().float() for n, p in model.named_parameters() if p.requires_grad} # Compile the TRAINING forward for throughput (fixed shape). The harness scores # val on the original `model` (variable lengths), so it never recompiles. fwd = model if str(device).startswith("cuda"): try: fwd = torch.compile(model) except Exception: fwd = model g = torch.Generator().manual_seed(int(_ctx_get(ctx, "seed", 0))) order = torch.randperm(n_windows, generator=g) t0 = time.time() last = 0.0 ema = 0.0 cursor = 0 step = 0 def _next_micro(n): nonlocal cursor, order if cursor + n > n_windows: order = torch.randperm(n_windows, generator=g) cursor = 0 idx = order[cursor : cursor + n] cursor += n return (x[idx].to(device, non_blocking=True).long(), y[idx].to(device, non_blocking=True).long()) while step < max_steps: guard() elapsed = time.time() - t0 if elapsed >= time_budget: break mult = _sched(elapsed, time_budget) lr = MUON_LR * mult opt_muon.param_groups[0]["lr"] = MUON_LR * mult opt_adam.param_groups[0]["lr"] = EMB_LR * mult opt_adam.param_groups[1]["lr"] = HEAD_LR * mult opt_adam.param_groups[2]["lr"] = SCALAR_LR * mult for o in optimizers: o.zero_grad(set_to_none=True) try: for a in range(accum): xb, yb = _next_micro(micro) bm = _doc_block_mask(xb) if use_bf16: with torch.autocast(device_type="cuda", dtype=torch.bfloat16): logits = fwd(xb, block_mask=bm) loss = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), yb.reshape(-1)) else: logits = fwd(xb, block_mask=bm) loss = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), yb.reshape(-1)) (loss / accum).backward() grad_norm = float(torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP)) for o in optimizers: o.step() with torch.no_grad(): for _n, _p in model.named_parameters(): if _n in ema_params: ema_params[_n].mul_(EMA_DECAY).add_(_p.detach().float(), alpha=1.0 - EMA_DECAY) except RuntimeError as exc: if "out of memory" in str(exc).lower() and micro > 1: micro = max(1, micro // 2) accum = max(1, EFFECTIVE_TOKENS_PER_STEP // (micro * block)) if str(device).startswith("cuda"): torch.cuda.empty_cache() continue raise last = float(loss.item()) ema = last if step == 0 else (0.98 * ema + 0.02 * last) step += 1 if step == 1 or step % REPORT_EVERY == 0: prism_telemetry.report( loss=last, step=step, grad_norm=grad_norm, layer_stats={"model": {"grad_norm": grad_norm}}, ) if step == 1 or step % 100 == 0: print(f"[train] step {step}/{max_steps} loss {last:.3f} ema {ema:.3f} lr {lr:.2e} " f"gnorm {grad_norm:.2f} t {time.time()-t0:.0f}s", file=sys.stderr, flush=True) if val_ids and (step == 1 or step % 1000 == 0): try: model.eval() vloss = [] with torch.no_grad(): for _ids in val_ids: _i = _ids.to(device).unsqueeze(0) if use_bf16: with torch.autocast(device_type="cuda", dtype=torch.bfloat16): _lg = model(_i[:, :-1]) _tg = _i[:, 1:][:, -_lg.shape[1]:] _l = F.cross_entropy(_lg.reshape(-1, _lg.shape[-1]).float(), _tg.reshape(-1)) else: _lg = model(_i[:, :-1]) _tg = _i[:, 1:][:, -_lg.shape[1]:] _l = F.cross_entropy(_lg.reshape(-1, _lg.shape[-1]), _tg.reshape(-1)) vloss.append(float(_l)) model.train() _ce = sum(vloss) / max(1, len(vloss)) print(f"[val] step {step} val_ce {_ce:.4f} bpb {_ce/0.6931471805599453:.4f} " f"(top 4.2881) t {time.time()-t0:.0f}s", file=sys.stderr, flush=True) except Exception: model.train() with torch.no_grad(): for _n, _p in model.named_parameters(): if _n in ema_params: _p.copy_(ema_params[_n].to(_p.dtype)) prism_telemetry.report(loss=last, step=step, grad_norm=0.0) prism_telemetry.finish_evaluation() return {"train_loss": last, "train_steps": step, "train_seconds": time.time() - t0}