Update nova10b_dense.py

nova10b_dense.py

@@ -19,16 +19,25 @@ Training stack:
   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)
+  D10 ctx-warmup        – sequence-length curriculum
+
+Fixes applied vs original:
+  FIX1  _has_fp8_support() checks major >= 9 (Hopper sm_90 = H100/H200)
+        old _is_blackwell() checked >= 10, silently disabled fp8 on H200
+  FIX2  build_optimizers() called BEFORE wrap_fsdp() on the raw CPU model
+        after FSDP wrap, named_parameters() walks flat-param structure and
+        the 2D-weight classifier sees nothing -> empty muon_p -> ValueError
+  FIX3  sync_module_states=False + identical CPU seeding replaces the giant
+        NCCL broadcast that crashed on both Blackwell and early Hopper runs
 
 Launch:
   torchrun --nproc_per_node=4 nova10b_dense.py
   torchrun --nproc_per_node=8 nova10b_dense.py
 
-Debug launch (verbose CUDA/NCCL diagnostics, slower):
+Debug:
   NOVA_DEBUG=1 torchrun --nproc_per_node=8 nova10b_dense.py
 
-NCCL-safe fallback launch (disables P2P/SHM fast paths):
+NCCL safe fallback:
   NOVA_NCCL_SAFE=1 torchrun --nproc_per_node=8 nova10b_dense.py
 """
 
@@ -116,7 +125,7 @@ torch.backends.cuda.matmul.allow_tf32 = True
 torch.backends.cudnn.allow_tf32       = True
 torch.backends.cudnn.benchmark        = True
 
-# ── logging ───────────────────────────────────────────────────────────────
+# ── logging ──────────────────────────────────────────────────────────────────
 class _RankFilter(logging.Filter):
     def filter(self, record):
         record.rank = dist.get_rank() if dist.is_initialized() else 0
@@ -174,7 +183,7 @@ def seed_for_model_init(seed: int):
     """
     random.seed(seed)
     np.random.seed(seed)
-    torch.manual_seed(seed)   # CPU generator only — no CUDA touch
+    torch.manual_seed(seed)  # CPU generator only — no CUDA touch
 
 
 def seed_for_runtime(seed: int, device: torch.device):
@@ -197,14 +206,16 @@ def seed_for_runtime(seed: int, device: torch.device):
 def _stage(name: str):
     """
     Context manager: logs entry, syncs CUDA before AND after, barriers all
-    ranks, and on exception logs rank/device/stage/traceback before re-raising.
+    ranks, on exception logs rank/device/stage/traceback before re-raising.
     """
     class _StageCtx:
         def __enter__(self):
             self.t0 = time.time()
             dev = torch.cuda.current_device() if torch.cuda.is_available() else -1
-            log.info(f"[stage:{name}] enter | rank={rank()} device=cuda:{dev} "
-                      f"host={socket.gethostname()}")
+            log.info(
+                f"[stage:{name}] enter | rank={rank()} device=cuda:{dev} "
+                f"host={socket.gethostname()}"
+            )
             if torch.cuda.is_available():
                 torch.cuda.synchronize()
             return self
@@ -216,9 +227,10 @@ def _stage(name: str):
                     f"[stage:{name}] FAILED | rank={rank()} device=cuda:{dev} | "
                     f"{exc_type.__name__}: {exc_val}"
                 )
-                log.error("".join(traceback.format_exception(exc_type, exc_val, exc_tb)))
+                log.error(
+                    "".join(traceback.format_exception(exc_type, exc_val, exc_tb))
+                )
                 return False
-
             if torch.cuda.is_available():
                 torch.cuda.synchronize()
             dt = time.time() - self.t0
@@ -234,8 +246,13 @@ def log_cuda_environment(local_rank: int):
     try:
         props = torch.cuda.get_device_properties(dev)
         cap_major, cap_minor = torch.cuda.get_device_capability(dev)
