Fix numerical instability in A_log during training (closes #72)

[KakashiTech]

Root cause

The SSM recurrence uses A = -exp(A_log) where A_log is a learnable parameter. During training with Adam, the gradient dL/dA_log = dL/dA · (-exp(A_log)) = dL/dA · A. Since A < 0 always, sign(dL/dA_log) = -sign(dL/dA).

When the model needs to propagate information over long distances, it requires weaker decay (A → 0). This creates dL/dA < 0, which translates to dL/dA_log > 0, pushing A_log toward -inf. Consequently A = -exp(A_log) → 0⁻, deltaA = exp(dt·A) → 1, and the SSM becomes a pure accumulator. The hidden state grows unbounded and overflows fp16 → NaN loss.

This reproduces after 100–750 epochs of normal fp16 training with Adam at default LR, matching all reports in #72.

Fix

A = torch.clamp(A, max=-1e-4)  # one line after A = -exp(A_log)

Guarantees |A| ≥ 1e-4 at all times, ensuring the SSM always has minimum decay. The gradient naturally stops flowing through the clamp when A hits the boundary.

Why -1e-4?

• At |A| = 1e-4 with dt ≈ 4: exp(dt·A) ≈ 0.9996, time constant ≈ 2500 steps, comfortably exceeding typical sequence length 2048.
• Smaller values (e.g. |A| = 1e-6) barely decay over 2048 steps — still practically an accumulator.
• Larger values (e.g. |A| = 1e-2) cap long-range memory unnecessarily.

Verification

• Adam drift test (50K steps): Without fix, A reaches [-2.4e-6, -3.0e-7] (essentially 0). With fix, A stays at -1e-4.
• Hidden state (L=2048): |A|=2.0 → max|h|=0.05 (safe). |A|=1e-4 → max|h|=4.3 (safe). |A|=1e-6 → hidden state grows unbounded.

Impact

• Zero impact on healthy models: A initializes to -[1..d_state], well below the -1e-4 boundary. Only catches pathological drift.
• Consistent with Mamba3: Mamba3 already applies torch.clamp(_A, max=-self.A_floor) for the same reason.
• No CUDA kernel changes: A is clamped in Python before being passed to the kernel.
• Backward compatible: Preserves initialization, API, and checkpoint loading.

Files changed

• mamba_ssm/modules/mamba_simple.py — forward + step (2 locations)
• mamba_ssm/modules/mamba2.py — forward + step (2 locations)
• mamba_ssm/modules/mamba2_simple.py — forward (1 location)

5 insertions total.

[KakashiTech]

@aakashlahoti Could you take a look when you have a chance? This fixes the A_log drift issue (#72) with the same clamp pattern Mamba-3 already uses for A_floor. Only 5 lines changed across 3 files.

[Chessing234]

curious if clamping a_log changes trained checkpoints — maybe gate it behind a flag or only apply when training from scratch?