650亿参数大模型预训练加速38%,类LLaMA基础大模型最佳实践已开源

原创
07/18 16:48
阅读数 611

「百模大战」正风起云涌,AIGC相关企业融资和并购金额也屡创新高,全球科技企业争相入局。然而,AI大模型风光无限的背后是成本极其高昂,单次预训练成本或高达上千万元。于LLaMA等现有开源大模型的微调,也难以满足企业打造核心竞争力和多样化商业使用等需求。因此,如何低成本量身打造预训练基础大模型,已成为AI大模型浪潮的关键瓶颈。

Colossal-AI作为全球最大、最活跃的大模型开发工具与社区,以当前被最广泛使用的LLaMA为例,提供开箱即用的650亿参数预训练方案,可提升训练速度38%,为大模型企业节省大量成本。

开源地址:https://github.com/hpcaitech/ColossalAI

LLaMA 点燃开源热情

Meta 开源的 7B~65B LLaMA 大模型进一步激发了打造类 ChatGPT 的热情,并由此衍生出 Alpaca, Vicuna,ColossalChat 等微调项目。

但 LLaMA 只开源了模型权重且限制商业使用,微调能够提升和注入的知识与能力也相对有限。对于真正投身大模型浪潮的企业来说,仍必须预训练自己的核心大模型。开源社区为此做了诸多努力:

  • RedPajama:开源可商用类 LLaMA 数据集,无训练代码和模型

  • OpenLLaMA:开源可商用类 LLaMA 7B,13B 模型,使用 EasyLM 基于 JAX 和 TPU 训练

  • Falcon:开源可商用类 LLaMA 7B,40B 模型,无训练代码

但对于最主流的 PyTorch + GPU 生态,仍缺乏高效、可靠、易用的类 LLaMA 基础大模型预训练方案。

最佳大模型预训练方案提速 38%

针对上述空白与需求,Colossal-AI 首个开源了 650 亿参数 LLaMA 低成本预训练方案,相比业界其他主流选择,可提升预训练速度 38%,仅需 32 张 A100/A800 即可使用,并且不限制商业使用。

而像原生 PyTorch、FSDP 等,则因显存溢出无法运行该任务。Hugging Face accelerate、DeepSpeed、Megatron-LM 也未对 LLaMA 预训练进行官方支持。

开箱即用

1. 安装 Colossal-AI

git clone -b example/llama https://github.com/hpcaitech/ColossalAI.git
cd ColossalAI
# install and enable CUDA kernel fusion
CUDA_EXT=1 pip install .

2. 安装其他依赖

cd examples/language/llama
# install other dependencies
pip install -r requirements.txt
# use flash attention
pip install xformers

3. 数据集

默认数据集 togethercomputer/RedPajama-Data-1T-Sample 将在首次运行时自动下载,也可通过 -d 或 --dataset 指定自定义数据集。

4. 运行命令

已提供 7B 和 65B 的测速脚本,仅需根据实际硬件环境设置所用多节点的 host name 即可运行性能测试。

cd benchmark_65B/gemini_auto
bash batch12_seq2048_flash_attn.sh

对于实际的预训练任务,使用与速度测试一致,启动相应命令即可,如使用 4 节点 * 8 卡训练 65B 的模型

colossalai run --nproc_per_node 8 --hostfile YOUR_HOST_FILE --master_addr YOUR_MASTER_ADDR pretrain.py -c '65b' --plugin "gemini" -l 2048 -g -b 8 -a

例如,使用 Colossal-AI gemini_auto 并行策略,可便捷实现多机多卡并行训练,降低显存消耗的同时保持高速训练。还可根据硬件环境或实际需求,选择流水并行 + 张量并行 + ZeRO1 等复杂并行策略组合

其中,通过 Colossal-AI 的 Booster Plugins,用户可以便捷自定义并行训练,如选择 Low Level ZeRO,Gemini,DDP 等并行策略。Gradient checkpointing 通过在反向传播时重新计算模型的 activation 来减少内存使用。通过引入 Flash     attention 机制加速计算并节省显存。用户可以通过命令行参数便捷控制数十个类似的自定义参数,在保持高性能的同时为自定义开发保持了灵活性。

Colossal-AI 最新的 ShardFormer 极大降低了使用多维并行训练 LLM 的上手成本。现已支持包括 LLaMA 的多种等主流模型,且原生支持 Huggingface/transformers 模型库。无需改造模型,即可支持多维并行 (流水、张量、ZeRO、DDP 等) 的各种配置组合,能够在各种硬件配置上都发挥卓越的性能。

AI 大模型系统基础设施 Colossal-AI

Colossal-AI 为该方案提供了核心系统优化与加速能力支持,它由加州伯克利大学杰出教授 James Demmel 和新加坡国立大学校长青年教授尤洋领导开发。Colossal-AI 基于 PyTorch,可通过高效多维并行、异构内存等,降低 AI 大模型训练 / 微调 / 推理的开发与应用成本,降低 GPU 需求等。

Colossal-AI上述解决方案已在某世界500强落地应用,在千卡集群性能优异,仅需数周即可完成千亿参数私有大模型预训练。上海AI Lab与商汤等新近发布的InternLM预训练代码也参考了Colossal-Al

自开源以来,Colossal-AI 已经多次在 GitHub 热榜位列世界第一,获得 GitHub Star 超 3 万颗,并成功入选 SC、AAAI、PPoPP、CVPR、ISC 等国际 AI 与 HPC 顶级会议的官方教程,已有上百家企业参与共建 Colossal-AI 生态。其背后的潞晨科技,近期获得数亿元 A 轮融资,已在成立 18 个月内已迅速连续完成三轮融资。

开源地址:https://github.com/hpcaitech/ColossalAI
 

如果您有兴趣了解更多关于潞晨科技的最新动态,请扫描下方二维码关注我们的微信公众号。

扫码添加潞晨官方微信小助手,加入潞晨官方分享群提问讨论!我们有专业的工程师答疑。

​​​​

展开阅读全文
加载中
点击引领话题📣 发布并加入讨论🔥
打赏
0 评论
0 收藏
0
分享
返回顶部
顶部