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

【免费下载链接】ColossalAI Making large AI models cheaper, faster and more accessible 项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI

本篇技术指南以 Colossal-AI 官方快速入门文档（docs/source/en/get_started/run_demo.md）为骨架，系统梳理该大规模深度学习系统在真实仓库中的可运行演示路径：如何在单张 GPU 上训练模型获得基线性能，以及如何通过几类并行化方案（数据并行、混合并行、MoE 并行、序列并行）在多卡集群上大幅加速训练。阅读完本文，你将能够定位仓库内对应的示例代码与启动命令，理解 Booster/Plugin 层的核心机制，并据此把任意模型快速跑通单卡与多卡训练。

一、Colossal-AI 的定位与演示总览

Colossal-AI 是一套集成了高效并行化技术的大规模深度学习系统。它通过在多 GPU 分布式系统上施加并行化技术来加速模型训练，同时也能在仅有单张 GPU 的机器上正常运行。这一"单卡可跑、多卡可扩"的设计，正是其快速演示（Quick Demo）的核心体验。

原文档给出的演示矩阵可归纳如下：

演示场景	目标模型 / 数据集	并行维度	仓库内示例位置
单 GPU	ResNet / CIFAR-10	无（单卡基线）	examples/images/resnet
多 GPU · 数据并行	ResNet / CIFAR-10	batch 维度	examples/images/resnet
多 GPU · 混合并行	GPT	数据 + 张量 + 流水	examples/language/gpt
多 GPU · MoE 并行	专家混合模型	专家路由维度	MoE 集成教程
多 GPU · 序列并行	BERT	序列长度维度	examples/tutorial/sequence_parallel

二、单 GPU：用 ResNet-CIFAR10 跑通基线

Colossal-AI 在仅有一张 GPU 的机器上即可训练深度学习模型并获得与原生 PyTorch 相当的基线性能。官方示例采用经典的 ResNet-18 + CIFAR-10 图像分类任务，全部代码位于 examples/images/resnet，包含训练脚本 train.py、评估脚本 eval.py、依赖清单 requirements.txt 以及 CI 验证脚本 test_ci.sh。

2.1 准备环境

pip install -r requirements.txt

示例依赖 torchvision、tqdm 等，并假定当前环境已安装好 PyTorch 与 Colossal-AI。训练时 CIFAR-10 数据集目录可通过环境变量 DATA 指定，默认落在 ./data，首次运行会自动下载（参见 examples/images/resnet/train.py 中的 build_dataloader）。

2.2 启动单卡训练

在 train.py 内部，第一处关键调用是 colossalai.launch_from_torch()，它从 torchrun/colossalai run 注入的环境变量中初始化分布式环境；随后创建 DistCoordinator 用于协调进程。这意味着即便只使用单卡，也应通过分布式启动器运行，例如：

torchrun --standalone --nproc_per_node=1 train.py -c ./ckpt-single

colossalai run 同样可用，其 --nproc_per_node 参数即"每节点使用的 GPU 数"，设置为 1 即等价于单卡训练：

colossalai run --nproc_per_node 1 train.py -c ./ckpt-single

2.3 训练脚本的核心参数

train.py 通过 argparse 暴露了如下可配置项（源码见 examples/images/resnet/train.py）：

参数	含义	默认值
-p, --plugin	使用的并行插件，可选 torch_ddp、torch_ddp_fp16、low_level_zero、gemini	torch_ddp
-r, --resume	从第几个 epoch 的 checkpoint 恢复训练（-1 表示不恢复）	-1
-c, --checkpoint	checkpoint 保存目录	./checkpoint
-i, --interval	每间隔多少个 epoch 保存一次 checkpoint；为 0 则不保存	5
--target_acc	目标测试准确率，训练结束未达标会抛异常（供 CI 断言）	None

2.4 单卡训练背后的 Booster 机制

读懂这段代码，就等于拿到了 Colossal-AI 新 API 的"最小范式"，其调用链如下（examples/images/resnet/train.py）：

