Search papers, labs, and topics across Lattice.
This paper introduces a novel fine-tuning method for transformer language models that utilizes sparse attention to optimize key-value (KV) cache selection, enabling efficient long-context inference on moderate hardware. The approach allows the model to adapt to any KV cache policy, often surpassing the performance of models trained with exact attention, even under constraints typical of a single Nvidia A100 GPU. Additionally, the authors present an efficient implementation of H2O sparse attention, complete with a new open-source library, KeysAndValues, facilitating practical applications in long-context scenarios.
Fine-tuning with sparse attention can outperform traditional exact attention models while running efficiently on modest hardware.
A lot of prior work addressed key-value (KV) cache selection and compression by sparse attention to enable long-context inference for transformer language models without excessive hardware budgets. We provide a new method for fine-tuning models with sparse attention. It works for any KV cache policy, runs on a moderate hardware budget (e.g., a single Nvidia A100 GPU with 40 GB RAM), and allows the model to co-adapt with the policy, often outperforming models trained with exact attention (sequence parallelism). We also provide an efficient implementation of H2O sparse attention (the leading policy in our experiments) with dedicated scaled dot product attention kernel support. KeysAndValues (https://github.com/awslabs/keys_values), a new open source library for long-context inference and fine-tuning, provides easy-to-use and performant code for all methods discussed here.
Amazon Science