明敏 发自 凹非寺

量子位 | 公众号 QbitAI

不得不说，为了让更多的人能够使用大模型，技术圈真的是各奇招！

模型不够开放？有人开始自己做免费开源版。

比如最近风靡全网DALL·E Mini，Meta开放的OPT-175B（Open Pretrained Transformer）。

都是通过复制，让原本不够open成为每个人都可以使用的大模型。

也有人认为模型太大，个人玩家很难承受天价。

因此，提出异构内存、并行计算等方法，加快大模型训练，降低成本。

例如，开源项目Colossal-AI，不久前，英伟达3090只能单挑180亿参数大模型。

这两天，他们又来了一波新的浪潮：

无缝支持Hugging Face社区模型可以通过添加几行代码来实现低成本的大型模型培训和微调。

要知道，Hugging Face作为目前最流行的AI提供了5万多个库AI实现模型有很多AI大模型玩家训练的首选。

而Colossal-AI这波操作使公开模型的训练微调更加实用。

而且训练效果也有所提高。

单张GPU与微软相比DeepSpeed，使用Colossal-AI实现40%加速的自动优化策略。

而PyTorch传统的深度学习框架，如单张GPU这么大的模型已经不能在上面运行了。

对于使用8张GPU并行训练只需在启动命令中添加-nprocs 8就能实现。

这波下来，可以说是个人AI所有玩家需要考虑的成本、效率和实际问题都被抓住了~

光说不练假把式。

下面就以OPT例如，详细展开Colossal-AI的新功能到底怎么用。

OPT，全称为Open Pretrained Transformer。

它由Meta AI发布，对标GPT-3.最大参数可达1750亿。

最大的特点是，GPT-3没有公开模型的权重，但是OPT开源所有代码和权重。

因此，每个开发者都可以在此基础上开发个性化的下游任务。

下面的举例，就是根据OPT因果语言模型提供的预训练权重（Casual Language Modelling）的微调。

主要分为两个步骤：

第一步是根据要执行的任务添加配置文件。

比如在一张GPU以异构训练为例，只需在配置文件中添加相关配置项，无需更改代码的训练逻辑。

比如，tensor_placement_policy决定异构训练策略的参数可以是CUDA、CPU及auto。

每种策略都有不同的优势和适应性。

CUDA：所有模型参数都放置在GPU上，适合不offload传统场景仍然可以训练。

CPU：把所有的模型参数都放在里面CPU仅在内存中GPU目前参与计算的权重保留在显存中，适合超大型训练。

auto：根据实时内存信息，自动决定保留GPU最大限度地利用显存中的参数GPU同时减少显存CPU-GPU数据传输之间。

对于普通用户，使用auto策略是最方便的。

这样可以由Colossal-AI实时动态选择最佳异构策略，最大限度地提高计算效率。

from colossalai.zero.shard_utils import TensorShardStrategyzero = 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可以初始化。

用在这里Hugging Face提供的OPTForCausalLM模型和预训练权重在Wikitext数据集上进行微调。

with ZeroInitContext(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内存使用。

具体工作原理，就是在前面几个step预热，收集PyTorch内存消耗信息在动态计算图中。

预热后，计算算子前，使用收集的内存使用记录，Gemini在计算设备上预留该算子所需的峰值内存，同时从GPU显存移动一些模型张量CPU内存。

Gemini内置内存管理器将状态信息标记在每信息，包括HOLD、COMPUTE、FREE等。

然后，根据动态查询的内存使用情况，不断动态转换张量状态，调整张量位置。

在硬件非常有限的情况下，模型容量和平衡训练速度可以最大化。

要知道行业的主流方法ZeRO （Zero Reduency Optimizer），尽管也利用CPU GPU异构内存的方法，但由于是静态划分，仍导致系统崩溃、不必要的通信量等问题。

而且，使用动态异构CPU GPU添加内存条的方法也可以扩展内存。

比买高端显卡划算多了。

目前，使用Colossal-AI的方法，RTX 2060 6GB普通游戏本能训练15亿参数模型；RTX 3090 24GB主机直接单挑180亿参数大模型；Tesla V100 32GB能够赢得240亿参数。

除了最大度地利用内存外，Colossal-AI还采用分布式并行法，不断提高训练速度。

它提出并行使用数据，并行流水.复杂的并行策略，如5维张量并行。

虽然方法复杂，但上手还是很傻瓜操作的，简单声明就能自动实现。

复杂的底层逻辑不需要像其他系统和框架那样侵入代码。

parallel = dict( pipeline=2, tensor=dict(mode='2.5d', depth = 1, size=4))Colossal-AI还能做什么？

事实上，自开源以来，Colossal-AI已经多次在GitHub及Papers With Code热榜位居世界第一，在技术界名声不大。

除上述使用单张外GPU训练大模型外，Colossal-AI扩展到数十张甚至数百张GPU与英伟达相比，大规模并行场景Megatron-LM在现有系统中，性能可以翻倍，资源的使用可以减少到十分之一以下。

换算，预训练GPT-3等超大AI在模型上，可以节省数百万元。

据透露，Colossal-AI自动驾驶、云计算、零售、医药、芯片等行业知名厂商采用了相关解决方案。

同时，他们也非常重视开源社区的建设，提供中文教程和开放用户社区论坛，并根据每个人的需求反馈不断更新迭代。

比如我们发现之前有粉丝留言问，Colossal-AI能否直接加载？Hugging Face一些模型？

嗯，这次更新来了。

因此，你认为大模型训练还有哪些困难需要解决？

欢迎在评论区留言讨论~

项目地址：https://github.com/hpcaitech/ColossalAI

参考链接：
[1]https://medium.com/@yangyou_berkeley/colossal-ai-seamlessly-accelerates-large-models-at-low-costs-with-hugging-face-4d1a887e500d
[2]https://arxiv.org/abs/2202.05924v2
[3]https://arxiv.org/abs/2205.11487
[4]https://github.com/features/copilot
[5]https://github.com/huggingface/transformers

— 完 —

量子位 QbitAI · 头条号签约

关注我们，第一时间了解前沿科技动态