使用 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主训练脚本解析参数、初始化分布式环境、构造 Plugin/Booster、执行训练与评估data.py定义GLUEDataBuilder负责加载 GLUE 各任务数据、tokenize 预处理、按 Plugin 生成 DataLoaderrequirements.txt运行依赖清单含colossalai、datasets、transformers等test_ci.shCI 冒烟脚本循环跑 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使用的并行 Plugintorch_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.txtrequirements.txt 声明的依赖包括colossalai datasets torch tqdm transformers scipy scikit-learn其中colossalai提供 Booster/Plugin 体系datasetstransformers负责 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, configconfig)见 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_policystatic, strict_ddp_modeTrue, initial_scale2**5) elif args.plugin low_level_zero: plugin LowLevelZeroPlugin(initial_scale2**5) booster Booster(pluginplugin, **booster_kwargs)逐一解读torch_ddp默认直接使用 TorchDDPPlugin等价于标准 PyTorchDistributedDataParallel数据并行FP32 精度。torch_ddp_fp16仍使用TorchDDPPlugin但额外向 Booster 传入mixed_precisionfp16。注意这里不是为 Plugin 传参而是注入 Booster 的混合精度配置。根据 booster.py 的逻辑TorchDDPPlugin不控制精度control_precision()为假因此字符串fp16会被mixed_precision_factory解析为 PyTorch AMP 策略并应用到模型与优化器上。geminiGeminiPlugin 是 ColossalAI 的异构内存管理系统本示例采用placement_policystatic静态放置策略、开启strict_ddp_modeTrue严格 DDP 语义保证与 DDP 数值行为一致并用initial_scale2**5设定 FP16 梯度缩放初值。Gemini 会在 GPU/CPU 之间动态或静态放置参数与优化器状态以更小显存训练更大模型。low_level_zeroLowLevelZeroPlugin 提供低层 ZeRO 优化ZeRO-1 分片优化器状态等同样以initial_scale2**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_SIZE32是全局总 batchGLUEDataBuilder会把该数值交给 Plugin 的prepare_dataloader由 Plugin 按进程数切分。默认 4 卡时每卡约 8 条样本。主训练流程源码解析finetune.py 的main()分为 6 个阶段下面按官方文档Quick Start的脉络逐段拆解。1. 启动分布式环境colossalai.launch_from_torch(seed42) 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选择具体 Plugingemini/low_level_zero会自带 FP16 精度控制torch_ddp则为 FP32。Booster(pluginplugin, **booster_kwargs)是统一入口其构造与boost()的职责可参见 colossalai/booster/booster.py若 Plugin 控制设备control_device()则忽略外部device参数由 Plugin 管理否则默认使用 CUDA 加速器。若 Plugin 控制精度control_precision()则忽略mixed_precisiongemini、low_level_zero即属此类。检查点读写由checkpoint_io统一抽象。3. 数据加载GLUEDataBuilderdata_builder GLUEDataBuilder( model_name, plugin, args.task, train_batch_sizeBATCH_SIZE, eval_batch_sizeBATCH_SIZE ) train_dataloader data_builder.train_dataloader() test_dataloader data_builder.test_dataloader()data.py 中的GLUEDataBuilder是理解数据流的关键稍后单独展开。4. 组装模型、优化器与调度器模型AutoConfig.from_pretrained(model_name, num_labelsdata_builder.num_labels)依据任务设置输出类别数data_builder.num_labels来自 data.py 的glue_task_num_labels映射如 MRPC 为 2 分类、STSB 为 1 个回归输出、MNLI 为 3 分类。优化器分组标准 BERT 微调技巧——把bias与LayerNorm.weight排除出权重衰减组其余参数施加weight_decay0.01no_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, lrlr, eps1e-8)这里选用的是 ColossalAI 实现的HybridAdamcolossalai/nn/optimizer/hybrid_adam.py。从源码看它融合了 CPUAdam 与 FusedAdam参数在 GPU 上时走 Fused Adam CUDA kernel参数在 CPU 上如 Gemini 将部分状态卸载到 CPU时走 CPUAdam天然适配 Gemini 异构训练。该优化器要求 CUDA 扩展在安装或运行时完成编译。学习率调度器使用transformers的线性 warmup 调度warmup 步数占总步数的WARMUP_FRACTION10%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_stepsnum_warmup_steps, num_training_stepstotal_steps, )5. 统一加速booster.boost()model, optimizer, _, _, lr_scheduler booster.boost(model, optimizer, lr_schedulerlr_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()评估函数evaluatefinetune.py利用datasets.load_metric(glue, task_name, process_idcoordinator.rank, num_processcoordinator.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/sst2sentence单句mrpc/qqp/stsb/rte/wnlisentence1sentence2句子对mnli/axpremisehypothesisqnliquestionsentenceconvert_to_featuresdata.py对单句与句子对分别组装输入并用batch_encode_plus(..., max_length128, paddingmax_length, truncationTrue)统一截断/填充到 128 token同时把原始label字段重命名为labels以匹配BertForSequenceClassification前向接口。多字段任务通过zip拼成句子对列表交给 tokenizer逻辑简洁且可覆盖 9 类 GLUE 子任务。标签数映射glue_task_num_labels与 DataLoader 三件套train_dataloader()plugin.prepare_dataloader(dataset[train], batch_size..., shuffleTrue, drop_lastTrue)由 Plugin 实现分布式 Sampler/切分保证每张卡拿到的数据互不重叠。val_dataloader()/test_dataloader()单一验证集任务返回单个 DataLoader存在多个validation*切分如 MNLI时返回 DataLoader 列表供evaluate循环处理。此外GLUEDataBuilder接收的plugin参数类型为DPPluginBasecolossalai/booster/plugin/dp_plugin_base.py这也从类型层面印证了四种 Plugin 都归属数据并行基类。期望的 F1 结果README 官方表格给出bert-base-uncased在 MRPC 上的期望 F1 分数供验证环境与代码正确性实际结果可能因随机种子、环境差异略有浮动ModelSingle-GPU Baseline FP32Booster DDP with FP32Booster DDP with FP16Booster GeminiBooster Low Level Zerobert-base-uncased0.860.880.870.880.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快速校验安装与代码是否正确而无需人工比对日志。常见问题与扩展建议切换 GLUE 子任务执行colossalai run --nproc_per_node 4 finetune.py -t qqp或cola、sst2、mnli、stsb、rte、qnli、wnli即可无需改动代码——文本字段、标签数、评估切分均由 data.py 自动适配。指标与任务的匹配MRPC/QQP 报告 F1COLA 报告 Matthews 相关系数STSB 是单标签回归任务num_labels1评估走logits.squeeze()分支。--target_f1只在存在f1键时生效对非 F1 任务请勿依赖该参数做验收。MNLI 的多验证集test_dataloader()会对validation_matched与validation_mismatched分别返回 DataLoader评估结果会拼接为带切分后缀的键名便于分项观察。更换模型把 finetune.py 的model_name bert-base-uncased换成任意兼容的 transformers 序列分类模型 ID如roberta-base通常即可但需确认 tokenizer 与模型家族匹配。单机单卡/多卡适配BATCH_SIZE32表示全局 batch由 Plugin 按进程数拆分若只想用 1 卡运行把启动命令中的--nproc_per_node改为1即可分布式代码无需改动。更系统的 Booster/Plugin 用法还可参考 examples/tutorial/new_api 目录下的其他教程如 cifar_resnet、cifar_vit它们共用同一套BoosterPluginDistCoordinator骨架便于横向对照学习新 API 在不同任务上的落地方式。【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考