Colossal-AI

在企业级GPU上即刻运行Colossal-AI

无需繁琐的配置。在HPC-AI云上，您即可访问强大且预配置好的Colossal-AI环境。

只需点击一下，即可训练您的模型并扩展AI工作负载！

• NVIDIA Blackwell B200s：体验下一代AI性能（查看基准测试结果）。现可在云端以低至每小时2.47美元的价格使用。
• 高性价比H200集群：按需租赁，仅需每小时1.99美元，即可享受顶级性能。

立即开始并领取免费额度 →

以半价畅享顶尖开源模型

省去麻烦。通过HPC-AI模型API，您可以无缝访问强大的长上下文LLM。

使用HPC-AI模型API构建您的AI智能体、聊天机器人和RAG应用吧！

• 最新最全模型：体验Kimi 2.5、MiniMax 2.5和GLM 5.1等最先进的性能。非常适合处理超过200万token的超大上下文窗口及复杂编码任务。

• 无与伦比的价格：不再为API端点支付过高费用。以比OpenRouter低至50%的价格获得顶级推理速度。

立即开始并领取4美元免费额度 →

Colossal-AI基准测试

为验证这些性能提升在实际应用中的效果，我们使用Colossal-AI对类似Llama的模型进行了大规模语言模型训练基准测试。测试分别在8卡和16卡配置下进行，对应7B和70B规模的模型。

GPU	GPUs	模型大小	并行策略	每个数据并行组的批量大小	序列长度	吞吐量	TFLOPS/GPU	峰值显存(MiB)
H200	8	7B	zero2(dp8)	36	4096	17.13 样本/秒	534.18	119040.02
H200	16	70B	zero2	48	4096	3.27 样本/秒	469.1	150032.23
B200	8	7B	zero1(dp2)+tp2+pp4	128	4096	25.83 样本/秒	805.69	100119.77
H200	16	70B	zero1(dp2)+tp2+pp4	128	4096	5.66 样本/秒	811.79	100072.02

Colossal-AI基准测试的结果提供了极具实用价值的洞察。对于8卡上的7B模型，B200的吞吐量高出50%，且每GPU的TFLOPS显著提升。而对于16卡上的70B模型，B200同样展现出明显优势，其吞吐量和每GPU的TFLOPS均高出70%以上。这些数据表明，B200的性能提升能够直接转化为大规模模型更短的训练时间。

最新消息

• [2025/02] DeepSeek 671B 微调指南曝光——一键解锁升级版 DeepSeek 套件，AI 爱好者欣喜若狂！
• [2024/12] 视频生成模型的开发成本节省了50%！现提供开源解决方案，并附赠 H200 GPU 代金券 [代码] [代金券]
• [2024/10] 如何构建低成本的 Sora 类应用？为您提供的解决方案
• [2024/09] 新加坡初创公司 HPC-AI Tech 获得 5000 万美元 A 轮融资，用于打造视频生成 AI 模型和 GPU 平台
• [2024/09] 通过 FP8 混合精度训练升级，仅需一行代码即可将 AI 大模型训练成本降低 30%
• [2024/06] Open-Sora 继续开源：一键生成任意 16 秒 720p 高清视频，模型权重即用
• [2024/05] 大型 AI 模型推理速度翻倍，Colossal-Inference 开源发布
• [2024/04] Open-Sora 全面升级：拥抱开源，支持单次生成 16 秒、720p 分辨率视频
• [2024/04] 针对 LLaMA3 系列量身定制的最具性价比的推理、微调和预训练解决方案

目录

• 为什么选择 Colossal-AI
• 功能特性
• Colossal-AI 在实际应用中的表现
  • Open-Sora：公开完整模型参数、训练细节及所有 Sora 类视频生成模型相关信息
  • Colossal-LLaMA-2：只需花费几百美元进行半天训练，效果即可媲美主流大模型；开源且无商业限制的领域专用 LLM 解决方案
  • ColossalChat：包含完整 RLHF 流程的开源 ChatGPT 克隆解决方案
  • AIGC：Stable Diffusion 加速
  • 生物医学：AlphaFold 蛋白质结构预测加速
• 并行训练演示
  • LLaMA 1/2/3
  • MoE
  • GPT-3
  • GPT-2
  • BERT
  • PaLM
  • OPT
  • ViT
  • 推荐系统模型
• 单 GPU 训练演示
  • GPT-2
  • PaLM
