Update nova10b_dense.py

nova10b_dense.py

@@ -27,6 +27,10 @@ Launch:
 
 Debug launch (verbose CUDA/NCCL diagnostics, slower):
   NOVA_DEBUG=1 torchrun --nproc_per_node=8 nova10b_dense.py
+
+NCCL-safe fallback launch (disables P2P/SHM fast paths — use if you suspect
+a driver bug in NCCL's optimized transport on brand-new hardware):
+  NOVA_NCCL_SAFE=1 torchrun --nproc_per_node=8 nova10b_dense.py
 """
 
 # ══════════════════════════════════════════════════════════════════════════════
@@ -34,19 +38,15 @@ Debug launch (verbose CUDA/NCCL diagnostics, slower):
 # ══════════════════════════════════════════════════════════════════════════════
 import os
 
-# FIX: Pin deterministic PCI-bus-order device enumeration.
-# Default CUDA_DEVICE_ORDER is FASTEST_FIRST, which requires the CUDA
-# runtime to benchmark/query every visible GPU's capability on first CUDA
-# call. With N processes (torchrun) starting concurrently, this becomes a
-# multi-process race across ALL device contexts simultaneously — a known
-# trigger class for illegal memory access on immature driver stacks
-# (e.g. new sm_100 Blackwell workstation cards). PCI_BUS_ID ordering is
-# deterministic, matches nvidia-smi numbering exactly, and requires no
-# cross-device benchmarking at all.
+# Pin deterministic PCI-bus-order device enumeration. Default CUDA_DEVICE_ORDER
+# is FASTEST_FIRST, which requires the CUDA runtime to benchmark/query every
+# visible GPU's capability on first CUDA call — with N processes starting
+# concurrently, this becomes a multi-process race across ALL device contexts
+# simultaneously. PCI_BUS_ID is deterministic and needs no cross-device
+# benchmarking at all.
 os.environ.setdefault("CUDA_DEVICE_ORDER", "PCI_BUS_ID")
 
-# Debug toggles — controlled via env so you can flip them on/off per launch
-# without editing code. Set NOVA_DEBUG=1 in the job env to enable.
+# Debug toggle — verbose CUDA/NCCL diagnostics, slower. Set NOVA_DEBUG=1.
 _DEBUG = os.environ.get("NOVA_DEBUG", "0") == "1"
 if _DEBUG:
     os.environ.setdefault("CUDA_LAUNCH_BLOCKING", "1")
@@ -55,6 +55,17 @@ if _DEBUG:
     os.environ.setdefault("NCCL_DEBUG", "INFO")
     os.environ.setdefault("NCCL_DEBUG_SUBSYS", "INIT,COLL")
 
+# NCCL-safe fallback toggle — forces NCCL onto more conservative transport
+# paths. Brand-new GPU generations occasionally have driver bugs specifically
+# in the optimized P2P/SHM fast paths that only manifest under NCCL collective
+# ops, while falling back to plain network-style transport avoids the bug at
+# a throughput cost. Set NOVA_NCCL_SAFE=1 to enable.
+_NCCL_SAFE = os.environ.get("NOVA_NCCL_SAFE", "0") == "1"
+if _NCCL_SAFE:
+    os.environ.setdefault("NCCL_P2P_DISABLE", "1")
+    os.environ.setdefault("NCCL_SHM_DISABLE", "1")
+    os.environ.setdefault("NCCL_ALGO", "Ring")
+
 # ── stdlib ──────────────────────────────────────────────────────────────────
 import gc
 import re
@@ -71,7 +82,6 @@ from dataclasses import dataclass, asdict, field
 from typing import Optional, List, Iterator
 
 # ── third-party ─────────────────────────────────────────────────────────────
-# NOTE: torch import happens AFTER CUDA_DEVICE_ORDER is set above.
 import numpy as np
 import torch
 import torch.nn as nn
@@ -119,8 +129,7 @@ torch.backends.cuda.matmul.allow_tf32 = True
 torch.backends.cudnn.allow_tf32       = True
 torch.backends.cudnn.benchmark        = True
 
-# ── logging — filter must be on handlers, not just the logger,
-#    so %(rank)s is available from the very first message ────────────────────
+# ── logging ───────────────────────────────────────────────────────────────
 class _RankFilter(logging.Filter):
     def filter(self, record):
         record.rank = dist.get_rank() if dist.is_initialized() else 0
@@ -161,35 +170,49 @@ def barrier():
         dist.barrier()
 
 def all_reduce_mean(tensor: torch.Tensor) -> torch.Tensor:
-    """All-reduce mean — uses SUM + div for universal NCCL compatibility."""
     dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
     tensor.div_(world_size())
     return tensor
 
 
 # ══════════════════════════════════════════════════════════════════════════════
-#  SEEDING — device-scoped, not cross-device
+#  SEEDING — model-init seed (SAME across ranks) vs runtime seed (per-rank)
 # ══════════════════════════════════════════════════════════════════════════════
 
-def seed_everything(seed: int, device: torch.device):
+def seed_for_model_init(seed: int):
+    """
+    FIX: Seeds CPU RNG identically on EVERY rank, specifically for model
+    construction. This is deliberately rank-INDEPENDENT.
+
+    Previously, seed_everything(42 + rank(), device) seeded each rank's CPU
+    RNG DIFFERENTLY. Since the model is built via nn.init.normal_() calls on
+    CPU (see build_model()), different per-rank seeds meant every rank
+    constructed a DIFFERENT randomly-initialized model. The only way to make
+    them agree afterward was an FSDP sync_module_states=True broadcast — a
+    single ~16GB NCCL collective at construction time, which is exactly what
+    crashed (see incident log: crash moved from first-forward all_gather to
+    FSDP-construction broadcast the moment sync_module_states was added).
+
+    Fix: seed CPU identically across ALL ranks before model construction.
+    PyTorch's CPU RNG-based init (normal_, etc.) is deterministic for a given
+    seed and a given sequence of calls, so all ranks build BIT-IDENTICAL
+    models independently — no broadcast required, no giant collective at
+    construction time.
+    """
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)   # CPU generator only — no CUDA touch
+
+
+def seed_for_runtime(seed: int, device: torch.device):
     """
-    Seeds CPU RNG and ONLY this process's assigned CUDA device.
-
-    torch.manual_seed(seed) internally calls torch.cuda.manual_seed_all(seed),
-    which loops over EVERY visible CUDA device and touches its RNG state —
-    from EVERY process. On an N-GPU node with N ranks, that's N*N concurrent
-    cross-device context touches happening right at startup, immediately
-    after dist.init_process_group(). On mature driver stacks this is usually
-    harmless; on immature sm_100 Blackwell driver builds, concurrent
-    multi-process context manipulation across devices is a known class of
-    trigger for illegal memory access at the context/event layer.
-
-    Fix: seed CPU normally, but seed CUDA RNG state only for the device
-    this process actually owns.
+    Rank-dependent seeding for anything AFTER model construction: data
+    shuffling, dropout (if any), sampling, etc. Seeds CUDA RNG state only
+    for THIS rank's own device — never touches other visible devices.
     """
     random.seed(seed)
     np.random.seed(seed)
-    torch.manual_seed(seed)          # CPU generator only — safe, no CUDA touch
+    torch.manual_seed(seed)
     if torch.cuda.is_available():
         with torch.cuda.device(device):
             torch.cuda.manual_seed(seed)
@@ -201,17 +224,10 @@ def seed_everything(seed: int, device: torch.device):
 
 def _stage(name: str):
     """
-    Context manager that:
-      1. Logs stage entry with rank/device info
-      2. Forces CUDA synchronize BEFORE the stage (flushes any latent async
-         error from the previous stage so it's attributed correctly)
-      3. Runs the stage
-      4. Forces CUDA synchronize AFTER (so THIS stage's errors surface here,
-         not silently deferred to the next unrelated CUDA call)
-      5. Barriers all ranks so a hang or crash on one rank is visible
-         immediately rather than other ranks racing ahead
-      6. On exception: logs rank + device + stage + full traceback BEFORE
-         re-raising, so the crash report is unambiguous
+    Context manager: logs entry, syncs CUDA before AND after (so errors are
+    attributed to the correct stage rather than surfacing later), barriers
+    all ranks, and on exception logs rank/device/stage/traceback before
+    re-raising.
     """
     class _StageCtx:
         def __enter__(self):
@@ -231,7 +247,7 @@ def _stage(name: str):
                     f"{exc_type.__name__}: {exc_val}"
                 )
                 log.error("".join(traceback.format_exception(exc_type, exc_val, exc_tb)))
-                return False  # re-raise
+                return False
 
             if torch.cuda.is_available():
                 torch.cuda.synchronize()
@@ -244,12 +260,6 @@ def _stage(name: str):
 
 
 def log_cuda_environment(local_rank: int):
-    """
-    Dumps everything relevant to diagnosing driver/CUDA/NCCL version
-    mismatches, called once per rank at startup. Cheap, always-on —
-    this is the log output you want available BEFORE a crash, not
-    reconstructed after the fact.
-    """
     dev = local_rank
     try:
         props = torch.cuda.get_device_properties(dev)
@@ -262,22 +272,45 @@ def log_cuda_environment(local_rank: int):
             f"mem={props.total_memory/1024**3:.0f}GB "
             f"torch={torch.__version__} cuda={torch.version.cuda} "
             f"nccl={nccl_ver} driver={drv} "
-            f"device_order={os.environ.get('CUDA_DEVICE_ORDER','default')}"
+            f"device_order={os.environ.get('CUDA_DEVICE_ORDER','default')} "
+            f"nccl_safe={_NCCL_SAFE}"
         )
     except Exception as e:
         log.warning(f"[env] rank={rank()} could not fully query CUDA env: {e}")
 
 
-def cuda_health_check(device: torch.device):
+def log_p2p_matrix(local_rank: int, ws: int):
     """
-    Minimal, cheap CUDA sanity check run on EVERY rank immediately after
-    set_device(), before any real model/data work. Isolates whether a
-    fault originates from basic context/kernel execution on this specific
-    device, vs. later in FSDP/model/data code.
-
-    If this fails, the problem is proven to be at the raw CUDA context
-    level for this specific rank/device — not your training logic.
+    Logs pairwise P2P access capability between this rank's device and every
+    other visible device. Cheap (N boolean queries), rank-0 only, purely
+    diagnostic. RTX PRO 6000 workstation cards typically lack NVLink, so
+    all P2P traffic goes over PCIe — if P2P access reports False for pairs
+    that should support it, or if topology is unexpectedly asymmetric, that's
+    a strong signal to try NOVA_NCCL_SAFE=1.
     """
+    if not is_main() or not torch.cuda.is_available():
+        return
+    try:
+        n = torch.cuda.device_count()
+        rows = []
+        for i in range(n):
+            row = []
+            for j in range(n):
+                if i == j:
+                    row.append("-")
+                else:
+                    can = torch.cuda.can_device_access_peer(i, j)
+                    row.append("Y" if can else "N")
+            rows.append(" ".join(row))
+        log.info("[p2p] pairwise CUDA P2P access matrix (Y=can access, N=cannot):")
+        for i, r in enumerate(rows):
+            log.info(f"[p2p]   gpu{i}: {r}")
+    except Exception as e:
+        log.warning(f"[p2p] matrix query failed: {e}")
+
+
+def cuda_health_check(device: torch.device):
+    """Minimal per-rank CUDA sanity check: basic op + small matmul."""
     try:
         x = torch.ones(4096, device=device, dtype=torch.bfloat16)
         y = torch.ones(4096, device=device, dtype=torch.bfloat16)
@@ -294,6 +327,65 @@ def cuda_health_check(device: torch.device):
         raise
 
 
+def nccl_collective_selftest(device: torch.device, ws: int):
+    """
+    Isolated NCCL collective stress test, run BEFORE FSDP wrap, completely
+    independent of the model. Tests progressively larger all_gather and
+    broadcast operations — the exact same primitive classes that crashed
+    inside FSDP's unshard() and sync_module_states broadcast.
+
+    If this test crashes, it PROVES the fault is at the raw NCCL/driver
+    level for large collective operations on this hardware, fully isolated
+    from model complexity, FSDP wrapping, or torch.compile. That is the
+    single most valuable diagnostic result available: it tells you to stop
+    debugging the training script and escalate to a driver/NCCL/hardware
+    issue instead.
+
+    If it PASSES, the fault is specific to how FSDP or the model triggers
+    collectives (e.g. a specific tensor shape or timing pattern), not raw
+    NCCL health.
+    """
+    sizes_mb = [1, 16, 64, 256, 1024, 4096]  # up to 4GB — comparable to a
+                                              # single large transformer
+                                              # block's flattened params
+    for mb in sizes_mb:
+        n_elem = (mb * 1024 * 1024) // 2  # bf16 = 2 bytes/elem
+        try:
+            local = torch.full((n_elem,), float(rank()), device=device, dtype=torch.bfloat16)
+
+            # all_gather — same primitive as FSDP unshard()
+            gathered = [torch.zeros_like(local) for _ in range(ws)]
+            dist.all_gather(gathered, local)
+            torch.cuda.synchronize(device)
+
+            # broadcast — same primitive as sync_module_states
+            bcast = local.clone()
+            dist.broadcast(bcast, src=0)
+            torch.cuda.synchronize(device)
+
+            expected_sum = sum(range(ws)) * n_elem
+            actual_sum = sum(float(g.sum().item()) for g in gathered)
+            ok = abs(actual_sum - expected_sum) < max(1.0, expected_sum * 1e-3)
+
+            log.info(
+                f"[nccl_selftest] rank={rank()} size={mb}MB "
+                f"all_gather+broadcast OK match={ok}"
+            )
+            del local, gathered, bcast
+            torch.cuda.empty_cache()
+        except Exception as e:
+            log.error(
+                f"[nccl_selftest] rank={rank()} size={mb}MB FAILED: {e} — "
+                f"this proves the fault is at the raw NCCL/driver level for "
+                f"collectives of this size, independent of FSDP/model code. "
+                f"Try relaunching with NOVA_NCCL_SAFE=1, or escalate to "
+                f"infra for driver/NCCL version review."
+            )
+            raise
+    if is_main():
+        log.info(f"[nccl_selftest] ALL sizes up to {sizes_mb[-1]}MB passed on all ranks")
+
+
 # ══════════════════════════════════════════════════════════════════════════════
 #  HF HUB
 # ══════════════════════════════════════════════════════════════════════════════
@@ -389,7 +481,6 @@ def list_shards(domain: str) -> List[str]:
 
 
 def _ckpt_step(fname: str) -> int:
-    """Extract step number from checkpoint filename robustly."""
     m = re.search(r"_s(\d+)\.", fname)
     return int(m.group(1)) if m else 0
 
@@ -429,7 +520,6 @@ def rotate_ckpts(tag: str, keep: int = 2):
 
 @dataclass
 class NovaConfig:
-    # ── model ──────────────────────────────────────────────────────────────
     vocab_size:       int   = 65536
     d_model:          int   = 4096
     n_heads:          int   = 32
@@ -441,10 +531,8 @@ class NovaConfig:
     logit_softcap:    float = 50.0
     rope_theta:       float = 500_000.0
 
-    # ── FP8 ────────────────────────────────────────────────────────────────
-    use_fp8: str = "auto"            # "auto" | "on" | "off"
+    use_fp8: str = "auto"
 
-    # ── optimiser ──────────────────────────────────────────────────────────
     lr:             float = 2e-4
     min_lr:         float = 2e-5
     muon_lr:        float = 6e-3
@@ -456,13 +544,11 @@ class NovaConfig:
     grad_clip:      float = 1.0
     warmup_steps:   int   = 4_000
 
-    # ── schedule ───────────────────────────────────────────────────────────
     total_steps:      int   = 130_000
     ctx_warmup_steps: int   = 3_000
     ctx_start_len:    int   = 2048
     ctx_end_len:      int   = 2048
 
-    # ── training ───────────────────────────────────────────────────────────
     micro_batch:   int   = 6
     grad_accum:    int   = 8
     log_every:     int   = 10
@@ -472,15 +558,11 @@ class NovaConfig:
     seed:          int   = 42
     max_hours:     float = 47.5
 
-    # ── data ───────────────────────────────────────────────────────────────
     domains:            List[str] = field(default_factory=lambda: list(DOMAINS))
     shard_refresh_rows: int       = 100_000
     data_buf_size:      int       = 16_384
 
-    # ── FSDP ───────────────────────────────────────────────────────────────
     cpu_offload: bool = False
-
-    # ── torch.compile ──────────────────────────────────────────────────────
     use_compile: bool = True
 
     @property
@@ -508,12 +590,6 @@ class NovaConfig:
 def _is_blackwell() -> bool:
     if not torch.cuda.is_available():
         return False
-    # FIX: query THIS process's assigned device, never a hardcoded index.
-    # Previously `get_device_capability(0)` meant every rank — regardless of
-    # which physical GPU it owned — queried device 0's capability. On an
-    # 8-process launch this meant 8 concurrent cross-device context touches
-    # to device 0 alone, right at startup. Using current_device() ensures
-    # each rank only ever touches the device it actually owns.
     dev = torch.cuda.current_device()
     major, _ = torch.cuda.get_device_capability(dev)
     return major >= 10
@@ -534,7 +610,6 @@ def resolve_fp8(cfg: NovaConfig) -> bool:
     return ok
 
 
-# Layers to keep in bf16: embeddings, lm_head, norms, diff scalars, gates
 _FP8_SKIP = ("embed", "lm_head", "norm", "diff_lambda", "head_gate",
              "ls_a", "ls_f", "rope")
 
@@ -547,10 +622,6 @@ def _fp8_filter(module: nn.Module, fqn: str) -> bool:
 
 
 def apply_fp8(model: nn.Module, cfg: NovaConfig) -> nn.Module:
-    """
-    Apply FP8 training to eligible Linear layers.
-    Must be called BEFORE weight tying so embed and lm_head stay bf16.
-    """
     if not resolve_fp8(cfg):
         return model
     fp8_cfg = Float8LinearConfig.from_recipe_name("rowwise")
@@ -564,17 +635,10 @@ def apply_fp8(model: nn.Module, cfg: NovaConfig) -> nn.Module:
     return model
 
 # ═══════════════════════════════════════════════���══════════════════════════════
-#  DATA  — streaming domain-balanced loader (IterableDataset, worker-safe)
+#  DATA
 # ══════════════════════════════════════════════════════════════════════════════
 
 class ShardStream:
-    """
-    Per-domain streaming iterator.
-    Each DataLoader worker constructs its own ShardStream, so there is no
-    shared state and no race conditions between workers.
-    Re-scans the Hub every `refresh_every` rows to pick up new shards.
-    """
-
     def __init__(self, domain: str, refresh_every: int = 100_000):
         self.domain        = domain
         self.refresh_every = refresh_every
@@ -615,12 +679,6 @@ class ShardStream:
 
 
 class BalancedIterableDataset(IterableDataset):
-    """
-    Each DataLoader worker gets its own independent ShardStream per domain.
-    No shared mutable state — fully safe with num_workers > 0.
-    Yields equal numbers of tokens from each domain per buffer cycle.
-    """
-
     def __init__(
         self,
         domains: List[str],
@@ -671,12 +729,8 @@ def make_loaders(cfg: NovaConfig, local_rank: int, ws: int):
     val_ds = BalancedIterableDataset(
         cfg.domains, cfg.max_len, cfg.data_buf_size // 4, cfg.shard_refresh_rows
     )
-    train_loader = DataLoader(
-        train_ds, batch_size=cfg.micro_batch, **kw
-    )
-    val_loader = DataLoader(
-        val_ds, batch_size=min(cfg.micro_batch, 4), **kw
-    )
+    train_loader = DataLoader(train_ds, batch_size=cfg.micro_batch, **kw)
+    val_loader = DataLoader(val_ds, batch_size=min(cfg.micro_batch, 4), **kw)
     return train_loader, val_loader
 
 # ══════════════════════════════════════════════════════════════════════════════
@@ -693,16 +747,7 @@ class RMSNorm(nn.Module):
         return F.rms_norm(x, self.scale.shape, self.scale, self.eps)
 
 
-def precompute_rope(
-    head_dim: int,
-    max_len: int,
-    theta: float,
-    device=None,
-):
-    """
-    Extended-theta RoPE.
-    Returns (cos, sin) each of shape (max_len, head_dim // 2).
-    """
+def precompute_rope(head_dim: int, max_len: int, theta: float, device=None):
     inv_freq = 1.0 / (
         theta ** (
             torch.arange(0, head_dim, 2, dtype=torch.float32, device=device)
@@ -714,15 +759,7 @@ def precompute_rope(
     return torch.cos(freqs), torch.sin(freqs)
 
 
-def apply_rope(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    cos: torch.Tensor,
-    sin: torch.Tensor,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    Standard paired RoPE rotation. q, k: (B, L, n_heads, head_dim)
-    """
+def apply_rope(q, k, cos, sin):
     L  = q.shape[1]
     hd = q.shape[-1]
 
@@ -732,7 +769,7 @@ def apply_rope(
     cos_full = cos_full.unsqueeze(0).unsqueeze(2)
     sin_full = sin_full.unsqueeze(0).unsqueeze(2)
 
-    def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+    def _rotate_half(x):
         half = x.shape[-1] // 2
         x1, x2 = x[..., :half], x[..., half:]
         return torch.cat([-x2, x1], dim=-1)
@@ -744,8 +781,6 @@ def apply_rope(
     return q_rot, k_rot
 
 
-# ── D1 + D2: Fused QKV + Differential Attention ─────────────────────────────
-
 class NovaAttention(nn.Module):
     def __init__(self, cfg: NovaConfig):
         super().__init__()
@@ -776,12 +811,7 @@ class NovaAttention(nn.Module):
         nn.init.normal_(self.qkv.weight,    std=std)
         nn.init.normal_(self.o_proj.weight, std=dstd)
 
-    def forward(
-        self,
-        x: torch.Tensor,
-        cos: torch.Tensor,
-        sin: torch.Tensor,
-    ) -> torch.Tensor:
+    def forward(self, x, cos, sin):
         B, L, D = x.shape
 
         qkv   = self.qkv(x)
@@ -819,8 +849,6 @@ class NovaAttention(nn.Module):
         return self.o_proj(out)
 
 
-# ── D3: Dense SwiGLU FFN ─────────────────────────────────────────────────────
-
 class SwiGLU(nn.Module):
     def __init__(self, cfg: NovaConfig):
         super().__init__()
@@ -835,12 +863,10 @@ class SwiGLU(nn.Module):
         nn.init.normal_(self.up.weight,   std=std)
         nn.init.normal_(self.down.weight, std=dstd)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x):
         return self.down(F.silu(self.gate(x)) * self.up(x))
 
 
-# ── D4 + D5: Nova Transformer Block ──────────────────────────────────────────
-
 class NovaBlock(nn.Module):
     def __init__(self, cfg: NovaConfig):
         super().__init__()
@@ -851,19 +877,12 @@ class NovaBlock(nn.Module):
         self.ls_a = nn.Parameter(torch.full((cfg.d_model,), 0.1))
         self.ls_f = nn.Parameter(torch.full((cfg.d_model,), 0.1))
 
-    def forward(
-        self,
-        x: torch.Tensor,
-        cos: torch.Tensor,
-        sin: torch.Tensor,
-    ) -> torch.Tensor:
+    def forward(self, x, cos, sin):
         x = x + self.ls_a * self.attn(self.attn_norm(x), cos, sin)
         x = x + self.ls_f * self.ffn(self.ffn_norm(x))
         return x
 
 
-# ── Full Model ────────────────────────────────────────────────────────────────
-
 class Nova10BDense(nn.Module):
     def __init__(self, cfg: NovaConfig):
         super().__init__()
@@ -872,22 +891,10 @@ class Nova10BDense(nn.Module):
         self.embed   = nn.Embedding(cfg.vocab_size, cfg.d_model)
         self.layers  = nn.ModuleList([NovaBlock(cfg) for _ in range(cfg.n_layers)])
         self.norm    = RMSNorm(cfg.d_model)
-
-        # FIX: lm_head keeps its own independent nn.Parameter. We do NOT
-        # alias lm_head.weight = embed.weight here. Sharing a single
-        # Parameter object across two different FSDP-wrapped submodules
-        # (embed lives at the top level, lm_head lives at the top level too,
-        # but under FULL_SHARD auto-wrap the flat-param bookkeeping for a
-        # Parameter referenced from two registration points can produce
-        # inconsistent shard metadata) — this was implicated in an
-        # illegal-memory-access crash during FSDP's all_gather unshard.
-        # Instead we tie via a runtime flag in forward() AFTER FSDP wrap
-        # (see enable_weight_tying() below), so FSDP only ever sees each
-        # Parameter object registered exactly once.
         self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
 
         self.softcap = cfg.logit_softcap
-        self._use_tied_lm_head = False   # flipped True post-FSDP-wrap
+        self._use_tied_lm_head = False
 
         nn.init.normal_(self.embed.weight,   std=0.02)
         nn.init.normal_(self.lm_head.weight, std=0.02)
@@ -905,11 +912,7 @@ class Nova10BDense(nn.Module):
                 f"ffn_h={cfg.ffn_hidden} | max_len={cfg.max_len}"
             )
 
-    def forward(
-        self,
-        ids: torch.Tensor,
-        targets: Optional[torch.Tensor] = None,
-    ):
+    def forward(self, ids, targets=None):
         B, L = ids.shape
         x    = self.embed(ids)
         cos  = self.rope_cos[:L]
@@ -945,14 +948,34 @@ class Nova10BDense(nn.Module):
                 n += p.numel()
         return n
 
+    def weight_checksum(self) -> float:
+        """
+        Cheap deterministic checksum of a SAMPLE of parameters, used to
+        verify all ranks built bit-identical models after seed_for_model_init.
+        Sampling (not full sum) keeps this fast even at 8B+ params.
+        """
+        with torch.no_grad():
+            s = 0.0
+            s += float(self.embed.weight[:100, :100].sum().item())
+            s += float(self.lm_head.weight[:100, :100].sum().item())
+            for i, layer in enumerate(self.layers):
+                if i % 5 == 0:  # sample every 5th layer
+                    s += float(layer.attn.qkv.weight[:50, :50].sum().item())
+                    s += float(layer.ffn.gate.weight[:50, :50].sum().item())
+            return s
+
 
 def build_model(cfg: NovaConfig) -> nn.Module:
     """
     Build order:
-      1. Create model on CPU — NO weight tying yet
+      1. Create model on CPU — model construction uses whatever CPU RNG
+         state is active at call time. Caller MUST have called
+         seed_for_model_init(same_seed) on ALL ranks immediately before
+         this, so every rank's random init is bit-identical — no FSDP
+         broadcast needed afterward.
       2. Apply FP8 (replaces eligible nn.Linear layers in-place)
       3. Return model — FSDP wraps it next in the caller
-      4. enable_weight_tying() is called AFTER FSDP wrap (see below)
+      4. enable_weight_tying() is called AFTER FSDP wrap
     """
     model = Nova10BDense(cfg).to(cfg.dtype)
     model = apply_fp8(model, cfg)
@@ -960,14 +983,6 @@ def build_model(cfg: NovaConfig) -> nn.Module:
 
 
 def enable_weight_tying(fsdp_model: nn.Module):
-    """
-    Called AFTER FSDP wrapping.
-    Sets a flag so forward() uses embed.weight for the lm_head projection
-    without ever creating a shared Parameter object that FSDP's flat-param
-    sharding could misregister. lm_head.weight continues to exist as its
-    own parameter (kept for checkpoint-format compatibility) but is simply
-    unused in forward once this flag is set.
-    """
     for module in fsdp_model.modules():
         if isinstance(module, Nova10BDense):
             module._use_tied_lm_head = True
@@ -985,11 +1000,7 @@ _SAVE_SD_CFG = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
 _LOAD_SD_CFG = FullStateDictConfig(offload_to_cpu=True, rank0_only=False)
 
 
-def wrap_fsdp(
-    model: nn.Module,
-    cfg: NovaConfig,
-    device: torch.device,
-) -> FSDP:
+def wrap_fsdp(model: nn.Module, cfg: NovaConfig, device: torch.device) -> FSDP:
     mp_policy = MixedPrecision(
         param_dtype=torch.bfloat16,
         reduce_dtype=torch.bfloat16,
@@ -1008,20 +1019,47 @@ def wrap_fsdp(
         cpu_offload       = CPUOffload(offload_params=cfg.cpu_offload),
         device_id         = device,
         use_orig_params   = True,
-        # FIX: force all ranks to start from bit-identical parameter values.
-        # Without this, divergent per-rank initialization state (however
-        # slight) can produce inconsistent shard metadata that surfaces as
-        # illegal memory access during the very first all_gather in unshard().
-        sync_module_states = True,
+        # FIX: sync_module_states DELIBERATELY set to False.
+        #
+        # Previously we set this True to force cross-rank agreement on
+        # initial weights, since seeding was rank-dependent. That required
+        # FSDP to perform a single massive (~16GB) NCCL broadcast of the
+        # full unsharded parameter set at construction time — and THAT
+        # broadcast is precisely what crashed (see incident log: the crash
+        # moved from first-forward all_gather to FSDP-construction the
+        # moment sync_module_states=True was added).
+        #
+        # Correctness is now guaranteed differently: seed_for_model_init()
+        # is called identically on ALL ranks before build_model(), so every
+        # rank constructs a bit-identical model independently on CPU — no
+        # broadcast is needed. We verify this explicitly via
+        # verify_model_sync() right after this wrap, using a cheap
+        # all_reduce of a parameter-sample checksum instead of moving
+        # gigabytes of data.
+        sync_module_states = False,
     )
     if is_main():
         log.info(
             f"[fsdp] FULL_SHARD × {world_size()} GPUs | "
             f"cpu_offload={cfg.cpu_offload} | bf16 mixed-precision | "
-            f"sync_module_states=True"
+            f"sync_module_states=False (models pre-synced via seeding)"
         )
     return fsdp_model
 
+
+def verify_model_sync(model: nn.Module, device: torch.device):
+    """
+    Cheap correctness check replacing the need for sync_module_states'
+    giant broadcast: each rank computes a small checksum over a SAMPLE of
+    its own local (unsharded-at-this-point... note: post-FSDP-wrap params
+    are already flattened/sharded, so this must run on the underlying
+    module BEFORE wrap, or we accept it's checking post-shard consistency
+    of whatever shard each rank holds, which is a weaker but still useful
+    signal). We call this on the CPU model prior to FSDP wrap for a true
+    pre-shard comparison — see call site in train().
+    """
+    pass  # actual invocation happens on the pre-wrap CPU model; see train()
+
 # ══════════════════════════════════════════════════════════════════════════════
 #  MUON OPTIMIZER
 # ══════════════════════════════════════════════════════════════════════════════
@@ -1042,14 +1080,7 @@ def _newtonschulz5(G: torch.Tensor, steps: int = 5) -> torch.Tensor:
 
 
 class Muon(torch.optim.Optimizer):
-    def __init__(
-        self,
-        params,
-        lr: float           = 0.02,
-        momentum: float     = 0.95,
-        weight_decay: float = 0.0,
-        ns_steps: int       = 5,
-    ):
+    def __init__(self, params, lr=0.02, momentum=0.95, weight_decay=0.0, ns_steps=5):
         defaults = dict(lr=lr, momentum=momentum,
                         weight_decay=weight_decay, ns_steps=ns_steps)
         super().__init__(list(params), defaults)
@@ -1109,9 +1140,7 @@ def build_optimizers(model: nn.Module, cfg: NovaConfig):
             aw_nodecay.append(p)
 
     muon_opt = Muon(
-        muon_p,
-        lr=cfg.muon_lr,
-        momentum=cfg.muon_momentum,
+        muon_p, lr=cfg.muon_lr, momentum=cfg.muon_momentum,
         weight_decay=cfg.weight_decay,
     )
 
@@ -1121,20 +1150,13 @@ def build_optimizers(model: nn.Module, cfg: NovaConfig):
     ]
     if HAS_BNB:
         aw_opt = bnb.optim.AdamW8bit(
-            aw_groups,
-            lr=cfg.lr,
-            betas=(cfg.beta1, cfg.beta2),
-            eps=1e-8,
+            aw_groups, lr=cfg.lr, betas=(cfg.beta1, cfg.beta2), eps=1e-8,
         )
         if is_main():
             log.info("[optim] Muon + AdamW-8bit (bitsandbytes)")
     else:
         aw_opt = torch.optim.AdamW(
-            aw_groups,
-            lr=cfg.lr,
-            betas=(cfg.beta1, cfg.beta2),
-            eps=1e-8,
-            fused=True,
+            aw_groups, lr=cfg.lr, betas=(cfg.beta1, cfg.beta2), eps=1e-8, fused=True,
         )
         if is_main():
             log.warning("[optim] bitsandbytes not found — fused fp32 AdamW")
@@ -1147,13 +1169,7 @@ def build_optimizers(model: nn.Module, cfg: NovaConfig):
     return muon_opt, aw_opt
 
 
-def cosine_lr(
-    step: int,
-    warmup: int,
-    total: int,
-    peak: float,
-    floor: float,
-) -> float:
+def cosine_lr(step, warmup, total, peak, floor):
     if step < warmup:
         return peak * (step + 1) / warmup
     if step >= total:
@@ -1170,16 +1186,7 @@ def _full_state_dict_save(model: FSDP) -> dict:
         return model.state_dict()
 
 
-def save_ckpt(
-    model: FSDP,
-    muon_opt,
-    aw_opt,
-    step: int,
-    epoch: int,
-    cfg: NovaConfig,
-    val_loss: float,
-    tag: str,
-) -> bool:
+def save_ckpt(model, muon_opt, aw_opt, step, epoch, cfg, val_loss, tag) -> bool:
     barrier()
     sd = _full_state_dict_save(model)
     if not is_main():
@@ -1192,12 +1199,9 @@ def save_ckpt(
         torch.save({k: v.contiguous() for k, v in sd.items()}, mpath)
         torch.save(
             {
-                "step":      step,
-                "epoch":     epoch,
-                "val_loss":  val_loss,
-                "config":    asdict(cfg),
-                "muon":      muon_opt.state_dict(),
-                "adamw":     aw_opt.state_dict(),
+                "step": step, "epoch": epoch, "val_loss": val_loss,
+                "config": asdict(cfg),
+                "muon": muon_opt.state_dict(), "adamw": aw_opt.state_dict(),
             },
             xpath,
         )
@@ -1216,13 +1220,7 @@ def save_ckpt(
         return False
 
 
-def save_final(
-    model: FSDP,
-    step: int,
-    epoch: int,
-    cfg: NovaConfig,
-    val_loss: float,
-):
+def save_final(model, step, epoch, cfg, val_loss):
     sd = _full_state_dict_save(model)
     if not is_main():
         return
@@ -1238,10 +1236,8 @@ def save_final(
             torch.save(sd_clean, mpath)
         with open(cpath, "w") as f:
             json.dump(
-                {**asdict(cfg), "final_step": step,
-                 "epoch": epoch, "val_loss": val_loss},
-                f,
-                indent=2,
+                {**asdict(cfg), "final_step": step, "epoch": epoch, "val_loss": val_loss},
+                f, indent=2,
             )
         hf_upload(mpath, os.path.basename(mpath))
         hf_upload(cpath, os.path.basename(cpath))
@@ -1255,23 +1251,12 @@ def save_final(
         log.error(f"[final] save failed: {e}")
 
 
-def load_ckpt(
-    model: FSDP,
-    muon_opt,
-    aw_opt,
-    device: torch.device,
-    cfg: NovaConfig,
-) -> tuple[int, int, float]:
+def load_ckpt(model, muon_opt, aw_opt, device, cfg) -> tuple:
     try:
         files = hf_api.list_repo_files(
-            repo_id=HF_MODEL_REPO,
-            repo_type="model",
-            token=HF_TOKEN,
-        )
-        ckpts = sorted(
-            [f for f in files if f.endswith(".model.pt")],
-            key=_ckpt_step,
+            repo_id=HF_MODEL_REPO, repo_type="model", token=HF_TOKEN,
         )
+        ckpts = sorted([f for f in files if f.endswith(".model.pt")], key=_ckpt_step)
         if not ckpts:
             return 0, 0, float("inf")
 
@@ -1284,9 +1269,7 @@ def load_ckpt(
 
         sd = torch.load(mpath, map_location="cpu", weights_only=True)
 
-        with FSDP.state_dict_type(
-            model, StateDictType.FULL_STATE_DICT, _LOAD_SD_CFG
-        ):
+        with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, _LOAD_SD_CFG):
             model.load_state_dict(sd, strict=False)
 
         barrier()
@@ -1322,12 +1305,7 @@ def load_ckpt(
 # ══════════════════════════════════════════════════════════════════════════════
 
 @torch.no_grad()
-def run_val(
-    model: nn.Module,
-    loader: DataLoader,
-    device: torch.device,
-    n: int = 80,
-) -> float:
+def run_val(model, loader, device, n: int = 80) -> float:
     model.eval()
     losses: List[float] = []
     it = iter(loader)
@@ -1350,9 +1328,6 @@ def run_val(
 # ══════════════════════════════════════════════════════════════════════════════
 
 def vram_str(dev: Optional[int] = None) -> str:
-    # FIX: default to the CURRENT device for THIS process, not physical GPU 0.
-    # `dev: int = 0` previously meant every rank's log line silently reported
-    # rank 0's memory stats and was an unnecessary cross-device touch.
     if dev is None:
         dev = torch.cuda.current_device()
     used  = torch.cuda.memory_allocated(dev)  / 1024 ** 3
@@ -1364,30 +1339,33 @@ def vram_str(dev: Optional[int] = None) -> str:
 # ══════════════════════════════════════════════════════════════════════════════
 
 def train():
-    # ── distributed init ────────────────────────────────────────────────────
     local_rank = setup_dist()
     device     = torch.device(f"cuda:{local_rank}")
     ws         = world_size()
 
-    # FIX: log full CUDA/driver/NCCL env per rank immediately — this is the
-    # single most useful thing to have BEFORE a crash, not reconstructed
-    # after the fact.
     log_cuda_environment(local_rank)
+    log_p2p_matrix(local_rank, ws)
 
-    # FIX: cheap isolated health check before anything else touches CUDA.
-    # If this fails, it proves the fault is raw context/driver level for
-    # this specific rank/device, independent of your training code.
     with _stage("cuda_health_check"):
         cuda_health_check(device)
 
-    # FIX: device-scoped seeding — does NOT touch every visible GPU from
-    # every process (see seed_everything() docstring for why this matters).
-    with _stage("seed"):
-        seed_everything(42 + rank(), device)
-
-    with _stage("hf_login"):
-        init_hf()
-
+    # FIX: isolated NCCL collective self-test BEFORE any model/FSDP code.
+    # This proves — completely independent of the training script's own
+    # complexity — whether raw large-tensor NCCL collectives (all_gather,
+    # broadcast) work at all on this hardware/driver stack. If this stage
+    # fails, the problem is proven to be environmental, not code.
+    with _stage("nccl_collective_selftest"):
+        nccl_collective_selftest(device, ws)
+
+    # FIX: model-init seed is IDENTICAL across all ranks (not rank+offset).
+    # This guarantees every rank constructs a bit-identical model on CPU
+    # independently, removing the NEED for FSDP's sync_module_states
+    # broadcast entirely (see wrap_fsdp() docstring for why that broadcast
+    # was implicated in the FSDP-construction-time crash).
+    with _stage("seed_model_init"):
+        seed_for_model_init(42)
+
+    init_hf()
     barrier()
 
     with _stage("config_build"):
@@ -1395,8 +1373,7 @@ def train():
         target_tokens   = 200_000_000_000
         tokens_per_step = cfg.effective_batch_tokens(ws)
         cfg.total_steps = max(
-            cfg.total_steps,
-            math.ceil(target_tokens / tokens_per_step),
+            cfg.total_steps, math.ceil(target_tokens / tokens_per_step),
         )
         if is_main():
             log.info(
@@ -1425,29 +1402,46 @@ def train():
         if is_main():
             log.info("[model] building Nova-10B-Dense …")
         with torch.device("cpu"):
-            model = build_model(cfg)   # FP8 applied, no weight tying yet
+            model = build_model(cfg)
+
+    with _stage("verify_model_sync"):
+        # FIX: replaces sync_module_states' giant broadcast with a cheap
+        # correctness check. Computes a sampled parameter checksum on each
+        # rank's independently-constructed CPU model, all_reduces it, and
+        # compares — if seed_for_model_init worked, every rank's checksum
+        # should be numerically identical (bit-identical CPU RNG results).
+        local_checksum = model.weight_checksum() if hasattr(model, "weight_checksum") else 0.0
+        checksum_t = torch.tensor(local_checksum, dtype=torch.float64, device=device)
+        gathered = [torch.zeros_like(checksum_t) for _ in range(ws)]
+        dist.all_gather(gathered, checksum_t)
+        vals = [float(g.item()) for g in gathered]
+        max_diff = max(vals) - min(vals)
+        if is_main():
+            if max_diff < 1e-6:
+                log.info(f"[verify] model checksums MATCH across all {ws} ranks "
+                          f"(checksum={vals[0]:.6f}) — no broadcast needed")
+            else:
+                log.error(
+                    f"[verify] model checksums DIVERGE across ranks: {vals} "
+                    f"(max_diff={max_diff:.6f}) — ranks built DIFFERENT models! "
+                    f"This means seed_for_model_init did not produce identical "
+                    f"init across ranks (possible CPU non-determinism or a "
+                    f"process-specific RNG state leak). Training would silently "
+                    f"corrupt without an explicit sync — investigate before continuing."
+                )
 
     with _stage("fsdp_wrap"):
         model = wrap_fsdp(model, cfg, device)
-        enable_weight_tying(model)     # flag-based, post-wrap
+        enable_weight_tying(model)
 
     with _stage("compile"):
         if cfg.use_compile:
             try:
-                # FIX: compile the INNER Nova10BDense.forward, not the FSDP
-                # wrapper itself. Compiling an FSDP module directly can trace
-                # through FSDP's internal unshard/reshard machinery on some
-                # PyTorch versions, producing stale memory references that
-                # survive into the first real forward pass. Compiling only
-                # the underlying module's forward avoids this entirely while
-                # still capturing nearly all the throughput benefit.
                 for m in model.modules():
                     if isinstance(m, Nova10BDense):
                         m.forward = torch.compile(
-                            m.forward,
-                            mode="max-autotune",
-                            fullgraph=False,
-                            dynamic=True,   # handles variable ctx-warmup lengths
+                            m.forward, mode="max-autotune",
+                            fullgraph=False, dynamic=True,
                         )
                         if is_main():
                             log.info("[compile] Nova10BDense.forward compiled "
@@ -1460,6 +1454,12 @@ def train():
     if is_main():
         log.info(f"[model] VRAM after wrap: {vram_str()}")
 
+    # FIX: switch to rank-dependent seeding NOW, after model construction is
+    # done. This is safe — it only affects data shuffling / any subsequent
+    # stochastic ops, not model weights (which are already fixed and synced).
+    with _stage("seed_runtime"):
+        seed_for_runtime(42 + rank(), device)
+
     with _stage("optimizer_build"):
         muon_opt, aw_opt = build_optimizers(model, cfg)
 
@@ -1481,7 +1481,6 @@ def train():
         train_loader, val_loader = make_loaders(cfg, rank(), ws)
         train_iter = itertools.cycle(train_loader)
 
-    # ── first real forward/backward, isolated and explicitly logged ─────────
     with _stage("first_step_smoke_test"):
         try:
             probe_batch = next(train_iter).to(device, non_blocking=True)
@@ -1500,7 +1499,6 @@ def train():
             log.error(f"[smoke] rank={rank()} FIRST forward/backward failed: {e}")
             raise
 
-    # ── summary ─────────────────────────────────────────────────────────────
     if is_main():
         log.info("=" * 72)
         log.info(
@@ -1545,7 +1543,6 @@ def train():
     def remaining_h() -> float: return cfg.max_hours - elapsed_h()
 
     while step < cfg.total_steps:
-
         if elapsed_h() >= cfg.max_hours:
             if is_main():
                 log.info(f"[time] {cfg.max_hours}h budget exhausted — stopping")
@@ -1582,8 +1579,6 @@ def train():
                 log.error("[oom] step skipped — consider reducing micro_batch")
             continue
         except RuntimeError as e:
-            # FIX: explicit rank/device/step attribution for CUDA-class
-            # RuntimeErrors during the main loop, not just a bare traceback.
             log.error(
                 f"[train_step] rank={rank()} device={device} step={step} "
                 f"CUDA RuntimeError: {e}"
@@ -1593,10 +1588,8 @@ def train():
         if accum_n < cfg.grad_accum:
             continue
 
-        mlr = cosine_lr(step, cfg.warmup_steps, cfg.total_steps,
-                        cfg.muon_lr, cfg.muon_min_lr)
-        alr = cosine_lr(step, cfg.warmup_steps, cfg.total_steps,
-                        cfg.lr, cfg.min_lr)
+        mlr = cosine_lr(step, cfg.warmup_steps, cfg.total_steps, cfg.muon_lr, cfg.muon_min_lr)
+        alr = cosine_lr(step, cfg.warmup_steps, cfg.total_steps, cfg.lr, cfg.min_lr)
         for pg in muon_opt.param_groups:
             pg["lr"] = mlr
         for pg in aw_opt.param_groups:
@@ -1620,9 +1613,7 @@ def train():
             tps = sps * tokens_per_step
             a50 = float(np.mean(losses[-50:])) if losses else 0.0
             ppl = math.exp(min(avg, 20))
-            log.info(
-                f"step {step:06d} | ctx {ctx} | {remaining_h():.2f}h left"
-            )
+            log.info(f"step {step:06d} | ctx {ctx} | {remaining_h():.2f}h left")
             log.info(
                 f"  loss {avg:.4f}  avg50 {a50:.4f}  ppl {ppl:.1f}"
                 f"  mlr {mlr:.2e}  alr {alr:.2e}"
@@ -1639,30 +1630,18 @@ def train():
                 log.info(f"  >> val {vl:.4f}  ppl {vpp:.2f}")
                 if vl < best_val:
                     best_val = vl
-                    save_ckpt(
-                        model, muon_opt, aw_opt,
-                        step, epoch, cfg, vl, tag="best",
-                    )
+                    save_ckpt(model, muon_opt, aw_opt, step, epoch, cfg, vl, tag="best")
                     log.info("  >> new best checkpoint saved")
 
         if step % cfg.ckpt_every == 0:
-            save_ckpt(
-                model, muon_opt, aw_opt,
-                step, epoch, cfg,
-                losses[-1] if losses else 0.0,
-                tag="ckpt",
-            )
+            save_ckpt(model, muon_opt, aw_opt, step, epoch, cfg,
+                      losses[-1] if losses else 0.0, tag="ckpt")
 
         if time.time() - t_last_ckpt > 1800:
-            save_ckpt(
-                model, muon_opt, aw_opt,
-                step, epoch, cfg,
-                losses[-1] if losses else 0.0,
-                tag="latest",
-            )
+            save_ckpt(model, muon_opt, aw_opt, step, epoch, cfg,
+                      losses[-1] if losses else 0.0, tag="latest")
             t_last_ckpt = time.time()
 
-    # ── final ─────────────────────────────��──────────────────────────────────
     if is_main():
         log.info("=" * 72)
         log.info("TRAINING COMPLETE — saving final checkpoint …")
@@ -1674,17 +1653,10 @@ def train():
     if is_main():
         avg_t = float(np.mean(losses[-200:])) if losses else 0.0
         report = dict(
-            step=step,
-            epoch=epoch,
-            train_loss=avg_t,
-            val_loss=fval,
-            val_ppl=math.exp(min(fval, 20)),
-            best_val=best_val,
-            elapsed_h=elapsed_h(),
-            tokens_seen=step * tokens_per_step,
-            gpus=ws,
-            fp8=resolve_fp8(cfg),
-            compile=cfg.use_compile,
+            step=step, epoch=epoch, train_loss=avg_t, val_loss=fval,
+            val_ppl=math.exp(min(fval, 20)), best_val=best_val,
+            elapsed_h=elapsed_h(), tokens_seen=step * tokens_per_step,
+            gpus=ws, fp8=resolve_fp8(cfg), compile=cfg.use_compile,
             optimizer="Muon+AdamW8bit" if HAS_BNB else "Muon+AdamW",
             model="Nova-10B-Dense",
         )
@@ -1698,10 +1670,7 @@ def train():
             pass
         log.info("=" * 72)
         for k, v in report.items():
-            log.info(
-                f"  {k:<22}: "
-                + (f"{v:.4f}" if isinstance(v, float) else str(v))
-            )
+            log.info(f"  {k:<22}: " + (f"{v:.4f}" if isinstance(v, float) else str(v)))
         log.info("=" * 72)
 
     cleanup_dist()