使用 ColossalAI Booster 新 API 在 GLUE 上分布式微调 BERT：四种 Plugin 模式实战指南

本指南基于 ColossalAI 官方示例 examples/tutorial/new_api/glue_bert/，完整讲解如何利用新版 Booster API + Plugin 插件体系，在 GLUE 基准任务上对 bert-base-uncased 进行多卡分布式微调。读完本文，你将掌握 torch_ddp、torch_ddp_fp16、gemini、low_level_zero 四种并行训练模式的选择与启动方式，理解 Booster.boost() 背后对模型、优化器、学习率调度器的统一包装流程，并能在 MRPC 等任务上复现文档给出的 F1 指标。

示例概览与项目定位

该示例位于仓库 examples/tutorial/new_api/glue_bert 目录，属于 ColossalAI 新版 API 教程系列（入口见 examples/tutorial/new_api/README.md）。它不同于旧版 colossalai.legacy 训练流程，而是以 Booster 与 Plugin 为核心的"统一加速接口"设计：用户只需要按常规 PyTorch 方式组装模型、优化器与 DataLoader，再通过一个 Booster 对象完成分布式并行改造。

目录内共 5 个文件，职责清晰：

文件	作用
finetune.py	主训练脚本：解析参数、初始化分布式环境、构造 Plugin/Booster、执行训练与评估
data.py	定义 GLUEDataBuilder：负责加载 GLUE 各任务数据、tokenize 预处理、按 Plugin 生成 DataLoader
requirements.txt	运行依赖清单（含 colossalai、datasets、transformers 等）
test_ci.sh	CI 冒烟脚本：循环跑 4 种 Plugin，并以 --target_f1 0.80 做自动化验收
README.md	官方使用说明，即本文展开的主体

命令行参数详解

训练脚本 finetune.py 通过 argparse 暴露 3 个核心参数，官方文档中列出的参数与含义如下：

参数	含义	取值 / 默认值
-t, --task	要运行的 GLUE 任务	默认 mrpc；其余可选值由 data.py 的 task_text_field_map 决定：cola、sst2、qqp、stsb、mnli、qnli、rte、wnli、ax 等
-p, --plugin	使用的并行 Plugin	torch_ddp、torch_ddp_fp16、gemini、low_level_zero，默认 torch_ddp
--target_f1	目标 F1 分数，若最终评估达不到则抛出异常（供 CI 验收用）	默认 None（不检查）

其中 --target_f1 的断言逻辑位于 finetune.py：仅当 coordinator.is_master() 且结果字典中存在 "f1" 键时才校验 results["f1"] >= args.target_f1。因此它适用于返回 F1 指标的 GLUE 任务（如 mrpc）；对以准确率为主的任务（如 cola、sst2）该参数不会触发。这一点在使用时需留意。

环境准备

安装依赖

进入示例目录后执行：

pip install -r requirements.txt

requirements.txt 声明的依赖包括：

colossalai
datasets
torch
tqdm
transformers
scipy
scikit-learn

其中 colossalai 提供 Booster/Plugin 体系，datasets + transformers 负责 GLUE 数据集与 BERT 预训练模型，scipy/scikit-learn 用于 GLUE 指标（如 Matthews 相关系数、F1）计算。tqdm 用于训练进度条展示。

数据与模型来源

• GLUE 数据通过 Hugging Face datasets 库按任务在线加载：datasets.load_dataset("glue", task_name)（见 data.py）。
• BERT 权重通过 transformers 从 bert-base-uncased 拉取：BertForSequenceClassification.from_pretrained(model_name, config=config)（见 finetune.py）。

两者均需网络可访问对应 Hub 资源，运行前请确保网络连通。

四种 Plugin 的运行方式与启动命令

官方 README 给出 4 条可直接运行的训练命令（默认在 MRPC 任务上、使用 4 个进程）：

# train with torch DDP with fp32
colossalai run --nproc_per_node 4 finetune.py

# train with torch DDP with mixed precision training
colossalai run --nproc_per_node 4 finetune.py -p torch_ddp_fp16

# train with gemini
colossalai run --nproc_per_node 4 finetune.py -p gemini

# train with low level zero
colossalai run --nproc_per_node 4 finetune.py -p low_level_zero