• 推理
  • Colossal-Inference：大型 AI 模型推理速度翻倍
  • Grok-1：PyTorch + HuggingFace 的 314B 模型推理
  • SwiftInfer：以 46% 的加速打破 LLM 多轮对话的长度限制
• 安装
  • PyPI
  • 从源码安装
• 使用 Docker
• 社区
• 贡献代码
• 引用我们

为什么选择 Colossal-AI

(返回顶部)

功能特性

Colossal-AI 为您提供一系列并行化组件。我们的目标是让您像在笔记本电脑上编写模型一样轻松地编写分布式深度学习模型。我们提供友好的工具，只需几行代码即可启动分布式训练和推理。

• 并行策略：

  • 数据并行
  • 流水线并行
  • 1D、2D、2.5D、3D 张量并行
  • 序列并行
  • Zero Redundancy Optimizer (ZeRO)
  • 自动并行

• 异构内存管理：

  • PatrickStar

• 友好易用：

  • 基于配置文件的并行化

(返回顶部)

Colossal-AI 在现实世界中的应用

Open-Sora

Open-Sora：揭秘完整模型参数、训练细节以及所有与Sora类似视频生成模型相关的内容 [代码] [博客] [模型权重] [演示] [GPU云平台] [OpenSora图像]

(返回顶部)

Colossal-LLaMA-2

[GPU云平台] [LLaMA3图像]

• 7B：仅需几百美元、半天的训练，即可获得与主流大模型相当的效果，是一款开源且无商业限制的领域专用LLM解决方案。 [代码] [博客] [HuggingFace模型权重] [Modelscope模型权重]

• 13B：仅需5000美元，即可构建出性能优异的13B规模私有模型。 [代码] [博客] [HuggingFace模型权重] [Modelscope模型权重]

模型	主干网络	消耗的token数	MMLU (5-shot)	CMMLU (5-shot)	AGIEval (5-shot)	GAOKAO (0-shot)	CEval (5-shot)
Baichuan-7B	-	1.2T	42.32 (42.30)	44.53 (44.02)	38.72	36.74	42.80
Baichuan-13B-Base	-	1.4T	50.51 (51.60)	55.73 (55.30)	47.20	51.41	53.60
Baichuan2-7B-Base	-	2.6T	46.97 (54.16)	57.67 (57.07)	45.76	52.60	54.00
Baichuan2-13B-Base	-	2.6T	54.84 (59.17)	62.62 (61.97)	52.08	58.25	58.10
ChatGLM-6B	-	1.0T	39.67 (40.63)	41.17 (-)	40.10	36.53	38.90
ChatGLM2-6B	-	1.4T	44.74 (45.46)	49.40 (-)	46.36	45.49	51.70
InternLM-7B	-	1.6T	46.70 (51.00)	52.00 (-)	44.77	61.64	52.80
Qwen-7B	-	2.2T	54.29 (56.70)	56.03 (58.80)	52.47	56.42	59.60
Llama-2-7B	-	2.0T	44.47 (45.30)	32.97 (-)	32.60	25.46	-
Linly-AI/Chinese-LLaMA-2-7B-hf	Llama-2-7B	1.0T	37.43	29.92	32.00	27.57	-
wenge-research/yayi-7b-llama2	Llama-2-7B	-	38.56	31.52	30.99	25.95	-
ziqingyang/chinese-llama-2-7b	Llama-2-7B	-	33.86	34.69	34.52	25.18	34.2
TigerResearch/tigerbot-7b-base	Llama-2-7B	0.3T	43.73	42.04	37.64	30.61	-
LinkSoul/Chinese-Llama-2-7b	Llama-2-7B	-	48.41	38.31	38.45	27.72	-
FlagAlpha/Atom-7B	Llama-2-7B	0.1T	49.96	41.10	39.83	33.00	-
IDEA-CCNL/Ziya-LLaMA-13B-v1.1	Llama-13B	0.11T	50.25	40.99	40.04	30.54	-
Colossal-LLaMA-2-7b-base	Llama-2-7B	0.0085T	53.06	49.89	51.48	58.82	50.2
Colossal-LLaMA-2-13b-base	Llama-2-13B	0.025T	56.42	61.80	54.69	69.53	60.3

ColossalChat

ColossalChat: 一个开源解决方案，用于克隆 ChatGPT，并配备完整的 RLHF 流程。 [代码] [博客] [演示] [教程]

