Upload App_fsdp.py

App_fsdp.py

@@ -3,7 +3,7 @@
 Orion Flagship Nano — T2.1 single-file base-pretraining trainer.
 
 Target hardware:
-    NVIDIA RTX PRO 6000 Blackwell, SM 120, ~96 GB VRAM.
+    NVIDIA H200 (SM 90) or RTX PRO 6000 Blackwell (SM 120), ~96 GB VRAM.
 
 T2.1 changes over the original Mini prototype:
     - ~135M-parameter Nano configuration.
@@ -31,7 +31,7 @@ Safety / resumability:
 
 This is a new architecture/run. It intentionally does NOT load Mini-T2 weights.
 
-Distributed launch (one node with eight H200s):
+Distributed launch (one node with eight supported GPUs):
     torchrun --standalone --nproc_per_node=8 App_fsdp.py
 
 The configured micro-batch is global (112 sequences); each rank processes 14
@@ -60,7 +60,7 @@ HF_REPO_ID = os.environ.get(
     "Project-Prism/Orion-Flagship-Nano-T2.1",
 )
 HF_PRIVATE = True
-HF_TOKEN = os.environ.get("HF_TOKEN", "hf_").strip()
+HF_TOKEN = os.environ.get("HF_TOKEN", "").strip()
 
 # 50B-token BASE corpus built by orion_nano_base_pretokenizer_v1.py.
 DATA_REPO_ID = "smilyai-large-team/nano-orion-ultra"
@@ -80,9 +80,11 @@ START_FROM_SCRATCH = False
 ALLOW_FRESH_START_WITH_EXISTING_REMOTE = False
 ALLOW_REMOTE_POINTER_ROLLBACK = False
 
-# Large physical batch on the 96GB Blackwell; base pretraining stays at 1K context.
+# Large physical batch on the target ~96GB-class GPUs; base pretraining stays at
+# 1K context.
 SEQUENCE_LENGTH = 1024
 GLOBAL_MICRO_BATCH_SIZE = 112
+# For OOMs, set this to 56: 7 sequences/GPU × 8 GPUs × 2 accumulation steps.
 # Set to the per-rank value after torch.distributed is initialized.  Keeping
 # the global value explicit prevents accidentally multiplying the update batch
 # by eight when launching with torchrun.
@@ -123,6 +125,17 @@ RUN_KERNEL_TESTS = True
 GROUPED_EXPERT_GEMM_ACTIVE = False
 AUTOCAST_CACHE = True
 
+# The fused kernels and attention path are intentionally gated to the two
+# architectures exercised by this run.  Do not silently run on a nearby
+# compute capability: Triton/PTX support and BF16/Flash behavior can differ.
+SUPPORTED_COMPUTE_CAPABILITIES = ((9, 0), (12, 0))
+HARDWARE_PROFILE_NAMES = {
+    (9, 0): "Hopper / H200",
+    (12, 0): "Blackwell",
+}
+KERNEL_TEST_TIMEOUT_SECONDS = 120.0
+ATTENTION_PROBE_TIMEOUT_SECONDS = 180.0
+
 CPU_THREADS = 8
 PREFETCH_BATCHES = 4
 DATA_SEED = 42
@@ -167,6 +180,17 @@ HF_API_RETRY_MAX_SECONDS = 120.0
 MAX_CONSECUTIVE_NONFINITE = 10
 AUTO_INSTALL_MISSING = True
 
+# Explicitly bounded validation path.  Production training is unchanged unless
+# this is enabled in the launch environment, for example:
+#   ORION_NANO_PREFLIGHT_SMOKE=1 torchrun --standalone --nproc_per_node=8 App_fsdp.py
+PREFLIGHT_SMOKE = os.environ.get("ORION_NANO_PREFLIGHT_SMOKE", "").strip().lower() in {
+    "1", "true", "yes", "on",
+}
+if PREFLIGHT_SMOKE:
+    # Keep the synthetic checkpoint and its metadata out of the production
+    # resume directory when both modes share ORION_NANO_WORK_DIR.
+    WORK_DIR = os.path.join(WORK_DIR, "preflight-smoke")
+
 # ===========================================================================
 # Internal runtime configuration — no external configuration needed
 # ===========================================================================
@@ -295,6 +319,82 @@ def rank0_print(*args, **kwargs):
         print(*args, **kwargs)
 
 
+def _format_capability(capability):
+    return f"SM {capability[0]}.{capability[1]}"
+
+
+def _hardware_inventory():
+    """Return visible GPU details without making a profile assumption."""
+    inventory = []
+    for index in range(torch.cuda.device_count()):
+        properties = torch.cuda.get_device_properties(index)
+        capability = tuple(torch.cuda.get_device_capability(index))
+        inventory.append({
+            "index": index,
+            "name": properties.name,
+            "memory_gib": properties.total_memory / 1024**3,
+            "capability": capability,
+        })
+    return inventory
+
+
+def _format_hardware_inventory(inventory):
+    return "\n".join(
+        f"  GPU {item['index']}: {item['name']} | "
+        f"{_format_capability(item['capability'])} | "
+        f"{item['memory_gib']:.1f} GiB"
+        for item in inventory
+    ) or "  (no visible CUDA GPUs)"
+
+
+def validate_hardware():
+    """Validate the visible rank devices before any real data is touched."""
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA GPU required; torch.cuda.is_available() is false.")
+
+    inventory = _hardware_inventory()
+    visible_count = len(inventory)
+    if visible_count < WORLD_SIZE:
+        raise RuntimeError(
+            f"GPU count mismatch: torchrun requested {WORLD_SIZE} ranks but "
+            f"only {visible_count} CUDA GPU(s) are visible.\n"
+            f"Visible devices:\n{_format_hardware_inventory(inventory)}"
+        )
+    if LOCAL_RANK >= visible_count:
+        raise RuntimeError(
+            f"LOCAL_RANK={LOCAL_RANK} is outside the {visible_count} visible "
+            "CUDA GPU(s).\n"
+            f"Visible devices:\n{_format_hardware_inventory(inventory)}"
+        )
+
+    participating = inventory[:WORLD_SIZE]
+    unsupported = [
+        item for item in participating
+        if item["capability"] not in SUPPORTED_COMPUTE_CAPABILITIES
+    ]
+    if unsupported:
+        supported = ", ".join(
+            _format_capability(capability)
+            for capability in SUPPORTED_COMPUTE_CAPABILITIES
+        )
+        raise RuntimeError(
+            "Unsupported CUDA compute capability for this run. Supported "
+            f"capabilities are {supported}; every participating rank must use "
+            "one of them.\n"
+            f"Visible devices:\n{_format_hardware_inventory(inventory)}"
+        )
+
+    capabilities = {item["capability"] for item in participating}
+    if len(capabilities) != 1:
+        raise RuntimeError(
+            "Mixed GPU compute capabilities are not supported for one FSDP run; "
+            "use homogeneous H200/SM 90 or Blackwell/SM 120 ranks.\n"
+            f"Visible devices:\n{_format_hardware_inventory(inventory)}"
+        )
+
+    return inventory[LOCAL_RANK], inventory
+
+
 def dist_barrier():
     if dist.is_initialized():
         dist.barrier()
