Moshi integration
Author: ylacombe
Pull request: yes
Labels: name: run-slow

# What does this PR do?

Moshi is the latest Kyutai model. It is a streaming speech-to-speech model, that can also do an inner dialogue (i.e it outputs text as well).

In particular, it means that Moshi deals with 3 streams of information:
1. The user's audio
2. Moshi's audio
3. Moshi's textual output

Similarly to `Musicgen`, audio is represented with audio codebooks, which can be interpreted like tokens. The main difference between text tokens and audio codebooks is that audio codebooks introduce an additional dimension of information.
Text tokens are typically of dim `(batch_size, sequence_length)` but audio tokens are of dim `(batch_size, num_codebooks, sequence_length)`.
![image](https://github.com/user-attachments/assets/e790311d-11dc-4292-936e-51683c2b3fe4)

--------

It's made of 3 components:

**1. The main decoder (Helium in the paper)**

Here, it corresponds to `MoshiForCausalLM`. It is strictly a classic text LLM, that uses an architecture similar to `Gemma`. In other words, it takes text tokens, embeds them, pass them through the decoder and a language head, to get text logits.

**2. The depth decoder**

On its own, it's also a classic LLM, but this time, instead of generating over the time dimension, it generates over the codebook dimension.

It also means that its context length is `num_codebooks` -> it can't generate more than `num_codebooks`.

Another interesting difference with a classic LLM is that each timestamp (here it correspond to each codebook) got its own set of Linear Layers and Embeddings.

**3. Mimi**

It's the audio encoder from Kyutai, that has recently been integrated to transformers, which is used to "tokenize" audio. It has the same use that `Encodec` has in `Musicgen`.

--------

## Architecture choice:

1. `MoshiForCausalLM` corresponds to the main decoder, it can be used as a textual LLM.


2. `MoshiDepthDecoder` is the depth decoder mentioned above


3. `MoshiForConditionalGeneration` encapsulates the main decoder, the depth decoder and the audio encoder.

Conceptually, `MoshiForConditionalGeneration` takes as input one stream of text and two streams of audio inputs - what the user has said so far, and what the model have generated so far - and generates two streams - a text stream and an audio stream.

**How does it work:**

-> The input streams are embedded and combined into `inputs_embeds`.

-> `inputs_embeds` is passed through the main decoder. There's nothing special done here, it's the same operation as Gemma or so on. 

-> The main decoder outputs `text logits` but also its `last hidden state` which is called `temporal context` in the picture above. 

-> the depth decoder switches the dimension on which we generate (codebooks instead of time). It uses the token generated from `text logits`  and the `temporal context` to auto-regressively generate audio codebooks.





Comments:
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).

cc @gante and @ArthurZucker , there's still a few failing tests, but I think it'd be good to have a first review, given the model complexity 

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

For the remaining failures: 

- quality issue seems linked to `VipLlavaForConditionalGeneration`

- tokenization tests - cc @ArthurZucker  - do you think it's linked to our offline discussion ?
<details><summary>Details tokenization</summary>
<p>
FAILED tests/models/moshi/test_tokenization_moshi.py::MoshiTokenizationTest::test_added_tokens_do_lower_case - AssertionError: '▁' == '▁'
FAILED tests/models/moshi/test_tokenization_moshi.py::MoshiTokenizationTest::test_alignement_methods - AssertionError: 0 != 5
FAILED tests/models/moshi/test_tokenization_moshi.py::MoshiTokenizationTest::test_training_new_tokenizer - AssertionError: 'This▁is▁the▁first▁sentence' != 'This is the first sentence'
FAILED tests/models/moshi/test_tokenization_moshi.py::MoshiTokenizationTest::test_training_new_tokenizer_with_special_tokens_change - AssertionError: 'This▁is▁the▁first▁sentence' != 'This is the first sentence'
</p>
</details> 

- sdpa and eager attention mismatch - I'm not really sure where it's coming from: the model (`MoshiDecoderTest`) is almost the same as `Gemma`

<details><summary>Details mismatch</summary>
<p>
FAILED tests/models/moshi/test_modeling_moshi.py::MoshiDecoderTest::test_eager_matches_sdpa_inference_0_float16 - AssertionError: False is not true : padding_side=left, use_mask=False, batch_size=1, enable_kernels=False: mean relative difference: 1.938e+00, torch atol = 0.005, torch rtol = 0.005
FAILED tests/models/moshi/test_modeling_moshi.py::MoshiDecoderTest::test_eager_matches_sdpa_inference_1_bfloat16 - AssertionError: False is not true : padding_side=left, use_mask=False, batch_size=1, enable_kernels=False: mean relative difference: 8.164e-01, torch atol = 0.01, torch rtol = 0.01
FAILED tests/models/moshi/test_modeling_moshi.py::MoshiDecoderTest::test_eager_matches_sdpa_inference_2_float32 - AssertionError: False is not true : padding_side=left, use_mask=False, batch_size=1, enable_kernels=False: mean relative difference: 6.824e-01, torch atol = 1e-06, torch rtol = 0....
</p>
</details> 



Hey @SunMarc, we could benefit from your help here: When using device_map="auto", I'm facing a lot of issues, see [here](https://github.com/huggingface/transformers/actions/runs/11294837205/job/31416241230?pr=33624) and [here](https://github.com/huggingface/transformers/actions/runs/11294837205/job/31416241531?pr=33624). Everything passes locally though.


Everything's green, I'll merge
Thanks @ArthurZucker and @SunMarc for the support!