Repository: huggingface/transformers
Pull request: Fix mamba regression

This fixes the sneaky regression introduced in #38086 causing loading errors for falcon_mamba:
```
 "line": "tests/models/falcon_mamba/test_modeling_falcon_mamba.py::FalconMambaModelTest::test_model_from_pretrained",
                "trace": "(line 2593)  RuntimeError: Error(s) in loading state_dict for FalconMambaMixer:"
            },
            {
                "line": "tests/models/falcon_mamba/test_modeling_falcon_mamba.py::FalconMambaIntegrationTests::test_batched_generation",
                "trace": "(line 2593)  RuntimeError: Error(s) in loading state_dict for FalconMambaMixer:"
            },
            {
                "line": "tests/models/falcon_mamba/test_modeling_falcon_mamba.py::FalconMambaIntegrationTests::test_generation_4bit",
                "trace": "(line 2593)  RuntimeError: Error(s) in loading state_dict for FalconMambaMixer:"
            },
            {
                "line": "tests/models/falcon_mamba/test_modeling_falcon_mamba.py::FalconMambaIntegrationTests::test_generation_fp16",
                "trace": "(line 2593)  RuntimeError: Error(s) in loading state_dict for FalconMambaMixer:"
            },
            {
                "line": "tests/models/falcon_mamba/test_modeling_falcon_mamba.py::FalconMambaIntegrationTests::test_generation_torch_compile",
                "trace": "(line 2593)  RuntimeError: Error(s) in loading state_dict for FalconMambaMixer:"
```

The gist of the problem is: modular forces the `super().__init__` call to be on top of FalconMambaConfig. However, before modular rewrite, it was at the bottom, which was critical for the `intermediate_size` property from the config file to take effect.

Second bug fixed: `tests/models/mamba/test_modeling_mamba.py::MambaIntegrationTests::test_compile_mamba_cache ` was failing due to a missplaced import.

Changed files:
filename: src/transformers/models/falcon_mamba/configuration_falcon_mamba.py
patch: @@ -141,7 +141,12 @@ def __init__(
         self.layer_norm_epsilon = layer_norm_epsilon
         self.conv_kernel = conv_kernel
         self.expand = expand
-        self.intermediate_size = int(expand * self.hidden_size)
+        # This is needed since mamba overrides the intermediate_size attribute
+        self.intermediate_size = (
+            int(expand * self.hidden_size)
+            if kwargs.get("intermediate_size") is None
+            else kwargs.get("intermediate_size")
+        )
         self.bos_token_id = bos_token_id
         self.eos_token_id = eos_token_id
         self.pad_token_id = pad_token_id

filename: src/transformers/models/falcon_mamba/modular_falcon_mamba.py
patch: @@ -192,6 +192,12 @@ def __init__(
             **kwargs,
         )
         self.mixer_rms_eps = mixer_rms_eps
+        # This is needed since mamba overrides the intermediate_size attribute
+        self.intermediate_size = (
+            int(expand * self.hidden_size)
+            if kwargs.get("intermediate_size") is None
+            else kwargs.get("intermediate_size")
+        )
 
 
 class FalconMambaCache(MambaCache):

filename: src/transformers/models/mamba/modeling_mamba.py
patch: @@ -39,6 +39,10 @@
 
 logger = logging.get_logger(__name__)
 
+if is_mambapy_available():
+    from mambapy.pscan import pscan
+else:
+    pscan = None
 
 if is_mamba_ssm_available():
     from mamba_ssm.ops.selective_scan_interface import mamba_inner_fn, selective_scan_fn
@@ -330,10 +334,6 @@ def cuda_kernels_forward(
 
     # fmt: off
     def slow_forward(self, input_states, cache_params: Optional[MambaCache]=None, cache_position:Optional[torch.LongTensor]=None, attention_mask: Optional[torch.LongTensor] = None):
-        if is_mambapy_available():
-            from mambapy.pscan import pscan
-        else:
-            pscan = None
         batch_size, seq_len, _ = input_states.shape
         dtype = input_states.dtype
         # 1. Gated MLP's linear projection