Missing `from mamba_ssm import Mamba` in `lit_gpt/model.py`

[lucianyao]

Description

In lit_gpt/model.py, the create_block function uses Mamba without importing it from mamba_ssm, causing a NameError: name 'Mamba' is not defined.

Steps to Reproduce

1. Clone the repository: git clone https://github.com/NVlabs/GatedDeltaNet.
2. Ensure mamba_ssm is installed: pip install mamba_ssm==2.2.4.
3. Run a test script:

   from lit_gpt.model import create_block
   block = create_block(d_model=512)
   print('Block created')

Error occurs

Traceback (most recent call last):
File "", line 2, in
File "/path/to/GatedDeltaNet/lit_gpt/model.py", line 50, in create_block
mixer_cls = partial(Mamba, layer_idx=layer_idx, **ssm_cfg, **factory_kwargs)
NameError: name 'Mamba' is not defined

Expected Behavior

The code should run without errors when mamba_ssm is installed.

Actual Behavior

NameError due to missing import.

Suggested Fix

Add from mamba_ssm import Mamba at the top of lit_gpt/model.py. After adding this line, the test script runs successfully:
Block created

Environment

Python: 3.10
mamba_ssm: 2.2.4
OS: Ubuntu

[kyleliang919]

Tried with docker, still doesn't work
#11 (comment)