We're releasing an optimized implementation of GPT2/GPT3 with FlashAttention🚀!
This trains 3-5x faster than the Huggingface version, reaching up to 189 TFLOPs/sec per A100, 60.6% (model) FLOPs util of the theoretical maximum. 1/6
github.com/HazyResearch/f…
This trains 3-5x faster than the Huggingface version, reaching up to 189 TFLOPs/sec per A100, 60.6% (model) FLOPs util of the theoretical maximum. 1/6
github.com/HazyResearch/f…
The main ingredient is FlashAttention, which computes attention fast (2-4x) and with less memory (10x), without any approximation. This means that we don't need to do any activation checkpointing
2/6
2/6
https://twitter.com/tri_dao/status/1531437619791290369
We also provide optimized implementations of other layers:
- Fused matmul + bias + gelu for the MLP (based on Apex and cuBLASLt)
- Optimized cross entropy loss (based on Apex)
- Fused rotary embedding
- Fused dropout + residual + LayerNorm (building on Apex's FastLayerNorm)
3/6
- Fused matmul + bias + gelu for the MLP (based on Apex and cuBLASLt)
- Optimized cross entropy loss (based on Apex)
- Fused rotary embedding
- Fused dropout + residual + LayerNorm (building on Apex's FastLayerNorm)
3/6
We include training scripts to train GPT2 on Openwebtext and GPT3 on The Pile as examples.
On 1-node, our implementation is 3-5x faster than Huggingface and 1.3-2x faster than Megatron-LM!
4/6
On 1-node, our implementation is 3-5x faster than Huggingface and 1.3-2x faster than Megatron-LM!
4/6
Of course Megatron-LM is great for large-scale distributed training, and Huggingface has such a friendly API to play with all the latest models. We're planning to integrate our work into these libraries
5/6
5/6
I'll be at #NeurIPS2022 presenting FlashAttention (poster Session 4 Hall J #917, Wednesday 4-6 PM). Drop by and say hi!
6/6
6/6
https://twitter.com/realDanFu/status/1597330278510342144
• • •
Missing some Tweet in this thread? You can try to
force a refresh
　
