Fix mamba regression

[manueldeprada]

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.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[manueldeprada]

run-slow: falcon_mamba

[github-actions]

This comment contains run-slow, running the specified jobs:

models: ['models/falcon_mamba']
quantizations: [] ...

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: falcon_mamba, mamba

[manueldeprada]

run-slow: falcon_mamba, mamba

[github-actions]

This comment contains run-slow, running the specified jobs:

models: ['models/falcon_mamba', 'models/mamba']
quantizations: [] ...

[ydshieh]

does this model require mambapy.pscan or it could work with pscan = None too?

[manueldeprada]

it uses pscan when the fast path (mamby library) is available. Otherwise it defaults to a slow (python code) forward pass.

The import was at the top before, I moved it for debugging and it slipped in 😅 , now I am reverting the change.

[ydshieh]

sure.

(FYI, so for the test of this model, we are testing the slow path I think)

[ydshieh]

Could you explain this part a bit more for me 🙏 , but I believe you are right.

[manueldeprada]

When using modular_converter, super.init() unravels MambaConfig.init which sets intermediate_size to int(expand * self.hidden_size), overriding any value passed via kwargs.

Before #38086, setting intermediate_size to int(expand * self.hidden_size) wasn't an issue because PretrainedConfig.__init__() was called last, and the kwargs value prevailed.

However, #38086 reversed that order due to modular_converter, which forces PretrainedConfig.__init__() to run first, thus overwriting the kwargsintermediate_size.

The new code explicitly assigns intermediate_size to ensure the kwargs value takes precedence again.

[ydshieh]

OK, I understand better. And here, since intermediate_size is stored in tiiuae/falcon-mamba-7b, during loading, it is passed as kwargs, and causing the issue.

Looks like this (issue) is something that would happen quite frequently and we have to be careful (as modular_converter force it as you mentioned)

[ydshieh]

One final nit question: do you know why we have a config (i.e. tiiuae/falcon-mamba-7b)

that have intermediate_size != int(expand * self.hidden_size)

sounds a bit strange

[manueldeprada]

Yeah, we have to be careful. In general, it is counterintuitive to have values hardcoded in the config initialization. I think kwargs should always take precedence there.

So to me, the good fix would be to change MambaConfig. As for the question, it is just their design choice. hidden size was very big, and they just chose to make the intermediates and the convs smaller.

[ydshieh]

OK for me from the explanation, but better for a second review.