Add Zamba
Author: pglorio
Pull request: yes
Labels: name: New model

name: run-slow

# What does this PR do?

Please include support for Zamba architecture created by Zyphra Technologies.

## Before submitting
- [ ] 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.
- [X] 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).
- [X] Did you write any new necessary tests?


## Who can review?

@ArthurZucker @younesbelkada

Comments:
Just for future reference, we measured latency of a single mamba layer in Zamba and compared it to that of a single layer in Mamba, which have very similar implementations (we have some reshapings in Zamba, but they should be a non-op, and a concatenation), and found that that the mamba layer in Zamba to have the same speed in a single forward pass, but to be slower on generation.

More specifically, we instantiated these two (random) models:

```
config = ZambaConfig(num_hidden_layers=81, hidden_size=3712, n_mamba_heads=1, use_cache=True)
model_1 = ZambaForCausalLM(config).cuda()
config = MambaConfig(num_hidden_layers=81, hidden_size=3712, use_cache=True)
model_2 = MambaForCausalLM(config).cuda()
```

(here n_mamba_heads=1 corresponds to the original Mamba architecture), and use this code for generation:

```
model.eval()
input_ids = torch.randint(1000, (1, 2048)).to(device=model.device)
with torch.no_grad():
    output = model.generate(input_ids, max_new_tokens=300, return_dict_in_generate=False, output_scores=False, use_cache=True, num_beams=1, do_sample=False)
```

We found that the total time spent computing this line
https://github.com/huggingface/transformers/blob/87ec8722c936b52ea33969623810c5234bd070ee/src/transformers/models/zamba/modeling_zamba.py#L1009
is 8.1s, and for this line
https://github.com/huggingface/transformers/blob/87ec8722c936b52ea33969623810c5234bd070ee/src/transformers/models/mamba/modeling_mamba.py#L341
is 6.3s.

cc @younesbelkada !

I tried running basic training script with gradient accumulation and without on this fork and am getting this error:
 File "/home/user/transformers_zamba/src/transformers/models/zamba/modeling_zamba.py", line 1051, in forward
    hidden_states = hidden_states + from_tf if from_tf is not None else hidden_states
RuntimeError: The size of tensor a (7424) must match the size of tensor b (3712) at non-singleton dimension 2

The from_tf is not well described in the doc strings. Not sure what is not working here.



> I tried running basic training script with gradient accumulation and without on this fork and am getting this error:
File "/home/user/transformers_zamba/src/transformers/models/zamba/modeling_zamba.py", line 1051, in forward
hidden_states = hidden_states + from_tf if from_tf is not None else hidden_states
RuntimeError: The size of tensor a (7424) must match the size of tensor b (3712) at non-singleton dimension 2

> The from_tf is not well described in the doc strings. Not sure what is not working here.

Thanks for spotting this. We fixed the issue in the most recent push. Please try again and let us know if you still encounter issues.

