【大模型】- DualPipe 并行

宋 算法
  2026-08-19 09:19 2026-08-19 09:19 创建   2026-08-21 16:20:11 2026-08-21 16:20:11 更新      1.6k 字  7 分钟

DualPipe 并行

优化流水线并行，减少气泡

类型： 学习　｜　语言： Python　｜　🏷 前置：《多 Token 预测》（本系列第 17 篇）

学习目标

  • 理解流水线并行的原理
  • 实现 DualPipe 并行
  • 减少流水线气泡
  • 优化通信开销
  • 分析 DualPipe 的优势

流水线并行概述

流水线并行将模型分成多个阶段，每个阶段在不同的设备上执行。

标准流水线并行的问题

  • 气泡：设备空闲等待数据
  • 内存峰值：需要存储多个微批次的激活
  • 通信开销：阶段间传输激活

DualPipe 的解决方案

DualPipe 通过双缓冲和重叠计算与通信，减少气泡和内存峰值。

DualPipe 实现

基本架构

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
import torch
import torch.nn as nn
import torch.distributed as dist

class DualPipe:
    def __init__(self, model, n_stages, n_microbatches):
        self.model = model
        self.n_stages = n_stages
        self.n_microbatches = n_microbatches
        
        # 将模型分块
        self.stages = self.split_model()
        
        # 通信缓冲区
        self.forward_buffers = [None] * n_microbatches
        self.backward_buffers = [None] * n_microbatches
        
        # 双缓冲
        self.double_buffer = {
            "forward": [None, None],
            "backward": [None, None]
        }
    
    def split_model(self):
        """将模型分块"""
        layers = list(self.model.children())
        chunk_size = len(layers) // self.n_stages
        
        stages = []
        for i in range(self.n_stages):
            start = i * chunk_size
            end = start + chunk_size if i < self.n_stages - 1 else len(layers)
            stage = nn.Sequential(*layers[start:end])
            stages.append(stage)
        
        return nn.ModuleList(stages)

前向传播

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
class DualPipeForward:
    def __init__(self, pipe, rank, world_size):
        self.pipe = pipe
        self.rank = rank
        self.world_size = world_size
        
        # 前向队列
        self.forward_queue = []
    
    def forward(self, input_ids):
        """前向传播"""
        batch_size = input_ids.shape[0]
        
        # 分割成微批次
        microbatches = torch.chunk(input_ids, self.pipe.n_microbatches)
        
        # 双缓冲前向传播
        outputs = []
        
        for i, microbatch in enumerate(microbatches):
            # 选择缓冲区
            buffer_idx = i % 2
            
            # 如果是当前阶段
            if self.rank < self.pipe.n_stages:
                # 计算
                if self.rank == 0:
                    # 第一个阶段：处理输入
                    output = self.pipe.stages[self.rank](microbatch)
                else:
                    # 接收前一个阶段的输出
                    if self.pipe.double_buffer["forward"][buffer_idx] is not None:
                        microbatch = self.pipe.double_buffer["forward"][buffer_idx]
                    else:
                        microbatch = self.receive_forward(buffer_idx)
                    
                    output = self.pipe.stages[self.rank](microbatch)
                
                # 发送到下一个阶段
                if self.rank < self.pipe.n_stages - 1:
                    self.pipe.double_buffer["forward"][buffer_idx] = output
                
                outputs.append(output)
        
        return torch.cat(outputs, dim=0)
    
    def receive_forward(self, buffer_idx):
        """接收前向数据"""
        if self.rank > 0:
            # 从 rank-1 接收
            recv_tensor = torch.zeros_like(self.pipe.stages[self.rank](torch.zeros(1, 768)))
            dist.recv(recv_tensor, src=self.rank - 1)
            return recv_tensor
        return None

反向传播

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
class DualPipeBackward:
    def __init__(self, pipe, rank, world_size):
        self.pipe = pipe
        self.rank = rank
        self.world_size = world_size
    
    def backward(self, outputs, labels):
        """反向传播"""
        # 反向传播顺序与前向相反
        total_loss = 0
        
        for i in range(len(outputs) - 1, -1, -1):
            output = outputs[i]
            
            # 计算损失
            loss = F.cross_entropy(
                output.view(-1, output.size(-1)),
                labels[i].view(-1)
            )
            
            total_loss += loss
        
        # 反向传播
        total_loss.backward()
        
        return total_loss

优化技术

