RedCoast：一种在任意GPU/TPU上自动化LLM分布式训练的轻量级工具

AI近期的进展在很大程度上可归因于大语言模型（LLM）。然而，其日益增长的内存需求给机器学习（ML）研究人员与工程师带来了挑战。为此，开发者需将大模型切分并分布到多个GPU或TPU上。使用现有的模型并行工具（如Megatron-LM、DeepSpeed和Alpa）需要大量编码与复杂配置工作。这些工具要求用户具备机器学习系统（MLSys）专业知识，形成了LLM开发的瓶颈，尤其对无MLSys背景的开发者而言。本工作中，我们提出RedCoast（Redco），一种轻量且用户友好的工具，旨在自动化LLM的分布式训练与推理，并简化ML流水线开发。Redco的设计强调两个关键方面。首先，为实现模型并行自动化，本研究识别出两条简单规则，可为任意给定LLM生成张量并行策略。将这些规则集成进Redco可轻松实现分布式LLM训练与推理，无需额外编码或复杂配置。我们通过在GPT-J、LLaMA、T5和OPT等一组LLM架构（最大达66B）上应用Redco证明了其有效性。其次，我们提出一种机制，仅通过定义三个函数即可定制多样化ML流水线，避免了多主机相关处理等冗余与公式化代码。该机制可适配从基础语言建模到元学习与强化学习等复杂算法的广泛ML算法。因此，Redco实现相比官方对应实现具有显著更少的代码行数。