• RLHF PPO Stage3 训练速度最高可提升至10倍

• 单服务器训练速度最高可提升至7.73倍，单GPU推理速度最高可提升至1.42倍

• 在单个GPU上，模型容量最高可增长10.3倍
• 一次小型演示训练过程仅需1.62GB显存（任何消费级GPU均可）

• 在单个GPU上，微调模型的容量最高可提升至3.7倍
• 同时保持足够高的运行速度

(返回顶部)

AIGC

加速AIGC（人工智能生成内容）模型，例如 Stable Diffusion v1 和 Stable Diffusion v2。

• 训练: 将Stable Diffusion的显存消耗降低至多5.6倍，硬件成本降低至多46倍（从A100降至RTX3060）。

• DreamBooth微调: 仅需3–5张目标对象的照片即可个性化您的模型。

• 推理: 将推理过程中的显存消耗减少2.5倍。

(返回顶部)

生物医药

加速 AlphaFold蛋白质结构预测

• FastFold: 加速GPU集群上的训练和推理，提升数据处理速度，支持超过10000个残基的序列推理。

• FastFold与Intel结合: 推理速度提升3倍，成本降低39%。

• xTrimoMultimer: 将蛋白质单体和多聚体的结构预测速度提升11倍。

(返回顶部)

并行训练演示

LLaMA3

• 700亿参数的LLaMA3模型训练加速18% [代码] [GPU云平台] [LLaMA3图像]

LLaMA2

• 700亿参数的LLaMA2模型训练加速195% [代码] [博客]

LLaMA1

• 650亿参数的大模型预训练加速38% [代码] [博客]

MoE

• 增强的MoE并行性，开源MoE模型训练效率可提高9倍 [代码] [博客]

GPT-3

• 节省50%的GPU资源，并实现10.7%的加速

GPT-2

• 显存消耗降低11倍，且采用张量并行时具有超线性扩展效率

• 在相同硬件条件下，模型规模扩大24倍
• 加速超过3倍

BERT

• 训练速度提升2倍，或序列长度延长50%

PaLM

• PaLM-colossalai：谷歌Pathways语言模型（PaLM）的可扩展实现。

OPT

• Open Pretrained Transformer (OPT)，由Meta发布的1750亿参数AI语言模型，其公开的预训练权重激发了AI开发者进行各种下游任务和应用部署。
• 以较低的代码成本实现OPT微调速度提升45%。[示例] [在线推理]

更多详情请访问我们的文档和示例。

ViT

• 对于张量并行度为64的情况，批量大小扩大14倍，训练速度提升5倍。

推荐系统模型

• Cached Embedding，利用软件缓存技术，在较小的GPU显存预算下训练更大的嵌入表。

(返回顶部)

单GPU训练演示

GPT-2

• 在相同硬件上，模型规模扩大20倍。

• 在相同硬件（RTX 3080）上，模型规模扩大120倍。

PaLM

• 在相同硬件上，模型规模扩大34倍。

(返回顶部)

推理

Colossal-Inference

• 在某些情况下，大型AI模型的推理速度相比vLLM的离线推理性能提升了一倍。 [代码] [博客] [GPU云平台] [LLaMA3图像]

Grok-1

• 3140亿参数的Grok-1推理加速3.8倍，提供易于使用的Python + PyTorch + HuggingFace版本用于推理。

[代码] [博客] [HuggingFace Grok-1 PyTorch模型权重] [ModelScope Grok-1 PyTorch模型权重]

SwiftInfer

• SwiftInfer：推理性能提升46%，开源解决方案突破了LLM在多轮对话中的长度限制。

(返回顶部)

安装

要求：

• PyTorch ≥ 2.2
• Python ≥ 3.7
• CUDA ≥ 11.0
• NVIDIA GPU计算能力 ≥ 7.0（V100/RTX20及以上）
• Linux操作系统

如果在安装过程中遇到任何问题，您可以在本仓库中提交issue。

通过PyPI安装

您可以使用以下命令轻松安装Colossal-AI。默认情况下，我们在安装时不会构建PyTorch扩展。

pip install colossalai

注意：目前仅支持Linux系统。

然而，如果您希望在安装时构建PyTorch扩展，可以设置BUILD_EXT=1。

BUILD_EXT=1 pip install colossalai

否则，CUDA内核将在您实际需要时于运行时构建。

