WIP: Mamba-Attention Hybrid — 18L (15 SSM + 3 GQA), First SSM Entry

[johnlennyt5]

WIP: Mamba-Attention Hybrid — 18L (15 Mamba-2 SSM + 3 GQA Attention)

Summary

First SSM-based submission to Parameter Golf. Replaces 15 of 18 layers with Mamba-2 selective state-space blocks while retaining 3
GQA attention layers for global context mixing.

Requesting RunPod credits for GPU training and evaluation.

Architecture

Parameter	Value
Total layers	18
Mamba-2 SSM layers	15 (layers 0–11, 15–17)
GQA Attention layers	3 (layers 12–14)
d_model	512
d_state (SSM)	32
d_conv	4
expand	1.5
Total params	27,191,065
Mamba params	19,249,920 (70.8%)
Estimated artifact	7.3 MB (well under 16 MB)

Current Status — CPU Validated, Awaiting GPU

• 115 unit/integration tests passing on CPU
• Full hybrid forward/backward verified (all params receive gradients)
• GPTQ int6 quantization pipeline validated end-to-end (quantize → LZMA → dequantize → forward)
• Training convergence verified on real FineWeb data
• BPB evaluation pipeline tested with actual validation tokens
• All planned ablation configurations pre-validated
• Artifact size confirmed under budget (7.3 MB LZMA, 8.7 MB headroom)

Why This Is Interesting

OpenAI specifically requested state-space model submissions. This is (to our knowledge) the first Mamba-based entry to Parameter Golf. The hybrid design addresses SSMs' known weakness on associative recall by keeping 3 attention layers at strategic positions,
while leveraging Mamba-2's O(n) sequence processing for the majority of layers.

Key technical contributions:

• Hybrid dispatch: Mamba and attention layers share a U-Net encoder-decoder with skip connections
• Muon optimizer for SSM: Mamba matrix params (in_proj, out_proj, dt_proj, c_proj) trained with Newton-Schulz Muon
• GPTQ for Mamba: Extended Hessian collection and int6 quantization to SSM parameters
• Gradient checkpointing: Selective checkpointing for Mamba layers (attention uses FA3's built-in memory management)

Test Plan

• Verify CUDA kernels on RunPod (mamba-ssm + causal-conv1d)
• Measure step time on 8xH100 (target: <100ms)
• Full 600s training run, seed 42
• Architecture ablation (layer count, attention positions, d_state)
• 3-seed evaluation with Welch's t-test vs SOTA
• Create submission package with logs

Run Command

MAMBA_LAYERS=0,1,2,3,4,5,6,7,8,9,10,11,15,16,17 NUM_LAYERS=18 \
MAMBA_D_STATE=32 MAMBA_D_CONV=4 MAMBA_EXPAND=1.5 \
MAMBA_MATRIX_LR=0.015 BIGRAM_VOCAB_SIZE=2048 BIGRAM_DIM=128 \
WARMDOWN_ITERS=3500 TARGET_MB=15.9 SEED=42 \
torchrun --standalone --nproc_per_node=8 train_gpt.py

Requirements

mamba-ssm>=2.2.0
causal-conv1d>=1.4.0
sentencepiece

[MatoTeziTanka]

Community Review — WIP: Mamba-Attention Hybrid — 18L (15 SSM + 3 GQA), First SSM Entry

BPB: (not parsed — see PR title) | Compliance: LOOKS CLEAN — score-first-per-chunk TTT (legal #1413 dexhunter pattern)

What I found in the code (head SHA 9b775e152041, file train_gpt.py):

The TTT path at line 1276 implements the score-first-per-chunk pattern: each chunk is scored under torch.no_grad() / inference_mode() before the base_model.train() + SGD adaptation runs on that same chunk, with an is_last_chunk guard so the final chunk gets no adaptation pass. This is the structural shape of the current leaderboard's legal frontier (PR #1413 dexhunter, the 1.0828 SP8192 + QK-Gain 5 + Legal TTT entry — verified at its head SHA against the is_last_chunk + torch.no_grad() score-first accumulator pattern).

Per Issue #402 and Issue #677, TTT is legal when each token is scored before the adapter updates on it, and that's what the code does here — chunk ci is scored under weights adapted only on chunks 0..ci-1. No prequant_ttt_adapt_adamw(val_tokens, ...) multi-epoch fine-tune, no scored-region SLOT, no target-in-key n-gram cache.

CPU smoke test (CT2038 proteus-engine, 2026-04-11): import OK in 1.27s, dim=512, layers=11, vocab=1024, code=128982 B, SMOKE_TEST_PASS

Verdict: LOOKS CLEAN.

Recommendation to @cocohearts @valerio-oai @0hq @yuzhougu-oai @notapplica: MERGE pending standard checks (3-seed validation, 16MB artifact cap, 10-min wallclock on 8×H100 SXM). The compliance picture matches the legal reference frontier and no flags were raised by the classification pass.

Auto-classification caveat: this review was drafted by the AST-based classifier against a template derived from manually-reviewed cluster PRs (#1420, #1450, #1487, #1541, #1529, #1533, #1518). If I've misread a subtlety in your eval path — e.g., multi-epoch TTT that I mistook for single-pass, or a target-in-key lookup I missed in a helper function — please flag it and I'll re-run the audit manually.

---

Reviewed by @MatoTeziTanka — The Agora. CPU smoke test (CT2038 proteus-engine, 2026-04-11): import OK in 1.27s, dim=512, layers=11, vocab=1024, code=128982 B, SMOKE_TEST_PASS. Classification via deterministic AST-based classify_prs.py (pattern bank derived from ~65 manually-reviewed PRs earlier in the 2026-04-11 sweep). This review was auto-drafted from a template and spot-checked before posting — if the template misread your code, please call it out so I can iterate the classifier.