Skip to content
Bud's Wiki

Mamba-2 dt_bias initialization bug in HuggingFace transformers and flash-linear-attention

TL;DR

Section titled “TL;DR”

A Mamba-2 implementation bug in the HuggingFace transformers and flash-linear-attention (FLA) repositories: the SSM time-step bias dt_bias was being incorrectly initialized. Two distinct root causes were identified — (1) a wrong constant init (torch.ones) hit when Mamba-2 layers were used standalone outside the Mamba2ForCausalLM wrapper, and (2) FSDP-2 + DTensor “meta device” init silently skipping the custom initialization. The author reports the loss gap at 7B-MoE scale is substantial — Mamba-2 is unusually sensitive to dt_bias init. Both issues have associated FLA PRs (one merged, one pending).

Key claims

Section titled “Key claims”
  • The HF transformers and fla-org/flash-linear-attention repos had a dt_bias init bug that incorrectly used torch.ones when Mamba-2 layers were instantiated in isolation (i.e. without going through Mamba2ForCausalLM); a fix has been merged into FLA [tweet body].
  • A second, distinct issue: under FSDP-2 with DTensors, the custom Mamba-2 init is skipped entirely because parameters are constructed on the meta device — so even with the correct init code path, no custom init actually runs [tweet body].
  • The training-quality impact is substantial at 7B-MoE scale (“Mamba-2 seems to be quite sensitive to the initialization”); the author links wandb runs comparing correct vs. buggy init [tweet body, linked W&B project mayank31398/ma...].
  • The upstream state-spaces/mamba README has long warned that frameworks with post-init hooks that zero nn.Linear biases will silently destroy the targeted dt_proj.bias range and require a framework-specific opt-out — i.e. the FLA/FSDP-2 failure mode is precisely the class of footgun the README flagged [tweet body, corroborated by upstream Mamba README].

Method

Section titled “Method”

Not a paper — a bug report tweet with two pull requests (one merged, one in review on fla-org/flash-linear-attention) and a W&B run set as evidence. The mechanism: Mamba-2’s time-step parameter Δ has a targeted log-uniform range set by initializing the bias dt_bias = log(uniform(time_step_min, time_step_max)) (softplus-inverted), which gives Δ a useful initial spread of timescales across heads. Any framework-level hook that overrides this — generic “zero all biases”, or simply not running the custom init because the parameter was constructed on the meta device under FSDP-2 + DTensor — replaces that careful spread with a constant, and downstream the SSM recurrence sees a degenerate timescale distribution.

The two specific failure paths reported:

  1. Standalone-layer path: when a Mamba2 mixer block is used outside the Mamba2ForCausalLM wrapper (e.g. as a sub-module in a custom architecture), the correct init was guarded by the wrapper class and the fallback was torch.ones. Fixed in FLA.
  2. FSDP-2 meta-init path: FSDP-2 constructs parameters as DTensors on the meta device for memory efficiency; the custom Mamba-2 init code did not handle the meta-device case and was silently a no-op. A PR fixing this in FLA was pending at the time of the tweet.

Results

Section titled “Results”

No formal evaluation numbers in the tweet — the author points at a W&B project (wandb.ai/mayank31398/ma...) showing the gap at the 7B-MoE scale. Reported qualitatively as “substantial.” The fact that both fixes were considered worth a public callout, and that one had to ship through FLA quickly, suggests the gap was large enough to invalidate prior 7B-scale results trained with the buggy code path.

Why it’s interesting

Section titled “Why it’s interesting”

This is the empirical mirror of the “stability via initialization / parametrization” thread the wiki has been tracking — concretely, it’s a real-world instance where init dominates outcomes at the 7B-MoE scale, complementing Training stability at scale‘s architecture-vs-optimizer-vs-parametrization framing with a fourth, mundane locus: implementation correctness of the init itself. Two specific connections: (a) it sharpens How to Set the Learning Rate for Large-Scale Pre-training?‘s claim that QK-Norm subsumes µP’s stability role — Mamba-2 has no QK-Norm and its stability really does depend on a careful targeted init, so “the architecture handles it” is architecture-specific; (b) the FSDP-2 meta-device failure mode is a systems-side analogue to the Controlled LLM Training on Spectral Sphere argument that activations only stay bounded end-to-end if both weights and updates are controlled — here the weights drift to a wrong init not from optimizer drift but from a silent framework hook. Worth flagging for any Luma team building on Mamba-2 (or hybrid Mamba-attention) blocks via HF + FSDP-2.

See also

Section titled “See also”