Error w/ training Mamba with ResNet Feature Extractor: RuntimeError: Function MambaInnerFnBackward returned an invalid gradient at index 6 - got [] but expected shape compatible with [10]

[DanaOsama]

Hello,

I want to use Mamba with a bunch of images, so I am using ResNet as a feature extractor. My mamba "wrapper class"'s init function looks like this:

self.feature_extractor = Feature_Extractor(feature_extractor)

        self.mamba = Mamba(
            d_model=feature_sizes[feature_extractor],  # Model dimension d_model
            d_state=d_state,  # SSM state expansion factor
            d_conv=d_conv,  # Local convolution width
            expand=expand,  # Block expansion factor
            bias=False,
        )

Feature_Extractor is another class that is essentially just a ResNet without the last linear layer.

I am getting the error below:

Traceback (most recent call last):
File "/home/dana/test/train.py", line 372, in
loss = train(model, optimizer, criterion, train_loader, DEVICE)
File "/home/dana/test/utils.py", line 66, in train
loss.backward()
File "/home/dana.conda/envs/multitask2/lib/python3.10/site-packages/torch/_tensor.py", line 522, in backward
torch.autograd.backward(
File "/home/dana/.conda/envs/multitask2/lib/python3.10/site-packages/torch/autograd/init.py", line 266, in backward
Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
RuntimeError: Function MambaInnerFnBackward returned an invalid gradient at index 6 - got [] but expected shape compatible with [10]

Note: Issue #37 may be relevant.

Is there a specific way I should train Mamba?
Thanks in advance!

[ZeLo-yAnG]

我找到问题的原因了，在mamba_ssm/ops/selective_scan_interface.py lin 278 and line 287
class MambaInnerFn(torch.autograd.Function): def backward(ctx, dout):
你只需要把dB_proj_bias = dB.sum(0) if not ctx.B_proj_bias_is_None else None 里面的sum(0)去掉就可以了
也就是
dB_proj_bias = dB if not ctx.B_proj_bias_is_None else None
对于dC_proj_bias也是一样的做法

I have found the cause of the problem. In mamba_ssm/ops/selective_scan_interface.py at lines 278 and 287, within the class MambaInnerFn(torch.autograd.Function): definition, specifically in the def backward(ctx, dout): method, you just need to remove the .sum(0) from the line:
dB_proj_bias = dB.sum(0) if not ctx.B_proj_bias_is_None else None
That is to say:
dB_proj_bias = dB if not ctx.B_proj_bias_is_None else None
The same approach should be applied to dC_proj_bias.