DeepSpeed 如何从存储层优化推理性能？

全文概览

随着支持AI的应用程序能够释放更多的商业价值，需要支持不同程度的生成式AI（Gen AI），涵盖从数据中心的基础模型训练到边缘和移动设备上的推理部署。

闪存存储，尤其是PCIe/NVMe存储，凭借其高密度和低成本的优势，在实现这一目标方面发挥着关键作用。要为Gen AI启用NVMe卸载，需要结合精心的机器学习（ML）模型设计及其在存储层上的高效部署。

以推理为例，通过Microsoft Deep Speed，突显了NVMe卸载的优势，并指出了NVMe存储可以针对的具体优化和改进方向，以展示改进的大型语言模型（LLM）推理指标。

阅读收获

• 深刻理解大规模训练和推理的重要性及其在当前技术环境中的必要性。
• 全面了解启用大型语言模型（LLM）的NVMe卸载需求，包括技术要求和资源配置。
• 重点关注NVMe闪存在提升LLM推理性能方面的潜在机会，探索其在加速计算过程中的作用。

关于演讲者

Suresh Rajgopal 在北卡罗来纳大学教堂山分校荣获了计算机科学博士学位。自2010年起，他一直效力于Micron Technology公司，致力于固态硬盘（SSD）系统架构、SSD高级封装技术、闪存及SSD控制器的专用集成电路（ASIC）设计，以及面向客户端、企业和云托管的NAND产品研发等领域。他在固态硬盘、NAND闪存、IO扩展器、计算机架构、ASIC实现和设计自动化等领域取得了丰硕的成果，发表了20余篇学术论文，并拥有20多项专利。

Katya Giannios是Micron Technology的机器学习工程师，负责TPG Pathfinding & Strategy Scalable Memory Systems Group。她专注于商业工作负载分析，特别是AI软硬件协同设计的应用。Katya拥有慕尼黑工业大学（TUM）与马克斯-普朗克等离子体物理研究所合作的应用数学博士学位。完成学业后，她转入数据科学领域，在安联全球投资和普渡大学等不同行业积累了丰富经验。

L1-模型训练的效率和成本观察

AI/ML 工作流的拆解，以往整理的材料中已屡见不鲜（Solidigm对AI工作流的梳理最为细致，详见下文），之所以仍摘出这张片子，还是想强调下基础设施建设对于上层应用的重要意义，这和后文要中重点分析的 NVMe SSD 也不谋而合。

• Solidigm：NVMe SSD 在AI存储中的价值

黄底高亮的数据湖和热数据缓存层，是支撑整个AI/ML 工作流的核心，国内头部互联网厂商过去几年，重点攻关建设的差不多，但其他行业的差距仍然非常明显。鹏弟24年和高校数据相关的课题组交流来看，学界对这一层缺乏感知，部分是缺乏海量数据的应用场景，部分是场景解决方案的认知差距。

模型训练的基本常识

基础模型的训练在成本、时间和功率上都有极大的挑战和影响。随着模型规模的扩大，训练所需的资源和成本会显著增加，只有少数公司能够承担这样的开销。因此，基础模型的训练将只能对少数几家公司可行。

--

成本（Cost）:

• GPT-3模型的每次训练成本为500万美元。
• 基础模型训练的成本超过1亿美元。
• 到2027年，最大的模型的训练成本可能会超过10亿美元。

时间（Time）:

• Meta的Llama2 70B模型训练时间为1.7M小时。
• Palm-540B模型的训练时间为8.8M小时。
• GPT-3模型训练所需时间：使用8个V100GPU需要36年，或使用512个GPU时需要7个月。
• GPU的利用率在最佳情况下为50%，通常低于此水平。

功率（Power）:

• GPT-2模型训练消耗28MWh（百万瓦时）的电量。
• GPT-3的功率消耗为GPT-2的10倍，达到284MWh，相当于500多个冰箱每年运行的电力。
• Google刚刚报告，因AI在数据中心的使用，温室气体排放增加了48%。

