diff --git a/recipes/ainic-gdr-flush-sdc.yaml b/recipes/ainic-gdr-flush-sdc.yaml index 5d2d4fb6..e4a10e79 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 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: 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..31604025 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,14 @@ 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, + # 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 4e47d0e4..b75a1adc 100644 --- a/src/aorta/race/config.py +++ b/src/aorta/race/config.py @@ -156,9 +156,45 @@ 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 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'. + """ + # ========================================================================= # Optimizer (used by modes that support it, e.g., DDP) # ========================================================================= @@ -255,6 +291,23 @@ 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). + # + # 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)}, " + f"got {self.compute_type!r}" + ) + @dataclass class ReproducerResult: @@ -281,6 +334,25 @@ 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).""" + + 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 eda477b8..bf247561 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 @@ -30,6 +39,8 @@ import torch import torch.distributed as dist +from aorta.models import BlockConfig, RepeatedTransformerBlock + from ..base import BaseReproducer from ..config import ReproducerConfig @@ -77,6 +88,21 @@ 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]] = [] + + # 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 + def _setup_compute(self) -> None: """ Override base compute setup -- FSDP manages its own per-layer compute. @@ -90,7 +116,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,27 +163,105 @@ 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( + 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 + # 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" + ) + # 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, + 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 + ) + # 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) + ) + 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 + # 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", ) - 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", - ) + # 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'}" + + ( + 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: @@ -166,30 +273,90 @@ 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 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. + """ + 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: """ 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) + 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. """ - # 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 +371,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 +463,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 +479,76 @@ 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 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 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 + 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 + layer_has_mismatch = False + 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], + }) + 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 def _verify_all_gather(self) -> bool: diff --git a/src/aorta/workloads/race.py b/src/aorta/workloads/race.py index ace41e9f..3a83c1d7 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,13 @@ 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, + "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.py b/tests/workloads/test_race.py index 7a38b276..251e2ba9 100644 --- a/tests/workloads/test_race.py +++ b/tests/workloads/test_race.py @@ -85,6 +85,31 @@ 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_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"): + 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..0242369c --- /dev/null +++ b/tests/workloads/test_race_checksums.py @@ -0,0 +1,127 @@ +"""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_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. + 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..a1f70309 --- /dev/null +++ b/tests/workloads/test_race_transformer_smoke.py @@ -0,0 +1,160 @@ +"""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, 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(shard) + comm_output = FSDPModeReproducer._checksum(full_param) + 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_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_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() + 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