Optimizing RLHF Training for Large Language Models with Stage Fusion
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_ | 1866908330554818560 |
|---|---|
| author | Zhong, Yinmin Zhang, Zili Wu, Bingyang Liu, Shengyu Chen, Yukun Wan, Changyi Hu, Hanpeng Xia, Lei Ming, Ranchen Zhu, Yibo Jin, Xin |
| author_facet | Zhong, Yinmin Zhang, Zili Wu, Bingyang Liu, Shengyu Chen, Yukun Wan, Changyi Hu, Hanpeng Xia, Lei Ming, Ranchen Zhu, Yibo Jin, Xin |
| contents | We present RLHFuse, an efficient training system with stage fusion for Reinforcement Learning from Human Feedback (RLHF). Due to the intrinsic nature of RLHF training, i.e., the data skewness in the generation stage and the pipeline bubbles in the training stage, existing RLHF systems suffer from low GPU utilization. RLHFuse breaks the traditional view of RLHF workflow as a composition of individual tasks, splitting each task into finer-grained subtasks, and performing stage fusion to improve GPU utilization. RLHFuse contains two key ideas. First, for generation and inference tasks, RLHFuse splits them into sample-level subtasks, enabling efficient inter-stage fusion to overlap the execution of generation and inference stages, thus mitigating the original generation bottleneck dominated by long-tailed samples. Second, for training tasks, RLHFuse breaks them into subtasks of micro-batches and performs intra-stage fusion to concurrently execute these subtasks in the training stage with a fused pipeline schedule, effectively mitigating the pipeline bubbles. The experiments show that RLHFuse increases the training throughput by up to $3.7\times$, compared to existing systems. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2409_13221 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | Optimizing RLHF Training for Large Language Models with Stage Fusion Zhong, Yinmin Zhang, Zili Wu, Bingyang Liu, Shengyu Chen, Yukun Wan, Changyi Hu, Hanpeng Xia, Lei Ming, Ranchen Zhu, Yibo Jin, Xin Machine Learning Computation and Language Distributed, Parallel, and Cluster Computing We present RLHFuse, an efficient training system with stage fusion for Reinforcement Learning from Human Feedback (RLHF). Due to the intrinsic nature of RLHF training, i.e., the data skewness in the generation stage and the pipeline bubbles in the training stage, existing RLHF systems suffer from low GPU utilization. RLHFuse breaks the traditional view of RLHF workflow as a composition of individual tasks, splitting each task into finer-grained subtasks, and performing stage fusion to improve GPU utilization. RLHFuse contains two key ideas. First, for generation and inference tasks, RLHFuse splits them into sample-level subtasks, enabling efficient inter-stage fusion to overlap the execution of generation and inference stages, thus mitigating the original generation bottleneck dominated by long-tailed samples. Second, for training tasks, RLHFuse breaks them into subtasks of micro-batches and performs intra-stage fusion to concurrently execute these subtasks in the training stage with a fused pipeline schedule, effectively mitigating the pipeline bubbles. The experiments show that RLHFuse increases the training throughput by up to $3.7\times$, compared to existing systems. |
| title | Optimizing RLHF Training for Large Language Models with Stage Fusion |
| topic | Machine Learning Computation and Language Distributed, Parallel, and Cluster Computing |
| url | https://arxiv.org/abs/2409.13221 |