Use Conv1d for TDNN

[gau-nernst]

What does this PR do?

Partially fixes #25476

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,
  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.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• 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.

[ArthurZucker]

cc @sanchit-gandhi

[sanchit-gandhi]

Very clean, like this solution a lot @gau-nernst! Do you know what kind of speed-ups we get already?

[gau-nernst]

I will try to run some benchmarks and let you know.
Since I'm also working on LoRA at my work, I realize doing this will bypass any "magic" to be applied on nn.Linear, since we don't call nn.Linear.forward(). Although LoRA is not likedly to be applied here, since the TDNN layers will most likely be trained from scratch, it may lead to unexpected behavior if the user wants to "hack" nn.Linear.
Another solution is to replace nn.Linear completely with nn.Conv1d, and apply some load_state_dict hook for backward compatability. But I think PyTorch's load state dict hook is only available recently. I will check on this.

[sanchit-gandhi]

Cool cool! That's an astute observation regarding LoRA. I think it's probably ok given that the linear TDNN layers are only a fraction of the Transformer encoder ones, and as you say they will be trained from scratch in a round of fine-tuning. Replacing nn.Linear with nn.Conv1d might work, but we'll have to propagate this change to Flax and TF to ensure we have cross platform compatibility (this could get quite involved). Overall, the priority is maintaining cross platform compatibility within transformers, over LoRA from third-party libraries

[gau-nernst]

On RTX 3090, 5s input

dtype	batch size	type	de13970	4ad9e29
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

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")

[sanchit-gandhi]

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

[gau-nernst]

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)

[sanchit-gandhi]

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):

   transformers/tests/models/wav2vec2/test_modeling_wav2vec2.py

   Line 1922 in 1c6f072

   def test_inference_speaker_verification(self):

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

[sanchit-gandhi]

Thanks @gau-nernst! Really elegant solution 👏 Nice job on bringing the speed-ups without breaking changes

If the slow test mentioned before passes, then this is good to merge for me! Requesting final review from @ArthurZucker

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.

[ArthurZucker]

Very elegant!
Regarding this:

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

TDNN Is not documented so not really part of our Public API, which means we could change it (but mostly only if we really have a significant gain in performance)
It's a breaking change for people who trained a model with LoRA no? That my only concern

[gau-nernst]

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.

[ArthurZucker]

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?) 😉

[gau-nernst]

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.

[sanchit-gandhi]

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

[ArthurZucker]

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

[sanchit-gandhi]

cc @younesbelkada as well: currently we have un-optimised TDNN code, where we call nn.functional.unfold to do a conv1d, then apply a linear layer self.kernel

@gau-nernst has proposed is a solution where we still use the weights of the linear self.kernel layer (rather than switching to a conv1d layer which would change the state dict), but then use nn.functional.conv1d to do the convolution

=> the benchmarks show there's about a 10% performance gain by doing this

However, we worry that since we don't call self.kernel.forward, the LoRA layers for this layer will no longer work for inference. Is this hunch correct?

[younesbelkada]

Correct, these changes will make LoRA unusable for the kernel weights on PEFT side as we call torch.nn.functional.conv1d instead of the forward of the module itself.
However, usually we only adapt self attention modules (query, key, value) and in rare cases MLP layers of the transformer block. So I am not sure users really adapt kernel layers.
I am not sure if we could make that PR BC (through a config variable maybe?)

cc @pacman100 & @BenjaminBossan as well

[BenjaminBossan]

If this is an edge case that is unlikely to be practically used, I would be fine with breaking backwards compatibility here, given the benefits.

If we are worried about this, however, we could make a check and raise an error/give a warning if a user happens to apply LoRA here. That check could be made here (check if self.kernel looks like a LoRA layer) or we could add a check in PEFT if we detect that a user tries to apply LoRA to this specific layer of this specific model.

Btw. this also affects other adapters such as IA³.

[sanchit-gandhi]

Awesome - thanks both for your input! It sounds like it's super low likelihood that a user is applying LoRA to these TDNN kernel layers, so I'd be happy proceeding with these changes if it's good for you @gau-nernst @ArthurZucker?

If not, and you'd like to preserve backwards comp for this edge case, I like @BenjaminBossan's idea of checking whether the weights look like a LoRA layer, and triggering a warning if so. There's also the option of having two implementations for the TDNN layer:

1. Old, un-optimised variant that is LoRA compatible
2. New, optimised variant that is not LoRA compatible

=> we could use 2 by default, but switch to 1 if the user has applied LoRA to the kernel

[ArthurZucker]

I'm fine with just warning the user if there are lora weights !

[sanchit-gandhi]

Okay so the problem here is that Wav2Vec2 doesn't have input embeddings: it works directly with the raw audio waveform and computes positional embeddings directly from the input. Since there is no notion of input embeddings, we don't implement the method .get_input_embeddings. This breaks the PEFT LoRA logic, which expects this method to be implemented.

This suggests that Wav2Vec2 is not compatible at all with PEFT LoRA as things currently stand. Would you mind indeed opening an issue on PEFT with the codesnippet you shared, and a copy-and-paste explanation of why it doesn't work?

This means though that we don't need to worry about breaking LoRA for Wav2Vec2ForXVector, since it doesn't work with it anyway at the moment.

[gau-nernst]

The error is due to setting gradient_checkpointing=True in Wav2Vec2, which is the default for facebook/wav2vec2-base. Other Wav2Vec2 checkpoints, like facebook/wav2vec2-xls-r-300m, have gradient_checkpointing=False, thus there is no error, and PEFT LoRA still works with them. Overriding gradient_checkpointing=False can also make facebook/wav2vec2-base works with PEFT LoRA.

[sanchit-gandhi]

Ah ok thanks for the clarification! Let's stick to our original strategy then of raising a warning if there are LoRA weights

[sanchit-gandhi]

Hey @gau-nernst - would you be interested in finishing this PR? Now that we've agreed on the logic here and you've done a great job at laying down the foundations, it should be quite fast to tidy it up and complete the feature!

[gau-nernst]

Hello @sanchit-gandhi. I'm happy to finish this. Do let me know which parts you want to be changed.

[sanchit-gandhi]

Thanks for adding the LoRA warnings, this PR looks great @gau-nernst 🙌 Requesting final review from @ArthurZucker!

[ArthurZucker]

Really sorry for the late review!
Can you rebase and I'll merge!

[ArthurZucker]

Great work here on finding an elegant solution 🤗

[gau-nernst]

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

[ArthurZucker]

Thanks! Merging 🤗