Update nova10b_dense.py

nova10b_dense.py

@@ -1,25 +1,30 @@
 #!/usr/bin/env python3
 """
-Nova-10B-Dense  ·  HuggingFace Jobs  ·  4-8× RTX PRO 6000 Blackwell
+Nova-10B-Dense  ·  HuggingFace Jobs  ·  4-8x RTX PRO 6000 Blackwell
 ~10B fully-dense transformer  (no MoE, no routing, no MoD)
 
 Architecture:
-  D1  Fused QKV         – single Linear for Q+K+V, split after projection
-  D2  Differential Attn – dual-stream diff attn + QK-Norm + head-gate (GQA/RoPE)
-  D3  SwiGLU FFN        – dense gated MLP, no routing
-  D4  RMSNorm           – pre-norm on attn + FFN
-  D5  LayerScale        – per-layer learnable residual scale (0.1 init)
-  D6  RoPE              – extended-theta (500k) rotary embeddings
-  D7  Logit soft-cap    – tanh(x/cap)*cap on final logits
-  D8  FP8               – torchao rowwise (Blackwell sm_100)
+  D1  Fused QKV         - single Linear for Q+K+V, split after projection
+  D2  Differential Attn - dual-stream diff attn + QK-Norm + head-gate (GQA/RoPE)
+  D3  SwiGLU FFN        - dense gated MLP, no routing
+  D4  RMSNorm           - pre-norm on attn + FFN
+  D5  LayerScale        - per-layer learnable residual scale (0.1 init)
+  D6  RoPE              - extended-theta (500k) rotary embeddings
+  D7  Logit soft-cap    - tanh(x/cap)*cap on final logits
+  D8  FP8               - torchao rowwise (Blackwell sm_100)
 
 Training stack:
-  O1  Muon              – Newton-Schulz orthogonalized momentum for 2-D weights
-  O2  AdamW-8bit        – bitsandbytes, for embed/norms/1-D params
-  O3  FSDP2             – PyTorch FullyShardedDataParallel
-  O4  Cosine LR         – separate schedules for Muon vs AdamW
-  D9  Balanced streaming – equal tokens per domain, new shards auto-detected
-  D10 ctx-warmup        – sequence-length curriculum (1024 -> 8192)
+  O1  Muon              - Newton-Schulz orthogonalized momentum for 2-D weights
+  O2  AdamW-8bit        - bitsandbytes, for embed/norms/1-D params
+  O3  FSDP2 SHARD_GRAD_OP (ZeRO-2) - grads + optimizer state sharded,
+                          FULL params replicated per GPU. Required because
+                          Muon's Newton-Schulz needs genuine full 2-D weight
+                          matrices at optimizer-step time; FULL_SHARD
+                          (ZeRO-3) only exposes 1-D local shards and
+                          silently disables Muon.
+  O4  Cosine LR         - separate schedules for Muon vs AdamW
+  D9  Balanced streaming - equal tokens per domain, new shards auto-detected
+  D10 ctx-warmup        - sequence-length curriculum (1024 -> 8192)
 
 Launch:
   torchrun --nproc_per_node=4 nova10b_dense.py
@@ -28,7 +33,6 @@ Launch:
 
 # ── stdlib ──────────────────────────────────────────────────────────────────
 import os
-import gc
 import re
 import json
 import time
@@ -38,7 +42,7 @@ import logging
 import itertools
 import functools
 from dataclasses import dataclass, asdict, field
-from typing import Optional, List, Iterator
+from typing import Optional, List, Iterator, Tuple
 
 # ── third-party ─────────────────────────────────────────────────────────────
 import numpy as np
@@ -86,8 +90,8 @@ 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 — rank-0 gets INFO, workers get ERROR only (prevents spam) ──────
+
 class _RankFilter(logging.Filter):
     def filter(self, record):
         record.rank = dist.get_rank() if dist.is_initialized() else 0
@@ -95,30 +99,40 @@ class _RankFilter(logging.Filter):
 
 _rank_filter = _RankFilter()
 
+_local_rank_for_log = int(os.environ.get("LOCAL_RANK", "0"))
+_log_level = logging.INFO if _local_rank_for_log == 0 else logging.ERROR
+
 logging.basicConfig(
-    level=logging.INFO,
+    level=_log_level,
     format="%(asctime)s | R%(rank)s | %(message)s",
     datefmt="%H:%M:%S",
 )
-# Attach to every handler so %(rank)s is always resolved
 for _h in logging.root.handlers:
     _h.addFilter(_rank_filter)
 
 log = logging.getLogger("nova10b")
 log.addFilter(_rank_filter)
 
+if _local_rank_for_log != 0:
+    logging.getLogger("huggingface_hub").setLevel(logging.WARNING)
+    logging.getLogger("filelock").setLevel(logging.WARNING)
+    logging.getLogger("urllib3").setLevel(logging.WARNING)
+    logging.getLogger("requests").setLevel(logging.WARNING)
+
 # ══════════════════════════════════════════════════════════════════════════════
 #  DISTRIBUTED HELPERS
 # ═════════════════��════════════════════════════════════════════════════════════
 
 def setup_dist() -> int:
-    dist.init_process_group(backend="nccl")
+    """set_device BEFORE init_process_group avoids the 'device unknown' NCCL warning."""
     local_rank = int(os.environ.get("LOCAL_RANK", 0))
     torch.cuda.set_device(local_rank)
+    dist.init_process_group(backend="nccl")
     return local_rank
 
 def cleanup_dist():
-    dist.destroy_process_group()
+    if dist.is_initialized():
+        dist.destroy_process_group()
 
 def is_main()    -> bool: return rank() == 0
 def rank()       -> int:  return dist.get_rank()      if dist.is_initialized() else 0
@@ -129,7 +143,7 @@ def barrier():
         dist.barrier()
 
 def all_reduce_mean(tensor: torch.Tensor) -> torch.Tensor:
-    """All-reduce mean — uses SUM + div for universal NCCL compatibility."""
+    """SUM + div for universal NCCL compatibility (ReduceOp.AVG not always supported)."""
     dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
     tensor.div_(world_size())
     return tensor
@@ -176,7 +190,7 @@ def hf_upload(local: str, remote: str) -> bool:
             token=HF_TOKEN,
         )
         if is_main():
-            log.info(f"[hf] ↑ {remote}")
+            log.info(f"[hf] uploaded {remote}")
         return True
     except Exception as e:
         if is_main():
@@ -229,7 +243,7 @@ def list_shards(domain: str) -> List[str]:
 
 
 def _ckpt_step(fname: str) -> int:
-    """Extract step number from checkpoint filename robustly."""
+    """Extract step number from checkpoint filename robustly via regex."""
     m = re.search(r"_s(\d+)\.", fname)
     return int(m.group(1)) if m else 0
 
@@ -259,7 +273,7 @@ def rotate_ckpts(tag: str, keep: int = 2):
                     )
                 except Exception:
                     pass
-            log.info(f"[hf] rotated {victim}")
+            log.info(f"[hf] rotated out {victim}")
     except Exception as e:
         log.warning(f"[hf] rotate failed: {e}")
 
@@ -278,11 +292,11 @@ class NovaConfig:
     max_len:          int   = 8192
     ffn_hidden:       int   = 14336
     diff_lambda_init: float = 0.8
-    logit_softcap:    float = 50.0   # loosened from 30 — less aggressive at 10B scale
+    logit_softcap:    float = 50.0
     rope_theta:       float = 500_000.0
 
     # ── FP8 ────────────────────────────────────────────────────────────────
-    use_fp8: str = "auto"            # "auto" | "on" | "off"
+    use_fp8: str = "auto"   # "auto" | "on" | "off"
 
     # ── optimiser ──────────────────────────────────────────────────────────
     lr:             float = 2e-4
@@ -357,17 +371,16 @@ def resolve_fp8(cfg: NovaConfig) -> bool:
         return False
     ok = HAS_TORCHAO and _is_blackwell()
     if cfg.use_fp8 == "on" and not ok:
-        log.warning("[fp8] forced ON but torchao missing or GPU < sm_100 — bf16 fallback")
+        log.warning("[fp8] forced ON but torchao missing or GPU < sm_100 - bf16 fallback")
         return False
     if cfg.use_fp8 == "auto" and is_main():
         log.info(
-            f"[fp8] auto → torchao={HAS_TORCHAO} "
-            f"blackwell={_is_blackwell()} → {'ON' if ok else 'OFF'}"
+            f"[fp8] auto -> torchao={HAS_TORCHAO} "
+            f"blackwell={_is_blackwell()} -> {'ON' if ok else 'OFF'}"
         )
     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")
 
@@ -380,10 +393,7 @@ 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.
-    """
+    """Must be called BEFORE weight tying so embed/lm_head stay untouched bf16."""
     if not resolve_fp8(cfg):
         return model
     fp8_cfg = Float8LinearConfig.from_recipe_name("rowwise")
