Fix Mamba2ForCausalLM weight tying

[Anri-Lombard]

What does this PR do?

Fixes #43206

Adds the _tied_weights_keys mapping to Mamba2ForCausalLM to enable proper weight tying when tie_word_embeddings=True.

The Bug

When tie_word_embeddings=True, the embedding weights should be shared with the lm_head. However, Mamba2ForCausalLM had:

_tied_weights_keys = {}  # Empty - weight tying never happens

This caused:

• Weights not actually tied (different tensors)
• resize_token_embeddings() not resizing lm_head properly

The Fix

_tied_weights_keys = {"lm_head.weight": "backbone.embeddings.weight"}

This is the standard pattern used by MambaForCausalLM (v1), GPT2, LLaMA, and other models.

Verification

from transformers import Mamba2ForCausalLM, Mamba2Config

config = Mamba2Config(vocab_size=1000, tie_word_embeddings=True)
model = Mamba2ForCausalLM(config)

# Weights are now properly tied
assert model.lm_head.weight.data_ptr() == model.backbone.embeddings.weight.data_ptr()

# Resize works correctly
model.resize_token_embeddings(1100)
assert model.lm_head.weight.shape[0] == 1100

[vasqu]

Can you add a fast test as regression test

[Anri-Lombard]

Thanks! Enabled tie_word_embeddings=True in the ModelTester to run the standard weight tying regression test. Are you happy with that? 🙏

[vasqu]

Can we instead create a small test for this? It seems that the original model did not have tied weights so this is more of an addition IMO and we check this as an explicit test (would be nice to link / mention your issue as well)

[github-actions]

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

run-slow: mamba2

[Anri-Lombard]

@vasqu Added an explicit regression test that checks both tie_word_embeddings=True and tie_word_embeddings=False 🙏

[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.

[vasqu]

Thank you, just checking with run-slow once just to be paranoid safe then merging :D

[vasqu]

Awesome, thanks for linking 🙏

[vasqu]

run-slow: mamba2

[github-actions]

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

models: ["models/mamba2"]
quantizations: []

[github-actions]

CI Results

Workflow Run ⚙️

✅ No failing test specific to this PR 🎉 !