-        nccl_ver = torch.cuda.nccl.version() if hasattr(torch.cuda, "nccl") else "unknown"
-        drv = torch.cuda.driver_version() if hasattr(torch.cuda, "driver_version") else "n/a"
+        nccl_ver = (
+            torch.cuda.nccl.version() if hasattr(torch.cuda, "nccl") else "unknown"
+        )
+        drv = (
+            torch.cuda.driver_version()
+            if hasattr(torch.cuda, "driver_version") else "n/a"
+        )
         log.info(
             f"[env] rank={rank()} local_rank={local_rank} "
             f"gpu={props.name} sm_{cap_major}{cap_minor} "
@@ -251,12 +268,9 @@ def log_cuda_environment(local_rank: int):
 
 def log_p2p_matrix(local_rank: int, ws: int):
     """
-    Logs pairwise P2P access capability between this rank's device and every
-    other visible device. H100/H200 SXM nodes typically ship with NVLink +
-    NVSwitch, so expect P2P access = Y across ALL pairs. If P2P reports False
-    anywhere on this hardware, that is a red flag worth escalating — it means
-    you are not getting the NVLink fabric you are paying for. Unlike PCIe-only
-    workstation cards, a False here is unexpected and should be investigated.
+    Logs pairwise P2P access capability. H200 SXM nodes have NVLink +
+    NVSwitch so expect all-Y. Any N on H200 is a red flag — you are not
+    getting the fabric you are paying for, escalate to infra.
     """
     if not is_main() or not torch.cuda.is_available():
         return
@@ -272,8 +286,7 @@ def log_p2p_matrix(local_rank: int, ws: int):
                     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):")
-        log.info("[p2p] expect all-Y on H200 SXM NVLink/NVSwitch nodes:")
+        log.info("[p2p] pairwise CUDA P2P access matrix (expect all-Y on H200 NVLink):")
         for i, r in enumerate(rows):
             log.info(f"[p2p]   gpu{i}: {r}")
     except Exception as e:
@@ -291,23 +304,27 @@ def cuda_health_check(device: torch.device):
         p = m @ n
         torch.cuda.synchronize(device)
         checksum = float(z.sum().item()) + float(p.sum().item())
-        log.info(f"[health] rank={rank()} device={device} basic op OK "
-                  f"(checksum={checksum:.1f})")
+        log.info(
+            f"[health] rank={rank()} device={device} basic op OK "
+            f"(checksum={checksum:.1f})"
+        )
     except Exception as e:
-        log.error(f"[health] rank={rank()} device={device} FAILED basic CUDA op: {e}")
+        log.error(
+            f"[health] rank={rank()} device={device} FAILED basic CUDA op: {e}"
+        )
         raise
 
 
 def nccl_collective_selftest(device: torch.device, ws: int):
     """
-    Isolated NCCL collective stress test, run BEFORE FSDP wrap.
-    Tests progressively larger all_gather and broadcast ops — same primitive
-    classes used by FSDP unshard() and sync_module_states.
+    Isolated NCCL collective stress test run BEFORE FSDP wrap or any model
+    code. Tests the exact same primitive classes (all_gather, broadcast) that
+    FSDP uses internally.
 
-    If this fails  → fault is environmental (NCCL/driver), not model code.
-    If this passes → fault is specific to FSDP/model collective patterns.
+    If this fails  -> fault is environmental (NCCL/driver), not model code.
+    If this passes -> fault is specific to FSDP/model collective patterns.
 
-    On H200 SXM with NVLink this should pass all sizes easily and quickly.
+    On H200 SXM with NVLink this should pass all sizes quickly.
     """
     sizes_mb = [1, 16, 64, 256, 1024, 4096]
     for mb in sizes_mb:
@@ -317,12 +334,10 @@ def nccl_collective_selftest(device: torch.device, ws: int):
                 (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)
@@ -352,7 +367,7 @@ def nccl_collective_selftest(device: torch.device, ws: int):
     if is_main():
         log.info(
             f"[nccl_selftest] ALL sizes up to {sizes_mb[-1]}MB passed "
-            f"on all {ws} ranks — NCCL/NVLink healthy"
+            f"on all {ws} ranks — NCCL/NVLink healthy ✓"
         )
 
 
