硬件预算最高直降46倍！低成本上手AIGC和千亿大模型，一行代码自动并行，Colossal-AI再升级
最近，AI大模型连续火爆出圈，人工智能生成模型（AIGC）的热度尚未褪去，聊天机器人ChatGPT便引发全网热议，两周吸引百万用户。还有卷趴一半程序员的AlphaCode，生成全新蛋白质的ESM2等，不断探索AI大模型落地的新领域。面对大模型带来的技术革命，连谷歌都拉响“红色警报”，担心ChatGPT砸掉自己搜索引擎的饭碗。
作为当下最火热的开源AI大模型解决方案，Colossal-AI已收获Github Star七千多颗，此前在Stable Diffusion、GPT-3、AlphaFold等大模型上展现卓越性能优势。针对AI大模型落地成本高昂这一痛点，Colossal-AI本次更新聚焦于降低大模型应用成本以及增强易用性，包括:
Stable Diffusion 2.0低成本训练/微调/推理，最多可降低显存消耗5.6倍，使硬件成本直降46倍，一行代码即可使用;
1750亿参数BLOOM模型单机推理，显存消耗降低4倍，使硬件成本降低10余倍;
一行代码实现自动搜索最佳并行策略，显著降低分布式训练上手门槛，原生支持Hugging Face，Timm等热门AI模型库。
开源地址：https://github.com/hpcaitech/ColossalAI
Stable Diffusion 2.0低成本训练/微调/推理
AIGC（人工智能生成内容）目前是AI领域最热门的话题之一，还被Science评为2022年度十大科学突破。短短数月，爆红AI圈的Stable Diffusion已升级到2.0版本，正如部分用户所言：“我从未见过任何技术发展如此之快。V1还没整透彻，V2就来了”。
然而，高昂的成本也一定程度上限制了AIGC产业的大规模普及。以Stable Diffusion为例，其背后的Stability AI维护超过 4000 个英伟达 A100 的 GPU 集群，并已为此支出超过 5000 万美元的运营成本。面对快速迭代的模型、算法和下游任务，如何降低应用成本成为AIGC真正走向落地的核心问题。
Stable Diffusion 2.0基于简单易用的PyTorch Lightning框架搭建。作为PyTorch Lightning的官方大模型解决方案，Colossal-AI在第一时间跟进，开源了更为高效，硬件门槛更低的全套训练/微调/推理方案:
可减少训练显存消耗5.6倍，硬件成本最高降低46倍；
支持DreamBooth单GPU快速个性化微调；
推理显存消耗降低2.5倍。
该方案也将于近期合并进入当下最火热的AI模型社区Hugging Face，进一步方便用户使用。
训练
为了加快训练速度，降低训练成本，使用更大的batch size已成为被广泛使用的有效手段。但GPU有限的显存容量，严重限制了batch size大小，推高了训练硬件门槛。
Colossal-AI通过一系列显存优化技术和支持Stable Diffusion 2.0，使Stable Diffusion平均在每个GPU上使用大batch size 16训练的显存需求从64.5GB降低5.6倍至11.6GB，还可扩展至单GPU或多GPU并行。相比使用最先进的A100 80GB，目前仅需3060等消费级显卡即可满足需求，硬件成本最高直降46倍。更多用户可以在消费级GPU上，低成本地开展Stable Diffusion的相关研究与应用落地。
背后显存优化
Flash Attention：早在Stable Diffusion 1.0版本，Colossal-AI就率先引入了Flash Attention技术，成功将attention的速度提升 104%，将端到端训练的峰值显存减少 23%。Flash Attention是针对长序列 attention 的加速版本，使用 Flatten 来减少 GPU 高带宽内存（HBM）之间的内存读 / 写次数, Flash attention 同时针对块状稀疏的 attention，设计了一个近似的注意力算法，比任何现有的近似 attention 方法都要快。在Stable Diffusion 1.0版本，整个Diffusion Model只有少量attention层，Flash attention还没有体现出其性能优势。在Stable Diffusion 2.0中，由于将大量卷积层替换为attention层，进一步发挥了Flash Attention的显存优化潜力。
ZeRO + Gemini：Colossal-AI支持使用零冗余优化器 (ZeRO)的方法来消除内存冗余，与经典的数据并行性策略相比，可极大提高内存使用效率，同时不牺牲计算粒度和通信效率。此外，Colossal-AI 还引入了Chunk机制进一步提升ZeRO的性能。运算顺序上连续的一组参数存入一个Chunk中（Chunk即一段连续的内存空间），每个Chunk的大小相同。Chunk方式组织内存可以保证PCI-e和GPU-GPU之间网络带宽的高效利用，减小了通信次数，同时避免潜在的内存碎片。
而Colossal-AI的异构内存空间管理器Gemini支持将优化器状态从GPU卸载到CPU，以节省GPU内存占用。可以同时利用GPU内存、CPU内存（由CPU DRAM或NVMe SSD内存组成）来突破单GPU内存墙的限制，进一步扩展了可训练模型规模。
一行代码快速上手
作为PyTorch Lightning的官方合作伙伴，仅需一行代码即可调用Colossal-AI的上述显存优化。
from lightning.pytorch import trainer, LightningModule
from lightning.pytorch.strategies import ColossalAIStrategy
Mystrategy = ColossalAIStrategy(use_chunk=True, enable_distributed_storage=True, placement_policy=auto)
trainer = Trainer(accelerator="gpu", devices=4, precision=16, strategy=Mystrategy)
trainer.fit(model)DreamBooth微调
DreamBooth是一种用于个性化文本到图像模型（例如Stable Diffusion）的方法，只需要几张（3~5张）期望主题的图像，即可生成指定主题的一系列图像。用户只需直接运行文件train_dreambooth_colossalai.py，即可在该微调任务上充分发挥Colossal-AI的显存优化，个性化快速微调自己的图文模型，极大地降低了使用门槛。
推理
由于模型推理对数值精度不敏感，这为实现低精度的低成本推理提供了可能。对于Stable Diffusion 2.0模型，可以通过添加一行代码，支持模型的Int8量化推理，显存消耗降低2.5倍，仅需3.1GB，且不造成显著性能损失。
model = replace_module(model)
1750亿BLOOM模型低成本推理
随着模型规模的不断增大，模型推理的显存及内存占用也变成了不可忽视的关键因素。以Hugging Face发布的1750亿参数开源模型BLOOM为例，如果直接使用常见的FP32/FP16进行推理，在单节点8张GPU使用模型并行，每张GPU需要消耗至少87.5GB/43.8GB的显存。如此大的显存占用，即使是最先进的8卡A100（80GB/40GB）服务器，也无法直接部署推理服务，而多节点推理又会带来沉重的额外成本和通信开销。
Colossal-AI实现了高效的Int8量化和模型并行推理，可以将1750亿参数的BLOOM等大模型的推理服务，部署到3090/4090等消费级显卡的8卡服务器，同时不产生显著的CPU内存占用提升及性能损耗。相比原有的A100方案，可降低硬件部署成本十余倍。
在Colossal-AI中，通过对模型进行了Int8量化，可将模型总体显存占用从352.3GB（FP16）降低到185.6GB， 同时使用Colossal-AI的模型并行技术，将每张显卡的占用减少到了23.2GB。在模型并行中，为了不增加CPU内存占用，Colossal-AI在主进程中进行模型的量化和切分工作，其余的每个进程中分别使用lazy_init获得几乎不占显存和内存的meta model，再通过gloo backend在进程之间传递模型参数。通过上述方案，在没有分段加载模型参数的情况下，便可以使得CPU内存占用峰值，达到理论上的较优水平。相较于将模型按层切分的“类流水线”分布方式，模型并行可以提高非密集请求下的显存使用效率。
一行代码自动并行
大模型的分布式混合部署是一个非常复杂的问题，目前常见的分布式大模型训练方案，都依赖用户人工反复尝试以及系统专家的经验来进行配置部署。然而这对绝大多数AI开发者来说十分不友好，因为不希望将时间精力花费在研究分布式系统和试错上。
Colossal-AI提供了业界急需的一套高效易用自动并行系统。相比现有其他手动配置复杂并行策略和修改模型的解决方案，Colossal-AI仅需增加一行代码，提供cluster信息以及单机训练模型即可获得分布式训练能力，并且原生支持包括Hugging Face，Timm等热门AI模型库。
# wrap the model using auto_engine
model, optimizer = auto_engine(model, optimizer, cluster_info)
# normal training loop
...因此，Colossal-AI可以极大地降低AI开发者的使用分布式技术训练和微调大模型门槛。同时，自动并行系统可以从更细粒度搜索并行策略，找到更加高效的并行方案。
Graph Tracing
Colossal-AI是首个基于PyTorch框架使用静态图分析的自动并行系统。PyTorch作为一个动态图框架，获取其静态的执行计划是机器学习系统领域被长期研究的问题。Colossal-AI使用基于torch.FX Tracer的ColoTracer，在tracing过程中推导并记录了每个tensor的元信息，例如tensor shape，dims，dtype等，可以为后续的自动并行策略搜索提供帮助。因此Colossal-AI具有更好的模型泛化能力，而不是依靠模型名或手动修改来适配并行策略。
细粒度分布式训练策略搜索
Colossal-AI会在满足内存预算的限制下，以最快运行时间为目标，为每个op进行策略搜索，最终得到真实训练时的策略，包括每个tensor的切分策略，不同计算节点间需要插入的通信算子类型，是否要进行算子替换等。现有系统中的张量并行，数据并行，NVIDIA在Megatron-LM等并行系统中使用的column切分和row切分并行等混合并行，都是自动并行可以搜索到的策略的子集。除了这些可以手动指定的并行方式外，Colossal-AI的自动并行系统有能力为每个op指定独特的并行方式，因此有可能找到比依赖专家经验和试错配置的手动切分更好的并行策略。
分布式tensor与shape consistency系统
与PyTorch最新发布的DTensor类似，Colossal-AI也使用了device mesh对集群进行了抽象管理。具体来说，Colossal-AI使用sharding spec对tensor的分布式存储状态进行标注，使用shape consistency manager自动地对同一tensor在不同sharding spec间进行转换。这让Colossal-AI的通用性和易用性极大地提升，借助shape consistency manager可以没有负担地切分tensor，而不用担心上游op的output与下游的input在集群中的存储方式不同。
相较于PyTorch DTensor，Colossal-AI有以下3个优势：
Colossal-AI的device mesh可以profiling到集群性能指标，对不同的通信算子进行耗时估算。
Colossal-AI的shape consistency会贪心地搜索sharding spec间的转换方式，而不是朴素地逐dimension进行转换，这样能找到更高效的转换路径，进而使得sharding spec间的转换通信开销更小。
加入了all_to_all操作，使得Colossal-AI的扩展性更强，这在大规模集群上进行训练时，可以展现出很大的优势。
与activation checkpoint结合
作为大模型训练中必不可少的显存压缩技术，Colossal-AI也提供了对于activation checkpoint的自动搜索功能。相比于大部分将最大显存压缩作为目标的技术方案，Colossal-AI的搜索目标是在显存预算以内，找到最快的activation checkpoint方案。同时，为了避免将activation checkpoint的搜索一起建模到SPMD solver中导致搜索时间爆炸，Colossal-AI做了2-stage search的设计，因此可以在合理的时间内搜索到有效可行的分布式训练方案。
关于Colossal-AI
面向大模型时代的通用深度学习系统Colossal-AI，可实现高效快速部署AI大模型训练和推理，降低AI大模型应用成本。自开源以来，Colossal-AI已经多次在GitHub热榜位列世界第一，获得GitHub Star超七千颗，并成功入选SC、AAAI、PPoPP等国际AI与HPC顶级会议的官方教程。
Colossal-AI相关解决方案已成功在自动驾驶、云计算、零售、医药、芯片等行业知名厂商落地应用，广受好评。例如，近期火爆的ChatGPT尚未开源，且不具备联网功能。Colossal-AI已成功帮助某世界500强企业，开发具备在线搜索引擎能力增强的聊天机器人模型。
开源地址：https://github.com/hpcaitech/ColossalAI
参考链接：https://www.hpc-ai.tech/blog/colossal-ai-0-2-0