训练成本趋势

训练的成本正以每年大约增长3倍的速度增加，且这种增长主要受到硬件和能源成本的推动。为了使更多的数据科学家能够访问并进行训练，如何降低这些成本变得越来越重要。图表显示了机器学习系统的训练成本呈指数增长，因此需要采取措施来“民主化”训练，使其不再是少数大型公司或研究机构的特权。

--

左图展示了根据价格-性能趋势，估计训练计算成本（以美元为单位）与机器学习系统发布时间之间的关系（图中显示了带有回归线的散点图），图表中显示124个机器学习系统（不仅限于LLMs，指大规模语言模型）。

L2-数据卸载在训练、推理场景的技术分析

在AI训练中的卸载（NVMe存储，CPU内存）

AI训练需要大量内存和存储来支持模型计算，特别是在处理大规模模型时。为了有效管理这些资源，当前采用的是多GPU内存汇聚和不同类型的并行计算。卸载技术（将数据分配到CPU、CXL和NVMe闪存中）被认为是提高效率、降低成本和能耗的关键。

--

AI训练依赖于将所有与训练相关的数据保持在GPU附近:

• 数据类型:
  • 模型参数（权重和偏置）
  • 优化器状态（训练批次之间）和梯度（参数调整）
  • 检查点数据（中间状态）
  • 工作内存（前向和反向传播时使用）
• 对于一个1T（万亿参数）的模型，GPU需要约30TB内存来训练数据，这被称为“内存墙”（Memory Wall）。
• 这些需求会随着模型规模和上下文大小的增加而增长。

今天:

• 模型扩展依赖于汇总多个GPU内存（跨多个GPU）。
• 3D并行：数据、张量或流水线并行。

卸载:

• 利用AI服务器中的异构性，将训练数据分配到CPU、CXL或NVMe闪存中。
• 有效的卸载可以提供显著的成本和功率优势。

计算卸载的机会

图展示了AI训练和推理中的卸载机会，重点是通过优化数据存储和计算资源管理（如ZeRO Infinity、Flash Neuron等）来提高性能。通过利用GPUDirect和SSD的特定属性，可以显著提高训练和推理的效率。该图强调了如何通过部分卸载和灵活的技术（如FlexGen、Hugging Face）来优化整个过程。

ZeRO (零冗余优化器) 介绍

ZeRO（Zero Redundancy Optimizer） 是由微软的DeepSpeed团队提出的一项技术，旨在优化大规模深度学习模型训练，特别是在处理非常大的模型时。ZeRO技术通过减少冗余数据和提高资源利用效率，显著降低训练过程中对内存和存储的需求，从而使得训练大模型变得更加可行。

ZeRO技术的核心思想是通过分布式训练和高效的数据管理方法来优化模型的训练过程，特别是当模型参数量极为庞大时。以下是ZeRO技术的几个关键特点：

1. ZeRO的基本概念

ZeRO技术的目标是解决大规模深度学习模型训练中的内存瓶颈问题。深度学习训练中，尤其是训练超大规模模型时，通常会遇到模型参数过多、梯度计算和优化步骤需要大量内存的问题。ZeRO解决方案通过分布式训练来拆分模型的不同部分，从而降低每个设备内存负担。

2. ZeRO的三个阶段

ZeRO技术有三个主要的阶段，每个阶段都专注于优化不同资源管理方面。

• ZeRO Stage 1:
  • 参数分散Stage 1将模型参数分散到多个设备上，而不是将所有参数保存在单个设备的内存中。每个设备只保存一部分模型参数，并负责部分计算任务。这可以减少内存占用，使得更大的模型可以在有限的硬件资源下进行训练。
• ZeRO Stage 2:
  • 梯度分散除了分散模型参数外，Stage 2还将梯度信息分散到不同设备上。在训练过程中，梯度是模型优化的重要部分，通过将梯度分散，可以减少内存需求并提高计算效率。
