Fix cuda configuration error. #570
Conversation
Port PriorLabs/TabPFN-private#46 to public. * remove flash attention package use * chunck scaled_dot_product_attention This might mean people who installed their own flash attention get slower inference, particularly anyone using PyTorch 2.1, where FA2 isn't built in. But I doubt anyone is doing this? Co-authored-by: Oscar Key <oscar@priorlabs.ai>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request removes the dependency on the flash-attn package and introduces a chunking mechanism for scaled_dot_product_attention to work around a CUDA kernel limitation with large batch sizes. This is a good simplification that improves robustness. The changes are well-implemented and include new tests. I've identified one edge case where an empty batch would cause a runtime error and have suggested a fix along with an additional test case to cover it. Otherwise, the changes look solid.
There was a problem hiding this comment.
LGTM thanks!
This might mean people who installed their own flash attention get slower inference. The difference should be non existent/small for more recent PyTorch, but would be big for anyone using PyTorch 2.1 where FA2 isn't built in. But probably no one is doing this?
Yes that seems fine to me. I'm wondering if we want to force torch 2.2. On the one hand it's quite recent (jan 2024), on the other hand people using torch 2.1 will have a very bad experience 🤔
|
I would be in favour of forcing 2.2, but I created https://linear.app/priorlabs/issue/RES-813/drop-torch-21warn-if-used-on-big-datasets |
Co-authored-by: mirror-bot <mirror-bot@users.noreply.github.com>
Port https://github.com/PriorLabs/TabPFN-private/pull/46 to public:
This might mean people who installed their own flash attention get slower inference. The difference should be non existent/small for more recent PyTorch, but would be big for anyone using PyTorch 2.1 where FA2 isn't built in. But probably no one is doing this?