RedCoast: A Lightweight Tool to Automate Distributed Training of LLMs on Any GPU/TPUs

TL;DR

RedCoast introduces a lightweight tool to automate distributed training and inference of large language models, addressing memory bottlenecks without requiring MLSys expertise. It achieves this through two design pillars: (i) automatic tensor parallelism built on two simple dimension-splitting rules applicable across transformer architectures, implemented on top of jax.pjit; and (ii) a three-function pipeline design (collate, loss, predict) that abstracts away data/data-parallelism, multi-host orchestration, and checkpointing. The approach is demonstrated across LLMs such as LLaMA-7B, T5-11B, and OPT-66B, with Redco delivering throughput comparable to state-of-the-art tools like Alpa and outperforming FSDP, while requiring far less engineering effort. Additionally, Redcoast provides a three-component library (Deployer, Trainer, Predictor) that enables multi-host support and standardized checkpointing, with a lightweight dependency footprint built on Flax/Jax. Collectively, the work lowers the barrier to distributed LLM development and extends pipeline flexibility to advanced AI paradigms like meta-learning and reinforcement learning, by significantly reducing code size and complexity without compromising performance.

Abstract

The recent progress of AI can be largely attributed to large language models (LLMs). However, their escalating memory requirements introduce challenges for machine learning (ML) researchers and engineers. Addressing this requires developers to partition a large model to distribute it across multiple GPUs or TPUs. This necessitates considerable coding and intricate configuration efforts with existing model parallel tools, such as Megatron-LM, DeepSpeed, and Alpa. These tools require users' expertise in machine learning systems (MLSys), creating a bottleneck in LLM development, particularly for developers without MLSys background. In this work, we present RedCoast (Redco), a lightweight and user-friendly tool crafted to automate distributed training and inference for LLMs, as well as to simplify ML pipeline development. The design of Redco emphasizes two key aspects. Firstly, to automate model parallelism, our study identifies two straightforward rules to generate tensor parallel strategies for any given LLM. Integrating these rules into Redco facilitates effortless distributed LLM training and inference, eliminating the need of additional coding or complex configurations. We demonstrate the effectiveness by applying Redco on a set of LLM architectures, such as GPT-J, LLaMA, T5, and OPT, up to the size of 66B. Secondly, we propose a mechanism that allows for the customization of diverse ML pipelines through the definition of merely three functions, avoiding redundant and formulaic code like multi-host related processing. This mechanism proves adaptable across a spectrum of ML algorithms, from foundational language modeling to complex algorithms like meta-learning and reinforcement learning. As a result, Redco implementations exhibit significantly fewer lines of code compared to their official counterparts.

Figures (6)

• Figure 1: With a number of shards specified by user, Redco automatically conduct the model partitioning and distribution across hosts and devices.
• Figure 2: The template code of using Redco to implement distributed training, where users only have to design a pipepine through three fucntions, without concerning data and model parallelism, multi-host related processing, randomness control, etc., which eliminates a lot of boilerplate coding.
• Figure 3: The comparison of throughput in the running of GPT-J-6B on a $4 \times$ A100 server. Redco's performance surpasses that of FSDP and is close to that of Alpa, the tool with state-of-the-art model parallel efficiency but intricate engineering.
• Figure 4: Pipeline design functions of image captioning under Redco.
• Figure 6: The comparison of code lines across a diverse set of ML algorithms. (There is no well-accepted official Flax implementations for FedAvg and MAML.)