5.2 流水线并行与微批次调度

摘要

随着大模型参数规模持续增长，单GPU计算能力成为性能瓶颈。流水线并行技术通过将模型分割为多个阶段并分布在多个GPU上，实现了计算资源的有效利用。结合微批次调度策略，可以显著提升推理吞吐量并优化GPU利用率。本文档深入探讨流水线并行的架构设计、通信优化、微批次调度算法以及容错机制，为大规模模型推理提供完整的并行化解决方案。

1. 流水线并行基础架构

1.1 模型分割与阶段划分

动态流水线平衡策略

import torch
import torch.nn as nn
import torch.distributed as dist
import math
import numpy as np
from typing import List, Dict, Tuple, Optional, Union
from dataclasses import dataclass
import asyncio
import time
from enum import Enum, auto

class PipelineStage:
    """流水线阶段类"""
    
    def __init__(self, 
                 stage_id: int,
                 layers: List[nn.Module],
                 device: torch.device):
        self.stage_id = stage_id
        self.layers = nn.ModuleList(layers).to(device)
        self.device = device
        self.input_buffer = None
        self.output_buffer = None
        self.forward_count = 0
        self.backward_count = 0
        self.computation_time = 0.0
        self.communication_time = 0.0
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """前向传播"""
        start_time = time.time()
        
        with torch.no_grad():
            for layer in self.layers:
                x = layer(x)
        
        self.computation_time += time.time() - start_time
        self.forward_count += 1
        
        return x
    
    def get_statistics(self) -> Dict:
        """获取阶段统计信息"""
        return {
            'stage_id': self.stage_id,
            'num_layers': len(self.layers),
            'computation_time': self.computation_time,
            'communication_time': self.communication_time,
            'forward_count': self.forward_count,
            'backward_count': self.backward_count,
            'avg_computation_time': self.computation_time / max(self.forward_count, 1),
            'avg_communication_time': self.communication_time / max(self.forward_count, 1)
        }

class PipelineParallelModel(nn.Module):
    """流水线并行模型"""
    
    def __init__(self, 
                 layers: List[nn.Module],
                 num_stages: int,
                 devices: List[torch.device]):
        super().__init__()
        self.num_stages = num_stages
        self.devices = devices
        self.stages = nn.ModuleList()
        
        # 平衡分配层到各个阶段
        stage_assignments = self._balance_layers(layers, num_stages)
        
        # 创建流水线阶段
        for i in range(num_stages):
            stage_layers = stage_assignments[i]
            device = devices[i]
            stage = PipelineStage(i, stage_layers, device)
            self.stages.append(stage)
    
    def _balance_layers(self, layers: List[nn.Module], num_stages: int) -> List[List[nn.Module]]:
        """平衡分配层到各个阶段"""
        # 计算每个层的计算量（简化估计）
        layer_complexities = []
        for layer in layers:
            # 简化的复杂度估计
            if isinstance(layer, nn.Linear):
                input_size = layer.weight.shape[1]
                output_size = layer.weight.shape[0]
                complexity = input_size * output_size
            elif isinstance(layer, nn.MultiheadAttention):
                d_model = layer.embed_dim
                complexity = d_model * d_model * 4  # 注意力机制复杂度
            else:
                complexity = 1000  # 默认复杂度
            
            layer_complexities.append(complexity)
        
        total_complexity = sum(layer_complexities)
        target_per_stage = total_complexity / num_stages
        
        # 动态分配算法
        assignments = [[] for _ in range(num_stages)]
        current_stage = 0
        current_sum = 0
        
        for i, layer in enumerate(layers):
            current_sum += layer_complexities[i]
            
            # 检查是否需要切换到下一个阶段
            if current_sum >= target_per_stage and current_stage < num_stages - 1:
                current_stage += 1
                current_sum = 0
            
            assignments[current_stage].append(layer)
        
        return assignments
    
    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        """流水线并行前向传播"""
        outputs = []
        current_output = input_ids
        
        # 简单的顺序执行（实际应该使用异步流水线）
        for stage in self.stages:
            current_output = current_output.to(stage.device)
            current_output = stage(current_output)
            outputs.append(current_output)
        
        return outputs[-1]  # 返回最后一个阶段的输出
    
    def get_pipeline_statistics(self) -> Dict:
        """获取流水线统计信息"""
        stats = {
            'num_stages': self.num_stages,
            'devices': [str(device) for device in self.devices],
            'stages': [stage.get_statistics() for stage in self.stages]
        }
        
        # 计算整体性能指标
        total_computation = sum(stat['computation_time'] for stat in stats['stages'])
        total_communication = sum(stat['communication_time'] for stat in stats['stages'])
        
        stats['total_computation_time'] = total_computation
        stats['total_communication_time'] = total_communication
        stats['total_time'] = total_computation + total_communication
        stats['communication_ratio'] = total_communication / total_computation if total_computation > 0 else 0
        
        return stats