@@ -780,13 +880,22 @@ class FusedCrossEntropy(torch.autograd.Function):
 # ===========================================================================
 
 
-def test_optimized_kernels(device="cuda", dtypes=None):
+def _check_preflight_deadline(deadline, label):
+    if deadline is not None and time.monotonic() >= deadline:
+        raise TimeoutError(
+            f"{label} exceeded its preflight time limit; "
+            "reduce the smoke-test scope or inspect the CUDA/Triton setup."
+        )
+
+
+def test_optimized_kernels(device="cuda", dtypes=None, deadline=None):
     # Test odd tails, routing collisions, cast-backward rounding, ignored
     # labels, large logits, and the actual vocabulary-sized reduction.
     if dtypes is None:
         dtypes = (torch.float32, torch.bfloat16)
     for dtype in dtypes:
         for d in (127, 2048):
+            _check_preflight_deadline(deadline, "Triton kernel preflight")
             n, k = 17, 2
             x = torch.randn(n, d, device=device, requires_grad=True)
             order = torch.randperm(n * k, device=device)
@@ -800,6 +909,7 @@ def test_optimized_kernels(device="cuda", dtypes=None):
             torch.testing.assert_close(dx, dxr, rtol=0, atol=0)
 
         for v in (127, 32768):
+            _check_preflight_deadline(deadline, "Triton kernel preflight")
             logits = (torch.randn(9, v, device=device) * 4 + 80).to(dtype)
             logits.requires_grad_(True)
             targets = torch.randint(v, (9,), device=device)
@@ -815,79 +925,92 @@ def test_optimized_kernels(device="cuda", dtypes=None):
                                        atol=2e-7)
 
 
-def test_kernels():
-    test_optimized_kernels()
-    for dtype in (torch.float32, torch.bfloat16):
-        tol = 0.04 if dtype == torch.bfloat16 else 3e-4
-
-        x = torch.randn(
-            8, 2048, device="cuda", dtype=dtype, requires_grad=True
-        )
-        g = torch.randn_like(x)
-        y = TritonRMSFunction.apply(x, 1e-6)
-        dx = torch.autograd.grad(y, x, g)[0]
-
-        xr = x.detach().float().requires_grad_(True)
-        yr = xr * torch.rsqrt(xr.square().mean(-1, keepdim=True) + 1e-6)
-        dxr = torch.autograd.grad(yr, xr, g.float())[0]
-
-        torch.testing.assert_close(y.float(), yr, rtol=tol, atol=tol)
-        torch.testing.assert_close(dx.float(), dxr, rtol=tol, atol=tol)
-
-        a = torch.randn(
-            4096, device="cuda", dtype=dtype, requires_grad=True
-        )
-        b = torch.randn_like(a, requires_grad=True)
-        g = torch.randn_like(a)
+def test_kernels(timeout_seconds=None):
+    if timeout_seconds is None:
+        timeout_seconds = KERNEL_TEST_TIMEOUT_SECONDS
+    deadline = time.monotonic() + timeout_seconds
+    try:
+        test_optimized_kernels(deadline=deadline)
+        for dtype in (torch.float32, torch.bfloat16):
+            _check_preflight_deadline(deadline, "Triton kernel preflight")
+            tol = 0.04 if dtype == torch.bfloat16 else 3e-4
 
-        z = TritonSwiGLUFunction.apply(a, b)
-        da, db = torch.autograd.grad(z, (a, b), g)
+            x = torch.randn(
+                8, 2048, device="cuda", dtype=dtype, requires_grad=True
+            )
+            g = torch.randn_like(x)
+            y = TritonRMSFunction.apply(x, 1e-6)
+            dx = torch.autograd.grad(y, x, g)[0]
 
-        ar = a.detach().float().requires_grad_(True)
-        br = b.detach().float().requires_grad_(True)
-        zr = F.silu(ar) * br
-        dar, dbr = torch.autograd.grad(zr, (ar, br), g.float())
+            xr = x.detach().float().requires_grad_(True)
+            yr = xr * torch.rsqrt(xr.square().mean(-1, keepdim=True) + 1e-6)
+            dxr = torch.autograd.grad(yr, xr, g.float())[0]
 
-        torch.testing.assert_close(z.float(), zr, rtol=tol, atol=tol)
-        torch.testing.assert_close(da.float(), dar, rtol=tol, atol=tol)
-        torch.testing.assert_close(db.float(), dbr, rtol=tol, atol=tol)
+            torch.testing.assert_close(y.float(), yr, rtol=tol, atol=tol)
+            torch.testing.assert_close(dx.float(), dxr, rtol=tol, atol=tol)
 
-        if USE_FUSED_MOE_COMBINE:
-            for d in (127, 2048):
-                n, k = 129, 2
-                grouped = torch.randn(
-                    n * k, d, device="cuda",
-                    dtype=dtype, requires_grad=True,
-                )
-                weights = torch.softmax(
-                    torch.randn(n, k, device="cuda"), dim=-1
-                ).detach().requires_grad_(True)
-                inverse = torch.randperm(n * k, device="cuda")
-                upstream = torch.randn(n, d, device="cuda")
-
-                actual = FusedMoECombine.apply(grouped, weights, inverse)
-                dg, dw = torch.autograd.grad(
-                    actual, (grouped, weights), upstream
-                )
+            a = torch.randn(
+                4096, device="cuda", dtype=dtype, requires_grad=True
+            )
+            b = torch.randn_like(a, requires_grad=True)
+            g = torch.randn_like(a)
+
+            z = TritonSwiGLUFunction.apply(a, b)
+            da, db = torch.autograd.grad(z, (a, b), g)
+
+            ar = a.detach().float().requires_grad_(True)
+            br = b.detach().float().requires_grad_(True)
+            zr = F.silu(ar) * br
+            dar, dbr = torch.autograd.grad(zr, (ar, br), g.float())
+
+            torch.testing.assert_close(z.float(), zr, rtol=tol, atol=tol)
+            torch.testing.assert_close(da.float(), dar, rtol=tol, atol=tol)
+            torch.testing.assert_close(db.float(), dbr, rtol=tol, atol=tol)
+
+            if USE_FUSED_MOE_COMBINE:
+                for d in (127, 2048):
+                    _check_preflight_deadline(deadline, "Triton kernel preflight")
+                    n, k = 129, 2
+                    grouped = torch.randn(
+                        n * k, d, device="cuda",
+                        dtype=dtype, requires_grad=True,
+                    )
+                    weights = torch.softmax(
+                        torch.randn(n, k, device="cuda"), dim=-1
+                    ).detach().requires_grad_(True)
+                    inverse = torch.randperm(n * k, device="cuda")
+                    upstream = torch.randn(n, d, device="cuda")
+
+                    actual = FusedMoECombine.apply(grouped, weights, inverse)
+                    dg, dw = torch.autograd.grad(
+                        actual, (grouped, weights), upstream
+                    )
 
