BitMamba 1B — training the first SSM model

Context

The TatuEngine’s hybrid CPU/GPU pipeline was already working — 252× speedup, Block Codec compressing 36 GB into 281 MB (128×), 41.9 µs per token. But there was a critical gap: the model couldn’t generate coherent text.

BitMamba-1B is a PyTorch implementation of Mamba-2, with 48 layers. ~1.4B parameters total (1.446B). Compared to a Transformer of similar size — 7B, 70B — it seems small, but for a pure SSM it’s still a 7-headed beast.

The problem is that Mamba doesn’t follow the same transfer learning path as Transformers. You don’t take a Llama-2 checkpoint and fine-tune it. There’s no from_pretrained that works — the representation space is fundamentally different.

The training timeline

Full Warmup — 200 pretrain steps

The first checkpoint, full_warmup, was trained with pure pretrain data (train_ids.pt, 96MB). 200 steps, batch 2, GPT-2 tokenizer. Result: a 1.94 GB checkpoint that knows token statistics but hasn’t formed language — perplexity ~444M on holdout.

# Loading checkpoint full_warmup/best
# Auto-detect: dims read from checkpoint
# Parameters: 1,446,747,136 total
# VRAM: 2.93 GB

Phase 2 — SFT with cold start

I took full_warmup and tried supervised fine-tuning with 1000 reasoning examples (code, systems, deduction). Three attempts:

An average loss of 100 means the model was guessing almost uniformly among 50K vocabulary tokens. log(50288) ≈ 10.8 would be the loss of a random model. Our loss of 100+ indicated that the model wasn’t just failing to learn — it was actively sabotaging itself, probably through exploding activations or NaN gradients.

# Sample generated at step 100 (Phase 2):
# "372366 � partName� startlingOpt elevate Gemuchs charitychelamer Ideally..."
# Format [THOUGHT]: 
# Format [ANSWER]:  

This isn’t language. It’s token noise with positional bias.

Phase 3 — Continuous SFT

I continued training from the best Phase 2 checkpoint. Evaluation result:

 EVALUATION REPORT
──────────────────────────────────
Average loss:  19.9131
Perplexity:  444,768,275.92
Tokens evaluated:  84,844
Examples:  201

By domain:
  code:  Loss 19.87  PPL 424M
  deduction: Loss 19.33  PPL 249M
  systems:  Loss 19.78  PPL 388M

Perplexity of 444 million. For reference: a random-guessing model has perplexity ~50K (log perplexity ~vocab size). 444M means the model’s token distribution is worse than random — it actively assigns low probability to the correct tokens.

Sample generated by the model:

"K Moon acoustic sign Ossvertis0 ParkM em anINa� usel�vertis scoutingl Pac..."

Diagnosis: what went wrong?

Three problems combined:

1. Under-trained pretrain

200 warmup steps isn’t enough for a 1.4B parameter model. For reference, Mamba-2 models in the literature are trained on hundreds of billions of tokens. We had maybe 10M tokens.

The warmup serves to stabilize activations — A_log, dt_bias, norms — but not to form linguistic representations. Semantics emerge at data scales 3-4 orders of magnitude larger.

2. SFT too early

I skipped the continued pretrain phase and went straight to SFT with 1000 examples. The result is catastrophic because:

• The model has no internal language representation to fine-tune
• The SFT loss pushes the model in directions that shallow pretrain can’t sustain
• With batch 1 (VRAM limitation), the gradient is extremely noisy

3. Mixing semantic data with a non-semantic model

It’s like trying to teach quantum physics to someone who hasn’t learned to count. SFT assumes the model already “understands” language — it just needs to be adjusted for a specific format (user: ...\nassistant: ...). Without the foundation, SFT destroys instead of builds.

The solution: Phase 4 — Hybrid Pipeline

The resurrection plan (Phase 4) attacks all 3 problems at once:

Phase 1 (steps 0000-0300): PURE PRETRAIN
  train_ids.pt, pretrain only
  Goal: bring down perplexity + stabilize cross-entropy

Phase 2 (steps 0301-2000): CO-TRAINING 3:1
  3 pretrain examples : 1 SFT example per batch
  Loss = w_pretrain * L_pretrain + w_sft * L_sft
  w_pretrain: 0.70 → 0.50 (linear decay)
  Natural language learning anchors the SFT

Phase 4 hyperparameters

The magic: 4 parameter groups with separate LR

The most important trick is treating each part of the model with different learning rates:

groups = [
  {'name': 'embed',  'lr': lr * 0.3,  'weight_decay': 0.01},  # vocabulary
  {'name': 'state',  'lr': lr * 10.0,  'weight_decay': 0.0},  # A_log, D, dt_bias
  {'name': 'residual', 'lr': lr * 0.5,  'weight_decay': 0.01},  # MLP + norm
  {'name': 'base',  'lr': lr,  'weight_decay': 0.01},  # mixer, proj
]

• Embeddings at 0.3×: The vocabulary was already reasonably trained in warmup — no need for drastic changes
• State params at 10×: This is the shock to the hippocampus. A_log, D, dt_bias are the parameters that control the SSM’s memory. If they’re not pushed hard, the model never learns to carry state between tokens
• Residual at 0.5×: The MLP + RMSNorm layers already have some structure — change slowly to avoid breaking
• Base at 1.0×: Linear projections (in_proj, out_proj, conv1d) — the model’s engine

