fix onnx export of speech foundation models
Author: nikosanto13
Pull request: yes

# What does this PR do?

This old issue https://github.com/huggingface/transformers/issues/10004 described an error in the onnx export of wav2vec2-base-960h.
The issue was partially fixed for some of the speech foundation models in this PR: https://github.com/huggingface/transformers/pull/16004.

This PR:
* applies the same change of expanding the attention mask in all modeling scripts that this occurs (e.g. I came across this when trying to export `wavlm-base-plus` (modeling_wavlm.py) to onnx),
* applies the same change of expanding the downsampled padding mask in the <Wav2Vec2, Hubert, ...>ForSequenceClassification modules' forward implementation, because the onnx export of these fails similarly, due to the same broadcasting-related error

## 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.
- [ ] 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?

Tagging @xenova, @ylacombe, @eustlb  (let me know if I should also add someone else)

Comments:
Maybe also cc @xenova for ONNX :)

@LysandreJik nobody responded so far, should I cc someone else?

The best suited to review is @ylacombe I believe; let's give Yoach a few days to answer, thanks for your PR @nikosanto13!

ok nice, thanks for the help @LysandreJik 

thanks for mentioning @ylacombe - I'll rebase and add an empty commit for the slow tests.

For your comment about this being a costly operation, I admit I pretty much adopted the same patch (with the previous pr), without giving it much consideration. My initial thought is that this shouldn't be much of a concern (given that the masks are boolean for the majority of cases).

Let me conduct a mini benchmark and I'll return with actual numbers.

> Let me conduct a mini benchmark and I'll return with actual numbers.

Great, looking forward to it

@ylacombe I conducted a mini benchmark using WavLM (`microsoft/wavlm-base`).

System info: A10 GPU (24GB VRAM)

Details on the benchmark: I run WavLM's forward pass for both train and eval modes. In the code below, you can see the details - a batch (batch_size = 16) with 10-sec audios is generated and fed through the model.

I created the following script:
```python

from contextlib import nullcontext

import numpy as np
import torch
from torch.autograd import profiler

from transformers import AutoFeatureExtractor, WavLMModel


# taken from torch.autograd.profiler
def _format_time(time_us):
    """Define how to format time in FunctionEvent."""
    US_IN_SECOND = 1000.0 * 1000.0
    US_IN_MS = 1000.0
    if time_us >= US_IN_SECOND:
        return f"{time_us / US_IN_SECOND:.3f}s"
    if time_us >= US_IN_MS:
        return f"{time_us / US_IN_MS:.3f}ms"
    return f"{time_us:.3f}us"

# taken from torch.autograd.profiler
def _format_memory(nbytes):
    """Return a formatted memory size string."""
    KB = 1024
    MB = 1024 * KB
    GB = 1024 * MB
    if abs(nbytes) >= GB:
        return f"{nbytes * 1.0 / GB:.2f} Gb"
    elif abs(nbytes) >= MB:
        return f"{nbytes * 1.0 / MB:.2f} Mb"
    elif abs(nbytes) >= KB:
        return f"{nbytes * 1.0 / KB:.2f} Kb"
    else:
        return str(nbytes) + " b"


def run_forward(model, feat_extractor, train=False):

    if train:
        model.train()
    else:
        model.eval()

    ctx = torch.no_grad() if not train else nullcontext()

    input_values = torch.randn(16, 10*feat.sampling_rate).to(device)
    attention_mask = torch.ones_like(input_values, dtype=torch.bool)

    with ctx, profiler.profile(use_cuda='cuda', profile_memory=True) as prof:
        output = model(input_values, attention_mask=attention_mask)

    events_list = prof.key_averages()
    mask_event = [elem for elem in events_list if elem.key == 'MASK_HIDDEN_STATES'][0]

    print(mask_event)
    return mask_event.device_time, mask_event.device_memory_usage


if __name__ == "__main__":

    device = 'cuda'
    model = WavLMModel.from_pretrained("microsoft/wavlm-base").to(device)
    feat = AutoFeatureExtractor.from_pretrained("microsoft/wavlm-base")

    times, mems = [], []
    for i in range(10):
        time, mem = run_forward(model, feat)
        # skipping first iteration due to additional overhead
        if i > 0:
            times.append(time)
            mems.append(mem)

    avg_time, std_time = np.mean(times), np.std(times)
    avg_mem, std_mem = np.mean(mems), np.std(mems)

    print(f"Average time: {_format_time(avg_time)} ± {_format_time(std_time)}")
    print(f"Average memory: {_format_memory(avg_mem)} ± {_format_memory(std_mem)}")
```

**NOTE**: In addition to that, you also have to wrap the masking operation with the `profiler.record_function` context manager. E.g. in modeling_wavlm.py and the `WavLMEncoder` forward, I replaced this

```python
if attention_mask is not None:
       expand_attention_mask = attention_mask.unsqueeze(-1).repeat(1, 1, hidden_states.shape[2])
       hidden_states[~expand_attention_mask] = 0
```

with this:

```python
# don't forget to import torch.autograd.profiler
# from torch.autograd import profiler
if attention_mask is not None:
            with profiler.record_function("MASK_HIDDEN_STATES"):
                expand_attention_mask = attention_mask.unsqueeze(-1).repeat(1, 1, hidden_states.shape[2])
                hidden_states[~expand_attention_mask] = 0
```


-----

The following measurements were obtained:

| branch/mode |  Time  | CUDA Memory Consumption |
|--------------|--------| ----------------------------- |
| main/train | 278.000us ± 7.333us  |  8.00 Kb ± 0.0 b  |
| onnx-export-bugfix-where-nodes/train | 435.222us ± 12.417us | 11.70 Mb ± 0.0 b |
| main/eval   |  145.444us ± 1.066us  |  0.0 b ± 0.0 b  |
| onnx-export-bugfix-where-nodes/eval |  251.333us ± 3.127us |  5.85 Mb ± 0.0 b | 





Hey @nikosanto13, thanks for conducting the benchmark! I'm not sure to understand what you benchmarked though. Did you compare the speed of this new operation as compared to how it was done previously?

Hey @ylacombe - well, based on your message I figured I should benchmark both speed and memory consumption.

As you can see in the attached code snippet I'm comparing those two measurements for a forward pass of the WavLM model, before (main) and after the changes (my branch).

Do you think of a more appropriate way to benchmark this? I'd be happy to extend this.

Thanks @nikosanto13, it's much clearer now. Looks like it's not a costly operation. This PR looks great to me.

cc @ArthurZucker  or @Rocketknight1 for a core maintainer review!