May 2024 Summaries

Flash Attention received the inaugural Stanford Open Source Software award for its significant impact, engagement, and adoption across the industry. FlashAttention is an algorithm that reorders attention computation to speed up Transformer training and inference by reducing memory usage from quadratic to linear in sequence length. Its variants, including FlashAttention-2, offer further improvements with speeds of up to 4x faster training and fine-tuning of Large Language Models (LLMs), achieving 72% model FLOPs utilization for training on NVIDIA A100s. The technology is now widely used by companies and researchers and has been integrated into popular frameworks such as PyTorch and Hugging Face, with its Github repo receiving over 11k stars. FlashAttention-2 is designed as a drop-in replacement for the original algorithm, offering a 2x speedup on core attention operations and achieving further improvements in training Transformers end-to-end.

We have developed a simple framework called ThunderKittens, an embedded Domain Specific Language (DSL) for AI kernels, aiming to make it easy to express key technical ideas in a clean and understandable code. The fundamental object of the DSL is a matrix that fits into tensor cores, which are 94% of the compute on an H100, to keep them busy. We've made the API PyTorch-like to be familiar to AI people, while still providing full power of the host and transparency in accelerators. ThunderKittens has been used by our team, resulting in impressive performance gains, including matching FA2 performance on 4090s and A100s, and being faster forward and backward than FA2 on H100s. The project is an art project, but we hope it makes key ideas clear and are open to feedback. We're releasing ThunderKittens now, integrated with Andrej's NanoGPT project, to make key concepts accessible to a wider audience.

The RedPajama-V2 dataset is a 30 trillion token web dataset designed for training large language models (LLMs). It's not intended to be used out of the box, but rather as a foundation for creating high-quality datasets. The dataset has a high-recall vs high-precision design, enabling researchers to experiment with different data selection techniques and discover recipes that produce downstream models with desired properties. To facilitate its use, the dataset comes with quality signals, duplication tags, and minhash signatures, which can be used to filter out low-quality documents, deduplicate data, and perform fuzzy deduplication. The dataset is partitioned into head, middle, and tail partitions based on perplexity, with the head and middle partitions having higher quality data than the tail partition. Users can use metadata fields extracted from the CCNet pipeline to filter the dataset by URL domain or date. The dataset also provides an example of how to remove documents containing ellipses and interpret the format of the quality signals. To ensure retries when downloading the dataset, users can pass a download config with the number of retries set to something larger than 1. The dataset's total size is approximately 260TB, comprising four components: text documents, quality signals, minhash signatures, and duplicate IDs. The dataset can be used as a starting point for creating high-quality datasets that resemble data from high-quality sources such as Wikipedia or OpenWebText.