-                gr = grouped.detach().clone().requires_grad_(True)
-                wr = weights.detach().clone().requires_grad_(True)
-                reference = combine_reference(gr, wr, inverse)
-                dgr, dwr = torch.autograd.grad(
-                    reference, (gr, wr), upstream
-                )
+                    gr = grouped.detach().clone().requires_grad_(True)
+                    wr = weights.detach().clone().requires_grad_(True)
+                    reference = combine_reference(gr, wr, inverse)
+                    dgr, dwr = torch.autograd.grad(
+                        reference, (gr, wr), upstream
+                    )
 
-                torch.testing.assert_close(
-                    actual, reference, rtol=1e-5, atol=1e-5
-                )
-                torch.testing.assert_close(
-                    dg.float(), dgr.float(), rtol=1e-5, atol=1e-5
-                )
-                torch.testing.assert_close(
-                    dw, dwr, rtol=3e-4, atol=3e-4
-                )
+                    torch.testing.assert_close(
+                        actual, reference, rtol=1e-5, atol=1e-5
+                    )
+                    torch.testing.assert_close(
+                        dg.float(), dgr.float(), rtol=1e-5, atol=1e-5
+                    )
+                    torch.testing.assert_close(
+                        dw, dwr, rtol=3e-4, atol=3e-4
+                    )
 
-    torch.cuda.synchronize()
+        _check_preflight_deadline(deadline, "Triton kernel preflight")
+        torch.cuda.synchronize()
+        _check_preflight_deadline(deadline, "Triton kernel preflight")
+    except Exception as exc:
+        raise RuntimeError(
+            "Triton kernel preflight failed on the active GPU. Check the "
+            "CUDA/Triton/PyTorch versions and the reported assertion or timeout."
+        ) from exc
     print("All enabled Triton forward/backward checks passed.", flush=True)
 
 
@@ -1018,13 +1141,18 @@ def attention_expand(self, x):
 Attention.forward = attention_gqa
 
 
-def select_attention(cfg):
+def select_attention(cfg, timeout_seconds=None):
+    if timeout_seconds is None:
+        timeout_seconds = ATTENTION_PROBE_TIMEOUT_SECONDS
     hd = cfg["d_model"] // cfg["n_heads"]
+    deadline = time.monotonic() + timeout_seconds
 
     def probe(native):