@@ -504,30 +519,30 @@ class NovaConfig:
 
     use_fp8: str = "auto"
 
-    lr:             float = 2e-4
-    min_lr:         float = 2e-5
-    muon_lr:        float = 6e-3
-    muon_min_lr:    float = 6e-4
-    weight_decay:   float = 0.1
-    beta1:          float = 0.9
-    beta2:          float = 0.95
-    muon_momentum:  float = 0.95
-    grad_clip:      float = 1.0
-    warmup_steps:   int   = 4_000
+    lr:            float = 2e-4
+    min_lr:        float = 2e-5
+    muon_lr:       float = 6e-3
+    muon_min_lr:   float = 6e-4
+    weight_decay:  float = 0.1
+    beta1:         float = 0.9
+    beta2:         float = 0.95
+    muon_momentum: float = 0.95
+    grad_clip:     float = 1.0
+    warmup_steps:  int   = 4_000
 
     total_steps:      int   = 130_000
     ctx_warmup_steps: int   = 3_000
     ctx_start_len:    int   = 2048
     ctx_end_len:      int   = 2048
 
-    micro_batch:   int   = 6
-    grad_accum:    int   = 8
-    log_every:     int   = 10
-    ckpt_every:    int   = 1_000
-    eval_every:    int   = 500
-    ckpt_keep:     int   = 2
-    seed:          int   = 42
-    max_hours:     float = 47.5
+    micro_batch: int   = 6
+    grad_accum:  int   = 8
+    log_every:   int   = 10
+    ckpt_every:  int   = 1_000
+    eval_every:  int   = 500
+    ckpt_keep:   int   = 2
+    seed:        int   = 42
+    max_hours:   float = 47.5
 
     domains:            List[str] = field(default_factory=lambda: list(DOMAINS))
     shard_refresh_rows: int       = 100_000
@@ -550,23 +565,23 @@ class NovaConfig:
     def current_ctx_len(self, step: int) -> int:
         if self.ctx_warmup_steps <= 0 or step >= self.ctx_warmup_steps:
             return self.ctx_end_len