1. 构造数据：plugin.prepare_dataloader(train_dataset, batch_size=100, shuffle=True, drop_last=True) —— DataLoader 的分布式切分交由 Plugin 完成；
2. 构造模型与优化器：模型使用 torchvision.models.resnet18(num_classes=10)，优化器使用 Colossal-AI 提供的 HybridAdam；
3. 线性学习率缩放：LEARNING_RATE *= coordinator.world_size（代码注释为 old_gpu_num / old_lr = new_gpu_num / new_lr），使单卡与多卡训练保持等效步长；
4. 统一加速入口：model, optimizer, criterion, _, lr_scheduler = booster.boost(...)，Booster 依据所选 Plugin 完成 DDP/ZeRO/混合精度等包装；
5. 反向传播统一走 booster.backward(loss, optimizer)；
6. checkpoint 保存与加载统一走 booster.save_model / save_optimizer / save_lr_scheduler 与 load_* 系列接口。

其中 Plugin 的选型直接决定了并行与显存优化策略：torch_ddp 对应 PyTorch 原生 DDP + FP32；torch_ddp_fp16 额外叠加混合精度（booster_kwargs["mixed_precision"] = "fp16"）；low_level_zero 与 gemini 则对应 Colossal-AI 的低层级 ZeRO 与 Gemini 显存管理。值得注意的是源码中有一处 FIXME(ver217): gemini is not supported resnet now 注释，提示 ResNet 下 Gemini 支持尚不完善。

三、数据并行：同一脚本，改一个参数即可多卡

由于 Colossal-AI 的并行逻辑被抽象进了 Plugin/Booster 层，示例代码本身无需改动，只需要把启动时的 GPU 数量设置为机器上可用的卡数，示例即从单卡模式无缝切换为数据并行（data parallel）模式。

# 机器上有 N 张 GPU，即以数据并行方式训练
colossalai run --nproc_per_node N train.py -c ./ckpt-dp

# 或使用 PyTorch 启动器
torchrun --standalone --nproc_per_node N train.py -c ./ckpt-dp

数据并行的核心是把同一个 batch 沿样本维度切分到各卡，各卡持有完整模型副本、计算各自梯度后再做全局梯度同步（AllReduce）。上述 build_dataloader 中 plugin.prepare_dataloader(...) 与训练循环中 booster.backward(...) 的调用方式，正是为了让 Plugin 在背后替你完成 DistributedSampler 式的数据切分与梯度规约。colossalai run 这一 CLI 工具在 colossalai/cli/launcher 中实现，支持 --nproc_per_node（每节点 GPU 数）等参数，并会解析 hostfile 以支持多节点扩展。

评估与精度基线

训练完成后可用评估脚本对指定 epoch 的 checkpoint 进行测试：

# 参数：-e 指定 epoch，-c 指定 checkpoint 目录
python eval.py -c ./ckpt-fp32 -e 80

该示例的 README（examples/images/resnet/README.md）给出了 ResNet-18 在 CIFAR-10 上运行 80 个 epoch（batch size 100、学习率 1e-3 并按卡数线性放大、MultiStepLR 在 [20,40,60,80] 处以 1/3 衰减）后的参考精度，可直接作为验证环境是否正确的"标尺"：

方案	测试准确率
Single-GPU Baseline FP32	85.85%
Booster DDP + FP32	84.91%
Booster DDP + FP16	85.46%
Booster Low Level Zero	84.50%
Booster Gemini	84.60%

四、混合并行：数据 + 张量 + 流水并行（GPT 示例）

当模型规模增长到单卡显存放不下、或数据并行通信占比过高时，需要混合并行（hybrid parallel）——即同时启用数据并行、张量并行与流水线并行。原文档指出 Colossal-AI 支持多种形态的张量并行切分方式（如 1D、2D、2.5D、3D），可仅通过修改配置在它们之间切换；需要说明的是，这一提法对应的是 Colossal-AI 早期的张量切分配置体系，而在当前仓库中，该能力已收敛为由用户可配置的并行维度参数驱动，具体落地路径建议以 examples/language/gpt 下的实际代码为准。

4.1 仓库内的 GPT 分布式训练方案