-        q = torch.randn(1, cfg["n_heads"], 256, hd,
+        _check_preflight_deadline(deadline, "Flash attention preflight")
+        probe_tokens = max(1, min(256, cfg["sequence_length"]))
+        q = torch.randn(1, cfg["n_heads"], probe_tokens, hd,
                         device="cuda", dtype=torch.bfloat16, requires_grad=True)
-        k = torch.randn(1, cfg["n_kv_heads"], 256, hd,
+        k = torch.randn(1, cfg["n_kv_heads"], probe_tokens, hd,
                         device="cuda", dtype=torch.bfloat16, requires_grad=True)
         v = torch.randn_like(k, requires_grad=True)
         with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
@@ -1042,14 +1170,22 @@ def select_attention(cfg):
                 )
             out.float().square().mean().backward()
         torch.cuda.synchronize()
+        _check_preflight_deadline(deadline, "Flash attention preflight")
 
     try:
         probe(True)
         eager = attention_gqa
         print("Native Flash GQA forward/backward supported.")
-    except Exception as exc:
-        print("Native GQA probe failed:", exc)
-        probe(False)
+    except Exception as native_exc:
+        print("Native GQA probe failed; trying explicit KV expansion:", native_exc)
+        try:
+            probe(False)
+        except Exception as fallback_exc:
+            raise RuntimeError(
+                "Flash attention preflight failed for both native GQA and "
+                "explicit KV expansion. Check GPU capability, BF16 support, "
+                "and the PyTorch CUDA attention backend."
+            ) from fallback_exc
         eager = attention_expand
         print("Using explicit KV expansion with Flash attention.")
 
@@ -1059,11 +1195,14 @@ def select_attention(cfg):
 
     trial = x = out = None
     try:
+        _check_preflight_deadline(deadline, "Attention compilation preflight")
         with torch.device("cuda"):
             trial = Attention(cfg)
         compiled = torch.compile(eager, fullgraph=True, dynamic=False)
+        probe_batch = max(1, min(MICRO_BATCH_SIZE, 2))
+        probe_tokens = max(1, min(SEQUENCE_LENGTH, 256))
         x = torch.randn(
-            MICRO_BATCH_SIZE, SEQUENCE_LENGTH, cfg["d_model"],
+            probe_batch, probe_tokens, cfg["d_model"],
             device="cuda", dtype=torch.float32, requires_grad=True,
         )
         with torch.autocast("cuda", dtype=torch.bfloat16,
@@ -1071,8 +1210,14 @@ def select_attention(cfg):
             out = compiled(trial, x)
         out.float().square().mean().backward()
         torch.cuda.synchronize()
+        _check_preflight_deadline(deadline, "Attention compilation preflight")
         Attention.forward = compiled
         print("Compiled attention startup forward/backward passed.")
+    except TimeoutError as exc:
+        raise RuntimeError(
+            "Attention compilation preflight exceeded its time limit; "
+            "inspect the CUDA/PyTorch compiler setup before training."
+        ) from exc
     except Exception as exc:
         Attention.forward = eager
         print("Attention compilation failed; using eager attention:", exc)
@@ -1513,45 +1658,6 @@ class OrionFlagshipNanoT21(nn.Module):
             "banks": banks,
         }
 
-    @torch.no_grad()
-    def run_telemetry_probe(self, input_ids, labels):
-        was_training = self.training
-        self.eval()
-        self.reset_telemetry()
-        self.set_telemetry(True)
-        try:
-            with torch.autocast(
-                "cuda", dtype=torch.bfloat16, cache_enabled=AUTOCAST_CACHE
-            ):
-                ce, auxiliary = self(input_ids, labels)
-            stats = self.telemetry_summary()
-            stats["probe_ce"] = float(ce)
-            stats["probe_aux"] = float(auxiliary)
-            return stats
-        finally:
-            self.set_telemetry(False)
-            if was_training:
-                self.train()
-
-    @torch.no_grad()
-    def run_ablation_probe(self, input_ids, labels):
-        was_training = self.training
-        self.eval()
-        try:
-            results = {}
-            with torch.autocast(
-                "cuda", dtype=torch.bfloat16, cache_enabled=AUTOCAST_CACHE
-            ):
-                base, _ = self(input_ids, labels)
-                results["full"] = float(base)
-                for name in ("vault", "state", "procedure", "deliberation"):
-                    ce, _ = self(input_ids, labels, ablate={name})
-                    results[f"no_{name}"] = float(ce)
-            return results
-        finally:
-            if was_training:
-                self.train()
-
     def initialize_weights(self):
         for module in self.modules():
             if isinstance(module, (nn.Linear, nn.Embedding)):
@@ -1688,6 +1794,48 @@ class OrionFlagshipNanoT21(nn.Module):
         return ce, auxiliary / max(1, executed_stages)
 
 
+@torch.no_grad()
+def run_telemetry_probe(model, input_ids, labels):
+    """All ranks must enter; forwards must go through the FSDP wrapper."""
+    module = model.module if isinstance(model, FSDP) else model
+    was_training = model.training
+    try:
+        model.eval()
+        module.reset_telemetry()
+        module.set_telemetry(True)
+        with torch.autocast(
+            "cuda", dtype=torch.bfloat16, cache_enabled=AUTOCAST_CACHE
+        ):
+            ce, auxiliary = model(input_ids, labels)
+        stats = module.telemetry_summary()
+        stats["probe_ce"] = float(ce)
+        stats["probe_aux"] = float(auxiliary)
+        return stats
+    finally:
+        module.set_telemetry(False)
+        model.train(was_training)
+
+
+@torch.no_grad()
+def run_ablation_probe(model, input_ids, labels):
+    """Run each ablation through the wrapper so FSDP gathers its parameters."""
+    was_training = model.training
+    try:
+        model.eval()
+        results = {}
+        with torch.autocast(
+            "cuda", dtype=torch.bfloat16, cache_enabled=AUTOCAST_CACHE
+        ):
+            base, _ = model(input_ids, labels)
+            results["full"] = float(base)
+            for name in ("vault", "state", "procedure", "deliberation"):
+                ce, _ = model(input_ids, labels, ablate={name})
+                results[f"no_{name}"] = float(ce)
+        return results
+    finally:
+        model.train(was_training)
+
+
 def parameter_counts(cfg):
     """Exact meta count + approximate per-stage active compute proxy."""
     with torch.device("meta"):
@@ -1711,37 +1859,35 @@ def parameter_counts(cfg):
 
 
 def gradient_health_summary(model):
-    groups = {
-        "cortex": [],
-        "procedure": [],
-        "vault": [],
-        "state": [],
-        "embedding": [],
-    }
+    # Original parameters expose only this rank's shard under FULL_SHARD.
+    # Keep every group in the collective even when this rank has no gradients.
+    groups = ("cortex", "procedure", "vault", "state", "embedding")
+    first_parameter = next(model.parameters(), None)
+    device = first_parameter.device if first_parameter is not None else DEVICE
+    squared_norms = torch.zeros(len(groups), device=device, dtype=torch.float32)
     for name, p in model.named_parameters():
         if p.grad is None:
             continue
+        name = ".".join(
+            part for part in name.split(".") if part != "_fsdp_wrapped_module"
+        )
         g = p.grad.detach().float()
         value = g.square().sum()
         if name.startswith("procedure_banks") or name.startswith("route_state_proj") \
                 or name.startswith("stage_embedding") or name.startswith("cycle_embedding"):
-            groups["procedure"].append(value)
+            group = "procedure"
         elif name.startswith("vault"):
-            groups["vault"].append(value)
+            group = "vault"
         elif name.startswith("working") or name.startswith("state_seed"):
-            groups["state"].append(value)
+            group = "state"
         elif name.startswith("embedding"):
-            groups["embedding"].append(value)
+            group = "embedding"
         else:
-            groups["cortex"].append(value)
-
-    result = {}
-    for key, values in groups.items():
-        if values:
-            result[key] = float(torch.stack(values).sum().sqrt())
-        else:
-            result[key] = 0.0
-    return result
+            group = "cortex"
+        squared_norms[groups.index(group)] += value
+    if dist.is_initialized():
+        dist.all_reduce(squared_norms, op=dist.ReduceOp.SUM)
+    return dict(zip(groups, squared_norms.sqrt().tolist()))
 
 
 def health_warnings(stats, grad_stats, step):
@@ -1788,7 +1934,10 @@ def run_startup_model_probe(model, vocab_size):
     with torch.autocast("cuda", dtype=torch.bfloat16, cache_enabled=AUTOCAST_CACHE):
         ce, aux = model(x, y)
         loss = ce + aux
-    if not bool(torch.isfinite(loss).item()):
+    finite_loss = torch.isfinite(loss.detach()).to(dtype=torch.int32)
+    if dist.is_initialized():
+        dist.all_reduce(finite_loss, op=dist.ReduceOp.MIN)
+    if not bool(finite_loss.item()):
         raise RuntimeError("Startup model probe produced nonfinite loss.")
     loss.backward()
     grads = gradient_health_summary(model)
@@ -1796,7 +1945,7 @@ def run_startup_model_probe(model, vocab_size):
         if not math.isfinite(grads[key]) or grads[key] <= 0.0:
             raise RuntimeError(f"Startup probe: invalid {key} gradient norm {grads[key]}")
     model.zero_grad(set_to_none=True)
-    stats = model.run_telemetry_probe(x, y)
+    stats = run_telemetry_probe(model, x, y)
     print(
         f"Startup probe PASS: ce={float(ce):.4f} aux={float(aux):.4f} "
         f"grads={grads} vault_gate={stats['vault_gate']:.3f} "
@@ -2278,10 +2427,24 @@ class PrefetchedStream:
 def distributed_batch(stream):
     """Fetch one global batch on rank zero and give every rank a disjoint slice."""
     local_batch = GLOBAL_MICRO_BATCH_SIZE // WORLD_SIZE
+    failure = None
+    status = None
     if RANK == 0:
-        x, y, cursor = stream.batch()
-        packed = torch.stack((x, y), dim=0)
-    else:
+        try:
+            x, y, cursor = stream.batch()
+            packed = torch.stack((x, y), dim=0)
+            status = "ok"
+        except Exception as exc:
+            failure = exc
+            status = "oom" if isinstance(exc, torch.cuda.OutOfMemoryError) else "failed"
+    # Never leave peers waiting for a batch when its producer has failed. All
+    # ranks take the same recovery path, using the last committed data cursor.
+    status = broadcast_object(status)
+    if status == "oom":
+        raise torch.cuda.OutOfMemoryError("Rank-zero batch preparation ran out of memory.") from failure
+    if status != "ok":
+        raise DataPipelineError("Rank-zero batch preparation failed.") from failure
+    if RANK != 0:
         packed = torch.empty(
             (2, GLOBAL_MICRO_BATCH_SIZE, SEQUENCE_LENGTH),
             device=DEVICE,
@@ -2405,13 +2568,6 @@ def save_local_checkpoint(
     model, optimizer, tokenizer, progress, data_state
 ):
     cfg = model.module.cfg if isinstance(model, FSDP) else model.cfg
-    root = WORK / "checkpoints"
-    root.mkdir(parents=True, exist_ok=True)
-
-    name = f"step-{progress['step']:09d}-{uuid.uuid4().hex[:8]}"
-    temporary = root / (".building-" + name)
-    final = root / name
-
     # FULL_STATE_DICT is collective: every rank must enter this context, but
     # rank zero is the only writer.  This keeps the existing portable checkpoint
     # layout while making it valid for sharded parameters.
@@ -2421,9 +2577,31 @@ def save_local_checkpoint(
                               state_cfg, optim_cfg):
         full_model = model.state_dict()
         full_optimizer = FSDP.optim_state_dict(model, optimizer)
-    if RANK != 0:
-        dist_barrier()
-        return None
+    final = None
+    if RANK == 0:
+        try:
+            final = _write_local_checkpoint(
+                full_model, full_optimizer, tokenizer, progress, data_state, cfg
+            )
+        except BaseException as exc:
+            # Peers wait for this outcome, including failures before file writing
+            # (directory creation, disk-space checks, CUDA synchronization).
+            broadcast_object(f"{type(exc).__name__}: {exc}")
+            raise
+    error = broadcast_object(None)
+    if error is not None:
+        raise RuntimeError(f"Rank-zero checkpoint save failed: {error}")
+    return final
+
+
+def _write_local_checkpoint(
+    full_model, full_optimizer, tokenizer, progress, data_state, cfg
+):
+    root = WORK / "checkpoints"
+    root.mkdir(parents=True, exist_ok=True)
+    name = f"step-{progress['step']:09d}-{uuid.uuid4().hex[:8]}"
+    temporary = root / (".building-" + name)
+    final = root / name
     required = nested_bytes(full_model) + nested_bytes(full_optimizer) + 2 * 1024**3
     free = shutil.disk_usage(root).free
     if free < required:
@@ -2494,7 +2672,6 @@ def save_local_checkpoint(
             "step": progress["step"],
             "tokens": progress["tokens"],
         })
-        dist_barrier()
         return final
 
     except BaseException:
@@ -2511,6 +2688,34 @@ def prune_local_checkpoints(keep):
                 shutil.rmtree(path)
 
 
+def _load_checkpoint_shards(path, manifest, key, load_shard):
+    filenames = manifest.get(key)
+    if not isinstance(filenames, list) or not filenames:
+        raise ValueError(f"Checkpoint manifest requires a non-empty {key} list.")
+    merged = {}
+    seen = set()
+    for filename in filenames:
+        if not isinstance(filename, str):
+            raise ValueError(f"Invalid filename in checkpoint {key}: {filename!r}")
+        filename = safe_relative_path(filename)
+        if filename in seen:
+            raise ValueError(f"Duplicate checkpoint shard in {key}: {filename}")
+        seen.add(filename)
+        shard = load_shard(path / filename)
+        if not isinstance(shard, dict):
+            raise ValueError(f"Checkpoint shard {filename} must contain a dictionary.")
+        duplicates = merged.keys() & shard.keys()
+        if duplicates:
+            raise ValueError(
+                f"Duplicate keys in checkpoint {key} shard {filename}: "
+                f"{sorted(map(str, duplicates))}"
+            )
+        # shard_items splits top-level entries, including optimizer state and
+        # param_groups; every listed shard contributes to the full dictionary.
+        merged.update(shard)
+    return merged
+
+
 def restore_checkpoint(path, model, optimizer):
     manifest = json.loads((path / "manifest.json").read_text())
     if manifest.get("run_id") != RUN_ID:
@@ -2521,16 +2726,19 @@ def restore_checkpoint(path, model, optimizer):
 
     if manifest.get("checkpoint_format") != "fsdp_full_state":
         raise ValueError("This FSDP trainer only accepts format-3 checkpoints.")
-    model_state = torch.load(
-        path / safe_relative_path(manifest["model_shards"][0]),
-        map_location="cpu", weights_only=False,
+    model_state = _load_checkpoint_shards(
+        path, manifest, "model_shards",
+        lambda shard: load_file(str(shard), device="cpu"),
     )
-    optimizer_state = torch.load(
-        path / safe_relative_path(manifest["optimizer_shards"][0]),
-        map_location="cpu", weights_only=False,
+    optimizer_state = _load_checkpoint_shards(
+        path, manifest, "optimizer_shards",
+        lambda shard: torch.load(shard, map_location="cpu", weights_only=False),
     )
-    state_cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
-    optim_cfg = FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=True)
+    # Every rank reads the complete state from the shared filesystem. Model
+    # loading needs the full dictionary on each rank, and optimizer conversion
+    # must shard each rank's dictionary rather than expect rank-zero-only input.
+    state_cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=False)
+    optim_cfg = FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=False)
     with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT,
                               state_cfg, optim_cfg):
         model.load_state_dict(model_state)
