显卡|1块显卡+几行代码:大模型训练提速40%!

不得不说,为了让更多人能用上大模型,技术圈真是各出奇招!
模型不够开放?有人自己上手搞免费开源版 。
比如最近风靡全网的DALL·E Mini,Meta开放的OPT-175B(Open Pretrained Transformer) 。
都是通过复刻的方式,让原本不够open的大模型,变成人人可用 。
显卡|1块显卡+几行代码:大模型训练提速40%!
文章图片

还有人觉得模型太大,个人玩家很难承受起天价成本 。
所以提出异构内存、并行计算等方法,让大模型训练加速又降本 。
比如开源项目Colossal-AI,前不久刚实现了让一块NVIDIA RTX 3090就能单挑180亿参数大模型 。
显卡|1块显卡+几行代码:大模型训练提速40%!
文章图片

而在这两天,他们又来了一波上新:
无缝支持Hugging Face社区模型,只需添加几行代码,就能实现大模型的低成本训练和微调 。
显卡|1块显卡+几行代码:大模型训练提速40%!
文章图片

要知道,Hugging Face作为当下最流行的AI库之一,提供了超过5万个AI模型的实现,是许多AI玩家训练大模型的首选 。
而Colossal-AI这波操作,是让公开模型的训练微调变得更加切实可行 。
并且在训练效果上也有提升 。
单张GPU上,相比于微软的DeepSpeed,使用Colossal-AI的自动优化策略,最快能实现40%的加速 。
而PyTorch等传统深度学习框架,在单张GPU上已经无法运行如此大的模型 。
对于使用8张GPU的并行训练,仅需在启动命令中添加-nprocs 8就能实现 。
显卡|1块显卡+几行代码:大模型训练提速40%!
文章图片

这波下来,可以说是把个人AI玩家需要考虑的成本、效率、实操问题,都拿捏住了~
无需修改代码逻辑
光说不练假把式 。
下面就以OPT为例,详细展开看看Colossal-AI的新功能到底怎么用 。
OPT,全称为Open Pretrained Transformer 。
它由Meta AI发布,对标GPT-3,最大参数量可达1750亿 。
最大特点就是,GPT-3没有公开模型权重,而OPT开源了所有代码及权重 。
因此,每一位开发者都能在此基础上开发个性化的下游任务 。
下面的举例,就是根据OPT提供的预训练权重,进行因果语言模型(Casual Language Modelling)的微调 。
主要分为两个步骤:
1、添加配置文件
2、运行启动
第一步,是根据想进行的任务添加配置文件 。
比如在一张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 TensorShardStrategy
zero = dict(model_config=dict(shard_strategy=TensorShardStrategy(),