这里的 colossalai run 是 ColossalAI 自带的多机多卡启动器（launcher 实现位于 colossalai/cli/launcher/run.py）。需要说明的是，仓库自带的 CI 脚本 test_ci.sh 使用了等价的 torchrun --standalone --nproc_per_node 4 启动方式，因此两种入口均可触发同一套分布式训练流程。

-p/--plugin 四种模式的本质区别，对应 finetune.py 中 Booster 的构造逻辑：

booster_kwargs = {}
if args.plugin == "torch_ddp_fp16":
    booster_kwargs["mixed_precision"] = "fp16"
if args.plugin.startswith("torch_ddp"):
    plugin = TorchDDPPlugin()
elif args.plugin == "gemini":
    plugin = GeminiPlugin(placement_policy="static", strict_ddp_mode=True, initial_scale=2**5)
elif args.plugin == "low_level_zero":
    plugin = LowLevelZeroPlugin(initial_scale=2**5)

booster = Booster(plugin=plugin, **booster_kwargs)

逐一解读：

1. torch_ddp（默认）：直接使用 TorchDDPPlugin，等价于标准 PyTorch DistributedDataParallel 数据并行，FP32 精度。
2. torch_ddp_fp16：仍使用 TorchDDPPlugin，但额外向 Booster 传入 mixed_precision="fp16"。注意这里不是为 Plugin 传参，而是注入 Booster 的混合精度配置。根据 booster.py 的逻辑，TorchDDPPlugin 不控制精度（control_precision() 为假），因此字符串 "fp16" 会被 mixed_precision_factory 解析为 PyTorch AMP 策略并应用到模型与优化器上。
3. gemini：GeminiPlugin 是 ColossalAI 的异构内存管理系统，本示例采用 placement_policy="static" 静态放置策略、开启 strict_ddp_mode=True（严格 DDP 语义，保证与 DDP 数值行为一致），并用 initial_scale=2**5 设定 FP16 梯度缩放初值。Gemini 会在 GPU/CPU 之间动态或静态放置参数与优化器状态，以更小显存训练更大模型。
4. low_level_zero：LowLevelZeroPlugin 提供低层 ZeRO 优化（ZeRO-1 分片优化器状态等），同样以 initial_scale=2**5 启用 FP16 混合精度与梯度缩放。

超参数约定（源码内定义）

学习率与 batch 等超参数直接写在 finetune.py 顶部：

NUM_EPOCHS = 1
BATCH_SIZE = 32
LEARNING_RATE = 2.4e-5
WEIGHT_DECAY = 0.01
WARMUP_FRACTION = 0.1

值得强调的两处设计（与官方文档呼应，可从源码印证）：

• 学习率按世界大小线性缩放：lr = LEARNING_RATE * coordinator.world_size（见 finetune.py），即每张卡上实际使用的 batch 累计效应被折算进学习率，多卡规模越大学习率越高。
• 全局 batch 语义：脚本注释 local_batch_size = BATCH_SIZE // coordinator.world_size 表明 BATCH_SIZE=32 是全局总 batch，GLUEDataBuilder 会把该数值交给 Plugin 的 prepare_dataloader，由 Plugin 按进程数切分。默认 4 卡时每卡约 8 条样本。

主训练流程源码解析

finetune.py 的 main() 分为 6 个阶段，下面按官方文档"Quick Start"的脉络逐段拆解。

1. 启动分布式环境

colossalai.launch_from_torch(seed=42)
coordinator = DistCoordinator()

launch_from_torch 从 torchrun/colossalai run 注入的 LOCAL_RANK、RANK、WORLD_SIZE 环境变量中读取进程拓扑并初始化分布式组（底层实现见 colossalai/initialize.py，默认 seed 为 1024，这里显式固定为 42）。随后创建的 DistCoordinator（定义于 colossalai/cluster/dist_coordinator.py）封装了 rank、world_size、is_master() 等常用查询，供日志与评估阶段判断主进程。

2. 实例化 Plugin 与 Booster

如上一节所述，按 --plugin 选择具体 Plugin（gemini/low_level_zero 会自带 FP16 精度控制，torch_ddp 则为 FP32）。Booster(plugin=plugin, **booster_kwargs) 是统一入口，其构造与 boost() 的职责可参见 colossalai/booster/booster.py：

• 若 Plugin 控制设备（control_device()），则忽略外部 device 参数，由 Plugin 管理；否则默认使用 CUDA 加速器。
• 若 Plugin 控制精度（control_precision()），则忽略 mixed_precision；gemini、low_level_zero 即属此类。
• 检查点读写由 checkpoint_io 统一抽象。