@@ -3219,15 +3427,22 @@ def main():
         dist.init_process_group("nccl")
     torch.cuda.set_device(LOCAL_RANK)
     if GLOBAL_MICRO_BATCH_SIZE % WORLD_SIZE:
-        raise ValueError("GLOBAL_MICRO_BATCH_SIZE must divide WORLD_SIZE.")
+        raise ValueError("GLOBAL_MICRO_BATCH_SIZE must be divisible by WORLD_SIZE.")
     MICRO_BATCH_SIZE = GLOBAL_MICRO_BATCH_SIZE // WORLD_SIZE
     GRAD_ACCUM_STEPS = TOKENS_PER_UPDATE // (
         GLOBAL_MICRO_BATCH_SIZE * SEQUENCE_LENGTH
     )
-    _fresh_start_banner()
+    if PREFLIGHT_SMOKE:
+        rank0_print(
+            "ORION FLAGSHIP NANO — bounded 8-rank preflight smoke "
+            "(no data stream, no Hugging Face upload)",
+            flush=True,
+        )
+    else:
+        _fresh_start_banner()
     STOP_REQUESTED = False
 
-    if not HF_TOKEN:
+    if not HF_TOKEN and not PREFLIGHT_SMOKE:
         raise ValueError(
             "Set HF_TOKEN to a Hugging Face token with read access to the data "
             "and write access to the checkpoint repo. Do not paste it into this file."
@@ -3237,24 +3452,37 @@ def main():
     if not torch.cuda.is_bf16_supported():
         raise RuntimeError("BF16 GPU support required.")
 
+    # Hardware and kernel preflight must complete before checkpoint discovery,
+    # tokenizer downloads, or construction of a real data stream.
+    hardware, inventory = validate_hardware()
+    capability = hardware["capability"]
+    rank0_print(
+        f"Visible CUDA GPUs: {len(inventory)}; participating ranks: {WORLD_SIZE}\n"
+        f"{_format_hardware_inventory(inventory)}\n"
+        f"Active hardware profile: {HARDWARE_PROFILE_NAMES[capability]} "
+        f"({_format_capability(capability)})",
+        flush=True,
+    )
+    rank0_print("Running bounded CUDA/Triton startup preflight...", flush=True)
+    if RUN_KERNEL_TESTS:
+        test_kernels()
+    select_attention(dict(ARCH))
+    dist_barrier()
+
     process_start = time.monotonic()
     hard_deadline = process_start + SESSION_HOURS * 3600
     nominal_train_deadline = process_start + MAX_TRAIN_HOURS * 3600
     WORK.mkdir(parents=True, exist_ok=True)
 
-    gpu = torch.cuda.get_device_properties(LOCAL_RANK)
-    capability = torch.cuda.get_device_capability(LOCAL_RANK)
     rank0_print(
-        f"GPU: {gpu.name} | {gpu.total_memory / 1024**3:.1f} GiB\n"
+        f"GPU: {hardware['name']} | {hardware['memory_gib']:.1f} GiB\n"
         f"SM: {capability[0]}.{capability[1]}\n"
         f"PyTorch: {torch.__version__} | CUDA: {torch.version.cuda}\n"
-        f"Batch: {MICRO_BATCH_SIZE} × {GRAD_ACCUM_STEPS} × "
+        f"Batch: {MICRO_BATCH_SIZE}/GPU × {WORLD_SIZE} GPUs × {GRAD_ACCUM_STEPS} × "
         f"{SEQUENCE_LENGTH} = {TOKENS_PER_UPDATE:,} tokens/update",
         flush=True,
     )
-    if capability != (12, 0):
-        rank0_print("Warning: this profile was tuned for SM 120 Blackwell.")
-    if gpu.total_memory / 1024**3 < 85:
+    if hardware["memory_gib"] < 85:
         rank0_print("Warning: defaults target a 96GB-class GPU.")
 
     random.seed(42)
@@ -3271,18 +3499,27 @@ def main():
     old_handlers = {}
 
     try:
-        hub = HFCheckpoints() if RANK == 0 else None
-        if RANK == 0:
-            resume_path, downloaded = choose_resume(hub)
-            resume_payload = {
-                "path": str(resume_path) if resume_path is not None else None,
-                "downloaded": downloaded,
-            }
+        if PREFLIGHT_SMOKE:
+            # Do not instantiate HFCheckpoints: its constructor creates the
+            # remote repository and the normal path may discover/download/upload
+            # production checkpoints.  The tokenizer read below is still needed
+            # to construct the production-shape model and is not an upload.
+            resume_path = None
+            downloaded = False
+            resume_payload = {"path": None, "downloaded": False}
         else:
-            resume_payload = None
-        resume_payload = broadcast_object(resume_payload)
-        resume_path = Path(resume_payload["path"]) if resume_payload["path"] else None
-        downloaded = bool(resume_payload["downloaded"])
+            hub = HFCheckpoints() if RANK == 0 else None
+            if RANK == 0:
+                resume_path, downloaded = choose_resume(hub)
+                resume_payload = {
+                    "path": str(resume_path) if resume_path is not None else None,
+                    "downloaded": downloaded,
+                }
+            else:
+                resume_payload = None
+            resume_payload = broadcast_object(resume_payload)
+            resume_path = Path(resume_payload["path"]) if resume_payload["path"] else None
+            downloaded = bool(resume_payload["downloaded"])
         dist_barrier()
 
         tokenizer_source = (
@@ -3290,7 +3527,7 @@ def main():
             if resume_path is not None else TOKENIZER_NAME
         )
         tokenizer = AutoTokenizer.from_pretrained(
-            tokenizer_source, token=HF_TOKEN, use_fast=True
+            tokenizer_source, token=HF_TOKEN or None, use_fast=True
         )
         if tokenizer.eos_token_id is None:
             raise ValueError("Tokenizer requires an EOS token.")
@@ -3314,43 +3551,44 @@ def main():
                 f"Nano parameter guard failed: expected ~135M, got {total:,}."
             )
 
