Colossal-AI 快速演示实战：从单卡基线到数据并行、混合并行、MoE 与序列并行

本文是 docs/source/zh-Hans/get_started/run_demo.md 的完整解读与实战扩展。作为集成式大规模深度学习系统，Colossal-AI 通过统一的并行化框架，让同一个训练脚本既能跑在单张 GPU 上，也能通过切换启动参数与配置平滑扩展到多 GPU 分布式集群，并在此基础上叠加数据并行、张量/流水线混合并行、MoE 专家并行与序列并行。读完本文，你将掌握基于当前仓库内真实示例（ResNet/CIFAR-10、GPT、BERT-序列并行）的一步步启动方法、关键命令行参数含义，以及并行配置背后的源码实现逻辑，能够直接把"单卡 demo"改造成多卡并行训练。

仓库环境与阅读准备

在动手运行示例前，建议先确认两件事：

1. Colossal-AI 已正确安装：安装步骤与 CUDA/PyTorch 版本要求请参考 docs/source/zh-Hans/get_started/installation.md。示例的加速能力由底层 kernel（见 extensions/README.md）提供，因此推荐从源码安装以获得完整算子支持。
2. 理解统一启动方式：本仓库示例统一采用 colossalai run 命令行入口来拉起分布式训练，而不是直接 python train.py。该入口的行为与 torchrun 相似，具体参数（如 --nproc_per_node、--hostfile、--master_addr/--master_port）见 docs/source/zh-Hans/basics/launch_colossalai.md 与 docs/source/zh-Hans/basics/command_line_tool.md。

对并行概念尚不熟悉的读者，可先阅读 docs/source/zh-Hans/concepts/paradigms_of_parallelism.md 了解数据、张量、流水线与序列并行等范式的基本划分，再回到本文对照示例验证。

单 GPU：以 ResNet/CIFAR-10 验证基线性能

Colossal-AI 并非只面向多卡场景——它可以在只有一个 GPU 的系统上完整运行深度学习训练，并达到与原生 PyTorch 相当的 baseline 性能。官方用于验证的示例是 examples/images/resnet，其 README.md 提供了训练与评估脚本的完整说明。

训练脚本的代码骨架

从 examples/images/resnet/train.py 可以看出该示例的"标准 Colossal-AI 写法"：

• 数据准备：CIFAR-10 通过 torchvision.datasets.CIFAR10 下载（可通过环境变量 DATA 指定存放目录），并调用 plugin.prepare_dataloader(...) 生成与并行策略适配的 DataLoader（train.py 第 30-50 行）。
• 模型与优化器：使用 torchvision.models.resnet18(num_classes=10)，优化器选用 Colossal-AI 的 HybridAdam，学习率调度采用 MultiStepLR(milestones=[20, 40, 60, 80], gamma=1/3)。
• 核心启动步骤：创建 Booster(plugin=plugin, ...) 后调用 model, optimizer, criterion, _, lr_scheduler = booster.boost(...)，即可让训练管线自动应用所选并行与精度策略（train.py 第 172-177 行）。
• 断点续训：--resume 指定从某个 epoch 恢复时，会依次加载 model/optimizer/lr_scheduler 三份 checkpoint（train.py 第 182-185 行）。

命令行参数详解

示例 README 明确定义了以下参数：

参数	含义	默认值
-p / --plugin	使用的插件，可选 torch_ddp、torch_ddp_fp16、low_level_zero	torch_ddp
-r / --resume	从指定 checkpoint epoch 恢复训练；-1 表示不恢复	-1
-c / --checkpoint	checkpoint 保存目录	./checkpoint
-i / --interval	每多少个 epoch 保存一次 checkpoint；0 表示不保存	5
--target_acc	目标精度；训练结束未达到则抛出异常（用于 CI 验证）	None

评估脚本 eval.py 则提供 -e/--epoch（指定评估哪个 epoch）与 -c/--checkpoint（指定 checkpoint 目录）。

运行方式

cd examples/images/resnet
pip install -r requirements.txt

# 单卡训练（nproc_per_node=1），fp32 baseline
colossalai run --nproc_per_node 1 train.py -c ./ckpt-fp32

训练完成后按同样的 checkpoint 路径评估：

python eval.py -c ./ckpt-fp32 -e 80

参考精度基准

示例 README.md 给出了在 CIFAR-10 上训练 80 epoch 后的可复现精度对照（说明该表基于 torchvision.models.resnet18 的实现）：

Model	单卡 Baseline FP32	Booster DDP FP32	Booster DDP FP16	Booster Low Level Zero	Booster Gemini
ResNet-18	85.85%	84.91%	85.46%	84.50%	84.60%

由此可见，无论采用哪种插件，训练精度都能稳定保持在 baseline 附近，验证了"加速但不牺牲精度"的设计目标。需要注意的是，上表是官方 README 记录的当时测试结果，实际数值会随 PyTorch/CUDA 版本与随机种子略有浮动。

多 GPU 数据并行：一行命令切换

