Fixing bug in Voxtral when merging text and audio embeddings

[rcogill]

What does this PR do?

This PR addresses two issues arising from an in-place operation where audio and text embeddings are merged in the Voxtral model. This fixes Issue #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, 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).
• Did you read the contributor guideline, Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case. -> Voxtral model fails with LoRA due to in-place operation error #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

[vasqu]

Not a fan of using a clone tbh, can we use a masked scatter here instead, e.g. see

transformers/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py

Line 1206 in 91b34be

inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)

[vasqu]

Should be out-of-place op then, so I would hope it resolves the inplace operation issue

(probably still needs the device movement)

[eustlb]

Thanks @rcogill for working on this and @vasqu for first review! Quick benchmarks show non-inplace masked_scatter is faster anyway. Can you add the suggested change and good to merge then.

[eustlb]

Agree here with @vasqu, let's not forget also to move the mask to the correct device and make it's broadcastable!
Can you please change to:

# replace text-audio token placeholders with audio embeddings
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))

[eustlb]

pushed directly to shortcut and merge! Thanks again, @vasqu and @rcogill

[github-actions]

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

run-slow: voxtral

[vasqu]

Adding a for patch labels here @eustlb? And glad to help ❤️

[eustlb]

Actually no, for-patch label is for what has been broken in last release, which is not the case here ;)

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

[rcogill]

@eustlb and @vasqu , thank you both!