，转载请联系源码转载。从AI绘画到NLP大型模型，AIGC的落地成本一下子降下来！话不多说，直接看结果：StableDiffusion2.0training/fine-tuning/inference最多可节省5.6倍的显存消耗，将硬件成本降低到1/46，一行代码即可启用；1750亿参数大模型BLOOM单机推理，显存消耗节省4倍，硬件成本降低到十分之一。一行代码实现自动搜索最佳并行策略，大幅降低分布式训练门槛，原生支持HuggingFace、Timm等热门AI模型库。要知道在AIGC爆发的另一面，高昂的成本困扰着整个行业。上周，首批人工智能绘画公司之一StockAI被迫宣布关闭其平台。没有别的原因，创始人说：驱动公司的成本太高，目前的收入难以为继。即使在ChatGPT背后有OpenAI和微软的支持，在平台开放几周后就发布公告，开始限制用户的日常使用数量。言下之意无非就是四个字：烧不起。总而言之，降低大型AI模型的实施成本是业界亟待解决的问题。与此同时，开源的AI大规模模型解决方案Colossal-AI在过去一年迅速走红，在GitHub上收获了7k+Stars。比如成本削减计划，就出自它的手。它是如何实现的？往下看，开源地址：https://github.com/hpcaitech/ColossalAIStableDiffusion2.0低成本训练/微调/推理与1.0版本相比，StableDiffusion2.0不仅提高了模型生成的图像像素，还引入了Depth2img模型、文本引导修复模型等，功能更加完备。这一波新的更新，着实让用户感到意外和措手不及。（毕竟1.0还没玩过，没了解过。）不过话又说回来，还是个老问题。实施AIGC模型的成本很高。以StableDiffusion为例，其背后的StabilityAI维护着超过4000个NVIDIAA100GPU集群，为此花费了超过5000万美元的运营成本。面对快速迭代的模型、算法和下游任务，如何降低应用成本成为AIGC真正走向落地的核心问题。StableDiffusion2.0建立在易于使用的PyTorchLightning框架之上。作为PyTorchLightning官方的大模型解决方案，Colossal-AI紧随其后。具体内容如下：显存消耗可节省5.6倍，硬件成本最多可降低到1/46。支持DreamBooth单GPU快速个性化微调推理，显存消耗可节省2.5倍，该方案近期也将合并到HuggingFace中，进一步方便用户使用。训练为了加快训练速度和降低训练成本，使用更大的batchsize成为一种被广泛使用的有效手段。然而，GPU有限的内存容量严重限制了batchsize并推高了训练硬件的门槛。通过一系列内存优化技术，Colossal-AI使StableDiffusion能够在每个GPU上使用largebatchsize16训练，将内存需求从64.5GB降低到11.6GB，节省5.6倍，并且可以扩展到单个GPU或多个GPU并行。与使用最先进的A10080GB相比，仅需3060等消费级显卡即可满足需求，硬件成本最高可降至1/46。这使得更多的用户能够以低成本在消费级GPU上开展StableDiffusion相关的研究和应用。内存优化FlashAttention背后早在StableDiffusion1.0版本，Colossal-AI就率先引入了FlashAttention技术，成功将attention速度提升104%，并将端到端训练的峰值内存降低了23%。FlashAttention是长序列注意力的加速版本。它使用Flatten来减少GPU高带宽内存(HBM)之间的内存读/写次数。Flashattention还为block-likesparseattention设计了一种近似attention算法。比任何现有的近似注意力方法都快。在StableDiffusion1.0版本中，整个DiffusionModel只有少量的attention层，Flashattention还没有表现出性能优势。在StableDiffusion2.0中，由于大量的卷积层被注意力层取代，FlashAttention的内存优化潜力被进一步挖掘。ZeRO+GeminiColossal-AI支持使用零冗余优化器（ZeRO）消除内存冗余的方法，在不牺牲计算粒度和通信效率的情况下，与经典的数据并行策略相比，可以极大地提高内存使用效率。此外，Colossal-AI还引入了Chunk机制，进一步提升ZeRO的性能。操作顺序中的一组连续参数存储在一个Chunk中（一个Chunk是一个连续的内存空间），每个Chunk的大小都是一样的。Chunk方式组织内存可以保证PCI-e和GPU-GPU之间网络带宽的高效利用，减少通信次数，避免潜在的内存碎片。Colossal-AI的异构内存空间管理器Gemini支持将优化器状态从GPU卸载到CPU，以节省GPU内存使用。GPU内存和CPU内存（由CPUDRAM或NVMeSSD内存组成）可以同时使用，突破单个GPU内存墙的限制，进一步扩大可训练模型的规模。一行代码快速上手作为PyTorchLightning官方合作伙伴，只需一行代码即可调用Colossal-AI的上述显存优化。来自lightning.pytorch导入培训师，来自lightning.pytorch.strategies的LightningModule导入ColossalAIStrategyMystrategy=ColossalAIStrategy（use_chunk=True，enable_distributed_storage=True，placement_policy=auto）trainer=Trainer（accelerator="gpu"，strdevices1ino=4，preciateMystrategy）培训师。fit(model)DreamBooth微调在推出StableDiffusion2.0加速方案的同时，Colossal-AI也“顺利”发布了DreamBooth模型的微调方案。这是谷歌今年8月发布的模型。只需要3-5张图片，加上一段文字说明，指定的对象就可以转移到其他场景或风格。与Dall-E2、Imagen等最大的区别在于DreamBooth可以忠实还原选中的对象。方案中，用户只需直接运行文件train_dreambooth_colossalai.py，即可在本次微调任务中充分发挥Colossal-AI的内存优化，以个性化的方式快速微调自己的图形模型。方式，大大降低了使用门槛。推理由于模型推理对数值精度不敏感，这为低精度和低成本推理提供了可能性。对于StableDiffusion2.0模型，可以添加一行代码支持模型的Int8量化推理，节省2.5倍的内存消耗，仅3.1GB，并且不会造成明显的性能损失。model=replace_module(model)使用RTX3090推断1750亿个BLOOM模型。在AI绘画大爆发的另一边，大型NLP模型的趋势仍在继续。今年7月，HuggingFace发布了1750亿参数开源模型BLOOM，使用384个A100进行细化。如果直接使用常见的FP32/FP16进行推理，在单节点8个GPU上使用模型并行，每个GPU至少需要消耗87.5GB/43.8GB显存。如此大的内存占用，即使是最先进的8卡A100（80GB/40GB）服务器也无法直接部署推理服务，多节点推理会带来沉重的额外成本和通信开销。基于此，Colossal-AI实现了高效的Int8量化和模型并行推理，无需3090/4090等消费级显卡即可将BLOOM等大型模型推理服务部署到3090/4090等消费级显卡的8卡服务器上。显着增加CPU内存使用和性能损失。与原来的A100方案相比，硬件部署成本可以降低到原来的十分之一。通过对模型进行Int8量化，Colossal-AI可以将模型的整体内存使用量从352.3GB（FP16）降低到185.6GB，并利用Colossal-AI的模型并行技术将每块显卡的使用量降低到23.2GB。在模型并行中，为了不增加CPU内存占用，Colossal-AI在主进程中进行模型量化和分割，其余各进程使用lazy_init获取几乎不占用显存和显存的元模型，然后通过gloo后端在进程之间传递模型参数。通过上述方案，在不分段加载模型参数的情况下，CPU内存占用峰值可以达到理论上的最优水平。与将模型分层的“流水线式”分布方式相比，模型并行可以提高非密集请求下的显存使用效率。一行代码实现自动并行大模型的分布式混合部署是一个非常复杂的问题。目前常见的分布式大规模模型训练方案，依赖于用户手动反复尝试，系统专家根据经验进行配置部署。但是，这对绝大多数AI开发者来说是很不友好的，因为大家都不想花太多的时间和精力去研究分布式系统和试错。因此，Colossal-AI高效易用的自动并行系统可以说是解决了大家的燃眉之急。只需添加一行代码，即可提供集群信息和单机训练模型，获得分布式训练能力，原生支持HuggingFace、Timm等流行AI模型库。#wrapthemodelusingauto_enginemodel,optimizer=auto_engine(model,optimizer,cluster_info)#normaltrainingloop...因此，Colossal-AI可以大大降低AI开发者使用分布式技术训练和微调大型模型的门槛。同时，自动并行系统可以更细粒度地搜索并行策略，寻找更高效的并行方案。GraphTracingColossal-AI是第一个基于PyTorch框架使用静态图分析的自动并行系统。作为动态图框架，获取其静态执行计划是机器学习系统领域长期研究的问题。Colossal-AI使用基于torch.FXTracer的ColoTracer，推导并记录每个张量在追踪过程中的元信息，如tensorshape、dims、dtype等，可为后续自动并行策略搜索提供辅助。因此，Colossal-AI具有更好的模型泛化能力，而不是依赖模型名称或手动修改来适应并行策略。细粒度分布式训练策略搜索Colossal-AI会在内存预算的约束下，以最快运行时间为目标，对每个op进行策略搜索，最终得到真正的训练策略，包括对每个张量的分割策略，不同计算节点之间需要插入的通信算子类型，是否需要替换算子等现有系统中的Tensor并行和数据并行，以及NVIDIA在并行系统中使用的列分段并行和行分段并行等混合并行，如Megatron-LM，都是可以自动并行搜索的策略子集。除了这些可以手动指定的并行方式外，Colossal-AI的自动并行系统还能够为每个op指定唯一的并行方式，因此有可能找到比依赖专家经验和试验的手动分割更好的并行策略。和错误配置。分布式张量和形状一致性系统类似于PyTorch发布的最新DTensor。Colossal-AI还使用设备网格来抽象集群管理。具体来说，Colossal-AI使用shardingspec标记张量的分布式存储状态，并使用shapeconsistencymanager自动在不同的shardingspec之间转换相同的张量。这极大地提高了Colossal-AI的通用性和易用性。借助shapeconsistencymanager，可以无负担的对tensors进行split，不用担心上游op的输出和下游op的输入在集群中的存储方式不同。与PyTorchDTensor相比，Colossal-AI具有以下三个优势：Colossal-AI的设备网格可以分析到集群性能指标，以及不同通信运营商的耗时估计。Colossal-AI的shapeconsistency会贪婪地寻找shardingspec之间的转换方式，而不是简单的逐维转换，这样可以找到更高效的转换路径，减少shardingspec之间转换的通信开销。添加all_to_all操作使Colossal-AI更具可扩展性，在大规模集群上训练时可以显示出巨大的优势。结合激活检查点作为大型模型训练中不可或缺的内存压缩技术，Colossal-AI还提供了激活检查点的自动搜索功能。与大多数以最大显存压缩为目标的技术方案相比，Colossal-AI的搜索目标是在显存预算内找到最快的激活检查点方案。同时，为了避免将激活检查点的搜索建模到SPMD求解器中造成的搜索时间爆炸，Colossal-AI设计了2阶段搜索，从而搜索到有效可行的分布式训练方案在合理的时间内。关于Colossal-AI通用深度学习系统Colossal-AI面向大模型时代。可实现AI大模型训练和推理的高效快速部署，降低AI大模型应用成本。自开源以来，Colossal-AI多次荣登GitHub热榜全球第一，获得7000多个GitHubStar，并成功入选国际顶级AI和HPC会议官方教程例如SC、AAAI和PPoPP。Colossal-AI相关解决方案已在自动驾驶、云计算、零售、医药、芯片等领域的知名厂商成功落地，获得广泛好评。比如最近爆火的ChatGPT，目前还没有开源，不具备联网功能。Colossal-AI已成功帮助一家财富500强公司开发了具有增强的在线搜索引擎功能的聊天机器人模型。传送门开源地址：https://github.com/hpcaitech/ColossalAI