3. 数据加载：GLUEDataBuilder

data_builder = GLUEDataBuilder(
    model_name, plugin, args.task, train_batch_size=BATCH_SIZE, eval_batch_size=BATCH_SIZE
)
train_dataloader = data_builder.train_dataloader()
test_dataloader = data_builder.test_dataloader()

data.py 中的 GLUEDataBuilder 是理解数据流的关键，稍后单独展开。

4. 组装模型、优化器与调度器

• 模型：AutoConfig.from_pretrained(model_name, num_labels=data_builder.num_labels) 依据任务设置输出类别数；data_builder.num_labels 来自 data.py 的 glue_task_num_labels 映射（如 MRPC 为 2 分类、STSB 为 1 个回归输出、MNLI 为 3 分类）。
• 优化器分组：标准 BERT 微调技巧——把 bias 与 LayerNorm.weight 排除出权重衰减组，其余参数施加 weight_decay=0.01：

no_decay = ["bias", "LayerNorm.weight"]
optimizer_grouped_parameters = [
    {
        "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
        "weight_decay": WEIGHT_DECAY,
    },
    {
        "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
        "weight_decay": 0.0,
    },
]
optimizer = HybridAdam(optimizer_grouped_parameters, lr=lr, eps=1e-8)

这里选用的是 ColossalAI 实现的 HybridAdam（colossalai/nn/optimizer/hybrid_adam.py）。从源码看，它融合了 CPUAdam 与 FusedAdam：参数在 GPU 上时走 Fused Adam CUDA kernel，参数在 CPU 上（如 Gemini 将部分状态卸载到 CPU）时走 CPUAdam，天然适配 Gemini 异构训练。该优化器要求 CUDA 扩展在安装或运行时完成编译。

• 学习率调度器：使用 transformers 的线性 warmup 调度，warmup 步数占总步数的 WARMUP_FRACTION=10%：

total_steps = len(train_dataloader) * NUM_EPOCHS
num_warmup_steps = int(WARMUP_FRACTION * total_steps)
lr_scheduler = get_linear_schedule_with_warmup(
    optimizer,
    num_warmup_steps=num_warmup_steps,
    num_training_steps=total_steps,
)

5. 统一加速：booster.boost()

model, optimizer, _, _, lr_scheduler = booster.boost(model, optimizer, lr_scheduler=lr_scheduler)

boost()（colossalai/booster/booster.py）会把传入对象交给 plugin.configure(...) 进行并行化包装，再叠加混合精度与加速器改造，最终返回包装后的 model/optimizer/lr_scheduler。经过这一步，模型可能被 DDP 包装（torch_ddp）、被 Gemini 静态分片（gemini）或被 ZeRO 优化器状态分片（low_level_zero），而外层代码无需感知差异——这正是新 API 的核心价值。

6. 训练与评估循环

训练单步的调用链（finetune.py）为：

booster.backward(loss, optimizer)   # 内部转调 optimizer.backward(loss)，适配 ZeRO/Gemini 的分片反向
optimizer.step()
optimizer.zero_grad()
lr_scheduler.step()

评估函数 evaluate（finetune.py）利用 datasets.load_metric("glue", task_name, process_id=coordinator.rank, num_process=coordinator.world_size) 做分布式指标聚合，并用 dist.all_reduce 汇总验证损失；只有主进程打印结果、执行 --target_f1 断言。多验证集任务（如 MNLI 的 validation_matched / validation_mismatched）会返回形如 {metric_matched: ..., metric_mismatched: ...} 的合并字典。

数据流水线 GLUEDataBuilder 详解

data.py 是官方文档没有展开、但支撑"换任务即可跑"的关键基础设施，体现在三张映射表与两类方法上。

任务字段映射 task_text_field_map 决定了每个 GLUE 任务用哪些列构造输入：

任务	文本字段	任务	文本字段
cola / sst2	sentence（单句）	mrpc/qqp/stsb/rte/wnli	sentence1+sentence2（句子对）
mnli / ax	premise+hypothesis	qnli	question+sentence

