Add custom ops for compatibility with PT Compile

[ani300]

This PR adds basic support for torch.compile() for the non-varlen variants of Flash Attention.

This essentially allows for models that use the flash_attn_qkvpacked_func, flash_attn_kvpacked_func, and flash_attn_func to compile without graph breaks.

I can add unit tests if it makes sense. I'll test it with our own training pipeline for performance measurements and I'll post them later.

This uses the new custom operators API in Pytorch 2.4. I can move to the older APIs if needed, or I can look up how to make both coexist

[ani300]

Added varlen functions too

[tridao]

Thank you! If this requires torch 2.4, is there some way to make it optional for older pytorch version? e.g. define torch.library.custom_op to be a no-op? Idk what people usually do.

[ani300]

There's a different API that was used in previous pytorch versions. It's more cumbersome to use, but let me think about how to use both APIs at the same time with a version selection. I can also do a noop wrapper based on version

[ani300]

Ok, I've updated the code to only wrap things if Pytorch is 2.4 or higher. Even though we're comparing strings, torch.__version__ implements semantic comparisons, so it should always work correctly. As soon as I have the training performance comparison I'll update the PR

[mayank31398]

aah, yikes
@ani300 I had started working on the same thing #1145 😓

Ill let you handle this 😃

[ani300]

@tridao I've updated the code after testing it with SFTTrainer to make sure everything was getting called correctly and compiling without graph breaks. I've also added a check for the pytorch version so this only activates if the APIs are available.

One question I have is if it's ok to change the C++ pytorch interface so the out/out_padded padding/unpadding can be done outside of the custom op. Right now, Pytorch custom ops don't allow outputs to be aliases of other outputs, which this pair is when head_dim % 8 == 0. This means I have to clone one of the outputs to unalias them, which is a waste of memory bandwidth.

[tridao]

Yes let's change the C++ code to not pad the output, and have the Python part handle that.

[GLivshits]

There's a different API that was used in previous pytorch versions. It's more cumbersome to use, but let me think about how to use both APIs at the same time with a version selection. I can also do a noop wrapper based on version

For older versions, you can use:

1. torch.library.define - just for definition of ops.
2. Then use this:
   OLD_TORCH_VERSION = version.parse(torch.version).base_version < "2.4.0"

if OLD_TORCH_VERSION:
    _torch_custom_op_wrapper = partial(torch.library.impl, types="cuda")
    _torch_register_fake_wrapper = torch.library.impl_abstract
else:
    _torch_custom_op_wrapper = partial(torch.library.custom_op, device_types="cuda", mutates_args=())
    _torch_register_fake_wrapper = torch.library.register_fake

Here I assume that mutates_args is empty. If it is not - I don't currently know how to handle that (and I didnt think much) without copy-pasting lots of code.

[mayank31398]

@GLivshits I dont think it can be handled in older versions of torch

[mayank31398]

@tridao @ani300 is there any progress/updates on this?
Its a pretty neat feature to have Flash Attention fully end-to-end traceable natively.

[ani300]

Hey, I'm almost done with the C++ portion, hopefully by Monday I can update the PR and rerun all the CI to make sure it's still good with all the changes. I'll also try the suggestions from @GLivshits to make it work with previous versions

[ani300]

@GLivshits the way the code is currently written, the backward functions need to have mutate_args, or we need to do an extra memory copy on GPU, which might significantly impact performance

[ani300]

@tridao this is ready for final review and/or merging. I tested with the whole unit test suite for both pytorch 2.3 and pytorch 2.4 on an A100

[raghukiran1224]

@tridao Any update on this?

[anijain2305]

Thanks @ani300 for taking this on. I started working on this here (#1209). But this is way further along. So, I closed mine.

cc @zou3519 for custom ops API usage review @Chillee

[anijain2305]

Nit - s/flashattn/flash_attn/

[anijain2305]

Perhaps use the same if condition

if torch.__version__ >= "2.4.0":

[anijain2305]

Is this equivalent of return func?

[anijain2305]

Ditto - return func

[anijain2305]

Nit - a more desciptive name - noop_custom_op_wrapper

And maybe a comment that we dont support < 2.4 versions.

[anijain2305]

return fn

[ani300]

@anijain2305 just pushed all the updates/fixes you suggested. I am running CI again to make sure it works for both 2.4 and 2.3. Will update the comments when it's done

[anijain2305]

Thanks @ani300

Custom ops changes look good to me. Hope @tridao can take a look at the other parts.

[ani300]

Tests all ran successfully again

[tridao]

Awesome, thanks!

[umarbutler]

I just built the flash-attention main branch from source and sadly I'm getting the same error message during training with the Hugging Face Trainer with torch_compile=True:

.venv/lib/python3.12/site-packages/torch/_dynamo/variables/functions.py:663: UserWarning: Graph break due to unsupported builtin flash_attn_2_cuda.PyCapsule.fwd. This function is either a Python builtin (e.g. _warnings.warn) or a third-party C/C++ Python extension (perhaps created with pybind). If it is a Python builtin, please file an issue on GitHub so the PyTorch team can add support for it and see the next case for a workaround. If it is a third-party C/C++ Python extension, please either wrap it into a PyTorch-understood custom operator (see https://pytorch.org/tutorials/advanced/custom_ops_landing_page.html for more details) or, if it is traceable, use torch.compiler.allow_in_graph.

[umarbutler]

It looks like I somehow ran python setup.py bdist_wheel and forgot I was creating a wheel (which I also ended up mistakenly deleting moments ago) 😅 hopefully second time's the charm

[ani300]

I was worried for a second, as I just used it with the Huggingface Trainer yesterday successfully 😥

[umarbutler]

@ani300 I can confirm that it works! And much faster than PyTorch SDPA, shaved off 50 training hours (400h -> 350h) 😆