Data preparation methods

[RonanKMcGovern]

I was somewhat surprised to see that Mamba's context length is finite.

I assume this means Mamba was trained on 2k token batches, but resetting the state of "h" in some kind of aggregate way.

Would it be better to feed the value of h at the end of one 2k sequence to be used as the starting value for the next sequence (which would appear in the next batch)?

In principle, this should result in the model being capable of almost infinite length. Or am I misunderstanding something?

[tridao]

That'd be interesting to try! Transformer-XL did something similar.

[radarFudan]

This stateful approach seems to be common in RNN models (https://arxiv.org/abs/1803.00144).

[radarFudan]

For example, the famous LSTM language model by salesforce is using this idea: https://github.com/salesforce/awd-lstm-lm/blob/master/model.py

[jzhang38]

I had similar ideas. Unfortunately, mamba currently does not support setting the initial hidden state during training.
#146
CUDA master Tri Dao said it will be implemented soon tho.
I plan to try it ASAP once this functionality is ready.
@RonanKMcGovern @tridao @radarFudan or anyone else wanna work together on this? My discord handle is peiyuan007
(btw I have already scaled the mamba fine-tuning length to 16384 and the test PPL keeps going down until ~40K on proof pile. I may write a note/blog soon regarding my findings so far. I optimistically believe we can realize infinite context length. This may also lead to something publishable.)

[radarFudan]

I am also interested in this! I have some slow solution to achieve the setting hidden states during training. (#51)

[jzhang38]

Yeah the simplest way is to just substitute selective_scan with a torch implementation, which is what I am currently doing.

[RonanKMcGovern]

@jzhang38 that's cool. How much memory did you need to train the 16k context. I tried on an A100 (80 GB) and ran OOM. Did you also train all modules or freeze some?

[CompRhys]

+1 for setting the initial state, setting the initial state is likely to be important when using the model in industrial control settings

[jzhang38]

https://github.com/jzhang38/LongMamba

[RonanKMcGovern]

@jzhang38 that is so good, thanks

[JulienSiems]

A related paper just came out: https://arxiv.org/pdf/2406.02080v1