重叠计算与通信

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
class OverlappedDualPipe:
    def __init__(self, pipe, rank, world_size):
        self.pipe = pipe
        self.rank = rank
        self.world_size = world_size
        
        # 异步通信
        self.send_handles = []
        self.recv_handles = []
    
    def overlapped_forward(self, input_ids):
        """重叠前向传播"""
        microbatches = torch.chunk(input_ids, self.pipe.n_microbatches)
        
        outputs = []
        
        for i, microbatch in enumerate(microbatches):
            buffer_idx = i % 2
            
            # 异步发送（不等待完成）
            if self.rank < self.pipe.n_stages - 1:
                if self.pipe.double_buffer["forward"][buffer_idx] is not None:
                    handle = dist.isend(
                        self.pipe.double_buffer["forward"][buffer_idx],
                        dst=self.rank + 1
                    )
                    self.send_handles.append(handle)
            
            # 异步接收（不等待完成）
            if self.rank > 0:
                recv_tensor = torch.zeros_like(microbatch)
                handle = dist.irecv(recv_tensor, src=self.rank - 1)
                self.recv_handles.append(handle)
            
            # 计算（与通信重叠）
            if self.rank == 0:
                output = self.pipe.stages[self.rank](microbatch)
            else:
                # 等待接收完成
                self.recv_handles[i].wait()
                output = self.pipe.stages[self.rank](recv_tensor)
            
            outputs.append(output)
        
        # 等待所有发送完成
        for handle in self.send_handles:
            handle.wait()
        
        return torch.cat(outputs, dim=0)

内存优化

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
class MemoryEfficientDualPipe:
    def __init__(self, pipe, rank, world_size):
        self.pipe = pipe
        self.rank = rank
        self.world_size = world_size
        
        # 激活检查点
        self.checkpoint_interval = 4
    
    def memory_efficient_forward(self, input_ids):
        """内存优化前向传播"""
        microbatches = torch.chunk(input_ids, self.pipe.n_microbatches)
        
        outputs = []
        
        for i, microbatch in enumerate(microbatches):
            # 检查是否需要检查点
            if i % self.checkpoint_interval == 0:
                # 保存激活
                self.save_activation(microbatch, i)
            
            # 前向传播
            output = self.pipe.stages[self.rank](microbatch)
            outputs.append(output)
            
            # 释放不需要的激活
            if i > 0 and (i - 1) % self.checkpoint_interval == 0:
                self.free_activation(i - 1)
        
        return torch.cat(outputs, dim=0)
    
    def save_activation(self, tensor, index):
        """保存激活"""
        # 使用 CPU 内存
        self.activations[index] = tensor.cpu()
    
    def free_activation(self, index):
        """释放激活"""
        if index in self.activations:
            del self.activations[index]

性能分析

气泡分析

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
def analyze_bubble_time(pipe, n_microbatches, n_stages):
    """分析气泡时间"""
    # 计算气泡时间
    bubble_time = 0
    
    for i in range(n_microbatches):
        for j in range(n_stages):
            # 气泡时间 = 等待时间
            if i < n_stages - 1:
                bubble_time += 1  # 简化表示
    
    total_time = n_microbatches * n_stages
    
    bubble_ratio = bubble_time / total_time
    
    print(f"气泡时间: {bubble_time}")
    print(f"总时间: {total_time}")
    print(f"气泡比例: {bubble_ratio:.2%}")
    
    return bubble_ratio

通信分析

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
def analyze_communication(pipe, n_microbatches, message_size):
    """分析通信开销"""
    # 前向通信
    forward_comm = n_microbatches * (pipe.n_stages - 1) * message_size
    
    # 反向通信
    backward_comm = n_microbatches * (pipe.n_stages - 1) * message_size
    
    total_comm = forward_comm + backward_comm
    
    print(f"前向通信: {forward_comm / 1e9:.2f}GB")
    print(f"反向通信: {backward_comm / 1e9:.2f}GB")
    print(f"总通信: {total_comm / 1e9:.2f}GB")
    
    return total_comm

与传统流水线比较

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
def compare_with_traditional(pipe_dual, pipe_traditional, n_microbatches):
    """与传统流水线比较"""
    
    # 双缓冲流水线
    time_dual = measure_time(pipe_dual, n_microbatches)
    memory_dual = measure_memory(pipe_dual)
    
    # 传统流水线
    time_traditional = measure_time(pipe_traditional, n_microbatches)
    memory_traditional = measure_memory(pipe_traditional)
    
    print(f"双缓冲流水线:")
    print(f"  时间: {time_dual:.4f}s")
    print(f"  内存: {memory_dual:.2f}MB")
    
    print(f"传统流水线:")
    print(f"  时间: {time_traditional:.4f}s")
    print(f"  内存: {memory_traditional:.2f}MB")
    
    print(f"加速比: {time_traditional / time_dual:.2f}x")
    print(f"内存减少: {(memory_traditional - memory_dual) / memory_traditional:.2%}")

