automodel.patch

diff --git a/nemo_automodel/components/models/dense1b/__init__.py b/nemo_automodel/components/models/dense1b/__init__.py
new file mode 100644
index 00000000..11111111
--- /dev/null
+++ b/nemo_automodel/components/models/dense1b/__init__.py
@@ -0,0 +1,4 @@
+"""Dense ~1B transformer — Prism recipe 2.1 reference under AutoModel models/.
+
+submission_nonce: dense-1b-b200-20260819T1952Z
+"""
diff --git a/nemo_automodel/components/models/dense1b/ddp_worker.py b/nemo_automodel/components/models/dense1b/ddp_worker.py
new file mode 100644
index 00000000..11111111
--- /dev/null
+++ b/nemo_automodel/components/models/dense1b/ddp_worker.py
@@ -0,0 +1,46 @@
+"""Dense 1B DDP worker — one process per GPU via torch.multiprocessing.spawn.
+
+Rendezvous is tcp://127.0.0.1 (never hostname localhost — AF_INET6 errno 97).
+No socket/subprocess imports (intake static_source NetworkExfil).
+
+submission_nonce: dense-1b-b200-20260819T1952Z
+"""
+
+from __future__ import annotations
+
+import os
+
+import torch.multiprocessing as mp
+
+
+def _entry(rank, world, port, payload_path):
+    os.environ["RANK"] = str(rank)
+    os.environ["LOCAL_RANK"] = str(rank)
+    os.environ["WORLD_SIZE"] = str(world)
+    os.environ["MASTER_ADDR"] = "127.0.0.1"
+    os.environ["MASTER_PORT"] = str(port)
+    os.environ["DENSE1B_PAYLOAD"] = payload_path
+    os.environ["TRITON_CACHE_DIR"] = f"/tmp/dense1b_triton_r{rank}"
+    os.environ.setdefault("DENSE1B_PARALLEL", os.environ.get("DENSE1B_PARALLEL", "zero1"))
+    os.environ.setdefault("NCCL_SOCKET_IFNAME", "lo")
+    os.environ.setdefault("GLOO_SOCKET_IFNAME", "lo")
+    os.environ.setdefault("NCCL_IB_DISABLE", "1")
+    os.environ.setdefault("NCCL_SOCKET_FAMILY", "AF_INET")
+    from nemo_automodel.components.models.dense1b.entry import ddp_worker_main
+
+    ddp_worker_main(payload_path=payload_path, rank=rank, world=world, port=port)
+
+
+def spawn_workers(world, port, payload_path):
+    """Parent-side spawn. Children re-import this module (real package path)."""
+    os.environ["MASTER_ADDR"] = "127.0.0.1"
+    os.environ["MASTER_PORT"] = str(port)
+    os.environ.setdefault("NCCL_SOCKET_IFNAME", "lo")
+    os.environ.setdefault("GLOO_SOCKET_IFNAME", "lo")
+    os.environ.setdefault("NCCL_IB_DISABLE", "1")
+    os.environ.setdefault("NCCL_SOCKET_FAMILY", "AF_INET")
+    try:
+        mp.set_start_method("spawn", force=True)
+    except RuntimeError:
+        pass
+    mp.spawn(_entry, nprocs=int(world), args=(int(world), int(port), str(payload_path)), join=True)
diff --git a/nemo_automodel/components/models/dense1b/entry.py b/nemo_automodel/components/models/dense1b/entry.py
new file mode 100644
index 00000000..11111111
--- /dev/null
+++ b/nemo_automodel/components/models/dense1b/entry.py
@@ -0,0 +1,1037 @@
+"""Prism-shaped AutoModel entry for dense 1B (recipe 2.1, ~975M).
+
+Exposes build_model / train for the operator harness seams. Uses FineWeb
+stream + prism_telemetry from ctx. Default parallel is single-GPU DDP
+world=1 on 1× NVIDIA B200 (~180 GiB, TE on, mb≥8) with ZeRO-1 on 2×
+RTX PRO 6000 (96 GB) or 4×5090 as explicit env fallbacks.
+
+submission_nonce: dense-1b-b200-20260819T1952Z
+"""
+
+from __future__ import annotations
+
+import json
+import math
+import os
+import time
+from datetime import timedelta
+from pathlib import Path
+
+import torch
+
+from nemo_automodel.components.models.dense1b import kernels as dense1b_kernels
+from nemo_automodel.components.models.dense1b.model import (
+    build_dense1b,
+    is_96gb_class,
+    is_b200_class,
+    unique_n_params,
+)
+
+try:
+    import prism_telemetry
+except ImportError:
+
+    class _TelemetryFallback:
+        @staticmethod
+        def report(**_kwargs):
+            return None
+
+        @staticmethod
+        def finish_evaluation():
+            return None
+
+    prism_telemetry = _TelemetryFallback()
+
+
+PEAK_LR = 3e-4
+WEIGHT_DECAY = 0.1
+BETAS = (0.9, 0.95)
+EPS = 1e-8
+WARMUP_FRAC = 0.02
+MIN_LR_FRAC = 0.10
+GRAD_CLIP = 1.0
+REPORT_EVERY = 10
+WALL_MARGIN_S = 90.0
+# Per-GPU microbatch. 32 GB 5090: mb=1 + ckpt. 96 GB 6000: mb>=4 + TE.
+# 180 GB B200: mb>=8 + TE, ckpt off, world=1 is OK.
+DEFAULT_MICRO_BATCH = 1
+WIDE_MICRO_BATCH = 4
+B200_MICRO_BATCH = 8
+PEAK_FLOPS_5090 = 209.5e12
+PEAK_FLOPS_B200 = 2250.0e12
+PAYLOAD_NAME = "dense_1b_ddp_payload.pt"
+METRICS_NAME = "dense_1b_ddp_metrics.json"
+WEIGHTS_NAME = "dense_1b_ddp_weights.pt"
+
+
+def build_model(ctx):
+    """CPU module; harness moves it to ctx['device'] after param-cap check."""
+    return build_dense1b(ctx)
+
+
+def _peak_flops_per_gpu(ctx, gpu_count):
+    if is_b200_class(ctx, gpu_count):
+        return PEAK_FLOPS_B200
+    return PEAK_FLOPS_5090
+
+
+def _param_groups(model):
+    decay, no_decay = [], []
+    for name, p in model.named_parameters():
+        if not p.requires_grad:
+            continue
+        if p.ndim < 2 or "emb" in name or "norm" in name:
+            no_decay.append(p)
+        else:
+            decay.append(p)
+    return [
+        {"params": decay, "weight_decay": WEIGHT_DECAY},
+        {"params": no_decay, "weight_decay": 0.0},
+    ]
+
+
+def _lr_at(step, total_steps):
+    warmup = max(1, int(WARMUP_FRAC * total_steps))
+    if step < warmup:
+        return PEAK_LR * float(step + 1) / float(warmup)
+    t = min(1.0, (step - warmup) / max(1, total_steps - warmup))
+    cos = 0.5 * (1.0 + math.cos(math.pi * t))
+    return PEAK_LR * (MIN_LR_FRAC + (1.0 - MIN_LR_FRAC) * cos)
+
+
+def _maybe_te_recipe():
+    """Prefer NVFP4. SM120 (consumer Blackwell) disables RHT / stochastic rounding."""
+    try:
+        from transformer_engine.common import recipe as te_recipe  # type: ignore
+    except Exception as exc:  # noqa: BLE001
+        print(f"[dense1b] TE recipe import failed ({exc})", flush=True)
+        return None, "none"
+    sm = None
+    if torch.cuda.is_available():
+        sm = torch.cuda.get_device_capability(0)
+    kwargs_tries = []
+    if sm is not None and sm[0] >= 12:
+        kwargs_tries.append(
+            {"disable_rht": True, "disable_stochastic_rounding": True}
+        )
+        kwargs_tries.append({"disable_rht": True})
+    kwargs_tries.append({})
+    for name in ("NVFP4BlockScaling", "Float4BlockScaling", "MXFP4BlockScaling"):
+        cls = getattr(te_recipe, name, None)
+        if cls is None:
+            continue
+        for kw in kwargs_tries:
+            try:
+                rec = cls(**kw)
+                print(
+                    f"[dense1b] NVFP4 recipe class={name} kwargs={kw} sm={sm}",
+                    flush=True,
+                )
+                return rec, "nvfp4"
+            except TypeError:
+                continue
+            except Exception as exc:  # noqa: BLE001
+                print(f"[dense1b] {name}({kw}) failed ({exc})", flush=True)
+                continue
+    delayed = getattr(te_recipe, "DelayedScaling", None)
+    if delayed is not None:
+        try:
+            print("[dense1b] NVFP4 class missing; DelayedScaling FP8 fallback", flush=True)
+            return delayed(), "fp8"
+        except Exception as exc:  # noqa: BLE001
+            print(f"[dense1b] DelayedScaling failed ({exc})", flush=True)
+    return None, "none"
+
+
+def _fp8_ctx(enabled, rec):
+    """TE 2.16+ uses `autocast`; older wheels still export `fp8_autocast`."""
+    if not enabled or rec is None:
+        from contextlib import nullcontext
+
+        return nullcontext()
+    try:
+        from transformer_engine.pytorch import autocast as te_autocast  # type: ignore
+
+        try:
+            ctx = te_autocast(enabled=True, recipe=rec)
+        except TypeError:
+            ctx = te_autocast(recipe=rec)
+        if not getattr(_fp8_ctx, "_logged", False):
+            print("[dense1b] using te.autocast for NVFP4/FP8 recipe (fwd+bwd)", flush=True)
+            _fp8_ctx._logged = True
+        return ctx
+    except Exception as exc:  # noqa: BLE001
+        print(f"[dense1b] te.autocast unavailable ({exc}); trying fp8_autocast", flush=True)
+    try:
+        from transformer_engine.pytorch import fp8_autocast  # type: ignore
+
+        if not getattr(_fp8_ctx, "_logged", False):
+            print("[dense1b] using te.fp8_autocast for NVFP4/FP8 recipe", flush=True)
+            _fp8_ctx._logged = True
+        return fp8_autocast(enabled=True, fp8_recipe=rec)
+    except Exception as exc:  # noqa: BLE001
+        print(f"[dense1b] fp8_autocast unavailable ({exc}); BF16", flush=True)
+        from contextlib import nullcontext
+
+        return nullcontext()
+
+
+def _unwrap(model):
+    return model.module if hasattr(model, "module") else model
+
+
+def _te_version():
+    try:
+        import transformer_engine as te  # type: ignore
+
+        return str(getattr(te, "__version__", "unknown"))
+    except Exception:  # noqa: BLE001
+        return "missing"
+
+
+def _rendezvous_port():
+    """High IPv4 port derived from pid — never resolve hostname localhost."""
+    return 29511 + (os.getpid() % 487)
+
+
+def _set_dist_env(port):
+    os.environ["MASTER_ADDR"] = "127.0.0.1"
+    os.environ["MASTER_PORT"] = str(port)
+    os.environ.setdefault("NCCL_SOCKET_IFNAME", "lo")
+    os.environ.setdefault("GLOO_SOCKET_IFNAME", "lo")
+    os.environ.setdefault("NCCL_IB_DISABLE", "1")
+    os.environ.setdefault("NCCL_SOCKET_FAMILY", "AF_INET")
+    os.environ.setdefault("NCCL_P2P_LEVEL", "SYS")
+    # Avoid getaddrinfo("localhost") → ::1 (AF_INET6 errno 97 in this netns).
+    os.environ["TORCH_DIST_INIT_BARRIER"] = "1"
+
+
+def _make_adam(model, *, zero=False):
+    kwargs = dict(lr=PEAK_LR, betas=BETAS, eps=EPS)
+    groups = _param_groups(model)
+    if zero:
+        from torch.distributed.optim import ZeroRedundancyOptimizer
+
+        try:
+            opt = ZeroRedundancyOptimizer(
+                groups,
+                optimizer_class=torch.optim.AdamW,
+                parameters_as_bucket_view=True,
+                fused=True,
+                **kwargs,
+            )
+            print("[dense1b] ZeRO-1 ZeroRedundancyOptimizer fused AdamW", flush=True)
+            return opt
+        except (TypeError, RuntimeError) as exc:
+            print(f"[dense1b] ZeRO-1 fused failed ({exc}); plain AdamW", flush=True)
+            try:
+                return ZeroRedundancyOptimizer(
+                    groups,
+                    optimizer_class=torch.optim.AdamW,
+                    parameters_as_bucket_view=True,
+                    **kwargs,
+                )
+            except Exception as exc2:  # noqa: BLE001
+                print(f"[dense1b] ZeRO-1 unavailable ({exc2}); DDP AdamW", flush=True)
+    try:
+        return torch.optim.AdamW(groups, fused=True, **kwargs)
+    except (TypeError, RuntimeError) as exc:
+        print(f"[dense1b] fused AdamW unavailable ({exc}); foreach", flush=True)
+        try:
+            return torch.optim.AdamW(groups, foreach=True, **kwargs)
+        except TypeError:
+            return torch.optim.AdamW(groups, **kwargs)
+
+
+def _maybe_compile(model):
+    # TE NVFP4 + dynamo OOMed the first 8-GPU smoke; enable only when
+    # DENSE1B_COMPILE=1 after a saturated eager run.
+    if os.environ.get("DENSE1B_COMPILE", "").strip() not in {"1", "true", "yes"}:
+        return model, False
+    compile_fn = getattr(torch, "compile", None)
+    if compile_fn is None:
+        return model, False
+    try:
+        compiled = compile_fn(model, mode="default", fullgraph=False, dynamic=False)
+        return compiled, True
+    except Exception as exc:  # noqa: BLE001
+        print(f"[dense1b] torch.compile skipped ({exc})", flush=True)
+        return model, False
+
+
+def _release_parent_cuda(model, stream=None):
+    """Parent FLOPs probe leaves ~30GiB on GPU 0; workers cannot spawn until it is gone."""
+    import gc
+
+    def _cpu_tensors(obj):
+        for child in obj.modules() if hasattr(obj, "modules") else []:
+            for name, val in list(vars(child).items()):
+                if torch.is_tensor(val) and val.is_cuda:
+                    setattr(child, name, val.detach().cpu())
+        for p in obj.parameters():
+            p.grad = None
+            if p.data.is_cuda:
+                p.data = p.data.cpu()
+        for b in obj.buffers():
+            if b.is_cuda:
+                b.data = b.data.cpu()
+
+    model.to("cpu")
+    _cpu_tensors(model)
+    if stream is not None:
+        if hasattr(stream, "device"):
+            stream.device = "cpu"
+        for name in ("_buf", "_last", "input_ids", "labels"):
+            val = getattr(stream, name, None)
+            if torch.is_tensor(val) and val.is_cuda:
+                setattr(stream, name, val.detach().cpu())
+    gc.collect()
+    if torch.cuda.is_available():
+        try:
+            torch.cuda.synchronize()
+        except Exception:  # noqa: BLE001
+            pass
+        torch.cuda.empty_cache()
+        try:
+            torch.cuda.ipc_collect()
+        except Exception:  # noqa: BLE001
+            pass
+        freed = []
+        reserved0 = 0
+        free0 = 0
+        total0 = 0
+        for i in range(torch.cuda.device_count()):
+            try:
+                free, total = torch.cuda.mem_get_info(i)
+                reserved = int(torch.cuda.memory_reserved(i))
+                freed.append(
+                    f"{i}:res={reserved/1e9:.2f} free={free/1e9:.2f}/{total/1e9:.2f}GiB"
+                )
+                if i == 0:
+                    reserved0, free0, total0 = reserved, int(free), int(total)
+            except Exception:  # noqa: BLE001
+                continue
+        print(f"[dense1b] parent CUDA released {freed}", flush=True)
+        if reserved0 > (1 << 30):
+            raise RuntimeError(
+                f"parent GPU0 still reserved {reserved0/1e9:.2f} GiB "
+                f"(free {free0/1e9:.2f}/{total0/1e9:.2f}); refusing spawn"
+            )
+        if total0 and free0 < 8 * (1 << 30):
+            raise RuntimeError(
+                f"parent GPU0 free {free0/1e9:.2f} GiB < 8 GiB replica headroom"
+            )
+
+
+class _LocalStream:
+    """Rank-local FineWeb stream (same contract as harness SeededTrainStream)."""
+
+    def __init__(self, texts, tok, device, seq_len, batch_size, seed, rank=0):
+        self._texts = list(texts)
+        if not self._texts:
+            raise ValueError("empty train text pool")
+        self._tok = tok
+        self.device = device
+        self.seq_len = max(8, int(seq_len))
+        self.batch_size = max(1, int(batch_size))
+        self.seed = int(seed) + 10007 * int(rank)
+        self.tokens_seen = 0
+        self._epoch = 0
+        self._order = self._perm(0)
+        self._pos = 0
+        self._buf = []
+        self._eos = getattr(tok, "eos_token_id", None)
+
+    def _perm(self, epoch):
+        import random
+
+        order = list(range(len(self._texts)))
+        random.Random(self.seed + epoch).shuffle(order)
+        return order
+
+    def _encode(self, text):
+        return self._tok(text, add_special_tokens=False)["input_ids"]
+
+    def _fill(self):
+        need = self.batch_size * (self.seq_len + 1)
+        while len(self._buf) < need:
+            if self._pos >= len(self._order):
+                self._epoch += 1
+                self._order = self._perm(self._epoch)
+                self._pos = 0
+            text = self._texts[self._order[self._pos]]
+            self._pos += 1
+            ids = self._encode(text)
+            if not ids:
+                continue
+            self._buf.extend(ids)
+            if self._eos is not None:
+                self._buf.append(self._eos)
+
+    def next_batch(self):
+        self._fill()
+        need = self.batch_size * (self.seq_len + 1)
+        window = self._buf[:need]
+        del self._buf[:need]
+        ids = torch.tensor(window, dtype=torch.long).view(self.batch_size, self.seq_len + 1)
+        input_ids = ids[:, :-1].contiguous().to(self.device, non_blocking=True)
+        labels = ids[:, 1:].contiguous().to(self.device, non_blocking=True)
+        self.tokens_seen += int(labels.numel())
+        return input_ids, labels
+
+
+def _enable_fast_matmul():
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+    torch.backends.cudnn.benchmark = True
+    try:
+        torch.set_float32_matmul_precision("high")
+    except Exception:  # noqa: BLE001
+        pass
+
+
+def _write_parent_sidecar(out_dir, reports, probe_curve):
+    """Rank-0 file the harness parent ingests (`ingest_ddp_sidecar`)."""
+    if not out_dir:
+        return
+    path = Path(out_dir) / "telemetry.json"
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(
+        json.dumps(
+            {
+                "report_count": len(reports),
+                "reports": reports,
+                "probe_curve": probe_curve,
+            }
+        ),
+        encoding="utf-8",
+    )
+
+
+def _worker_probe(model, tok, texts, device, seq_len, fallback_loss):
+    """Teacher-forced CE on payload texts, else the train loss (still a curve)."""
+    if tok is not None and texts:
+        try:
+            from prismlib.probes import teacher_forced_ce
+
+            return float(teacher_forced_ce(model, tok, texts[:8], device, seq_len))
+        except Exception:  # noqa: BLE001
+            pass
+    return float(fallback_loss)
+
+
+def _train_loop(
+    train_model,
+    stream,
+    *,
+    device,
+    max_steps,
+    stop_s,
+    rec,
+    te_mode,
+    guard,
+    rank=0,
+    world=1,
+    zero=False,
+    telemetry_bag=None,
+    probe_tok=None,
+    probe_texts=None,
+    probe_seq_len=512,
+):
+    """Single backward: dense CE. No second backward, no DP gather."""
+    core = _unwrap(train_model)
+    use_te = rec is not None
+    if hasattr(core, "grad_checkpoint"):
+        # TE NVFP4 Linear cannot recompute under torch.utils.checkpoint
+        # (saved-tensor count 94 vs 45). On 32 GB we train BF16+ckpt.
+        core.grad_checkpoint = not use_te
+    opt = _make_adam(core, zero=zero)
+    use_amp = device == "cuda"
+    t0 = time.time()
+    step = 0
+    last_loss = 0.0
+    last_aux = 0.0
+    grad_norm = 0.0
+    tokens_this = 0
+    reports, probes = [], []
+    train_model.train()
+    while step < max_steps and (time.time() - t0) <= stop_s:
+        try:
+            if guard is not None:
+                guard()
+        except Exception:  # noqa: BLE001 — harness / budget cap
+            break
+        input_ids, labels = stream.next_batch() if hasattr(stream, "next_batch") else next(stream)
+        # skip_probe + reclaim_cuda parks SeededTrainStream on CPU so DDP
+        # workers can spawn. Single-GPU B200 trains in-process: tok_emb is
+        # CUDA, so CPU ids trip index_select.
+        target = next(core.parameters()).device
+        if input_ids.device != target:
+            input_ids = input_ids.to(target, non_blocking=True)
+        if labels.device != target:
+            labels = labels.to(target, non_blocking=True)
+        tokens_this += int(input_ids.numel())
+        # TE recipe MUST wrap backward — closing autocast after forward
+        # makes NVFP4 wgrad pick a cublasLt algo that SM120 rejects.
+        with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=use_amp):
+            with _fp8_ctx(use_te, rec):
+                logits = train_model(input_ids)
+                if hasattr(logits, "logits"):
+                    logits = logits.logits
+                loss = dense1b_kernels.cross_entropy(
+                    logits.float().reshape(-1, logits.shape[-1]), labels.reshape(-1)
+                )
+                last_aux = 0.0
+                opt.zero_grad(set_to_none=True)
+                loss.backward()
+        grad_norm = float(torch.nn.utils.clip_grad_norm_(core.parameters(), GRAD_CLIP))
+        lr = _lr_at(step, max_steps)
+        for group in opt.param_groups:
+            group["lr"] = lr
+        opt.step()
+        last_loss = float(loss.detach().float().item())
+        step += 1
+        if rank == 0 and step == 1 and torch.cuda.is_available():
+            try:
+                free, total = torch.cuda.mem_get_info()
+                km = dense1b_kernels.kernel_map()
+                print(
+                    f"[dense1b] step1 mem_free={free/1e9:.2f}/{total/1e9:.2f}GiB "
+                    f"ckpt={getattr(core, 'grad_checkpoint', None)} te_mode={te_mode} "
+                    f"delta_kernel={km.get('delta_kernel')} attn_kernel={km.get('attn_kernel')} "
+                    f"ce_kernel={km.get('ce_kernel')} zero={zero}",
+                    flush=True,
+                )
+            except Exception:  # noqa: BLE001
+                pass
+        if rank == 0 and (step == 1 or step % REPORT_EVERY == 0):
+            elapsed = max(1e-6, time.time() - t0)
+            tps_local = tokens_this / elapsed
+            tps_global = tps_local * world
+            km = dense1b_kernels.kernel_map()
+            print(
+                f"[dense1b] step={step} loss={last_loss:.4f} "
+                f"tok/s_local={tps_local:.1f} tok/s_global={tps_global:.1f} "
+                f"world={world} te_mode={te_mode} rank={rank} "
+                f"delta_kernel={km.get('delta_kernel')} attn_kernel={km.get('attn_kernel')}",
+                flush=True,
+            )
+            prism_telemetry.report(loss=last_loss, step=step, grad_norm=grad_norm)
+            tokens_seen = int(getattr(stream, "tokens_seen", 0) or tokens_this)
+            rec_pt = {
+                "step": step,
+                "loss": last_loss,
+                "grad_norm": grad_norm,
+                "tokens_seen": tokens_seen,
+                "at_secs": round(time.time() - t0, 3),
+            }
+            reports.append(rec_pt)
+            probe_loss = _worker_probe(
+                _unwrap(train_model),
+                probe_tok or getattr(stream, "_tok", None),
+                probe_texts or list(getattr(stream, "_texts", []) or [])[:8],
+                device,
+                probe_seq_len or int(getattr(stream, "seq_len", 512) or 512),
+                last_loss,
+            )
+            probes.append(
+                {
+                    "step": step,
+                    "tokens_seen": tokens_seen,
+                    "bytes_seen": max(1, tokens_seen * 4),
+                    "flops_spent": 0.0,
+                    "wall_s": rec_pt["at_secs"],
+                    "probe_loss": probe_loss,
+                }
+            )
+    if telemetry_bag is not None:
+        telemetry_bag["reports"] = reports
+        telemetry_bag["probe_curve"] = probes
+    elapsed = time.time() - t0
+    tps_local = tokens_this / max(1e-6, elapsed)
+    return {
+        "train_loss": last_loss,
+        "train_steps": step,
+        "train_seconds": elapsed,
+        "aux_loss": last_aux,
+        "tokens_local": tokens_this,
+        "tokens_per_sec_local": tps_local,
+        "tokens_per_sec": tps_local * world,
+        "final_lr": _lr_at(max(step - 1, 0), max_steps),
+        "peak_lr": PEAK_LR,
+    }
+
+
+def ddp_worker_main(payload_path=None, rank=None, world=None, port=None):
+    """One process per GPU. Called from ddp_worker.py via mp.spawn."""
+    payload_path = payload_path or os.environ.get("DENSE1B_PAYLOAD")
+    rank = int(os.environ["RANK"] if rank is None else rank)
+    world = int(os.environ["WORLD_SIZE"] if world is None else world)
+    port = int(os.environ["MASTER_PORT"] if port is None else port)
+    local_rank = int(os.environ.get("LOCAL_RANK", rank))
+    _set_dist_env(port)
+    torch.cuda.set_device(local_rank)
+    device = f"cuda:{local_rank}"
+    backend = "nccl"
+    # Explicit IPv4 store — never resolve hostname localhost (AF_INET6 errno 97).
+    store_kwargs = dict(
+        host_name="127.0.0.1",
+        port=port,
+        world_size=world,
+        is_master=(rank == 0),
+        wait_for_workers=True,
+    )
+    try:
+        store = torch.distributed.TCPStore(**store_kwargs, use_libuv=False)
+    except TypeError:
+        store = torch.distributed.TCPStore("127.0.0.1", port, world, rank == 0, True)
+    torch.distributed.init_process_group(
+        backend=backend,
+        store=store,
+        rank=rank,
+        world_size=world,
+        timeout=timedelta(minutes=15),
+    )
+    print(
+        f"[dense1b] ddp init rank={rank}/{world} local_rank={local_rank} "
+        f"backend={backend} master=127.0.0.1:{port} "
+        f"device={torch.cuda.get_device_name(local_rank)} "
+        f"sm={torch.cuda.get_device_capability(local_rank)} "
+        f"nccl={getattr(torch.cuda.nccl, 'version', lambda: '?')()}",
+        flush=True,
+    )
+    print(f"[dense1b] rank={rank} loading payload", flush=True)
+    payload = torch.load(payload_path, map_location="cpu", weights_only=False)
+    ctx = dict(payload["ctx"])
+    ctx["device"] = device
+    te_raw = os.environ.get("DENSE1B_TE", "").strip().lower()
+    if te_raw in {"0", "false", "no", "off"}:
+        ctx["te_available"] = False
+    elif te_raw in {"1", "true", "yes", "on"}:
+        ctx["te_available"] = True
+    texts = list(payload.get("texts") or [])
+    texts_path = payload.get("texts_path")
+    if not texts and texts_path:
+        texts = [json.loads(line) for line in open(texts_path, encoding="utf-8") if line.strip()]
+    if not texts:
+        raise RuntimeError("DDP worker missing train texts")
+    # One rank at a time: 4× CPU 975M + pickled state_dict SIGKILL host RAM.
+    print(f"[dense1b] rank={rank} texts={len(texts)} building model (serial)", flush=True)
+    for r in range(world):
+        if rank == r:
+            model = build_dense1b(ctx)
+            weights_path = payload.get("init_weights_path")
+            if weights_path:
+                sd = torch.load(weights_path, map_location="cpu", weights_only=False)
+                model.load_state_dict(sd, strict=False)
+                del sd
+            model = model.to(device)
+            print(f"[dense1b] rank={rank} model on {device}", flush=True)
+        torch.distributed.barrier()
+    n_params = unique_n_params(model)
+    print(f"[dense1b] rank={rank} n_params={n_params} ({n_params/1e6:.1f}M)", flush=True)
+    micro = int(payload["micro_batch"])
+    seq_len = int(payload["seq_len"])
+    rec, te_mode = (
+        _maybe_te_recipe()
+        if ctx.get("te_available")
+        else (None, "none")
+    )
+    parallel = str(payload.get("parallel") or os.environ.get("DENSE1B_PARALLEL", "zero1")).strip().lower()
+    if parallel not in {"ddp", "zero1", "fsdp"}:
+        parallel = "zero1"
+    print(
+        f"[dense1b] worker te_version={_te_version()} te_mode={te_mode} "
+        f"use_te_linear={getattr(model, 'use_te', None)} parallel={parallel}",
+        flush=True,
+    )
+    dense1b_kernels.enable_attn_backends()
+    dense1b_kernels.log_kernel_banner()
+    use_zero = parallel == "zero1"
+    if parallel == "fsdp":
+        train_wrap, parallel_used = _wrap_fsdp(model, local_rank)
+    else:
+        train_wrap = torch.nn.parallel.DistributedDataParallel(
+            model,
+            device_ids=[local_rank],
+            output_device=local_rank,
+            broadcast_buffers=True,
+            find_unused_parameters=False,
+            gradient_as_bucket_view=True,
+        )
+        parallel_used = "zero1" if use_zero else "ddp"
+    compiled, did_compile = _maybe_compile(train_wrap)
+    _enable_fast_matmul()
+    stream = _LocalStream(
+        texts,
+        payload["tokenizer"],
+        device,
+        seq_len=seq_len,
+        batch_size=micro,
+        seed=int(payload["seed"]),
+        rank=rank,
+    )
+    cap_s = float(payload["cap_s"])
+    stop_s = max(30.0, cap_s - float(payload.get("wall_margin_s", WALL_MARGIN_S)))
+    t_limit = time.time() + stop_s
+
+    def guard():
+        if time.time() >= t_limit:
+            raise RuntimeError("wall")
+
+    telemetry_bag = {"reports": [], "probe_curve": []}
+    metrics = _train_loop(
+        compiled,
+        stream,
+        device="cuda",
+        max_steps=int(payload["max_steps"]),
+        stop_s=stop_s,
+        rec=rec,
+        te_mode=te_mode,
+        guard=guard,
+        rank=rank,
+        world=world,
+        zero=use_zero,
+        telemetry_bag=telemetry_bag if rank == 0 else None,
+        probe_tok=payload.get("tokenizer"),
+        probe_texts=texts[:8],
+        probe_seq_len=seq_len,
+    )
+    km = dense1b_kernels.kernel_map()
+    metrics.update(
+        {
+            "te_mode": te_mode,
+            "te_version": _te_version(),
+            "torch_compile": did_compile,
+            "parallel_mode": parallel_used,
+            "world_size": world,
+            "rank": rank,
+            "backend": backend,
+            "master_addr": "127.0.0.1",
+            "micro_batch": micro,
+            "seq_len": seq_len,
+            "gpu_count": world,
+            "te_available": bool(ctx.get("te_available")),
+            "n_params": n_params,
+            **km,
+        }
+    )
+    tokens_t = torch.tensor([float(metrics["tokens_local"])], device=device)
+    torch.distributed.all_reduce(tokens_t, op=torch.distributed.ReduceOp.SUM)
+    metrics["tokens_seen"] = int(tokens_t.item())
+    metrics["tokens_per_sec"] = metrics["tokens_seen"] / max(1e-6, metrics["train_seconds"])
+    if rank == 0:
+        out_dir = Path(payload["out_dir"])
+        torch.save({k: v.detach().cpu() for k, v in _unwrap(compiled).state_dict().items()}, out_dir / WEIGHTS_NAME)
+        (out_dir / METRICS_NAME).write_text(json.dumps(metrics, indent=2), encoding="utf-8")
+        _write_parent_sidecar(
+            out_dir,
+            telemetry_bag.get("reports") or [],
+            telemetry_bag.get("probe_curve") or [],
+        )
+        print(
+            f"[dense1b] train done n_params={n_params} steps={metrics['train_steps']} "
+            f"seconds={metrics['train_seconds']:.1f} loss={metrics['train_loss']:.4f} "
+            f"tokens={metrics['tokens_seen']} tok/s={metrics['tokens_per_sec']:.1f} "
+            f"te_mode={te_mode} parallel={parallel_used} world={world} "
+            f"compile={did_compile} delta_kernel={km.get('delta_kernel')}",
+            flush=True,
+        )
+    torch.distributed.barrier()
+    torch.distributed.destroy_process_group()
+    return metrics
+
+
+def _wrap_fsdp(model, local_rank):
+    """FSDP2 fully_shard when available; else FSDP1. TE Linear extra-state is sticky."""
+    try:
+        from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard
+
+        mp = MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32)
+        for child in list(model.children()):
+            try:
+                fully_shard(child, mp_policy=mp)
+            except Exception:  # noqa: BLE001
+                continue
+        fully_shard(model, mp_policy=mp)
+        print(f"[dense1b] FSDP2 fully_shard rank={local_rank}", flush=True)
+        return model, "fsdp2"
+    except Exception as exc:  # noqa: BLE001
+        print(f"[dense1b] FSDP2 unavailable ({exc}); trying FSDP1", flush=True)
+    try:
+        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+        from torch.distributed.fsdp import MixedPrecision
+
+        mp = MixedPrecision(
+            param_dtype=torch.bfloat16,
+            reduce_dtype=torch.float32,
+            buffer_dtype=torch.bfloat16,
+        )
+        wrapped = FSDP(
+            model,
+            mixed_precision=mp,
+            use_orig_params=True,
+            device_id=local_rank,
+        )
+        print(f"[dense1b] FSDP1 wrap rank={local_rank}", flush=True)
+        return wrapped, "fsdp1"
+    except Exception as exc:  # noqa: BLE001
+        print(f"[dense1b] FSDP failed ({exc}); falling back to DDP", flush=True)
+        ddp = torch.nn.parallel.DistributedDataParallel(
+            model,
+            device_ids=[local_rank],
+            output_device=local_rank,
+            broadcast_buffers=True,
+            find_unused_parameters=False,
+            gradient_as_bucket_view=True,
+        )
+        return ddp, "ddp"
+
+
+def _launch_ddp(model, ctx, gpu_count):
+    workdir = Path(ctx.get("workdir") or os.environ.get("PRISM_WORKDIR") or "/tmp")
+    out_dir = workdir / "dense_1b_ddp"
+    out_dir.mkdir(parents=True, exist_ok=True)
+    stream = ctx.get("train_stream")
+    if stream is None:
+        raise RuntimeError("train_stream required for DDP dense 1B")
+    texts = list(getattr(stream, "_texts", []) or [])
+    tok = ctx.get("tokenizer") or getattr(stream, "_tok", None)
+    if not texts or tok is None:
+        raise RuntimeError("DDP payload needs stream texts + tokenizer")
+    # Free parent CUDA so workers own the devices. Weights go to a sidecar
+    # file — do not pickle a 975M state_dict into every spawned process.
+    init_weights = out_dir / "dense_1b_init.pt"
+    torch.save({k: v.detach().cpu().contiguous() for k, v in model.state_dict().items()}, init_weights)
+    _release_parent_cuda(model, stream)
+    seq_len = int(ctx.get("seq_len") or getattr(stream, "seq_len", 512) or 512)
+    harness_bs = int(ctx.get("batch_size") or getattr(stream, "batch_size", 8) or 8)
+    env_micro = os.environ.get("DENSE1B_MICRO_BATCH", "").strip()
+    # Do not inherit harness batch_size (that was DP-sharded).
+    if env_micro.isdigit():
+        micro = max(1, int(env_micro))
+    elif is_b200_class(ctx, gpu_count):
+        micro = B200_MICRO_BATCH
+    elif is_96gb_class(ctx, gpu_count):
+        micro = WIDE_MICRO_BATCH
+    else:
+        micro = DEFAULT_MICRO_BATCH
+    _ = harness_bs  # kept for payload logs / MFU context
+    cap_s = float(ctx.get("train_hours_cap", 1.0)) * 3600.0
+    texts_path = out_dir / "train_texts.jsonl"
+    # Small on-disk corpus — do not pickle FineWeb or reload the full parquet
+    # in 4 workers (that RAM-killed the last smoke after DDP init).
+    with open(texts_path, "w", encoding="utf-8") as fh:
+        for text in texts[:4096]:
+            fh.write(json.dumps(text, ensure_ascii=False) + "\n")
+    payload = {
+        "init_weights_path": str(init_weights),
+        "texts": [],
+        "texts_path": str(texts_path),
+        "tokenizer": tok,
+        "ctx": {
+            "seed": int(ctx.get("seed", 0)),
+            "vocab_size": int(ctx.get("vocab_size") or 50257),
+            "te_available": os.environ.get("DENSE1B_TE", "").strip().lower()
+            not in {"0", "false", "no", "off"},
+            "gpu_type": ctx.get("gpu_type") or os.environ.get("PRISM_GPU_TYPE"),
+            "gpu_count": gpu_count,
+            "arch": ctx.get("arch"),
+            "prism_width_multiplier": ctx.get("prism_width_multiplier", 1.0),
+        },
+        "seed": int(ctx.get("seed", 0)),
+        "seq_len": seq_len,
+        "micro_batch": micro,
+        "max_steps": int(ctx.get("max_train_steps", 20000)),
+        "cap_s": cap_s,
+        "wall_margin_s": WALL_MARGIN_S,
+        "out_dir": str(out_dir),
+        "parallel": os.environ.get("DENSE1B_PARALLEL", "zero1").strip().lower(),
+    }
+    payload_path = out_dir / PAYLOAD_NAME
+    torch.save(payload, payload_path)
+    port = _rendezvous_port()
+    _set_dist_env(port)
+    print(
+        f"[dense1b] launching dist spawn world={gpu_count} master=127.0.0.1:{port} "
+        f"micro_batch={micro} seq={seq_len} parallel={payload['parallel']}",
+        flush=True,
+    )
+    from nemo_automodel.components.models.dense1b.ddp_worker import spawn_workers
+
+    spawn_workers(gpu_count, port, str(payload_path))
+    metrics_path = out_dir / METRICS_NAME
+    weights_path = out_dir / WEIGHTS_NAME
+    if not metrics_path.is_file() or not weights_path.is_file():
+        raise RuntimeError("DDP workers did not write metrics/weights")
+    metrics = json.loads(metrics_path.read_text(encoding="utf-8"))
+    trained = torch.load(weights_path, map_location="cpu", weights_only=False)
+    model.load_state_dict(trained, strict=False)
+    device = ctx.get("device") or "cuda"
+    if device != "cpu":
+        model.to(device)
+    # Authoritative harness counter + FLOPs spend so MFU is real.
+    tokens = int(metrics.get("tokens_seen") or 0)
+    if hasattr(stream, "tokens_seen"):
+        stream.tokens_seen = int(getattr(stream, "tokens_seen", 0)) + tokens
+        if getattr(stream, "flops_per_token", 0.0):
+            stream.flops_spent = float(stream.flops_per_token) * float(stream.tokens_seen)
+        stream.batches_yielded = int(getattr(stream, "batches_yielded", 0)) + int(
+            metrics.get("train_steps") or 0
+        )
+    return metrics
+
+
+def train(model, ctx):
+    device = ctx.get("device") or "cuda"
+    seed = int(ctx.get("seed", 0))
+    torch.manual_seed(seed)
+    if device == "cuda" and torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+    _enable_fast_matmul()
+    guard = ctx.get("guard")
+    gpu_count = int(ctx.get("gpu_count") or 0)
+    if gpu_count <= 0 and torch.cuda.is_available():
+        gpu_count = int(torch.cuda.device_count())
+    te_available = bool(ctx.get("te_available", False))
+    if not te_available:
+        try:
+            te_available = __import__("importlib").util.find_spec("transformer_engine") is not None
+        except Exception:  # noqa: BLE001
+            te_available = False
+
+    wide = is_96gb_class(ctx, gpu_count)
+    b200 = is_b200_class(ctx, gpu_count)
+    te_raw = os.environ.get("DENSE1B_TE", "").strip().lower()
+    if te_raw in {"0", "false", "no", "off"}:
+        te_want = False
+    elif te_raw in {"1", "true", "yes", "on"}:
+        te_want = True
+    else:
+        te_want = wide or b200
+    rec, te_mode = _maybe_te_recipe() if (te_available and te_want) else (None, "none")
+    n_params = unique_n_params(model)
+    print(
+        f"[dense1b] train start n_params={n_params} "
+        f"({n_params/1e6:.1f}M; floor 850e6 cap 1e9) "
+        f"gpu_count={gpu_count} te_available={te_available} "
+        f"te_mode={te_mode} te_version={_te_version()} "
+        f"cuda_devices={torch.cuda.device_count() if torch.cuda.is_available() else 0} "
+        f"use_te_linear={getattr(model, 'use_te', None)} "
+        f"te_want={te_want} wide96={wide} b200={b200} "
+        f"parallel={os.environ.get('DENSE1B_PARALLEL', 'zero1')}",
+        flush=True,
+    )
+
+    # Never train on an 8× fallback. Cap at profile width (1 / 2 / 4).
+    if b200:
+        default_max = "1"
+    elif wide:
+        default_max = "2"
+    else:
+        default_max = "4"
+    max_gpus = int(os.environ.get("DENSE1B_MAX_GPUS", default_max) or default_max)
+    if gpu_count > max_gpus:
+        print(f"[dense1b] capping visible GPUs {gpu_count} -> {max_gpus}", flush=True)
+        gpu_count = max_gpus
+    if gpu_count <= 1 and device != "cpu" and torch.cuda.is_available():
+        model.to(device)
+    if gpu_count > 1 and torch.cuda.is_available():
+        os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(str(i) for i in range(gpu_count))
+        metrics = _launch_ddp(model, ctx, gpu_count)
+        fpt = float(ctx.get("flops_per_token_probe") or 0.0)
+        elapsed = float(metrics.get("train_seconds") or 1.0)
+        tokens = float(metrics.get("tokens_seen") or 0.0)
+        n_params = float(n_params)
+        loop_f = float(getattr(model, "prism_loop_factor", 1.0) or 1.0)
+        fpt_analytic = 6.0 * n_params * loop_f
+        fpt_source = "probe"
+        if fpt <= 0.0:
+            fpt = fpt_analytic
+            fpt_source = "analytic_6n_loops"
+        mfu = (tokens * fpt) / (_peak_flops_per_gpu(ctx, gpu_count) * gpu_count * elapsed) if fpt > 0 else 0.0
+        metrics["mfu_est"] = mfu
+        metrics["flops_per_token_probe"] = fpt
+        metrics["flops_per_token_analytic"] = fpt_analytic
+        metrics["flops_per_token_source"] = fpt_source
+        metrics["n_params"] = int(n_params)
+        print(
+            f"[dense1b] ddp parent metrics n_params={int(n_params)} "
+            f"world={metrics.get('world_size')} "
+            f"te_mode={metrics.get('te_mode')} tok/s={metrics.get('tokens_per_sec')} "
+            f"mfu_est={mfu*100:.2f}% fpt_src={fpt_source} "
+            f"parallel={metrics.get('parallel_mode')} compile={metrics.get('torch_compile')}",
+            flush=True,
+        )
+        prism_telemetry.finish_evaluation()
+        return metrics
+
+    stream = ctx.get("train_stream")
+    if stream is None:
+        raise RuntimeError("train_stream required for live AutoModel dense 1B")
+    env_micro = os.environ.get("DENSE1B_MICRO_BATCH", "").strip()
+    if env_micro.isdigit():
+        micro = max(1, int(env_micro))
+    elif b200:
+        micro = B200_MICRO_BATCH
+    elif wide:
+        micro = WIDE_MICRO_BATCH
+    else:
+        micro = DEFAULT_MICRO_BATCH
+    if hasattr(stream, "batch_size"):
+        stream.batch_size = max(1, int(micro))
+    if device != "cpu" and hasattr(stream, "device"):
+        stream.device = device
+    compiled, did_compile = _maybe_compile(model)
+    max_steps = int(ctx.get("max_train_steps", 20000))
+    cap_s = float(ctx.get("train_hours_cap", 1.0)) * 3600.0
+    stop_s = max(60.0, cap_s - WALL_MARGIN_S)
+    metrics = _train_loop(
+        compiled,
+        stream,
+        device=device,
+        max_steps=max_steps,
+        stop_s=stop_s,
+        rec=rec,
+        te_mode=te_mode,
+        guard=guard,
+        rank=0,
+        world=1,
+    )
+    elapsed = float(metrics.get("train_seconds") or 1.0)
+    tokens = float(int(getattr(stream, "tokens_seen", 0)) or int(metrics.get("tokens_local") or 0))
+    loop_f = float(getattr(model, "prism_loop_factor", 1.0) or 1.0)
+    fpt_analytic = 6.0 * float(n_params) * loop_f
+    fpt = float(ctx.get("flops_per_token_probe") or 0.0)
+    fpt_source = "probe"
+    if fpt <= 0.0:
+        fpt = fpt_analytic
+        fpt_source = "analytic_6n_loops"
+    mfu = (tokens * fpt) / (_peak_flops_per_gpu(ctx, gpu_count) * max(1, gpu_count) * elapsed) if fpt > 0 else 0.0
+    metrics.update(
+        {
+            "te_mode": te_mode,
+            "te_version": _te_version(),
+            "te_available": te_available,
+            "torch_compile": did_compile,
+            "parallel_mode": "single",
+            "world_size": 1,
+            "backend": "none",
+            "gpu_count": gpu_count,
+            "micro_batch": micro,
+            "tokens_seen": int(tokens),
+            "n_params": n_params,
+            "mfu_est": mfu,
+            "flops_per_token_probe": fpt,
+            "flops_per_token_analytic": fpt_analytic,
+            "flops_per_token_source": fpt_source,
+        }
+    )
+    print(
+        f"[dense1b] train done n_params={n_params} steps={metrics['train_steps']} "
+        f"seconds={metrics['train_seconds']:.1f} te_mode={te_mode} parallel=single "
+        f"micro_batch={micro} tok/s={metrics.get('tokens_per_sec')} mfu_est={mfu*100:.2f}%",
+        flush=True,
+    )
+    prism_telemetry.finish_evaluation()
+    return metrics
diff --git a/nemo_automodel/components/models/dense1b/kernels.py b/nemo_automodel/components/models/dense1b/kernels.py
new file mode 100644
index 00000000..11111111
--- /dev/null
+++ b/nemo_automodel/components/models/dense1b/kernels.py
@@ -0,0 +1,134 @@
+"""Attention / RMS / RoPE / CE helpers for the dense 1B reference.
+
+submission_nonce: dense-1b-b200-20260819T1952Z
+
+Hot-path train (seq=512) is one Flash/SDPA call — no Python token loop.
+"""
+
+from __future__ import annotations
+
+import os
+
+import torch
+import torch.nn.functional as F
+
+ATTN_KERNEL = "sdpa"
+RMS_KERNEL = "torch"
+CE_KERNEL = "torch"
+SWIGLU_KERNEL = "eager"
+ROPE_KERNEL = "torch"
+
+
+def kernel_map():
+    return {
+        "attn_kernel": ATTN_KERNEL,
+        "rmsnorm_kernel": RMS_KERNEL,
+        "ce_kernel": CE_KERNEL,
+        "swiglu_kernel": SWIGLU_KERNEL,
+        "rope_kernel": ROPE_KERNEL,
+    }
+
+
+def enable_attn_backends():
+    global ATTN_KERNEL
+    if not torch.cuda.is_available():
+        ATTN_KERNEL = "math"
+        return
+    try:
+        torch.backends.cuda.enable_flash_sdp(True)
+        torch.backends.cuda.enable_mem_efficient_sdp(True)
+        torch.backends.cuda.enable_math_sdp(True)
+    except Exception:  # noqa: BLE001
+        pass
+    attn = os.environ.get("DENSE1B_ATTN_KERNEL", "").strip().lower()
+    if attn:
+        ATTN_KERNEL = attn
+        return
+    try:
+        import flash_attn  # noqa: F401
+
+        ATTN_KERNEL = "fa2"
+        return
+    except Exception:  # noqa: BLE001
+        pass
+    try:
+        import transformer_engine.pytorch as te  # noqa: F401
+
+        if hasattr(te, "DotProductAttention"):
+            ATTN_KERNEL = "te_avail"
+    except Exception:  # noqa: BLE001
+        pass
+    ATTN_KERNEL = "sdpa"
+
+
+def sdpa(q, k, v, *, is_causal=False, attn_mask=None):
+    """q/k/v: (b, h, t, d). Uses the fastest enabled SDPA backend."""
+    return F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, is_causal=is_causal)
+
+
+def rms_norm(x, weight, eps=1e-6):
+    global RMS_KERNEL
+    RMS_KERNEL = "torch"
+    return F.rms_norm(x, (x.shape[-1],), weight=weight, eps=eps)
+
+
+def apply_rope(x, cos, sin):
+    global ROPE_KERNEL
+    ROPE_KERNEL = "torch"
+    half = x.shape[-1] // 2
+    x1, x2 = x[..., :half], x[..., half:]
+    c = cos.unsqueeze(0).unsqueeze(0)
+    s = sin.unsqueeze(0).unsqueeze(0)
+    return torch.cat([x1 * c - x2 * s, x1 * s + x2 * c], dim=-1)
+
+
+def rope_tables(t, head_dim, theta, device, dtype):
+    inv_freq = 1.0 / (
+        theta ** (torch.arange(0, head_dim, 2, device=device, dtype=torch.float32) / head_dim)
+    )
+    pos = torch.arange(t, device=device, dtype=torch.float32)
+    freqs = torch.outer(pos, inv_freq)
+    return freqs.cos().to(dtype), freqs.sin().to(dtype)
+
+
+_CE_FN = None
+_CE_PROBED = False
+
+
+def _probe_ce():
+    global _CE_FN, _CE_PROBED, CE_KERNEL
+    if _CE_PROBED:
+        return _CE_FN
+    _CE_PROBED = True
+    try:
+        from liger_kernel.transformers.cross_entropy import LigerCrossEntropyLoss  # type: ignore
+
+        _CE_FN = LigerCrossEntropyLoss(reduction="mean")
+        CE_KERNEL = "liger"
+    except Exception:  # noqa: BLE001
+        _CE_FN = None
+        CE_KERNEL = "torch"
+    return _CE_FN
+
+
+def cross_entropy(logits, labels):
+    """logits (N, V) float, labels (N,)."""
+    fn = _probe_ce()
+    if fn is not None:
+        try:
+            return fn(logits, labels)
+        except Exception as exc:  # noqa: BLE001
+            print(f"[dense1b] liger CE failed ({exc}); torch", flush=True)
+    global CE_KERNEL
+    CE_KERNEL = "torch"
+    return F.cross_entropy(logits, labels)
+
+
+def log_kernel_banner():
+    enable_attn_backends()
+    _probe_ce()
+    print(
+        f"[dense1b] kernel_map attn={ATTN_KERNEL} rmsnorm={RMS_KERNEL} "
+        f"ce={CE_KERNEL} swiglu={SWIGLU_KERNEL} rope={ROPE_KERNEL}",
+        flush=True,
+    )
diff --git a/nemo_automodel/components/models/dense1b/model.py b/nemo_automodel/components/models/dense1b/model.py
new file mode 100644
index 00000000..11111111
--- /dev/null
+++ b/nemo_automodel/components/models/dense1b/model.py
@@ -0,0 +1,336 @@
+"""Dense ~975M transformer — Prism recipe 2.1 reference under models/.
+
+submission_nonce: dense-1b-b200-20260819T1952Z
+
+GQA + RMSNorm + SwiGLU + RoPE + QK-norm. Tied embeddings. No MoE, no
+routed experts, no LoopMoE core. Fine-grained MoE at 1B wastes MFU
+(tiny expert GEMMs, irregular routing, no NVFP4 wgrad).
+
+On 32 GB (4×5090) Linear is BF16 + activation ckpt; `DENSE1B_TE=1` opts in.
+On 96 GB-class (2× RTX PRO 6000) and 180 GB-class (1× B200) TE NVFP4
+defaults on and ckpt is off.
+"""
+
+from __future__ import annotations
+
+import math
+import os
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.checkpoint import checkpoint as _activation_checkpoint
+
+try:
+    from nemo_automodel.components.models.dense1b import kernels as _k
+except ImportError:  # local pack / unit tests
+    from . import kernels as _k
+
+DEFAULTS = {
+    "vocab_size": 50257,
+    # ~975.4M unique total (tied embeddings once). Floor 850M / cap 1B.
+    "d_model": 2048,
+    "n_layers": 16,
+    "n_heads": 16,
+    "n_kv_heads": 4,
+    "mlp_hidden": 7168,
+    "window": 2048,
+    "rope_theta": 50000.0,
+    "init_std": 0.02,
+    # 32 GB: ckpt on. 96/180 GB-class / TE: ckpt off unless DENSE1B_CKPT=1.
+    "grad_checkpoint": True,
+}
+
+_OVERRIDE_KEYS = tuple(DEFAULTS.keys())
+_TE_LINEAR = None
+_TE_PROBED = False
+
+
+def is_b200_class(ctx=None, gpu_count=None):
+    """True for ~180–192 GiB NVIDIA B200 (name or device memory)."""
+    ctx = ctx if isinstance(ctx, dict) else {}
+    _ = gpu_count
+    name = str(ctx.get("gpu_type") or os.environ.get("PRISM_GPU_TYPE") or "")
+    if "B200" in name.upper():
+        return True
+    try:
+        if torch.cuda.is_available():
+            mem_gib = torch.cuda.get_device_properties(0).total_memory / float(1024**3)
+            if mem_gib >= 170.0:
+                return True
+    except Exception:  # noqa: BLE001
+        pass
+    return False
+
+
+def is_96gb_class(ctx=None, gpu_count=None):
+    """True for ~96 GiB cards, 2× RTX PRO 6000, or B200 (~180 GiB)."""
+    if is_b200_class(ctx, gpu_count):
+        return True
+    ctx = ctx if isinstance(ctx, dict) else {}
+    count = int(gpu_count if gpu_count is not None else ctx.get("gpu_count") or 0)
+    name = str(ctx.get("gpu_type") or os.environ.get("PRISM_GPU_TYPE") or "")
+    upper = name.upper()
+    if "RTX PRO 6000" in upper or "BLACKWELL SERVER" in upper:
+        return True
+    if count == 2 and "6000" in upper:
+        return True
+    try:
+        if torch.cuda.is_available():
+            mem_gib = torch.cuda.get_device_properties(0).total_memory / float(1024**3)
+            if mem_gib >= 90.0:
+                return True
+    except Exception:  # noqa: BLE001
+        pass
+    return False
+
+
+def _env_flag(name, default=None):
+    raw = os.environ.get(name, "").strip().lower()
+    if raw in {"1", "true", "yes", "on"}:
+        return True
+    if raw in {"0", "false", "no", "off"}:
+        return False
+    return default
+
+
+def unique_n_params(model):
+    """Total unique parameters (tied embeddings counted once)."""
+    seen = set()
+    n = 0
+    for p in model.parameters():
+        key = id(p)
+        if key in seen:
+            continue
+        seen.add(key)
+        n += p.numel()
+    return int(n)
+
+
+def _probe_te_linear():
+    global _TE_LINEAR, _TE_PROBED
+    if _TE_PROBED:
+        return _TE_LINEAR
+    _TE_PROBED = True
+    try:
+        import transformer_engine.pytorch as te  # type: ignore
+
+        _TE_LINEAR = te.Linear
+    except Exception:  # noqa: BLE001 — optional acceleration
+        _TE_LINEAR = None
+    return _TE_LINEAR
+
+
+def _linear(in_f, out_f, *, bias=False, use_te=False):
+    """TE Linear when requested+available; else nn.Linear (BF16-safe).
+
+    NVFP4 block size is 16 — TE Linear with an axis not divisible by 16
+    dies at quantize time.
+    """
+    te_cls = _probe_te_linear() if use_te else None
+    if te_cls is not None and (int(in_f) % 16 == 0) and (int(out_f) % 16 == 0):
+        try:
+            return te_cls(in_f, out_f, bias=bias)
+        except Exception:  # noqa: BLE001
+            pass
+    return nn.Linear(in_f, out_f, bias=bias)
+
+
+class ModelOutput:
+    __slots__ = ("logits",)
+
+    def __init__(self, logits):
+        self.logits = logits
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x):
+        return _k.rms_norm(x, self.weight, eps=self.eps)
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, d_model, hidden, out_dim=None, use_te=False):
+        super().__init__()
+        out_dim = out_dim or d_model
+        self.w1 = _linear(d_model, hidden, use_te=use_te)
+        self.w3 = _linear(d_model, hidden, use_te=use_te)
+        self.w2 = _linear(hidden, out_dim, use_te=use_te)
+
+    def forward(self, x):
+        # NVFP4 block=16; cublasLt SM120 wgrad wants a larger tile (64).
+        n = int(x.shape[0])
+        pad = (64 - n % 64) % 64
+        if pad:
+            x = torch.cat([x, x.new_zeros(pad, *x.shape[1:])], dim=0)
+        y = self.w2(F.silu(self.w1(x)) * self.w3(x))
+        return y[:n] if pad else y
+
+
+class GQAAttention(nn.Module):
+    """Grouped-query attention + RoPE + QK-norm. One SDPA, no token loop."""
+
+    def __init__(self, d_model, n_heads, n_kv_heads, window, rope_theta, use_te=False):
+        super().__init__()
+        if d_model % n_heads != 0:
+            raise ValueError("d_model must divide n_heads")
+        if n_heads % n_kv_heads != 0:
+            raise ValueError("n_heads must divide by n_kv_heads")
+        self.n_heads = int(n_heads)
+        self.n_kv_heads = int(n_kv_heads)
+        self.head_dim = d_model // n_heads
+        self.repeats = self.n_heads // self.n_kv_heads
+        self.window = int(window)
+        self.rope_theta = float(rope_theta)
+        self.wq = _linear(d_model, self.n_heads * self.head_dim, use_te=use_te)
+        self.wk = _linear(d_model, self.n_kv_heads * self.head_dim, use_te=use_te)
+        self.wv = _linear(d_model, self.n_kv_heads * self.head_dim, use_te=use_te)
+        self.wo = _linear(self.n_heads * self.head_dim, d_model, use_te=use_te)
+        self.q_norm = RMSNorm(self.head_dim)
+        self.k_norm = RMSNorm(self.head_dim)
+        self._cos = None
+        self._sin = None
+
+    def _rope(self, q, k):
+        t = q.shape[-2]
+        if self._cos is None or self._cos.shape[0] < t or self._cos.device != q.device:
+            cos, sin = _k.rope_tables(2 * t, self.head_dim, self.rope_theta, q.device, q.dtype)
+            self._cos, self._sin = cos, sin
+        return _k.apply_rope(q, self._cos[:t], self._sin[:t]), _k.apply_rope(
+            k, self._cos[:t], self._sin[:t]
+        )
+
+    def forward(self, x):
+        b, t, d = x.shape
+        hd = self.head_dim
+        q = self.wq(x).view(b, t, self.n_heads, hd)
+        k = self.wk(x).view(b, t, self.n_kv_heads, hd)
+        v = self.wv(x).view(b, t, self.n_kv_heads, hd)
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+        q, k = self._rope(q, k)
+        if self.repeats > 1:
+            k = k.repeat_interleave(self.repeats, dim=1)
+            v = v.repeat_interleave(self.repeats, dim=1)
+        # Train seq=512 <= window=2048: one flash/mem-efficient SDPA.
+        o = _k.sdpa(q, k, v, is_causal=True)
+        o = o.transpose(1, 2).reshape(b, t, d)
+        return self.wo(o)
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, cfg, use_te=False):
+        super().__init__()
+        d = int(cfg["d_model"])
+        self.norm1 = RMSNorm(d)
+        self.attn = GQAAttention(
+            d,
+            int(cfg["n_heads"]),
+            int(cfg["n_kv_heads"]),
+            int(cfg["window"]),
+            float(cfg["rope_theta"]),
+            use_te=use_te,
+        )
+        self.norm2 = RMSNorm(d)
+        self.mlp = SwiGLU(d, int(cfg["mlp_hidden"]), use_te=use_te)
+
+    def forward(self, x):
+        x = x + self.attn(self.norm1(x))
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class DenseTransformer(nn.Module):
+    def __init__(self, cfg, use_te=False):
+        super().__init__()
+        self.cfg = dict(cfg)
+        self.use_te = bool(use_te)
+        d = int(cfg["d_model"])
+        self.tok_emb = nn.Embedding(int(cfg["vocab_size"]), d)
+        self.layers = nn.ModuleList(
+            TransformerBlock(cfg, use_te=use_te) for _ in range(int(cfg["n_layers"]))
+        )
+        self.norm = RMSNorm(d)
+        self.head = nn.Linear(d, int(cfg["vocab_size"]), bias=False)
+        self.head.weight = self.tok_emb.weight
+        self.logits = None
+        self.grad_checkpoint = bool(cfg.get("grad_checkpoint", False)) and not use_te
+        self.prism_loop_factor = 1.0
+        self.prism_active_param_fraction = 1.0
+        self._init_weights(float(cfg["init_std"]))
+
+    def _init_weights(self, std):
+        n_eff = max(1, len(self.layers))
+        for name, p in self.named_parameters():
+            if p.ndim >= 2:
+                if name.endswith(("wo.weight", "w2.weight")):
+                    nn.init.normal_(p, mean=0.0, std=std / math.sqrt(2 * n_eff))
+                else:
+                    nn.init.normal_(p, mean=0.0, std=std)
+
+    def _run_block(self, block, x):
+        ckpt = self.grad_checkpoint and torch.is_grad_enabled()
+        if ckpt:
+            return _activation_checkpoint(block, x, use_reentrant=False)
+        return block(x)
+
+    def forward(self, ids):
+        x = self.tok_emb(ids)
+        for block in self.layers:
+            x = self._run_block(block, x)
+        x = self.norm(x)
+        logits = self.head(x)
+        self.logits = logits
+        return logits
+
+
+def _config_from_ctx(ctx):
+    cfg = dict(DEFAULTS)
+    if isinstance(ctx, dict):
+        overrides = ctx.get("arch")
+        if isinstance(overrides, dict):
+            cfg.update({k: v for k, v in overrides.items() if k in cfg})
+        for k in _OVERRIDE_KEYS:
+            if k in ctx:
+                cfg[k] = ctx[k]
+        mult = float(ctx.get("prism_width_multiplier", 1.0) or 1.0)
+        if abs(mult - 1.0) > 1e-12:
+            if mult <= 0:
+                raise ValueError("prism_width_multiplier must be > 0")
+            for key in ("d_model", "mlp_hidden"):
+                cfg[key] = max(1, int(round(int(cfg[key]) * mult)))
+            head_dim = int(DEFAULTS["d_model"]) // int(DEFAULTS["n_heads"])
+            if head_dim > 0 and cfg["d_model"] % head_dim == 0:
+                cfg["n_heads"] = cfg["d_model"] // head_dim
+            elif cfg["d_model"] % int(cfg["n_heads"]) != 0:
+                h = min(int(cfg["n_heads"]), cfg["d_model"])
+                while h > 1 and cfg["d_model"] % h != 0:
+                    h -= 1
+                cfg["n_heads"] = h
+            kv = int(cfg["n_kv_heads"])
+            if cfg["n_heads"] % kv != 0:
+                while kv > 1 and cfg["n_heads"] % kv != 0:
+                    kv -= 1
+                cfg["n_kv_heads"] = max(1, kv)
+    return cfg
+
+
+def build_dense1b(ctx):
+    ctx = ctx if isinstance(ctx, dict) else {}
+    torch.manual_seed(int(ctx.get("seed", 0)))
+    wide = is_96gb_class(ctx)
+    b200 = is_b200_class(ctx)
+    # 32 GB: TE off unless DENSE1B_TE=1. 96/180 GB-class: TE on unless off.
+    env_te = _env_flag("DENSE1B_TE", default=True if (wide or b200) else False)
+    te_flag = bool(env_te) and _probe_te_linear() is not None
+    cfg = _config_from_ctx(ctx)
+    ckpt = _env_flag("DENSE1B_CKPT", default=False if (wide or b200 or te_flag) else True)
+    cfg["grad_checkpoint"] = bool(ckpt) and not te_flag
+    return DenseTransformer(cfg, use_te=te_flag)
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 00000000..11111111
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,4 @@
+# 32 GB: BF16 + ckpt. 96 GB-class / 1× B200: TE NVFP4 default (DENSE1B_TE=0 to off).
+# Do not install AutoModel pyproject (debian blinker has no RECORD).
+# Do not put pip option lines here.
+# submission_nonce: dense-1b-b200-20260819T1952Z