Colossal-AI大模型并行訓練系統(tǒng)

ColossalAI 是一個具有高效并行化技術的綜合大規(guī)模模型訓練系統(tǒng)。旨在無縫整合不同的并行化技術范式，包括數據并行、管道并行、多張量并行和序列并行。

Colossal-AI 的目標是支持人工智能社區(qū)以與他們正常編寫模型相同的方式編寫分布式模型。這使得他們可以專注于開發(fā)模型架構，并將分布式訓練的問題從開發(fā)過程中分離出來。

ColossalAI 提供了一組并行訓練組件。旨在支持用戶編寫分布式深度學習模型，就像編寫單 GPU 模型一樣。提供友好的工具，只需幾行即可啟動分布式培訓。

• 數據并行
• 管道并行
• 1D、2D、2.5D、3D 和序列并行
• 友好的 trainer 和 engine
• 可擴展以實現(xiàn)新的并行性
• 混合精度訓練
• 零冗余優(yōu)化器 (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倍序列長度