race/fsdp: real transformer compute + observable per-layer checksum detector

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:
   # ------------------------------------------------------------------

src/aorta/race/base.py

@@ -78,15 +78,24 @@ 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()
 
     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),
         )
 
 

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)