我们每周还会向PyPI发布夜间版本，使您能够体验主分支中尚未发布的功能和错误修复。 可通过以下命令进行安装：

pip install colossalai-nightly

从源码下载

Colossal-AI的版本将与仓库的主分支保持一致。如遇任何问题，请随时提出issue。:)

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

# 安装colossalai
pip install .

默认情况下，我们不会编译CUDA/C++内核。ColossalAI会在运行时构建它们。 如果您希望安装并启用CUDA内核融合（使用融合优化器时必须安装）：

BUILD_EXT=1 pip install .

对于使用CUDA 10.2的用户，仍然可以从源码构建ColossalAI。不过，您需要手动下载cub库并将其复制到相应目录。

# 克隆仓库
git clone https://github.com/hpcaitech/ColossalAI.git
cd ColossalAI

# 下载cub库
wget https://github.com/NVIDIA/cub/archive/refs/tags/1.8.0.zip
unzip 1.8.0.zip
cp -r cub-1.8.0/cub/ colossalai/kernel/cuda_native/csrc/kernels/include/

# 安装
BUILD_EXT=1 pip install .

(返回顶部)

使用Docker

从DockerHub拉取

您可以直接从我们的DockerHub页面拉取Docker镜像。每次发布时，镜像都会自动上传。

Build On Your Own

Run the following command to build a docker image from Dockerfile provided.

Building Colossal-AI from scratch requires GPU support, you need to use Nvidia Docker Runtime as the default when doing docker build. More details can be found here. We recommend you install Colossal-AI from our project page directly.

cd ColossalAI
docker build -t colossalai ./docker

Run the following command to start the docker container in interactive mode.

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

(back to top)

Community

Join the Colossal-AI community on Forum, Slack, and WeChat(微信) to share your suggestions, feedback, and questions with our engineering team.

Contributing

Referring to the successful attempts of BLOOM and Stable Diffusion, any and all developers and partners with computing powers, datasets, models are welcome to join and build the Colossal-AI community, making efforts towards the era of big AI models!

You may contact us or participate in the following ways:

1. Leaving a Star ⭐ to show your like and support. Thanks!
2. Posting an issue, or submitting a PR on GitHub follow the guideline in Contributing
3. Send your official proposal to email contact@hpcaitech.com

Thanks so much to all of our amazing contributors!

(back to top)

CI/CD

We leverage the power of GitHub Actions to automate our development, release and deployment workflows. Please check out this documentation on how the automated workflows are operated.

Cite Us

This project is inspired by some related projects (some by our team and some by other organizations). We would like to credit these amazing projects as listed in the Reference List.

To cite this project, you can use the following BibTeX citation.

