MosaicML-面向生成式AI的机器学习平台

前段时间，大数据巨头 Databricks 宣布已签署最终协议，将以13亿美元的价格，收购位于旧金山的人工智能初创公司MosaicML 。这篇文章来自 MosaicML官方的技术博客，是对 MosaicML 大模型训练平台的一个简单介绍。
AIGC领域最大收购：Databricks 13亿美元买下MosaicML，成立仅2年员工60人-36氪
36kr.com/p/2319199142134407

https://www.mosaicml.com/blog/train-custom-gpt-diffusion-models
www.mosaicml.com/blog/train-custom-gpt-diffusion-models
MosaicML 平台旨在应对训练 ChatGPT、LaMDA 和 Stable Diffusion 等大模型的挑战。这篇文章详细分析了训练这些模型的困难，并展示了 MosaicML 平台如何使 AI 大模型的训练变得更加轻松。

如今，像 ChatGPT、LaMDA 和 Stable Diffusion 这样的 AI 大模型为初创企业和成熟企业提供了新的发展或者腾飞机会。

然而，对于许多公司和机构来说，训练这些模型过于复杂和昂贵，需要特定的专业知识和工具。因此，只有少数几家公司拥有训练这些模型的能力。我们建立了 MosaicML 平台，以使大规模的模型训练更加容易。现在，各种规模的公司和机构都可以训练自己行业的特定大模型，拥有完整的模型所有权，并且不用担心数据隐私问题。

AI大模型训练面临的挑战

基础设施方面的挑战：如GPU可用性不足、软件工具栈的复杂性、数百个GPU的扩展以及容错能力，使得训练大模型变得困难
什么使得训练大模型如此困难？

GPU 的可用性

训练大模型需要大量先进的 GPU。举例来说，Meta 公司的拥有 1750 亿参数的 OPT-175 模型在 1024 个 NVIDIA A100 GPU 上进行了为期 33 天的训练。类似地，谷歌的拥有 1370 亿参数的 LaMDA 模型在 1024 个 TPU 上进行了长达 57 天的训练。

GPU 往往一卡难求。并且，一旦选择了某家云服务商并最终获得所需的 GPU 访问权限，专有工具如 AWS SageMaker Model Parallel Library 会使您难以将模型训练迁移到另一家云服务商（微软云或者谷歌云）。

工具栈的复杂性

一旦有了 GPU 集群，就需要一个工具栈来协调数千个 GPU 的分布式训练任务。以下是一个大模型工具栈的简介（为简单起见，省略了与 ML 训练无关的部分）。

一个大模型工具栈的简介
配置这种技术工具栈中每一部分的组件都是棘手且容易出错的。比如，错误配置 InfiniBand 驱动程序，可能会导致训练速度减慢 5 倍！

横向扩展

训练拥有数十亿参数的大模型往往需要在至少数百个 GPU 上进行操作，以使大模型能够放进 GPU 显存并在合理的时间内完成训练。

在如此多的 GPU 上进行分布式训练是一项具有挑战性的任务，它涉及到需要选择合适的并行策略（数据并行、张量并行和流水线并行），并集成实现对应并行策略的工具库（DeepSpped 或者 Megatron），调整训练的超参数如 global batch size 等，并确保不会因为 GPU 显存不足或其他错误导致训练崩溃。

大模型分布式训练是非常脆弱的，一个错误的选择可能导致运行失败或生成无用的模型。

容错性

部署软件工具栈和配置分布式训练只是个开始。在庞大的计算基础设施上运行如此复杂的软件工具栈涉及许多操作上的挑战，需要处理各种奇怪的故障。

在 Meta AI 的 OPT-175B 训练日志中，他们描述了训练过程中发生的 100 多次重启。他们的日志揭示了训练过程中遇到的运维挑战：

2021年5月12日…从 checkpoint_15_45000 恢复训练需要 50 分钟！ 2021年11月18日…无法连续训练超过 1-2 天… 许多故障需要手动检测和修复，浪费了计算资源和研究人员的时间…"
清华大学 KEG 团队负责训练大语言模型 GLM-130B，也分享了类似的经验：

2022年1月…频繁的随机硬件故障… 2022年3月…无法启动超过 2,000 个计算节点 => 通过调整 Linux 内核 TCP 参数克服这一问题并支持 6,000 个节点的训练…
正如上文所示（以及我们自己的经验所学），训练大语言模型极易出现诸如 GPU 或网络故障、软件工具中的错误和错误配置、影响模型收敛的loss波动等错误。这些错误会导致显著的延迟，并消耗大量宝贵的时间和金钱进行调试和故障排除。

MosaicML Platform: 为 AI 大模型训练而生
为了解决这些挑战，我们建立了 MosaicML 平台，该平台已被斯坦福大学 CRFM 和 StabilityAI 等组织使用，用于在各自安全环境中使用自己的数据训练 LLMs。

2.1 一个稳定可靠的大模型工具栈
我们从头开始设计了 MosaicML 平台，以解决训练工具栈内各个部分的问题：从驱动程序和工具包一直到 K8S 作业编排和分布式训练。
在这里插入图片描述

MosaicML 平台提供了一系列大模型工具栈
我们的平台提供了必要的基础设施，可以部署和编排跨任意数量 GPU 的训练作业，同时处理故障检测和自动恢复，实现快速恢复。它还包括一个 Runtime，提供经过性能测试的 Docker 镜像，配置了最新的 GPU 驱动程序，一个分布式训练框架（Composer）和一个流式数据加载器（StreamingDataset），可与多种数据格式和云存储配合使用。

