Repository: huggingface/transformers
Pull request: Granite speech speedups

# What does this PR do?
This PR speeds up Granite speech by the following changes:
1. Register attention_dist to avoid cpu-to-gpu transfer every layer.
2. `pad_sequence` is much faster than per-sample-padding + concat.
3. Avoid returning audio to cpu when using a compute device.


## 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.
- [ ] 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).
- [ ] Did you write any new necessary tests?


## Who can review?

@ArthurZucker @eustlb can you give that a look 🙏 
CC: @gsaon @alex-jw-brooks @avishaiElmakies 

Changed files:
filename: src/transformers/models/granite_speech/feature_extraction_granite_speech.py
patch: @@ -117,8 +117,6 @@ def _extract_mel_spectrograms(self, audio: "torch.Tensor", device="cpu"):
             # stacking and skipping by 2
             audio = logmel.reshape(bsz, -1, 2 * logmel.shape[-1])
 
-        if audio.device != "cpu":
-            return audio.detach().cpu()
         return audio
 
     def _get_num_audio_features(self, audio_lengths: Sequence[int]) -> Sequence[int]:
@@ -178,11 +176,8 @@ def _get_audios_and_audio_lengths(self, audios: AudioInput) -> Sequence["torch.T
             if not torch.is_floating_point(audios[0]):
                 raise ValueError("Invalid audio provided. Audio should be a floating point between 0 and 1")
             lengths = [audio.shape[-1] for audio in audios]
-            padding = [max(lengths) - length for length in lengths]
-            # ensure all audios have a batch dimension:
-            audios = [audio.view(1, -1) for audio in audios]
-            padded = [torch.nn.functional.pad(audio, (0, pad)) for audio, pad in zip(audios, padding)]
-            audios = torch.cat(padded, dim=0)
+            audios = [audio.squeeze(0) for audio in audios]
+            audios = torch.nn.utils.rnn.pad_sequence(audios, batch_first=True, padding_value=0.0)
             return audios, lengths
 
         raise TypeError("Invalid audio provided. Audio should be a one or more torch tensors or numpy arrays")

filename: src/transformers/models/granite_speech/modeling_granite_speech.py
patch: @@ -157,8 +157,7 @@ def forward(self, hidden_states: torch.Tensor, attention_dists: torch.Tensor) ->
         value_states = value_states.reshape(bsz, num_blocks, self.context_size, self.num_heads, -1).transpose(2, 3)
 
         # shaw's relative positional embedding
-        dist = attention_dists.to(hidden_states.device)
-        rel_pos_emb = self.rel_pos_emb(dist)
+        rel_pos_emb = self.rel_pos_emb(attention_dists)
         # alternative computation of `pos_attn` - for readability
         # rel_pos_emb_expanded = rel_pos_emb.view([1, 1, 1] + list(rel_pos_emb.shape))
         # pos_attn = torch.sum(query_states.unsqueeze(-2) * rel_pos_emb_expanded, dim=-1) * self.scale
@@ -258,8 +257,8 @@ def __init__(self, config: GraniteSpeechEncoderConfig):
         # Precompute clamped relative positional encoding distances
         seq = torch.arange(config.context_size)
         relpos_dist = seq.view(-1, 1) - seq.view(1, -1)
-        self.attention_dists = torch.clamp(relpos_dist, -config.context_size, config.context_size) + config.max_pos_emb
-
+        attention_dists = torch.clamp(relpos_dist, -config.context_size, config.context_size) + config.max_pos_emb
+        self.register_buffer("attention_dists", attention_dists, persistent=False)
         self.input_linear = nn.Linear(config.input_dim, config.hidden_dim, bias=True)
         self.layers = nn.ModuleList([GraniteSpeechConformerBlock(config) for _ in range(config.num_layers)])