using ssm_state and conv_state during training about mamba

[albertfgu]

Agree it's a useful feature, it's on the roadmap!

[albertfgu]

Unfortunately not right now. It will hopefully be ready within a few weeks.

[ejmejm]

Unfortunately not right now. It will hopefully be ready within a few weeks.

Just checking to see if this is still in progress. I was going to start working on this myself, but I would hate to spend a couple weeks figuring it out just for an official implementation to come out at the same time.

[sentialx]

@albertfgu Would states be differentiable? I would rather like to not stop the gradient during training unlike @robflynnyh

[albertfgu]

Not sure, it's a lot of extra work to make it differentiable through the final state. And that's also not necessary for the main use case that we wanted to support (continuing training on the next state). What use case are you thinking of?

[sentialx]

I'm thinking of backpropagation through time on multiple chunks (e.g. 4) instead of fitting full sequence in one huge window

[tridao]

I'm thinking of backpropagation through time on multiple chunks (e.g. 4) instead of fitting full sequence in one huge window

Would that take the same amount of activation memory as computing the full sequence?

[sentialx]

Hmm yeah, my bad

[sustcsonglin]

Are there any convenient ways to set up the initial state for mamba? I wanna use TBPTT to train mamba on longer ctx size, so there is no need to make initial/final states of each chunk differentiable.

[woominsong]

Not sure, it's a lot of extra work to make it differentiable through the final state. And that's also not necessary for the main use case that we wanted to support (continuing training on the next state). What use case are you thinking of?

@albertfgu I think differentiable initial states is a critical feature for investigating state initialization. For example, we would be able to perform prefix tuning for mamba (similar to tuning KV-cache in Transformers).