www.ydisp.cn/oss/202207/13/d85680468180b2c1952886efaf6237186??48f3d.png"alt="1个GPU+几行代码，大模型训练提速40%！无缝支持HuggingFace"title="1个GPU+几行代码，大模型训练提速40%！无缝支持HuggingFace"style="visibility:visible;宽度：1080px；height:223px;"data-type="inline">有人认为模型太大，个人玩家难以承受天价成本，因此提出了异构内存、并行计算等方法来加速和降低大型模型训练的成本，比如开源项目Colossal-AI刚刚实现了一块NVIDIA3090可以单选出180亿参数的大型模型，这两天又来了一波ofnew：无缝支持HuggingFace社区模型，只需添加几行代码，即可实现大型模型的低成本训练和微调。要知道，HuggingFace作为目前最流行的AI库之一当下，提供了超过50000个AI模型的实现，是很多AI玩家训练大模型的首选。Colossal-AI的运行使得公共模型的训练和微调变得更加可行。而且训练效果也有了Leaflet在GPU上，与微软的DeepSpeed相比，使用Colossal-AI的自动优化策略最快可以实现40%的加速。然而，PyTorch等传统深度学习框架无法在单个GPU上运行如此大的模型。对于使用8个GPU的并行训练，只需将-nprocs8添加到启动命令即可实现。这一波下来，可以说是把个人AI玩家需要考虑的成本、效率、实际问题都掌握了~不用修改代码逻辑，都是空谈，不练假动作。下面以OPT为例，详细看看Colossal-AI的新功能如何使用。OPT，全称是OpenPretrainedTransformer。由MetaAI发布，对标GPT-3，最大参数数可达1750亿。最大的特点是GPT-3不公开模型权重，而OPT开源所有代码和权重。因此，每个开发者都可以在此基础上开发个性化的下游任务。以下示例是根据OPT提供的预训练权重对CasualLanguageModeling进行微调。主要分为两步：添加配置文件，运行启动。第一步是根据要执行的任务添加配置文件。比如在GPU上，以异构训练为例，只需要在配置文件中添加相关的配置项，而不需要改变代码的训练逻辑。例如tensor_placement_policy决定了异构训练策略，参数可以是CUDA、CPU、auto。每种策略都有不同的优势，适用于不同的情况。CUDA：所有模型参数都放在GPU上，适用于传统场景不卸载依然可以训练。CPU：所有模型参数都放在CPU内存中，GPU内存中只保留当前参与计算的权重，适合超大型模型的训练。auto：根据实时显存信息，自动判断GPU显存保留的参数量，可以最大限度地利用GPU显存，减少CPU与GPU之间的数据传输。对于普通用户来说，使用自动策略是最方便的。这样，Colossal-AI就可以实时自动动态选择最佳异构策略，实现计算效率最大化。fromcolossalai.zero.shard_utilsimportTensorShardStrategyzero=dict(model_config=dict(shard_strategy=TensorShardStrategy(),tensor_placement_policy="auto"),optimizer_config=dict(gpu_margin_mem_ratio=0.8))第二步插入几行代码开始新功能。首先，使用配置文件一行代码启动Colossal-AI。Colossal-AI会自动初始化分布式环境，读取相关配置，然后自动将配置中的函数注入到模型、优化器等组件中。colossalai.launch_from_torch(config='./configs/colossalai_zero.py')然后，像往常一样定义数据集、模型、优化器、损失函数等。比如直接使用原生的PyTorch代码，在定义模型的时候，只需要在ZeroInitContext下初始化模型即可。在这里，使用HuggingFace提供的OPTForCausalLM模型和预训练权重在Wikitext数据集上进行微调。withZeroInitContext(target_device=torch.cuda.current_device(),shard_strategy=shard_strategy,shard_param=True):model=OPTForCausalLM.from_pretrained('facebook/opt-1.3b'config=config)接下来，只需调用colossalai.initialize，然后可以将配置文件中定义的异构内存函数统一注入训练引擎，并启动相应的函数。engine,train_dataloader,eval_dataloader,lr_scheduler=colossalai.initialize(model=model,optimizer=optimizer,criterion=criterion,train_dataloader=train_dataloader,test_dataloader=eval_dataloader,lr_scheduler=lr_scheduler)仍然依赖GPU+CPU异构性让用户实现以上“傻瓜式”操作的关键在于AI系统本身要足够智能。发挥核心作用的是Colossal-AI系统的高效异构内存管理子系统Gemini。它就像系统中的管理员。收集计算所需的信息后，动态分配CPU和GPU的内存使用。具体工作原理是在前几步进行预热，收集PyTorch动态计算图中的内存消耗信息。预热结束后，在计算一个算子之前，Gemini会利用收集到的内存使用记录在计算设备上预留算子所需的峰值内存，同时将部分模型张量从GPU内存移至CPU记忆。Gemini内置的内存管理器为每个张量标记状态信息，包括HOLD、COMPUTE、FREE等，然后根据动态查询到的内存使用情况动态转换张量的状态，不断调整张量的位置。直接的好处是可以在硬件非常有限的情况下最大化模型容量和平衡训练速度。要知道，目前业界主流的ZeRO（ZeroReduencyOptimizer）方式，也是采用CPU+GPU异构显存的方式，但是因为静态分区，还是会造成系统崩溃和不必要的流量等问题。而且采用动态异构CPU+GPU显存的方式，还可以通过加内存条的方式扩展显存。比买高端显卡划算多了。目前使用Colossal-AI方法，RTX20606GB普通游戏本能训练出15亿参数模型；RTX309024GB主机直接单挑一个180亿参数的大模型；TeslaV10032GB甚至可以拿下240亿参数。Colossal-AI除了最大限度地利用内存，还采用了分布式并行的方式不断提升训练速度。提出同时使用数据并行、流水线并行、2.5维张量并行等复杂并行策略。方法虽然复杂，但上手还是很“傻瓜式操作”的。它可以用一个简单的语句自动实现。无需像其他系统和框架那样侵入代码，手动处理复杂的底层逻辑。parallel=dict(pipeline=2,tensor=dict(mode='2.5d',depth=1,size=4))Colossal-AI还能做什么？事实上，自开源以来，Colossal-AI已多次在GitHub和PapersWithCode热榜上位列全球第一，在技术圈享有盛誉。除了上面提到的使用单个GPU训练大型模型外，当Colossal-AI扩展到拥有数十个甚至数百个GPU的大规模并行场景时，其性能相比NVIDIAMegatron等现有系统可以提升一倍-LM。次，资源使用可以减少到不到十分之一。也就是说，在预训练GPT-3等超大型AI模型上，节省可达数百万。据介绍，Colossal-AI相关解决方案已被自动驾驶、云计算、零售、医药、芯片等领域的知名厂商采用。同时，他们也非常注重开源社区的建设，提供中文教程，开设用户社区论坛，并根据大家的反馈不断更新迭代。比如我们发现有粉丝留言问，Colossal-AI可以直接在HuggingFace上加载一些模型吗？好吧，这个更新就在这里。那么，对于大型模型的训练，您认为有哪些难点亟待解决？欢迎在评论区留言讨论~传送门项目地址：https://github.com/hpcaitech/ColossalAI
