快速上手：从单卡到 8 卡的最小例子

ColossalAI 的最大特色是"换一行 plugin 就换一种并行策略"。本章用一份 30 行训练脚本， 演示从单卡到 8 卡 ZeRO / Gemini / 3D 并行 的切换体验。重点不是看每个 plugin 内部细节， 而是体会"plugin 抽象"是怎么把分布式细节藏起来的。

3.1训练脚本骨架（30 行）

所有 ColossalAI 训练脚本都长这个样子：

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from torch.optim import AdamW
import colossalai
from colossalai.booster import Booster
from colossalai.booster.plugin import GeminiPlugin

# 1. 初始化 colossalai（默认 torch.distributed 后端）
colossalai.launch_from_torch()

# 2. 选 plugin（这一行决定走哪种并行）
plugin = GeminiPlugin(
    placement_policy="auto",     # 自动 GPU/CPU 卸载
    precision="bf16",
)

# 3. 构造 Booster
booster = Booster(plugin=plugin)

# 4. 加载模型 / 数据 / optimizer（写法和普通 PyTorch 一样）
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-4B-Base", torch_dtype=torch.bfloat16)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-4B-Base")
optimizer = AdamW(model.parameters(), lr=5e-5)
dataloader = build_my_dataloader(...)        # 你自己的 dataloader

# 5. 一行"加速"
model, optimizer, _, dataloader, _ = booster.boost(model, optimizer, None, dataloader, None)

# 6. 训练循环：写法和单卡 PyTorch 一模一样
for step, batch in enumerate(dataloader):
    outputs = model(**batch)
    loss = outputs.loss
    booster.backward(loss, optimizer)     # ★ 用 booster.backward，不是 loss.backward()
    optimizer.step()
    optimizer.zero_grad()

# 7. 保存
booster.save_model(model, "ckpt", shard=True)

关键的三处变化对比纯 PyTorch：

除了这三处，其他代码跟单卡 PyTorch 一模一样。这就是 Booster 的"封装哲学"。

3.2启动命令：colossalai run

所有 plugin 的启动都用同一个命令：

# 单卡
colossalai run --nproc_per_node 1 train.py

# 单机 8 卡
colossalai run --nproc_per_node 8 train.py

# 2 节点 × 8 卡
colossalai run --nproc_per_node 8 --hostfile hostfile train.py
# hostfile 内容：
# worker-0 slots=8
# worker-1 slots=8

colossalai run 是 torchrun 的包装，多加了一些环境变量设置（如 NCCL_DEBUG、CUDA_DEVICE_MAX_CONNECTIONS）。 你也可以直接 torchrun --nproc_per_node 8 train.py，等价。

3.3切 plugin = 切并行策略

同一份 train.py，只换第 2 行的 plugin，就能跑不同并行模式：

切换示例：

# GeminiPlugin（默认场景）
plugin = GeminiPlugin(placement_policy="auto", precision="bf16")

# ZeRO-2
from colossalai.booster.plugin import LowLevelZeroPlugin
plugin = LowLevelZeroPlugin(stage=2, precision="bf16")

# 3D 并行
from colossalai.booster.plugin import HybridParallelPlugin
plugin = HybridParallelPlugin(
    tp_size=2, pp_size=2, zero_stage=1,
    precision="bf16",
    enable_fused_normalization=True,
    enable_flash_attention=True,
)

# PyTorch DDP
from colossalai.booster.plugin import TorchDDPPlugin
plugin = TorchDDPPlugin()

用户脚本的其他代码不变。这就是 ColossalAI 的核心承诺。

3.4启动流程做了什么

3.54 个 plugin 跑同一份代码的对比（13B 模型，8×A100 80G）

同一份训练代码不动，只换一行 plugin 就能扫不同方案。这种"灰度对比"在 ColossalAI 里成本接近为 0。

3.6看一份完整能跑的脚本

仓库 examples/language/llama/finetune.py 是 ColossalAI 官方维护的微调入门：

# 简化展示（约 100 行）
import torch, argparse
from transformers import AutoModelForCausalLM, AutoTokenizer, get_scheduler
from torch.optim import AdamW
from datasets import load_dataset
import colossalai
from colossalai.booster import Booster
from colossalai.booster.plugin import GeminiPlugin, HybridParallelPlugin

def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--model_type", default="meta-llama/Llama-3.2-1B")
    p.add_argument("--plugin", default="gemini", choices=["gemini", "hybrid_parallel", "ddp"])
    p.add_argument("--batch_size", type=int, default=4)
    p.add_argument("--max_length", type=int, default=2048)
    p.add_argument("--lr", type=float, default=5e-5)
    p.add_argument("--epoch", type=int, default=3)
    return p.parse_args()

def main():
    args = parse_args()
    colossalai.launch_from_torch()

    # 按命令行选 plugin
    if args.plugin == "gemini":
        plugin = GeminiPlugin(placement_policy="auto", precision="bf16")
    elif args.plugin == "hybrid_parallel":
        plugin = HybridParallelPlugin(tp_size=2, pp_size=1, zero_stage=1, precision="bf16")
    elif args.plugin == "ddp":
        from colossalai.booster.plugin import TorchDDPPlugin
        plugin = TorchDDPPlugin()

    booster = Booster(plugin=plugin)

    tokenizer = AutoTokenizer.from_pretrained(args.model_type)
    model     = AutoModelForCausalLM.from_pretrained(args.model_type)
    optimizer = AdamW(model.parameters(), lr=args.lr)

    dataset = load_dataset("yahma/alpaca-cleaned", split="train")
    # ... tokenize + collator + dataloader ...

    lr_scheduler = get_scheduler("cosine", optimizer, num_training_steps=len(dataloader)*args.epoch)

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

    for epoch in range(args.epoch):
        for step, batch in enumerate(dataloader):
            outputs = model(**batch)
            loss = outputs.loss
            booster.backward(loss, optimizer)
            optimizer.step()
            lr_scheduler.step()
            optimizer.zero_grad()

    booster.save_model(model, "out", shard=True)

if __name__ == "__main__":
    main()

启动

# 单卡
colossalai run --nproc_per_node 1 finetune.py --plugin gemini

# 8 卡 ZeRO
colossalai run --nproc_per_node 8 finetune.py --plugin gemini

# 8 卡 TP=2 + ZeRO-1
colossalai run --nproc_per_node 8 finetune.py --plugin hybrid_parallel

3.7训练时盯什么

ColossalAI 的 logging 跟 HuggingFace Trainer 不太一样，需要自己写。常用监控：

因为 ColossalAI不强制 trainer 模式（不像 HF Trainer / LLaMA-Factory），logging 要自己写一点点胶水代码。

3.8这章你需要带走的

• ColossalAI 训练脚本骨架 = "launch + plugin + Booster + boost + backward"；
• 切并行策略 = 换一行 plugin；
• 必须用 booster.backward(loss, optimizer)，不能 loss.backward()；
• 启动用 colossalai run --nproc_per_node N 或 torchrun；
• 常见组合：DDP（baseline）→ ZeRO-2 → Gemini → HybridParallel；
• logging 要自己写（无 trainer 自动）。