Making big AI models cheaper, easier, and more scalable https://github.com/hpcaitech/ColossalAI

天问 773a26e6f8 Update 'README.md' 10 months ago
README.md 773a26e6f8 Update 'README.md' 10 months ago

README.md

Colossal-AI

Colossal-AI: 让AI大模型更低成本、方便易用、高效扩展

特点

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

(返回顶端)

Colossal-AI 成功案例

ColossalChat

ColossalChat: 完整RLHF流程0门槛克隆 ChatGPT [代码] [博客] [在线样例]

  • 最高可提升单机训练速度7.73倍,单卡推理速度1.42倍

  • 单卡模型容量最多提升10.3倍
  • 最小demo训练流程最低仅需1.62GB显存 (任意消费级GPU)

  • 提升单卡的微调模型容量3.7倍
  • 同时保持高速运行

(back to top)

AIGC

加速AIGC(AI内容生成)模型,如Stable Diffusion v1Stable Diffusion v2

  • 训练: 减少5.6倍显存消耗,硬件成本最高降低46倍(从A100到RTX3060)

  • 推理: GPU推理显存消耗降低2.5倍

(返回顶端)

生物医药

加速 AlphaFold 蛋白质结构预测

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

(返回顶端)

并行训练样例展示

GPT-3

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

GPT-2

  • 降低11倍 GPU 显存占用,或超线性扩展(张量并行)

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

BERT

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

PaLM

OPT

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

ViT

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

推荐系统模型

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

(返回顶端)

单GPU训练样例展示

GPT-2

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

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

PaLM

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

(返回顶端)

推理 (Energon-AI) 样例展示

  • BLOOM: 降低1760亿参数BLOOM模型部署推理成本超10倍

(返回顶端)

安装

环境要求:

如果你遇到安装问题,可以向本项目 反馈

从PyPI安装

您可以用下面的命令直接从PyPI上下载并安装Colossal-AI。我们默认不会安装PyTorch扩展包。

pip install colossalai

注:目前只支持Linux。

但是,如果你想在安装时就直接构建PyTorch扩展,您可以设置环境变量CUDA_EXT=1.

CUDA_EXT=1 pip install colossalai

否则,PyTorch扩展只会在你实际需要使用他们时在运行时里被构建。

与此同时,我们也每周定时发布Nightly版本,这能让你提前体验到新的feature和bug fix。你可以通过以下命令安装Nightly版本。

pip install colossalai-nightly

从源码安装

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

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

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

# install colossalai
pip install .

我们默认在pip install时不安装PyTorch扩展,而是在运行时临时编译,如果你想要提前安装这些扩展的话(在使用融合优化器时会用到),可以使用一下命令。

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