nvFuser integration for operation fusion

[yuanandonly]

What does this PR do?

**redirect to #357 **

Uses AOTAutograd and NVfuser to created fused operator patterns layers

• Bias + Activation + Dropout
• Bias + Dropout + Residual
• Bias + Dropout + Residual + Layernorm

Integrated BiasActivationDropout pattern into MLP feedforward

Wrote benchmarking against pytorch eager and triton fused patterns

Before submitting

• Did you have fun?
  • Make sure you had fun coding 🙃
• Did you read the contributor guideline?
• Was this discussed/approved via a Github issue? (no need for typos, doc improvements)
  • N/A
• Did you make sure to update the docs?
  • N/A
• Did you write any new necessary tests?
  • N/A
• Did you update the changelog? (if needed)
  • [] N/A

PR review

Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

[dianaml0]

Thanks for the initial work here @yuanandonly! Left a few comments

[blefaudeux]

Nice, and the perf is super good ! This is pretty great @yuanandonly, thanks already !

I would agree with @dianaml0 and it will be simpler for you actually, probably better not to register a new layer type ? It can be automatically (conditionally on nvfuser being present and in the cases which are just faster) applied to the existing layers, that would be even better I feel -less complexity and automatically improving the situation for all users-. You can still expose a macro flag, in the form of a bool that users can flip -possibly at runtime-, for debugging for instance

[fmassa]

Thanks for the PR!

On top of what Diana and Benjamin mentioned, I have also a couple of comments.

[Chillee]

Usage of AOTAutograd looks fine to me.

I do wonder about the API though. One of the nice things about using AOTAutograd is that it gives the user a lot of flexibility in how to construct their function. It seems like this API sacrifices that advantage.

Also, it might be worth trying out the TorchInductor backend as well.

[yuanandonly]

@Chillee What sort of flexibility are you suggesting for the API? Do you mean keeping the nvfuser component more general instead of having different components for the different fused patterns?

[dianaml0]

Perhaps we can pull out a general method that takes in a function and outputs the nvfused version. Then we can use that to provide these various fusions

[dianaml0]

It could remove some redundancy in these files

[dianaml0]

How come we skip here? Should this error out if a user tries it?

[yuanandonly]

This is because the layernormstyle argument doesn't apply to the other patterns, so when layer_norm_style == LayerNormStyle.Pre all the possible configurations have already been tested once

[dianaml0]

So should this just be if fused_pattern != NVFusedBiasDropoutResLayerNorm ?

[fmassa]

Thanks for the changes Chris!

I've left some more comments. I think that ideally we would want the nvfuser model to be a drop-in replacement for the vanilla model, which would imply that the weight layout should be the same. I understand this might complicate a few things for now, so let's not make this a merge-blocking requirement, but it would be good to get this fixed.

[fmassa]

What is the desired way to let users have the nvfuser path enabled? Will they need to clone and modify xformers for that?

[yuanandonly]

I think currently we have it as something that users will need to modify. I haven't looked into it a great deal, perhaps users who haven't cloned xformers can do something like python -c "import xformers; xformers._is_functorch_available=True;

[fmassa]

The fact that the bias live in a different module means that you won't be able to easily swap to use nvfuser on a pre-trained model which used the standard layers. This is IMO a drawback which would be good to be solved.

Using this optimization means that this is not just a transparent change to the user, and the snippet like I proposed before with state_dict won't work out of the box.