-        if RUN_KERNEL_TESTS:
-            test_kernels()
-        select_attention(cfg)
-
-        rank0_print("Allocating FP32 T2.1 model...", flush=True)
-        with torch.device("meta"):
-            model = OrionFlagshipNanoT21(cfg)
-        model.to_empty(device=DEVICE)
-        for block in model.blocks:
-            block.attn.reset_rope()
-        if resume_path is None:
-            model.initialize_weights()
-        model.train()
-        assert sum(p.numel() for p in model.parameters()) == total
+        def allocate_model(initialize_weights):
+            rank0_print("Allocating FP32 T2.1 model...", flush=True)
+            with torch.device("meta"):
+                candidate = OrionFlagshipNanoT21(cfg)
+            candidate.to_empty(device=DEVICE)
+            for block in candidate.blocks:
+                block.attn.reset_rope()
+            if initialize_weights:
+                candidate.initialize_weights()
+            candidate.train()
+            assert sum(p.numel() for p in candidate.parameters()) == total
+            return FSDP(
+                candidate,
+                sharding_strategy=ShardingStrategy.FULL_SHARD,
+                use_orig_params=True,
+                device_id=DEVICE,
+                sync_module_states=False,
+                limit_all_gathers=True,
+            )
 
-        model = FSDP(
-            model,
-            sharding_strategy=ShardingStrategy.FULL_SHARD,
-            use_orig_params=True,
-            device_id=DEVICE,
-            sync_module_states=False,
-            limit_all_gathers=True,
-        )
+        model = allocate_model(resume_path is None)
 
         # Probe before creating/consuming a real data stream. This catches wiring,
         # gradient, NaN, or memory failures before the 50B run starts.
         if resume_path is None:
             run_startup_model_probe(model, len(tokenizer))
 
