尧图网页设计 
Colossal-AI 流水并行Pipeline Parallelism实战1F1B 调度原理与基于 Booster 的 Bert 微调【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI流水并行Pipeline Parallelism是训练超大模型时把模型按层切分到不同设备上的核心技术与张量并行互补可在数据并行之外进一步扩大单模型规模。本文以 Colossal-AI 文档《流水并行》为主体结合仓库内pipeline/schedule调度器实现与 Bert/GLUE 微调示例讲透 GPipe 与 1F1B 的差异、非交错/交错两种调度方式的工作原理并给出可直接运行的HybridParallelPlugin Booster.execute_pipeline完整代码帮助你掌握在 Colossal-AI 中落地流水并行的配置方法、约束条件与调试手段。为什么需要流水并行从 GPipe 说起当单卡显存放不下完整模型时除了张量并行将单个算子切到多卡另一种直观思路是按层把模型纵向切开设备 1 计算层 1~4设备 2 计算层 5~8前一个设备算完才把中间结果传给下一个设备。这就是流水并行。为了让不同设备能“同时”工作而不是串行空等训练数据会被切成多个microbatch微批次依次灌入流水线。经典 GPipe 的做法是先把一个批次所有 microbatch 的前向全部算完再统一开始反向因此设备之间天然存在大量“气泡”bubble同时为了重算反向需要把多个 microbatch 的前向激活值长时间驻留在显存中内存压力也更大。认识 1F1B 流水线一次前向紧接着一次后向在 Colossal-AI 中默认采用的是1F1BOne Forward One Backward调度即每个设备交替执行“一次前向”和“一次后向”从而让前向激活值的生命周期大幅缩短。相对于 GPipe1F1B 在内存占用上明显更优前向一结束马上反向不必长期保留全部中间激活一轮迭代的总耗时也可通过更紧凑的流水调度得到改善。仓库中调度器统一继承自抽象基类PipelineSchedule见 colossalai/pipeline/schedule/base.py对外暴露三种具体实现见 colossalai/pipeline/schedule/init.pyOneForwardOneBackwardSchedule非交错 1F1BInterleavedSchedule交错 1F1BZeroBubbleVPipeSchedulerZero BubblezbvV 型调度。非交错 Schedule非交错式 1F1B 把一个完整训练步切分为三个阶段这一点可以直接从源码的执行函数run_forward_backwardcolossalai/pipeline/schedule/one_f_one_b.py#L359-L441中对照阅读热身阶段warmup靠近流水线尾部的设备需要先多跑若干个前向。源码用公式num_warmup_microbatches num_stages - stage - 1并与 microbatch 总数取较小值计算各自需要预热的前向次数one_f_one_b.py#L376-L377。稳态阶段1F1B设备每做一次前向紧接着就 pop 出最早保存的一对(input_obj, output_obj)做一次后向即 1F1B 节奏one_f_one_b.py#L403-L423。冷却阶段cooldown前向全部完成各设备把剩余积压的前向逐一补做后向one_f_one_b.py#L426-L432。由于前向与后向交替进行、中间激活被及时消费释放这种模式比 GPipe 更省显存但整条流水线所有设备完成一轮总工作量的时间与非交错调度的固有特性有关气泡仍随流水级数增加而放大。交错 ScheduleInterleaved交错调度的核心思想是不再让每个设备只负责一段连续层而是把模型切成num_model_chunks 个模型块model chunks按循环方式交错分配给各设备。例如原来设备 1 持有层 1~4、设备 2 持有层 5~8交错后设备 1 持有层 1、2、9、10设备 2 持有层 3、4、11、12。每个设备都同时扮演多个“虚拟流水级”的角色从而细化流水粒度、填满更多空闲官方文档的表述是既省内存又省时间气泡被显著压缩设备利用率更高。使用交错调度有一个硬性前提源码以断言形式固化在 colossalai/pipeline/schedule/interleaved_pp.py#L88-L90num_microbatch % stage_manager.num_stages 0即microbatch 的数量必须是流水级数的整数倍。交错调度在源码里通过get_model_chunk_id计算当前 microbatch 属于哪个模型块前向顺序取模、后向逆序见 interleaved_pp.py#L118-L136并按模型块分别维护input_objs/output_objs的 FIFO 队列来完成交错的 1F1B 稳态执行interleaved_pp.py#L419-L582。由于一个设备要切换处理多个模型块显存中同一时刻保存的激活自然增多所以交错调度通常需要配合梯度检查点gradient checkpointing等显存优化手段一起使用。Colossal-AI 中的流水并行实现Scheduler × Shardformer × HybridParallelPlugin流水并行不是独立存在的功能它由三部分协同完成调度器Scheduler决定每个设备在每个时刻该做前向、后向还是通信即上文的两类 schedule。Shardformer负责对模型按层做切分层分割并把模型的forward函数改写为与调度器兼容的形式参见 Shardformer 教程。HybridParallelPlugin这是新 Booster API 下管理流水并行的一站式入口。它内部会建立流水并行通信组在ProcessGroupMesh中沿流水轴构建PipelineStageManagercolossalai/pipeline/stage_manager.py依据pp_style参数实例化对应的 schedulecolossalai/booster/plugin/hybrid_parallel_plugin.py#L1150-L1176在boost时调用 Shardformer 对模型执行optimize层切分并通过execute_pipeline驱动 scheduler 逐个执行模型各流水段。关于pp_style早期版本文档描述插件“暂时只支持 1F1B、交错调度即将支持”而在当前仓库代码中HybridParallelPlugin已经支持三种风格[1f1b, interleaved, zbv]hybrid_parallel_plugin.py#L1124。相关约束如下pp_style1f1b时num_model_chunks必须等于 1hybrid_parallel_plugin.py#L1125-L1127pp_styleinterleaved时num_model_chunks必须大于 1hybrid_parallel_plugin.py#L1150-L1151只要开启流水pp_size 1就必须显式给出num_microbatches或microbatch_size二者之一hybrid_parallel_plugin.py#L1131-L1133为避免流水并行下梯度同步开销过大ZeRO 阶段只能取 0 或 1zero_stage in (0, 1, 2)且流水开启时限制为 0/1hybrid_parallel_plugin.py#L1122-L1136世界总卡数必须能被tp_size * pp_size整除hybrid_parallel_plugin.py#L1054。其余 Booster 插件的通用细节tp_size、zero_stage、precision、enable_all_optimization等可参考 Booster 插件文档与 Booster API 文档。实战使用流水并行微调 Bert文档与仓库示例均以Bert GLUEMRPC 等而非更重的 ViT/ImageNet 为例便于在较小规模上观察流水行为。仓库内可直接运行/对照的完整脚本是 examples/language/bert/finetune.py数据构造逻辑在 examples/language/bert/data.py 的GLUEDataBuilder中。第 1 步准备训练组件与普通训练一致需要先定义超参数、模型、数据加载器、优化器与学习率调度器import argparse from typing import Callable, List, Union import torch import torch.nn as nn from data import GLUEDataBuilder from torch.optim import Adam, Optimizer from torch.optim.lr_scheduler import _LRScheduler as LRScheduler from torch.utils.data import DataLoader from tqdm import tqdm from transformers import ( AlbertForSequenceClassification, AutoConfig, BertForSequenceClassification, get_linear_schedule_with_warmup, ) import colossalai from colossalai.booster import Booster from colossalai.booster.plugin import HybridParallelPlugin from colossalai.cluster import DistCoordinator from colossalai.nn.optimizer import HybridAdam # Define some config NUM_EPOCHS 3 BATCH_SIZE 32 LEARNING_RATE 2.4e-5 WEIGHT_DECAY 0.01 WARMUP_FRACTION 0.1 coordinator DistCoordinator() def move_to_cuda(batch): return {k: v.cuda() for k, v in batch.items()}其中值得注意的几点criterion 必须写成接收两个参数的可调用对象_criterion(outputs, inputs)因为它会被传给execute_pipeline由调度器在流水线末级last stage统一调用# Define criterion function with two inputs, which will be passed to execute_pipeline. def _criterion(outputs, inputs): return outputs.loss优化器使用 Colossal-AI 的HybridAdam融合优化器可配合 ZeRO 分片使用并按“是否包含bias/LayerNorm.weight”分组设置weight_decaylr LEARNING_RATE 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, lrlr, eps1e-8)学习率采用带 warmup 的线性衰减调度总步数由len(train_dataloader) * NUM_EPOCHS决定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, )模型加载 Hugging Face 预训练权重注意脚本还支持albert-xxlarge-v2通过--model_type albert切换cfg AutoConfig.from_pretrained(model_name, num_labelsdata_builder.num_labels) model BertForSequenceClassification.from_pretrained(model_name, configcfg)数据加载器由仓库封装的GLUEDataBuilder构造它会依据所选 plugin 自动处理 batch 的组织方式data_builder GLUEDataBuilder( model_name, plugin, args.task, train_batch_sizeBATCH_SIZE, eval_batch_sizeBATCH_SIZE, ) train_dataloader data_builder.train_dataloader()第 2 步用 HybridParallelPlugin 初始化 Booster文档给出的是典型的 1F1B 配置2 个流水段、tp_size1一个 batch 切分成 1 个 microbatch实际请按显存与卡数自行调整plugin HybridParallelPlugin(tp_size1, pp_size2, num_microbatchesNone, microbatch_size1, enable_all_optimizationTrue, zero_stage1, precisionfp16, initial_scale1) booster Booster(pluginplugin)各关键参数的含义与约束如下参数含义说明与约束tp_size张量并行度本例为 1即只开流水、不做张量并行pp_size流水并行度本例为 2对应 2 个流水段总卡数需满足世界大小 % (tp_size * pp_size) 0num_microbatchesmicrobatch 数量与microbatch_size二选一给出后会自动推导microbatch_size batch_size / num_microbatches要求能整除microbatch_size每个 microbatch 的样本数给出后会自动推导num_microbatches batch_size / microbatch_size训练模式下 1F1B 还要求num_microbatches num_stagesone_f_one_b.py#L93-L95enable_all_optimization是否启用 Shardformer 的全部优化算子融合、层切分等一般置Truezero_stageZeRO 分片阶段开启流水后只允许0或1precision混合精度策略如fp16另有bf16、fp8等可参考 混合精度训练initial_scaleFP16 初始动态 loss scale用于梯度缩放随后把优化特性注入训练组件。booster.boost返回顺序为(model, optimizer, criterion, dataloader, lr_scheduler)其中 criterion 用占位符_接收因为它由 plugin 包装进执行管线model, optimizer, _criterion, _, lr_scheduler booster.boost(model, optimizer, criterion_criterion, lr_schedulerlr_scheduler)第 3 步用 execute_pipeline 训练流水训练的关键差异在于不再手动执行loss model(batch)booster.backward而是把数据迭代器、模型、criterion、optimizer 一起交给booster.execute_pipeline由内部 scheduler 负责跨设备收发中间张量与梯度def train_epoch(epoch: int, model: nn.Module, optimizer: Optimizer, _criterion: Callable, lr_scheduler: LRScheduler, train_dataloader: DataLoader, booster: Booster, coordinator: DistCoordinator): is_pp_last_stage booster.plugin.stage_manager.is_last_stage() total_step len(train_dataloader) model.train() optimizer.zero_grad() # convert train_dataloader to a iterator train_dataloader_iter iter(train_dataloader) with tqdm(range(total_step), descfEpoch [{epoch 1}/{NUM_EPOCHS}], disablenot (is_pp_last_stage)) as pbar: # Forward pass for _ in pbar: outputs booster.execute_pipeline(train_dataloader_iter, model, _criterion, optimizer, return_lossTrue) # Backward and optimize if is_pp_last_stage: loss outputs[loss] pbar.set_postfix({loss: loss.item()}) optimizer.step() optimizer.zero_grad() lr_scheduler.step() # Train model for epoch in range(NUM_EPOCHS): train_epoch(epoch, model, optimizer, _criterion, lr_scheduler, train_dataloader, booster, coordinator)这段代码里有几个与流水并行强相关的细节is_pp_last_stage booster.plugin.stage_manager.is_last_stage()只有流水线末级设备才能拿到聚合后的真实 loss来自 scheduler 返回字典的outputs[loss]因此 loss 打印/进度条要按末级 gating避免每个 rank 都去读一个本不存在的 tensor。optimizer.step() / zero_grad() / lr_scheduler.step()放在每个 rank调度器只在内部负责正反向与 p2p 通信参数更新本身由各 rank 的优化器完成若开启 ZeRO/张量并行还会有相应的梯度规约逻辑。需要反向时torch.enable_grad()生效、优化器非空execute_pipeline走run_forward_backward而在无梯度torch.no_grad()的推理场景下会自动走run_forward_only见 one_f_one_b.py#L465-L472因此同一套 pipeline 接口也能用于 eval。第 4 步仓库内可运行版本与交错调度配置本文档示例在仓库中的可执行版本是 examples/language/bert/finetune.py它通过命令行参数选择 pluginpython finetune.py -t mrpc -p hybrid_parallel --model_type bert-p/--plugin支持torch_ddp、torch_ddp_fp16、gemini、low_level_zero、hybrid_parallel、torch_fsdp等选项对应不同的 Booster 插件便于在相同数据与模型下横向对比多种并行/优化策略。当你用-p hybrid_parallel运行时脚本内部实际构造的是交错调度配置finetune.py#L223-L237plugin HybridParallelPlugin( tp_size1, pp_size2, num_microbatchesNone, pp_styleinterleaved, num_model_chunks2, microbatch_size16, enable_all_optimizationTrue, zero_stage1, precisionfp16, initial_scale1, fp8_communicationFalse, )与文档示例相比这里把microbatch_size加大到 16并额外指定pp_styleinterleaved与num_model_chunks2。由于pp_size2、num_model_chunks2调度器内部实际会构建 4 个“虚拟流水级”每个设备交替处理两个模型块。脚本还示范了流水并行下的评估写法finetune.py#L58-L96判断isinstance(booster.plugin, HybridParallelPlugin) and booster.plugin.pp_size 1进入 pipeline 分支用booster.execute_pipeline(batch, model, criterion, return_lossTrue, return_outputsTrue)获取outputs[outputs][logits]与outputs[loss]并以stage_manager.is_last_stage(ignore_chunkTrue)识别末级设备交错调度下忽略 chunk 维度判断真正的流水末端最后由末级 rank 通过broadcast_object_list把预测结果广播给同一流水组内其它 rank 参与 GLUE metric 统计。--target_f1参数还可设定达标阈值用于 CI 回归验证。由于finetune.py内部调用的是colossalai.launch_from_torch(seed42)finetune.py#L206运行时需要通过外部分布式启动器如torchrun按tp_size * pp_size的规模拉起进程并在data.py的GLUEDataBuilder之外保证数据集可访问。常见问题与排错要点断言 “World size is not divisible by tp_size * pp_size”卡数与tp_size * pp_size不匹配请调整启动进程数。断言 “Either num_microbatches or microbatch_size should be provided”开启流水后必须二选一显式指定hybrid_parallel_plugin.py#L1131-L1133。断言 “Number of microbatch should be larger than number of stages”1F1B 非交错调度下 microbatch 数不能小于流水级数否则设备无法进入稳态one_f_one_b.py#L93-L95。断言 “Number of microbatch should be an integer multiple of number of pipeline parallel devices”交错调度要求 microbatch 数是流水级数的整数倍interleaved_pp.py#L88-L90。batch 无法整除num_microbatches/microbatch_size与 batch size 之间须满足整除关系load_batch中会校验见 one_f_one_b.py#L82-L91。loss 只在末级可见不要在非末级 rank 上读取outputs[loss]应按stage_manager.is_last_stage()判断后再打印或记录。显存受限交错调度相对更吃激活显存建议搭配梯度检查点ZeRO 阶段在流水下最多取 1若需更大优化空间可结合enable_all_optimizationShardformer 算子级优化实现。更完整的并行范式背景可继续阅读并行技术总览若想进一步探索 V 型/Zero Bubble 流水调度的原理与使用可参考仓库内的 Zero Bubble 流水并行文档与 rl_example_zero_bubble.py。【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 Colossal-AI 流水并行（Pipeline Parallelism）实战：1F1B 调度原理与基于 Booster 的 Bert 微调
Colossal-AI 流水并行Pipeline Parallelism实战1F1B 调度原理与基于 Booster 的 Bert 微调【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI流水并行Pipeline Parallelism是训练超大模型时把模型按层切分到不同设备上的核心技术与张量并行互补可在数据并行之外进一步扩大单模型规模。本文以 Colossal-AI 文档《流水并行》为主体结合仓库内pipeline/schedule调度器实现与 Bert/GLUE 微调示例讲透 GPipe 与 1F1B 的差异、非交错/交错两种调度方式的工作原理并给出可直接运行的HybridParallelPlugin Booster.execute_pipeline完整代码帮助你掌握在 Colossal-AI 中落地流水并行的配置方法、约束条件与调试手段。为什么需要流水并行从 GPipe 说起当单卡显存放不下完整模型时除了张量并行将单个算子切到多卡另一种直观思路是按层把模型纵向切开设备 1 计算层 1~4设备 2 计算层 5~8前一个设备算完才把中间结果传给下一个设备。这就是流水并行。为了让不同设备能“同时”工作而不是串行空等训练数据会被切成多个microbatch微批次依次灌入流水线。经典 GPipe 的做法是先把一个批次所有 microbatch 的前向全部算完再统一开始反向因此设备之间天然存在大量“气泡”bubble同时为了重算反向需要把多个 microbatch 的前向激活值长时间驻留在显存中内存压力也更大。认识 1F1B 流水线一次前向紧接着一次后向在 Colossal-AI 中默认采用的是1F1BOne Forward One Backward调度即每个设备交替执行“一次前向”和“一次后向”从而让前向激活值的生命周期大幅缩短。相对于 GPipe1F1B 在内存占用上明显更优前向一结束马上反向不必长期保留全部中间激活一轮迭代的总耗时也可通过更紧凑的流水调度得到改善。仓库中调度器统一继承自抽象基类PipelineSchedule见 colossalai/pipeline/schedule/base.py对外暴露三种具体实现见 colossalai/pipeline/schedule/init.pyOneForwardOneBackwardSchedule非交错 1F1BInterleavedSchedule交错 1F1BZeroBubbleVPipeSchedulerZero BubblezbvV 型调度。非交错 Schedule非交错式 1F1B 把一个完整训练步切分为三个阶段这一点可以直接从源码的执行函数run_forward_backwardcolossalai/pipeline/schedule/one_f_one_b.py#L359-L441中对照阅读热身阶段warmup靠近流水线尾部的设备需要先多跑若干个前向。源码用公式num_warmup_microbatches num_stages - stage - 1并与 microbatch 总数取较小值计算各自需要预热的前向次数one_f_one_b.py#L376-L377。稳态阶段1F1B设备每做一次前向紧接着就 pop 出最早保存的一对(input_obj, output_obj)做一次后向即 1F1B 节奏one_f_one_b.py#L403-L423。冷却阶段cooldown前向全部完成各设备把剩余积压的前向逐一补做后向one_f_one_b.py#L426-L432。由于前向与后向交替进行、中间激活被及时消费释放这种模式比 GPipe 更省显存但整条流水线所有设备完成一轮总工作量的时间与非交错调度的固有特性有关气泡仍随流水级数增加而放大。交错 ScheduleInterleaved交错调度的核心思想是不再让每个设备只负责一段连续层而是把模型切成num_model_chunks 个模型块model chunks按循环方式交错分配给各设备。例如原来设备 1 持有层 1~4、设备 2 持有层 5~8交错后设备 1 持有层 1、2、9、10设备 2 持有层 3、4、11、12。每个设备都同时扮演多个“虚拟流水级”的角色从而细化流水粒度、填满更多空闲官方文档的表述是既省内存又省时间气泡被显著压缩设备利用率更高。使用交错调度有一个硬性前提源码以断言形式固化在 colossalai/pipeline/schedule/interleaved_pp.py#L88-L90num_microbatch % stage_manager.num_stages 0即microbatch 的数量必须是流水级数的整数倍。交错调度在源码里通过get_model_chunk_id计算当前 microbatch 属于哪个模型块前向顺序取模、后向逆序见 interleaved_pp.py#L118-L136并按模型块分别维护input_objs/output_objs的 FIFO 队列来完成交错的 1F1B 稳态执行interleaved_pp.py#L419-L582。由于一个设备要切换处理多个模型块显存中同一时刻保存的激活自然增多所以交错调度通常需要配合梯度检查点gradient checkpointing等显存优化手段一起使用。Colossal-AI 中的流水并行实现Scheduler × Shardformer × HybridParallelPlugin流水并行不是独立存在的功能它由三部分协同完成调度器Scheduler决定每个设备在每个时刻该做前向、后向还是通信即上文的两类 schedule。Shardformer负责对模型按层做切分层分割并把模型的forward函数改写为与调度器兼容的形式参见 Shardformer 教程。HybridParallelPlugin这是新 Booster API 下管理流水并行的一站式入口。它内部会建立流水并行通信组在ProcessGroupMesh中沿流水轴构建PipelineStageManagercolossalai/pipeline/stage_manager.py依据pp_style参数实例化对应的 schedulecolossalai/booster/plugin/hybrid_parallel_plugin.py#L1150-L1176在boost时调用 Shardformer 对模型执行optimize层切分并通过execute_pipeline驱动 scheduler 逐个执行模型各流水段。关于pp_style早期版本文档描述插件“暂时只支持 1F1B、交错调度即将支持”而在当前仓库代码中HybridParallelPlugin已经支持三种风格[1f1b, interleaved, zbv]hybrid_parallel_plugin.py#L1124。相关约束如下pp_style1f1b时num_model_chunks必须等于 1hybrid_parallel_plugin.py#L1125-L1127pp_styleinterleaved时num_model_chunks必须大于 1hybrid_parallel_plugin.py#L1150-L1151只要开启流水pp_size 1就必须显式给出num_microbatches或microbatch_size二者之一hybrid_parallel_plugin.py#L1131-L1133为避免流水并行下梯度同步开销过大ZeRO 阶段只能取 0 或 1zero_stage in (0, 1, 2)且流水开启时限制为 0/1hybrid_parallel_plugin.py#L1122-L1136世界总卡数必须能被tp_size * pp_size整除hybrid_parallel_plugin.py#L1054。其余 Booster 插件的通用细节tp_size、zero_stage、precision、enable_all_optimization等可参考 Booster 插件文档与 Booster API 文档。实战使用流水并行微调 Bert文档与仓库示例均以Bert GLUEMRPC 等而非更重的 ViT/ImageNet 为例便于在较小规模上观察流水行为。仓库内可直接运行/对照的完整脚本是 examples/language/bert/finetune.py数据构造逻辑在 examples/language/bert/data.py 的GLUEDataBuilder中。第 1 步准备训练组件与普通训练一致需要先定义超参数、模型、数据加载器、优化器与学习率调度器import argparse from typing import Callable, List, Union import torch import torch.nn as nn from data import GLUEDataBuilder from torch.optim import Adam, Optimizer from torch.optim.lr_scheduler import _LRScheduler as LRScheduler from torch.utils.data import DataLoader from tqdm import tqdm from transformers import ( AlbertForSequenceClassification, AutoConfig, BertForSequenceClassification, get_linear_schedule_with_warmup, ) import colossalai from colossalai.booster import Booster from colossalai.booster.plugin import HybridParallelPlugin from colossalai.cluster import DistCoordinator from colossalai.nn.optimizer import HybridAdam # Define some config NUM_EPOCHS 3 BATCH_SIZE 32 LEARNING_RATE 2.4e-5 WEIGHT_DECAY 0.01 WARMUP_FRACTION 0.1 coordinator DistCoordinator() def move_to_cuda(batch): return {k: v.cuda() for k, v in batch.items()}其中值得注意的几点criterion 必须写成接收两个参数的可调用对象_criterion(outputs, inputs)因为它会被传给execute_pipeline由调度器在流水线末级last stage统一调用# Define criterion function with two inputs, which will be passed to execute_pipeline. def _criterion(outputs, inputs): return outputs.loss优化器使用 Colossal-AI 的HybridAdam融合优化器可配合 ZeRO 分片使用并按“是否包含bias/LayerNorm.weight”分组设置weight_decaylr LEARNING_RATE 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, lrlr, eps1e-8)学习率采用带 warmup 的线性衰减调度总步数由len(train_dataloader) * NUM_EPOCHS决定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, )模型加载 Hugging Face 预训练权重注意脚本还支持albert-xxlarge-v2通过--model_type albert切换cfg AutoConfig.from_pretrained(model_name, num_labelsdata_builder.num_labels) model BertForSequenceClassification.from_pretrained(model_name, configcfg)数据加载器由仓库封装的GLUEDataBuilder构造它会依据所选 plugin 自动处理 batch 的组织方式data_builder GLUEDataBuilder( model_name, plugin, args.task, train_batch_sizeBATCH_SIZE, eval_batch_sizeBATCH_SIZE, ) train_dataloader data_builder.train_dataloader()第 2 步用 HybridParallelPlugin 初始化 Booster文档给出的是典型的 1F1B 配置2 个流水段、tp_size1一个 batch 切分成 1 个 microbatch实际请按显存与卡数自行调整plugin HybridParallelPlugin(tp_size1, pp_size2, num_microbatchesNone, microbatch_size1, enable_all_optimizationTrue, zero_stage1, precisionfp16, initial_scale1) booster Booster(pluginplugin)各关键参数的含义与约束如下参数含义说明与约束tp_size张量并行度本例为 1即只开流水、不做张量并行pp_size流水并行度本例为 2对应 2 个流水段总卡数需满足世界大小 % (tp_size * pp_size) 0num_microbatchesmicrobatch 数量与microbatch_size二选一给出后会自动推导microbatch_size batch_size / num_microbatches要求能整除microbatch_size每个 microbatch 的样本数给出后会自动推导num_microbatches batch_size / microbatch_size训练模式下 1F1B 还要求num_microbatches num_stagesone_f_one_b.py#L93-L95enable_all_optimization是否启用 Shardformer 的全部优化算子融合、层切分等一般置Truezero_stageZeRO 分片阶段开启流水后只允许0或1precision混合精度策略如fp16另有bf16、fp8等可参考 混合精度训练initial_scaleFP16 初始动态 loss scale用于梯度缩放随后把优化特性注入训练组件。booster.boost返回顺序为(model, optimizer, criterion, dataloader, lr_scheduler)其中 criterion 用占位符_接收因为它由 plugin 包装进执行管线model, optimizer, _criterion, _, lr_scheduler booster.boost(model, optimizer, criterion_criterion, lr_schedulerlr_scheduler)第 3 步用 execute_pipeline 训练流水训练的关键差异在于不再手动执行loss model(batch)booster.backward而是把数据迭代器、模型、criterion、optimizer 一起交给booster.execute_pipeline由内部 scheduler 负责跨设备收发中间张量与梯度def train_epoch(epoch: int, model: nn.Module, optimizer: Optimizer, _criterion: Callable, lr_scheduler: LRScheduler, train_dataloader: DataLoader, booster: Booster, coordinator: DistCoordinator): is_pp_last_stage booster.plugin.stage_manager.is_last_stage() total_step len(train_dataloader) model.train() optimizer.zero_grad() # convert train_dataloader to a iterator train_dataloader_iter iter(train_dataloader) with tqdm(range(total_step), descfEpoch [{epoch 1}/{NUM_EPOCHS}], disablenot (is_pp_last_stage)) as pbar: # Forward pass for _ in pbar: outputs booster.execute_pipeline(train_dataloader_iter, model, _criterion, optimizer, return_lossTrue) # Backward and optimize if is_pp_last_stage: loss outputs[loss] pbar.set_postfix({loss: loss.item()}) optimizer.step() optimizer.zero_grad() lr_scheduler.step() # Train model for epoch in range(NUM_EPOCHS): train_epoch(epoch, model, optimizer, _criterion, lr_scheduler, train_dataloader, booster, coordinator)这段代码里有几个与流水并行强相关的细节is_pp_last_stage booster.plugin.stage_manager.is_last_stage()只有流水线末级设备才能拿到聚合后的真实 loss来自 scheduler 返回字典的outputs[loss]因此 loss 打印/进度条要按末级 gating避免每个 rank 都去读一个本不存在的 tensor。optimizer.step() / zero_grad() / lr_scheduler.step()放在每个 rank调度器只在内部负责正反向与 p2p 通信参数更新本身由各 rank 的优化器完成若开启 ZeRO/张量并行还会有相应的梯度规约逻辑。需要反向时torch.enable_grad()生效、优化器非空execute_pipeline走run_forward_backward而在无梯度torch.no_grad()的推理场景下会自动走run_forward_only见 one_f_one_b.py#L465-L472因此同一套 pipeline 接口也能用于 eval。第 4 步仓库内可运行版本与交错调度配置本文档示例在仓库中的可执行版本是 examples/language/bert/finetune.py它通过命令行参数选择 pluginpython finetune.py -t mrpc -p hybrid_parallel --model_type bert-p/--plugin支持torch_ddp、torch_ddp_fp16、gemini、low_level_zero、hybrid_parallel、torch_fsdp等选项对应不同的 Booster 插件便于在相同数据与模型下横向对比多种并行/优化策略。当你用-p hybrid_parallel运行时脚本内部实际构造的是交错调度配置finetune.py#L223-L237plugin HybridParallelPlugin( tp_size1, pp_size2, num_microbatchesNone, pp_styleinterleaved, num_model_chunks2, microbatch_size16, enable_all_optimizationTrue, zero_stage1, precisionfp16, initial_scale1, fp8_communicationFalse, )与文档示例相比这里把microbatch_size加大到 16并额外指定pp_styleinterleaved与num_model_chunks2。由于pp_size2、num_model_chunks2调度器内部实际会构建 4 个“虚拟流水级”每个设备交替处理两个模型块。脚本还示范了流水并行下的评估写法finetune.py#L58-L96判断isinstance(booster.plugin, HybridParallelPlugin) and booster.plugin.pp_size 1进入 pipeline 分支用booster.execute_pipeline(batch, model, criterion, return_lossTrue, return_outputsTrue)获取outputs[outputs][logits]与outputs[loss]并以stage_manager.is_last_stage(ignore_chunkTrue)识别末级设备交错调度下忽略 chunk 维度判断真正的流水末端最后由末级 rank 通过broadcast_object_list把预测结果广播给同一流水组内其它 rank 参与 GLUE metric 统计。--target_f1参数还可设定达标阈值用于 CI 回归验证。由于finetune.py内部调用的是colossalai.launch_from_torch(seed42)finetune.py#L206运行时需要通过外部分布式启动器如torchrun按tp_size * pp_size的规模拉起进程并在data.py的GLUEDataBuilder之外保证数据集可访问。常见问题与排错要点断言 “World size is not divisible by tp_size * pp_size”卡数与tp_size * pp_size不匹配请调整启动进程数。断言 “Either num_microbatches or microbatch_size should be provided”开启流水后必须二选一显式指定hybrid_parallel_plugin.py#L1131-L1133。断言 “Number of microbatch should be larger than number of stages”1F1B 非交错调度下 microbatch 数不能小于流水级数否则设备无法进入稳态one_f_one_b.py#L93-L95。断言 “Number of microbatch should be an integer multiple of number of pipeline parallel devices”交错调度要求 microbatch 数是流水级数的整数倍interleaved_pp.py#L88-L90。batch 无法整除num_microbatches/microbatch_size与 batch size 之间须满足整除关系load_batch中会校验见 one_f_one_b.py#L82-L91。loss 只在末级可见不要在非末级 rank 上读取outputs[loss]应按stage_manager.is_last_stage()判断后再打印或记录。显存受限交错调度相对更吃激活显存建议搭配梯度检查点ZeRO 阶段在流水下最多取 1若需更大优化空间可结合enable_all_optimizationShardformer 算子级优化实现。更完整的并行范式背景可继续阅读并行技术总览若想进一步探索 V 型/Zero Bubble 流水调度的原理与使用可参考仓库内的 Zero Bubble 流水并行文档与 rl_example_zero_bubble.py。【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考