• ZeRO Stage 3:
  • 优化器状态分散在ZeRO的最终阶段，优化器的状态（如动量、学习率等）也被分散到多个设备上。这意味着每个设备不仅仅承担模型参数和梯度计算，还包括优化器相关信息，从而大大降低了内存冗余。

  3. ZeRO Infinity

ZeRO Infinity是ZeRO技术的进一步扩展，它通过引入更多的计算资源和存储资源（例如NVMe和其他高效的存储解决方案）来支持更大规模的模型训练。ZeRO Infinity可以在极大规模的训练任务中进行数据卸载，即将训练数据存储到高效存储介质（如NVMe SSD）上，进一步减轻GPU内存压力。

4. ZeRO的优势

• 内存效率通过减少内存中的冗余数据（如模型参数、梯度和优化器状态），ZeRO显著提高了内存的利用率，从而能够训练更大模型。
• 提高扩展性ZeRO支持分布式训练，能够通过增加设备数量来扩展训练规模，从而训练超大规模模型。
• 成本效益通过优化资源使用，ZeRO降低了硬件成本，使得训练大型模型的门槛降低。
• 更快的训练速度通过减少内存瓶颈，ZeRO使得训练过程更加高效，进而加速了训练过程。

5. 实际应用

ZeRO技术广泛应用于各种大规模深度学习任务，特别是在训练语言模型（如GPT、BERT等）时。通过使用ZeRO技术，研究人员和工程师能够训练更多参数、更大规模的模型，而不必依赖昂贵的硬件资源。

ZeRO 技术原理

关键卸载原则:

• 将分割的模型状态（Pi、Gi、Ai）卸载到CPU-DRAM或NVMe存储：
  • 通过将模型状态分割并卸载到不同的存储设备（如CPU-DRAM或NVMe），可以提高内存和存储资源利用效率。
• 启用并行内存访问：
  • 这种方法允许几乎无限制异构内存带宽，尤其是当使用NVMe时（NVMe的带宽可以达到每秒100GB左右，如DGX-H100）。
• 动态预取（Dynamic Prefetching）：
  • 在计算时，预取操作可以在从NVMe到CPU的数据读取/写入时进行重叠，减少数据传输对训练过程的影响。

图示说明:

• 该图展示了数据流和模型状态的分配。
  • 模型状态分为三类：参数（P₀）、梯度（G₀）和优化器状态（A₀）。
  • 数据通过网络（AllGather和ReduceScatter操作）在不同GPU之间传输。
  • 数据从慢速内存（CPU + NVMe）加载并传输到GPU，以优化训练效率。

结论：

在上述卸载设计下的计算系统，GPU显存将不再是内存瓶颈。可以利用CPU DRAM内存和NVMe存储来微调万亿参数模型！

推理环节拆解

推理阶段的计算工作负载有两部分：预填充阶段依赖计算处理，而解码阶段则更多依赖于内存带宽。

为了提高效率，分区模型参数和KV-Cache卸载是可能的优化方法。然而，将KV-Cache卸载到NVMe存储中可能导致性能下降，特别是在解码阶段，由于对带宽的高需求。

DeepSpeed MSFT ZeRO推理与训练的对比

• 在推理过程中，通过将模型参数卸载到NVMe存储并将KV缓存卸载到CPU内存，可以减轻GPU负担，从而提高推理效率，尤其是在处理大规模模型时。
• 在训练过程中，除了卸载模型参数和优化器状态外，还需要处理梯度和更多模型状态，这使得训练需要更多的资源来进行高效分布式训练。

更多关于微软 DeepSpeed 技术的报道，可参考： Micron：SSD替换主存，加速AI落地

推理场景应用

• 在线推理通常要求低延迟，关键指标是首次响应时间和总延迟，适用于需要实时反馈的应用，如聊天机器人。
• 离线推理更注重吞吐量，适用于不需要即时响应的应用，通过利用NVMe存储可以提升吞吐量效率，适用于例如总结和推荐系统等需要高处理能力的应用。