convert_to_features（data.py）对单句与句子对分别组装输入，并用 batch_encode_plus(..., max_length=128, padding="max_length", truncation=True) 统一截断/填充到 128 token，同时把原始 label 字段重命名为 labels，以匹配 BertForSequenceClassification 前向接口。多字段任务通过 zip 拼成句子对列表交给 tokenizer，逻辑简洁且可覆盖 9 类 GLUE 子任务。

标签数映射 glue_task_num_labels 与 DataLoader 三件套：

• train_dataloader()：plugin.prepare_dataloader(dataset["train"], batch_size=..., shuffle=True, drop_last=True)，由 Plugin 实现分布式 Sampler/切分，保证每张卡拿到的数据互不重叠。
• val_dataloader() / test_dataloader()：单一验证集任务返回单个 DataLoader；存在多个 validation* 切分（如 MNLI）时返回 DataLoader 列表，供 evaluate 循环处理。

此外，GLUEDataBuilder 接收的 plugin 参数类型为 DPPluginBase（colossalai/booster/plugin/dp_plugin_base.py），这也从类型层面印证了四种 Plugin 都归属数据并行基类。

期望的 F1 结果

README 官方表格给出 bert-base-uncased 在 MRPC 上的期望 F1 分数（供验证环境与代码正确性，实际结果可能因随机种子、环境差异略有浮动）：

Model	Single-GPU Baseline FP32	Booster DDP with FP32	Booster DDP with FP16	Booster Gemini	Booster Low Level Zero
bert-base-uncased	0.86	0.88	0.87	0.88	0.89

可以看到 Booster 各并行模式在 MRPC 上的 F1 均达到或超过单卡 FP32 基线。请勿将其理解为性能宣称或精确复现值——它更多是文档作者为便于用户核对训练正确性给出的参考量级。仓库对结果采用更稳妥的自动化校验方式：CI 脚本把阈值放宽到 0.80。

CI 自动化验收：test_ci.sh

仓库提供了 test_ci.sh，可作为一键复现脚本，其核心是串行跑遍四种 Plugin 并对 F1 做断言：

#!/bin/bash
set -xe

pip install -r requirements.txt

for plugin in "torch_ddp" "torch_ddp_fp16" "gemini" "low_level_zero"; do
    torchrun --standalone --nproc_per_node 4  finetune.py --target_f1 0.80 --plugin $plugin
done

要点解读：

• set -xe 保证任一步失败即中止（-e）并回显命令（-x），便于 CI 定位失败的是哪种 Plugin。
• 脚本使用 torchrun --standalone --nproc_per_node 4，与 README 中的 colossalai run --nproc_per_node 4 在单机场景下作用等价——两种启动器都会设置标准分布式环境变量，供 launch_from_torch 读取。
• --target_f1 0.80 即把官方 README 表格的最低参考值下调后作为自动通过阈值；这也提示你在自己机器上运行时，可直接用 --target_f1 快速校验安装与代码是否正确，而无需人工比对日志。

常见问题与扩展建议

1. 切换 GLUE 子任务：执行 colossalai run --nproc_per_node 4 finetune.py -t qqp（或 cola、sst2、mnli、stsb、rte、qnli、wnli）即可，无需改动代码——文本字段、标签数、评估切分均由 data.py 自动适配。
2. 指标与任务的匹配：MRPC/QQP 报告 F1；COLA 报告 Matthews 相关系数；STSB 是单标签回归任务（num_labels=1，评估走 logits.squeeze() 分支）。--target_f1 只在存在 f1 键时生效，对非 F1 任务请勿依赖该参数做验收。
3. MNLI 的多验证集：test_dataloader() 会对 validation_matched 与 validation_mismatched 分别返回 DataLoader，评估结果会拼接为带切分后缀的键名，便于分项观察。
4. 更换模型：把 finetune.py 的 model_name = "bert-base-uncased" 换成任意兼容的 transformers 序列分类模型 ID（如 roberta-base）通常即可，但需确认 tokenizer 与模型家族匹配。
5. 单机单卡/多卡适配：BATCH_SIZE=32 表示全局 batch，由 Plugin 按进程数拆分；若只想用 1 卡运行，把启动命令中的 --nproc_per_node 改为 1 即可，分布式代码无需改动。

更系统的 Booster/Plugin 用法还可参考 examples/tutorial/new_api 目录下的其他教程（如 cifar_resnet、cifar_vit），它们共用同一套 Booster + Plugin + DistCoordinator 骨架，便于横向对照学习新 API 在不同任务上的落地方式。