@inproceedings{10.1145/3605573.3605613,
author = {Li, Shenggui and Liu, Hongxin and Bian, Zhengda and Fang, Jiarui and Huang, Haichen and Liu, Yuliang and Wang, Boxiang and You, Yang},
title = {Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel Training},
year = {2023},
isbn = {9798400708435},
publisher = {Association for Computing Machinery},
address = {New York, NY, USA},
url = {https://doi.org/10.1145/3605573.3605613},
doi = {10.1145/3605573.3605613},
abstract = {The success of Transformer models has pushed the deep learning model scale to billions of parameters, but the memory limitation of a single GPU has led to an urgent need for training on multi-GPU clusters. However, the best practice for choosing the optimal parallel strategy is still lacking, as it requires domain expertise in both deep learning and parallel computing. The Colossal-AI system addressed the above challenge by introducing a unified interface to scale your sequential code of model training to distributed environments. It supports parallel training methods such as data, pipeline, tensor, and sequence parallelism and is integrated with heterogeneous training and zero redundancy optimizer. Compared to the baseline system, Colossal-AI can achieve up to 2.76 times training speedup on large-scale models.},
booktitle = {Proceedings of the 52nd International Conference on Parallel Processing},
pages = {766–775},
numpages = {10},
keywords = {datasets, gaze detection, text tagging, neural networks},
location = {Salt Lake City, UT, USA},
series = {ICPP '23}
}

Colossal-AI has been accepted as official tutorial by top conferences NeurIPS, SC, AAAI, PPoPP, CVPR, ISC, NVIDIA GTC ,etc.

(back to top)

ColossalAI 快速上手指南

ColossalAI 是一个旨在让大型 AI 模型训练更便宜、更快、更易用的开源系统。它提供了一套并行组件，让你能像在本机上编写模型一样轻松编写分布式深度学习模型。

环境准备

在开始之前，请确保你的开发环境满足以下要求：

• 操作系统: Linux (推荐 Ubuntu 18.04/20.04/22.04)
• Python: 3.8 - 3.10
• CUDA: 11.0 或更高版本 (根据显卡驱动和 PyTorch 版本匹配)
• PyTorch: 1.12 或更高版本
• 硬件: 至少一张 NVIDIA GPU (支持多卡及集群分布式训练)

提示：国内开发者若遇到网络问题，建议在安装依赖时配置清华源或阿里源加速。

安装步骤

你可以通过 PyPI 直接安装，或者从源码安装以获取最新功能。

方式一：通过 PyPI 安装（推荐）

这是最快捷的安装方式。

pip install colossalai

国内加速安装：

pip install colossalai -i https://pypi.tuna.tsinghua.edu.cn/simple

方式二：从源码安装

如果你需要最新的功能或进行二次开发，建议从源码安装。

git clone https://github.com/hpcaitech/ColossalAI.git
cd ColossalAI
pip install -e .

国内加速克隆：

git clone https://gitee.com/hpcaitech/ColossalAI.git  # 如果存在镜像
# 或者使用 git 代理加速

方式三：使用 Docker（最省心）

为了避免环境配置冲突，可以直接拉取预配置好的 Docker 镜像。

docker pull hpcaitech/colossalai:latest
docker run --gpus all -it hpcaitech/colossalai:latest

基本使用

ColossalAI 的核心优势在于只需几行代码即可启动分布式训练。以下是一个基于配置文件启动并行训练的最简示例。

1. 准备配置文件 (config.py)

ColossalAI 允许通过配置文件定义并行策略（如数据并行、流水线并行、张量并行等）。

from colossalai.context import ParallelMode
from colossalai.core import global_context as gpc
from colossalai.utils import get_dataloader

# 定义并行配置
CONFIG = dict(
    parallel=dict(
        pipeline=1,
        tensor=dict(
            size=2,
            mode='1d',
        ),
        data=1,
    ),
    fp16=dict(
        mode='auto',
    ),
    gradient_accumulation=1,
    clip_grad_norm=1.0,
)

2. 编写训练脚本 (train.py)

在你的训练脚本中，引入 colossalai 并初始化上下文，即可自动应用并行策略。

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision.datasets import CIFAR10
from torchvision.transforms import ToTensor, Normalize, Compose

import colossalai
from colossalai.core import global_context as gpc
from colossalai.logging import disable_existing_loggers
from colossalai.utils import get_dataloader
from colossalai.trainer import Trainer, TrainerHook
from colossalai.context import ParallelMode
from colossalai.nn.optimizer import HybridAdam

# 简单的模型定义
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.linear = nn.Linear(3072, 10)

    def forward(self, x):
        x = self.flatten(x)
        return self.linear(x)

def main():
    # 1. 初始化 ColossalAI 上下文 (自动读取环境变量中的 rank, world_size 等)
    colossalai.launch(config=CONFIG, rank=gpc.get_global_rank(), world_size=gpc.get_world_size(), host='localhost', port=29500, backend='nccl')
    
    # 2. 构建模型、优化器和数据加载器
    model = SimpleModel().cuda()
    optimizer = HybridAdam(model.parameters(), lr=1e-3)
    
    transform = Compose([ToTensor(), Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
    train_dataset = CIFAR10(root='./data', train=True, download=True, transform=transform)
    train_dataloader = get_dataloader(dataset=train_dataset, batch_size=64, shuffle=True, drop_last=True)

    # 3. 创建 Trainer
    trainer = Trainer(model=model, optimizer=optimizer, criterion=torch.nn.CrossEntropyLoss())

    # 4. 开始训练
    trainer.fit(train_dataloader=train_dataloader, epochs=1)

if __name__ == '__main__':
    main()

3. 启动训练

使用 colossalai run 命令启动分布式任务。以下示例是在单机 2 张 GPU 上运行：

colossalai run --nproc_per_node 2 train.py

如果是多机多卡环境，可以使用 colossalai run 配合主机列表，或使用 Slurm 等调度系统启动。

---

现在你已经成功运行了第一个 ColossalAI 分布式训练任务！你可以前往 官方文档 探索更多高级功能，如大模型微调 (LLaMA, ChatGLM 等) 和推理加速。