我们不断更新和测试 MosaicML 平台，以确保提供的工具栈内所有部分都经过优化，并能无缝协同工作。

2.2 在多云上进行训练
MosaicML 平台通过使训练作业能够在任何主要云服务商上运行而无需更改代码，从而使客户摆脱了对特定云服务商的依赖。我们使客户能够在他们自己的云环境中训练模型，以确保我们永远不会看到他们的训练数据。

为了实现在多个云上的无缝使用和部署，我们设计了具有三部分架构的平台：客户端界面、控制平面和计算平面。

The MosaicML platform has three parts: the client interfaces, the control plane, and the compute plane.
客户端界面
在这里插入图片描述

包括用于启动和管理训练作业的 Python API 和命令行界面。
提供 Web 控制台以管理访问、用户和团队，以及进行使用账户、配额管理和计费。
控制界面

托管 MosaicML 多节点编排、故障检测和恢复逻辑。
编排跨本地集群或云服务商（如 AWS、OCI、GCP、Azure 和 CoreWeave）的训练。
管理应用逻辑和元数据，如运行配置、日志等。
计算界面

支持分布式训练。
与云服务商无关：可部署到任何 Kubernetes 集群，包括本地。
保护数据隐私：客户的训练数据永远不会离开他们的虚拟私有云（VPC）。
无状态且无需在不同集群之间同步数据集和代码。数据集在流式传输的同时，用户代码、Docker 镜像和其他对象会在运行时动态加载。
借助 MosaicML平台，在不同云服务商间进行训练就像在运行提交中更改一个参数一样简单。以下是使用MosaicML CLI（MCLI）的示例：

动图封面
Change cloud providers with a single parameter in your run submission.
2.3 无缝扩展训练模型的规模
跨多个节点和更大模型规模的训练扩展不仅仅是生成一堆作业。为了使扩展“正常运作”，我们解决了整个工具栈中的一系列挑战，包括编排、优化的并行配置和稳定的数据加载模块。我们使客户能够在少数几个 GPU 上运行实验，然后扩展到数百个GPU进行完整的模型训练，无需进行代码更改！

动图
Scale from 1B to 30B parameters with our optimized configurations.
我们使用 Kubernetes 构建了 MosaicML 编排工具，因为它使我们能够实现多云操作，并且非常适合管理大量计算工作负载。然而，Kubernetes 缺乏 LLM 训练所需的功能，如跨节点（“gang-scheduling”）和跨集群调度。为解决此问题，我们实现了自己的调度程序，以优化大规模训练作业的编排。

在 MosaicML 调度到作业后，MosaicML 训练运行时会进行以下操作：

通过配置优化网络和设置多进程来处理在多个节点间分配作业。
自动确定如何将模型和数据批次（data batch）适应 GPU 显存，并实现自动梯度累积，无需调整 batch size 以避免 CUDA 内存溢出错误。
使用我们的流式数据集处理大规模的流式训练数据 - 只需指向您的云存储桶并开始训练。

Auto-gradient accumulation figures out how much of the data batch can fit in memory at once.
2.4 运行监控：自动故障检测与快速恢复
训练大模型容易出现错误，经常会遇到从硬件崩溃到 NCCL 超时再到 loss 激增等各种失败。我们在整个工具栈上进行全面的监控，快速检测问题并使用自动优雅恢复从最新的 Checkpoint 重新开始训练，让用户无需盯着监视并手动重新运行失败的作业。

重新启动作业涉及几个耗时（和浪费金钱！）的步骤，比如重新拉取 Docke r镜像、下载最新的 Checkpoint 和重新运行数据加载器。为了最大程度减少 GPU 空闲时间并在从故障中恢复时为客户节省成本，我们采用了几种加速作业恢复的技术，比如 Docker 镜像预取和缓存以及数据加载器的 fast-forwarding。

MosaicML 平台会自动处理 LLM 训练的所涉及到的所有方面，让我们的客户可以专注于训练他们的模型。

使用 MosaicML
使用MosaicML平台，训练大语言模型就像将我们的平台指向一个 S3 存储桶并启动一个“训练运行”一样简单。
首先，安装一些SDK工具。

$ pip install mosaicml-cli
Collecting mosaicml-cli
…
Successfully installed mosaicml-cli-0.2.36
2. 安装 SDK 包到本地后，从 MosaicML 控制台获取您的 API 密钥，然后使用 Python SDK 在本地环境中设置您的 API 密钥。

import mcli.sdk as msdk
msdk.initialize(api_key=‘your-api-key’)
3. 接下来，您需要配置用于训练所需的集成。这些集成包括 GitHub 用于将代码拉取到训练集群；对象存储，比如 AWS S3 用于训练数据、访问和 Checkpoint 的存储；容器注册表，比如 Docker Hub 用于拉取容器镜像；以及实验跟踪工具，比如 CometML 或 Weights & Biases。了解更多支持的集成和如何在我们的文档中设置它们。

We integrate with all of your favorite tools.
4. 添加了所有配置后，现在是时候使用 MCLI 或者 Python SDK 启动训练任务了。下面的代码片段使用 Python SDK 和这个示例配置。

from mcli.api.runs import RunConfig, create_run