examples/language/gpt 是了解混合并行的主入口，其 README.md 归纳了两条主流路线：

• Gemini 路线：通过 run_gemini.sh 启动 train_gpt_demo.py，在无需修改 HuggingFace 模型结构的前提下组合 Gemini / ZeRO DDP 与张量并行，可快速接入新模型；
• HybridParallelism 插件路线：位于 examples/language/gpt/hybridparallelism，以"插件"方式统一配置多种并行方法，训练与推理均可复用。

4.2 使用 hybridparallelism 插件一键微调 GPT-2

当前仓库内最适合快速跑通混合并行的是 examples/language/gpt/hybridparallelism，它提供 finetune.py（微调主脚本）、benchmark.py（基准测试）、data.py（数据加载）与现成启动脚本 run.sh：

cd examples/language/gpt/hybridparallelism

# run.sh 内容（4 卡，hybrid_parallel 插件，加载 HuggingFace gpt2）
torchrun --standalone --nproc_per_node 4 --master_port 29800 \
  finetune.py --target_f1 0.6 --plugin hybrid_parallel --model_type "gpt2"

若已有本地预训练权重，可追加 --pretrained_path "your/path/to/pretrained_model" 从本地加载。finetune.py 内部将 --plugin 限定为 torch_ddp、torch_ddp_fp16、gemini、low_level_zero、hybrid_parallel 之一（源码见 finetune.py），其中 hybrid_parallel 即对应混合并行插件。

4.3 并行维度的显式配置入口

更精细地控制各并行维度，可以参考同目录下 benchmark.py 暴露的命令行参数，其对应关系非常直观：

参数	含义	默认值
--tp	张量并行（tensor parallel）大小	1
--pp	流水线并行（pipeline parallel）大小	1
--sp	序列并行（sequence parallel）大小	1
--sp_mode	序列并行模式	ring_attn
--extra_dp	额外数据并行大小（供 Gemini 使用）	1

例如 --tp 2 --pp 2 即启用"张量并行 ×2 + 流水线并行 ×2"，再叠加数据并行构成 3D 混合并行。由代码可见，当启用序列并行时插件会设置 enable_sequence_parallelism=True 并传入 sequence_parallelism_mode，说明序列并行能力已被整合进新的混合并行插件体系。

五、MoE 并行：让模型容量与计算量解耦

混合专家（Mixture-of-Experts, MoE）并行面向的是"模型总参数量极大、但每个 token 只需激活一小部分专家"的场景。原文档以 ViT-MoE（WideNet）为例，说明 MoE 能让模型在参数量增大的同时保持接近恒定的计算量。

需要留意的是，原文档链接的 ViT-MoE 演示位于独立的 ColossalAI-Examples 仓库，并未随本仓库一起提供；当前仓库内与 MoE 相关的权威阅读材料与可运行资产包括：

• 官方教程 Tutorial: Integrate Mixture-of-Experts Into Your Model —— 原文档明确指向的进阶阅读，一步步讲解如何把 MoE 层集成进自己的模型，是理解"MoE 并行从哪来"的第一手材料；
• MoE 底层算子实现 colossalai/moe，包含专家路由等核心 _operation.py；
• 上层示例应用 applications/ColossalMoE，提供了可直接运行/推理的 train.sh、infer.sh；
• MoE 模型的端到端示例 examples/language/mixtral。

从仓库源码结构可以推断，MoE 并行的工程重点在于专家分配与 All-to-All 通信调度：MoE 层将 token 依据门控路由结果发送到对应专家所在的设备，再聚合各专家的输出。若想深入了解底层实现细节，建议从 colossalai/moe/_operation.py 的通信原语入手，并结合上述教程逐步实现一个带 MoE 层的模型。

六、序列并行：突破长序列训练的内存与长度上限

序列并行（sequence parallel）专门针对 NLP 任务中的两个痛点：显存效率与序列长度上限。其思想是沿序列（sequence）维度切分输入张量与中间激活，从而允许在更大 batch 与更长序列下训练。原文档给出的 BERT 示例在仓库中的对应位置是 examples/tutorial/sequence_parallel，该目录包含 config.py（配置）、train.py（训练入口）、model/（BERT 实现）、loss_func/、lr_scheduler/ 等完整结构，并配有独立的 README.md。