-        t = step / self.ctx_warmup_steps
+        t   = step / self.ctx_warmup_steps
         raw = int(self.ctx_start_len + t * (self.ctx_end_len - self.ctx_start_len))
         return ((raw + 63) // 64) * 64
 
 
 # ══════════════════════════════════════════════════════════════════════════════
-#  FP8 HELPERS  — Hopper (sm_90, H100/H200) + Blackwell (sm_100+)
+#  FP8 HELPERS
+#  FIX1: _has_fp8_support() checks major >= 9 so H100/H200 (sm_90) get fp8.
+#        Old _is_blackwell() checked >= 10, silently killed fp8 on Hopper.
 # ══════════════════════════════════════════════════════════════════════════════
 
 def _has_fp8_support() -> bool:
     """
-    Native fp8 tensor cores + torchao rowwise scaling are supported on:
-      - Hopper  sm_90  (H100, H200)      ← primary target here
-      - Blackwell sm_100+ (RTX PRO 6000) ← also supported
-    Ampere sm_80 (A100) and below do NOT have fp8 tensor cores at all.
-    The old _is_blackwell() checked major >= 10 which silently returned
-    False on H100/H200 (sm_90) — this is the fixed version.
+    FP8 tensor cores + torchao rowwise scaling are available on:
+      Hopper   sm_90  (H100, H200)       <- primary target
+      Blackwell sm_100+ (RTX PRO 6000)   <- also supported
+    Ampere sm_80 (A100) and below: NO fp8 tensor cores.
     """
     if not torch.cuda.is_available():
         return False
@@ -581,14 +596,14 @@ def resolve_fp8(cfg: NovaConfig) -> bool:
     ok = HAS_TORCHAO and _has_fp8_support()
     if cfg.use_fp8 == "on" and not ok:
         log.warning(
-            "[fp8] forced ON but torchao missing or GPU < sm_90 (Hopper) "
-            "— bf16 fallback"
+            "[fp8] forced ON but torchao missing or GPU < sm_90 — bf16 fallback"
         )
         return False
     if cfg.use_fp8 == "auto" and is_main():
-        dev = torch.cuda.current_device() if torch.cuda.is_available() else -1
         if torch.cuda.is_available():
-            major, minor = torch.cuda.get_device_capability(dev)
+            major, minor = torch.cuda.get_device_capability(
+                torch.cuda.current_device()
+            )
         else:
             major, minor = 0, 0
         log.info(
@@ -600,8 +615,10 @@ def resolve_fp8(cfg: NovaConfig) -> bool:
     return ok
 
 
-_FP8_SKIP = ("embed", "lm_head", "norm", "diff_lambda", "head_gate",
-             "ls_a", "ls_f", "rope")
+_FP8_SKIP = (
+    "embed", "lm_head", "norm", "diff_lambda",
+    "head_gate", "ls_a", "ls_f", "rope",
+)
 
 
 def _fp8_filter(module: nn.Module, fqn: str) -> bool:
@@ -633,19 +650,24 @@ def apply_fp8(model: nn.Module, cfg: NovaConfig) -> nn.Module:
 # ══════════════════════════════════════════════════════════════════════════════
 
 class ShardStream:
+    """
+    Per-domain infinite shard iterator. Periodically re-queries HF for new
+    shards so training automatically picks up newly uploaded data files
+    without restart — general and code shards trickling in are handled.
+    """
     def __init__(self, domain: str, refresh_every: int = 100_000):
         self.domain        = domain
         self.refresh_every = refresh_every
         self._seen:  List[str] = []
         self._queue: List[str] = []
-        self._since_refresh = 0
+        self._since_refresh    = 0
         self._refresh()
 
     def _refresh(self):
         all_shards = list_shards(self.domain)
         new        = [s for s in all_shards if s not in self._seen]
         if new and is_main():
-            log.info(f"[data] {self.domain}: +{len(new)} new shards")
+            log.info(f"[data] {self.domain}: +{len(new)} new shards discovered")
         random.shuffle(new)
         self._queue.extend(new)
         self._seen.extend(new)
@@ -687,10 +709,7 @@ class BalancedIterableDataset(IterableDataset):
         self.refresh_every = refresh_every
 
     def __iter__(self) -> Iterator[torch.Tensor]:
-        streams = {
-            d: iter(ShardStream(d, self.refresh_every))
-            for d in self.domains
-        }
+        streams    = {d: iter(ShardStream(d, self.refresh_every)) for d in self.domains}
         per_domain = max(1, self.buf_size // len(self.domains))
 
         while True:
@@ -755,11 +774,10 @@ def precompute_rope(head_dim: int, max_len: int, theta: float, device=None):
 
 
 def apply_rope(q, k, cos, sin):
-    L  = q.shape[1]
+    L = q.shape[1]
 
     cos_full = torch.cat([cos[:L], cos[:L]], dim=-1)
     sin_full = torch.cat([sin[:L], sin[:L]], dim=-1)
-
     cos_full = cos_full.unsqueeze(0).unsqueeze(2)
     sin_full = sin_full.unsqueeze(0).unsqueeze(2)
 
@@ -819,7 +837,6 @@ class NovaAttention(nn.Module):
 
         q = self.q_norm(q)
         k = self.k_norm(k)
-
         q, k = apply_rope(q, k, cos, sin)
 
         q = q.transpose(1, 2)
@@ -944,8 +961,8 @@ class Nova10BDense(nn.Module):
     def weight_checksum(self) -> float:
         """
         Cheap deterministic checksum over a sample of parameters.
-        Used to verify all ranks built bit-identical models after
-        seed_for_model_init — cheaper than a full sync_module_states broadcast.
+        Verifies all ranks built bit-identical models after seed_for_model_init
+        without needing a full sync_module_states broadcast.
         """
         with torch.no_grad():
             s = 0.0
@@ -960,8 +977,8 @@ class Nova10BDense(nn.Module):
 
 def build_model(cfg: NovaConfig) -> nn.Module:
     """
-    Build on CPU — caller must have called seed_for_model_init(same_seed)
-    on ALL ranks immediately before this, so every rank's random init is
+    Build on CPU. 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.
     """
     model = Nova10BDense(cfg).to(cfg.dtype)
@@ -1007,10 +1024,11 @@ def wrap_fsdp(model: nn.Module, cfg: NovaConfig, device: torch.device) -> FSDP:
         cpu_offload       = CPUOffload(offload_params=cfg.cpu_offload),
         device_id         = device,
         use_orig_params   = True,
-        # sync_module_states=False: correctness guaranteed by identical CPU
-        # seeding via seed_for_model_init() on all ranks. Verified cheaply
-        # by verify_model_sync() checksum all_reduce instead of a giant
-        # parameter broadcast.
+        # FIX3: sync_module_states=False — correctness guaranteed by identical
+        # CPU seeding in seed_for_model_init(). Verified cheaply via
+        # verify_model_sync() checksum all_reduce. The old True value caused
+        # a giant NCCL broadcast at construction time that crashed on both
+        # Blackwell and early Hopper runs.
         sync_module_states = False,
     )
     if is_main():
@@ -1043,8 +1061,9 @@ def _newtonschulz5(G: torch.Tensor, steps: int = 5) -> torch.Tensor:
 
 class Muon(torch.optim.Optimizer):
     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)
+        defaults = dict(
+            lr=lr, momentum=momentum, weight_decay=weight_decay, ns_steps=ns_steps
+        )
         super().__init__(list(params), defaults)
 
     @torch.no_grad()
@@ -1084,6 +1103,19 @@ class Muon(torch.optim.Optimizer):
 
 
 def build_optimizers(model: nn.Module, cfg: NovaConfig):
+    """
+    FIX2: MUST be called on the raw unwrapped CPU model BEFORE wrap_fsdp().
+
+    After FSDP wraps the model, named_parameters() walks FSDP's internal
+    flat-param structure. Parameter names look different and the 2D-weight
+    classification logic can see an empty set, producing:
+      ValueError: optimizer got an empty parameter list
+
+    With use_orig_params=True FSDP keeps references to the original param
+    objects alive, so optimizer references built pre-wrap remain valid and
+    correct after wrapping. Always call: build_optimizers() -> wrap_fsdp(),
+    never the other way around.
+    """
     muon_p     = []
     aw_decay   = []
     aw_nodecay = []
@@ -1093,7 +1125,9 @@ def build_optimizers(model: nn.Module, cfg: NovaConfig):
         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:
             muon_p.append(p)
         elif p.ndim >= 2:
@@ -1101,8 +1135,26 @@ def build_optimizers(model: nn.Module, cfg: NovaConfig):
         else:
             aw_nodecay.append(p)
 
+    if is_main():
+        log.info(
+            f"[optim] param classification: "
+            f"muon={len(muon_p)} tensors | "
+            f"aw_decay={len(aw_decay)} | "
+            f"aw_nodecay={len(aw_nodecay)}"
+        )
+
+    # Hard guard — if this fires, call order is wrong
+    if len(muon_p) == 0:
+        raise RuntimeError(
+            "[optim] muon_p is empty — no 2D non-embed/lm_head params found.\n"
+            "build_optimizers() MUST be called on the unwrapped model BEFORE "
+            "wrap_fsdp(). Check call order in train()."
+        )
+
     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,
     )
 