Note

区分在线和离线推理不同业务场景，对于推动 NVMe SSD 落地有直接帮助，否则找不到场景，试错成本太高。

L3-数据卸载多维度测试数据

测试机配置及参数

• 使用高性能工作站配置，包括多个最新的硬件组件，如Nvidia L40s GPU和Micron 9550 Pro NVMe SSD，来进行ZeRO推理评估。
• 测试了不同规模OPT模型和不同批量大小，以评估在不同条件下的推理性能。
• 通过CPU DRAM卸载和NVMe卸载测试，分别使用PCIe Gen4和Gen5的Micron SSD，评估了存储设备对推理性能的影响。

推理场景的内存卸载对比数据 (参数规模)

• 粉红色：表示CPU DRAM卸载。
  • 吞吐量（Tokens/s）CPU DRAM卸载的吞吐量随着批量大小和模型的增大有明显的提升，但也出现了内存不足的情况（如图中的星标部分），在较大模型（OPT-66b）时运行出现内存溢出。
  • 延迟（Latency）CPU DRAM卸载提供了最低的延迟，尤其是在小型模型（OPT-13b）时。
• 蓝色：表示PCIe Gen4的NVMe卸载。
  • 吞吐量（Tokens/s）随着模型大小的增加，NVMe卸载的吞吐量呈线性增长，表现出较好的扩展性。
  • 延迟（Latency）虽然吞吐量提高，延迟随着模型的增大也逐渐增加。
• 深蓝色：表示PCIe Gen5的NVMe卸载。
  • 吞吐量（Tokens/s）相比Gen4，Gen5的NVMe卸载提供了更高的吞吐量，特别是在处理更大规模的模型时。
  • 延迟（Latency）与Gen4相比，Gen5的延迟更低。

如何理解 Tokens/s 指标在推理场景的参考价值？

Tokens/s（每秒处理的标记数） 是在推理场景中衡量系统性能的一个重要指标，特别是在处理大规模语言模型（如GPT、BERT等）时。这个指标表示在推理过程中，系统每秒能够处理多少个“标记”（Tokens）。

标记（Tokens）是什么？

在自然语言处理中，标记（Token）是模型输入和输出的基本单位。标记可以是：

• 一个单词
• 一个子词（subword）
• 或者是一个字符。

例如，“I love AI”这个句子可能会被拆分为 [I, love, AI] 三个标记。对于不同的语言模型和分词方法，标记的定义和拆分方式可能会有所不同。

Tokens/s在推理中的参考价值

1. 吞吐量的衡量：
   • Tokens/s用来衡量模型每秒能够处理多少个标记。这个指标是推理性能的直接反映，吞吐量越高，表示模型在单位时间内处理的文本或数据量越大。
   • 对于需要处理大量文本输入的场景（如新闻推荐系统、搜索引擎等），较高的吞吐量能够确保系统响应更快，处理更多的请求。
2. 推理系统的效能：
   • 吞吐量（Tokens/s）可以帮助评估系统在给定的硬件配置下，如何应对不同规模的输入数据。吞吐量的提升通常表明系统能够更高效地处理更大的数据量。
   • 在大规模AI应用中，如机器翻译、自动文本生成、问答系统等，吞吐量越高意味着系统可以在短时间内完成更多任务，进而提升服务的响应速度和用户体验。
3. 硬件资源的利用：
   • Tokens/s 也能反映出系统如何有效地利用硬件资源，如GPU的计算能力、内存带宽和存储设备（如NVMe SSD）的效率。随着硬件提升（比如使用更快的NVMe存储或更强的GPU），吞吐量通常也会显著提升。

Tokens/s 是推理性能中非常重要的指标，尤其在处理自然语言模型时。它直接衡量了模型的处理能力，即每秒可以生成或处理多少个标记。在需要高吞吐量的应用场景（如大规模文本处理、实时推荐系统）中，这个指标至关重要。而在一些对响应时间要求更为严格的场景（如聊天机器人、即时翻译）中，吞吐量和延迟的平衡也需要被重点考虑。

