An Adaptive Placement and Parallelism Framework for Accelerating RLHF Training

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: Xiao, Youshao, Zhou, Zhenglei, Mao, Fagui, Wu, Weichang, Zhao, Shangchun, Ju, Lin, Liang, Lei, Zhang, Xiaolu, Zhou, Jun
Formato: Preprint
Publicado: 2023
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866909347093676032
author Xiao, Youshao
Zhou, Zhenglei
Mao, Fagui
Wu, Weichang
Zhao, Shangchun
Ju, Lin
Liang, Lei
Zhang, Xiaolu
Zhou, Jun
author_facet Xiao, Youshao
Zhou, Zhenglei
Mao, Fagui
Wu, Weichang
Zhao, Shangchun
Ju, Lin
Liang, Lei
Zhang, Xiaolu
Zhou, Jun
contents Recently, ChatGPT or InstructGPT like large language models (LLM) has made a significant impact in the AI world. Many works have attempted to reproduce the complex InstructGPT's training pipeline, namely Reinforcement Learning with Human Feedback (RLHF). However, the mainstream distributed RLHF training methods typically adopt a fixed model placement strategy, referred to as the Co-located strategy. This strategy treats all four interdependent models involved in RLHF as a single entity, distributing them across all devices and applying parallelism techniques designed for a single model, regardless of the workload heterogeneity inherent to each model. As a result, this strategy exacerbates the generation bottlenecks in the RLHF training and degrades the overall training efficiency. To address these issues, we propose a flexible model placement framework that offers two general and agile model placement strategies. The Interleaving strategy helps reduce memory redundancy and communication costs of RLHF training by placing models without dependencies on exclusive devices with careful orchestration. On the other hand, the Disaggregated strategy improves the throughput of model training by separating the training and inference runtime of the RLHF pipeline with additional shadow models. Furthermore, our framework provides a simple user interface and guidelines to easily and flexibly configure these strategies in various training scenarios. Our experiments have shown that our strategy can achieve notable improvements up to 11x, compared to the current state-of-the-art (SOTA) approaches. The results highlight the effectiveness and adaptability of our methods in accelerating the training of distributed RLHF.
format Preprint
id arxiv_https___arxiv_org_abs_2312_11819
institution arXiv
publishDate 2023
record_format arxiv
spellingShingle An Adaptive Placement and Parallelism Framework for Accelerating RLHF Training
Xiao, Youshao
Zhou, Zhenglei
Mao, Fagui
Wu, Weichang
Zhao, Shangchun
Ju, Lin
Liang, Lei
Zhang, Xiaolu
Zhou, Jun
Machine Learning
Artificial Intelligence
Computation and Language
Recently, ChatGPT or InstructGPT like large language models (LLM) has made a significant impact in the AI world. Many works have attempted to reproduce the complex InstructGPT's training pipeline, namely Reinforcement Learning with Human Feedback (RLHF). However, the mainstream distributed RLHF training methods typically adopt a fixed model placement strategy, referred to as the Co-located strategy. This strategy treats all four interdependent models involved in RLHF as a single entity, distributing them across all devices and applying parallelism techniques designed for a single model, regardless of the workload heterogeneity inherent to each model. As a result, this strategy exacerbates the generation bottlenecks in the RLHF training and degrades the overall training efficiency. To address these issues, we propose a flexible model placement framework that offers two general and agile model placement strategies. The Interleaving strategy helps reduce memory redundancy and communication costs of RLHF training by placing models without dependencies on exclusive devices with careful orchestration. On the other hand, the Disaggregated strategy improves the throughput of model training by separating the training and inference runtime of the RLHF pipeline with additional shadow models. Furthermore, our framework provides a simple user interface and guidelines to easily and flexibly configure these strategies in various training scenarios. Our experiments have shown that our strategy can achieve notable improvements up to 11x, compared to the current state-of-the-art (SOTA) approaches. The results highlight the effectiveness and adaptability of our methods in accelerating the training of distributed RLHF.
title An Adaptive Placement and Parallelism Framework for Accelerating RLHF Training
topic Machine Learning
Artificial Intelligence
Computation and Language
url https://arxiv.org/abs/2312.11819