Colossal-AIAIGC和自动并行的硬件节约高达46倍

新版Colossal-AI中AIGC和自动并行的硬件节约高达46倍

最近, 大规模人工智能模型获得了极大的关注和采用. AIGC一直很受欢迎, 而聊天机器人ChatGPT获得了广泛的关注, 并在短短两周内吸引了数百万用户, 甚至在谷歌引起了"红色警报", 因为它可能会颠覆搜索巨头的业务. 其他强大的人工智能应用, 如用于软件开发的AlphaCode和用于药物研究的ESM2, 也获得了关注. 这些和其他人工智能应用正在不断推动大型人工智能模型的可能性, 探索新的实施和应用领域.

新的Colossal-AI发布的0.2.0版本包括几个关键功能, 以大大减少实施这种大型模型的成本, 并加快其上市时间, 同时增强开发人员的易用性.

• 稳定扩散2.0模型配方: Colossal-AI中稳定 Diffusion 2.0的开箱即用配置可以实现低成本的训练, 微调和推理, 同时还可以将GPU内存消耗减少5.6倍, 硬件成本减少46倍. 所有这些只需一行代码就可以实现.
• BLOOM模型配方: 现成的包含1750亿个参数的BLOOM模型提供独立的推理, GPU内存消耗减少4倍, 硬件成本减少10倍以上.
• 自动并行化: 只需一行代码, 用户就可以自动搜索最佳并行策略, 使分布式训练更容易, 并原生支持Hugging Face和Timm等流行的AI模型库.

Colossal-AI在全球拥有AWS, Meta, BioMap等用户, 以及超过7,000颗Github星星, 并被证明具有无与伦比的性能, 对于那些希望利用大规模AI力量的人来说, Colossal-AI是一个值得信赖的资源.

稳定扩散2.0的优化功能

AIGC在人工智能领域是一个非常受欢迎的课题, 并被<科学>杂志评为"2022年年度突破". 稳定扩散技术最近也升级到了2.0版本. 一些用户评论说, 这项技术的发展速度很快, 指出第2版在第1版还没有完全开发出来时就已经发布了.

然而, 实施AIGC的高成本在一定程度上限制了其广泛采用. 例如, 稳定扩散背后的稳定AI维持着超过4000个英伟达A100 GPU集群, 并产生了超过5000万美元的运营成本. 随着AIGC继续快速发展, 迭代模型, 算法和下游任务, 降低成本已成为其成功实施的关键问题.

稳定扩散2.0建立在用户友好的PyTorch Lightning框架上, Colossal-AI及时发布了一个全面的配方, 包含开箱即用的开源训练, 微调和推理解决方案, 作为PyTorch Lightning的官方大规模模型解决方案, 效率更高, 硬件要求更低, 该配方包括:

• 将训练的GPU内存消耗减少5.6倍, 硬件成本减少高达46倍.
• 支持DreamBooth, 在单个GPU上快速进行个性化微调.
• 将推理的GPU内存消耗降低了2.5倍.

在不久的将来, 这一解决方案还将被整合到最受欢迎的人工智能模型社区Hugging Face中, 以进一步提升用户体验.

训练

增加批处理量是一种广泛使用的方法, 可以加快训练速度并降低成本.然而, GPU有限的内存容量限制了批量大小, 提高了硬件要求.

Colossal-AI的内存优化技术和对稳定扩散2.0的支持, 将在每个GPU上使用大批量16个训练的稳定扩散的内存需求从64.5GB减少到11.6GB, 减少5.6倍. 这些技术也可以扩展到单个或多个GPU并行使用. 这使得用户只需使用3060这样的消费级显卡就能满足他们的要求, 而不是最先进的A100 80GB, 从而节省硬件成本达46倍. 有了Colossal-AI, 更多的用户可以在消费级GPU上经济实惠地进行稳定扩散相关的研究和实施.

优化战略

Flash Attention

Colossal-AI是第一个为稳定扩散1.0引入Flash Attention技术的公司, 将注意力速度提高了104%, 将端到端训练内存峰值降低了23%. Flash Attention是对长序列任务注意力的有效实现, 它使用flatten来减少GPU高带宽内存(HBM)之间的内存读/写, 并使用近似注意力算法来阻止稀疏的注意力, 使其比现有的近似注意力方法更快. 虽然稳定扩散1.0只有少量的注意力层, 但在稳定扩散2.0中, 通过用注意力层取代许多卷积层, 进一步证明了Flash注意力的内存优化潜力.

