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 中默认采用的是1F1B（One Forward One Backward）调度，即每个设备交替执行“一次前向”和“一次后向”，从而让前向激活值的生命周期大幅缩短。相对于 GPipe，1F1B 在内存占用上明显更优（前向一结束马上反向，不必长期保留全部中间激活），一轮迭代的总耗时也可通过更紧凑的流水调度得到改善。

仓库中调度器统一继承自抽象基类PipelineSchedule（见 colossalai/pipeline/schedule/base.py），对外暴露三种具体实现（见 colossalai/pipeline/schedule/init.py）：

• OneForwardOneBackwardSchedule：非交错 1F1B；
• InterleavedSchedule：交错 1F1B；
• ZeroBubbleVPipeScheduler：Zero Bubble（zbv）V 型调度。

非交错 Schedule

非交错式 1F1B 把一个完整训练步切分为三个阶段，这一点可以直接从源码的执行函数run_forward_backward（colossalai/pipeline/schedule/one_f_one_b.py#L359-L441）中对照阅读：

1. 热身阶段（warmup）：靠近流水线尾部的设备需要先多跑若干个前向。源码用公式num_warmup_microbatches = num_stages - stage - 1（并与 microbatch 总数取较小值）计算各自需要预热的前向次数（one_f_one_b.py#L376-L377）。
2. 稳态阶段（1F1B）：设备每做一次前向，紧接着就 pop 出最早保存的一对(input_obj, output_obj)做一次后向，即 1F1B 节奏（one_f_one_b.py#L403-L423）。
3. 冷却阶段（cooldown）：前向全部完成，各设备把剩余积压的前向逐一补做后向（one_f_one_b.py#L426-L432）。

由于前向与后向交替进行、中间激活被及时消费释放，这种模式比 GPipe 更省显存；但整条流水线所有设备完成一轮总工作量的时间与非交错调度的固有特性有关，气泡仍随流水级数增加而放大。

交错 Schedule（Interleaved）

交错调度的核心思想是不再让每个设备只负责一段连续层，而是把模型切成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-L90：

num_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

流水并行不是独立存在的功能，它由三部分协同完成：

1. 调度器（Scheduler）：决定每个设备在每个时刻该做前向、后向还是通信，即上文的两类 schedule。
2. Shardformer：负责对模型按层做切分（层分割），并把模型的forward函数改写为与调度器兼容的形式（参见 Shardformer 教程）。
3. HybridParallelPlugin：这是新 Booster API 下管理流水并行的一站式入口。它内部会：
   • 建立流水并行通信组（在ProcessGroupMesh中沿流水轴构建PipelineStageManager，colossalai/pipeline/stage_manager.py）；
   • 依据pp_style参数实例化对应的 schedule（colossalai/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_style="1f1b"时num_model_chunks必须等于 1（hybrid_parallel_plugin.py#L1125-L1127）；
• pp_style="interleaved"时num_model_chunks必须大于 1（hybrid_parallel_plugin.py#L1150-L1151）；
• 只要开启流水（pp_size > 1），就必须显式给出num_microbatches或microbatch_size二者之一（hybrid_parallel_plugin.py#L1131-L1133）；
• 为避免流水并行下梯度同步开销过大，ZeRO 阶段只能取 0 或 1（zero_stage in (0, 1, 2)且流水开启时限制为 0/1，hybrid_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 + GLUE（MRPC 等）而非更重的 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_decay：

  lr = 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, lr=lr, eps=1e-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_steps=num_warmup_steps, num_training_steps=total_steps, )

• 模型加载 Hugging Face 预训练权重（注意脚本还支持albert-xxlarge-v2，通过--model_type albert切换）：

  cfg = AutoConfig.from_pretrained(model_name, num_labels=data_builder.num_labels) model = BertForSequenceClassification.from_pretrained(model_name, config=cfg)

• 数据加载器由仓库封装的GLUEDataBuilder构造，它会依据所选 plugin 自动处理 batch 的组织方式：

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

第 2 步：用 HybridParallelPlugin 初始化 Booster

文档给出的是典型的 1F1B 配置：2 个流水段、tp_size=1，一个 batch 切分成 1 个 microbatch（实际请按显存与卡数自行调整）：

plugin = HybridParallelPlugin(tp_size=1, pp_size=2, num_microbatches=None, microbatch_size=1, enable_all_optimization=True, zero_stage=1, precision='fp16', initial_scale=1) booster = Booster(plugin=plugin)

各关键参数的含义与约束如下：