数据并行是最直观的并行方式：每个 GPU 持有一份完整模型副本，各自处理不同 batch 的数据。基于与单 GPU 完全相同的 ResNet 示例，只需把 --nproc_per_node 设为本机 GPU 数量即可启用：

# 4 卡数据并行训练
colossalai run --nproc_per_node 4 train.py -c ./ckpt-ddp

# 4 卡 + 混合精度（torch_ddp_fp16）
colossalai run --nproc_per_node 4 train.py -p torch_ddp_fp16 -c ./ckpt-ddp-fp16

# 4 卡 + ZeRO（low_level_zero）
colossalai run --nproc_per_node 4 train.py -p low_level_zero -c ./ckpt-zero

其底层原理是：colossalai run 为每个 rank 拉起独立进程，脚本内部的 Booster 会根据所选 plugin 自动完成进程组初始化、梯度同步与数据切分。torch_ddp 对应 PyTorch 原生 DistributedDataParallel；torch_ddp_fp16 额外叠加混合精度；low_level_zero 则使用 Colossal-AI 的 ZeRO 优化器在数据并行组内分片优化器状态（相关插件实现位于 colossalai/booster/plugin/low_level_zero_plugin.py）。更深入的 ZeRO 原理可参考 docs/source/zh-Hans/features/zero_with_chunk.md。

混合并行：在 GPT 示例中组合数据 + 张量 + 流水线并行

当模型规模大到单卡放不下、或需要进一步压榨集群算力时，就需要混合并行。Colossal-AI 的混合并行策略由数据并行、张量并行与流水线并行组合而成，其中张量并行进一步细分为 1D、2D、2.5D 与 3D 四种模式——开发者无需改动模型代码，只需修改 config.py 中的配置即可切换。

官方推荐用于验证该能力的 GPT 示例位于 examples/language/gpt，其 README.md 给出两条稳定路线：

• Gemini + DDP/ZeRO + 张量并行：对 huggingface GPT-2 模型改动最小，推荐先用它快速跑通分布式。直接执行 bash run_gemini.sh 即可，脚本 run_gemini.sh 中的环境变量 CAI_Gemini 对应 Gemini + ZeRO DDP，另有 CAI_ZeRO1、CAI_ZeRO2、Pytorch_DDP、Pytorch_ZeRO 等方案可选。
• Titans（张量并行）+ ZeRO + 流水线并行：进阶路线，使用 examples/language/gpt/titans 中基于分布式算子构建的自定义 GPT 模型，通过配置文件在 TP/PP/ZeRO 之间自由切换，追求更极致的扩展性，但对新模型结构的适配成本更高。

示例目录下的 hybridparallelism 子目录还提供了一个便于上手的插件式 finetune 用例（cd hybridparallelism && bash run.sh）。在更早期的 legacy 配置风格中，张量并行模式在 config.py 的 parallel 字段中声明，例如：

parallel = dict(
    tensor=dict(size=4, mode='1d'),   # 可切换为 '2d' / '2.5d' / '3d'
    pipeline=dict(size=2),
)

四种张量并行的切分与通信细节可分别参考 1D_tensor_parallel.md、2D_tensor_parallel.md、2p5D_tensor_parallel.md 与 3D_tensor_parallel.md；流水线并行请参考 pipeline_parallel.md。期望深入原理的读者，还可以配合仓库内 hybrid parallel 教程 docs/source/zh-Hans/advanced_tutorials/train_gpt_using_hybrid_parallelism.md 一起阅读。

MoE 并行：把专家混合整合进模型

Mixture of Experts（MoE）通过稀疏激活的方式在不显著增加计算量的前提下扩大模型容量，Colossal-AI 为其提供了专有的并行支持：将不同专家（expert）分布到不同的 GPU 上，每个进程只持有部分专家，从而把 MoE 与数据并行、模型并行结合使用。

当前仓库中对应的概念教程是 docs/source/zh-Hans/advanced_tutorials/integrate_mixture_of_experts_into_your_model.md，其中给出了启用 MoE 并行的最小配置：

MOE_MODEL_PARALLEL_SIZE = ...   # 例如 4
parallel = dict(
    moe=dict(size=MOE_MODEL_PARALLEL_SIZE)
)

配置含义如下：

• size 指定一个 MoE 模型并行组的规模，例如设为 4 时，4 个进程被分配到 4 张（通常是连续的）GPU 上组成一个专家并行组，每个进程各持有部分专家。
• 增大 size：降低专家间的通信成本，但每张 GPU 的计算量与被切分的模型参数/activation 存储开销随之上升。
• 数据并行组大小无需手动声明，默认按 GPU 总数自动推导。

从仓库源码看，MoE 的运行时算子与通信原语集中在 colossalai/moe（如 _operation.py），与张量并行的融合则体现在低层实现中（可在 colossalai/_analyzer/_subclasses 的 flop/通信分析逻辑中窥见对 MoE 流量的统计）；对应的验证用例见 tests/test_moe（如 test_moe_ep_tp.py、test_moe_ep_zero.py，分别覆盖"专家并行 + 张量并行"与"专家并行 + ZeRO"的组合），可作为集成 MoE 到自定义模型时的参照。