# 演示流水线并行模型
def demonstrate_pipeline_parallel():
    """演示流水线并行模型"""
    print("流水线并行模型演示")
    print("-" * 50)
    
    # 创建测试模型层
    layers = []
    d_model = 512
    for i in range(12):  # 12层
        if i % 4 == 0:  # 每4层加入一个注意力层
            layers.append(nn.MultiheadAttention(d_model, 8))
        else:
            layers.append(nn.Linear(d_model, d_model))
    
    # 设备配置
    num_devices = 4
    devices = [torch.device(f'cuda:{i}') for i in range(num_devices)]
    
    # 创建流水线并行模型
    pipeline_model = PipelineParallelModel(layers, num_devices, devices)
    
    # 创建输入
    batch_size = 8
    seq_len = 64
    input_ids = torch.randint(0, 10000, (batch_size, seq_len))
    
    # 前向传播
    output = pipeline_model(input_ids)
    
    print(f"输入形状: {input_ids.shape}")
    print(f"输出形状: {output.shape}")
    print(f"流水线阶段数: {pipeline_model.num_stages}")
    
    # 统计信息
    stats = pipeline_model.get_pipeline_statistics()
    print("\n流水线统计信息:")
    print(f"总计算时间: {stats['total_computation_time']:.4f}s")
    print(f"总通信时间: {stats['total_communication_time']:.4f}s")
    print(f"通信占比: {stats['communication_ratio']:.2%}")
    
    for stage_stat in stats['stages']:
        print(f"阶段 {stage_stat['stage_id']}: "
              f"计算 {stage_stat['computation_time']:.4f}s, "
              f"通信 {stage_stat['communication_time']:.4f}s")

# 运行演示
demonstrate_pipeline_parallel()
​

智能层分配算法

class LayerBalancer:
    """智能层分配器"""
    
    def __init__(self, layers: List[nn.Module], num_stages: int):
        self.layers = layers
        self.num_stages = num_stages
        self.layer_complexities = self._calculate_layer_complexities()
        self.stage_assignments = None
    
    def _calculate_layer_complexities(self) -> List[float]:
        """计算各层复杂度"""
        complexities = []
        
        for layer in self.layers:
            # 基于层数类型和参数数量估算复杂度
            if isinstance(layer, nn.Linear):
                input_size = layer.weight.shape[1]
                output_size = layer.weight.shape[0]
                complexity = input_size * output_size * 2  # 矩阵乘法复杂度
            elif isinstance(layer, nn.MultiheadAttention):
                d_model = layer.embed_dim
                nhead = layer.num_heads
                complexity = d_model * d_model * nhead * 4  # 注意力复杂度
            elif isinstance(layer, nn.LayerNorm):
                complexity = layer.normalized_shape[0] * 2  # 归一化复杂度
            elif isinstance(layer, nn.GELU):
                complexity = sum(layer.normalized_shape)  # 激活函数复杂度
            else:
                complexity = 1000  # 默认复杂度
            
            complexities.append(complexity)
        
        return complexities
    
    def balance_layers(self, balance_strategy: str = 'dynamic') -> List[List[nn.Module]]:
        """平衡分配层到各个阶段"""
        if balance_strategy == 'dynamic':
            return self._dynamic_balance()
        elif balance_strategy == 'greedy':
            return self._greedy_balance()
        elif balance_strategy == 'optimal':
            return self._optimal_balance()
        else:
            raise ValueError(f"Unknown balance strategy: {balance_strategy}")
    
    def _dynamic_balance(self) -> List[List[nn.Module]]:
        """动态平衡算法"""
        assignments = [[] for _ in range(self.num_stages)]
        stage_loads = [0.0] * self.num_stages
        
        for i, (layer, complexity) in enumerate(zip(self.layers, self.layer_complexities)):
            # 找到当前负载最小的阶段
            min_load_stage = np.argmin(stage_loads)
            
            # 将层分配到最小负载的阶段
            assignments[min_load_stage].append(layer)
            stage_loads[min_load_stage] += complexity
            
            # 如果负载不均衡，重新平衡
            if self._load_imbalance(stage_loads) > 0.3:  # 30%的不平衡阈值
                self._rebalance_stages(assignments, stage_loads)
        
        return assignments
    
    def _greedy_balance(self) -> List[List[nn.Module]]:
        """贪心平衡算法"""
        assignments = [[] for _ in range(self.num_stages)]
        stage_loads = [0.0] * self.num_stages
        
        # 按复杂度排序层（从大到小）
        layer_indices = sorted(range(len(self.layers)), 
                              key=lambda i: self.layer_complexities[i], 
                              reverse=True)
        
        for i in layer_indices:
            layer = self.layers[i]
            complexity = self.layer_complexities[i]
            
            # 找到可以容纳该层的阶段
            candidates = []
            for stage_idx in range(self.num_stages):
                if stage_loads[stage_idx] + complexity <= self._get_stage_capacity(stage_idx):
                    candidates.append((stage_idx, stage_loads[stage_idx]))
            
            if candidates:
                # 选择负载最小的候选阶段
                min_load_stage = min(candidates, key=lambda x: x[1])[0]
                assignments[min_load_stage].append(layer)
                stage_loads[min_load_stage] += complexity
            else:
                # 如果没有合适的阶段，分配到负载最小的阶段
                min_load_stage = np.argmin(stage_loads)
                assignments[min_load_stage].append(layer)
                stage_loads[min_load_stage] += complexity
        
        return assignments
    
    def _optimal_balance(self) -> List[List[nn.Module]]:
        """最优平衡算法（简化版）"""
        # 使用动态规划寻找最优分配
        total_complexity = sum(self.layer_complexities)
        target_per_stage = total_complexity / self.num_stages
        
        assignments = [[] for _ in range(self.num_stages)]
        used = [False] * len(self.layers)
        
        # 迭代分配
        for stage_idx in range(self.num_stages):
            current_load = 0.0
            best_assignment = []
            
            # 尝试找到最优的层组合
            for layer_idx in range(len(self.layers)):
                if not used[layer_idx]:
                    test_assignment = best_assignment + [self.layers[layer_idx]]
                    test_load = current_load + self.layer_complexities[layer_idx]
                    
                    # 如果不超过目标负载，考虑加入
                    if test_load <= target_per_stage * 1.2:  # 20%的容差
                        best_assignment = test_assignment
                        current_load = test_load
            
            # 分配找到的层
            for layer in best_assignment:
                layer_idx = self.layers.index(layer)
                assignments[stage_idx].append(layer)
                used[layer_idx] = True
        
        return assignments
    
    def _load_imbalance(self, stage_loads: List[float]) -> float:
        """计算负载不平衡度"""
        max_load = max(stage_loads)
        min_load = min(stage_loads)
        avg_load = sum(stage_loads) / len(stage_loads)
        
        # 使用标准差作为不平衡度量
        std_load = np.std(stage_loads)
        imbalance = std_load / avg_load if avg_load > 0 else 0
        
        return imbalance
    
    def _get_stage_capacity(self, stage_idx: int) -> float:
        """获取阶段容量"""
        # 简化处理，实际应该基于GPU内存容量
        return float('inf')
    
    def _rebalance_stages(self, assignments: List[List[nn.Module]], stage_loads: List[float]):
        """重新平衡阶段负载"""
        # 简化的重新平衡算法
        # 实际实现可能需要更复杂的算法
        pass
    
    def get_assignment_statistics(self, assignments: List[List[nn.Module]]) -> Dict:
        """获取分配统计信息"""
        stage_loads = []
        stage_layer_counts = []
        
        for stage_idx, stage_layers in enumerate(assignments):
            load = sum(self.layer_complexities[i] for i, layer in enumerate(self.layers) 
                      if layer in stage_layers)
            layer_count = len(stage_layers)
            
            stage_loads.append(load)
            stage_layer_counts.append(layer_count)
        
        avg_load = sum(stage_loads) / len(stage_loads)
        max_load = max(stage_loads)
        min_load = min(stage_loads)
        imbalance = self._load_imbalance(stage_loads)
        
        return {
            'stage_loads': stage_loads,
            'stage_layer_counts': stage_layer_counts,
            'avg_load': avg_load,
            'max_load': max_load,
            'min_load': min_load,
            'load_imbalance': imbalance,
            'max_layer_count': max(stage_layer_counts),
            'min_layer_count': min(stage_layer_counts)
        }

