From f711cf74405d299a5f9098a24135ab72f9f2c62a Mon Sep 17 00:00:00 2001 From: oyazdanb Date: Fri, 5 Jun 2026 11:32:28 -0400 Subject: [PATCH 1/6] race/fsdp: real transformer compute + observable per-layer checksum detector Builds on PR #210 (shared-weight transformer + int16 checksums). Three changes: 1. compute_type=transformer now runs a REAL transformer block borrowed from aorta.models.RepeatedTransformerBlock (the model llm_determinism uses: real MHA + GLU FFN + LayerNorm) instead of a single gelu(W @ x) matmul. One shared block (fixed seed, rank-identical via fork_rng) + a fixed 3D reference input keeps every layer's output byte-identical, preserving the per-layer checksum invariant while supplying real-model L2/HBM pressure. Explicit all_gather/reduce_scatter (the AINIC data path) are untouched. 2. Observability: a clean run previously emitted no checksum signal, so green was indistinguishable from a no-op. Track layers_verified and layer_checksum_mismatches and surface them (plus compute_type) in WorkloadResult.metrics; the FSDP startup log now names the active compute path so a silent GEMM fallback is greppable. 3. Validation: reject unknown compute_type values (a typo like "transfomer" used to silently fall back to GEMM = false green) and warn when shared_layer_weights is set without compute_type=transformer. New config fields: num_heads, ffn_size, seq_len, batch_size. Tests (CPU-only, no GPU/dist): test_race_transformer_smoke.py drives the real block + checksum verifier end to end (clean pass + injected-corruption catch); test_race_checksums.py covers the verifier; test_race.py adds compute_type validation cases. 15 race tests pass. Co-Authored-By: Claude Opus 4 --- .../race-transformer-compute-support.md | 123 ++++++++++ recipes/ainic-gdr-flush-sdc.yaml | 25 +- src/aorta/race/base.py | 11 + src/aorta/race/config.py | 41 ++++ src/aorta/race/modes/fsdp.py | 232 ++++++++++++++++-- src/aorta/workloads/race.py | 18 +- tests/workloads/test_race.py | 17 ++ tests/workloads/test_race_checksums.py | 108 ++++++++ .../workloads/test_race_transformer_smoke.py | 125 ++++++++++ 9 files changed, 666 insertions(+), 34 deletions(-) create mode 100644 docs/design/race-transformer-compute-support.md create mode 100644 tests/workloads/test_race_checksums.py create mode 100644 tests/workloads/test_race_transformer_smoke.py diff --git a/docs/design/race-transformer-compute-support.md b/docs/design/race-transformer-compute-support.md new file mode 100644 index 00000000..a72af358 --- /dev/null +++ b/docs/design/race-transformer-compute-support.md @@ -0,0 +1,123 @@ +# Plan — Fully support `compute_type: transformer` + shared-weight checksums (PR #210) + +**Branch:** `users/oyazdanb/race-transformer-compute` (off `main`) +**Status:** DRAFT — for review before implementation +**Related:** PR #210 (`users/mycpuorg/shared-weight-transformer-checksums`), task `race__cross-rank-and-iter0-detection.md` + +## TL;DR + +PR #210's transformer + per-layer-checksum code appears **complete and correctly wired on its own branch** (config field, `setup_buffers` shared-weight path, `_forward_layer` dual dispatch, `_verify_layer_checksums` defined AND called). Yet the 2026-06-05T01-12-30 run — which *accepted* `compute_type: transformer` + `shared_layer_weights: true` (no "ignoring unknown key" warning) — showed **0.0 ms step time and emitted no checksum signal**, so we cannot confirm the path actually executed. The detector is **unobservable on success**: it logs/raises only on mismatch and writes no metric, so a clean run looks identical to a no-op. + +**This plan does NOT rewrite the compute path** (the branch already has it). It (1) resolves the did-it-run ambiguity, (2) makes the detector self-evidencing, and (3) hardens config validation so silent fallback is impossible. Work lands as a small change set that can either be merged into #210 or layered on top. + +## UPDATE (2026-06-05): PR #210's "transformer" is not a transformer — Option A + +Reading the PR branch source: `compute_type: transformer` does **`out = gelu(weight_matrix @ reference_input)`** (`fsdp.py:240-241`) — a single matmul, **no attention, no FFN, no softmax**. The "transformer" label refers to the *checksum scheme*, not the compute. So: +- **Real transformer model?** No (mm+gelu). +- **Real `torch.distributed.fsdp`?** No (simulated shards + REAL explicit `all_gather`/`reduce_scatter`). +- Both `gemm` and `transformer` compute_types are effectively a matmul. + +**Consequence for PR D / L2 pressure:** PR #210's transformer does NOT supply real-model L2/HBM pressure (the documented missing ingredient for small-scale repro). A single matmul ≠ a real block's memory traffic. So #210 is *not* the cheap step toward PR D I first assumed. + +**Option A (chosen):** make `compute_type: transformer` run a REAL transformer block by borrowing `RepeatedTransformerBlock` from `aorta.models.repeated_block` — the same model `llm_determinism` uses (real MHA + GLU FFN + LayerNorm, torch-only, no deps). This gives real L2/HBM pressure *inside* the existing race harness (explicit collectives preserved), without the full `torch.distributed.fsdp` refactor that is PR D. It is the recommended-for-PR-D model per the task file, used here as a compute kernel rather than under FSDP2. + +### Option A — implementation spec (source-grounded; refs on PR #210 branch unless noted) + +**Compute site:** `_forward_layer` shared path `fsdp.py:240-241` (fwd), `_backward_layer` `fsdp.py:279-281` (bwd). The `full_param` (all_gather output) is only *checksummed*, never fed to compute — keep that true. + +**Shape adaptation:** current compute is 2D `[dim,dim] @ [dim,dim]`. `RepeatedTransformerBlock.forward` (`repeated_block.py:108`) needs 3D `[batch, seq, hidden]`, `hidden==model_dim`, `hidden % num_heads == 0`. So `reference_input` becomes 3D `[batch, seq, model_dim]`. + +**Shared-weight invariant (the crux):** PR #210 already shares ONE weight across layers — `self.weight_matrices = [shared_w] * num_layers` (`fsdp.py:148-153`), `reference_input` seed 1. To keep "all layers byte-identical" with a real block, build **ONE** `RepeatedTransformerBlock` (not `RepeatedBlockModel` — we want a single block, no embed/LM-head) under a fixed seed and call it for every layer: +```python +# in setup_buffers, transformer+shared path +with torch.random.fork_rng(devices=["cuda"]): + torch.cuda.manual_seed(0) # identical weights on EVERY rank + self.shared_block = RepeatedTransformerBlock(block_cfg).to("cuda").to(self.dtype) +self.shared_block.eval() +g = torch.Generator(device="cuda"); g.manual_seed(1) +self.reference_input = torch.randn(batch, seq, model_dim, dtype=self.dtype, device="cuda", generator=g) +``` +In `_forward_layer` shared branch replace mm+gelu with `out = self.shared_block(self.reference_input)` under `torch.no_grad()`. The 4 checksums map unchanged: `comm_input`=param shard, `comm_output`=full_param post-all_gather, `compute_input`=`_checksum(reference_input)` (`_checksum` `fsdp.py:200` is shape-agnostic), `compute_output`=`_checksum(out)`. + +**Config mapping → `BlockConfig`:** `hidden_size <- model_dim`; build a single block (loop num_layers ourselves). **New `ReproducerConfig` fields** (`config.py` transformer section ~line 158): `num_heads: int = 16` (validate `model_dim % num_heads == 0`, else BlockConfig raises `repeated_block.py:44`), `ffn_size: int = 0` (0 ⇒ derive `4*model_dim`), `seq_len: int = 512`, `batch_size: int = 1`. **Reused:** `model_dim`, `num_layers`, `dtype`, `shared_layer_weights`, `compute_type`, `simulate_compute`, `include_backward_compute`. + +**Backward:** current is manual `mm(W.T, grad)` (`fsdp.py:279-281`), NOT autograd; `reference_input.requires_grad=False`. **Keep backward as a timing-only proxy** — when `include_backward_compute`, call `self.shared_block(self.reference_input)` a second time under `no_grad`. Do NOT add autograd (breaks no-grad determinism, balloons memory, scope creep). + +**Stays untouched (exercises AINIC):** all_gather (`fsdp.py:228,272`), reduce_scatter (`:284`), `_fill_patterns`, `_verify_all_gather` (`:439`), `_verify_reduce_scatter` (`:473`) — they read `full_param`/`grad_shard`, never `reference_input`. Only the compute kernel changes. + +**Concrete edits:** +- `config.py` ~L158: add `num_heads`, `ffn_size`, `seq_len`, `batch_size`. +- `fsdp.py:30`: `from aorta.models import BlockConfig, RepeatedTransformerBlock`. +- `fsdp.py:80-82` `__init__`: `self.shared_block = None`. +- `fsdp.py:144-159` `setup_buffers`: block construction + 3D reference_input (transformer+shared path); leave non-shared 2D path as-is. +- `fsdp.py:236-251` `_forward_layer` shared branch: real block forward under `no_grad`. +- `fsdp.py:277-281` `_backward_layer`: block-forward proxy on shared path. + +**Pitfalls:** `RepeatedTransformerBlock` defaults float32 → `.to(self.dtype)` mandatory (race=bf16). LayerNorm/softmax upcast (`repeated_block.py:119`) is deterministic — fine. Memory: real block activations `[batch,seq,hidden]` + attention `[batch,heads,seq,seq]` ≫ 2D matmul → default `batch_size=1, seq_len=512`; block built once (shared) so param memory ≈ one layer regardless of `num_layers`. + +**Preserves corruption localization: YES** — `_verify_layer_checksums` (`fsdp.py:419-436`) is checksum-key-agnostic: comm mismatch ⇒ RCCL/NIC, compute mismatch ⇒ GPU compute, verbatim. Rank-fill checks on `full_param`/`grad_shard` are untouched (LOW risk). + +**One real setup risk:** if ranks build the block with diverging RNG, `compute_output` differs *across ranks* but stays identical *across layers within a rank* → intra-rank `_verify_layer_checksums` still PASSES, masking a false setup bug. Mitigate: seed global CUDA RNG identically on all ranks before block construction (the `fork_rng` above), and optionally add a one-time cross-rank `compute_output` all-reduce equality assert at setup. + +**Effort:** ~medium / half-day. ~4 edited functions + 4 config fields + 1 import. No new files. The earlier Steps 1–3 (observability, config validation, detector unit test) still apply on top. + +## Step 0 — Resolve the contradiction FIRST (no code yet) + +Before writing anything, confirm what actually ran on the cluster. On the cluster: +``` +cd /it-share/oyazdanb/aorta && git branch --show-current +python - <<'PY' +import inspect, aorta.race.modes.fsdp as f +src = inspect.getsource(f) +for tok in ("_verify_layer_checksums","_checksum","reference_input","use_shared","shared_w"): + print(tok, tok in src) +PY +``` +- All `True` → branch code IS installed; the path is *present* but unobservable → go to Steps 1–3. +- Any `False` → a stale/main build is installed; `pip install -e` didn't take → reinstall, rerun, re-check. (This alone may explain the 0.0 ms / no-metric run.) + +Also: 0.0 ms step is suspicious even for GEMM. Confirm `simulate_compute` is actually doing work — a real 24-layer transformer fwd+bwd at model_dim=1024 cannot be 0.0 ms. If it is 0.0 ms with the branch installed, compute is being skipped or mistimed — that's a real bug to find in `_forward_layer`/timing, not just observability. + +## Step 1 — Make the detector self-evidencing (observability) + +Root problem: a green result does not prove the checksum verifier ran. Fix so it is impossible to miss. + +- **Startup log (once per trial):** in `setup_buffers`/`run` when `compute_type=="transformer"` and `shared_layer_weights`, emit e.g. + `log.info("race: compute=transformer shared_layer_weights=ON layers=%d dim=%d; per-layer checksum verify ENABLED", num_layers, model_dim)`. + And in the GEMM branch: `log.info("race: compute=gemm")`. Now fallback is one grep away. +- **Result metric (always, not just on failure):** add to `WorkloadResult.metrics` (mapped in `workloads/race.py`): + - `layers_verified` (int, per step or total), + - `layer_checksum_mismatches` (int, 0 on clean), + - `compute_type` (echo the effective value). + A clean run then shows `layer_checksum_mismatches: 0, layers_verified: >0` — provably ran. Today `metrics` is only `{avg_step_time_ms, mode, rank, world_size}`. +- **Per-step debug counter** behind `log_interval` so long runs show progress. + +## Step 2 — Harden config so silent fallback is impossible + +- **Validate `compute_type`** against `{"gemm","transformer"}` in `workloads/race.py` (mirror the existing `_VALID_DTYPES` pattern). Today any string is accepted, so a typo (`transfomer`) silently runs GEMM with no warning. +- **Warn on inert combo:** if `shared_layer_weights=True` but `compute_type!="transformer"`, log a WARNING — currently it silently no-ops via the `use_shared` AND-gate. +- Keep accepting the keys (they already round-trip on the branch) — this is validation, not new surface. + +## Step 3 — Confirm the checksum actually catches corruption (test the detector) + +A detector that never fires on a clean cluster is unproven. Add a unit/integration check: +- **Unit:** forge a layer activation buffer so one layer's checksum differs; assert `_verify_layer_checksums` reports a mismatch with the right layer index. (No GPU/dist needed.) +- **Negative:** identical layers → 0 mismatches. +- This proves the second signal works regardless of whether the AINIC bug reproduces. + +## Files + +- `src/aorta/race/modes/fsdp.py` — startup log; mismatch counter; ensure `_verify_layer_checksums` returns counts (verified, mismatches) rather than only logging. +- `src/aorta/workloads/race.py` — `_VALID_COMPUTE_TYPES` validation; warn on inert `shared_layer_weights`; map new metrics into `WorkloadResult.metrics`. +- `tests/workloads/test_race.py` (or race-mode test) — checksum-catches-forged-corruption + clean-pass + config-validation cases. +- `recipes/ainic-gdr-flush-sdc.yaml` — already correct on #210 (dtype `bfloat16`, transformer, shared_layer_weights). No change beyond confirming. + +## Relationship to PR #210 and PR D + +- **PR #210:** this is *complementary* — it makes #210's detector observable, validated, and tested. Decide at review: fold these commits into #210, or land as a follow-up on this branch. +- **PR D (real torch.distributed.fsdp):** **deferred — do NOT start.** Per the task file, the 5-node no-repro triggers **PR C first** (cold-QP/iter-0), and PR D only if real-model L2 pressure is *confirmed* the missing ingredient. PR #210's transformer compute is the cheap way to raise model-realistic L2/HBM pressure *inside* the existing harness — test that hypothesis before any PR D refactor. Also: PR D would **not** close the iter-0/cold-QP gap (both reuse the launcher's warm PG); only PR C does. + +## Open questions for review +1. Step 0 result: is the branch actually installed on the cluster (path present) or is the 0.0 ms / no-metric run a stale-build artifact? This decides whether Steps 1–3 are the whole job or whether there's also a real compute-skip bug. +2. Fold into #210 vs land as follow-up branch? +3. Does `_verify_layer_checksums` currently return counts, or only log? (Determines how much of Step 1's metric plumbing is new.) diff --git a/recipes/ainic-gdr-flush-sdc.yaml b/recipes/ainic-gdr-flush-sdc.yaml index 5d2d4fb6..80b77a4a 100644 --- a/recipes/ainic-gdr-flush-sdc.yaml +++ b/recipes/ainic-gdr-flush-sdc.yaml @@ -91,18 +91,25 @@ workload_config: # 1M elements x 2 B = 2 MB shard; 2 MB x (N-1) all_gather fan-out per layer. fsdp_shard_size: 1000000 - # GEMM compute between collectives keeps the reproducer near production - # per-step timing (~500ms) so the L2->HBM writeback race window matches real - # training rather than running orders of magnitude faster. gemm_layers ~36 at - # gemm_size 5120 (~14ms/layer on MI300X) approximates the documented ~500ms. - # model_dim / num_layers carried over from the original model: block. + # Transformer compute between collectives keeps the reproducer near production + # per-step timing so the L2->HBM writeback race window matches real training. + # model_dim=1024, num_layers=24: attention + FFN per layer, forward + backward, + # approximates ~500ms/step on MI355X. + # + # shared_layer_weights=true: all 24 layers share one weight matrix (fixed seed=0) + # and each layer independently receives the same fixed reference input (seed=1) + # rather than chaining activations layer-to-layer. This makes every layer's + # forward output analytically identical so _verify_layer_activations() can serve + # as a second independent corruption signal alongside the collective-buffer + # pattern check: a mismatch across layers means an all_gather corruption on that + # layer propagated through GEMM compute -- something the rank-fill pattern check + # on full_param alone cannot catch. simulate_compute: true - compute_type: gemm - gemm_size: 5120 - gemm_layers: 36 - include_backward_compute: true + compute_type: transformer model_dim: 1024 num_layers: 24 + include_backward_compute: true + shared_layer_weights: true cells: # ------------------------------------------------------------------ diff --git a/src/aorta/race/base.py b/src/aorta/race/base.py index 9d9a180f..2b17db16 100644 --- a/src/aorta/race/base.py +++ b/src/aorta/race/base.py @@ -78,6 +78,12 @@ def __init__(self, config: ReproducerConfig, rank: int, world_size: int): self.in_verification_phase: bool = False self.corruption_details: List[Dict] = [] + # Detector observability: a clean run is otherwise indistinguishable + # from a no-op. Subclasses that do per-layer checksum verification + # increment these so the result proves the detector executed. + self.layers_verified: int = 0 + self.layer_checksum_mismatches: int = 0 + # Dtype self.dtype = self._get_dtype() @@ -85,8 +91,11 @@ def _get_dtype(self) -> torch.dtype: """Get torch dtype from config string.""" dtype_map = { "bfloat16": torch.bfloat16, + "bf16": torch.bfloat16, "float16": torch.float16, + "fp16": torch.float16, "float32": torch.float32, + "fp32": torch.float32, } return dtype_map.get(self.config.dtype, torch.bfloat16) @@ -528,6 +537,8 @@ def run(self) -> ReproducerResult: corruption_details=self.corruption_details, elapsed_time_sec=elapsed, avg_step_time_ms=avg_step_ms, + layers_verified=self.layers_verified, + layer_checksum_mismatches=self.layer_checksum_mismatches, ) diff --git a/src/aorta/race/config.py b/src/aorta/race/config.py index 4e47d0e4..5eee1d32 100644 --- a/src/aorta/race/config.py +++ b/src/aorta/race/config.py @@ -156,9 +156,43 @@ class ReproducerConfig: 4 layers at model_dim=2048 gives ~300ms/step on MI350X. """ + num_heads: int = 0 + """ + Number of attention heads (transformer compute type). + + 0 means auto-derive as model_dim // 128. Must divide model_dim evenly. + """ + + ffn_size: int = 0 + """ + FFN intermediate (hidden) size (transformer compute type). + + 0 means auto-derive as model_dim * 4. + """ + + seq_len: int = 512 + """Sequence length of the reference input (transformer compute type).""" + + batch_size: int = 1 + """Batch size of the reference input (transformer compute type).""" + include_backward_compute: bool = True """Also simulate backward pass (doubles compute time).""" + shared_layer_weights: bool = False + """ + Use a single shared weight matrix for all layers (transformer compute type). + + When True, all num_layers layers use the same weight matrix initialized with + a fixed seed, and each layer receives the same fixed reference input rather + than chaining activations. This makes per-layer forward outputs analytically + identical so that cross-layer activation comparison can serve as a secondary + corruption signal: any all_gather corruption that propagates through GEMM + compute will produce a mismatch between layer outputs. + + Only meaningful when compute_type == 'transformer'. + """ + # ========================================================================= # Optimizer (used by modes that support it, e.g., DDP) # ========================================================================= @@ -281,6 +315,13 @@ class ReproducerResult: avg_step_time_ms: float """Average time per step in milliseconds.""" + layers_verified: int = 0 + """Cross-layer checksum comparisons performed (shared-weight transformer). + >0 proves the per-layer checksum detector actually ran; 0 means it did not.""" + + layer_checksum_mismatches: int = 0 + """Per-layer checksum mismatches found (0 on a clean run).""" + # ============================================================================= # Race Injection Config (broader aorta system) diff --git a/src/aorta/race/modes/fsdp.py b/src/aorta/race/modes/fsdp.py index eda477b8..4720a032 100644 --- a/src/aorta/race/modes/fsdp.py +++ b/src/aorta/race/modes/fsdp.py @@ -30,6 +30,8 @@ import torch import torch.distributed as dist +from aorta.models import BlockConfig, RepeatedTransformerBlock + from ..base import BaseReproducer from ..config import ReproducerConfig @@ -77,6 +79,13 @@ def __init__(self, config: ReproducerConfig, rank: int, world_size: int): self.activation: Optional[torch.Tensor] = None self.grad_buffer: Optional[torch.Tensor] = None + # Shared-weight transformer: fixed reference input + per-layer checksums + self.reference_input: Optional[torch.Tensor] = None + self.layer_checksums: List[Optional[dict]] = [] + + # Real transformer block shared across all layers (shared-weight path) + self.shared_block: Optional[RepeatedTransformerBlock] = None + def _setup_compute(self) -> None: """ Override base compute setup -- FSDP manages its own per-layer compute. @@ -90,7 +99,10 @@ def _setup_compute(self) -> None: if not self.config.simulate_compute: return - # Validate buffer sizes based on compute type + # Validate buffer sizes based on compute type. + # NOTE: for the shared-weight transformer path the real compute size is + # governed by batch_size × seq_len × model_dim (the block's activation), + # not dim²; this min only sizes the H2D staging buffer and stays harmless. dim = self._dim min_h2d_size = dim * dim if self.config.h2d_tensor_size < min_h2d_size: @@ -134,13 +146,56 @@ def setup_buffers(self) -> None: # per-layer, unlike the base compute simulator which runs all layers at once. if cfg.simulate_compute: dim = self._dim - self.weight_matrices = [ - torch.randn( - dim, dim, - dtype=self.dtype, device="cuda", + use_shared = ( + cfg.shared_layer_weights and cfg.compute_type == "transformer" + ) + if use_shared: + # All layers share ONE real transformer block with deterministic, + # rank-identical weights so block(reference_input) is analytically + # identical for every layer and every rank. Any divergence across + # layers indicates compute-path corruption. + hidden = cfg.model_dim + num_heads = cfg.num_heads or (hidden // 128) + if num_heads < 1: + num_heads = 1 + ffn = cfg.ffn_size or (hidden * 4) + if hidden % num_heads != 0: + raise ValueError( + f"model_dim ({hidden}) must be divisible by num_heads " + f"({num_heads}) for shared-weight transformer compute" + ) + block_cfg = BlockConfig( + hidden_size=hidden, + num_heads=num_heads, + num_layers=1, + ffn_size=ffn, + num_experts=1, + seq_len=cfg.seq_len, + vocab_size=16, # embed is unused on this path; keep tiny ) - for _ in range(self.num_layers) - ] + # fork_rng so seeding the block's init RNG to a fixed value gives + # bit-identical weights on every rank without perturbing global RNG. + with torch.random.fork_rng(devices=["cuda"]): + torch.cuda.manual_seed(0) + self.shared_block = ( + RepeatedTransformerBlock(block_cfg).to("cuda").to(self.dtype) + ) + self.shared_block.eval() + + # Fixed reference input, same seed across all ranks and iterations. + g = torch.Generator(device="cuda") + g.manual_seed(1) + self.reference_input = torch.randn( + cfg.batch_size, cfg.seq_len, hidden, + dtype=self.dtype, device="cuda", generator=g, + ) + self.weight_matrices = [] + self.layer_checksums = [None] * self.num_layers + else: + self.weight_matrices = [ + torch.randn(dim, dim, dtype=self.dtype, device="cuda") + for _ in range(self.num_layers) + ] self.activation = torch.randn( dim, dim, dtype=self.dtype, device="cuda", @@ -150,11 +205,18 @@ def setup_buffers(self) -> None: dtype=self.dtype, device="cuda", ) + # Startup line names the active compute path so a silent fallback + # (e.g. transformer requested but GEMM ran) is greppable in logs. + shared_active = self.shared_block is not None log.info( f"Allocated FSDP buffers: layers={self.num_layers}, " f"shard_size={self.shard_size}, " f"full_param_size={self.shard_size * ws}, " - f"compute={'enabled' if cfg.simulate_compute else 'disabled'}" + f"compute={'enabled' if cfg.simulate_compute else 'disabled'}, " + f"compute_type={cfg.compute_type}, " + f"shared_layer_weights={cfg.shared_layer_weights}, " + f"transformer_block={'active' if shared_active else 'none'}, " + f"layer_checksum_verify={'ON' if shared_active else 'OFF'}" ) def _fill_patterns(self) -> None: @@ -166,30 +228,82 @@ def _fill_patterns(self) -> None: # Each rank fills full_grad with rank + 1 (for reduce_scatter verification) self.full_grad.fill_(float(self.rank + 1)) + @staticmethod + def _checksum(tensor: torch.Tensor) -> int: + """ + Bitwise checksum: reinterpret-cast to int16 and sum. + + bf16 (or any 16-bit dtype) is viewed as int16 so every bit pattern + contributes to the checksum with zero information loss -- no float + rounding, no abs(), and NaN / denorm bit patterns are included. + Accumulation is done in int64 to avoid overflow. + """ + return tensor.view(torch.int16).to(torch.int64).sum().item() + def _forward_layer(self, layer_idx: int) -> None: """ Forward pass for a single FSDP layer. 1. all_gather: reconstruct full parameter from shards across ranks 2. GEMM: compute with full parameter (if enabled) + + Shared-weight path: every layer receives the same fixed reference_input + so that outputs are analytically identical. Input/output checksums are + recorded for both the comm kernel (all_gather) and the compute kernel + (GEMM + GELU) so _verify_layer_checksums() can pinpoint whether + corruption entered during communication or compute. + + Chained path (default): layer 0 seeds activation from batch_gpu (H2D race + opportunity) and each subsequent layer receives the previous layer's output. """ - # all_gather: each rank contributes its shard → full_param + use_shared = ( + self.config.shared_layer_weights + and self.config.compute_type == "transformer" + and self.reference_input is not None + ) + + # ── comm kernel: all_gather ────────────────────────────────── + if use_shared: + comm_input_cksum = self._checksum(self.param_shards[layer_idx]) + dist.all_gather_into_tensor( self.full_param, self.param_shards[layer_idx] ) - # GEMM forward (if compute enabled) - if self.config.simulate_compute and self.weight_matrices: - # Use batch_gpu for data dependency on first layer (H2D race opportunity) - if layer_idx == 0: - dim = self._dim - batch_slice = self.batch_gpu[:dim * dim] - self.activation = batch_slice.view(dim, dim) + if use_shared: + comm_output_cksum = self._checksum(self.full_param) - self.activation = torch.mm( - self.weight_matrices[layer_idx], self.activation - ) - self.activation = torch.nn.functional.gelu(self.activation) + # ── compute kernel: transformer block (shared) or GEMM + GELU ── + # Gate admits the shared-block path even though weight_matrices is now + # empty on that path; the GEMM/chained else still requires weight_matrices. + if self.config.simulate_compute and ( + self.shared_block is not None or self.weight_matrices + ): + if use_shared: + compute_input_cksum = self._checksum(self.reference_input) + + with torch.no_grad(): + out = self.shared_block(self.reference_input) + + compute_output_cksum = self._checksum(out) + + self.layer_checksums[layer_idx] = { + "comm_input": comm_input_cksum, + "comm_output": comm_output_cksum, + "compute_input": compute_input_cksum, + "compute_output": compute_output_cksum, + } + self.activation = out + else: + # Use batch_gpu for data dependency on first layer (H2D race opportunity) + if layer_idx == 0: + dim = self._dim + batch_slice = self.batch_gpu[:dim * dim] + self.activation = batch_slice.view(dim, dim) + self.activation = torch.mm( + self.weight_matrices[layer_idx], self.activation + ) + self.activation = torch.nn.functional.gelu(self.activation) def _backward_layer(self, layer_idx: int) -> None: """ @@ -204,9 +318,15 @@ def _backward_layer(self, layer_idx: int) -> None: self.full_param, self.param_shards[layer_idx] ) - # GEMM backward (if compute enabled) - if self.config.simulate_compute and self.weight_matrices: - if self.config.include_backward_compute: + # Backward compute (if enabled) + if self.config.simulate_compute and self.config.include_backward_compute: + if self.shared_block is not None: + # Shared-transformer path: re-run forward as a backward timing + # proxy (we don't train, so an exact bwd kernel isn't needed — + # only the comm/compute overlap timing matters here). + with torch.no_grad(): + _ = self.shared_block(self.reference_input) + elif self.weight_matrices: self.grad_buffer = torch.mm( self.weight_matrices[layer_idx].T, self.grad_buffer ) @@ -290,7 +410,8 @@ def _run_iteration_prefetch(self, iteration: int) -> bool: return result def _verify(self, iteration: int) -> bool: - """Verify H2D, last all_gather, and last reduce_scatter results.""" + """Verify H2D, last all_gather, last reduce_scatter, and (if shared-weight + transformer) cross-layer activation consistency.""" all_correct = True # Check H2D result @@ -305,6 +426,69 @@ def _verify(self, iteration: int) -> bool: if not self._verify_reduce_scatter(): all_correct = False + # Cross-layer checksum comparison (shared-weight transformer only) + if ( + self.config.shared_layer_weights + and self.config.compute_type == "transformer" + and self.layer_checksums + ): + if not self._verify_layer_checksums(iteration): + all_correct = False + + return all_correct + + def _verify_layer_checksums(self, iteration: int) -> bool: + """ + Verify that per-kernel int16 checksums are identical across all layers. + + With shared weights and a fixed reference input every layer runs the + same comm kernel (all_gather of rank-filled shard) and the same compute + kernel (GEMM + GELU with shared W and fixed reference_input). Both the + input and output of each kernel are checksummed via reinterpret-cast to + int16 → int64 sum, so every bit contributes with zero information loss. + + Four checksums per layer: + comm_input -- param shard before all_gather (should be identical: + every shard is filled with float(rank)) + comm_output -- full_param after all_gather + compute_input -- reference_input fed to GEMM (constant across layers) + compute_output-- activation after GELU + + If comm_output diverges but comm_input matches, corruption is in the + collective (RCCL / NIC path). If compute_output diverges but + comm_output matches, corruption is in the compute kernel (GPU ALU). + """ + ref = self.layer_checksums[0] + if ref is None: + return True + + all_correct = True + for i in range(1, len(self.layer_checksums)): + cmp = self.layer_checksums[i] + if cmp is None: + continue + # Count every cross-layer comparison so a clean (green) run still + # proves the detector ran: layers_verified > 0. + self.layers_verified += 1 + for key in ("comm_input", "comm_output", "compute_input", "compute_output"): + if cmp[key] != ref[key]: + log.error( + f"LAYER_CHECKSUM_MISMATCH ({key}): " + f"rank={self.rank} iter={iteration} " + f"layer_0={ref[key]} layer_{i}={cmp[key]}" + ) + self.corruption_details.append({ + "type": f"layer_checksum_mismatch_{key}", + "rank": self.rank, + "iteration": iteration, + "layer_ref": 0, + "layer_cmp": i, + "ref_checksum": ref[key], + "cmp_checksum": cmp[key], + }) + self.layer_checksum_mismatches += 1 + all_correct = False + return all_correct def _verify_all_gather(self) -> bool: diff --git a/src/aorta/workloads/race.py b/src/aorta/workloads/race.py index ace41e9f..26575a0c 100644 --- a/src/aorta/workloads/race.py +++ b/src/aorta/workloads/race.py @@ -30,7 +30,8 @@ log = logging.getLogger(__name__) _VALID_MODES = {"default", "ddp", "fsdp"} -_VALID_DTYPES = {"bfloat16", "float16", "float32"} +_VALID_DTYPES = {"bfloat16", "bf16", "float16", "fp16", "float32", "fp32"} +_VALID_COMPUTE_TYPES = {"gemm", "transformer"} # Platform-injected config keys that are NOT ReproducerConfig fields but are # always present (the dispatcher writes `steps` into every workload config; @@ -60,6 +61,18 @@ def _race_config_from_dict(self, d: dict[str, Any]) -> ReproducerConfig: raise ValueError(f"mode must be one of {sorted(_VALID_MODES)}, got {cfg.mode!r}") if cfg.dtype not in _VALID_DTYPES: raise ValueError(f"dtype must be one of {sorted(_VALID_DTYPES)}, got {cfg.dtype!r}") + if cfg.compute_type not in _VALID_COMPUTE_TYPES: + # Reject typos (e.g. "transfomer") that would silently fall back to + # the GEMM path and produce a false green. + raise ValueError( + f"compute_type must be one of {sorted(_VALID_COMPUTE_TYPES)}, got {cfg.compute_type!r}" + ) + if cfg.shared_layer_weights and cfg.compute_type != "transformer": + log.warning( + "race: shared_layer_weights=True has no effect with compute_type=%r " + "(only applies to compute_type='transformer')", + cfg.compute_type, + ) return cfg def setup(self) -> None: @@ -90,6 +103,9 @@ def run(self) -> WorkloadResult: metrics={ "avg_step_time_ms": res.avg_step_time_ms, "mode": self._cfg.mode, + "compute_type": self._cfg.compute_type, + "layers_verified": res.layers_verified, + "layer_checksum_mismatches": res.layer_checksum_mismatches, "rank": self._rank, "world_size": self._world, }, diff --git a/tests/workloads/test_race.py b/tests/workloads/test_race.py index 7a38b276..64f0b0f6 100644 --- a/tests/workloads/test_race.py +++ b/tests/workloads/test_race.py @@ -85,6 +85,23 @@ def test_race_config_from_dict_rejects_bad_dtype(): wl._race_config_from_dict({"dtype": "int8"}) +def test_race_config_from_dict_rejects_bad_compute_type(): + wl = RaceWorkload({}) + # A typo like "transfomer" must error, not silently fall back to GEMM. + with pytest.raises(ValueError, match="compute_type must be one of"): + wl._race_config_from_dict({"compute_type": "transfomer"}) + + +def test_race_config_warns_shared_weights_without_transformer(caplog): + wl = RaceWorkload({}) + with caplog.at_level("WARNING"): + cfg = wl._race_config_from_dict( + {"compute_type": "gemm", "shared_layer_weights": True} + ) + assert cfg.compute_type == "gemm" + assert any("shared_layer_weights" in r.message for r in caplog.records) + + def test_race_workload_maps_result(monkeypatch): """run() maps every ReproducerResult field onto WorkloadResult.""" stub_result = ReproducerResult( diff --git a/tests/workloads/test_race_checksums.py b/tests/workloads/test_race_checksums.py new file mode 100644 index 00000000..e807141c --- /dev/null +++ b/tests/workloads/test_race_checksums.py @@ -0,0 +1,108 @@ +"""Unit tests for FSDPModeReproducer._verify_layer_checksums. + +These are CPU-only: no GPU, no torch.distributed. The method under test only +reads `self.layer_checksums`, `self.rank`, and appends to +`self.corruption_details` -- it never touches CUDA. So we bypass __init__ with +object.__new__ and set just those three attributes. + +Contract (read from src/aorta/race/modes/fsdp.py): + _verify_layer_checksums(iteration) -> bool + - uses layer_checksums[0] as the reference dict + - returns True if every later layer matches the reference on all four keys + (comm_input, comm_output, compute_input, compute_output) + - returns True (no false positive) when the reference is None + - returns False on any mismatch, logging LAYER_CHECKSUM_MISMATCH () + and appending a {"type": "layer_checksum_mismatch_", ...} record + - None entries among later layers are skipped +""" + +from aorta.race.modes.fsdp import FSDPModeReproducer + + +def _make_reproducer(layer_checksums): + """Build an FSDPModeReproducer with only the attrs the method reads. + + object.__new__ skips __init__, so no CUDA buffers are allocated. + """ + r = object.__new__(FSDPModeReproducer) + r.layer_checksums = layer_checksums + r.rank = 0 + r.corruption_details = [] + # Observability counters normally set in __init__ (skipped by object.__new__). + r.layers_verified = 0 + r.layer_checksum_mismatches = 0 + return r + + +def _checksums(comm_in=10, comm_out=20, compute_in=30, compute_out=40): + return { + "comm_input": comm_in, + "comm_output": comm_out, + "compute_input": compute_in, + "compute_output": compute_out, + } + + +def test_clean_layers_pass(): + """Identical checksum dicts across all layers -> pass, no corruption recorded.""" + r = _make_reproducer([_checksums(), _checksums(), _checksums(), _checksums()]) + assert r._verify_layer_checksums(iteration=0) is True + assert r.corruption_details == [] + # Observability: a clean run must still PROVE the detector ran. + assert r.layers_verified == 3 # layers 1..3 compared against layer 0 + assert r.layer_checksum_mismatches == 0 + + +def test_compute_corruption_detected(): + """One layer with a divergent compute_output is flagged and localized to COMPUTE. + + comm_* still match, so the recorded mismatch type must be the compute key and + the offending layer index must be exposed in corruption_details. + """ + bad_layer = 2 + layers = [_checksums() for _ in range(4)] + layers[bad_layer] = _checksums(compute_out=999) # only compute_output differs + + r = _make_reproducer(layers) + assert r._verify_layer_checksums(iteration=5) is False + + assert len(r.corruption_details) == 1 + detail = r.corruption_details[0] + assert detail["type"] == "layer_checksum_mismatch_compute_output" + assert detail["layer_cmp"] == bad_layer + assert detail["layer_ref"] == 0 + # localized to compute, NOT comm + assert "comm" not in detail["type"] + assert r.layers_verified == 3 + assert r.layer_checksum_mismatches == 1 + + +def test_comm_corruption_detected(): + """One layer with a divergent comm_output is flagged and localized to COMM/NIC.""" + bad_layer = 1 + layers = [_checksums() for _ in range(3)] + layers[bad_layer] = _checksums(comm_out=777) # only comm_output differs + + r = _make_reproducer(layers) + assert r._verify_layer_checksums(iteration=9) is False + + assert len(r.corruption_details) == 1 + detail = r.corruption_details[0] + assert detail["type"] == "layer_checksum_mismatch_comm_output" + assert detail["layer_cmp"] == bad_layer + assert detail["cmp_checksum"] == 777 + # localized to comm, NOT compute + assert "compute" not in detail["type"] + + +def test_single_layer_or_empty(): + """1 layer or empty -> nothing to compare against -> pass, no false positive.""" + # Single layer: loop over range(1, 1) never runs. + single = _make_reproducer([_checksums()]) + assert single._verify_layer_checksums(iteration=0) is True + assert single.corruption_details == [] + + # Reference is None (e.g. compute disabled for layer 0) -> early True return. + none_ref = _make_reproducer([None, _checksums()]) + assert none_ref._verify_layer_checksums(iteration=0) is True + assert none_ref.corruption_details == [] diff --git a/tests/workloads/test_race_transformer_smoke.py b/tests/workloads/test_race_transformer_smoke.py new file mode 100644 index 00000000..c6e7b546 --- /dev/null +++ b/tests/workloads/test_race_transformer_smoke.py @@ -0,0 +1,125 @@ +"""CPU smoke test for the shared-weight transformer compute + checksum detector. + +Runs the REAL pieces end-to-end without a GPU or torch.distributed: + * the borrowed RepeatedTransformerBlock (same model llm_determinism uses), + * the shared-block + fixed-reference-input invariant from setup_buffers, + * the real FSDPModeReproducer._checksum and _verify_layer_checksums. + +It proves three things a green cluster run alone cannot: + 1. a real transformer block forward actually runs (not mm+gelu, not a no-op), + 2. with shared weights + same input, every layer's output is byte-identical, + 3. the per-layer checksum detector PASSES when clean and FIRES (localized to + compute) when a layer's output is corrupted. + +Run: python -m pytest tests/workloads/test_race_transformer_smoke.py -v +""" + +import pytest + +torch = pytest.importorskip("torch") + +from aorta.models import BlockConfig, RepeatedTransformerBlock +from aorta.race.modes.fsdp import FSDPModeReproducer + + +HIDDEN = 64 +NUM_LAYERS = 4 +NUM_HEADS = 4 +SEQ = 8 +BATCH = 1 +DTYPE = torch.bfloat16 + + +def _build_shared_block_and_input(): + """Mirror setup_buffers' shared-weight transformer construction (CPU).""" + cfg = BlockConfig( + hidden_size=HIDDEN, + num_heads=NUM_HEADS, + num_layers=1, + ffn_size=HIDDEN * 4, + num_experts=1, + ) + # Fixed seed -> deterministic, reproducible weights (CPU analogue of the + # fork_rng + cuda.manual_seed(0) used on device). + torch.manual_seed(0) + block = RepeatedTransformerBlock(cfg).to(DTYPE) + block.eval() + g = torch.Generator() + g.manual_seed(1) + reference_input = torch.randn(BATCH, SEQ, HIDDEN, dtype=DTYPE, generator=g) + return block, reference_input + + +def _run_layers(block, reference_input): + """Per-layer forward + 4 checksums, exactly like _forward_layer's shared path.""" + layer_checksums = [] + for _ in range(NUM_LAYERS): + comm_input = FSDPModeReproducer._checksum(reference_input) + comm_output = comm_input # no real all_gather on CPU; identical by construction + compute_input = FSDPModeReproducer._checksum(reference_input) + with torch.no_grad(): + out = block(reference_input) + compute_output = FSDPModeReproducer._checksum(out) + layer_checksums.append( + { + "comm_input": comm_input, + "comm_output": comm_output, + "compute_input": compute_input, + "compute_output": compute_output, + } + ) + return layer_checksums + + +def _verifier(layer_checksums): + """Minimal FSDPModeReproducer carrying just what _verify_layer_checksums reads.""" + r = object.__new__(FSDPModeReproducer) + r.layer_checksums = layer_checksums + r.rank = 0 + r.corruption_details = [] + r.layers_verified = 0 + r.layer_checksum_mismatches = 0 + return r + + +def test_real_transformer_block_runs_on_cpu(): + """A real RepeatedTransformerBlock forward executes and returns the right shape.""" + block, ref = _build_shared_block_and_input() + with torch.no_grad(): + out = block(ref) + assert out.shape == (BATCH, SEQ, HIDDEN) + assert out.dtype == DTYPE + # Not a trivial no-op: output differs from input. + assert FSDPModeReproducer._checksum(out) != FSDPModeReproducer._checksum(ref) + + +def test_shared_weights_make_layers_identical_and_detector_passes(): + """Clean path: shared block + same input -> identical layers -> detector PASSES.""" + block, ref = _build_shared_block_and_input() + layer_checksums = _run_layers(block, ref) + + # Invariant: every layer's compute_output checksum is identical. + outs = {c["compute_output"] for c in layer_checksums} + assert len(outs) == 1, "shared weights + same input must yield identical layer outputs" + + r = _verifier(layer_checksums) + assert r._verify_layer_checksums(iteration=0) is True + assert r.layer_checksum_mismatches == 0 + assert r.layers_verified == NUM_LAYERS - 1 # layers 1..N compared to layer 0 + + +def test_injected_compute_corruption_is_detected_and_localized(): + """Corrupt one layer's compute_output -> detector FIRES, localized to compute.""" + block, ref = _build_shared_block_and_input() + layer_checksums = _run_layers(block, ref) + + bad_layer = 2 + layer_checksums[bad_layer]["compute_output"] += 1 # flip the checksum + + r = _verifier(layer_checksums) + assert r._verify_layer_checksums(iteration=0) is False + assert r.layer_checksum_mismatches == 1 + detail = r.corruption_details[0] + assert detail["type"] == "layer_checksum_mismatch_compute_output" + assert detail["layer_cmp"] == bad_layer + assert "comm" not in detail["type"] # localized to compute, not the NIC path From fa7600056c2e31903e7afb2fb1452f123ce85b8a Mon Sep 17 00:00:00 2001 From: oyazdanb Date: Fri, 5 Jun 2026 12:06:53 -0400 Subject: [PATCH 2/6] race/fsdp: clarify this harness is NOT torch.distributed FSDP Module docstring note: mode=fsdp simulates the FSDP comm pattern with explicit all_gather/reduce_scatter; it does not use FSDP1/FSDP2. RepeatedTransformerBlock is reused as a compute kernel only, not under fully_shard. Pre-empts the "why not FSDP2 like llm_determinism?" question; real-FSDP is the deferred PR D. Co-Authored-By: Claude Opus 4 --- src/aorta/race/modes/fsdp.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/aorta/race/modes/fsdp.py b/src/aorta/race/modes/fsdp.py index 4720a032..607bdfbd 100644 --- a/src/aorta/race/modes/fsdp.py +++ b/src/aorta/race/modes/fsdp.py @@ -1,6 +1,15 @@ """ FSDP mode reproducer (Fully Sharded Data Parallel pattern). +NOTE: This *simulates* the FSDP communication pattern with EXPLICIT +``all_gather`` / ``reduce_scatter`` calls. It does NOT use ``torch.distributed`` +FSDP (neither FSDP1 ``FullyShardedDataParallel`` nor FSDP2 ``fully_shard``). +When ``compute_type=transformer``, ``RepeatedTransformerBlock`` is reused only as +a compute kernel between the explicit collectives -- it is not wrapped in +``fully_shard``. Explicit collectives are what give the clean rank-fill + +shared-input checksum invariant. Real-FSDP coverage is a separate, deferred +workload (PR D). + This mode simulates an FSDP-style workload with: - H2D transfer for batch data (single- or double-buffered via --prefetch) - Per-layer all_gather to reconstruct full parameters before compute From e72910ef3a0ca4a8eeffb23a8bd748e33cbd5d75 Mon Sep 17 00:00:00 2001 From: oyazdanb Date: Fri, 5 Jun 2026 12:17:38 -0400 Subject: [PATCH 3/6] race/fsdp: record resolved transformer shape (heads/ffn/seq/batch) num_heads and ffn_size auto-derive when 0 (model_dim//128, model_dim*4), so the config value alone doesn't record what actually ran. Capture the resolved values and surface them in the FSDP startup log (resolved_shape=...) and in WorkloadResult.metrics (eff_num_heads/eff_ffn_size/eff_seq_len/eff_batch_size), so the transformer shape a run used is provable from the result JSON. getattr in base.run keeps it safe for non-FSDP modes. Adds an auto-derive unit test. Co-Authored-By: Claude Opus 4 --- src/aorta/race/base.py | 6 ++++++ src/aorta/race/config.py | 12 ++++++++++++ src/aorta/race/modes/fsdp.py | 18 ++++++++++++++++++ src/aorta/workloads/race.py | 4 ++++ tests/workloads/test_race_transformer_smoke.py | 13 +++++++++++++ 5 files changed, 53 insertions(+) diff --git a/src/aorta/race/base.py b/src/aorta/race/base.py index 2b17db16..31604025 100644 --- a/src/aorta/race/base.py +++ b/src/aorta/race/base.py @@ -539,6 +539,12 @@ def run(self) -> ReproducerResult: avg_step_time_ms=avg_step_ms, layers_verified=self.layers_verified, layer_checksum_mismatches=self.layer_checksum_mismatches, + # Resolved transformer shape (FSDP shared-weight path only; getattr + # keeps this safe for modes that don't set these). + eff_num_heads=getattr(self, "eff_num_heads", None), + eff_ffn_size=getattr(self, "eff_ffn_size", None), + eff_seq_len=getattr(self, "eff_seq_len", None), + eff_batch_size=getattr(self, "eff_batch_size", None), ) diff --git a/src/aorta/race/config.py b/src/aorta/race/config.py index 5eee1d32..ca1be2cf 100644 --- a/src/aorta/race/config.py +++ b/src/aorta/race/config.py @@ -322,6 +322,18 @@ class ReproducerResult: layer_checksum_mismatches: int = 0 """Per-layer checksum mismatches found (0 on a clean run).""" + eff_num_heads: Optional[int] = None + """Resolved attention heads actually used (auto-derived when config=0).""" + + eff_ffn_size: Optional[int] = None + """Resolved FFN intermediate size actually used (auto-derived when config=0).""" + + eff_seq_len: Optional[int] = None + """Resolved sequence length of the reference input (transformer compute).""" + + eff_batch_size: Optional[int] = None + """Resolved batch size of the reference input (transformer compute).""" + # ============================================================================= # Race Injection Config (broader aorta system) diff --git a/src/aorta/race/modes/fsdp.py b/src/aorta/race/modes/fsdp.py index 607bdfbd..85453c7b 100644 --- a/src/aorta/race/modes/fsdp.py +++ b/src/aorta/race/modes/fsdp.py @@ -92,6 +92,14 @@ def __init__(self, config: ReproducerConfig, rank: int, world_size: int): self.reference_input: Optional[torch.Tensor] = None self.layer_checksums: List[Optional[dict]] = [] + # Effective (resolved) transformer block shape. num_heads/ffn_size are + # auto-derived when 0, so the config value alone doesn't record what + # actually ran -- store the resolved values for the startup log + metrics. + self.eff_num_heads: Optional[int] = None + self.eff_ffn_size: Optional[int] = None + self.eff_seq_len: Optional[int] = None + self.eff_batch_size: Optional[int] = None + # Real transformer block shared across all layers (shared-weight path) self.shared_block: Optional[RepeatedTransformerBlock] = None @@ -173,6 +181,11 @@ def setup_buffers(self) -> None: f"model_dim ({hidden}) must be divisible by num_heads " f"({num_heads}) for shared-weight transformer compute" ) + # Record the resolved shape (num_heads/ffn may be auto-derived). + self.eff_num_heads = num_heads + self.eff_ffn_size = ffn + self.eff_seq_len = cfg.seq_len + self.eff_batch_size = cfg.batch_size block_cfg = BlockConfig( hidden_size=hidden, num_heads=num_heads, @@ -226,6 +239,11 @@ def setup_buffers(self) -> None: f"shared_layer_weights={cfg.shared_layer_weights}, " f"transformer_block={'active' if shared_active else 'none'}, " f"layer_checksum_verify={'ON' if shared_active else 'OFF'}" + + ( + f", resolved_shape=heads:{self.eff_num_heads} ffn:{self.eff_ffn_size} " + f"seq:{self.eff_seq_len} batch:{self.eff_batch_size}" + if shared_active else "" + ) ) def _fill_patterns(self) -> None: diff --git a/src/aorta/workloads/race.py b/src/aorta/workloads/race.py index 26575a0c..3a83c1d7 100644 --- a/src/aorta/workloads/race.py +++ b/src/aorta/workloads/race.py @@ -106,6 +106,10 @@ def run(self) -> WorkloadResult: "compute_type": self._cfg.compute_type, "layers_verified": res.layers_verified, "layer_checksum_mismatches": res.layer_checksum_mismatches, + "eff_num_heads": res.eff_num_heads, + "eff_ffn_size": res.eff_ffn_size, + "eff_seq_len": res.eff_seq_len, + "eff_batch_size": res.eff_batch_size, "rank": self._rank, "world_size": self._world, }, diff --git a/tests/workloads/test_race_transformer_smoke.py b/tests/workloads/test_race_transformer_smoke.py index c6e7b546..6ec969c6 100644 --- a/tests/workloads/test_race_transformer_smoke.py +++ b/tests/workloads/test_race_transformer_smoke.py @@ -82,6 +82,19 @@ def _verifier(layer_checksums): return r +def test_num_heads_auto_derived_when_zero(): + """num_heads=0 must resolve to model_dim//128 (the recipe relies on this).""" + # mirrors fsdp.setup_buffers derivation + hidden = 1024 + cfg_num_heads = 0 + resolved = cfg_num_heads or (hidden // 128) + assert resolved == 8 + # and the block accepts it + cfg = BlockConfig(hidden_size=hidden, num_heads=resolved, num_layers=1, + ffn_size=hidden * 4, num_experts=1) + assert cfg.hidden_size % cfg.num_heads == 0 + + def test_real_transformer_block_runs_on_cpu(): """A real RepeatedTransformerBlock forward executes and returns the right shape.""" block, ref = _build_shared_block_and_input() From 785ef89270e20dafef826164208138910bdab3a8 Mon Sep 17 00:00:00 2001 From: oyazdanb Date: Fri, 5 Jun 2026 12:22:14 -0400 Subject: [PATCH 4/6] =?UTF-8?q?race/fsdp:=20address=20PR=20review=20?= =?UTF-8?q?=E2=80=94=20CPU=20RNG=20seed=20+=20stale=20GEMM=20docs?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - fsdp.py: also torch.manual_seed(0) inside fork_rng. RepeatedTransformerBlock inits params on CPU before .to("cuda"), so seeding only the CUDA RNG left the shared block's weights dependent on each rank's CPU RNG state -- ranks could diverge while the intra-rank per-layer checksum still passed (false green). Seeding CPU RNG too restores the rank-identical invariant the comment claims. - config.py / fsdp.py / recipe: update shared_layer_weights + _verify_layer_checksums docstrings and the recipe comment that still described the old GEMM (gelu(W@x)) path and the wrong function name (_verify_layer_activations); the shared path now runs a real RepeatedTransformerBlock. - remove docs/design/race-transformer-compute-support.md from the PR (it was a pre-implementation draft; the PR description carries the summary). Co-Authored-By: Claude Opus 4 --- .../race-transformer-compute-support.md | 123 ------------------ recipes/ainic-gdr-flush-sdc.yaml | 16 +-- src/aorta/race/config.py | 18 +-- src/aorta/race/modes/fsdp.py | 26 ++-- 4 files changed, 35 insertions(+), 148 deletions(-) delete mode 100644 docs/design/race-transformer-compute-support.md diff --git a/docs/design/race-transformer-compute-support.md b/docs/design/race-transformer-compute-support.md deleted file mode 100644 index a72af358..00000000 --- a/docs/design/race-transformer-compute-support.md +++ /dev/null @@ -1,123 +0,0 @@ -# Plan — Fully support `compute_type: transformer` + shared-weight checksums (PR #210) - -**Branch:** `users/oyazdanb/race-transformer-compute` (off `main`) -**Status:** DRAFT — for review before implementation -**Related:** PR #210 (`users/mycpuorg/shared-weight-transformer-checksums`), task `race__cross-rank-and-iter0-detection.md` - -## TL;DR - -PR #210's transformer + per-layer-checksum code appears **complete and correctly wired on its own branch** (config field, `setup_buffers` shared-weight path, `_forward_layer` dual dispatch, `_verify_layer_checksums` defined AND called). Yet the 2026-06-05T01-12-30 run — which *accepted* `compute_type: transformer` + `shared_layer_weights: true` (no "ignoring unknown key" warning) — showed **0.0 ms step time and emitted no checksum signal**, so we cannot confirm the path actually executed. The detector is **unobservable on success**: it logs/raises only on mismatch and writes no metric, so a clean run looks identical to a no-op. - -**This plan does NOT rewrite the compute path** (the branch already has it). It (1) resolves the did-it-run ambiguity, (2) makes the detector self-evidencing, and (3) hardens config validation so silent fallback is impossible. Work lands as a small change set that can either be merged into #210 or layered on top. - -## UPDATE (2026-06-05): PR #210's "transformer" is not a transformer — Option A - -Reading the PR branch source: `compute_type: transformer` does **`out = gelu(weight_matrix @ reference_input)`** (`fsdp.py:240-241`) — a single matmul, **no attention, no FFN, no softmax**. The "transformer" label refers to the *checksum scheme*, not the compute. So: -- **Real transformer model?** No (mm+gelu). -- **Real `torch.distributed.fsdp`?** No (simulated shards + REAL explicit `all_gather`/`reduce_scatter`). -- Both `gemm` and `transformer` compute_types are effectively a matmul. - -**Consequence for PR D / L2 pressure:** PR #210's transformer does NOT supply real-model L2/HBM pressure (the documented missing ingredient for small-scale repro). A single matmul ≠ a real block's memory traffic. So #210 is *not* the cheap step toward PR D I first assumed. - -**Option A (chosen):** make `compute_type: transformer` run a REAL transformer block by borrowing `RepeatedTransformerBlock` from `aorta.models.repeated_block` — the same model `llm_determinism` uses (real MHA + GLU FFN + LayerNorm, torch-only, no deps). This gives real L2/HBM pressure *inside* the existing race harness (explicit collectives preserved), without the full `torch.distributed.fsdp` refactor that is PR D. It is the recommended-for-PR-D model per the task file, used here as a compute kernel rather than under FSDP2. - -### Option A — implementation spec (source-grounded; refs on PR #210 branch unless noted) - -**Compute site:** `_forward_layer` shared path `fsdp.py:240-241` (fwd), `_backward_layer` `fsdp.py:279-281` (bwd). The `full_param` (all_gather output) is only *checksummed*, never fed to compute — keep that true. - -**Shape adaptation:** current compute is 2D `[dim,dim] @ [dim,dim]`. `RepeatedTransformerBlock.forward` (`repeated_block.py:108`) needs 3D `[batch, seq, hidden]`, `hidden==model_dim`, `hidden % num_heads == 0`. So `reference_input` becomes 3D `[batch, seq, model_dim]`. - -**Shared-weight invariant (the crux):** PR #210 already shares ONE weight across layers — `self.weight_matrices = [shared_w] * num_layers` (`fsdp.py:148-153`), `reference_input` seed 1. To keep "all layers byte-identical" with a real block, build **ONE** `RepeatedTransformerBlock` (not `RepeatedBlockModel` — we want a single block, no embed/LM-head) under a fixed seed and call it for every layer: -```python -# in setup_buffers, transformer+shared path -with torch.random.fork_rng(devices=["cuda"]): - torch.cuda.manual_seed(0) # identical weights on EVERY rank - self.shared_block = RepeatedTransformerBlock(block_cfg).to("cuda").to(self.dtype) -self.shared_block.eval() -g = torch.Generator(device="cuda"); g.manual_seed(1) -self.reference_input = torch.randn(batch, seq, model_dim, dtype=self.dtype, device="cuda", generator=g) -``` -In `_forward_layer` shared branch replace mm+gelu with `out = self.shared_block(self.reference_input)` under `torch.no_grad()`. The 4 checksums map unchanged: `comm_input`=param shard, `comm_output`=full_param post-all_gather, `compute_input`=`_checksum(reference_input)` (`_checksum` `fsdp.py:200` is shape-agnostic), `compute_output`=`_checksum(out)`. - -**Config mapping → `BlockConfig`:** `hidden_size <- model_dim`; build a single block (loop num_layers ourselves). **New `ReproducerConfig` fields** (`config.py` transformer section ~line 158): `num_heads: int = 16` (validate `model_dim % num_heads == 0`, else BlockConfig raises `repeated_block.py:44`), `ffn_size: int = 0` (0 ⇒ derive `4*model_dim`), `seq_len: int = 512`, `batch_size: int = 1`. **Reused:** `model_dim`, `num_layers`, `dtype`, `shared_layer_weights`, `compute_type`, `simulate_compute`, `include_backward_compute`. - -**Backward:** current is manual `mm(W.T, grad)` (`fsdp.py:279-281`), NOT autograd; `reference_input.requires_grad=False`. **Keep backward as a timing-only proxy** — when `include_backward_compute`, call `self.shared_block(self.reference_input)` a second time under `no_grad`. Do NOT add autograd (breaks no-grad determinism, balloons memory, scope creep). - -**Stays untouched (exercises AINIC):** all_gather (`fsdp.py:228,272`), reduce_scatter (`:284`), `_fill_patterns`, `_verify_all_gather` (`:439`), `_verify_reduce_scatter` (`:473`) — they read `full_param`/`grad_shard`, never `reference_input`. Only the compute kernel changes. - -**Concrete edits:** -- `config.py` ~L158: add `num_heads`, `ffn_size`, `seq_len`, `batch_size`. -- `fsdp.py:30`: `from aorta.models import BlockConfig, RepeatedTransformerBlock`. -- `fsdp.py:80-82` `__init__`: `self.shared_block = None`. -- `fsdp.py:144-159` `setup_buffers`: block construction + 3D reference_input (transformer+shared path); leave non-shared 2D path as-is. -- `fsdp.py:236-251` `_forward_layer` shared branch: real block forward under `no_grad`. -- `fsdp.py:277-281` `_backward_layer`: block-forward proxy on shared path. - -**Pitfalls:** `RepeatedTransformerBlock` defaults float32 → `.to(self.dtype)` mandatory (race=bf16). LayerNorm/softmax upcast (`repeated_block.py:119`) is deterministic — fine. Memory: real block activations `[batch,seq,hidden]` + attention `[batch,heads,seq,seq]` ≫ 2D matmul → default `batch_size=1, seq_len=512`; block built once (shared) so param memory ≈ one layer regardless of `num_layers`. - -**Preserves corruption localization: YES** — `_verify_layer_checksums` (`fsdp.py:419-436`) is checksum-key-agnostic: comm mismatch ⇒ RCCL/NIC, compute mismatch ⇒ GPU compute, verbatim. Rank-fill checks on `full_param`/`grad_shard` are untouched (LOW risk). - -**One real setup risk:** if ranks build the block with diverging RNG, `compute_output` differs *across ranks* but stays identical *across layers within a rank* → intra-rank `_verify_layer_checksums` still PASSES, masking a false setup bug. Mitigate: seed global CUDA RNG identically on all ranks before block construction (the `fork_rng` above), and optionally add a one-time cross-rank `compute_output` all-reduce equality assert at setup. - -**Effort:** ~medium / half-day. ~4 edited functions + 4 config fields + 1 import. No new files. The earlier Steps 1–3 (observability, config validation, detector unit test) still apply on top. - -## Step 0 — Resolve the contradiction FIRST (no code yet) - -Before writing anything, confirm what actually ran on the cluster. On the cluster: -``` -cd /it-share/oyazdanb/aorta && git branch --show-current -python - <<'PY' -import inspect, aorta.race.modes.fsdp as f -src = inspect.getsource(f) -for tok in ("_verify_layer_checksums","_checksum","reference_input","use_shared","shared_w"): - print(tok, tok in src) -PY -``` -- All `True` → branch code IS installed; the path is *present* but unobservable → go to Steps 1–3. -- Any `False` → a stale/main build is installed; `pip install -e` didn't take → reinstall, rerun, re-check. (This alone may explain the 0.0 ms / no-metric run.) - -Also: 0.0 ms step is suspicious even for GEMM. Confirm `simulate_compute` is actually doing work — a real 24-layer transformer fwd+bwd at model_dim=1024 cannot be 0.0 ms. If it is 0.0 ms with the branch installed, compute is being skipped or mistimed — that's a real bug to find in `_forward_layer`/timing, not just observability. - -## Step 1 — Make the detector self-evidencing (observability) - -Root problem: a green result does not prove the checksum verifier ran. Fix so it is impossible to miss. - -- **Startup log (once per trial):** in `setup_buffers`/`run` when `compute_type=="transformer"` and `shared_layer_weights`, emit e.g. - `log.info("race: compute=transformer shared_layer_weights=ON layers=%d dim=%d; per-layer checksum verify ENABLED", num_layers, model_dim)`. - And in the GEMM branch: `log.info("race: compute=gemm")`. Now fallback is one grep away. -- **Result metric (always, not just on failure):** add to `WorkloadResult.metrics` (mapped in `workloads/race.py`): - - `layers_verified` (int, per step or total), - - `layer_checksum_mismatches` (int, 0 on clean), - - `compute_type` (echo the effective value). - A clean run then shows `layer_checksum_mismatches: 0, layers_verified: >0` — provably ran. Today `metrics` is only `{avg_step_time_ms, mode, rank, world_size}`. -- **Per-step debug counter** behind `log_interval` so long runs show progress. - -## Step 2 — Harden config so silent fallback is impossible - -- **Validate `compute_type`** against `{"gemm","transformer"}` in `workloads/race.py` (mirror the existing `_VALID_DTYPES` pattern). Today any string is accepted, so a typo (`transfomer`) silently runs GEMM with no warning. -- **Warn on inert combo:** if `shared_layer_weights=True` but `compute_type!="transformer"`, log a WARNING — currently it silently no-ops via the `use_shared` AND-gate. -- Keep accepting the keys (they already round-trip on the branch) — this is validation, not new surface. - -## Step 3 — Confirm the checksum actually catches corruption (test the detector) - -A detector that never fires on a clean cluster is unproven. Add a unit/integration check: -- **Unit:** forge a layer activation buffer so one layer's checksum differs; assert `_verify_layer_checksums` reports a mismatch with the right layer index. (No GPU/dist needed.) -- **Negative:** identical layers → 0 mismatches. -- This proves the second signal works regardless of whether the AINIC bug reproduces. - -## Files - -- `src/aorta/race/modes/fsdp.py` — startup log; mismatch counter; ensure `_verify_layer_checksums` returns counts (verified, mismatches) rather than only logging. -- `src/aorta/workloads/race.py` — `_VALID_COMPUTE_TYPES` validation; warn on inert `shared_layer_weights`; map new metrics into `WorkloadResult.metrics`. -- `tests/workloads/test_race.py` (or race-mode test) — checksum-catches-forged-corruption + clean-pass + config-validation cases. -- `recipes/ainic-gdr-flush-sdc.yaml` — already correct on #210 (dtype `bfloat16`, transformer, shared_layer_weights). No change beyond confirming. - -## Relationship to PR #210 and PR D - -- **PR #210:** this is *complementary* — it makes #210's detector observable, validated, and tested. Decide at review: fold these commits into #210, or land as a follow-up on this branch. -- **PR D (real torch.distributed.fsdp):** **deferred — do NOT start.** Per the task file, the 5-node no-repro triggers **PR C first** (cold-QP/iter-0), and PR D only if real-model L2 pressure is *confirmed* the missing ingredient. PR #210's transformer compute is the cheap way to raise model-realistic L2/HBM pressure *inside* the existing harness — test that hypothesis before any PR D refactor. Also: PR D would **not** close the iter-0/cold-QP gap (both reuse the launcher's warm PG); only PR C does. - -## Open questions for review -1. Step 0 result: is the branch actually installed on the cluster (path present) or is the 0.0 ms / no-metric run a stale-build artifact? This decides whether Steps 1–3 are the whole job or whether there's also a real compute-skip bug. -2. Fold into #210 vs land as follow-up branch? -3. Does `_verify_layer_checksums` currently return counts, or only log? (Determines how much of Step 1's metric plumbing is new.) diff --git a/recipes/ainic-gdr-flush-sdc.yaml b/recipes/ainic-gdr-flush-sdc.yaml index 80b77a4a..e4a10e79 100644 --- a/recipes/ainic-gdr-flush-sdc.yaml +++ b/recipes/ainic-gdr-flush-sdc.yaml @@ -96,14 +96,14 @@ workload_config: # model_dim=1024, num_layers=24: attention + FFN per layer, forward + backward, # approximates ~500ms/step on MI355X. # - # shared_layer_weights=true: all 24 layers share one weight matrix (fixed seed=0) - # and each layer independently receives the same fixed reference input (seed=1) - # rather than chaining activations layer-to-layer. This makes every layer's - # forward output analytically identical so _verify_layer_activations() can serve - # as a second independent corruption signal alongside the collective-buffer - # pattern check: a mismatch across layers means an all_gather corruption on that - # layer propagated through GEMM compute -- something the rank-fill pattern check - # on full_param alone cannot catch. + # shared_layer_weights=true: all 24 layers run one shared transformer block + # (fixed seed=0) on the same fixed reference input (seed=1) rather than chaining + # activations layer-to-layer. This makes every layer's forward output + # analytically identical so _verify_layer_checksums() can serve as a second + # independent corruption signal alongside the collective-buffer pattern check: + # a per-layer compute_output checksum mismatch means that layer's compute path + # was corrupted -- something the rank-fill pattern check on full_param alone + # cannot catch. simulate_compute: true compute_type: transformer model_dim: 1024 diff --git a/src/aorta/race/config.py b/src/aorta/race/config.py index ca1be2cf..8b14a7fc 100644 --- a/src/aorta/race/config.py +++ b/src/aorta/race/config.py @@ -181,14 +181,16 @@ class ReproducerConfig: shared_layer_weights: bool = False """ - Use a single shared weight matrix for all layers (transformer compute type). - - When True, all num_layers layers use the same weight matrix initialized with - a fixed seed, and each layer receives the same fixed reference input rather - than chaining activations. This makes per-layer forward outputs analytically - identical so that cross-layer activation comparison can serve as a secondary - corruption signal: any all_gather corruption that propagates through GEMM - compute will produce a mismatch between layer outputs. + Use a single shared transformer block for all layers (transformer compute type). + + When True, all num_layers layers run the SAME RepeatedTransformerBlock + (one block built with a fixed seed) on the SAME fixed reference input, rather + than chaining activations. This makes every layer's forward output + analytically identical, so cross-layer comparison of the per-kernel checksums + becomes a secondary corruption signal: if a layer's compute_output checksum + diverges from the others, that layer's compute path was corrupted. (The + collective path is checked separately via the comm_input/comm_output + checksums on each layer's all_gather.) Only meaningful when compute_type == 'transformer'. """ diff --git a/src/aorta/race/modes/fsdp.py b/src/aorta/race/modes/fsdp.py index 85453c7b..13e835b6 100644 --- a/src/aorta/race/modes/fsdp.py +++ b/src/aorta/race/modes/fsdp.py @@ -195,9 +195,15 @@ def setup_buffers(self) -> None: seq_len=cfg.seq_len, vocab_size=16, # embed is unused on this path; keep tiny ) - # fork_rng so seeding the block's init RNG to a fixed value gives - # bit-identical weights on every rank without perturbing global RNG. + # fork_rng so we can fix the seed without perturbing global RNG. + # RepeatedTransformerBlock initializes its params on CPU (nn.Linear + # / LayerNorm use the CPU RNG) BEFORE .to("cuda"), so we must seed + # the CPU RNG too -- seeding only CUDA would leave weights dependent + # on each rank's CPU RNG state and break the rank-identical invariant + # (every layer would still match within a rank, so the per-layer + # checksum would falsely pass while ranks silently diverged). with torch.random.fork_rng(devices=["cuda"]): + torch.manual_seed(0) torch.cuda.manual_seed(0) self.shared_block = ( RepeatedTransformerBlock(block_cfg).to("cuda").to(self.dtype) @@ -468,18 +474,20 @@ def _verify_layer_checksums(self, iteration: int) -> bool: """ Verify that per-kernel int16 checksums are identical across all layers. - With shared weights and a fixed reference input every layer runs the - same comm kernel (all_gather of rank-filled shard) and the same compute - kernel (GEMM + GELU with shared W and fixed reference_input). Both the - input and output of each kernel are checksummed via reinterpret-cast to - int16 → int64 sum, so every bit contributes with zero information loss. + With a shared transformer block and a fixed reference input every layer + runs the same comm kernel (all_gather of rank-filled shard) and the same + compute kernel (the shared RepeatedTransformerBlock on reference_input). + Both the input and output of each kernel are checksummed via + reinterpret-cast to int16 → int64 sum, so every bit contributes with zero + information loss. Four checksums per layer: comm_input -- param shard before all_gather (should be identical: every shard is filled with float(rank)) comm_output -- full_param after all_gather - compute_input -- reference_input fed to GEMM (constant across layers) - compute_output-- activation after GELU + compute_input -- reference_input fed to the transformer block + (constant across layers) + compute_output-- transformer block output If comm_output diverges but comm_input matches, corruption is in the collective (RCCL / NIC path). If compute_output diverges but From 9a8b16e67d751a63e63ad4ae35127cd7c1fbd355 Mon Sep 17 00:00:00 2001 From: oyazdanb Date: Fri, 5 Jun 2026 12:54:33 -0400 Subject: [PATCH 5/6] race/fsdp: address 2nd PR review round (counter, validation, buffers, doc) - layer_checksum_mismatches now counts once per CORRUPTED LAYER, not once per checksum key. A single bad layer could previously inflate it up to 4x (comm_in/comm_out/compute_in/compute_out), contradicting the per-layer docstring. corruption_details still records every key for localization. - Validate compute_type in ReproducerConfig.__post_init__ so EVERY entry point is covered (the aorta.race CLI / direct construction), not just the RaceWorkload adapter -- a typo like "transfomer" now raises instead of a silent GEMM false-green. - Skip the unused dim x dim activation/grad_buffer allocations on the shared-weight transformer path (forward sets activation to the block output; backward re-runs the block) -- avoids needless GPU memory / OOM risk at large model_dim. GEMM/chained path keeps them. - Fix the _forward_layer docstring that still said "GEMM + GELU" on the shared path (now runs RepeatedTransformerBlock). Tests: per-layer-count (4 keys bad -> counter==1), direct-construction compute_type rejection. 18 race tests pass. Co-Authored-By: Claude Opus 4 --- src/aorta/race/config.py | 12 +++++++ src/aorta/race/modes/fsdp.py | 48 ++++++++++++++++---------- tests/workloads/test_race.py | 8 +++++ tests/workloads/test_race_checksums.py | 19 ++++++++++ 4 files changed, 68 insertions(+), 19 deletions(-) diff --git a/src/aorta/race/config.py b/src/aorta/race/config.py index 8b14a7fc..616b532d 100644 --- a/src/aorta/race/config.py +++ b/src/aorta/race/config.py @@ -291,6 +291,18 @@ class ReproducerConfig: - 4+: Full parallelism (exposes timing-sensitive bugs) """ + def __post_init__(self) -> None: + # Validate here (not only in the RaceWorkload adapter) so EVERY entry + # point is covered -- the aorta.race CLI and any direct reproducer + # construction. A typo like "transfomer" must error, not silently fall + # back to the GEMM path (false green). + valid_compute = {"gemm", "transformer"} + if self.compute_type not in valid_compute: + raise ValueError( + f"compute_type must be one of {sorted(valid_compute)}, " + f"got {self.compute_type!r}" + ) + @dataclass class ReproducerResult: diff --git a/src/aorta/race/modes/fsdp.py b/src/aorta/race/modes/fsdp.py index 13e835b6..67ca9dd6 100644 --- a/src/aorta/race/modes/fsdp.py +++ b/src/aorta/race/modes/fsdp.py @@ -219,19 +219,22 @@ def setup_buffers(self) -> None: ) self.weight_matrices = [] self.layer_checksums = [None] * self.num_layers + # activation/grad_buffer below are unused on the shared path + # (forward sets activation to the block output; backward re-runs + # the block), so skip those dim x dim allocations. else: self.weight_matrices = [ torch.randn(dim, dim, dtype=self.dtype, device="cuda") for _ in range(self.num_layers) ] - self.activation = torch.randn( - dim, dim, - dtype=self.dtype, device="cuda", - ) - self.grad_buffer = torch.randn( - dim, dim, - dtype=self.dtype, device="cuda", - ) + self.activation = torch.randn( + dim, dim, + dtype=self.dtype, device="cuda", + ) + self.grad_buffer = torch.randn( + dim, dim, + dtype=self.dtype, device="cuda", + ) # Startup line names the active compute path so a silent fallback # (e.g. transformer requested but GEMM ran) is greppable in logs. @@ -278,16 +281,18 @@ def _forward_layer(self, layer_idx: int) -> None: Forward pass for a single FSDP layer. 1. all_gather: reconstruct full parameter from shards across ranks - 2. GEMM: compute with full parameter (if enabled) - - Shared-weight path: every layer receives the same fixed reference_input - so that outputs are analytically identical. Input/output checksums are - recorded for both the comm kernel (all_gather) and the compute kernel - (GEMM + GELU) so _verify_layer_checksums() can pinpoint whether - corruption entered during communication or compute. - - Chained path (default): layer 0 seeds activation from batch_gpu (H2D race - opportunity) and each subsequent layer receives the previous layer's output. + 2. compute with the reconstructed parameter (if enabled) + + Shared-weight transformer path: every layer runs the same shared + RepeatedTransformerBlock on the same fixed reference_input, so outputs are + analytically identical. Input/output checksums are recorded for both the + comm kernel (all_gather) and the compute kernel (the transformer block) so + _verify_layer_checksums() can pinpoint whether corruption entered during + communication or compute. + + Chained path (default, GEMM): layer 0 seeds activation from batch_gpu (H2D + race opportunity) and each subsequent layer receives the previous layer's + output through a GEMM + GELU. """ use_shared = ( self.config.shared_layer_weights @@ -505,6 +510,7 @@ def _verify_layer_checksums(self, iteration: int) -> bool: # Count every cross-layer comparison so a clean (green) run still # proves the detector ran: layers_verified > 0. self.layers_verified += 1 + layer_has_mismatch = False for key in ("comm_input", "comm_output", "compute_input", "compute_output"): if cmp[key] != ref[key]: log.error( @@ -521,8 +527,12 @@ def _verify_layer_checksums(self, iteration: int) -> bool: "ref_checksum": ref[key], "cmp_checksum": cmp[key], }) - self.layer_checksum_mismatches += 1 + layer_has_mismatch = True all_correct = False + # Count once per CORRUPTED LAYER, not once per key -- a single bad + # layer must not inflate the metric up to 4x (one per checksum key). + if layer_has_mismatch: + self.layer_checksum_mismatches += 1 return all_correct diff --git a/tests/workloads/test_race.py b/tests/workloads/test_race.py index 64f0b0f6..251e2ba9 100644 --- a/tests/workloads/test_race.py +++ b/tests/workloads/test_race.py @@ -92,6 +92,14 @@ def test_race_config_from_dict_rejects_bad_compute_type(): wl._race_config_from_dict({"compute_type": "transfomer"}) +def test_reproducer_config_rejects_bad_compute_type_directly(): + """Validation lives in ReproducerConfig.__post_init__, so even direct + construction (bypassing the RaceWorkload adapter, e.g. the aorta.race CLI) + rejects a typo instead of silently running GEMM (false green).""" + with pytest.raises(ValueError, match="compute_type must be one of"): + ReproducerConfig(compute_type="transfomer") + + def test_race_config_warns_shared_weights_without_transformer(caplog): wl = RaceWorkload({}) with caplog.at_level("WARNING"): diff --git a/tests/workloads/test_race_checksums.py b/tests/workloads/test_race_checksums.py index e807141c..0242369c 100644 --- a/tests/workloads/test_race_checksums.py +++ b/tests/workloads/test_race_checksums.py @@ -95,6 +95,25 @@ def test_comm_corruption_detected(): assert "compute" not in detail["type"] +def test_mismatch_counter_is_per_layer_not_per_key(): + """A single corrupted layer with MULTIPLE bad keys counts ONCE, not 4x. + + layer_checksum_mismatches is a per-layer counter; a layer that diverges on + several checksum keys must not inflate it (was previously +1 per key). + """ + bad_layer = 1 + layers = [_checksums() for _ in range(3)] + # All four keys differ on the one bad layer. + layers[bad_layer] = _checksums(comm_in=1, comm_out=2, compute_in=3, compute_out=4) + + r = _make_reproducer(layers) + assert r._verify_layer_checksums(iteration=0) is False + # corruption_details still records each key (full localization detail)... + assert len(r.corruption_details) == 4 + # ...but the per-layer metric counts the layer ONCE. + assert r.layer_checksum_mismatches == 1 + + def test_single_layer_or_empty(): """1 layer or empty -> nothing to compare against -> pass, no false positive.""" # Single layer: loop over range(1, 1) never runs. From 577b73cce88b9ba804d9b96fd04fc0a2be67f7c5 Mon Sep 17 00:00:00 2001 From: oyazdanb Date: Fri, 5 Jun 2026 14:20:06 -0400 Subject: [PATCH 6/6] race/fsdp: address 3rd PR review round (checksum dtype, fallback, registry, test) - _checksum is now element-size aware: 2-byte dtypes -> int16, 4-byte (fp32) -> int32, etc. Previously hard-coded torch.int16, which would crash/miscount on float32 (an allowed dtype). bf16 path unchanged. - compute_type='transformer' with shared_layer_weights=False now warns loudly that it runs the GEMM path (only the shared transformer path is implemented) -- no more silent transformer->GEMM fallback. - ReproducerConfig.__post_init__ validates compute_type against the pluggable COMPUTE_REGISTRY (single source of truth; "transformer" is already registered) instead of a hard-coded set, so custom register_compute() backends stay valid. Lazy import avoids a circular import (compute.py only TYPE_CHECKING-imports config). - Smoke test: comm_input/comm_output checksums now use distinct rank-fill shard / all_gather-result buffers (mirroring real _forward_layer) instead of reference_input, so the test actually exercises the comm path. Added a multi-dtype _checksum test (bf16/fp16/fp32). 19 race tests pass. Co-Authored-By: Claude Opus 4 --- src/aorta/race/config.py | 7 ++++- src/aorta/race/modes/fsdp.py | 29 +++++++++++++----- .../workloads/test_race_transformer_smoke.py | 30 ++++++++++++++++--- 3 files changed, 54 insertions(+), 12 deletions(-) diff --git a/src/aorta/race/config.py b/src/aorta/race/config.py index 616b532d..b75a1adc 100644 --- a/src/aorta/race/config.py +++ b/src/aorta/race/config.py @@ -296,7 +296,12 @@ def __post_init__(self) -> None: # point is covered -- the aorta.race CLI and any direct reproducer # construction. A typo like "transfomer" must error, not silently fall # back to the GEMM path (false green). - valid_compute = {"gemm", "transformer"} + # + # Use the pluggable COMPUTE_REGISTRY as the single source of truth so + # custom backends registered via register_compute() remain valid (don't + # hard-code the list). Imported lazily to avoid a circular import. + from .compute import COMPUTE_REGISTRY + valid_compute = set(COMPUTE_REGISTRY) if self.compute_type not in valid_compute: raise ValueError( f"compute_type must be one of {sorted(valid_compute)}, " diff --git a/src/aorta/race/modes/fsdp.py b/src/aorta/race/modes/fsdp.py index 67ca9dd6..bf247561 100644 --- a/src/aorta/race/modes/fsdp.py +++ b/src/aorta/race/modes/fsdp.py @@ -166,6 +166,15 @@ def setup_buffers(self) -> None: use_shared = ( cfg.shared_layer_weights and cfg.compute_type == "transformer" ) + if cfg.compute_type == "transformer" and not cfg.shared_layer_weights: + # Only the shared-weight transformer path is implemented; without + # shared weights we fall back to GEMM. Warn loudly so this is not + # a silent transformer->GEMM fallback (the thing this PR fixes). + log.warning( + "race: compute_type='transformer' but shared_layer_weights=False " + "-- the non-shared transformer path is not implemented; running " + "the GEMM compute path instead." + ) if use_shared: # All layers share ONE real transformer block with deterministic, # rank-identical weights so block(reference_input) is analytically @@ -267,14 +276,20 @@ def _fill_patterns(self) -> None: @staticmethod def _checksum(tensor: torch.Tensor) -> int: """ - Bitwise checksum: reinterpret-cast to int16 and sum. - - bf16 (or any 16-bit dtype) is viewed as int16 so every bit pattern - contributes to the checksum with zero information loss -- no float - rounding, no abs(), and NaN / denorm bit patterns are included. - Accumulation is done in int64 to avoid overflow. + Bitwise checksum: reinterpret-cast to an int of the SAME element size + and sum. + + Every bit pattern contributes to the checksum with zero information loss + -- no float rounding, no abs(), and NaN / denorm bit patterns are + included. The int view must match the dtype's byte width: 2-byte dtypes + (bf16/fp16) -> int16, 4-byte (fp32) -> int32, 1-byte -> int8. Accumulation + is done in int64 to avoid overflow. """ - return tensor.view(torch.int16).to(torch.int64).sum().item() + itemsize = tensor.element_size() + int_view = {1: torch.int8, 2: torch.int16, 4: torch.int32, 8: torch.int64}.get(itemsize) + if int_view is None: + raise ValueError(f"_checksum: unsupported element size {itemsize} bytes") + return tensor.view(int_view).to(torch.int64).sum().item() def _forward_layer(self, layer_idx: int) -> None: """ diff --git a/tests/workloads/test_race_transformer_smoke.py b/tests/workloads/test_race_transformer_smoke.py index 6ec969c6..a1f70309 100644 --- a/tests/workloads/test_race_transformer_smoke.py +++ b/tests/workloads/test_race_transformer_smoke.py @@ -50,12 +50,20 @@ def _build_shared_block_and_input(): return block, reference_input -def _run_layers(block, reference_input): - """Per-layer forward + 4 checksums, exactly like _forward_layer's shared path.""" +def _run_layers(block, reference_input, world_size=4, rank=0): + """Per-layer forward + 4 checksums, mirroring _forward_layer's shared path. + + comm checksums use distinct buffers like the real code (NOT reference_input): + comm_input = rank-filled param shard, comm_output = the all_gather result + (here built deterministically as the concatenation of every rank's shard). + """ + shard = torch.full((SEQ * HIDDEN,), float(rank), dtype=DTYPE) # this rank's shard + full_param = torch.cat([torch.full((SEQ * HIDDEN,), float(r), dtype=DTYPE) + for r in range(world_size)]) # all_gather result layer_checksums = [] for _ in range(NUM_LAYERS): - comm_input = FSDPModeReproducer._checksum(reference_input) - comm_output = comm_input # no real all_gather on CPU; identical by construction + comm_input = FSDPModeReproducer._checksum(shard) + comm_output = FSDPModeReproducer._checksum(full_param) compute_input = FSDPModeReproducer._checksum(reference_input) with torch.no_grad(): out = block(reference_input) @@ -95,6 +103,20 @@ def test_num_heads_auto_derived_when_zero(): assert cfg.hidden_size % cfg.num_heads == 0 +def test_checksum_handles_multiple_dtypes(): + """_checksum must work for 2-byte AND 4-byte dtypes (fp32 is an allowed dtype). + + Guards against the int16-only view crashing on float32. + """ + for dt in (torch.bfloat16, torch.float16, torch.float32): + t = torch.randn(4, 4, dtype=dt) + # returns an int and is deterministic for identical data + c1 = FSDPModeReproducer._checksum(t) + c2 = FSDPModeReproducer._checksum(t.clone()) + assert isinstance(c1, int) + assert c1 == c2 + + def test_real_transformer_block_runs_on_cpu(): """A real RepeatedTransformerBlock forward executes and returns the right shape.""" block, ref = _build_shared_block_and_input()