Support for SSM models (Mamba, Mamba2)

Hi, Was wondering if adding support for SSM models (Mamba, Mamba2) would be something you'd consider having in litgpt? https://arxiv.org/abs/2405.21060

Thanks!