Fixes and Enhancements for Mamba Inference and Reference Implementations #743
Conversation
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
This pull request addresses several bugs and limitations within the Mamba codebase, primarily aimed at improving inference robustness in the Mamba2 module and increasing the accuracy of reference implementations. Key changes include: In mamba_ssm/modules/mamba2.py: Resolved an issue in _get_states_from_cache to correctly handle dynamic batch sizes during inference, ensuring proper state re-initialization when batch sizes change. Removed the batch == 1 assertion in the forward method for variable-length sequence inference, enabling batched processing for these inputs. Updated the fallback path in the step method to support ngroups > 1, allowing grouped SSM inference even if Triton kernels are not available. In mamba_ssm/ops/selective_scan_interface.py: Added optional RMS Normalization for B, C, and delta tensors within mamba_inner_ref to better match the main MambaInnerFn's behavior and improve numerical consistency. Corrected a shape comment in selective_scan_ref for clarity. In mamba_ssm/models/mixer_seq_simple.py: Removed a redundant comment in the _init_weights function. In mamba_ssm/utils/hf.py: Addressed a bug in load_state_dict_hf to ensure correct dtype conversion and device placement when loading Hugging Face model weights. These modifications enhance the stability, flexibility, and correctness of the Mamba library.
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.
This pull request addresses several bugs and limitations within the Mamba codebase, primarily aimed at improving inference robustness in the Mamba2 module and increasing the accuracy of reference implementations.
Key changes include:
batch == 1assertion in the forward method for variable-length sequence inference, enabling batched processing for these inputs.ngroups > 1, allowing grouped SSM inference even if Triton kernels are not available.mamba_ssm/ops/selective_scan_interface.py:mamba_inner_refto better match the mainMambaInnerFn's behavior and improve numerical consistency.selective_scan_reffor clarity.mamba_ssm/models/mixer_seq_simple.py:_init_weightsfunction.mamba_ssm/utils/hf.py:load_state_dict_hfto ensure correct dtype conversion and device placement when loading Hugging Face model weights.These modifications enhance the stability, flexibility, and correctness of the Mamba library.