参数	含义	说明与约束
tp_size	张量并行度	本例为 1，即只开流水、不做张量并行
pp_size	流水并行度	本例为 2，对应 2 个流水段；总卡数需满足世界大小 % (tp_size * pp_size) == 0
num_microbatches	microbatch 数量	与microbatch_size二选一；给出后会自动推导microbatch_size = batch_size / num_microbatches，要求能整除
microbatch_size	每个 microbatch 的样本数	给出后会自动推导num_microbatches = batch_size / microbatch_size；训练模式下 1F1B 还要求num_microbatches >= num_stages（one_f_one_b.py#L93-L95）
enable_all_optimization	是否启用 Shardformer 的全部优化（算子融合、层切分等）	一般置True
zero_stage	ZeRO 分片阶段	开启流水后只允许0或1
precision	混合精度策略	如'fp16'；另有'bf16'、'fp8'等，可参考 混合精度训练
initial_scale	FP16 初始动态 loss scale	用于梯度缩放

随后把优化特性注入训练组件。booster.boost返回顺序为(model, optimizer, criterion, dataloader, lr_scheduler)，其中 criterion 用占位符_接收（因为它由 plugin 包装进执行管线）：

model, optimizer, _criterion, _, lr_scheduler = booster.boost(model, optimizer, criterion=_criterion, lr_scheduler=lr_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), desc=f'Epoch [{epoch + 1}/{NUM_EPOCHS}]', disable=not (is_pp_last_stage)) as pbar: # Forward pass for _ in pbar: outputs = booster.execute_pipeline(train_dataloader_iter, model, _criterion, optimizer, return_loss=True) # 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，它通过命令行参数选择 plugin：

python 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-L237）：

plugin = HybridParallelPlugin( tp_size=1, pp_size=2, num_microbatches=None, pp_style="interleaved", num_model_chunks=2, microbatch_size=16, enable_all_optimization=True, zero_stage=1, precision="fp16", initial_scale=1, fp8_communication=False, )

与文档示例相比，这里把microbatch_size加大到 16，并额外指定pp_style="interleaved"与num_model_chunks=2。由于pp_size=2、num_model_chunks=2，调度器内部实际会构建 4 个“虚拟流水级”，每个设备交替处理两个模型块。

脚本还示范了流水并行下的评估写法（finetune.py#L58-L96）：判断isinstance(booster.plugin, HybridParallelPlugin) and booster.plugin.pp_size > 1进入 pipeline 分支，用booster.execute_pipeline(batch, model, criterion, return_loss=True, return_outputs=True)获取outputs["outputs"]["logits"]与outputs["loss"]，并以stage_manager.is_last_stage(ignore_chunk=True)识别末级设备（交错调度下忽略 chunk 维度判断真正的流水末端），最后由末级 rank 通过broadcast_object_list把预测结果广播给同一流水组内其它 rank 参与 GLUE metric 统计。--target_f1参数还可设定达标阈值用于 CI 回归验证。

由于finetune.py内部调用的是colossalai.launch_from_torch(seed=42)（finetune.py#L206），运行时需要通过外部分布式启动器（如torchrun）按tp_size * pp_size的规模拉起进程，并在data.py的GLUEDataBuilder之外保证数据集可访问。

常见问题与排错要点

1. 断言 “World size is not divisible by tp_size * pp_size”：卡数与tp_size * pp_size不匹配，请调整启动进程数。
2. 断言 “Either num_microbatches or microbatch_size should be provided”：开启流水后必须二选一显式指定（hybrid_parallel_plugin.py#L1131-L1133）。
3. 断言 “Number of microbatch should be larger than number of stages”：1F1B 非交错调度下 microbatch 数不能小于流水级数，否则设备无法进入稳态（one_f_one_b.py#L93-L95）。
4. 断言 “Number of microbatch should be an integer multiple of number of pipeline parallel devices”：交错调度要求 microbatch 数是流水级数的整数倍（interleaved_pp.py#L88-L90）。
5. batch 无法整除：num_microbatches/microbatch_size与 batch size 之间须满足整除关系（load_batch中会校验，见 one_f_one_b.py#L82-L91）。
6. loss 只在末级可见：不要在非末级 rank 上读取outputs['loss']，应按stage_manager.is_last_stage()判断后再打印或记录。
7. 显存受限：交错调度相对更吃激活显存，建议搭配梯度检查点；ZeRO 阶段在流水下最多取 1，若需更大优化空间可结合enable_all_optimization（Shardformer 算子级优化）实现。

更完整的并行范式背景可继续阅读并行技术总览，若想进一步探索 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