6.1 一键启动

export PYTHONPATH=$PWD   # 使示例内的 model/ 等包可被导入

# 使用 4 张 GPU 以序列并行方式训练 BERT（数据为合成数据）
colossalai run --nproc_per_node 4 train.py

6.2 关键配置解读

序列并行的开关全部集中在 examples/tutorial/sequence_parallel/config.py：

from colossalai.legacy.amp import AMP_TYPE

TRAIN_ITERS = 10
GLOBAL_BATCH_SIZE = 32      # dp world size × 每卡句子数
SEQ_LENGTH = 128            # 序列长度
DEPTH = 4
NUM_ATTENTION_HEADS = 4
HIDDEN_SIZE = 128           # 小规模 BERT，便于快速验证

# 核心：并行配置
parallel = dict(pipeline=1, tensor=dict(size=2, mode="sequence"))
fp16 = dict(mode=AMP_TYPE.NAIVE, verbose=True)
gradient_handler = [dict(type="SequenceParallelGradientHandler")]

• parallel.tensor.mode="sequence" 是开启序列并行的标志，size=2 表示把序列维度切分成 2 份（切到多卡）；如需在 8 卡上沿序列维并行，将 size 改为 8 即可；
• parallel.pipeline 控制是否叠加流水线并行，pipeline > 1 时会启用 NUM_MICRO_BATCHES 微批次调度（config 中默认 NUM_MICRO_BATCHES = 4，并注释"only enabled when pipeline > 1"）；
• 序列并行的梯度同步不能套用普通 DDP 的梯度处理策略，因此必须注册专用处理器 SequenceParallelGradientHandler；
• 同时叠加 AMP_TYPE.NAIVE 的 FP16 混合精度训练以进一步省显存。

6.3 变更并行规模

该 config.py 是一个纯 Python 配置对象，改完即生效，无需改动模型代码：

• 单机多卡：colossalai run --nproc_per_node <num_gpus> --master_addr localhost --master_port 29500 train.py；
• 多机多卡：README 建议改用 colossalai.launch_from_slurm 或 colossalai.launch_from_openmpi 对接 SLURM / OpenMPI 启动器；若使用自有启动器，可回退到默认的 colossalai.launch 函数。

这种"改配置、不改模型"的体验，与第四节中 hybridparallelism 插件把 --tp/--pp/--sp 暴露为命令行参数的设计一脉相承——并行策略与模型定义解耦，正是 Colossal-AI 的核心设计取向。

七、小结：从文档到可运行示例的对照路径

回顾原文档与当前仓库，可将"想用并行 → 找哪个示例 → 怎么启动"归纳如下：

1. 先跑通单卡：使用 examples/images/resnet，以 torchrun --standalone --nproc_per_node=1 train.py 获得基线精度（约 85%+）；
2. 扩展为数据并行：把 --nproc_per_node 改为卡数即可，代码零改动；
3. 需要张量/流水并行：进入 examples/language/gpt/hybridparallelism，用 --plugin hybrid_parallel 配合 --tp/--pp 配置维度；
4. 需要长序列 NLP 训练：参考 examples/tutorial/sequence_parallel/config.py 的 mode="sequence" 配置，并按 examples/tutorial/sequence_parallel/README.md 启动；
5. 需要 MoE 稀疏化容量扩展：研读 MoE 集成教程 并参考 applications/ColossalMoE。

无论选择哪种方案，统一的运行前提都是：已安装对应版本 PyTorch 与 Colossal-AI、具备 CUDA 设备，并按示例目录中的 requirements.txt 安装依赖。借助 Booster/Plugin 抽象，一套训练代码即可在单卡与多卡、不同并行策略之间平滑切换，这也是快速上手 Colossal-AI 最值得先掌握的思维模型。

【免费下载链接】ColossalAI Making large AI models cheaper, faster and more accessible 项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI