StableDiffusion，风靡全球的预训练微调成本，正式发布！现在，预训练速度提升至6.5倍，成本降低85%。仅RTX2070/3050即可微调，硬件成本降低近90%。这是国内项目Colossal-AI最新的开源加速方案。开源地址：https://github.com/hpcaitech/ColossalAI今年以来，StableDiffusion可谓圈内外火爆，近期更是凭借免费开源、效果好的特点，更进一步走出圈子，让人看到了AIGC的巨大潜力和前景。艺术圈、AI圈的大咖争先恐后地用它来创作，GitHubstar数已达33.5k，背后公司的估值更是飙升至10亿美元。StableDiffusion也因此成为扩散模型中最受欢迎的项目。BUT，它唯一的缺点就是训练成本实在是太高了！只有StableDiffusionv1版本的模型单次训练需要150,000A100GPUHours；微调至少需要RTX3090/4090等高端消费级显卡。其背后的公司StabilityAI为了维持运营和算力需求，拥有一个由4000个A100组成的GPU集群，仅成本就超过5000万美元（折合人民币约3.6亿元）。事实上，由于扩散模型本身特殊的原理和机制，对硬件内存和计算速度都有很高的要求，所用的数据库也在千亿量级。因此，要提高训练速度和降低成本，确实不是一件容易的事。那么，Colossal-AI究竟是如何实现的呢？从显存说起，先说明为什么扩散模型对硬件要求高。扩散模型主要包括前向扩散和后向生成两个过程。前向扩散就是逐渐给图像加上高斯噪声，直到变成随机噪声。反向生成是一个去噪过程。使用多个U-Net逐渐对随机噪声进行去噪，直到生成图像。这也是扩散模型的训练过程。与传统的端到端深度学习模型相比，扩散模型的训练过程显然更加复杂。其中，除了扩散模型本身，StableDiffusion还包括一个FrozenCLIPTextcoder来输入文本提示（textprompts），一个Autoencoder将高分辨率图像压缩到一个潜在空间（LatentSpace），它还需要每A时间步添加一次以计算损失。这样，训练方案对内存开销和计算速度提出了很大的挑战。针对以上问题，Colossal-AI从预训练和微调两个步骤进行了优化。在预训练阶段，主要要做的是减少内存开销。一般来说，batchsize越大，训练速度越快，扩散模型也是如此。通过ZeRO、Gemini、基于Chunk的内存管理等策略，以及优化Cross-attention计算的FlashAttention模块，Colossal-AI允许用户在8GB显存的普通消费级显卡上训练扩散模型（例如RTX2070/3050）。如果是在A100上，单卡BatchSize可以直接扩展到256，相比Stable-diffusion-v1FP32的DistributedDataParallel(DDP)，训练速度可以提升6.5倍。因此，培训时间可以大大减少，从而减少数百万美元的培训成本。在微调方面，对于大多数AIGC玩家来说，他们倾向于选择使用开源的预训练模型权重来微调个性化的下游任务。一方面是由于扩散模型本身的复杂性，另一方面，StableDiffusion预训练使用的是LAION-5B数据集，包含5850亿个图文对，需要240TB的存储空间空间。然而，在现有的很多开源微调方案中，使用的训练并行方式主要是DDP，导致训练过程中占用大量显存。即使微调也至少需要最高端的消费级显卡，如RTX3090/4090。另外，现阶段很多开源训练框架并没有提供完整的训练配置和脚本，需要大家自己做一些补全和调试。基于现状，Colossal-AI提出了第一个同时完整的训练配置参数和训练脚本的开源方案。一方面，该方案让用户不再局限于现有的权重推理，更灵活地为下游新任务训练最新版本的细分模型，扩大应用范围；另一方面，由于内存优化等技术的引入，该方案支持在普通个人电脑的单显卡（如GeForceRTX2070/30508GB）上进行快速微调任务。要知道，目前RTX2070的价格比3090/4090便宜了近90%。此外，对于速度不敏感的任务，还可以进一步使用Colossal-AINVMe，即使用低成本的硬盘空间来减少显存消耗。什么是优化作弊？那么使用了哪些技术来降低内存开销呢？主要基于三个方面：ZeRO+GeminiFlashAttentionFP16、acitvationcheckpoint等常用优化技术我们先来看ZeRO+Gemini。ZeRO是一个零冗余优化器。该方法主要通过对优化器状态、梯度和模型参数进行分段，从而消除数据并行和模型并行中的内存冗余，使GPU只保存当前计算所需的部分。与经典的数据并行策略相比，在不牺牲计算粒度和通信效率的情况下，可以大大提高内存使用效率。同时，Colossal-AI还引入了Chunk机制，进一步提升ZeRO的性能。将一组运算顺序连续的参数存储到一个Chunk中（一个Chunk是一个连续的内存空间），每个Chunk的大小相同。Chunk方式组织内存可以保证PCI-e和GPU-GPU之间网络带宽的高效利用，减少通信次数，避免潜在的内存碎片。△Chunk机制Gemini是Colossal-AI的异构内存空间管理器。它支持将优化器从GPU卸载到CPU以节省GPU内存占用。通过打破GPU内存墙的限制，可以同时利用GPU内存和CPU内存（由CPUDRAM或NVMeSSD内存组成），从而进一步扩大可训练模型的规模。△通过ZeRO+Gemini提升硬件的模型容量第二个方面是FlashAttention机制。LDM（LatentDiffusionModels）是一种提高扩散模型效率的方法。它通过在模型框架中引入交叉注意力机制（Cross-Attention）来实现多模态训练，使得扩散模型可以更灵活地实现对class-condition、text-to-image、layout-to-image的支持.但是，Cross-Attention层相对于原始扩散模型的CNN层增加了额外的计算开销，这大大增加了训练成本。相比之下，FlashAttention是针对长序列的Attention的加速版本。它使用Flatten来减少GPU高带宽内存(HBM)之间的内存读/写次数。FlashAttention还为block-likesparseAttention设计了一种近似的attention算法，比现有的任何attention方法都快。可以将Attention的速度提升104%，将end-to-endtraining的peakmemory降低23%。△FlashAttention除了以上两方面，Colossal-AI还集成了FP16、activationcheckpoint等常用优化技术。例如，通过用计算换取内存来激活检查点。它避免了存储整个计算图的所有中间激活用于反向计算，并且不在检查点部分保存中间激活，而是在反向传递中重新计算它们，进一步减少显存。而FP16在不影响精度的情况下，将原来的32位浮点计算转换为16位，减少了显存的使用，提高了计算效率。它也非常容易上手。在实际操作中，Colossal-AI只需少量代码即可快速启动扩散模型训练。与常见的PyTorch开源项目不同，StableDiffusion是基于PyTorchLighting构建的。PyTorchLightning为流行的深度学习框架PyTorch提供简单易用、灵活高效的高层接口，为AI研究人员提供简单易用的高层抽象，让深度学习实验更简单阅读和重现。Colossal-AI已经集成为PyTorchLightning的官方大模型解决方案。以StableDiffusion为例，启动代码如下：fromcolossalai.nn.optimizerimportHybridAdamfromlightning.pytorchimporttrainerclassMyDiffuser(LightningModule):...defconfigure_sharded_model(self)->None:#createyourmodelhereself.model=construct_diffuser_model(...)...defconfigure_optimizers(self):#使用指定的优化器optimizer=HybridAdam(self.model.parameters(),self.lr)...model=MyDiffuser()trainer=Trainer(accelerator="gpu",devices=1,precisenotallow=16,strategy="colossalai")trainer.fit(model)另外Colossal-AI和PyTorchLightning也对OPT、HuggingFace等流行的算法提供了很好的支持和优化模型和社区。低成本微调在微调方面，Colossal-AI提供了基于HuggingFace上开源StableDiffusion模型的权重进行微调的功能。用户只需简单修改Dataloader加载自己的微调数据集并读取预训练权重，然后在个人电脑上修改参数配置yaml文件并运行训练脚本即可微调自己的个性化模型。模型：目标：ldm.models.diffusion.ddpm.LatentDiffusion参数：your_sub_module_config：目标：your.model.import.path参数：from_pretrained：'your_file_path/unet/diffusion_pytorch_model.bin'...闪电：训练师：策略：目标：pytorch_lightning.strategies.ColossalAIStrategy参数：...pythonmain.py--logdir/your_log_dir-t-bconfig/train_colossalai.yaml快速推理Colossal-AI还支持原生稳定扩散推理管道。训练或微调后，只需要直接调用diffuser库，加载自己保存的模型参数，就可以直接进行inference，无需其他改动。便于新用户熟悉推理过程，也便于习惯使用原有框架的用户快速上手。fromdiffusersimportStableDiffusionPipelinepipe=StableDiffusionPipeline.from_pretrained("your_ColoDiffusion_checkpoint_path").to("cuda")image=pipe('你的提示',num_inference_steps=50)["sample"][0]image.save('文件路径')△OneMoreThing，上述推理过程产生的工作实际上，上面提到的Colossal-AI不仅可以用来加速训练扩散模型，而是面向各种大型模型的。Colossal-AI通过多维自动并行、异构内存管理、大规模优化库、自适应任务调度等方式，高效快速部署AI大模型训练和推理，降低应用成本。自开源以来，Colossal-AI在GitHub和PaperWithCode热榜上多次位居全球第一。经过国际专家的严格评审，Colossal-AI已成功入选SC、AAAI、PPoPP等国际顶级AI和HPC会议的官方课程。目前，其相关解决方案已成功应用于自动驾驶、云计算、零售、医药、芯片等领域的知名厂商。例如，针对生物医药行业的蛋白质结构预测模型AlphaFold，基于Colossal-AI的优化方案，FastFold成功将单个GPU的合理氨基酸序列长度扩展到10000，覆盖99.9999%的蛋白质。仅使用笔记本电脑上的消费级显卡即可解析90%的蛋白质。而且，它还可以并行加速训练和推理的全过程，帮助很多新药研发企业缩短了研发流程，降低了研发成本。传送门开源地址：https://github.com/hpcaitech/ColossalAI
