明敏 发自 凹非寺

量子位 | 公众号 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 · 头条号签约

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