Add Zamba2
Author: pglorio
Pull request: yes

# What does this PR do?

Please include support for Zamba2 architecture created by Zyphra Technologies.

- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [X] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

@ArthurZucker


Comments:
Hey @arthur,

Thank you again for your help in getting Zamba2 into `transformers`! The PR is now finally ready to be reviewed. I added the documentation and all unit tests pass, including slow tests.

A few remarks, mostly related to `modular transformers`:

1. To generate modeling and configuration I used `utils/modular_model_converter.py` from a [previous commit](https://github.com/huggingface/transformers/blob/9e3d704e2340fe9b306b5bd6b12605e4341c012b/utils/modular_model_converter.py) because the most [recent version](https://github.com/huggingface/transformers/blob/e2ac16b28a0b8b900e136750309ca40c49d975c5/utils/modular_model_converter.py) of this script that followed from a large refactoring produces an error that I was not able to fix:
```
Converting src/transformers/models/zamba2/modular_zamba2.py to a single model single file format
Traceback (most recent call last):
  File "/workspace/transformers_zamba/utils/modular_model_converter.py", line 1510, in <module>
    converted_files = convert_modular_file(file_name, args.old_model_name, args.new_model_name)
  File "/workspace/transformers_zamba/utils/modular_model_converter.py", line 1447, in convert_modular_file
    for file, module in create_modules(cst_transformers).items():
  File "/workspace/transformers_zamba/utils/modular_model_converter.py", line 1387, in create_modules
    nodes_to_add, file_type, new_imports = get_class_node_and_dependencies(modular_mapper, class_name, node, files)
  File "/workspace/transformers_zamba/utils/modular_model_converter.py", line 1337, in get_class_node_and_dependencies
    new_node_dependencies, new_imports = check_dependencies_and_create_import_node(
  File "/workspace/transformers_zamba/utils/modular_model_converter.py", line 1283, in check_dependencies_and_create_import_node
    class_dependencies = {dep for dep in new_dependencies if m.matches(mapper.global_nodes[dep], m.ClassDef())}
  File "/workspace/transformers_zamba/utils/modular_model_converter.py", line 1283, in <setcomp>
    class_dependencies = {dep for dep in new_dependencies if m.matches(mapper.global_nodes[dep], m.ClassDef())}
KeyError: 'Zamba2Config'
```

I carefully compared `Zamba2Config` with classes of other models that also use `modular` (such as `Gemma2Config`) and they appear to have consistent format. Relatedly, the `utils/modular_model_converter.py` in the current PR ([path](https://github.com/Zyphra/transformers_zamba/blob/ffc2a58f3dfdee943c8343c73e6155f025de780f/utils/modular_model_converter.py)) is the version from the [previous commit](https://github.com/huggingface/transformers/blob/9e3d704e2340fe9b306b5bd6b12605e4341c012b/utils/modular_model_converter.py) mentioned above.

2. After running `utils/modular_model_converter.py`, the modeling and configuration files generated contain unintended code that I had to update. All these modifications are in [this commit](https://github.com/Zyphra/transformers_zamba/commit/9adf85e00b42a7dca00ff13c8646f504dad6c2c4). In particular, the produced modeling file contains `Zamba2DynamicCache`, which is the correct cache of Zamba2 as well as `HybridMambaAttentionDynamicCache`, which is the cache of Zamba and is not relevant to Zamba2, so I deleted `HybridMambaAttentionDynamicCache` and related references.

3. I ran `make fixup` and all zamba-related tests pass, with the exception of `python utils/check_modular_conversion.py`. This test doesn't pass due to the modifications mentioned in the previous point.

4. I slightly edited the Zamba2MambaMixer compared to the original Mamba2Mixer of mamba2, the main difference is that I added [these lines](https://github.com/Zyphra/transformers_zamba/blob/4da8d5ff6dc5aed96db4d7ec6f2dbf4a3255688c/src/transformers/models/zamba2/modular_zamba2.py#L1138-L1142), which was necessary to appropriately process the mamba2 cache (note this step already existed in the `torch forward` in [these lines](https://github.com/Zyphra/transformers_zamba/blob/4da8d5ff6dc5aed96db4d7ec6f2dbf4a3255688c/src/transformers/models/zamba2/modular_zamba2.py#L1213-L1217)).

Looking forward to your feedback. Thanks so much!

Thank you @ArthurZucker! 

To follow up on your previous message:

- We are using LoRA modules inside the pretrained model. More explicitly, inside the shared transformer blocks, we added un-shared LoRA layers to give some amount of non-shared parameters inside the transformer forward (see e.g. [here](https://github.com/Zyphra/transformers_zamba/blob/10a0b1e1bcc8d8b541ea93b9bddeafc20e695858/src/transformers/models/zamba2/modular_zamba2.py#L1232)). We call these layers LoRA because they have the same structure as the usual LoRA, but they are unrelated to finetuning, they are part of the base model and should be included in the modular/modeling files.
- 

> mamba2 should be used for the mamba chunk scan implementation

I may not be understanding this, do you mean to replace [these lines](https://github.com/Zyphra/transformers_zamba/blob/10a0b1e1bcc8d8b541ea93b9bddeafc20e695858/src/transformers/models/zamba2/modular_zamba2.py#L55-L66) with:

```
from ..mamba2.modeling_mamba2 import (
    mamba_chunk_scan_combined, 
    mamba_split_conv1d_scan_combined, 
    selective_state_update, 
    causal_conv1d_fn, 
    causal_conv1d_update,
)
```
? 
Just to be sure: we did not import directly the mamba2 layer (i.e., `mamba2.Mamba2Mixer`) because both the init and the forward methods are slightly different. As mentioned in a [previous msg](https://github.com/huggingface/transformers/pull/34517#issuecomment-2467369443), we made a few edits inside of the forward of this layer, mostly related to the `attention_mask`. In the init, some of the attributes are different, e.g. we set `self.num_heads = self.config.n_mamba_heads` instead of `self.num_heads = config.num_heads` because in the config we distinguish between `num_attention_heads` and `n_mamba_heads`, while `num_heads` would be ambiguous.


Also, below are the updates I made:

- I inherited the `Zamba2HybridDynamicCache` from `ZambaHybridDynamicCache` as you requested. In the `__init__` method, I omitted `super().__init__()` as discussed with @Cyrilvallez, because this would add unwanted lines involving `self.conv_states` and `self.ssm_states` that come from `ZambaHybridDynamicCache` (`Zamba2HybridDynamicCache` manages the SSM cache the same way as mamba2, which uses a slightly different structure for `self.conv_states` and `self.ssm_states` compared to `ZambaHybridDynamicCache`).
- Regarding the config, I didn't inherit it from `ZambaConfig` because some flags of `ZambaConfig` do not exist in `Zamba2Config`, such as `use_mamba_kernels`. Also many default values of `Zamba2Config` are different from those of `ZambaConfig`. I noticed that Gemma2Config doesn't inherit from the GemmaConfig, so because of these reasons I thought it is a simpler option to just define the config from scratch. We're happy to revisit this if needed.
- I imported `MambaRMSNormGated` from mamba2
- Inherited `Zamba2AttentionDecoderLayer` (I overrode the forward method because `hidden_states = self.feed_forward(hidden_states, layer_idx)` takes the `layer_idx` in [Zamba2](https://github.com/Zyphra/transformers_zamba/blob/10a0b1e1bcc8d8b541ea93b9bddeafc20e695858/src/transformers/models/zamba2/modular_zamba2.py#L1343), this was not the case for [Zamba](https://github.com/Zyphra/transformers_zamba/blob/10a0b1e1bcc8d8b541ea93b9bddeafc20e695858/src/transformers/models/zamba/modeling_zamba.py#L847))
- Inherited `Zamba2MLP`. There are three differences with `ZambaMLP`: 1. added [pretrained LoRA](https://github.com/Zyphra/transformers_zamba/blob/10a0b1e1bcc8d8b541ea93b9bddeafc20e695858/src/transformers/models/zamba2/modular_zamba2.py#L1251) (as described above) and 2. replaced `self.up_proj`, and `self.gate_proj` with their [concatenation](https://github.com/Zyphra/transformers_zamba/blob/10a0b1e1bcc8d8b541ea93b9bddeafc20e695858/src/transformers/models/zamba2/modular_zamba2.py#L1248) `self.gate_up_proj` (this is necessary as the pretrained LoRA was added collectively to the concat of up and gate), 3. modified the [activation](https://github.com/Zyphra/transformers_zamba/blob/10a0b1e1bcc8d8b541ea93b9bddeafc20e695858/src/transformers/models/zamba2/modular_zamba2.py#L1243) as it should take the output of self.gate_up_proj as input, instead of the output of self.gate_proj.
- I renamed shared_transf -> shared_transformer
- In `Zamba2HybridLayer`, I [overrode the forward](https://github.com/Zyphra/transformers_zamba/blob/10a0b1e1bcc8d8b541ea93b9bddeafc20e695858/src/transformers/models/zamba2/modular_zamba2.py#L1371) of `ZambaHybridLayer` because the latter has `self.shared_transf` in its forward method (I cannot change `shared_transf` -> `shared_transformer` in Zamba1 because this would require replacing ~200 checkpoints stored in the hub). The only way to avoid overriding the forward would be to keep `shared_transf` instead of renaming it to `shared_transformer`.
- As agreed with @Cyrilvallez, I did not inherit `Zamba2PreTrainedModel`. The reason is that if I inherit `Zamba2PreTrainedModel` I need to overload <this method> in `ZambaPreTrainedModel` with <this method> in `PretrainedModel`. Calling `super(PreTrainedModel, cls)._check_and_enable_flash_attn_2` didn't work, so I found it easier to simply not inherit `Zamba2PreTrainedModel` (in fact, this approach makes the class in modular shorter).
- I re-ran `make fixup`, all tests pass except two that are unrelated to zamba:
```
python utils/check_modular_conversion.py

Differences found between the generated code and src/transformers/models/gemma/configuration_gemma.py:

   1 --- src/transformers/models/gemma/configuration_gemma.py_generated
   2 +++ src/transformers/models/gemma/configuration_gemma.py
   3 @@ -19,6 +19,7 @@
   4  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   5  # See the License for the specific language governing permissions and
   6  # limitations under the License.
   7 +
   8  
   9  from ...configuration_utils import PretrainedConfig
  10  
Traceback (most recent call last):
  File "/workspace/transformers_zamba/utils/check_modular_conversion.py", line 78, in <module>
    raise ValueError("Some diff and their modeling code did not match.")
ValueError: Some diff and their modeling code did not match.
```
and
```
Traceback (most recent call last):
  File "/workspace/transformers_zamba/utils/check_repo.py", line 1211, in <module>
    check_repo_quality()
  File "/workspace/transformers_zamba/utils/check_repo.py", line 1195, in check_repo_quality
    check_all_objects_are_documented()
  File "/workspace/transformers_zamba/utils/check_repo.py", line 1076, in check_all_objects_are_documented
    check_public_method_exists(documented_methods_map)
  File "/workspace/transformers_zamba/utils/check_repo.py", line 1125, in check_public_method_exists
    raise Exception("\n".join(failures))
Exception: The following public method is explicitly documented but not defined in the corresponding class. class: MusicgenMelodyFeatureExtractor, method: _extract_stem_indices. If the method is defined, this error can be due to lacking dependencies.
```


Thank you!

cc @Cyrilvallez can you provide help here? 🤗 

Hey @pglorio, so sorry for the delay!! I think the first error you encounter comes from a wrong `ruff` version. Our CI uses `ruff==0.5.1`, so could you make sure to install the same and run it again?

EDIT: but actually, the only failing tests I see are Zamba2 tests

Thank you @Cyrilvallez  for the super helpful feedback! Please see above the replies to each comment.

It seems zamba2-related tests all pass locally for us, but not in the circle/ci setting of the PR. The error is always 
```
   1 failed because `IndexError` -> Dimension out of range (expected to be in range of [-2, 1], but got 2)
   9 failed because `RuntimeError` -> Sizes of tensors must match except in dimension 2. Expected size 8 but got size 1 for tensor number 1 in the list.
```
for all tests that fail. We investigated the discrepancy by examinating more closely the test `tests/models/zamba2/test_modeling_zamba2.py::Zamba2ModelTest::test_beam_sample_generate`. We tried using `pytest-flakefinder` as suggested [here](https://github.com/huggingface/transformers/pull/26304#issuecomment-1819382300) and using this command `python3 -m pytest -m 'generate' -n 8 -vvv --dist=loadfile --flake-finder --flake-runs=1000 --max-worker-restart=0 tests/models/zamba2/test_modeling_zamba2.py::Zamba2ModelTest::test_beam_sample_generate` but the test still passes. This is the closest reproduction of the circle/ci test that we could reproduce to the best of our knowledge. We'd appreciate if we can get a hint of what we can do to fix this discrepancy!

Also, we are happy to rename lora layers as you suggested [here](https://github.com/huggingface/transformers/pull/34517#pullrequestreview-2476481472) and agree it would clarify the role of this module. Would `pretrained_lora` be also good as an option (`residual` might for some people refer to the identity stream)?

Thank you so much again!

We renamed lora -> adapter [here](https://github.com/huggingface/transformers/pull/34517/commits/337faed6afa70c244357c44b2076bd3920728d04) and, following your suggestion in [this comment](https://github.com/huggingface/transformers/pull/34517#discussion_r1868257901) we simplified the adapter using `nn.Sequential` in [this commit](https://github.com/huggingface/transformers/pull/34517/commits/f1b31a13289255e55f612041dea9df4a467831e6).

We re-ran the model tests and they all pass locally, but it looks like they still fail in this PR environment.

I cannot reproduce your issue with tests though, they are failing locally for me. Running `CUDA_VISIBLE_DEVICES=2 pytest tests/models/zamba2/test_modeling_zamba2.py` gives:
<img width="1340" alt="Screenshot 2024-12-18 at 17 47 59" src="https://github.com/user-attachments/assets/f6d00761-acfb-4dd1-84fc-0f0431a3911a" />


Thank you so much for [this feedback](https://github.com/huggingface/transformers/pull/34517#issuecomment-2551816849) @Cyrilvallez. I realized that the issue with unit tests was inside the `torch_forward` method of the mamba2 mixer (when i ran locally, the unit tests used `cuda_kernels` method instead). I fixed that method here: [1](https://github.com/huggingface/transformers/pull/34517/commits/5d0a5d46c79f83b97b5bf24eb42a9c47750decce) [2](https://github.com/huggingface/transformers/pull/34517/commits/ef055c90de487771d5ee70f9fac148c752c99c96) [3](https://github.com/huggingface/transformers/pull/34517/commits/b993a7895e7e42be6a3f9647759a19320f6bd73e).

By the way, we originally took the the `torch_forward` from the mamba2 model, so the same issues hold there. In particular, running this:

```python
config = Mamba2Config(num_heads=8,
        n_groups=8,
        state_size=2,
        head_dim=8,
        conv_kernel=4,
        chunk_size=8,
        vocab_size=99,
        hidden_size=32,
        num_hidden_layers=4,
        hidden_act="silu",
        hidden_dropout_prob=0.1,
        max_position_embeddings=512,
                      )
model = Mamba2ForCausalLM(config)

inputs = {'input_ids': torch.tensor([[86,  6, 51,  3, 12, 15, 33, 18,  4, 92],
         [69, 66, 49, 45, 48, 44, 61, 56, 68, 85]]),
 'attention_mask': torch.tensor([[0, 0, 1, 1, 1, 0, 0, 0, 1, 0],
         [0, 1, 1, 0, 1, 1, 1, 0, 1, 1]])}

outputs_cpu = model.generate(**inputs, max_new_tokens=5, return_dict_in_generate=False, output_scores=False, use_cache=True, num_beams=1, do_sample=False)
model = model.to('cuda')
inputs = {key: tensor.to(device=model.device) for key, tensor in inputs.items()}
outputs_cuda = model.generate(**inputs, max_new_tokens=5, return_dict_in_generate=False, output_scores=False, use_cache=True, num_beams=1, do_sample=False)
print(torch.all(outputs_cpu == outputs_cuda.cpu()).item())
```
returns `False`.

Hi @Cyrilvallez and @ArthurZucker,

I updated the attention forward to the new standard of `transformers` [here](https://github.com/huggingface/transformers/pull/34517/commits/cb90bb4efbab0243bc86ada0203dacfb60080c8b) and [here](https://github.com/huggingface/transformers/pull/34517/commits/1dbc8c73e54e71345cfe676eea704aa862b7fc84). 

I ran all final tests, including `@slow` tests, and everything appears to pass!

Thank you @Cyrilvallez for the review. I addressed the comments above, although there are a couple of pending points.

All zamba-related tests appear to pass.

Hello @Cyrilvallez, I ran all model tests on two GPUs and after a couple of minor fixes everything appears to work now. I'm skipping [this](https://github.com/huggingface/transformers/pull/34517/commits/9afb57ec0dd2e877c7977a188caaaa0b64dc3115#diff-67d87c36da1310eb4bb253eb029c63484c1a052ae06224eb5e69e72755ec46b2R336-R339) test as it gives an error related to mamba2 kernels. I indeed verified that mamba2 skips that test [here](https://github.com/huggingface/transformers/blob/9afb57ec0dd2e877c7977a188caaaa0b64dc3115/tests/models/mamba2/test_modeling_mamba2.py#L260).

Separately, when running `utils/check_modular_conversion.py` I get the following error:

```
Differences found between the generated code and src/transformers/models/zamba2/modeling_zamba2.py:

   1 --- src/transformers/models/zamba2/modeling_zamba2.py_generated
   2 +++ src/transformers/models/zamba2/modeling_zamba2.py
   3 @@ -313,6 +313,13 @@
   4      return attn_output, attn_weights
   5  
   6  
   7 +def rotate_half(x):
   8 +    """Rotates half the hidden dims of the input."""
   9 +    x1 = x[..., : x.shape[-1] // 2]
  10 +    x2 = x[..., x.shape[-1] // 2 :]
  11 +    return torch.cat((-x2, x1), dim=-1)
  12 +
  13 +
  14  def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
  15      """Applies Rotary Position Embedding to the query and key tensors.
  16  
  17 @@ -338,13 +345,6 @@
  18      q_embed = (q * cos) + (rotate_half(q) * sin)
  19      k_embed = (k * cos) + (rotate_half(k) * sin)
  20      return q_embed, k_embed
  21 -
  22 -
  23 -def rotate_half(x):
  24 -    """Rotates half the hidden dims of the input."""
  25 -    x1 = x[..., : x.shape[-1] // 2]
  26 -    x2 = x[..., x.shape[-1] // 2 :]
  27 -    return torch.cat((-x2, x1), dim=-1)
```
which I was not getting before despite this part was identical. 

Awesome, sounds good!

Thank you @ArthurZucker! I think all your comments have been addressed. All zamba-related tests appear to pass!

Thanks @ArthurZucker, I replied to your comments above.