# 演示智能层分配
def demonstrate_layer_balancing():
    """演示智能层分配算法"""
    print("\n智能层分配算法演示")
    print("-" * 50)
    
    # 创建测试模型层
    layers = []
    d_model = 512
    
    # 混合不同类型的层
    for i in range(24):
        if i % 8 == 0:
            layers.append(nn.MultiheadAttention(d_model, 8))
        elif i % 8 == 4:
            layers.append(nn.Linear(d_model, d_model))
        else:
            layers.append(nn.LayerNorm(d_model))
    
    # 测试不同的平衡策略
    num_stages = 4
    strategies = ['dynamic', 'greedy', 'optimal']
    
    for strategy in strategies:
        print(f"\n{strategy.upper()} 平衡策略:")
        balancer = LayerBalancer(layers, num_stages)
        assignments = balancer.balance_layers(strategy)
        stats = balancer.get_assignment_statistics(assignments)
        
        print(f"阶段负载: {[f'{load:.0f}' for load in stats['stage_loads']]}")
        print(f"层数量: {[count for count in stats['stage_layer_counts']]}")
        print(f"负载不平衡度: {stats['load_imbalance']:.3f}")
        print(f"最大层数量: {stats['max_layer_count']}")
        print(f"最小层数量: {stats['min_layer_count']}")

# 运行演示
demonstrate_layer_balancing()
​

1.2 异步流水线通信

高效的通信优化策略

class AsyncPipeline:
    """异步流水线实现"""
    
    def __init__(self, 
                 pipeline_model: PipelineParallelModel,
                 microbatch_size: int,
                 overlap_communication: bool = True):
        self.pipeline_model = pipeline_model
        self.microbatch_size = microbatch_size
        self.overlap_communication = overlap_communication
        self.stages = pipeline_model.stages
        self.num_stages = len(stages)
​