Use Conv1d for TDNN
Author: gau-nernst
Pull request: yes

# What does this PR do?

Partially fixes #25476 

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Fixes # (issue)


## 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](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      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?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker and @younesbelkada
- vision models: @amyeroberts
- speech models: @sanchit-gandhi
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @gante
- pipelines: @Narsil
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerz and @pacman100

Integrations:

- deepspeed: HF Trainer/Accelerate: @pacman100
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc and @younesbelkada

Documentation: @stevhliu and @MKhalusova

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Comments:
cc @sanchit-gandhi 

On RTX 3090, 5s input

dtype | batch size | type | de139702a | 4ad9e29a5
--|--|--|--|--
torch.float32 | 1 | forward | 61.8698 it/s | 65.4886 it/s
torch.float32 | 1 | forward+backward | 15.7255 it/s | 15.3732 it/s
torch.float32 | 8 | forward | 12.1235 it/s | 13.3573 it/s
torch.float32 | 8 | forward+backward | 4.3197 it/s | 4.4702 it/s
torch.float16 | 1 | forward | 70.7141 it/s | 73.2518 it/s
torch.float16 | 1 | forward+backward | 17.1505 it/s | 17.1628 it/s
torch.float16 | 8 | forward | 20.3633 it/s | 24.0503 it/s
torch.float16 | 8 | forward+backward | 7.5402 it/s | 8.0300 it/s

The improvement is around 5%. It's small but not insignificant.

Another optimization we can do is to make `TDNNLayer` accepts channels-first input instead of channels-last (current approach) to avoid transposing data twice every layer. However, this would break backward compatibility if someone relies on `TDNNLayer` externally (though I think it's unlikely).

Benchmark script

```python
import time
from itertools import product
import torch
from transformers.models.wav2vec2.modeling_wav2vec2 import Wav2Vec2ForXVector

torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.backends.cudnn.benchmark = True

device = "cuda"
batch_sizes = (1, 8)
dtypes = (torch.float32, torch.float16)

m = Wav2Vec2ForXVector.from_pretrained("facebook/wav2vec2-xls-r-300m").to(device)

for dtype, bsize in product(dtypes, batch_sizes):
    m.to(dtype)
    x = torch.randn(bsize, 16_000 * 5, device=device, dtype=dtype)

    m.eval()
    with torch.no_grad():
        m(x)  # warmup
    torch.cuda.synchronize()

    N = 100

    with torch.no_grad():
        time0 = time.perf_counter()
        for _ in range(N):
            out = m(x)
            torch.cuda.synchronize()
    print(f"{dtype}, {bsize=}, forward: {N / (time.perf_counter() - time0):.4f} it/s")

    m.train()
    m(x)[0].sum().backward()  # warmup

    time0 = time.perf_counter()
    for _ in range(N):
        out = m(x)
        out[0].sum().backward()
        torch.cuda.synchronize()
    print(f"{dtype}, {bsize=}, forward+backward: {N / (time.perf_counter() - time0):.4f} it/s")
```

Very cool! Thanks for the results @gau-nernst and nice benchmark script! I think it's worth pursuing in this case: 5% is the lower end of what we'd expect to get for `torch.compile`, and this has proven to be quite high demand for optimising PyTorch models, so a 5% improvement gain would be welcomed by the Wav2Vec2 community

I would be against adding a breaking change with the channel-last approach, since this is in contradiction to the scope of the PR where we set out to optimise in an entirely non-breaking way

Cool! Anything else you want me to add to the PR, perhaps apart from removing the commented out old code? Is there a correctness test that we can add?

I didn't benchmark with torch.compile() since Wav2Vec2 did not work with it yet (I hope the relevant PR will be merged soon). Would you want to see benchmark results with torch.compile() for TDNNLayer alone? Perhaps it can even optimize the old way (unfold + linear)

It looks super clean already @gau-nernst - I think all that's left to do is:
1. Remove the commented out code
2. Run the slow tests to check we still have equivalence (this will check for correctness): https://github.com/huggingface/transformers/blob/1c6f072db0c17c7d82bb0d3b7529d57ebc9a0f2f/tests/models/wav2vec2/test_modeling_wav2vec2.py#L1922

No worries about benchmarking with torch compile - the results you provided previously were more than enough to justify this code change. We won't expect all users to be using PT 2.0, so this speed-up will definitely hold in many cases

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_25728). All of your documentation changes will be reflected on that endpoint.

I think it's unlikely that someone will apply LoRA on the TDNN layers. Most likely they will apply LoRA on the transformer backbone only, and train the TDNN layers from scratch.

My question is rather, what if someone has already apply LoRA on the TDNN layers? (Though it might not be recommended, someone who does not specifically know the model could have just quantized all the linear layers no?) 😉 

I can't answer that, though I doubt that people who know how to use LoRA, and use it on TDNN layers specifically (most LoRA guides only show how to use it with transformers as far as I know), will face problems with this change. You guys from HF should decide whether this outweighs the cost i.e. speed-up vs annoy people who use LoRA on TDNN layers, and what kind of guarantee the library should provide i.e. a model surgery technique (e.g. LoRA) should still work after an update.

That is true: LORA will now not apply to the linear layers. Just to clarify, is this just concerned with applying lora weights to the full-precision model and doing a new fine-tuning run? Or, would it actually break inference if someone wants to use lora fine-tuned weights with the TDNN model? If it's the former, I think it's fine since TDNN is not documented and the linear layers are a fraction of the transformer ones. If it's the latter, then we maybe need to think of a workaround

Pretty sure it will break inference, but we can probably have some kind of fix / deprecation cycle for this 😓 

Great work here on finding an elegant solution 🤗 

@ArthurZucker i have merged master and make some small changes accordingly

Thanks! Merging 🤗