Complete reasoning dataset structure

For Phase 4, the SFT dataset has 1000 examples balanced by domain using sample weights:

Each example follows the format:

user: <prompt>
assistant: <response with [THOUGHT] and [ANSWER]>

Co-training loss weighting

The key insight of co-training is that the model can’t forget pretrain while learning SFT:

co_progress = (global_step - PHASE1_STEPS) / max(1, PHASE2_STEPS)
w_pretrain = 0.70 + (0.50 - 0.70) * co_progress  # 70% → 50%
w_sft = 1.0 - w_pretrain  # 30% → 50%
loss = w_pretrain * loss_pretrain + w_sft * loss_sft

It starts with 70% weight on pretrain (to maintain the foundation), and gradually drops to 50% (giving more room for SFT to shape the output).

Gradient telemetry at step 1

An addition I paid dearly to learn: measure the gradient on the first step.

if global_step == 1:
  raw_norm = clip_grad_norm_(model.parameters(), max_norm=inf)
  zero_count = sum((p.grad.abs() < 1e-8).all() for p in model.parameters()...)
  print(f"Grad norm: {raw_norm:.6f}")
  print(f"{zero_count}/{total_grad} params with grad ≈ 0")

This immediately detects if:

• The gradient is exploding (norm > 100)
• The gradient is dead (80%+ of params with grad ≈ 0)
• Some subgroup isn’t receiving gradients (e.g., frozen state params)

A_log telemetry at step 50

A_log controls the SSM’s decay — essentially the model’s “memory”:

if global_step == 50:
  mean_A = exp(mean_a_log)
  mem_tau = 1.0 / max(1e-10, mean_A)
  print(f"Mean A_log: {mean_a_log:.4f} | |A|: {mean_A:.4f} | τ: ~{mem_tau:.0f} tok")

If τ < 1, the model forgets everything between consecutive tokens — impossible to learn long-range dependencies. The ideal is τ ≈ 100-500 tokens.

Lessons learned

1. Mamba is more sensitive to hyperparameters than Transformer

A Transformer with SGD + simple warmup converges to a reasonable loss even with insufficient data. Mamba SSM has 3 coupled subsystems (temporal conv1d + recurrent SSM + linear projections) that need to be balanced. The warmup needs to be longer, the LR needs to be lower, and gradient clipping is mandatory.

2. Batch size 1 is a waste of time

With 12GB of VRAM, the maximum I could achieve was batch 4 (seq_len 768) or batch 8 (seq_len 512). Batch 1, which I used in Phase 2, produces such noisy gradients that loss doesn’t descend — even with grad_accum 4.

3. Parameter groups with separate LR aren’t optional — they’re what saves you

Without the state group at 10× LR, the SSM’s memory never learns to propagate information. Without embed at 0.3× LR, the vocabulary corrupts in 50 steps. This separation by function (vocab, memory, reasoning, projection) was the most expensive insight — it cost ~3 weeks of failed training.

4. Perplexity of 444M is worse than random — but it’s diagnosable

Log perplexity > log vocab_size = the model is actively wrong. Almost certain cause: diverging activations combined with softmax that pushes probability toward specific tokens (usually the first ones in the vocabulary, positions 0-1000). Solution: activation clipping + monitoring mean A_log.

5. Co-training is fragile — needs an isolated pretrain checkpoint

If the model enters co-training without finishing pretrain (stable loss, stabilized A_log, gradient norm < 10), the SFT contaminates the pretrain and both suffer. Pure pretrain needs to reach loss < 5-6 before any SFT.

The numbers (so far)

What’s next

The tatu_phase4_ressurection.py script is ready, linted, and waiting for execution. The next step is:

1. Run Phase 4 with 2000 steps (~3-4 days on RTX 3060)
2. Validate whether perplexity drops below 500 (intermediate milestone)
3. If it works, expand the pretrain dataset with more data (code, docs, wikis)
4. Try generation with the post-Phase 4 checkpoint

If it doesn’t work… well, then I’m going back to Transformer. But the SSM premise still seems right for edge devices: 1.4B params in 2.93GB of VRAM, inference at 41.9µs with Block Codec. The potential exists — it’s the training that’s missing.

# The lesson that will stick:
# Training SSM is 修真 — every advance costs blood, sweat, and CUDA errors.
# But when the gradient passes cleanly on step 1, you know the path exists.

TL;DR: I trained a Mamba-2 with 1.4B parameters in 12GB of VRAM. Warmup works, SFT fails, the 444M perplexity is the diagnosis that the pretrain was too shallow. Phase 4 attacks with 2000 steps of hybrid co-training — 300 pure pretrain + 1700 with weighted loss 70%→50%. If it works, we open the path to trainable SSMs on consumer hardware. If not, at least the logs are documented for the next attempt.

╔══════════════════════════════════════╗
║  Samuel Medeiros                    ║
║  Senior Software Engineer           ║
║  Stack: Python · TypeScript · Rust  ║
║  Projetos: Arachne, Dogwalk,        ║
║            Capivara, TatuEngine      ║
╚══════════════════════════════════════╝