MoE Mamba

Implementation of MoE Mamba from the paper: "MoE-Mamba: Efficient Selective State Space Models with Mixture of Experts" in Pytorch and Zeta.

PAPER LINK

Install

pip install moe-mamba

Usage

`MoEMambaBlock`

import torch 
from moe_mamba import MoEMambaBlock

x = torch.randn(1, 10, 512)
model = MoEMambaBlock(
    dim=512,
    depth=6,
    d_state=128,
    expand=4,
    num_experts=4,
)
out = model(x)
print(out)

Code Quality 🧹

make style to format the code
make check_code_quality to check code quality (PEP8 basically)
black .
ruff . --fix

Citation

@misc{pióro2024moemamba,
    title={MoE-Mamba: Efficient Selective State Space Models with Mixture of Experts}, 
    author={Maciej Pióro and Kamil Ciebiera and Krystian Król and Jan Ludziejewski and Sebastian Jaszczur},
    year={2024},
    eprint={2401.04081},
    archivePrefix={arXiv},
    primaryClass={cs.LG}
}

License

MIT

Metadata

Release files for moe-mamba 0.0.3

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for moe-mamba 0.0.3
File	Size	Uploaded
moe_mamba-0.0.3.tar.gz	5.6 kB	Jan 22, 2024	Details

Built distribution (wheel)

Table of built distributions (wheels) for moe-mamba 0.0.3
File	Interpreter	ABI	Platform	Reset
moe_mamba-0.0.3-py3-none-any.whl 5.4 kB Jan 22, 2024	Python 3	none	any	Details

Total release size: 10.9 kB

Release files / moe_mamba-0.0.3.tar.gz

Download URL	moe_mamba-0.0.3.tar.gz
Size	5.6 kB
Tags	Source
SHA-256 checksum How to use checksums	`abf077f44d29ef046973d47d5559525d8ddff5962a79b0590ffcb80d7438fc36`
BLAKE2b-256 checksum How to use checksums	`c46224bd10853d3843c06556a28116bd4db38f01089dcc9c4db7bd5b1152cbd9`
Upload date	Jan 22, 2024
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`poetry/1.3.2 CPython/3.11.0 Darwin/22.4.0`

Release files / moe_mamba-0.0.3-py3-none-any.whl

Download URL	moe_mamba-0.0.3-py3-none-any.whl
Size	5.4 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`937d8d3cf2c65058f74761228b5879a4eb7c985b15e47a572a4048ddbbc7e913`
BLAKE2b-256 checksum How to use checksums	`550fc5c5ba01552ebfb72618a7d612f69602a51de4333b2f9b3b76927f61383d`
Upload date	Jan 22, 2024
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`poetry/1.3.2 CPython/3.11.0 Darwin/22.4.0`