Repository navigation
fix(mamba): apply Falcon-Mamba B/C/dt RMS norm once, no per-call ones - #1574
Merged
Merged
Conversation
When use_bcdt_rms is set (every falcon_mamba checkpoint), MambaBlock::ssm_step normalized delta, B and C twice per call via self.mixer_norm(&self.mixer_norm(&x)), and the old rms_norm_no_scale helper allocated a fresh ones weight array on every one of those six calls per token per layer. The upstream mlx-lm reference (mlx_lm/models/mamba.py) applies the weight-less RMS norm exactly once to each of delta, B and C when use_bcdt_rms is true, using mx.fast.rms_norm(x, weight=None, eps=mixer_rms_eps). This change makes mixer_norm call the bridge's existing fast_rms_norm_no_weight directly instead of allocating ones and calling fast_rms_norm, removes the now-unused rms_norm_no_scale helper, and applies mixer_norm once per tensor in ssm_step instead of nesting it. For a T-token prefill this drops 3T redundant norm launches and 6T ones allocations per layer (64 layers on the 7B checkpoint), with no change to the normalized values beyond the removal of the second, spurious application. Added two unit tests in mamba_tests.rs that construct a minimal MambaBlock and pin the fix at the mixer_norm boundary: one asserts mixer_norm(x) is bit-exact to a single fast_rms_norm_no_weight call and bit-distinct from a second application on top of it (array_equal is exact, not tolerance-based, which is what actually catches a regression back to the double call since the double application is numerically close to a no-op), and one asserts mixer_norm is an identity pass-through when use_bcdt_rms is false. Validated: cargo test --profile test-fast --features metal,accelerate --lib models::mamba (18 passed, 2 pre-existing ignored), cargo clippy --profile test-fast --lib --tests --features metal,accelerate -- -D warnings, cargo fmt --all -- --check. Refs #1333
inureyes
added
type:bug
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
When
use_bcdt_rmsis set (everyfalcon_mambacheckpoint),MambaBlock::ssm_stepnormalizeddelta,BandCtwice per call and allocated a freshonesweight array on every one of those six calls. This PR applies the weight-less RMS norm exactly once per tensor through the bridge's existingfast_rms_norm_no_weight, matching the upstream mlx-lm reference, and removes the now-unusedones-allocating helper.What changed
src/models/mamba.rs: deletedrms_norm_no_scale(it allocated a freshonesarray per call just to reuse the weightedfast_rms_normkernel).src/models/mamba.rs:MambaBlock::mixer_normnow callsmlxcel_core::fast_rms_norm_no_weight(x, self.mixer_rms_eps)directly.src/models/mamba.rs:ssm_stepappliesmixer_normonce per tensor (delta,B,C) instead ofself.mixer_norm(&self.mixer_norm(&x)), and the stale "applies mixer_norm TWICE" comment is replaced with a one-line description of the corrected behavior.src/models/mamba_tests.rs: addedtiny_mamba_block(a minimal, functionally-inertMambaBlockbuilder for unit tests) and two tests:mamba_mixer_norm_applies_the_bridge_norm_exactly_once, which assertsmixer_norm(x)is bit-exact to onefast_rms_norm_no_weightcall and bit-distinct from a second application on top of it, andmamba_mixer_norm_is_identity_when_bcdt_rms_disabled, which asserts the pass-through path whenuse_bcdt_rmsis false.For a
T-token prefill this drops3Tredundant norm launches and6Tonesallocations per layer (64 layers on the falcon-mamba-7b checkpoint), with no change to the normalized values beyond removing the second, spurious application.mamba2.rs,falcon_h1.rs,jamba.rs,granitemoehybrid.rsandplamo2.rscarry nouse_bcdt_rmsfield and are untouched.Test plan
cargo test --profile test-fast --features metal,accelerate --lib models::mamba(18 passed, 2 pre-existing ignored)cargo clippy --profile test-fast --lib --tests --features metal,accelerate -- -D warningscargo fmt --all -- --checkmodels/falcon-mamba-7b-4bitbefore/after this change (orchestrator follow-up)Closes #1333