Colossal-AI大模型并行训练系统

ColossalAI 是一个具有高效并行化技术的综合大规模模型训练系统。旨在无缝整合不同的并行化技术范式，包括数据并行、管道并行、多张量并行和序列并行。

Colossal-AI 的目标是支持人工智能社区以与他们正常编写模型相同的方式编写分布式模型。这使得他们可以专注于开发模型架构，并将分布式训练的问题从开发过程中分离出来。

ColossalAI 提供了一组并行训练组件。旨在支持用户编写分布式深度学习模型，就像编写单 GPU 模型一样。提供友好的工具，只需几行即可启动分布式培训。

• 数据并行
• 管道并行
• 1D、2D、2.5D、3D 和序列并行
• 友好的 trainer 和 engine
• 可扩展以实现新的并行性
• 混合精度训练
• 零冗余优化器 (ZeRO)

import colossalai
from colossalai.engine import Engine
from colossalai.trainer import Trainer
from colossalai.core import global_context as gpc

model, train_dataloader, test_dataloader, criterion, optimizer, schedule, lr_scheduler = colossalai.initialize()
engine = Engine(
    model=model,
    criterion=criterion,
    optimizer=optimizer,
    lr_scheduler=lr_scheduler,
    schedule=schedule
)

trainer = Trainer(engine=engine,
                  hooks_cfg=gpc.config.hooks,
                  verbose=True)
trainer.fit(
    train_dataloader=train_dataloader,
    test_dataloader=test_dataloader,
    max_epochs=gpc.config.num_epochs,
    display_progress=True,
    test_interval=5
)

展示样例

ViT

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

GPT-3

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

GPT-2

• 降低11倍GPU显存占用，或超线性扩展

BERT

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