配置优化

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
def optimize_config(model_size, n_gpus, batch_size):
    """优化配置参数"""
    
    # 计算合适的阶段数
    if n_gpus >= 8:
        n_stages = 8
    elif n_gpus >= 4:
        n_stages = 4
    else:
        n_stages = 2
    
    # 计算合适的微批次数
    if batch_size >= 32:
        n_microbatches = 8
    elif batch_size >= 16:
        n_microbatches = 4
    else:
        n_microbatches = 2
    
    # 计算 chunk 大小
    chunk_size = model_size // n_stages
    
    print(f"配置:")
    print(f"  阶段数: {n_stages}")
    print(f"  微批次数: {n_microbatches}")
    print(f"  Chunk 大小: {chunk_size / 1e6:.2f}M 参数")
    
    return {
        "n_stages": n_stages,
        "n_microbatches": n_microbatches,
        "chunk_size": chunk_size
    }

实际应用

大模型训练

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
def train_large_model(model, dataset, config):
    """使用 DualPipe 训练大模型"""
    
    # 初始化分布式训练
    dist.init_process_group("nccl")
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    
    # 创建 DualPipe
    pipe = DualPipe(model, world_size, config["n_microbatches"])
    
    # 训练循环
    for epoch in range(config["epochs"]):
        for batch in dataset:
            # 前向传播
            outputs = pipe.forward(batch["input_ids"])
            
            # 计算损失
            loss = F.cross_entropy(outputs, batch["labels"])
            
            # 反向传播
            loss.backward()
            
            # 更新参数
            optimizer.step()
            optimizer.zero_grad()

多节点训练

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
class MultiNodeDualPipe:
    def __init__(self, model, n_nodes, n_gpus_per_node):
        self.model = model
        self.n_nodes = n_nodes
        self.n_gpus_per_node = n_gpus_per_node
        
        # 节点内通信
        self.intra_pipe = DualPipe(model, n_gpus_per_node)
        
        # 节点间通信
        self.inter_comm = InterNodeComm(n_nodes)
    
    def train(self, input_ids):
        """多节点训练"""
        # 节点内流水线
        outputs = self.intra_pipe.forward(input_ids)
        
        # 节点间同步
        self.inter_comm.sync(outputs)
        
        return outputs

总结

DualPipe 通过双缓冲和重叠计算与通信，显著减少了流水线气泡和内存峰值。关键优化包括异步通信、内存优化和配置调整。DualPipe 特别适合训练大型语言模型。

下一步

下一课将深入讲解 DeepSeek-V3 的架构。

📚 本文改编自 AI Engineering from Scratch（MIT License · 作者 Rohit Ghumare），中文内容来自官方中文镜像。原课程共 503 课 · 20 阶段 · 免费开源，教程网站见 aiengineeringfromscratch.com。

  • 标题: 【大模型】- DualPipe 并行
  • 作者: 宋
  • 创建于 : 2026-08-19 09:19:00
  • 更新于 : 2026-08-21 16:20:11
  • 链接: https://sxl-space.tk/2026/08/19/010_LLM/010_LLM-19-DualPipe/
  • 版权声明: 版权所有 © 宋，禁止转载。
目录
【大模型】- DualPipe 并行
  1. 1. DualPipe 并行
  2. 2. 学习目标
  3. 3. 流水线并行概述
    1. 3.1. 标准流水线并行的问题
    2. 3.2. DualPipe 的解决方案
  4. 4. DualPipe 实现
    1. 4.1. 基本架构
    2. 4.2. 前向传播
    3. 4.3. 反向传播
  5. 5. 优化技术
    1. 5.1. 重叠计算与通信
    2. 5.2. 内存优化
  6. 6. 性能分析
    1. 6.1. 气泡分析
    2. 6.2. 通信分析
  7. 7. 与传统流水线比较
  8. 8. 配置优化
  9. 9. 实际应用
    1. 9.1. 大模型训练
    2. 9.2. 多节点训练
  10. 10. 总结
  11. 11. 下一步
  12. 12. 下一课将深入讲解 DeepSeek-V3 的架构。