@@ -402,10 +412,8 @@ def apply_fp8(model: nn.Module, cfg: NovaConfig) -> nn.Module:
 
 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.
+    Per-domain streaming iterator. Each DataLoader worker constructs its own
+    instance, so there is no shared mutable state between workers.
     """
 
     def __init__(self, domain: str, refresh_every: int = 100_000):
@@ -434,7 +442,6 @@ class ShardStream:
     def __iter__(self) -> Iterator[np.ndarray]:
         while True:
             if not self._queue:
-                # Restart from full seen list and rescan for new shards
                 q = list(self._seen)
                 random.shuffle(q)
                 self._queue = q
@@ -450,9 +457,8 @@ 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.
+    Each DataLoader worker builds its own independent ShardStream per domain
+    inside __iter__ — no shared state, no race conditions with num_workers>0.
     """
 
     def __init__(
@@ -469,7 +475,6 @@ class BalancedIterableDataset(IterableDataset):
         self.refresh_every = refresh_every
 
     def __iter__(self) -> Iterator[torch.Tensor]:
-        # Each worker (or main process) builds its own streams independently
         streams = {
             d: iter(ShardStream(d, self.refresh_every))
             for d in self.domains
@@ -484,7 +489,6 @@ class BalancedIterableDataset(IterableDataset):
             random.shuffle(buf)
             for row in buf:
                 tokens = row[: self.seq_len].astype(np.int64)
-                # pad if shard rows are shorter than seq_len
                 if len(tokens) < self.seq_len:
                     tokens = np.pad(
                         tokens,
@@ -507,15 +511,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
     )
-    # IterableDataset does not use a sampler — worker splitting is handled
-    # inside __iter__ via get_worker_info() if needed (currently each worker
-    # draws independently, giving statistical balance across ranks).
-    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
 
 # ══════════════════════════════════════════════════════════════════════════════
@@ -537,13 +534,8 @@ def precompute_rope(
     max_len: int,
     theta: float,
     device=None,
-):
-    """
-    Extended-theta RoPE.
-    Returns (cos, sin) each of shape (max_len, head_dim // 2).
-    The caller is responsible for expanding to full head_dim via
-    torch.cat([cos, cos], dim=-1).
-    """
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Returns (cos, sin), each shape (max_len, head_dim // 2)."""
     inv_freq = 1.0 / (
         theta ** (
             torch.arange(0, head_dim, 2, dtype=torch.float32, device=device)
@@ -551,7 +543,7 @@ def precompute_rope(
         )
     )
     t     = torch.arange(max_len, dtype=torch.float32, device=device)
-    freqs = torch.outer(t, inv_freq)  # (max_len, head_dim // 2)
+    freqs = torch.outer(t, inv_freq)
     return torch.cos(freqs), torch.sin(freqs)
 
 
@@ -560,37 +552,25 @@ def apply_rope(
     k: torch.Tensor,
     cos: torch.Tensor,
     sin: torch.Tensor,
-) -> tuple[torch.Tensor, torch.Tensor]:
+) -> Tuple[torch.Tensor, torch.Tensor]:
     """
-    Standard paired RoPE rotation.
-
-    Args:
-        q, k  : (B, L, n_heads, head_dim)
-        cos, sin: (L, head_dim // 2)  — from precompute_rope
-
-    The rotation maps [x0, x1, ..., x_{hd-2}, x_{hd-1}] to
-    [-x_{hd/2}, ..., -x_{hd-1}, x_0, ..., x_{hd/2 - 1}].
-    This is the standard "split-half" formulation used by LLaMA/Mistral.
-    Works for any even head_dim, including non-powers-of-two.
+    Standard split-half RoPE (LLaMA/Mistral style).
+    q, k    : (B, L, n_heads, head_dim)
+    cos, sin: (L, head_dim // 2)
+    Works correctly for any even head_dim.
     """
-    L  = q.shape[1]
-    hd = q.shape[-1]
+    L = q.shape[1]
 
-    # Expand half-dim cos/sin to full head_dim: (L, hd)
     cos_full = torch.cat([cos[:L], cos[:L]], dim=-1)   # (L, hd)
     sin_full = torch.cat([sin[:L], sin[:L]], dim=-1)   # (L, hd)
-
-    # Broadcast over batch and heads: (1, L, 1, hd)
-    cos_full = cos_full.unsqueeze(0).unsqueeze(2)
+    cos_full = cos_full.unsqueeze(0).unsqueeze(2)      # (1, L, 1, hd)
     sin_full = sin_full.unsqueeze(0).unsqueeze(2)
 
     def _rotate_half(x: torch.Tensor) -> torch.Tensor:
-        """[-x2, x1] split at hd // 2."""
         half = x.shape[-1] // 2
         x1, x2 = x[..., :half], x[..., half:]
         return torch.cat([-x2, x1], dim=-1)
 
-    # Upcast to float32 for numerical stability, then cast back
     qf = q.float()
     kf = k.float()
     q_rot = (qf * cos_full + _rotate_half(qf) * sin_full).to(q.dtype)
@@ -602,12 +582,11 @@ def apply_rope(
 
 class NovaAttention(nn.Module):
     """
-    D1 — Fused QKV: single matmul for Q, K, V; split after projection.
-    D2 �� Differential Attention:
-         Each head's head_dim is split into two half-heads (half = head_dim // 2).
-         Two independent SDPA operations are run; their outputs are differenced
-         with a learned per-head scalar lambda, then normed and gated.
-         v is also split so each SDPA sees (q_half, k_half, v_half) — shapes match.
+    D1 - Fused QKV: single matmul for Q, K, V; split after projection.
+    D2 - Differential Attention: head_dim split into two half-streams;
+         Q, K, AND V are all split so each SDPA call sees matching shapes.
+         Outputs recombined via concat (not subtraction) after per-head
+         lambda scaling of the second stream, then normed and gated.
     """
 
     def __init__(self, cfg: NovaConfig):
@@ -615,24 +594,19 @@ class NovaAttention(nn.Module):
         self.nh   = cfg.n_heads
         self.nkv  = cfg.n_kv_heads
         self.hd   = cfg.head_dim
-        self.ng   = cfg.n_heads // cfg.n_kv_heads   # query groups per KV head
-        self.half = cfg.head_dim // 2               # per-stream head dim
+        self.ng   = cfg.n_heads // cfg.n_kv_heads
+        self.half = cfg.head_dim // 2
 
         q_dim  = cfg.n_heads    * cfg.head_dim
         kv_dim = cfg.n_kv_heads * cfg.head_dim
 
-        # D1: fused QKV projection
         self.qkv    = nn.Linear(cfg.d_model, q_dim + 2 * kv_dim, bias=False)
         self.o_proj = nn.Linear(q_dim, cfg.d_model, bias=False)
 
-        # QK-Norm: stabilises attention logits at scale (applied per head)
         self.q_norm = RMSNorm(cfg.head_dim)
         self.k_norm = RMSNorm(cfg.head_dim)
-
-        # Post-differential norm: operates on the recombined full head_dim output
         self.h_norm = RMSNorm(cfg.head_dim)
 
-        # D2: per-head differential scalar and gate
         self.diff_lambda = nn.Parameter(
             torch.full((cfg.n_heads,), cfg.diff_lambda_init)
         )
@@ -652,51 +626,38 @@ class NovaAttention(nn.Module):
     ) -> torch.Tensor:
         B, L, D = x.shape
 
-        # ── D1: fused QKV ────────────────────────────────────────────────
         qkv   = self.qkv(x)
         q_dim = self.nh  * self.hd
         kv_d  = self.nkv * self.hd
         q, k, v = qkv.split([q_dim, kv_d, kv_d], dim=-1)
 
-        # Reshape to (B, L, n_heads, head_dim)
         q = q.view(B, L, self.nh,  self.hd)
         k = k.view(B, L, self.nkv, self.hd)
         v = v.view(B, L, self.nkv, self.hd)
 
-        # QK-Norm before RoPE (over head_dim dimension)
         q = self.q_norm(q)
         k = self.k_norm(k)
 
-        # RoPE
         q, k = apply_rope(q, k, cos, sin)
 
-        # Transpose to (B, n_heads, L, head_dim) for SDPA
         q = q.transpose(1, 2)
-        # GQA: repeat KV heads to match query head count
         k = k.transpose(1, 2).repeat_interleave(self.ng, dim=1)
         v = v.transpose(1, 2).repeat_interleave(self.ng, dim=1)
 
-        # ── D2: differential attention ───────────────────────────────────
-        # Split Q, K, V along head_dim into two equal streams.
-        # Each SDPA sees (B, nh, L, half) for all three tensors — shapes match.
+        # Split Q, K, V along head_dim -> two matching-shape streams
         q1, q2 = q[..., :self.half], q[..., self.half:]
         k1, k2 = k[..., :self.half], k[..., self.half:]
         v1, v2 = v[..., :self.half], v[..., self.half:]
 
-        # Two independent causal attention computations
-        out1 = F.scaled_dot_product_attention(q1, k1, v1, is_causal=True)  # (B,nh,L,half)
-        out2 = F.scaled_dot_product_attention(q2, k2, v2, is_causal=True)  # (B,nh,L,half)
+        out1 = F.scaled_dot_product_attention(q1, k1, v1, is_causal=True)
+        out2 = F.scaled_dot_product_attention(q2, k2, v2, is_causal=True)
 
-        # Learned per-head differential: out1 - lambda * out2
         lam = self.diff_lambda.view(1, self.nh, 1, 1).to(out1.dtype)
 
-        # Concatenate the two streams back to full head_dim, then norm + gate
-        # Shape: (B, nh, L, half) cat (B, nh, L, half) -> (B, nh, L, hd)
-        out = torch.cat([out1, lam * out2], dim=-1)
+        out = torch.cat([out1, lam * out2], dim=-1)     # (B, nh, L, hd)
         out = self.h_norm(out) * (1.0 - self.lambda_init)
         out = out * self.head_gate.view(1, self.nh, 1, 1).to(out.dtype)
 
-        # Merge heads and project
         out = out.transpose(1, 2).contiguous().view(B, L, -1)
         return self.o_proj(out)
 
@@ -724,15 +685,12 @@ class SwiGLU(nn.Module):
 # ── D4 + D5: Nova Transformer Block ──────────────────────────────────────────
 
 class NovaBlock(nn.Module):
-    """Pre-norm block with LayerScale (D5). ls_a / ls_f init = 0.1."""
-
     def __init__(self, cfg: NovaConfig):
         super().__init__()
         self.attn_norm = RMSNorm(cfg.d_model)
         self.ffn_norm  = RMSNorm(cfg.d_model)
         self.attn      = NovaAttention(cfg)
         self.ffn       = SwiGLU(cfg)
-        # D5: LayerScale — small init for deep-network stability
         self.ls_a = nn.Parameter(torch.full((cfg.d_model,), 0.1))
         self.ls_f = nn.Parameter(torch.full((cfg.d_model,), 0.1))
 
@@ -754,19 +712,14 @@ class Nova10BDense(nn.Module):
         super().__init__()
         self.cfg  = cfg
 
-        # Embedding and lm_head stay bf16 (excluded from FP8 by _fp8_filter)
         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)
         self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
 
-        # Weight tying is applied AFTER FP8 conversion (see build_model())
-        # so both tensors remain nn.Parameter / bf16.
         self.softcap = cfg.logit_softcap
-
         nn.init.normal_(self.embed.weight, std=0.02)
 
-        # RoPE buffers — recomputed from config, not saved in checkpoint
         cos, sin = precompute_rope(cfg.head_dim, cfg.max_len, cfg.rope_theta)
         self.register_buffer("rope_cos", cos, persistent=False)
         self.register_buffer("rope_sin", sin, persistent=False)
@@ -794,7 +747,6 @@ class Nova10BDense(nn.Module):
             x = layer(x, cos, sin)
 
         logits = self.lm_head(self.norm(x))
-        # D7: logit soft-cap
         logits = self.softcap * torch.tanh(logits / self.softcap)
 
         if targets is None:
@@ -818,14 +770,15 @@ class Nova10BDense(nn.Module):
 
 def build_model(cfg: NovaConfig) -> nn.Module:
     """
-    Correct build order:
+    Build order matters:
       1. Create model on CPU
-      2. Apply FP8 (replaces eligible nn.Linear layers in-place)
-      3. Tie weights AFTER FP8 so both embed and lm_head stay bf16
+      2. Apply FP8 (replaces eligible nn.Linear in-place)
+      3. Tie lm_head.weight = embed.weight AFTER FP8 conversion, so both
+         stay plain bf16 nn.Linear/Embedding (both are in _FP8_SKIP anyway,
+         but this order also guards against future skip-list edits).
     """
     model = Nova10BDense(cfg).to(cfg.dtype)
     model = apply_fp8(model, cfg)
-    # Tie weights after FP8 — both remain nn.Linear / bf16 (skipped by _fp8_filter)
     model.lm_head.weight = model.embed.weight
     return model
 
@@ -833,9 +786,6 @@ def build_model(cfg: NovaConfig) -> nn.Module:
 #  FSDP WRAPPING
 # ══════════════════════════════════════════════════════════════════════════════
 
-# State-dict configs:
-#   SAVE: rank0_only=True  — only rank-0 holds the full dict (saves memory)
-#   LOAD: rank0_only=False — all ranks must participate in scatter
 _SAVE_SD_CFG = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
 _LOAD_SD_CFG = FullStateDictConfig(offload_to_cpu=True, rank0_only=False)
 
@@ -845,6 +795,27 @@ def wrap_fsdp(
     cfg: NovaConfig,
     device: torch.device,
 ) -> FSDP:
+    """
+    Uses SHARD_GRAD_OP (ZeRO-2), NOT FULL_SHARD (ZeRO-3).
+
+    WHY THIS MATTERS: Muon's Newton-Schulz orthogonalization operates on
+    genuine 2-D weight matrices. Under FULL_SHARD, each rank only ever
+    holds a 1-D flat LOCAL SHARD of every parameter except during the
+    brief window of an active forward/backward all-gather. Since
+    build_optimizers() runs once, before any forward pass, AND
+    Muon.step() runs after backward (when params have been re-sharded
+    back to 1-D), FULL_SHARD would either:
+      (a) crash immediately — optimizer sees ndim==1 params and empty
+          muon_p list at construction time, or
+      (b) silently reduce Muon to plain momentum-SGD every step (if using
+          the ndim-fallback safety net), defeating the entire point of
+          using Muon.
+
+    SHARD_GRAD_OP keeps full-shape parameters replicated on every GPU
+    (only gradients + optimizer state are sharded), so Muon always sees
+    real 2-D matrices. Memory cost: ~18GB/GPU for a 9B-param bf16 model,
+    which is trivial on 96GB RTX PRO 6000 cards.
+    """
     mp_policy = MixedPrecision(
         param_dtype=torch.bfloat16,
         reduce_dtype=torch.bfloat16,
@@ -856,18 +827,19 @@ def wrap_fsdp(
     )
     fsdp_model = FSDP(
         model,
-        auto_wrap_policy  = wrap_policy,
-        mixed_precision   = mp_policy,
-        sharding_strategy = ShardingStrategy.FULL_SHARD,
-        backward_prefetch = BackwardPrefetch.BACKWARD_PRE,
-        cpu_offload       = CPUOffload(offload_params=cfg.cpu_offload),
-        device_id         = device,
-        use_orig_params   = True,   # required: Muon needs .grad on original params
+        auto_wrap_policy   = wrap_policy,
+        mixed_precision    = mp_policy,
+        sharding_strategy  = ShardingStrategy.SHARD_GRAD_OP,   # ZeRO-2, not ZeRO-3
+        backward_prefetch  = BackwardPrefetch.BACKWARD_PRE,
+        cpu_offload        = CPUOffload(offload_params=cfg.cpu_offload),
+        device_id          = device,
+        use_orig_params    = True,
     )
     if is_main():
         log.info(
-            f"[fsdp] FULL_SHARD × {world_size()} GPUs | "
-            f"cpu_offload={cfg.cpu_offload} | bf16 mixed-precision"
+            f"[fsdp] SHARD_GRAD_OP (ZeRO-2) x {world_size()} GPUs | "
+            f"cpu_offload={cfg.cpu_offload} | bf16 mixed-precision | "
+            f"(full params replicated - required by Muon)"
         )
     return fsdp_model
 
@@ -893,11 +865,13 @@ def _newtonschulz5(G: torch.Tensor, steps: int = 5) -> torch.Tensor:
 
 class Muon(torch.optim.Optimizer):
     """
-    Muon: SGD with Newton-Schulz orthogonalised momentum.
-    Applied only to 2-D weight matrices.
-    With FSDP use_orig_params=True, original param tensors are exposed but
-    their .grad may be a sharded flat tensor during backward.
-    We handle non-2-D grads gracefully with a plain momentum SGD fallback.
+    Muon: SGD with Newton-Schulz orthogonalised momentum, for 2-D weights.
+
+    With SHARD_GRAD_OP + use_orig_params=True, gradients on 2-D params
+    should also be genuine 2-D tensors (matching the replicated full
+    parameter shape). We still keep a defensive fallback to plain
+    momentum-SGD for any unexpected non-2-D grad, so a shape surprise
+    degrades gracefully instead of crashing training.
     """
 
     def __init__(
@@ -908,9 +882,19 @@ class Muon(torch.optim.Optimizer):
         weight_decay: float = 0.0,
         ns_steps: int       = 5,
     ):
+        params = list(params)
+        if len(params) == 0:
+            raise ValueError(
+                "Muon received an empty parameter list. This almost always "
+                "means parameter classification ran on a model whose "
+                "params were not yet in their true (full) shape — check "
+                "that build_optimizers() classifies by NAME, not by "
+                "p.ndim, and that FSDP sharding_strategy is SHARD_GRAD_OP "
+                "(not FULL_SHARD)."
+            )
         defaults = dict(lr=lr, momentum=momentum,
                         weight_decay=weight_decay, ns_steps=ns_steps)
-        super().__init__(list(params), defaults)
+        super().__init__(params, defaults)
 
     @torch.no_grad()
     def step(self):
@@ -926,10 +910,9 @@ class Muon(torch.optim.Optimizer):
                 grad  = p.grad
                 state = self.state[p]
 
-                # ── FSDP may expose sharded / flat grads during backward.
-                #    If grad is not 2-D, fall back to plain momentum SGD
-                #    so we never crash on the ndim==2 assertion. ────────
                 if grad.ndim != 2:
+                    # Defensive fallback - should not normally trigger
+                    # under SHARD_GRAD_OP, but keeps training alive if it does.
                     if "buf" not in state:
                         state["buf"] = torch.zeros_like(grad)
                     buf = state["buf"]
@@ -939,12 +922,11 @@ class Muon(torch.optim.Optimizer):
                     p.add_(buf, alpha=-lr)
                     continue
 
-                # ── normal Muon path ─────────────────────────────────
                 if "buf" not in state:
                     state["buf"] = torch.zeros_like(grad)
                 buf = state["buf"]
                 buf.mul_(mom).add_(grad)
-                grad_ns = grad.add(buf, alpha=mom)       # Nesterov look-ahead
+                grad_ns = grad.add(buf, alpha=mom)
                 u = _newtonschulz5(grad_ns, steps=ns)
                 if wd > 0:
                     p.mul_(1.0 - lr * wd)
@@ -952,30 +934,62 @@ class Muon(torch.optim.Optimizer):
                 p.add_(u.to(p.dtype), alpha=-lr * scale)
 
 
+# Name substrings identifying the 2-D hidden weight matrices that should
+# receive Muon updates. Everything else (embed, lm_head, norms, gates,
+# LayerScale) goes to AdamW.
+_MUON_NAME_PATTERNS = (
+    "qkv.weight",
+    "o_proj.weight",
+    "gate.weight",
+    "up.weight",
+    "down.weight",
+)
+
+
 def build_optimizers(model: nn.Module, cfg: NovaConfig):
     """
-    Split parameters:
-      2-D hidden matrices (qkv, o_proj, gate, up, down) -> Muon
-      embed / lm_head (tied, 2-D but special)           -> AdamW with decay
-      1-D (RMSNorm scales, ls_*, diff_lambda, head_gate) -> AdamW no-decay
+    Parameter classification is done PURELY BY NAME, never by p.ndim.
+
+    WHY: this function is called once, right after wrap_fsdp(), BEFORE any
+    forward/backward pass has run. Under FSDP (any sharding strategy that
+    shards parameters), a parameter's exposed .shape / .ndim at this point
+    may not reflect its true logical shape - e.g. under FULL_SHARD every
+    parameter reports ndim==1 (flat local shard) until the first unshard().
+    Relying on p.ndim here would silently misclassify every 2-D weight
+    matrix as 1-D, starving Muon of all its parameters and crashing with
+    "optimizer got an empty parameter list".
+
+    Classifying by parameter NAME (which is stable regardless of sharding
+    state) avoids this entirely.
     """
-    muon_p    = []
-    aw_decay  = []
-    aw_nodecay = []
-    seen: set = set()
+    muon_p, aw_decay, aw_nodecay, seen = [], [], [], set()
 
     for name, p in model.named_parameters():
         if id(p) in seen or not p.requires_grad:
             continue
         seen.add(id(p))
+
         is_embed_or_head = ("embed" in name) or ("lm_head" in name)
-        if p.ndim == 2 and not is_embed_or_head:
+        is_muon_target = (
+            any(pat in name for pat in _MUON_NAME_PATTERNS)
+            and not is_embed_or_head
+        )
+
+        if is_muon_target:
             muon_p.append(p)
-        elif p.ndim >= 2:
+        elif is_embed_or_head:
             aw_decay.append(p)
         else:
             aw_nodecay.append(p)
 
+    if not muon_p:
+        raise RuntimeError(
+            "[optim] Muon parameter list is empty after NAME-based "
+            "classification. This means NovaAttention/SwiGLU submodule "
+            "attribute names no longer match _MUON_NAME_PATTERNS "
+            f"{_MUON_NAME_PATTERNS}. Check for renamed layers."
+        )
+
     muon_opt = Muon(
         muon_p,
         lr=cfg.muon_lr,
@@ -1005,12 +1019,15 @@ def build_optimizers(model: nn.Module, cfg: NovaConfig):
             fused=True,
         )
         if is_main():
-            log.warning("[optim] bitsandbytes not found — fused fp32 AdamW")
+            log.warning("[optim] bitsandbytes not found - fused fp32 AdamW")
 
     if is_main():
         nm = sum(p.numel() for p in muon_p)
         na = sum(p.numel() for p in aw_decay + aw_nodecay)
-        log.info(f"[optim] Muon {nm / 1e6:.0f}M params | AdamW {na / 1e6:.0f}M params")
+        log.info(
+            f"[optim] Muon {nm / 1e6:.0f}M params ({len(muon_p)} tensors) | "
+            f"AdamW {na / 1e6:.0f}M params ({len(aw_decay) + len(aw_nodecay)} tensors)"
+        )
 
     return muon_opt, aw_opt
 
@@ -1034,7 +1051,6 @@ def cosine_lr(
 # ══════════════════════════════════════════════════════════════════════════════
 
 def _full_state_dict_save(model: FSDP) -> dict:
-    """Gather full state dict to rank-0 CPU only."""
     with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, _SAVE_SD_CFG):
         return model.state_dict()
 
@@ -1049,8 +1065,8 @@ def save_ckpt(
     val_loss: float,
     tag: str,
 ) -> bool:
-    barrier()                            # all ranks finish backward first
-    sd = _full_state_dict_save(model)    # non-empty only on rank-0
+    barrier()
+    sd = _full_state_dict_save(model)
     if not is_main():
         return True
 
@@ -1130,12 +1146,8 @@ def load_ckpt(
     aw_opt,
     device: torch.device,
     cfg: NovaConfig,
-) -> tuple[int, int, float]:
-    """
-    Resume from the latest checkpoint on the Hub.
-    All ranks must call this — FSDP scatters the state dict across GPUs.
-    Returns (step, epoch, best_val_loss).
-    """
+) -> Tuple[int, int, float]:
+    """All ranks must call this — FSDP scatters the full state dict."""
     try:
         files = hf_api.list_repo_files(
             repo_id=HF_MODEL_REPO,
@@ -1158,13 +1170,12 @@ def load_ckpt(
 
         sd = torch.load(mpath, map_location="cpu", weights_only=True)
 
-        # LOAD config: rank0_only=False so all ranks receive weights
         with FSDP.state_dict_type(
             model, StateDictType.FULL_STATE_DICT, _LOAD_SD_CFG
         ):
             model.load_state_dict(sd, strict=False)
 
-        barrier()   # ensure all ranks finish scatter before optimizer load
+        barrier()
 
         meta = torch.load(xpath, map_location="cpu", weights_only=False)
         for opt, key in [(muon_opt, "muon"), (aw_opt, "adamw")]:
@@ -1174,7 +1185,6 @@ def load_ckpt(
                 except Exception as e:
                     log.warning(f"[ckpt] {key} optimizer not restored: {e}")
 
-        # Restore schedule params from saved config
         saved_cfg = meta.get("config", {})
         for field_name in [
             "total_steps", "warmup_steps",
@@ -1190,7 +1200,7 @@ def load_ckpt(
 
     except Exception as ex:
         if is_main():
-            log.error(f"[ckpt] resume failed: {ex} — fresh start")
+            log.error(f"[ckpt] resume failed: {ex} - fresh start")
         return 0, 0, float("inf")
 
 # ══════════════════════════════════════════════════════════════════════════════
@@ -1235,7 +1245,6 @@ def vram_str(dev: int = 0) -> str:
 # ══════════════════════════════════════════════════════════════════════════════
 
 def train():
-    # ── distributed init ────────────────────────────────────────────────────
     local_rank = setup_dist()
     device     = torch.device(f"cuda:{local_rank}")
     ws         = world_size()
@@ -1244,14 +1253,11 @@ def train():
     np.random.seed(42 + rank())
     random.seed(42 + rank())
 
-    # HF login on all ranks (data loading happens on every rank)
     init_hf()
     barrier()
 
-    # ── config ──────────────────────────────────────────────────────────────
     cfg = NovaConfig()
 
-    # Scale total_steps to target ~200B tokens regardless of GPU count
     target_tokens   = 200_000_000_000
     tokens_per_step = cfg.effective_batch_tokens(ws)
     cfg.total_steps = max(
@@ -1261,11 +1267,10 @@ def train():
     if is_main():
         log.info(
             f"[schedule] {ws} GPUs | {tokens_per_step / 1e6:.1f}M tok/step | "
-            f"{cfg.total_steps} steps ≈ "
+            f"{cfg.total_steps} steps ~= "
             f"{cfg.total_steps * tokens_per_step / 1e9:.0f}B tokens"
         )
 
-    # Pull vocab_size from dataset metadata if available (rank-0 only)
     if is_main():
         try:
             meta_p = hf_download_dataset("metadata.json")
@@ -1275,23 +1280,20 @@ def train():
             cfg.vocab_size = (actual + 63) // 64 * 64
             log.info(f"[data] vocab {actual} -> {cfg.vocab_size}")
         except Exception as e:
-            log.info(f"[data] no metadata.json ({e}) — default vocab={cfg.vocab_size}")
+            log.info(f"[data] no metadata.json ({e}) - default vocab={cfg.vocab_size}")
 
-    # Broadcast vocab_size to all ranks
     vs = torch.tensor(cfg.vocab_size, device=device)
     dist.broadcast(vs, src=0)
     cfg.vocab_size = int(vs.item())
 
-    # ── model ────────────────────────────────────────────────────────────────
     if is_main():
-        log.info("[model] building Nova-10B-Dense …")
+        log.info("[model] building Nova-10B-Dense ...")
 
     with torch.device("cpu"):
-        model = build_model(cfg)   # FP8 applied before weight tying
+        model = build_model(cfg)
 
     model = wrap_fsdp(model, cfg, device)
 
-    # Optional torch.compile — substantial throughput gain on Blackwell
     if cfg.use_compile and torch.__version__ >= "2.7.0":
         model = torch.compile(model, mode="max-autotune", fullgraph=False)
         if is_main():
@@ -1300,10 +1302,8 @@ def train():
     if is_main():
         log.info(f"[model] VRAM after wrap: {vram_str(local_rank)}")
 
-    # ── optimizers ──────────────────────────────────────────────────────────
     muon_opt, aw_opt = build_optimizers(model, cfg)
 
-    # ── resume ──────────────────────────────────────────────────────────────
     start_step, start_epoch, best_val = 0, 0, float("inf")
     try:
         start_step, start_epoch, best_val = load_ckpt(
@@ -1311,23 +1311,20 @@ def train():
         )
     except Exception as e:
         if is_main():
-            log.error(f"[ckpt] resume error: {e} — fresh start")
+            log.error(f"[ckpt] resume error: {e} - fresh start")
 
-    # Broadcast step / epoch so all ranks agree
     info = torch.tensor([start_step, start_epoch], dtype=torch.long, device=device)
     dist.broadcast(info, src=0)
     start_step, start_epoch = int(info[0]), int(info[1])
 
-    # ── data ────────────────────────────────────────────────────────────────
     train_loader, val_loader = make_loaders(cfg, rank(), ws)
     train_iter = itertools.cycle(train_loader)
 
-    # ── summary ─────────────────────────────────────────────────────────────
     if is_main():
         log.info("=" * 72)
         log.info(
-            f"Nova-10B-Dense | Blackwell × {ws} | "
-            f"FP8={resolve_fp8(cfg)} | FSDP FULL_SHARD | "
+            f"Nova-10B-Dense | Blackwell x {ws} | "
+            f"FP8={resolve_fp8(cfg)} | FSDP SHARD_GRAD_OP | "
             f"Muon + {'AdamW-8bit' if HAS_BNB else 'AdamW-fp32'}"
         )
         log.info(
@@ -1336,13 +1333,13 @@ def train():
             f"ffn_h={cfg.ffn_hidden} max_len={cfg.max_len}"
         )
         log.info(
-            f"  micro={cfg.micro_batch} × accum={cfg.grad_accum} × "
-            f"{ws} GPUs × seq={cfg.max_len} = "
+            f"  micro={cfg.micro_batch} x accum={cfg.grad_accum} x "
+            f"{ws} GPUs x seq={cfg.max_len} = "
             f"{tokens_per_step / 1e6:.1f}M tok/step"
         )
         log.info(
             f"  {cfg.total_steps} steps -> "
-            f"≈{cfg.total_steps * tokens_per_step / 1e9:.0f}B tokens"
+            f"~={cfg.total_steps * tokens_per_step / 1e9:.0f}B tokens"
         )
         log.info(
             f"  muon_lr={cfg.muon_lr}  aw_lr={cfg.lr}  "
@@ -1368,13 +1365,11 @@ def train():
 
     while step < cfg.total_steps:
 
-        # ── budget guard ────────────────────────────────────────────────────
         if elapsed_h() >= cfg.max_hours:
             if is_main():
-                log.info(f"[time] {cfg.max_hours}h budget exhausted — stopping")
+                log.info(f"[time] {cfg.max_hours}h budget exhausted - stopping")
             break
 
-        # ── fetch batch ─────────────────────────────────────────────────────
         try:
             batch = next(train_iter)
         except Exception as e:
@@ -1384,7 +1379,6 @@ def train():
 
         batch = batch.to(device, non_blocking=True)
 
-        # ── ctx-warmup: broadcast length from rank-0 so all ranks agree ────
         ctx_raw = cfg.current_ctx_len(step)
         ctx_t   = torch.tensor(ctx_raw, device=device)
         dist.broadcast(ctx_t, src=0)
@@ -1392,7 +1386,6 @@ def train():
         if ctx < batch.shape[1]:
             batch = batch[:, :ctx]
 
-        # ── forward + backward ──────────────────────────────────────────────
         try:
             with torch.amp.autocast("cuda", dtype=cfg.dtype):
                 _, loss = model(batch, targets=batch)
@@ -1405,13 +1398,12 @@ def train():
             aw_opt.zero_grad(set_to_none=True)
             accum_loss = accum_n = 0
             if is_main():
-                log.error("[oom] step skipped — consider reducing micro_batch")
+                log.error("[oom] step skipped - consider reducing micro_batch")
             continue
 
         if accum_n < cfg.grad_accum:
             continue
 
-        # ── optimizer step ──────────────────────────────────────────────────
         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,
@@ -1434,7 +1426,6 @@ def train():
         accum_loss = accum_n = 0
         step += 1
 
-        # ── logging ─────────────────────────────────────────────────────────
         if is_main() and step % cfg.log_every == 0:
             sps = (step - start_step) / max(time.time() - t_start, 1e-6)
             tps = sps * tokens_per_step
@@ -1452,7 +1443,6 @@ def train():
                 f"  vram {vram_str(local_rank)}"
             )
 
-        # ── validation ──────────────────────────────────────────────────────
         if step % cfg.eval_every == 0:
             vl  = run_val(model, val_loader, device, n=80)
             vpp = math.exp(min(vl, 20))
@@ -1466,7 +1456,6 @@ def train():
                     )
                     log.info("  >> new best checkpoint saved")
 
-        # ── periodic checkpoint ─────────────────────────────────────────────
         if step % cfg.ckpt_every == 0:
             save_ckpt(
                 model, muon_opt, aw_opt,
@@ -1475,7 +1464,6 @@ def train():
                 tag="ckpt",
             )
 
-        # ── 30-min heartbeat checkpoint ─────────────────────────────────────
         if time.time() - t_last_ckpt > 1800:
             save_ckpt(
                 model, muon_opt, aw_opt,
@@ -1485,10 +1473,9 @@ def train():
             )
             t_last_ckpt = time.time()
 
-    # ── final ────────────────────────────────────────────────────────────────
     if is_main():
         log.info("=" * 72)
-        log.info("TRAINING COMPLETE — saving final checkpoint …")
+        log.info("TRAINING COMPLETE - saving final checkpoint ...")
 
     fval = run_val(model, val_loader, device, n=200)
     barrier()