-        optimizer = torch.optim.AdamW(
-            model.parameters(),
-            lr=PEAK_LR,
-            betas=(0.9, 0.95),
-            eps=1e-8,
-            weight_decay=WEIGHT_DECAY,
-            fused=True,
-        )
+        def create_optimizer(model):
+            return torch.optim.AdamW(
+                model.parameters(),
+                lr=PEAK_LR,
+                betas=(0.9, 0.95),
+                eps=1e-8,
+                weight_decay=WEIGHT_DECAY,
+                fused=True,
+            )
+
+        optimizer = create_optimizer(model)
 
         progress = {
             "step": 0,
@@ -3359,6 +3597,86 @@ def main():
             "skipped_updates": 0,
             "consecutive_skips": 0,
         }
+
+        if PREFLIGHT_SMOKE:
+            rank0_print("Running one bounded synthetic optimizer update...", flush=True)
+            optimizer.zero_grad(set_to_none=True)
+            x = torch.randint(
+                0, len(tokenizer), (MICRO_BATCH_SIZE, SEQUENCE_LENGTH),
+                device=DEVICE, dtype=torch.long,
+            )
+            y = torch.randint(
+                0, len(tokenizer), (MICRO_BATCH_SIZE, SEQUENCE_LENGTH),
+                device=DEVICE, dtype=torch.long,
+            )
+            model.set_specialization_scale(
+                min(1.0, 1.0 / max(1, SPEC_WARMUP_STEPS))
+            )
+            with torch.autocast(
+                "cuda", dtype=torch.bfloat16, cache_enabled=AUTOCAST_CACHE
+            ):
+                ce, auxiliary = model(x, y)
+                loss = ce + auxiliary
+            finite_loss = torch.isfinite(loss.detach()).to(dtype=torch.int32)
+            dist.all_reduce(finite_loss, op=dist.ReduceOp.MIN)
+            if not bool(finite_loss.item()):
+                raise RuntimeError("Preflight smoke produced a nonfinite loss.")
+            loss.backward()
+            norm = model.clip_grad_norm_(GRAD_CLIP)
+            finite_grad = torch.isfinite(norm).to(dtype=torch.int32)
+            dist.all_reduce(finite_grad, op=dist.ReduceOp.MIN)
+            if not bool(finite_grad.item()):
+                raise RuntimeError("Preflight smoke produced nonfinite gradients.")
+            lr = learning_rate(1)
+            for group in optimizer.param_groups:
+                group["lr"] = lr
+            optimizer.step()
+            optimizer.zero_grad(set_to_none=True)
+            progress.update({
+                "step": 1,
+                "tokens": TOKENS_PER_UPDATE,
+                "sessions": 1,
+            })
+            del x, y, ce, auxiliary, loss, norm
+            gc.collect()
+            torch.cuda.empty_cache()
+
+            # The existing checkpoint format is exercised locally.  The data
+            # state is deliberately absent because this path never opens the
+            # 50B-token stream.
+            smoke_data_state = {"mode": "preflight_smoke", "cursor": None}
+            smoke_path = save_local_checkpoint(
+                model, optimizer, tokenizer, progress, smoke_data_state
+            )
+            smoke_path = broadcast_object(
+                str(smoke_path) if RANK == 0 else None
+            )
+            dist_barrier()
+            rank0_print("Local preflight checkpoint saved:", smoke_path, flush=True)
+
+            del model, optimizer
+            gc.collect()
+            torch.cuda.empty_cache()
+
+            fresh_model = allocate_model(False)
+            fresh_optimizer = create_optimizer(fresh_model)
+            restored = restore_checkpoint(
+                Path(smoke_path), fresh_model, fresh_optimizer
+            )
+            if (
+                restored.get("progress", {}).get("step") != 1
+                or restored.get("progress", {}).get("tokens") != TOKENS_PER_UPDATE
+            ):
+                raise RuntimeError("Preflight smoke restore returned unexpected progress.")
+            dist_barrier()
+            rank0_print(
+                "PREFLIGHT SMOKE PASS: startup probe, one optimizer update, "
+                "local save, and fresh model/optimizer restore completed.",
+                flush=True,
+            )
+            del fresh_model, fresh_optimizer, restored
+            return
+
         saved_data = None
         if resume_path is not None:
             rank0_print("Restoring Nano model + optimizer...", flush=True)
@@ -3453,28 +3771,45 @@ def main():
         def save_checkpoint(reason):
             nonlocal latest_local, last_save_time, dirty, checkpoint_size_estimate
             if RANK == 0:
-                hub.wait()
-                staging_root = WORK / "upload-staging"
-                if staging_root.exists():
-                    shutil.rmtree(staging_root)
+                try:
+                    hub.wait()
+                    staging_root = WORK / "upload-staging"
+                    if staging_root.exists():
+                        shutil.rmtree(staging_root)
+                except BaseException as exc:
+                    broadcast_object(f"Checkpoint preparation failed: {type(exc).__name__}")
+                    raise
+            preparation_error = broadcast_object(None)
+            if preparation_error is not None:
+                raise RuntimeError(preparation_error)
             optimizer.zero_grad(set_to_none=True)
             start = time.monotonic()
             rank0_print(f"\nSaving ({reason}) at step {progress['step']:,}...", flush=True)
             path = save_local_checkpoint(
                 model, optimizer, tokenizer, progress, committed_data
             )
+            saved = None
             if RANK == 0:
