johnma2006 / mamba-minimal Goto Github PK

Hi， great work！
Could you please explain why in selective_scan the "x = torch.zeros((b, d_in, n), device=deltaA.device)"?
In addition, I am confusing on u and x.

Thanks.

I note that the convolution used in the original paper is causal convolution, but I don't seem to see an implementation of causal convolution in this project.
A common grouping convolution is used in model.py. I wonder if this is correct

Nice work! I'm considering porting this to MLX, but I wonder if I can copy parts of your code. Could you please add a license to this repo?

Thanks

Hello.

Error 'AttributeError:' NoneType 'object has no attribute' seek '. You can only torch. load from a file that is searchable. Please pre load the data into a buffer like io.BytesIO and try to load from it install.'.

I have assigned weights and the file path is also correct. I don't know how to solve it. May I ask for your advice

Hi,
Thanks for providing the Mamba implementation. I would like to know if there is any workaround in the efficient computation of deltaA and deltaB_u that can avoid the GPU memory running out issue. The following are the parameters I used to create the Mamba instance:

d_model: 1024
n_layer: 4   
d_state: int = 1024
expand: int = 2

The other parameters are set to their default values.

It results in a model of ~60M parameters. However, I run out of memory (max GPU memory= 24 GB) when I train with a batch size of 256 or even as low as 64 and this probably happens due to large matrix computations for deltaA and deltaB_u.

Thanks for the clear implementation!

Can you explain the discretization of $B$ in selective scan?

Equation 4 in section 2 of the paper states $$\overline{B} = (\Delta A)^{-1} (exp((\Delta A) - I) \cdot \Delta B$$

In your implementation, the input is mapped into the hidden state by the following:

 deltaB_u = einsum(delta, B, u, 'b l d_in, b l n, b l d_in -> b d_in l n')

which if I understand correctly, implies that $\overline{B} = \Delta B$?

Sorry, i can`t connect to huggingface directly.

Maybe i could load the pre-trainned weights loaclly?(If someone could TEACH me!!!Please!!!)

I understand that d_inner is d_model * expansion (E=2) . But why is self.in_proj = nn.Linear(args.d_model, args.d_inner * 2 ...).

Why is the in projection expanded a second time by 2 ?

I can't seem to find the answer in the appointed paper section 3.4.

Any clarification would be appreciated.

Trying to run it on MPS backend on Mac M1 Max 64Gb modifying demo.ipynb like this:

Changes to notebook:

device = 'mps'

model = Mamba.from_pretrained(pretrained_model_name).to(device)

input_ids = tokenizer(prompt, return_tensors='pt').input_ids.to(device)

Getting gibberish and dead python kernel in dozen of tokens:

for t in generate(model, tokenizer, 'Mamba is the'): print(t, end='')

ocardial goal fibrobl ( triglycer ા judgement extraordinacebook CURIAM

and then it crashes.

Hey! Thanks so much for putting this together. You mention that the parameter initialization isn't correct, can you point me towards what it should be? I've tried looking at the official repo and didn't see any out-of-the-ordinary stuff.

why add "-" in this code?why not use A directly instead of A_log?
A = -torch.exp(self.A_log.float())

I have tried several versions of Transformers, but none of them work. May I ask what version of Torch and Transformers the blogger has installed

Nice implementation!

I thought that Mamba was somewhat recurrent, like keeping an internal state and then outputting one token at a time. But your code shows that for each new output token, the entire sentence must be passed to the model, including the last outputted token.

Is this the only mode of operation?

It also has the l parameter (sequence length). Does it mean it has a maximum sequence length? The paper shows up to 1 million, so I was expecting it to be recurrent and without limit

From the paper:

Table 2: (Induction Heads.) Models are trained on sequence length 2^8 = 256, and tested on increasing sequence lengths of 2^6 = 64 up to 2^20 = 1048576

It is not clear how they do that

I haven't run the official version of Mamba, but I've run your implementation, and it seems that the training speed of this model is much slower than that of the Transformer.

Uploading CyberTech services.zip…

hi John,

Is there any plan of implemention with https://github.com/ml-explore/mlx

thank you for your simple implementation, which is really helpful to learn for noobs like me

ignore this...

There is a way to perform the selective scan with two cumulative sums or torch.cumsum, which is effectively like a parallel scan but supported by pytorch natively.

I made a minimal commit in my fork here PeaBrane@2908f50. The correctness and functionality are tested, and I could observe an inference speed up of ~14x on an A30. But not sure how close it is to the original impl with parallel scan still. More details are here.

If intersted, it would be nice if someone could review this change, and discuss whether this could be merged here, albiet the explicitness of the code may suffer (as I understand the repo is meant to be pedagogical).

Nice implementation! I notice a small detour below and wonder if it's necessary:
In this line you define self.A_log = nn.Parameter(torch.log(A)) which is only used here A = -torch.exp(self.A_log.float()) to exp it back. What's the reason for not defining A directly in the class parameter?

In your file: https://github.com/johnma2006/mamba-minimal/blob/master/model.py, you have:

A = repeat(torch.arange(1, args.d_state + 1), 'n -> d n', d=args.d_inner)
self.A_log = nn.Parameter(torch.log(A))
A = -torch.exp(self.A_log.float()) # shape (d_in, n)

Does matrix A need to be diagonal?

Thanks!

torch.load without weights_only is a potential security issue:

Found using https://github.com/pytorch-labs/torchfix:

model.py:117:20: TOR102 [*] torch.load without weights_only parameter is unsafe. Explicitly set weights_only to False only if you trust the data you load and full pickle functionality is needed, otherwise set weights_only=True.

Hi, thanks for the great work! I would like to know does mamba-minimal has GPU op optimization? It seems to me that it doesn't have. I want to train a large-scale mamba and am currently considering the original mamba and mamba-minimal.