R ed C oast: A Lightweight Tool to Automate Distributed Training of LLM s on Any GPU / TPU s (2024)

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. RedCoast (Redco) has been released under Apache 2.0 license at https://github.com/tanyuqian/redco.

## introduction in recent years, the field of ai has witnessed profound advancements, predominantly attributed to the advent of llms with an impressive number of parameters, spanning from billions to hundreds of billions @xcite . notable examples include @xcite and llama @xcite . yet, the size of these llms presents distinct challenges in terms of model deployment for ml researchers and engineers. the primary challenge arises from the substantial memory requirements of llms, often exceed the capability of a single gpu or tpu. this necessitates the use of model parallelism, a technique that partitions the llms into various shards, subsequently distributing them across multiple devices or even different hosts. however, achieving this partitioning requires intricate engineering, including the formulation of a tensor-specific splitting strategy. while several specialized tools like deepspeed @xcite , alpa @xcite , and fsdp @xcite provide diverse model parallelism solutions, but they demand significant additional coding and intricate configurations based on model architecture and hardware specifics, requiring in-depth understanding of mlsys. such additional efforts make the deployment of llms particularly challenging, especially for users without mlsys expertise, such as algorithm developers or researchers. at times, the intricacy of coding for model parallelism proves to be even more daunting than the algorithm design itself. in this work, we introduce redcoast (redco) foot_0 , a lightweight and user-friendly tool designed to automate the distributed training and inference of llms, thereby users without mlsys expertise can also effortlessly use the tool without additional coding or intricate configurations. furthermore, we propose a novel and neat mechanism to implement ml algorithms. this method necessitates users to define merely three functions as their pipeline design, with redco managing all the remaining details in execution, such as data parallelism, multihost related processing, checkpointing, etc. redco's design emphasizes two key aspects. the first is the automatic model parallelism. we identify two straightforward rules to generate the model parallel strategy for arbitrarily given transformer architecture, and integrate them into redco. unlike tools such as megatron @xcite and deepspeed @xcite which require users to manipulate model forward function for different architecture and system specifics, redco automates the process, where users only need to specify the desired number of shards to partition the model. we verified the effectiveness of redco's model parallel strategy on multiple llms including llama-7b @xcite , t5-11b @xcite , and opt-66b @xcite . moreover, pipelines driven by redco demonstrate efficiency superior to those implemented with fsdp @xcite , and closely matching the performance of alpa @xcite , the tool with state-of-the-art model parallel efficiency. another pivotal feature of redco is the neat mechanism for ml pipeline development. with redco, users only need to write three intuitive functions to define a ml pipeline: a collate function to convert raw data examples into model inputs (e.g., text tokenization); a loss function to execute the model and compute loss (e.g., cross-entropy); and a predict function to run the model and deliver outcomes (e.g., beam search). with the defined pipeline from a user, redco automates all the remaining of pipeline execution such as data parallelism, multi-host related processing, checkpointing, log maintenance, and so forth. we demonstrate this neat mechanism is applicable to various ml paradigms, spanning from basic language modeling and sequence-to-sequence (seq2seq), to more complex algorithms like meta-learning and reinforcement learning. redco-based implementations consistently exhibit substantially fewer lines of code compared to their official counterparts. pipeline development tools. in the development process using neural network libraries such as pytorch @xcite and flax @xcite , certain boilerplate code is consistently present. common operations, such as backpropagation, gradient application, and batch iteration, recur in nearly every ml pipeline. a variety of tools aim to streamline pipeline development by eliminating repetitive code while maintaining as much development flexibility as possible. pytorchlightning @xcite offers a default training loop within pytorch, allowing users to customize their pipelines by inheriting a trainer class and modifying hook functions such as loss function and checkpoint saving. however, this mechanism may not be intuitive for all users. for some people, it requires a learning curve to become comfortable. furthermore, it may be unclear how to implement these hook functions for more complex algorithms, such as federated learning. huggingface-transformers @xcite provides a trainer for pytorch models, but it heavily relies on models defined in its specific transformer classes and primarily focuses on natural language processing pipelines. keras @xcite delivers higher-level apis on top of tensorflow @xcite , enabling users to specify data, model, and loss functions. however, it is not wellsuited for handling complex pipelines. our proposed redco is based on flax, and uses a more intuitive and flexible approach for users to design their pipelines. this mechanism can be applied to a wide array of ml algorithms together with the automatic support for distributed training, including complex algorithms such as federated learning, meta-learning, and reinforcement learning. ## automatic model parallelism for llms model parallelism refers to distributing the computation of a large model across multiple gpus or tpus, in order to address the memory limitations of a single device. two sub-paradigms within model parallelism are pipeline parallelism and tensor parallelism. pipeline parallelism partitions the layers of the model across different devices, and tensor parallelism distributes every tensor in the model across multiple devices. model parallel tools like megatron or alpa require a bunch of intricate configurations and extensively modifying users' code based on the model architecture and the hardware setting. for exam-ple, megatron requires users rewriting the model forward function to customize the tensor sharding for tensor parallelism and annotate breakpoints for pipeline parallelism. this demands substantial mlsys expertise, which is not possessed by most algorithm developers or researchers. in this work, we develop an automatic model parallel strategy in redco that applies across llms without requiring users' mlsys expertise or extra coding efforts. ## redcoast: library design in addition to the complexities of implementing model parallelism, ml pipelines often contain repetitive boilerplate code that demands significant effort from developers. examples of such code include back-propagation, gradient application, batch iteration, and so forth. furthermore, the hardware upgrades usually require patches on existing codebase. for example, a code developed within a single-gpu setting needs data parallelism and multi-host related processing to be added when adapted to multi-gpu machines or clusters. in redco, we design a neat and user-friendly mechanism to simplify ml pipeline developments. users only have to define their pipeline through three design functions, and redco handles all the remaining of the pipeline execution. in this section, we will introduce the software design of redco that enables this mechanism. ## conclusion in this work, we present a lightweight and userfriendly toolkit, redcoast (redco), designed to automate the distributed training of llms and simplify the ml pipeline development process. redco incorporates an automatic model parallelism strategy, fundamentally based on two intuitive rules, without requiring additional coding efforts or ml-sys expertise from the users. we evaluate its effectiveness on an array of llms, such as llama-7b, t5-11b and opt-66b. furthermore, redco has a novel and neat pipeline development mechanism. this mechanism requires users to specify only three intuitive pipeline design functions to implement a distributed ml pipeline. remarkably, this mechanism is general enough to accommodate various ml algorithms and needs significantly fewer lines of code compared to their official implementations.