Upload GITHUB_ISSUE.md

GITHUB_ISSUE.md

@@ -0,0 +1,169 @@
+# Bug Report: 4 Critical Issues Found in NAS / GPU Path (branch `do-not-merge-gpu-tests-from-23`)
+
+**Reporter:** ML Intern  
+**Branch audited:** `do-not-merge-gpu-tests-from-23`  
+**Fix branch:** `fix-gpu-tests-from-23-v2`  
+**Severity:** High — NAS trials have been wasted on dead parameters; GPU training fails with OOM / NaN / compile crashes.
+
+---
+
+## Summary
+
+During a neural-architecture-search audit of `nas_helixlm.py` v2.3 on the `do-not-merge-gpu-tests-from-23` branch, four interrelated bugs were found in `graph.py`, `mamba2.py`, `nodes.py`, `trainer.py`, and `nas_helixlm.py`.  All four have been fixed in branch `fix-gpu-tests-from-23-v2` (see commit `d6619f4`).
+
+| # | Bug | File(s) | Impact |
+|---|-----|---------|--------|
+| 1 | `nodes_per_column` is ignored by `_build_node_spec()` | `graph.py` | NAS permutes a dead parameter; topology is not reproducible by `seed` alone |
+| 2 | Mamba2 / SSM scan builds a 256-step autograd chain → OOM | `mamba2.py`, `nodes.py` | GPU OOM on SSM configs; ~1.1 GB saved tensors per mamba2 node |
+| 3 | `fp16` AMP forced for `seq_len > 128` → NaN | `nas_helixlm.py`, `trainer.py` | Immediate NaN on small models where there is no memory pressure |
+| 4 | `torch.compile` crashes on custom Python loops | `mamba2.py`, `nodes.py`, `nas_helixlm.py` | Inductor graph-breaks / invalid kernels on GPU; compile silently skipped for SSM/Titans configs |
+
+---
+
+## 1. `nodes_per_column` is dead — `_build_node_spec()` ignores it
+
+### Evidence
+
+`config.py` docstrings and presets (lines 22-28, 311-380) advertise `nodes_per_column` tuples like `(2,2)`, `(2,3,2)`, `(3,4,4,3)`.  Validation logic (lines 249-256) even pads/truncates the tuple to match `n_columns`.  
+
+**However**, `graph.py` lines 88-148 hardcode every column to:
+
+```python
+column = [
+    ("linear_attn" | "full_attn", {…}),
+    ("swiglu", {…}),
+    ("mamba2" | "ssm", {…})       # if use_ssm
+    ("titans", {…})                # if use_titans and ci==0
+    ("gate", {…})                  # always appended
+]
+```
+
+`nodes_per_column` is **never read** after validation.
+
+### Fix
+
+Wire `nodes_per_column` into `_build_node_spec()` by repeating the `[attention, swiglu]` base pattern until the target count is reached. Optional SSM/Titans nodes consume one slot each.  A gate is appended when there are multiple compute nodes or when `ci > 0`.
+
+```python
+for ci in range(cfg.n_columns):
+    target = cfg.nodes_per_column[ci]  # e.g. 3
+    # Build base: attn + swiglu
+    # Insert optional SSM/Titans if room
+    # Repeat [attn, swiglu] to fill remaining slots
+    # Append gate for aggregation
+```
+
+The RNG seed already controls lateral/vertical wiring; with this fix the **node count** is also deterministic and reproducible.
+
+---
+
+## 2. OOM from mamba2 scan's 256-step autograd chain
+
+### Evidence
+
+`_ssd_chunked_scan` (and `_ssm_chunked_scan` in `nodes.py`) keeps every intermediate `h` tensor for backward:
+
+```python
+h = A_c[:, t] * h + B_c[:, t] * x_c[:, t].unsqueeze(-1)
+# h is (B, d_inner, d_state) — kept for ALL timesteps
+```
+
+For `B=32, d_inner=768, d_state=64`, each mamba2 node alone stores **~1.1 GB** of saved tensors.  Multiple columns and loops multiply this.
+
+### Fix
+
+Wrap each chunk's inner loop in `torch.utils.checkpoint`:
+
+```python
+def _chunk_scan(h_in, A_c_in, ...):
+    h = h_in
+    for t in range(chunk_size):
+        h = ...
+        ys_c.append(y_t)
+    return h, torch.stack(ys_c, dim=1)
+
+h, ys_chunk = torch.utils.checkpoint.checkpoint(
+    _chunk_scan, h, A_c, B_c, x_c, C_c,
+    use_reentrant=False,
+)
+```
+
+Only the **chunk boundary** `h` states are materialised for backward.  Trade ~10–20 % extra compute for an order-of-magnitude memory reduction.
+
+---
+
+## 3. fp16 AMP causes NaN on `seq_len ≥ 256`
+
+### Evidence
+
+`nas_helixlm.py` line 395 forced `dtype_str = "float32"` everywhere because fp16 caused immediate NaN on `d ≥ 256` with `LR=3e-3`.  The root cause is not fp16 universally, but this architecture's SSM scan, Titans memory updates, and `ELU+1.0` feature maps — all of which underflow in fp16's narrow dynamic range.
+
+### Fix
+
+Use **bfloat16** instead of fp16 on GPU.  bf16 shares the same 8-bit exponent range as fp32, so it does not underflow on the scan or memory updates.  It is natively supported on Ampere+ (A100, H100, L4) and is typically as fast as fp16.
+
+```python
+if torch.cuda.is_available():
+    dtype_str = "bfloat16"
+    use_amp = True
+else:
+    dtype_str = "float32"
+    use_amp = False
+```
+
+The `Trainer` was also updated to skip `GradScaler` when the AMP dtype is bf16 (no loss-scaling needed) and to pass `torch.bfloat16` to `torch.amp.autocast`.
+
+---
+
+## 4. `torch.compile` breaks on custom Python loops
+
+### Evidence
+
+The inductor backend cannot compile the `for t in range(chunk_size)` loops with in-place `h` mutations inside `_ssd_chunked_scan` and `TitansMemoryNode.forward`.  The old workaround in `nas_helixlm.py` simply **skipped compilation entirely** for any SSM/Titans config:
+
+```python
+if use_ssm or use_titans:
+    return model, False, "skipped: SSM/Titans autograd not compile-safe"
+```
+
+### Fix
+
+Decorate the loop functions with `@torch.compiler.disable`.  This tells the inductor backend to treat them as opaque ops — the rest of the model (embeddings, linear projections, attention, SwiGLU) still gets compiled.
+
+Decorated functions:
+- `mamba2._ssd_chunked_scan`
+- `nodes._ssm_chunked_scan`
+- `nodes.TitansMemoryNode.forward`
+
+`try_compile_model` in `nas_helixlm.py` now removes the SSM/Titans skip and attempts compilation for all configs.
+
+---
+
+## Reproducing the Baseline
+
+The sacred CPU baseline (1000 samples, seq_len=96, 14 epochs) should yield:
+- Train perplexity **~ 23**
+- Val perplexity **~ 85–86**
+- Throughput **~ 1,892 tok/s**
+
+With these fixes applied, GPU configs should hit the same numbers (or better, thanks to bf16 + compile) without NaN or OOM.
+
+---
+
+## Patched Files
+
+| File | Lines changed | Nature of change |
+|------|---------------|------------------|
+| `helix_lm/graph.py` | +56 / −20 | `_build_node_spec()` now reads `nodes_per_column` |
+| `helix_lm/mamba2.py` | +25 / −6 | `@torch.compiler.disable` + `checkpoint` on chunk scan |
+| `helix_lm/nodes.py` | +23 / −7 | Same for `_ssm_chunked_scan` + `TitansMemoryNode.forward` |
+| `helix_lm/trainer.py` | +29 / −11 | bf16 autocast, conditional GradScaler |
+| `nas_helixlm.py` | +25 / −12 | bf16 dtype selection, remove compile skip for SSM/Titans |
+
+---
+
+## Branch
+
+`fix-gpu-tests-from-23-v2` (forked from `do-not-merge-gpu-tests-from-23`)
+
+**Note:** I do not have write access to push this branch to GitHub.  The commit `d6619f4` is ready in the local workspace; please pull / cherry-pick / review before merging to `main`.