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

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

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

本指南基于 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 通过argparse暴露 3 个核心参数，官方文档中列出的参数与含义如下：

其中--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 Facedatasets库按任务在线加载：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，等价于标准 PyTorchDistributedDataParallel数据并行，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 任务用哪些列构造输入：

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 分数（供验证环境与代码正确性，实际结果可能因随机种子、环境差异略有浮动）：

可以看到 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 在不同任务上的落地方式。

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