colossalai

ColossalAI 分布式训练：用 Gemini/ZeRO、张量并行和自动并行策略训练大模型

补充说明

ColossalAI 是一个分布式训练框架，提供 Gemini（显存管理）、ZeRO、张量并行、流水线并行、序列并行和自动并行策略。它适合在有限 GPU 上训练超出单卡显存的模型，也提供面向 Hugging Face 模型的轻量封装。和 DeepSpeed/FSDP 的目标重叠，但插件抽象和自动并行能力不同。

ColossalAI 的入口是 colossalai run 和 booster/plugin。选择插件前先明确瓶颈：显存不足优先 Gemini/ZeRO，单层算不动考虑张量并行，层数太多考虑流水线并行。不要一开始就开满所有并行策略，先用小模型验证正确性。

常用示例

# 安装
pip install colossalai

# 8 卡启动
colossalai run --nproc_per_node 8 train.py

# 多机：指定 hostfile
colossalai run --hostfile hosts --nproc_per_node 8 train.py

import colossalai
from colossalai.booster import Booster
from colossalai.booster.plugin import GeminiPlugin, HybridParallelPlugin

# 初始化分布式环境
colossalai.launch_from_torch()

# Gemini：面向显存优化的 ZeRO 风格插件
plugin = GeminiPlugin(placement_policy="auto", precision="bf16")
booster = Booster(plugin=plugin)

model, optimizer, criterion, dataloader, lr_scheduler = booster.boost(
    model, optimizer, criterion, dataloader, lr_scheduler
)

# 训练循环按 booster 的接口执行

# 混合并行：张量并行 + 数据并行
plugin = HybridParallelPlugin(
    tp_size=2, pp_size=1, num_microbatches=4, precision="bf16"
)
booster = Booster(plugin=plugin)

不同版本的 API 和插件参数会变化；以当前版本文档为准，先用单机 2 卡跑通再扩到多机。

易错点与排错

• 插件选型与瓶颈不匹配：显存不足用 Gemini/ZeRO，单层太大用张量并行，层数太多用流水线；混用前先定位瓶颈。
• 版本兼容性：ColossalAI 与 PyTorch/CUDA/Transformers 版本耦合较强；升级前锁定组合并跑回归。
• batch size 语义变化：张量/流水线并行后，单卡 batch 与全局 batch 不同；学习率和梯度累积要按全局 batch 调整。
• checkpoint 格式不通用：ColossalAI 保存的分片/格式与原生 PyTorch 不同；上线前验证保存、加载和恢复训练。
• 多机通信配置错误：hostfile、MASTER_ADDR/PORT、网卡和 NCCL 环境变量要一致；先跑 nccl-tests。
• 自动并行策略开销：自动搜索会花时间，且可能生成不稳定的方案；生产训练优先固定可复现的配置。
• 与 Hugging Face Trainer 集成：不要同时让 Trainer 和 booster 管理分布式/优化器；按官方示例选择一种入口。
• 只测吞吐不测收敛：并行策略可能改变数值行为；用固定种子和小数据集对比 loss 曲线。

相关命令

• deepspeed — ZeRO 系列显存优化。
• megatron-lm — 张量/流水线并行。
• fsdp — PyTorch 原生分片训练。
• torchrun — 分布式启动器。
• ray-train — Ray 上的分布式训练编排。

参考链接

• ColossalAI 文档
• ColossalAI GitHub
• Gemini 插件

读完自测

合上页面，先回答，再展开对照。

1. Gemini 插件主要解决什么问题？
2. 张量并行、流水线并行和数据并行分别适合什么瓶颈？
3. 为什么并行训练后学习率/梯度累积要重新检查？