ColossalAI/README-zh-Hans.md

Colossal-AI

Colossal-AI: 一个面向大模型时代的通用深度学习系统

论文 | 文档 | 例程 | 论坛 | 博客

| English | 中文 |

目录

• 为何选择 Colossal-AI
• 特点
• 并行训练样例展示
  • ViT
  • GPT-3
  • GPT-2
  • BERT
  • PaLM
  • OPT
  • 推荐系统模型
• 单GPU训练样例展示
  • GPT-2
  • PaLM
• 推理 (Energon-AI) 样例展示
  • GPT-3
  • 1750亿参数OPT在线推理服务
• Colossal-AI 成功案例
  • AIGC: 加速 Stable Diffusion
  • 生物医药: 加速AlphaFold蛋白质结构预测
• 安装
  • PyPI
  • 从源代码安装
• 使用 Docker
• 社区
• 做出贡献
• 快速预览
• 几行代码开启分布式训练
• 构建一个简单的2维并行模型
• 引用我们

为何选择 Colossal-AI

James Demmel 教授 (加州大学伯克利分校): Colossal-AI 让分布式训练高效、易用、可扩展。

(返回顶端)

特点

Colossal-AI 为您提供了一系列并行组件。我们的目标是让您的分布式 AI 模型像构建普通的单 GPU 模型一样简单。我们提供的友好工具可以让您在几行代码内快速开始分布式训练和推理。

• 并行化策略
  • 数据并行
  • 流水线并行
  • 1维, 2维, 2.5维, 3维 张量并行
  • 序列并行
  • 零冗余优化器 (ZeRO)
• 异构内存管理
  • PatrickStar
• 使用友好
  • 基于参数文件的并行化
• 推理
  • Energon-AI
• Colossal-AI 成功案例
  • 生物医药: FastFold 加速蛋白质结构预测 AlphaFold 训练与推理

(返回顶端)

并行训练样例展示

ViT

• 14倍批大小和5倍训练速度（张量并行=64）

GPT-3

• 释放 50% GPU 资源占用, 或 10.7% 加速

GPT-2

• 降低11倍 GPU 显存占用，或超线性扩展（张量并行）

• 用相同的硬件训练24倍大的模型
• 超3倍的吞吐量

BERT

• 2倍训练速度，或1.5倍序列长度

PaLM

• PaLM-colossalai: 可扩展的谷歌 Pathways Language Model (PaLM) 实现。

OPT

• Open Pretrained Transformer (OPT), 由Meta发布的1750亿语言模型，由于完全公开了预训练参数权重，因此促进了下游任务和应用部署的发展。
• 加速45%，仅用几行代码以低成本微调OPT。[样例] [在线推理]

请访问我们的 文档 和 例程 以了解详情。

推荐系统模型

• Cached Embedding, 使用软件Cache实现Embeddings，用更少GPU显存训练更大的模型。

(返回顶端)

单GPU训练样例展示

GPT-2

• 用相同的硬件训练20倍大的模型

• 用相同的硬件训练120倍大的模型 (RTX 3080)

PaLM

• 用相同的硬件训练34倍大的模型

(返回顶端)

推理 (Energon-AI) 样例展示

• Energon-AI ：用相同的硬件推理加速50%

• OPT推理服务: 无需注册，免费体验1750亿参数OPT在线推理服务

(返回顶端)

Colossal-AI 成功案例

AIGC

加速AIGC(AI内容生成)模型，如Stable Diffusion

• Colossal-AI优化Stable Diffusion: 6.5倍训练加速和预训练成本降低, 微调硬件成本下降约7倍(从RTX3090/4090到RTX3050/2070)

(返回顶端)

生物医药

加速 AlphaFold 蛋白质结构预测

• FastFold: 加速AlphaFold训练与推理、数据前处理、推理序列长度超过10000残基

• xTrimoMultimer: 11倍加速蛋白质单体与复合物结构预测

(返回顶端)

安装

从官方安装

您可以访问我们下载页面来安装Colossal-AI，在这个页面上发布的版本都预编译了CUDA扩展。

从源安装

此文档将与版本库的主分支保持一致。如果您遇到任何问题，欢迎给我们提 issue :)

git clone https://github.com/hpcaitech/ColossalAI.git
cd ColossalAI

# install dependency
pip install -r requirements/requirements.txt

# install colossalai
pip install .

如果您不想安装和启用 CUDA 内核融合（使用融合优化器时强制安装）：

NO_CUDA_EXT=1 pip install .

(返回顶端)

使用 Docker

从DockerHub获取镜像

您可以直接从我们的DockerHub主页获取最新的镜像，每一次发布我们都会自动上传最新的镜像。

本地构建镜像

运行以下命令从我们提供的 docker 文件中建立 docker 镜像。

在Dockerfile里编译Colossal-AI需要有GPU支持，您需要将Nvidia Docker Runtime设置为默认的Runtime。更多信息可以点击这里。 我们推荐从项目主页直接下载Colossal-AI.

cd ColossalAI
docker build -t colossalai ./docker

运行以下命令从以交互式启动 docker 镜像.

docker run -ti --gpus all --rm --ipc=host colossalai bash

(返回顶端)

社区

欢迎通过论坛, Slack, 或微信加入 Colossal-AI 社区，与我们分享你的建议和问题。

做出贡献

欢迎为该项目做出贡献，请参阅贡献指南。

真诚感谢所有贡献者！

贡献者头像的展示顺序是随机的。

(返回顶端)

快速预览

几行代码开启分布式训练

parallel = dict(
    pipeline=2,
    tensor=dict(mode='2.5d', depth = 1, size=4)
)

几行代码开启异构训练

zero = dict(
    model_config=dict(
        tensor_placement_policy='auto',
        shard_strategy=TensorShardStrategy(),
        reuse_fp16_shard=True
    ),
    optimizer_config=dict(initial_scale=2**5, gpu_margin_mem_ratio=0.2)
)

(返回顶端)

引用我们

@article{bian2021colossal,
  title={Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel Training},
  author={Bian, Zhengda and Liu, Hongxin and Wang, Boxiang and Huang, Haichen and Li, Yongbin and Wang, Chuanrui and Cui, Fan and You, Yang},
  journal={arXiv preprint arXiv:2110.14883},
  year={2021}
}

(返回顶端)