Registering (forward/backward) hooks on intermediate Mamba layers about mamba

[albertfgu]

In the fused fast implementations, the inner modules might not be getting called directly. Instead, we use their weights (e.g. backbone.layers.0.mixer.conv1d.weight) and pass it into a different function (e.g. from the causal-conv1d package).

[SrGonao]

I tried setting the fast path variable of the layers to false but I still couldn't get it to work. Is it even possible in the pretained versions?

[albertfgu]

Is it even possible in the pretained versions?

It should be. Pretrained versions just give a set of weights which are the same no matter the computation path.

I tried setting the fast path variable of the layers to false but I still couldn't get it to work.

You should double check that it's running the path you intend and that the modules are being called directly. For example, you want to hit the following line to call the conv module, instead of using the fast causal_conv1d_fn. Perhaps it's not hitting this path if you've already pip installed the module because it will automatically use the fast conv1d function.

https://github.com/state-spaces/mamba/blob/main/mamba_ssm/modules/mamba_simple.py#L169

[SrGonao]

Thank you I hadn't seen that logic there. Because it does not seem feasible to make sure that causal-conv1d is not installed, I will probably have to do some hacky work around to go around that

[albertfgu]

It should be fairly easy to add another flag to the init function and change the line

            if causal_conv1d_fn is None:

to something like

            if causal_conv1d_fn is None or not self.use_fast_conv:

[SrGonao]

Yes, that's what I was thinking