序列并行：为长序列训练破解显存与长度限制

序列并行针对的是 NLP 长序列训练中的两大痛点：中间激活的显存开销随序列长度线性增长，以及单卡可承载的序列长度上限受限。其思路是把输入张量与中间激活沿序列（sequence）维度切分到多张 GPU 上，从而支持更大的 batch 与更长的序列。

示例结构与启动命令

仓库中的演示位于 examples/tutorial/sequence_parallel，它以 BERT 为例实现了序列并行训练。其 README.md 给出了完整的运行步骤：

cd examples/tutorial/sequence_parallel
pip install -r requirements.txt
export PYTHONPATH=$PWD

# 默认配置：序列并行 size=2、pipeline size=1，使用合成数据
colossalai run --nproc_per_node 4 train.py

若希望把序列维度切分得更细，或叠加流水线并行，直接修改 config.py 中的 parallel 配置即可，例如将 size 从 2 改为 8（对应 8 张 GPU 切分序列），或把 pipeline 从 1 改为 2。

config.py 关键配置解读

示例的 config.py 内容及其含义如下：

from colossalai.legacy.amp import AMP_TYPE

# 训练超参
TRAIN_ITERS = 10
DECAY_ITERS = 4
WARMUP_FRACTION = 0.01
GLOBAL_BATCH_SIZE = 32   # dp world size * 每 GPU 的样本数
LR = 0.0001
MIN_LR = 1e-05
WEIGHT_DECAY = 0.01
SEQ_LENGTH = 128

# BERT 模型配置
DEPTH = 4
NUM_ATTENTION_HEADS = 4
HIDDEN_SIZE = 128

# 随机种子
SEED = 1234

# 仅当 pipeline > 1 时启用
NUM_MICRO_BATCHES = 4

# Colossal-AI 并行配置：将 size 改为 8 即在 8 卡上切分序列
parallel = dict(pipeline=1, tensor=dict(size=2, mode="sequence"))

# 混合精度：Naive AMP
fp16 = dict(mode=AMP_TYPE.NAIVE, verbose=True)

# 序列并行的梯度规约 Handler
gradient_handler = [dict(type="SequenceParallelGradientHandler")]

要点说明：

• parallel 字段：tensor=dict(size=2, mode="sequence") 是开启序列并行的关键——mode="sequence" 表示沿序列维度切分，而普通张量并行模式则为 "1d"/"2d" 等。
• fp16 字段：mode=AMP_TYPE.NAIVE 对应 Naive 混合精度策略（即"模型整体 fp16 + 主权重 fp32"），verbose=True 会在训练时打印精度相关信息；混合精度的其他模式可参考 docs/source/zh-Hans/features/mixed_precision_training_with_booster.md。
• gradient_handler：声明 SequenceParallelGradientHandler，这是因为序列切分后，跨 GPU 的激活依赖导致反向传播中需要对张量并行维度的梯度做规约，该 handler 负责在 backward 后执行对应通信。相关实现位于 colossalai/legacy/engine/gradient_handler。

多机扩展的启动方式

若使用多机多卡集群，除了 colossalai run --nproc_per_node <每机GPU数>（单机场景，配 --master_addr localhost --master_port 29500），序列并行示例 README 还建议：

• 使用 SLURM 调度时改用 colossalai launch_from_slurm 接口；
• 使用 OpenMPI 时改用 colossalai.launch_from_openmpi 接口；
• 若已有自定义启动器，则可回退到默认的 colossalai.launch 函数完成进程组初始化。

对应 CLI 层面的 hostfile 配置方式见 docs/source/zh-Hans/basics/launch_colossalai.md。序列并行的完整机制说明可进一步阅读 docs/source/zh-Hans/features/sequence_parallelism.md。

小结与扩展路线

围绕 Colossal-AI 的 Quick demos，可以总结出一条清晰的实战路径：

1. 先在单卡上跑通 ResNet/CIFAR-10（--nproc_per_node 1），确认安装正确、达到 baseline 精度；
2. 扩展到多卡数据并行（--nproc_per_node N），并用 -p 切换 torch_ddp / torch_ddp_fp16 / low_level_zero 插件；
3. 需要更大模型时，在 GPT 示例中叠加张量（1D/2D/2.5D/3D）与流水线并行，只需改 config.py；
4. 追求大容量稀疏模型时，为自定义模型接入 MoE 并行（parallel = dict(moe=...)）；
5. 处理超长序列 NLP 任务时，启用序列并行（tensor=dict(mode="sequence")）。

后续可继续阅读 docs/source/zh-Hans/get_started/reading_roadmap.md 规划学习顺序，或直接参考仓库测试目录（如 tests/test_booster、tests/test_moe）中的自动化用例，以理解各并行插件在真实分布式环境下的断言方式与期望行为。