Added mamba.py backend

[alxndrTL]

What does this PR do?

As discussed here, this PR ports the mamba.py backend to transformers. The mamba.py backend is used as a fallback for training when the official CUDA implementation is not found.
Currently, the naive version was used (which is a for loop over time). The mamba.py uses a parallel scan (just like the official implementation) which treats all inputs in parallel (technically, in log2(L) steps if enough parallel cores).
Here is the performance comparison of the three version :

Please see the mamba.py repo for a detailed analysis of the performance boost of this version.

Two things to note :

• this does not modify the behavior at inference, since the parallel scan is only useful at training where you can treat all your inputs at once.
• I've left the possibility to use the naive version via the use_mambapy argument in the MambaConfig (default to True). I have done so because the PyTorch parallel scan implemented in pscan.py is quite memory-hungry, so users with low memory capacity may want to stick with the naive version (although much slower).

You can check that the two versions (naive and mambapy) indeed gives the same numerical output for both the forward and backward pass with this code snippet :

import torch
from src.transformers.models.mamba import MambaForCausalLM, MambaConfig

torch.manual_seed(34567)
config = MambaConfig(vocab_size=60, hidden_size=64, num_hidden_layers=4, use_mambapy=False)
model_orig = MambaForCausalLM(config)

torch.manual_seed(34567)
config = MambaConfig(vocab_size=60, hidden_size=64, num_hidden_layers=4, use_mambapy=True)
model = MambaForCausalLM(config)

x = torch.randint(low=0, high=config.vocab_size-1, size=(16, 120))

y_orig = model_orig(x)
J_orig = y_orig.logits.sum()
J_orig.backward()

y = model(x)
J = y.logits.sum()
J.backward()

print(f"Is the forward pass the same ? {torch.allclose(y_orig.logits, y.logits, atol=0.01)} (difference: {torch.norm(y_orig.logits - y.logits)})")

gradients_same = True
max_diff = 0
for param1, param2 in zip(model_orig.parameters(), model.parameters()):
    diff = torch.norm(param1.grad - param2.grad)
    if diff > max_diff:
        max_diff = diff

    if not torch.allclose(param1.grad, param2.grad, atol=0.01):
        gradients_same = False
        break

print(f"Is the backward pass the same ? {gradients_same} (max difference: {max_diff})")

This is the output I got :

The fast path is not available because one of `(selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)` is None. Falling back to the implementation determined by the argument config `use_mambapy` for training. To install follow https://github.com/state-spaces/mamba/#installation and https://github.com/Dao-AILab/causal-conv1d
Is the forward pass the same ? True (difference: 5.619883449980989e-05)
Is the backward pass the same ? True (max difference: 0.0024947605561465025)

Hope my PR is clear
Bests

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

@ArthurZucker

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[ArthurZucker]

Sorry for the late reply and thanks for your hard work, will review in a bit!

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[ArthurZucker]

We rarely have full files like pscan.py in transformers, and would rather drive people to use your package 🤗 sorry for the delay in reviewing 😢

[ArthurZucker]

For such a dependency we need to protect the import! Similarly to what we do for the mamba kernels!

[ArthurZucker]

very nice and lightweight!

[ArthurZucker]

It might be better to recommend installing from your original package WDYT?
(Would drive usage there! 🤗 )

[ArthurZucker]

ps: it looks really good to me!

[alxndrTL]

Hello @ArthurZucker thanks for the review.
So I made a package version of mamba.py, that one can install with pip install mambapy.
I thus protected the imports in modeling_mamba.py using the is_mambapy_available that I defined in import_utils, just like what's being done with the CUDA version.
Also updated the warnings, telling the user that he can install mambapy if the naive version is set.
I hope this is more clear now :)

[ArthurZucker]

LGTM, just needs some cleanup, we should also add a test!

[ArthurZucker]

let's raise rather than asser!

[ArthurZucker]

let's remove this form the PR now!

[ArthurZucker]

Thanks a lot @alxndrTL. For the test, you can protect it, and run it locally. Not necessary to run it in the ci !

[ArthurZucker]

Hey! Should I take this over? Would be nice to ahve fore the release next week

[alxndrTL]

Sorry I was busy and forgot about it, I will work on it right away! For the test however, I don't know what to do precisely

[alxndrTL]

So I just removed the pscan.py file that I had forgotten to remove, and raised an error rather than using an assert.

[ArthurZucker]

Cool! Only thing left should be rebasing, and cis should go green if we correctly protected the import!

[ArthurZucker]

(Given that mamba 2 is around the corner, this will also be good for there!)

[alxndrTL]

Yes mamba-2 looks very promising I'm hoping to integrate it in mamba.py as well.
Sorry wdym by rebasing ? As for the cis, idk how to check what's the issue. (first time doing a PR on transformers haha)

[ArthurZucker]

Sure ! Errors are ERROR tests/utils/test_file_utils.py - RuntimeError: Failed to import transformers.models.mra.modeling_mra because of the following error (look up to see its traceback): cannot import name 'tarfile' from 'backports' (/home/circleci/.pyenv/versions/3.8.12/lib/python3.8/site-packages/backports/__init__.py) not sure why, let's do something like git pull upstream main or git rebase main

[ArthurZucker]

this will get you the latest changes from main and most probably fix tests! 🤗

[alxndrTL]

Hello @ArthurZucker, I rebased the branch and still get some weird errors in the tests.
check_code_quality gives me Would reformat: src/transformers/models/mamba/modeling_mamba.py.

tests_hub gives me an error that appears not related to the changes I've made.

tests_tokenization and tests_torch both yield :

56 failed because `ImportError` -> use_mambapy is set to True but the mambapy package is not installed. To install it follow https://github.com/alxndrTL/mamba.py.
Number of failed tests: 56

as if the import isn't protected but it is.

I don't really know what to do from here.
What I'm sure about is that the mambapy backends works properly (the tests between the CUDA and mambapy backend both yield the same results)

[ArthurZucker]

Hey!

• the hub is unrelated don't worry
• check code quality means you don't have the correct ruff version We pinned v1.5 I believe
• for that it just means that you activated it by default but the docker does not install it
  This should be pretty much it I believe

[alxndrTL]

@ArthurZucker so I ran make style with the correct ruff version and the check_code_quality now passes

[ArthurZucker]

Cool IMO we should default to false but leave liberty to users to select it!

[alxndrTL]

Yes you're right, done!

[ArthurZucker]

All good for me! 🚀 thanks for your great contribution I'll try to communicate about it!

[alxndrTL]

Thanks, happy to help contribute!