ZeRO Gemini

Colossal-AI使用零冗余优化器(ZeRO)来消除内存冗余, 与经典的数据并行相比, 在不牺牲计算粒度和通信效率的情况下, 大大提高了内存使用效率. Colossal-AI还加入了基于Chunk的内存管理, 这进一步提高了ZeRO的性能. 基于Chunk的内存管理将连续的参数集按操作顺序存储在一个称为Chunk的连续内存空间中, 所有的内存大小都是一样的, 以有效利用PCI-e和GPU-GPU之间的网络带宽, 减少通信, 并避免潜在的内存碎片化.

Colossal-AI的异构内存管理器Gemini通过将优化器状态卸载到CPU来减少GPU的内存占用, 允许同时使用GPU内存和CPU内存(包括CPU DRAM或NVMe SSD内存), 以增加可用模型的规模, 超出单个GPU的内存限制.

一行代码就能快速入门

作为PyTorch Lightning的官方合作伙伴, Colossal-AI的内存优化只需一行代码就可以轻松访问.

DreamBooth微调

DreamBooth是一种方法, 只需使用所需主题的3-5张图片, 就可以将文本到图像模型(如稳定扩散)个性化, 并生成该主题的一系列图像. Colossal-AI的内存优化可以通过train_dreambooth_colossalai.py文件轻松访问, 实现快速, 个性化的模型微调, 增强易用性.

推论

由于模型推理对数值精度不敏感, 这为实现低精度的低成本推理提供了可能. 稳定扩散2.0模型可以通过一行代码进行Int8量化推理, 从而使内存消耗降低2.5倍(需要3.1GB内存), 性能损失最小.

model = replace_module(model)

175亿BLOOM模型的低成本推理

随着模型变得更大, 推理过程中的内存消耗成为需要考虑的一个关键因素. 例如, 由Hugging Face发布的开源的1750亿个参数的BLOOM模型, 在单节点的8个GPU上使用普通的FP32/FP16进行推理时, 每个GPU至少需要87.5GB/43.8GB内存. 即使是最先进的8 * A100(80GB/40GB)服务器也无法支持这么大的内存占用, 而且多节点推理会带来额外的成本和通信开销.

Colossal-AI现在提供了一个专用的BLOOM配方, 可以实现高效的Int8量化模型并行推理, 允许在使用3090/4090等消费级显卡的8-GPU服务器上部署像BLOOM(1750亿个参数)这样的大型模型推理服务, 而没有明显的CPU内存占用增加或性能损失. 与使用A100解决方案相比, 这可以*减少10倍以上的硬件部署成本.

通过用Int8量化模型, 使用Colossal-AI, 模型的总体内存占用可以从352.3GB(FP16)减少到185.6GB, 而其模型并行技术将每个GPU的内存需求减少到23.2GB. 模型并行化通过在主进程中使用lazy_init来避免增加CPU的内存占用, 从而获得一个几乎没有内存占用的元模型. Colossal-AI在主进程中对模型进行量化和切片, 并在其余每个进程中使用lazy_init, 以获得几乎不占用内存的元模型, 然后通过gloo后端在进程之间传递模型参数. 这导致CPU内存使用峰值达到理论上的最佳水平, 而不需要分段加载模型参数, 与按层切分模型的"流水线式"分布相比, 在非密集型请求下提高了内存使用效率.

一行代码就能实现自动并行

近年来, 大规模机器学习模型的部署已变得越来越重要. 然而, 分布式训练系统往往需要手动并行化计划, 这可能很复杂, 需要系统工程和配置方面的专家知识. 这对大多数没有必要技能的人工智能开发者来说是一个挑战. 对人工并行化的需求会使大规模机器学习模型的部署变得困难和昂贵.

Colossal-AI的单线自动并行系统简化了AI开发者部署大规模机器学习模型的过程. 与其他需要手动配置复杂的并行策略和模型修改的解决方案相比, Colossal-AI只需要用户的一行代码, 以及集群信息和模型配置, 就可以实现分布式训练. 它可以与流行的人工智能模型框架(如Hugging Face和Timm)无缝集成.

