USP: A Unified Sequence Parallelism Approach for Long Context Generative AI
Fuente:
arXiv
Salvato in:
| Autori principali: | , |
|---|---|
| Natura: | Preprint |
| Pubblicazione: |
2024
|
| Soggetti: | |
| Accesso online: | |
| Tags: |
Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
|
| _version_ | 1866911939682107392 |
|---|---|
| author | Fang, Jiarui Zhao, Shangchun |
| author_facet | Fang, Jiarui Zhao, Shangchun |
| contents | Sequence parallelism (SP), which divides the sequence dimension of input tensors across multiple computational devices, is becoming key to unlocking the long-context capabilities of generative AI models. This paper investigates the state-of-the-art SP approaches, i.e. DeepSpeed-Ulysses and Ring-Attention, and proposes a unified SP approach, which is more robust to transformer model architectures and network hardware topology. This paper compares the communication and memory cost of SP and existing parallelism, including data/tensor/zero/pipeline parallelism, and discusses the best practices for designing hybrid 4D parallelism involving SP. We achieved 47% MFU on two 8xA800 nodes using SP for the LLAMA3-8B model training using sequence length 208K. Our code is publicly available at https://github.com/feifeibear/long-context-attention. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2405_07719 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | USP: A Unified Sequence Parallelism Approach for Long Context Generative AI Fang, Jiarui Zhao, Shangchun Machine Learning Artificial Intelligence Sequence parallelism (SP), which divides the sequence dimension of input tensors across multiple computational devices, is becoming key to unlocking the long-context capabilities of generative AI models. This paper investigates the state-of-the-art SP approaches, i.e. DeepSpeed-Ulysses and Ring-Attention, and proposes a unified SP approach, which is more robust to transformer model architectures and network hardware topology. This paper compares the communication and memory cost of SP and existing parallelism, including data/tensor/zero/pipeline parallelism, and discusses the best practices for designing hybrid 4D parallelism involving SP. We achieved 47% MFU on two 8xA800 nodes using SP for the LLAMA3-8B model training using sequence length 208K. Our code is publicly available at https://github.com/feifeibear/long-context-attention. |
| title | USP: A Unified Sequence Parallelism Approach for Long Context Generative AI |
| topic | Machine Learning Artificial Intelligence |
| url | https://arxiv.org/abs/2405.07719 |