-                latest_local = path
-                checkpoint_size_estimate = folder_bytes(path)
-                prune_local_checkpoints(path)
-                last_save_time = time.monotonic()
-                dirty = False
-                rank0_print(
-                    f"Local save finished in {time.monotonic() - start:.1f}s | "
-                    f"{checkpoint_size_estimate / 1024**3:.2f} GiB\n"
-                    f"Local checkpoint: {path}", flush=True,
-                )
-                hub.upload_async(path, progress)
+                try:
+                    size = folder_bytes(path)
+                    prune_local_checkpoints(path)
+                    rank0_print(
+                        f"Local save finished in {time.monotonic() - start:.1f}s | "
+                        f"{size / 1024**3:.2f} GiB\n"
+                        f"Local checkpoint: {path}", flush=True,
+                    )
+                    hub.upload_async(path, progress)
+                    saved = {"path": str(path), "size": size, "time": time.monotonic()}
+                except BaseException as exc:
+                    broadcast_object({"error": f"Checkpoint finalization failed: {type(exc).__name__}"})
+                    raise
+            saved = broadcast_object(saved)
+            if "error" in saved:
+                raise RuntimeError(saved["error"])
+            latest_local = Path(saved["path"])
+            checkpoint_size_estimate = saved["size"]
+            last_save_time = saved["time"]
+            dirty = False
 
         metrics = torch.zeros(2, device=DEVICE, dtype=torch.float32)
         telemetry_probe = None
@@ -3521,8 +3856,8 @@ def main():
                             loss = (ce + auxiliary) / GRAD_ACCUM_STEPS
                     except torch.cuda.OutOfMemoryError as exc:
                         raise RuntimeError(
-                            "Nano turbo profile OOM. Set MICRO_BATCH_SIZE=56 "
-                            "(grad_accum becomes 2 while preserving 114,688 tokens/update) "
+                            "Nano turbo profile OOM. Set GLOBAL_MICRO_BATCH_SIZE=56 "
+                            "(7 sequences/GPU, grad_accum=2, 114,688 tokens/update) "
                             "before changing the architecture."
                         ) from exc
 
@@ -3611,7 +3946,7 @@ def main():
                     and progress["step"] % T2_TELEMETRY_EVERY_STEPS == 0
                 ):
                     probe_x, probe_y = telemetry_probe
-                    t2 = model.run_telemetry_probe(probe_x, probe_y)
+                    t2 = run_telemetry_probe(model, probe_x, probe_y)
                     bank_text = " ".join(
                         f"B{b['bank']}:H={b['router_entropy']:.2f},"
                         f"load={100*b['min_load']:.1f}-{100*b['max_load']:.1f}%,"
@@ -3647,7 +3982,7 @@ def main():
                 ):
                     probe_x, probe_y = telemetry_probe
                     n = min(ABLATION_PROBE_BATCH, probe_x.shape[0])
-                    a = model.run_ablation_probe(probe_x[:n], probe_y[:n])
+                    a = run_ablation_probe(model, probe_x[:n], probe_y[:n])
                     rank0_print(
                         "ABLATE "
                         f"step={progress['step']:,} full={a['full']:.4f} "
@@ -3658,28 +3993,30 @@ def main():
                         flush=True,
                     )
 
-                now = time.monotonic()
-                upload_estimate = (
-                    hub.estimate_upload_seconds(checkpoint_size_estimate)
-                    if hub is not None else 0.0
-                )
-                final_guard = max(
-                    FINAL_SAVE_GUARD_MINUTES * 60,
-                    2 * upload_estimate + SAVE_RESERVE_MINUTES * 60,
-                )
-                if (
-                    dirty and not STOP_REQUESTED
-                    and (hub is None or not hub.busy())
-                    and now - last_save_time >= CHECKPOINT_EVERY_MINUTES * 60
-                    and train_deadline() - now > final_guard
-                ):
+                # Only rank zero owns the upload state and scheduling clock.
+                # All ranks must nevertheless enter every collective save.
+                should_save = False
+                if RANK == 0:
+                    now = time.monotonic()
+                    upload_estimate = hub.estimate_upload_seconds(checkpoint_size_estimate)
+                    final_guard = max(
+                        FINAL_SAVE_GUARD_MINUTES * 60,
+                        2 * upload_estimate + SAVE_RESERVE_MINUTES * 60,
+                    )
+                    should_save = (
+                        dirty and not STOP_REQUESTED
+                        and not hub.busy()
+                        and now - last_save_time >= CHECKPOINT_EVERY_MINUTES * 60
+                        and train_deadline() - now > final_guard
+                    )
+                if broadcast_object(should_save):
                     save_checkpoint("periodic")
                     metrics.zero_()
                     log_updates = 0
                     log_tokens = 0
                     log_start = time.monotonic()
 
-            if dirty:
+            if broadcast_object(dirty if RANK == 0 else None):
                 save_checkpoint("final")
             elif RANK == 0 and latest_local is not None:
                 pointer = hub.pointer
@@ -3713,7 +4050,7 @@ def main():
                 "Data pipeline failed. Preserving the last complete optimizer boundary.",
                 flush=True,
             )
-            if dirty and not in_optimizer_step:
+            if broadcast_object(dirty and not in_optimizer_step if RANK == 0 else None):
                 try:
                     save_checkpoint("data pipeline recovery")
                 except Exception as exc:
@@ -3726,8 +4063,9 @@ def main():
             torch.cuda.empty_cache()
             rank0_print(
                 "\nCUDA OOM. Existing checkpoints remain valid.\n"
-                "Set MICRO_BATCH_SIZE=56 first; GRAD_ACCUM_STEPS will become 2 "
-                "and keep 114,688 tokens/update. If necessary try 28/4 next.\n"
+                "Set GLOBAL_MICRO_BATCH_SIZE=56 first (7 sequences/GPU); "
+                "GRAD_ACCUM_STEPS becomes 2, keeping 114,688 tokens/update. "
+                "If necessary try GLOBAL_MICRO_BATCH_SIZE=16 (2/GPU, accumulation=7).\n"
                 "No checkpoint is taken from a partial optimizer update.",
                 flush=True,
             )
@@ -3751,7 +4089,8 @@ def main():
         if hub is not None:
             rank0_print("Waiting for any pending checkpoint upload...", flush=True)
             hub.close()
-        dist_barrier()
+        # Do not enter a new collective during exception unwinding: another
+        # rank may still be in a different collective when torchrun aborts it.
         if dist.is_initialized():
             dist.destroy_process_group()