Repository: huggingface/transformers
Pull request: fix tensor device placement issue of 2 UT cases

If run below 2 cases on 2 accelerators:

> pytest -rA tests/models/speech_to_text/test_modeling_speech_to_text.py::Speech2TextModelIntegrationTests::test_generation_librispeech

>pytest -rA tests/models/speech_to_text/test_modeling_speech_to_text.py::Speech2TextModelIntegrationTests::test_generation_librispeech_batched

Will fail w/ error as below:

> self = Speech2TextEncoderLayer(
>   (self_attn): Speech2TextAttention(
>     (k_proj): Linear(in_features=256, out_features=256, ...atures=2048, out_features=256, bias=True)
>   (final_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
> )
> hidden_states = tensor([[[-313.5520,  -90.0677,   48.3815,  ...,  -22.3849,  -48.9418,
>            -49.9119],
>          [-285.8702,  -94...2.9486],
>          [  25.0123,  -37.5042,   13.0347,  ...,  -58.4456,  -16.1031,
>             45.5035]]], device='xpu:1')
> attention_mask = tensor([[[[0., 0., 0.,  ..., 0., 0., 0.],
>           [0., 0., 0.,  ..., 0., 0., 0.],
>           [0., 0., 0.,  ..., 0., 0....., 0., 0., 0.],
>           [0., 0., 0.,  ..., 0., 0., 0.],
>           [0., 0., 0.,  ..., 0., 0., 0.]]]], device='xpu:0')
> output_attentions = False
> 
>     def forward(
>         self,
>         hidden_states: torch.Tensor,
>         attention_mask: torch.Tensor,
>         output_attentions: bool = False,
>     ) -> torch.Tensor:
>         """
>         Args:
>             hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
>             attention_mask (`torch.FloatTensor`): attention mask of size
>                 `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
>             output_attentions (`bool`, *optional*):
>                 Whether or not to return the attentions tensors of all attention layers. See `attentions` under
>                 returned tensors for more detail.
>         """
>         residual = hidden_states
>         hidden_states = self.self_attn_layer_norm(hidden_states)
>         hidden_states, attn_weights = self.self_attn(
>             hidden_states=hidden_states,
>             attention_mask=attention_mask,
>             output_attentions=output_attentions,
>         )
>         hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
>         hidden_states = residual + hidden_states
> 
>         residual = hidden_states
>         hidden_states = self.final_layer_norm(hidden_states)
>         hidden_states = self.activation_fn(self.fc1(hidden_states))
>         hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
>         hidden_states = self.fc2(hidden_states)
>         hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
>        hidden_states = residual + hidden_states
>                         ^^^^^^^^^^^^^^^^^^^^^^^^
> E       RuntimeError: Expected all tensors to be on the same device, but found at least two devices, xpu:0 and xpu:1!
> 
> src/transformers/models/speech_to_text/modeling_speech_to_text.py:369: RuntimeError

This PR fixed the issue.


Changed files:
filename: src/transformers/models/speech_to_text/modeling_speech_to_text.py
patch: @@ -544,6 +544,8 @@ class Speech2TextEncoder(Speech2TextPreTrainedModel):
         embed_tokens (nn.Embedding): output embedding
     """
 
+    _no_split_modules = ["Speech2TextEncoderLayer"]
+
     def __init__(self, config: Speech2TextConfig):
         super().__init__(config)