ssm_enhancement #689
ssm_enhancement #689
Conversation
These enhancements provide additional flexibility and options for implementing and experimenting with different recurrence methods in the Mamba and Jamba models, potentially improving performance and accuracy for various tasks.
|
Hey @markblee , Could you please take a look at my PR when you get a chance? Thanks! |
|
Hey @swiseman , Could you please take a look at my PR when you get a chance? Thanks! |
|
|
||
|
|
||
|
|
||
| class HybridMambaRecurrence(BaseMambaRecurrence): |
There was a problem hiding this comment.
Thanks for these new classes. Do people use either the hybrid recurrences or alternative recurrences defined below? Is there evidence that they are useful empirically? If not, I think it would be simpler to leave these classes out for now, and if necessary let people define them in downstream experiment files which import axlearn.common.ssm.
There was a problem hiding this comment.
Hey @swiseman , Thank you for your valuable input. I've reviewed the hybrid recurrences and alternative recurrences, and it seems that they haven't been used extensively in practice. Based on your benchmarking results, it appears that the AssociativeScanMambaRecurrence is more efficient than the HybridMambaRecurrence.
Given the lack of empirical evidence and the performance advantage of the AssociativeScanMambaRecurrence, I agree that it's reasonable to remove the HybridMambaRecurrence and other less-used recurrences from the core axlearn.common.ssm module for now.
This will simplify the codebase and make it easier for users to understand and use. If there's a strong need for these recurrences in the future, they can be defined in downstream experiment files as you suggested.
-Vishesh
- fixed functions redundant definitions - fixed Incorrect Module Import in layers.py
Hey @ruomingp we can close this PR, I am working on these hybrid structures and it will take a while... Thanks for reading; |
Pull Request: Enhancements to Mamba and Jamba State-space Models
Summary
This pull request introduces several enhancements to the Mamba and Jamba state-space models (SSMs) implementation, including new recurrence methods, hybrid approaches, and comprehensive testing.
Changes
1. New Recurrence Methods
2. Enhancements to
ssm.pyHybridMambaRecurrenceandAlternativeMambaRecurrenceclasses.MambaMixerLayerandJambaMixerLayerto integrate the new recurrence methods.3. Comprehensive Testing in
ssm_test.pyHybridMambaRecurrenceandAlternativeMambaRecurrenceinMambaMixerLayerTest.StackedMambaTest.StackedMixedSSMTransformerTest.(Documentation and Examples : Updated docstrings and comments to reflect the new features and changes.)
Testing
All new features have been thoroughly tested with the following configurations:
jnp.float32,jnp.bfloat16).MambaBlock,JambaMambaBlock, andStackedMixedSSMTransformerLayer.Conclusion
These enhancements provide additional flexibility and options for implementing and experimenting with different recurrence methods in the Mamba and Jamba models, potentially improving performance and accuracy for various tasks.