Context Parallelism for Scalable Million-Token Inference
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2024
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866910914749399040 |
|---|---|
| author | Yang, Amy Yang, Jingyi Ibrahim, Aya Xie, Xinfeng Tang, Bangsheng Sizov, Grigory Reizenstein, Jeremy Park, Jongsoo Huang, Jianyu |
| author_facet | Yang, Amy Yang, Jingyi Ibrahim, Aya Xie, Xinfeng Tang, Bangsheng Sizov, Grigory Reizenstein, Jeremy Park, Jongsoo Huang, Jianyu |
| contents | We present context parallelism for long-context large language model inference, which achieves near-linear scaling for long-context prefill latency with up to 128 H100 GPUs across 16 nodes. Particularly, our method achieves 1M context prefill with Llama3 405B model in 77s (93% parallelization efficiency, 63% FLOPS utilization) and 128K context prefill in 3.8s. We develop two lossless exact ring attention variants: pass-KV and pass-Q to cover a wide range of use cases with the state-of-the-art performance: full prefill, persistent KV prefill and decode. Benchmarks on H100 GPU hosts inter-connected with RDMA and TCP both show similar scalability for long-context prefill, demonstrating that our method scales well using common commercial data center with medium-to-low inter-host bandwidth. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2411_01783 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | Context Parallelism for Scalable Million-Token Inference Yang, Amy Yang, Jingyi Ibrahim, Aya Xie, Xinfeng Tang, Bangsheng Sizov, Grigory Reizenstein, Jeremy Park, Jongsoo Huang, Jianyu Distributed, Parallel, and Cluster Computing Artificial Intelligence Machine Learning We present context parallelism for long-context large language model inference, which achieves near-linear scaling for long-context prefill latency with up to 128 H100 GPUs across 16 nodes. Particularly, our method achieves 1M context prefill with Llama3 405B model in 77s (93% parallelization efficiency, 63% FLOPS utilization) and 128K context prefill in 3.8s. We develop two lossless exact ring attention variants: pass-KV and pass-Q to cover a wide range of use cases with the state-of-the-art performance: full prefill, persistent KV prefill and decode. Benchmarks on H100 GPU hosts inter-connected with RDMA and TCP both show similar scalability for long-context prefill, demonstrating that our method scales well using common commercial data center with medium-to-low inter-host bandwidth. |
| title | Context Parallelism for Scalable Million-Token Inference |
| topic | Distributed, Parallel, and Cluster Computing Artificial Intelligence Machine Learning |
| url | https://arxiv.org/abs/2411.01783 |