对比数据 (Batch Size)

• CPU DRAM卸载在处理批量（128）时提供了较高的吞吐量，但随着批量大小的增加，其性能增长逐渐减缓，最终趋于平稳。
• NVMe卸载特别是Gen5 NVMe卸载，提供了更高的吞吐量，尤其在较大批量的情况下，吞吐量表现更为突出，并且随着批量大小的增大，性能持续提升。
• 结论当批量大小增加时，NVMe卸载能够提高吞吐量，并在高批量的推理任务中发挥更大作用，而CPU DRAM卸载的性能则逐渐趋于饱和。

能耗效率比较

• NVMe卸载在较大的批量大小下，能够显著提高功率效率，减少每生成一个标记所消耗的功率。
• CPU DRAM卸载的功率消耗减少幅度较小，无法像NVMe卸载那样随着批量增大而提升效率。

PCIe Gen4/5 时延对比

改进：

• 28%的延迟减少：从Gen4到Gen5在facebook/opt-30b模型下，解码延迟减少了28%。
• 30%的延迟减少：在facebook/opt-66b模型下，Gen5相比Gen4减少了30%的解码延迟。
• 随着PCIe Gen5的引入，解码延迟显著改善，这与NVMe带宽的提升密切相关。
• Gen5不仅提升了带宽，而且在大规模模型的推理中减少了解码延迟，这直接提升了推理的功率效率，使得系统能够在更短的时间内处理更多的数据。
• 总体而言，随着PCIe Gen5的使用，系统不仅在吞吐量上有所提升，同时解码延迟也有了显著降低，从而提高了整体的推理效率和功率效率。

PCIe Gen4/5 能耗对比

• Micron 9550 Gen5 SSD相比Gen4提供了20%到30%的能效提升，能够显著改善推理性能。
• 在吞吐量方面，Gen5相比Gen4提高了25%，同时在能效上也有了显著改善，特别是在SSD能耗和系统能耗方面，Gen5能够显著降低功耗。
• Gen5的性能提升与能效优化使得其在处理更大模型时，能够提供更高的吞吐量，同时降低能源消耗，提升了整体系统的功率效率。

卸载的成本效益

推理模型：在一个30b参数模型（OPT-30b）上进行推理，批量大小为200。

硬件配置：

无卸载（No Offload）：使用6个L40S GPU，256GB的DRAM。 CPU卸载（CPU Offload）：使用1个L40S GPU，512GB的DRAM。 NVMe卸载（NVMe Offload）：使用1个L40S GPU，256GB的DRAM。

图表解读：

• 无卸载吞吐量相对较低，每1K美元的硬件投资产生的回报较少。
• CPU DRAM卸载回报提升了4.6倍，相比无卸载提高了吞吐量。
• NVMe存储卸载回报提升了4.8倍，表现出更高的效率和吞吐量。

L4-总结

AI大规模部署中的功率和成本考量：

NVMe卸载可以成为一种具有成本效益和功率效率的替代方案： - 支持更大规模的模型，提供更好的响应质量。 - 支持更大的批量大小，在单位时间内提供更多的推理请求，更好的GPU利用率。

卸载库：如 ZeRO Inference 应该被广泛利用。

启用NVMe卸载需要：

• 仔细的模型优化，以隐藏存储延迟并提高计算效率。
• 大块数据的使用和多线程的利用可以进一步加速SSD性能。

Cite

通过支持更大的数据块（如16KB或更大），可以显著提高存储效率，减少写放大因子（WAF），并优化垃圾回收（GC）过程。这种技术在QLC和高容量SSD中尤为重要。 详可参考：Samsung：从QLC应用生态来看大容量SSD前景

AI存储 – 行动呼吁：

转向更快的PCIe接口的SSD：如Gen4、Gen5等。 聚焦读取性能，优化带宽。 了解上层软件栈，以构建容忍延迟的解决方案。