# wrap the model using auto_engine
  model, optimizer = auto_engine(model, optimizer, cluster_info)
#  normal training loop
  ...

新的Colossal-AI版本大大降低了人工智能开发者使用分布式训练技术对大型模型进行预训练和微调的要求. 同时, 自动并行系统可以在更细的颗粒度上搜索最佳的并行策略.

图形追踪

Colossal-AI是第一个自动平行化系统, 使用了基于PyTorch框架的静态图分析. 为PyTorch这个动态图框架获得静态执行计划, 长期以来一直是机器学习系统领域的一个研究领域. Colossal-AI的并行计算系统使用ColoTracer, 即torch.FX Tracer的分叉版本, 来指导寻找最佳的并行化策略. 每个张量的元信息, 如张量的形状, dims, dtype等, 在追踪过程中被计算和记录. 这种方法的优点是具有更好的通用性, 因为它不与特定的模型或配置相联系.

细粒度并行搜索

Colossal-AI的自动并行性搜索关于每个操作数的策略, 目的是在满足内存预算限制的同时实现最快的运行时间. 它最终决定了实际的训练时间策略, 包括每个张量的张量分割策略, 不同计算节点之间插入的通信运算符类型, 是否替换运算符等. 英伟达公司在Megatron-LM中使用的张量, 数据和混合并行, 如列和行分割, 以及其他并行系统, 都是Colossal-AI可以搜索到的策略子集. 除了这些可以手动指定的并行, Colossal-AI还可以为每个操作指定一个独特的并行方法, 并且, 有可能找到比人类专家所能提供的更好的并行策略.

分布式张量和形体一致性系统

Colossal-AI系统使用一个设备网格, 类似于PyTorch最新发布的DTensor, 来管理其集群. Colossal-AI使用sharding-spec来注释每个张量的存储状态, 并促进它们在集群中的分布. 该系统还采用了一个形状一致性管理器, 在不同的分片规格之间自动转换张量, 允许张量的无缝切片和切块, 同时形状一致性管理器确保上游操作数的输出一致地存储在集群中, 无论下游操作数的输入如何存储. 这使得Colossal-AI具有高度的通用性和易用性, 用户在对其进行操作时不必担心张量的存储状态.

以下是Colossal-AI与PyTorch DTensor相比的一些主要优势:

1. Colossal-AI的设备网格使用集群性能指标和剖析结果来估计不同通信操作的时间消耗. 这有助于Colossal-AI优化节点之间的通信, 提高整体系统效率.
2. Colossal-AI的形状一致性管理器使用一种贪婪的搜索算法来寻找在不同分片规格之间转换张量的相对有效的方法, 而不是简单地逐个转换维度. 这可以带来更高效和有效的转换.
3. 在Colossal-AI中整合所有对所有的操作, 通过在节点之间实现更有效的通信, 提高了系统的可扩展性. 这对于需要在节点间传输大量数据的大规模机器学习任务特别有用.

与激活检查点整合

Colossal-AI使用激活检查点来压缩大型机器学习模型训练期间的内存使用. 激活检查点不是存储整个模型, 而是存储中间层的激活. Colossal-AI的激活检查点自动搜索功能在给定的内存预算内找到最有效的检查点, 而不是仅仅为了最大限度地压缩内存. 为了避免漫长的寻找最佳激活检查点的过程, Colossal-AI实施了一个两阶段的搜索过程. 这使得系统能够在合理的时间内找到一个可行的分布式训练方案, 同时仍然受益于激活检查点的内存管理. Colossal-AI中激活检查点的整合提高了大型模型训练的效率和效果.

关于Colossal-AI

Colossal-AI是一个面向大模型时代的统一深度学习系统, 支持高效快速地部署AI大模型训练和推理, 降低AI大模型应用的成本. 自成为开源项目以来, Colossal-AI多次在GitHub趋势榜上排名第一, 星级超过7000. 它已成功被SC, AAAI和PPoPP等国际顶级人工智能和高性能计算会议接受为官方教程.

相关解决方案已被知名科技巨头成功应用于自动驾驶, 云计算, 零售, 医药, 芯片等领域, 并受到广泛好评. 例如, 最近大热的ChatGPT还没有开放源代码, 也没有互联网接入. Colossal-AI已经成功帮助一家财富500强企业开发了具有增强在线搜索引擎功能的聊天机器人模型.