Repository: huggingface/transformers
Pull request: Fixing bug in Voxtral when merging text and audio embeddings

# What does this PR do?

This PR addresses two issues arising from an [in-place operation where audio and text embeddings are merged](https://github.com/huggingface/transformers/blob/main/src/transformers/models/voxtral/modeling_voxtral.py#L512) in the Voxtral model. This fixes [Issue #40488](https://github.com/huggingface/transformers/issues/40488) and an unreported issue resulting from using `device_map="auto"` with the Voxtral model.

More detail about the two issues and their resolution:
- In the issue reported in [#40488](https://github.com/huggingface/transformers/issues/40488), the `forward` method of `VoxtralForConditionalGeneration` fails when using LoRA. The underlying issue is that when the text embedding layer is frozen, the `inputs_embeds` tensor extracted from the embedding layer is a leaf tensor. When `inputs_embeds` is a leaf tensor that requires gradients, values of this tensor cannot be reassigned. To address this, `inputs_embeds` will be now be cloned to enable tracking of gradients when it is a leaf tensor that requires gradients.
- The second unreported issue arises when using `device_map="auto"` with Voxtral. When using `device_map="auto"`, audio and text layers might be distributed across different devices. When this is the case, `inputs_embeds` and `audio_embeds` might be on different devices. In this case, attempting to assign values of `audio_embeds` to `inputs_embeds` will fail since both tensors are expected to be on the same device. To address this, `audio_embeds` is moved to the same device as `inputs_embeds` before updating values of `inputs_embeds`.


## 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?
- [x] 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. -> https://github.com/huggingface/transformers/issues/40488
- [ ] Did you make sure to update the documentation with your changes? -> No documentation changes necessary.
- [ ] Did you write any new necessary tests? -> No new features added.


## Who can review?

@eustlb 

Changed files:
filename: src/transformers/models/voxtral/modeling_voxtral.py
patch: @@ -504,12 +504,14 @@ def forward(
         if inputs_embeds is None:
             inputs_embeds = self.get_input_embeddings()(input_ids)
 
-        if input_features is not None:
+        if input_features is not None and input_ids is not None:
             audio_embeds = self.get_audio_embeds(input_features)
 
             # replace text-audio token placeholders with audio embeddings
-            audio_token_mask = input_ids == self.config.audio_token_id
-            inputs_embeds[audio_token_mask] = audio_embeds
+            audio_token_mask = (input_ids == self.config.audio_token_id).unsqueeze(-1)
+            inputs_embeds = inputs_embeds.masked_scatter(
+                audio_token_mask.to(inputs_embeds.device), audio_embeds.to(inputs_embeds.device)
+            )
 
         outputs: BaseModelOutputWithPast = self.language_model(
             attention_mask=attention_mask,

filename: src/transformers/models/voxtral/modular_voxtral.py
patch: @@ -239,12 +239,14 @@ def forward(
         if inputs_embeds is None:
             inputs_embeds = self.get_input_embeddings()(input_ids)
 
-        if input_features is not None:
+        if input_features is not None and input_ids is not None:
             audio_embeds = self.get_audio_embeds(input_features)
 
             # replace text-audio token placeholders with audio embeddings
-            audio_token_mask = input_ids == self.config.audio_token_id
-            inputs_embeds[audio_token_mask] = audio_embeds
+            audio_token_mask = (input_ids == self.config.audio_token_id).unsqueeze(-1)
+            inputs_embeds = inputs_embeds.masked_scatter(
+                audio_token_mask.to(inputs_embeds.device), audio_embeds.to(inputs_embeds.device)
+            )
 
         outputs: BaseModelOutputWithPast = self.language_model(
             attention_mask=attention_mask,