@@ -1166,10 +1218,12 @@ def save_ckpt(model, muon_opt, aw_opt, step, epoch, cfg, val_loss, tag) -> bool:
         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,
         )
@@ -1314,6 +1368,20 @@ def vram_str(dev: Optional[int] = None) -> str:
 
 # ══════════════════════════════════════════════════════════════════════════════
 #  TRAIN
+#
+#  Stage order (IMPORTANT — do not reorder without reading FIX2 comment):
+#    seed_model_init
+#    config_build
+#    vocab_metadata
+#    model_build_cpu
+#    verify_model_sync
+#    optimizer_build       <- BEFORE fsdp_wrap (FIX2)
+#    fsdp_wrap             <- AFTER optimizer_build (FIX2)
+#    compile
+#    seed_runtime
+#    checkpoint_resume
+#    data_loaders
+#    first_step_smoke_test
 # ══════════════════════════════════════════════════════════════════════════════
 
 def train():
@@ -1378,12 +1446,10 @@ def train():
 
     with _stage("verify_model_sync"):
         # Cheap correctness check replacing sync_module_states broadcast.
-        # Each rank computes a sampled parameter checksum, we all_reduce,
-        # and compare. If seed_for_model_init worked correctly every rank's
-        # checksum is numerically identical.
+        # Sampled param checksum all_reduce — if all ranks seeded identically
+        # every checksum is bit-identical.
         local_checksum = (
-            model.weight_checksum()
-            if hasattr(model, "weight_checksum") else 0.0
+            model.weight_checksum() if hasattr(model, "weight_checksum") else 0.0
         )
         checksum_t = torch.tensor(
             local_checksum, dtype=torch.float64, device=device
@@ -1402,10 +1468,17 @@ def train():
                 log.error(
                     f"[verify] model checksums DIVERGE across ranks: {vals} "
                     f"(max_diff={max_diff:.6f}) — ranks built DIFFERENT models! "
-                    f"seed_for_model_init may not have produced identical init. "
-                    f"Investigate before continuing — training would silently corrupt."
+                    f"Investigate seed_for_model_init before continuing."
                 )
 
+    # ── FIX2: optimizer built HERE on raw CPU model, BEFORE fsdp_wrap ────────
+    # After wrap_fsdp(), named_parameters() walks FSDP's flat-param structure
+    # and the 2D-weight classifier finds nothing -> empty muon_p -> ValueError.
+    # With use_orig_params=True the optimizer's param references stay valid
+    # after wrapping, so this ordering is both correct and safe.
+    with _stage("optimizer_build"):
+        muon_opt, aw_opt = build_optimizers(model, cfg)
+
     with _stage("fsdp_wrap"):
         model = wrap_fsdp(model, cfg, device)
         enable_weight_tying(model)
@@ -1437,9 +1510,6 @@ def train():
     with _stage("seed_runtime"):
         seed_for_runtime(42 + rank(), device)
 
-    with _stage("optimizer_build"):
-        muon_opt, aw_opt = build_optimizers(model, cfg)
-
     with _stage("checkpoint_resume"):
         start_step, start_epoch, best_val = 0, 0, float("inf")
         try:
@@ -1561,9 +1631,8 @@ def train():
             accum_loss = accum_n = 0
             if is_main():
                 log.error(
-                    "[oom] step skipped — consider reducing micro_batch. "
-                    "Note: 141GB HBM3e per H200 means this is unlikely "
-                    "unless micro_batch is very large."
+                    "[oom] step skipped — reduce micro_batch. "
+                    "(Unlikely on 141GB H200 HBM3e unless micro_batch is huge.)"
                 )
             continue
         except RuntimeError as e:
@@ -1611,8 +1680,9 @@ def train():
                 f"  mlr {mlr:.2e}  alr {alr:.2e}"
             )
             log.info(
-                f"  {tps / 1e3:6.0f}K tok/s  {1 / max(sps, 1e-9):.2f}s/step"
-                f"  vram {vram_str()}"
+                f"  {tps / 1e3:6.0f}K tok/s  "
+                f"{1 / max(sps, 1e-9):.2f}s/step  "
+                f"vram {vram_str()}"
             )
 
         if step % cfg.eval_every == 0:
@@ -1624,7 +1694,7 @@ def train():
                     best_val = vl
                     save_ckpt(
                         model, muon_opt, aw_opt,
-                        step, epoch, cfg, vl, tag="best"
+                        step, epoch, cfg, vl, tag="best",
                     )
                     log.info("  >> new best checkpoint saved")
 
@@ -1632,14 +1702,14 @@ def train():
             save_ckpt(
                 model, muon_opt, aw_opt,
                 step, epoch, cfg,
-                losses[-1] if losses else 0.0, tag="ckpt"
+                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"
+                losses[-1] if losses else 0.0, tag="latest",
             )
             t_last_ckpt = time.time()