We are adding more docstrings to explain various parts of the architecture. We will add the description below for `from_tf` around [this line](https://github.com/Zyphra/transformers_zamba/blob/0b1db69682847e82a76277f10910487ebe162d8a/src/transformers/models/zamba/modeling_zamba.py#L1016):

`from_tf is the output of shared transformer + linear layer (these layers are shown in fig. 2 in https://arxiv.org/pdf/2405.16712). from_tf is then added to the input to the mamba layer (as described in eq. (6) of https://arxiv.org/pdf/2405.16712, where y_l in that equation is from_tf).`

> Thanks a lot for this PR ! I left some minor suggestions in the modeling code for general improvements
> Can you also make sure to rebase with main and make sure make fixup pass locally ? Let me know if you need any assistance!

Thank you for the thorough review!

We ran `make fixup` and `make fix-copies`. Running again `make fixup` gives this output:

```
Checking/fixing src/transformers/models/bigbird_pegasus/modeling_bigbird_pegasus.py src/transformers/models/zamba/configuration_zamba.py src/transformers/models/zamba/modeling_zamba.py tests/models/roc_bert/test_tokenization_roc_bert.py
All checks passed!
4 files left unchanged
python utils/custom_init_isort.py
python utils/sort_auto_mappings.py
python utils/check_doc_toc.py --fix_and_overwrite
running deps_table_update
updating src/transformers/dependency_versions_table.py
python utils/check_copies.py
Traceback (most recent call last):
  File "/workspace/transformers_zamba/utils/check_copies.py", line 1106, in <module>
    check_copies(args.fix_and_overwrite, args.file)
  File "/workspace/transformers_zamba/utils/check_copies.py", line 856, in check_copies
    raise Exception(
Exception: Found the following copy inconsistencies:
- tests/models/roc_bert/test_tokenization_roc_bert.py: copy does not match models.bert.test_tokenization_bert.BertTokenizationTest.test_is_whitespace at line 167
Run `make fix-copies` or `python utils/check_copies.py --fix_and_overwrite` to fix them.
```

It looks like `make fix-copies` is trying to correct parts of the code that are outside of our PR, and some of those fixes still fail. However, we now do not seem to get errors related to our PR.

We pushed the fixes done by `make fix-copies`.

Tried training again and am now getting this:

/trainer.py", line 3250, in training_step
    self.accelerator.backward(loss)
  File "/home/user/mambaforge/envs/zamba/lib/python3.10/site-packages/accelerate/accelerator.py", line 2127, in backward
    loss.backward(**kwargs)
  File "/home/user/mambaforge/envs/zamba/lib/python3.10/site-packages/torch/_tensor.py", line 525, in backward
    torch.autograd.backward(
  File "/home/user/mambaforge/envs/zamba/lib/python3.10/site-packages/torch/autograd/__init__.py", line 267, in backward
    _engine_run_backward(
  File "/home/user/mambaforge/envs/zamba/lib/python3.10/site-packages/torch/autograd/graph.py", line 744, in _engine_run_backward
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/home/user/mambaforge/envs/zamba/lib/python3.10/site-packages/torch/autograd/function.py", line 301, in apply
    return user_fn(self, *args)
  File "/home/user/mambaforge/envs/zamba/lib/python3.10/site-packages/torch/utils/checkpoint.py", line 320, in backward
    torch.autograd.backward(outputs_with_grad, args_with_grad)
  File "/home/user/mambaforge/envs/zamba/lib/python3.10/site-packages/torch/autograd/__init__.py", line 267, in backward
    _engine_run_backward(
  File "/home/user/mambaforge/envs/zamba/lib/python3.10/site-packages/torch/autograd/graph.py", line 744, in _engine_run_backward
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: Trying to create tensor with negative dimension -274340773632: [-274340773632]


> Tried training again and am now getting this:
> 
> /trainer.py", line 3250, in training_step self.accelerator.backward(loss) File "/home/user/mambaforge/envs/zamba/lib/python3.10/site-packages/accelerate/accelerator.py", line 2127, in backward loss.backward(**kwargs) File "/home/user/mambaforge/envs/zamba/lib/python3.10/site-packages/torch/_tensor.py", line 525, in backward torch.autograd.backward( File "/home/user/mambaforge/envs/zamba/lib/python3.10/site-packages/torch/autograd/**init**.py", line 267, in backward _engine_run_backward( File "/home/user/mambaforge/envs/zamba/lib/python3.10/site-packages/torch/autograd/graph.py", line 744, in _engine_run_backward return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass File "/home/user/mambaforge/envs/zamba/lib/python3.10/site-packages/torch/autograd/function.py", line 301, in apply return user_fn(self, *args) File "/home/user/mambaforge/envs/zamba/lib/python3.10/site-packages/torch/utils/checkpoint.py", line 320, in backward torch.autograd.backward(outputs_with_grad, args_with_grad) File "/home/user/mambaforge/envs/zamba/lib/python3.10/site-packages/torch/autograd/**init**.py", line 267, in backward _engine_run_backward( File "/home/user/mambaforge/envs/zamba/lib/python3.10/site-packages/torch/autograd/graph.py", line 744, in _engine_run_backward return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass RuntimeError: Trying to create tensor with negative dimension -274340773632: [-274340773632]

Hey there! We've been successfully using: https://github.com/huggingface/alignment-handbook/compare/main...Zyphra:alignment-handbook:zamba-instruct here recently to do sft of Zamba. Does your setup meaningfully differ from this? I can't seem to reproduce, can you provide us a reproducer?

cc @younesbelkada should I review this or do you want to do another pass? 🤗 

I am trying to extend the max context length. 
{
  "max_position_embeddings": 32768,
   "rope_theta": 192144,
}

also tried at 16k.

I tried running in your fork of alignment handbook and saw the same results.

> Thanks very much for your great work on this ! I left few minor improvements to address and some file changes to revert - can you make sure to make our CI happy (by making sure `make fixup` command passes + the tests pass `pytest tests/models/zamba/` pass ) let me know if you need any help or have any question

Thank you for your help, @younesbelkada!

We believe we have addressed most of the concerns you raised; we still have two pending questions:

- `pytest tests/models/zamba/`: Pytest flags only test_initialization as failing. The specific issue arises with `x_proj_weight` and `dt_proj_weight`, where their mean is approximately 10^-2, contrary to the expected 10^-9. This discrepancy is expected, it is due to the initialization scheme using a variance of (d_input)^(-0.5), where d_input is approximately 100 in the test configuration. We implemented nn.Parameter(torch.rand(...)) for initialization of these parameters, which we verified is equivalent to the Kaiming initialization typically used for nn.Linear. It seems that `transformers` may apply additional steps for the initialization of various layer types, which might not extend to parameters such as `x_proj_weight`. We have adjusted the tolerance for these parameters to 10^-2 in the initialization test in [this line](https://github.com/Zyphra/transformers_zamba/blob/main/tests/models/zamba/test_modeling_zamba.py#L353) of the test script. Please let us know if additional steps are required.

- `make fixup`: After running it, we have this output:

```xception: Found the following copy inconsistencies:
- src/transformers/models/bigbird_pegasus/modeling_bigbird_pegasus.py: copy does not match models.big_bird.modeling_big_bird.BigBirdBlockSparseAttention at line 720
- tests/models/roc_bert/test_tokenization_roc_bert.py: copy does not match models.bert.test_tokenization_bert.BertTokenizationTest.test_chinese at line 76
- tests/models/roc_bert/test_tokenization_roc_bert.py: copy does not match models.bert.test_tokenization_bert.BertTokenizationTest.test_basic_tokenizer_lower at line 85
- tests/models/roc_bert/test_tokenization_roc_bert.py: copy does not match models.bert.test_tokenization_bert.BertTokenizationTest.test_basic_tokenizer_lower_strip_accents_false at line 94
- tests/models/roc_bert/test_tokenization_roc_bert.py: copy does not match models.bert.test_tokenization_bert.BertTokenizationTest.test_basic_tokenizer_lower_strip_accents_true at line 103
- tests/models/roc_bert/test_tokenization_roc_bert.py: copy does not match models.bert.test_tokenization_bert.BertTokenizationTest.test_basic_tokenizer_lower_strip_accents_default at line 112
Run `make fix-copies` or `python utils/check_copies.py --fix_and_overwrite` to fix them.
make: *** [Makefile:38: repo-consistency] Error 1
```

all the lines are related to files outside of our PR, so we did not change those files, although indeed I do see that the CircleCI tests performed in this PR still fail. Please, let us know if further action is needed here and what would be the steps we'd need to take.

Thank you so much for your time and help!

Thank you so much for your guidance. We tried to rebase our PR and ran into an error related to model generation. It looks like the rebased `GenerationMixin.generate` method instantiates Zamba's cache as a `DynamicCache` class https://github.com/huggingface/transformers/blob/main/src/transformers/generation/utils.py#L1775. This is different from `HybridMambaAttentionDynamicCache` which would be the expected class for Zamba's cache (defined here https://github.com/Zyphra/transformers_zamba/blob/main/src/transformers/models/zamba/modeling_zamba.py#L130). In the non-rebased fork, the cache is instantiated by `GenerationMixin.generate` in this line: https://github.com/Zyphra/transformers_zamba/blob/main/src/transformers/generation/utils.py#L2379, which correctly instantiates cache as `HybridMambaAttentionDynamicCache`.

For reference, these are the calls performed from `model.generate` to the instantiation of the cache object:
using the rebased fork:
```
-> output = model.generate(**tokenized_prompt, max_new_tokens=300, return_dict_in_generate=False, output_scores=False, use_cache=True, num_beams=1, do_sample=False)
  /usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py(115)decorate_context()
-> return func(*args, **kwargs)
  /workspace/transformers_zamba_rebased/src/transformers/generation/utils.py(1775)generate()
-> model_kwargs["past_key_values"] = DynamicCache()
> /workspace/transformers_zamba_rebased/src/transformers/cache_utils.py(305)__init__()
```

and using the fork before rebasing:
```
-> output = model.generate(**tokenized_prompt, max_new_tokens=300, return_dict_in_generate=False, output_scores=False, use_cache=True, num_beams=1, do_sample=False)
  /usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py(115)decorate_context()
-> return func(*args, **kwargs)
  /workspace/transformers_zamba/src/transformers/generation/utils.py(1743)generate()
-> result = self._sample(
  /workspace/transformers_zamba/src/transformers/generation/utils.py(2379)_sample()
-> model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)  /workspace/transformers_zamba/src/transformers/models/zamba/modeling_zamba.py(1588)prepare_inputs_for_generation()
-> past_key_values = HybridMambaAttentionDynamicCache(
> /workspace/transformers_zamba/src/transformers/models/zamba/modeling_zamba.py(146)__init__()
```
Could you please let us know how we can force Zamba's cache to be `HybridMambaAttentionDynamicCache` in the rebased fork?

Thanks very much!

He! Super late in coming back to you, you should be able to force it by setting `cache_class = ""` in the ZambaPreTrainedModel class!

> He! Super late in coming back to you, you should be able to force it by setting cache_class = "" in the ZambaPreTrainedModel class!

Hello @ArthurZucker, thank you for the suggestion! We tried adding `cache_class = ""` to [this line](https://github.com/Zyphra/transformers_zamba/blob/cc7f7aaaf65e59aa7fbbd1322c7126c943af661b/src/transformers/models/zamba/modeling_zamba.py#L1106) but we still couldn't make generation work. As an alternative fix, we added `"zamba"` to [this line](https://github.com/Zyphra/transformers_zamba/blob/cc7f7aaaf65e59aa7fbbd1322c7126c943af661b/src/transformers/generation/utils.py#L1500), similarly to what was done with Jamba, in which case generation works fine. 

We are happy to either keep this fix or to use the one you suggested, in which case we would appreciate if you could say a few more words on how to implement it.

Meanwhile, we rebased our fork. All the local tests with `make fixup` have passed, except for a few warnings shown below which are unrelated to the updates we implemented:

```
/workspace/transformers_zamba/src/transformers/deepspeed.py:24: FutureWarning: transformers.deepspeed module is deprecated and will be removed in a future version. Please import deepspeed modules directly from transformers.integrations
/workspace/transformers_zamba/utils/check_repo.py:376: UserWarning: Full repo consistency checks require all backends to be installed (with `pip install -e '.[dev]'` in the Transformers repo, the following are missing: TensorFlow, Flax. While it's probably fine as long as you didn't make any change in one of those backends modeling files, you should probably execute the command above to be on the safe side.
/workspace/transformers_zamba/src/transformers/models/deit/image_processing_deit.py:87: DeprecationWarning: BICUBIC is deprecated and will be removed in Pillow 10 (2023-07-01). Use Resampling.BICUBIC instead.
  resample: PILImageResampling = PIL.Image.BICUBIC,
/workspace/transformers_zamba/src/transformers/models/chameleon/image_processing_chameleon.py:116: DeprecationWarning: LANCZOS is deprecated and will be removed in Pillow 10 (2023-07-01). Use Resampling.LANCZOS instead.
  resample: PILImageResampling = PIL.Image.LANCZOS,
/workspace/transformers_zamba/src/transformers/models/efficientnet/image_processing_efficientnet.py:92: DeprecationWarning: NEAREST is deprecated and will be removed in Pillow 10 (2023-07-01). Use Resampling.NEAREST or Dither.NONE instead.
  resample: PILImageResampling = PIL.Image.NEAREST,
```

 I see most of the CircleCI tests for this PR are still failing, please let me know if more needs to be done to fix those.

Thank you so much!

I'll have a look! 

Thank you so much for the extensive review @ArthurZucker, this was extremely helpful! I believe we addressed the points raised above. 

I rebased the repository and all Circle/CI tests seem to pass except this one which seems unrelated to Zamba:

```
self = Dataset({
    features: ['image', 'labels'],
    num_rows: 50
})
task = 'image-classification', id = 0

    @deprecated()
    def prepare_for_task(self, task: Union[str, TaskTemplate], id: int = 0) -> "Dataset":
        """
        Prepare a dataset for the given task by casting the dataset's [`Features`] to standardized column names and types as detailed in [`datasets.tasks`](./task_templates).
    
        Casts [`datasets.DatasetInfo.features`] according to a task-specific schema. Intended for single-use only, so all task templates are removed from [`datasets.DatasetInfo.task_templates`] after casting.
    
        Args:
            task (`Union[str, TaskTemplate]`):
                The task to prepare the dataset for during training and evaluation. If `str`, supported tasks include:
    
                - `"text-classification"`
                - `"question-answering"`
    
                If [`TaskTemplate`], must be one of the task templates in [`datasets.tasks`](./task_templates).
            id (`int`, defaults to `0`):
                The id required to unambiguously identify the task template when multiple task templates of the same type are supported.
        """
        # TODO(lewtun): Add support for casting nested features like answers.text and answers.answer_start in SQuAD
        if isinstance(task, str):
            tasks = [template.task for template in (self.info.task_templates or [])]
            compatible_templates = [template for template in (self.info.task_templates or []) if template.task == task]
            if not compatible_templates:
>               raise ValueError(
                    f"Task {task} is not compatible with this dataset! Available tasks: {list(unique_values(tasks))}"
                )
E               ValueError: Task image-classification is not compatible with this dataset! Available tasks: []

/usr/local/lib/python3.10/site-packages/datasets/arrow_dataset.py:2814: ValueError
```

Please let us know if more needs to be done. Thank you so much again!


Just merged and fixed a minor merge issue. Should be good to review from our end @ArthurZucker since the test failures don't appear related to our changes:

**tests_torch**:
```
FAILED tests/models/dpt/test_modeling_dpt_auto_backbone.py::DPTModelTest::test_attention_outputs - AttributeError: 'NoneType' object has no attribute 'shape'
FAILED tests/models/dpt/test_modeling_dpt_auto_backbone.py::DPTModelTest::test_retain_grad_hidden_states_attentions - AttributeError: 'NoneType' object has no attribute 'retain_grad'
```

**examples_torch**:
```
/usr/local/lib/python3.10/site-packages/albumentations/augmentations/blur/functional.py:9: in <module>
    from albucore.utils import clipped, maybe_process_in_chunks, preserve_channel_dim
E   ImportError: cannot import name 'preserve_channel_dim' from 'albucore.utils' (/usr/local/lib/python3.10/site-packages/albucore/utils.py)
```

@ArthurZucker -- Paolo and I have responded or addressed your comments. We just merged again. We're ready to kick this back to you. We're getting close!

The zamba tests are passing. There appears to be an unrelated error in `tests_torch` with paligemma:

```
FAILED tests/models/paligemma/test_modeling_paligemma.py::PaliGemmaForConditionalGenerationModelTest::test_static_cache_matches_dynamic - AssertionError: False is not true
```

Just merged again. Test failures again don't appear related to us?

### tests_hub:
```
FAILED tests/trainer/test_trainer.py::TrainerIntegrationWithHubTester::test_push_to_hub_in_organization - huggingface_hub.errors.HfHubHTTPError: 500 Server Error: Internal Server Error for url: https://hub-ci.huggingface.co/valid_org/test-trainer-org.git/info/lfs/objects/batch (Request ID: Root=1-66f46931-2271680b5d9f004f31f064c1;855ff866-65d6-4d2c-9b19-a1823475f1fc)

Cannot read properties of undefined (reading 'storageClass') (internal error hidden in production)
FAILED tests/trainer/test_trainer.py::TrainerIntegrationWithHubTester::test_push_to_hub_tags - huggingface_hub.errors.HfHubHTTPError: 500 Server Error: Internal Server Error for url: https://hub-ci.huggingface.co/__DUMMY_TRANSFORMERS_USER__/test-trainer-tags.git/info/lfs/objects/batch (Request ID: Root=1-66f46933-59f3b6205417227376459499;ff4fe7f7-ebbe-4c51-9719-a4b7e825918a)

Cannot read properties of undefined (reading 'storageClass') (internal error hidden in production)
FAILED tests/trainer/test_trainer.py::TrainerIntegrationWithHubTester::test_push_to_hub_with_revision - huggingface_hub.errors.HfHubHTTPError: 500 Server Error: Internal Server Error for url: https://hub-ci.huggingface.co/__DUMMY_TRANSFORMERS_USER__/test-trainer-revision.git/info/lfs/objects/batch (Request ID: Root=1-66f46933-53e3df2636b2308d6961e76b;a80b6a39-d5e8-4c88-817b-393c5581fc7f)

Cannot read properties of undefined (reading 'storageClass') (internal error hidden in production)
FAILED tests/trainer/test_trainer.py::TrainerIntegrationWithHubTester::test_push_to_hub_with_saves_each_epoch - AssertionError: no logs of level WARNING or higher triggered on root
======= 4 failed, 34 passed, 14 skipped, 17 warnings in 63.98s (0:01:03) =======
```

### tests_torch:
This is a timeout when loading the zamba-7B-v1 model. I don't get the issue on my local setup so I think it's an issue with the CI hardware not keeping up. Do we need to increase the timeout for a model of this size, or should this just be ignored?

```
self = <ssl.SSLSocket [closed] fd=-1, family=AddressFamily.AF_INET, type=SocketKind.SOCK_STREAM, proto=6>
len = 188416, buffer = <memory at 0x7fe605054880>

    def read(self, len=1024, buffer=None):
        """Read up to LEN bytes and return them.
        Return zero-length string on EOF."""
    
        self._checkClosed()
        if self._sslobj is None:
            raise ValueError("Read on closed or unwrapped SSL socket.")
        try:
            if buffer is not None:
>               return self._sslobj.read(len, buffer)
E               Failed: Timeout >120.0s

/usr/local/lib/python3.10/ssl.py:1163: Failed
```

### tests_generate:
Looks related to the llama3.2 release today. I was able to get the same error with the latest `huggingface:main`, so this is also not due to our changes:
```
FAILED tests/models/mllama/test_modeling_mllama.py::MllamaForConditionalGenerationModelTest::test_new_cache_format_0 - RuntimeError: The size of tensor a (1808) must match the size of tensor b (904) at non-singleton dimension 3
```

Yep don't worry about unrelated ones!

Merged again. Now the only failing CI test is the 7B loading timeout. Still not reproducible on my local machine, so it's likely a CI hardware issue (either remove the 7B or increase the timeout?):

```
self = <ssl.SSLSocket [closed] fd=-1, family=AddressFamily.AF_INET, type=SocketKind.SOCK_STREAM, proto=6>
len = 6365184, buffer = <memory at 0x7f87589ec580>

    def read(self, len=1024, buffer=None):
        """Read up to LEN bytes and return them.
        Return zero-length string on EOF."""
    
        self._checkClosed()
        if self._sslobj is None:
            raise ValueError("Read on closed or unwrapped SSL socket.")
        try:
            if buffer is not None:
>               return self._sslobj.read(len, buffer)
E               Failed: Timeout >120.0s

/usr/local/lib/python3.10/ssl.py:1163: Failed
```

Yeah 7B tests should be wrapped with `@slow` as we only run fast tests on the CIs!

@ArthurZucker yes, we were able to run the slow tests in https://github.com/Zyphra/transformers_zamba/blob/d0c1bc102b6f9b5e078a53b477e96e7729703762/tests/models/zamba/test_modeling_zamba.py (`test_simple_generate` and `test_simple_batched_generate_with_padding`) and they passed. We didn't run the flash attention-related tests (`test_flash_attn_2_fp32_ln`, `test_flash_attn_2_generate_padding_right`, `test_flash_attn_2_generate_use_cache`, `test_flash_attn_2_inference_equivalence_right_padding`) as these do not apply.

@ArthurZucker -- Green CI 👀 

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_30950). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Hey! 🤗 Thanks for your contribution to the `transformers` library!

Before merging this pull request, slow tests CI should be triggered. To enable this:
- Add the `run-slow` label to the PR
- When your PR is ready for merge and all reviewers' comments have been addressed, push an empty commit with the command `[run-slow]` followed by a comma separated list of all the models to be tested, i.e. `[run_slow] model_to_test_1, model_to_test_2`
  - If the pull request affects a lot of models, put at most 10 models in the commit message
- A `transformers` maintainer will then approve the workflow to start the tests

(For maintainers) The documentation for slow tests CI on PRs is [here](https://www.notion.so/huggingface2/CI-for-pull-requests-8335bd9217d24d1e8ed1d3a016f39804).

Does this include support for Zamba2?

@hg0428 thanks for asking! Support for Zamba2 will be added in a follow-up PR. Meanwhile, you can install Zyphra's local transformers as described in the Zamba2's [model card](https://huggingface.co/Zyphra/Zamba2-2.7B-instruct).

> @hg0428 thanks for asking! Support for Zamba2 will be added in a follow-up PR. Meanwhile, you can install Zyphra's local transformers as described in the Zamba2's [model card](https://huggingface.co/Zyphra/Zamba2-2.7B-instruct).

Unfortunately, that does not work on my device. Zamba2 transformers runs on mamba_ssm, which requires an NVIDIA GPU. I have Apple Silicon. See my issue: https://github.com/Zyphra/transformers_zamba2/issues/3

Just waiting for https://github.com/huggingface/transformers/actions/runs/11116137341/job/30897696145?pr=30950#step:12:64 to be fixed! (related to accelerate and auto device, good that we have this test!)

Hi @arthur, thank you again for reviewing.

The test mentioned above `test_multi_gpu_data_parallel_forward` now passes. We had to change some of the shared layers logic for it to work. Previously, `self.mamba_layers` and `self.linear_layers` were both `nn.ModuleList` objects and `self.layers` was not, which prevented most of the layers from being scattered across devices. Now only `self.layers` is `nn.ModuleList` and everything seems to work.

Additionally, we updated all the model's checkpoints on the hub since this involved changing some of the weight keys related to the shared layers. Separately, given we updated the checkpoints, we also swapped `up`<->`gate` in the MLP weight keys as well as in the [forward pass](https://github.com/huggingface/transformers/blob/1504774123b7d56b64334593f58220f14acbf4a8/src/transformers/models/zamba/modeling_zamba.py#L860) so [this issue](https://github.com/huggingface/transformers/pull/30950#discussion_r1770715099) is now addressed.

All tests related to zamba appear to pass. Thank you!

> Hi @arthur, thank you again for reviewing.
> 
> 
> 
> The test mentioned above `test_multi_gpu_data_parallel_forward` now passes. We had to change some of the shared layers logic for it to work. Previously, `self.mamba_layers` and `self.linear_layers` were both `nn.ModuleList` objects and `self.layers` was not, which prevented most of the layers from being scattered across devices. Now only `self.layers` is `nn.ModuleList` and everything seems to work.
> 
> 
> 
> Additionally, we updated all the model's checkpoints on the hub since this involved changing some of the weight keys related to the shared layers. Separately, given we updated the checkpoints, we also swapped `up`<->`gate` in the MLP weight keys as well as in the [forward pass](https://github.com/huggingface/transformers/blob/1504774123b7d56b64334593f58220f14acbf4a8/src/transformers/models/zamba/modeling_zamba.py#L860) so [this issue](https://github.com/huggingface/transformers/pull/30950#discussion_r1770715099) is now addressed.
> 
> 
> 
> All tests related to zamba appear to pass. Thank you!

Does this Zamba support work on Apple Silicon?

> > Hi @Arthur, thank you again for reviewing.
> > The test mentioned above `test_multi_gpu_data_parallel_forward` now passes. We had to change some of the shared layers logic for it to work. Previously, `self.mamba_layers` and `self.linear_layers` were both `nn.ModuleList` objects and `self.layers` was not, which prevented most of the layers from being scattered across devices. Now only `self.layers` is `nn.ModuleList` and everything seems to work.
> > Additionally, we updated all the model's checkpoints on the hub since this involved changing some of the weight keys related to the shared layers. Separately, given we updated the checkpoints, we also swapped `up`<->`gate` in the MLP weight keys as well as in the [forward pass](https://github.com/huggingface/transformers/blob/1504774123b7d56b64334593f58220f14acbf4a8/src/transformers/models/zamba/modeling_zamba.py#L860) so [this issue](https://github.com/huggingface/transformers/pull/30950#discussion_r1770715099) is now addressed.
> > All tests related to zamba appear to pass. Thank you!
> 
> Does this Zamba support work on Apple Silicon?

I don't believe so. We're working on MLX support in a separate (private for now) vein of work from this PR, which just seeks to get basic GPU integration into upstream HuggingFace Transformers.

🎉 🎉 🎉 🎉 🎉 🎉 🎉 🎉 🎉 🎉 

🚀 

Hi,
Are there any plans to add Zamba2 to Transformers?
Thanks!

I think the Zyphra team is already working on it!

Hopefully we get Apple Silicon support for Zamba and Zamba2 soon.