分布式训练深度解析：Data / Tensor / Pipeline / Sequence Parallelism

1
┌─────────────────────────────────────────────────────────────┐
2
│               LLM 参数规模增长趋势                             │
3
├─────────────────────────────────────────────────────────────┤
4
│                                                              │
5
│  2018: ELMo          ~ 94M 参数                             │
6
│  2018: BERT-Large    ~ 340M 参数                            │
7
│  2019: GPT-2         ~ 1.5B 参数                           │
8
│  2020: GPT-3         ~ 175B 参数                           │
9
│  2021: PaLM          ~ 540B 参数                           │
10
│  2023: GPT-4        ~ 1.8T 参数 (估计)                     │
11
│  2024: Gemini Ultra ~ 1.5T 参数                            │
12
│                                                              │
13
│  内存需求（FP16）：                                          │
14
│  ─────────────────────────────────────────────────────────  │
15
│  1B  参数  →  2 GB                                         │
16
│  7B  参数  →  14 GB                                        │
17
│  70B 参数  →  140 GB                                       │
18
│  175B 参数 →  350 GB                                       │
19
│  1T  参数  →  2 TB                                         │
20
│                                                              │
21
│  单卡 A100 80GB：仅能放下 ~35B 参数的模型                    │
22
│                                                              │
23
└─────────────────────────────────────────────────────────────┘

1
class ComputeRequirements:
2
    """
3
    计算需求分析
4
    """
5

6
    def flops_estimation(self):
7
        """
8
        训练 LLM 的 FLOPs 估算
9

10
        对于 decoder-only transformer:
11
        FLOPs ≈ 2 × N × D × G
12

13
        其中:
14
        - N: 模型参数量
15
        - D: 训练 tokens 数
16
        - G: 训练轮数
17
        """
18
        return {
19
            "gpt3_175b": {
20
                "params": "175B",
21
                "tokens": "300B",
22
                "flops": "~3.1 × 10^23 FLOPs",
23
                "a100_petaflops": "~3.1 × 10^5 PF",
24
                "a100_8x_90_days": "~1.5 × 10^6 PF-days",
25
            },
26
            "llama2_70b": {
27
                "params": "70B",
28
                "tokens": "2T (预训练)",
29
                "flops": "~1.4 × 10^24 FLOPs",
30
            },
31
        }
32

33
    def time_estimation(self):
34
        """
35
        训练时间估算
36
        """
37
        return """
38
        单卡 A100 (312 TFLOPS BF16):
39

40
        LLaMA-2 70B, 2T tokens:
41

42
        Theoretical time = FLOPs / FLOPS
43
                           = 1.4 × 10^24 / 3.12 × 10^14
44
                           = 4.5 × 10^9 seconds
45
                           = ~144 years
46

47
        结论: 单卡训练不现实
48

49
        需要的 GPU 数量:
50

51
        1000 GPUs → ~53 天
52
        4000 GPUs → ~13 天
53
        """

1
┌─────────────────────────────────────────────────────────────┐
2
│                   分布式训练并行策略                           │
3
├─────────────────────────────────────────────────────────────┤
4
│                                                              │
5
│  Data Parallelism (DP)                                       │
6
│  ─────────────────────────────────────────────────────────  │
7
│  → 复制模型，并行处理不同数据分片                             │
8
│  → 最简单，提升吞吐量                                         │
9
│  → 通信: AllReduce (梯度)                                    │
10
│                                                              │
11
│  Tensor Parallelism (TP)                                     │
12
│  ─────────────────────────────────────────────────────────  │
13
│  → 切分单个层的权重到多卡                                     │
14
│  → 用于超大型单层 (如 LLM 的 FFN/Attention)                  │
15
│  → 通信: AllReduce / AllGather                              │
16
│                                                              │
17
│  Pipeline Parallelism (PP)                                   │
18
│  ─────────────────────────────────────────────────────────  │
19
│  → 按层切分模型到多卡                                         │
20
│  → 减少单卡显存需求                                          │
21
│  → 通信: P2P (点对点)                                        │
22
│                                                              │
23
│  Sequence Parallelism (SP)                                   │
24
│  ─────────────────────────────────────────────────────────  │
25
│  → 按序列维度切分                                             │
26
│  → 配合 TP 处理超长上下文                                     │
27
│  → 通信: AllReduce                                          │
28
│                                                              │
29
│  组合: 2D / 3D Parallelism                                   │
30
│  → TP × PP × DP                                             │
31
│  → 千卡级别训练的必备                                       │
32
│                                                              │
33
└─────────────────────────────────────────────────────────────┘

1
class DataParallelism:
2
    """
3
    Data Parallelism
4
    """
5

6
    def basic_idea(self):
7
        """
8
        核心思想
9

10
        将训练数据分成多个 shard，每个 GPU 持有完整模型副本，
11
        并行处理不同数据，最后同步梯度。
12
        """
13
        return """
14
        Data Parallelism 示意:
15

16
        GPU 0: [Model Copy] → Batch 0 → Forward → Backward → Grad 0
17
        GPU 1: [Model Copy] → Batch 1 → Forward → Backward → Grad 1
18
        GPU 2: [Model Copy] → Batch 2 → Forward → Backward → Grad 2
19
        GPU 3: [Model Copy] → Batch 3 → Forward → Backward → Grad 3
20

21
        ↓ Gradient AllReduce ↓
22

23
        Avg(Grad 0, Grad 1, Grad 2, Grad 3)
24

25
        ↓ 每个 GPU 用平均梯度更新 ↓
26

27
        特点:
28
        → 每个 GPU 有完整模型
29
        → 只通信梯度
30
        → 增加 GPU 数量可线性提升吞吐量
31
        """
32

33
    def gradient_synchronization(self):
34
        """
35
        梯度同步
36
        """
37
        return '''
38
        class DataParallelGradientSync:
39
            """梯度同步实现"""
40

41
            def allreduce_sync(self, gradients, world_size):
42
                """
43
                同步 SGD (所有 GPU 同步等待)
44

45
                每个 GPU 计算梯度
46
                → AllReduce 同步梯度
47
                → 所有 GPU 都得到相同的平均梯度
48
                → 各自更新
49

50
                优点: 收敛行为与单卡相同
51
                缺点: 最慢的 GPU 拖慢整体
52
                """
53
                import torch.distributed as dist
54

55
                # 对每个参数梯度执行 AllReduce
56
                for param in model.parameters():
57
                    dist.all_reduce(param.grad, op=dist.ReduceOp.SUM)
58
                    param.grad /= world_size
59

60
            def distributed_sampler(self, dataset, num_gpus):
61
                """数据分片"""
62
                return torch.utils.data.DistributedSampler(
63
                    dataset,
64
                    num_replicas=num_gpus,
65
                    rank=self.rank,
66
                )
67
        '''
68

69
    def memory_analysis(self):
70
        """
71
        显存分析
72
        """
73
        return {
74
            "model_weights": "N GB (完整副本)",
75
            "optimizer_states": "2N GB (Adam: m + v, FP32)",
76
            "gradients": "2N GB (FP32)",
77
            "activations": "取决于 batch size",
78
            "total_per_gpu": "~6N GB + activations",
79

80
            "example_7b_fp16": {
81
                "model": "14 GB",
82
                "optimizer": "28 GB",
83
                "gradients": "28 GB",
84
                "activations_8": "16 GB (batch=8)",
85
                "total": "~86 GB per GPU",
86
                "a100_80gb": "不够，需要进一步优化",
87
            },
88
        }

1
class DDPImplementation:
2
    """
3
    PyTorch DDP 实现
4
    """
5

6
    def basic_usage(self):
7
        """
8
        基本用法
9
        """
10
        return '''
11
        import torch
12
        import torch.distributed as dist
13
        import torch.multiprocessing as mp
14
        from torch.nn.parallel import DistributedDataParallel as DDP
15

16
        def setup(rank, world_size):
17
            os.environ["MASTER_ADDR"] = "localhost"
18
            os.environ["MASTER_PORT"] = "12355"
19
            dist.init_process_group("nccl", rank=rank, world_size=world_size)
20
            torch.cuda.set_device(rank)
21

22
        def train(rank, world_size):
23
            setup(rank, world_size)
24

25
            # 创建模型
26
            model = MyModel().cuda()
27

28
            # 包装为 DDP
29
            model = DDP(model, device_ids=[rank])
30

31
            # 数据
32
            dataset = MyDataset()
33
            sampler = DistributedSampler(dataset)
34
            dataloader = DataLoader(dataset, sampler=sampler)
35

36
            # 训练循环
37
            for epoch in range(num_epochs):
38
                for batch in dataloader:
39
                    loss = model(batch).sum()
40
                    loss.backward()
41
                    optimizer.step()
42
                    optimizer.zero_grad()
43

44
            dist.destroy_process_group()
45

46
        # 启动
47
        mp.spawn(train, args=(world_size,), nprocs=world_size)
48
        '''
49

50
    def gradient_bucket(self):
51
        """
52
        Gradient Bucketing 优化
53

54
        DDP 使用 bucket 机制减少通信开销
55
        """
56
        return """
57
        问题:
58
        → 逐参数同步梯度效率低
59
        → 每个小梯度都需要完整的 AllReduce 通信
60

61
        解决: Gradient Bucketing
62

63
        ┌─────────────────────────────────────────────────────┐
64
        │  Gradients                                          │
65
        │  [g1][g2][g3][g4][g5][g6][g7][g8]...             │
66
        └─────────────────────────────────────────────────────┘
67
                      ↓ 分桶
68
        ┌─────────────────────────────────────────────────────┐
69
        │  Bucket 0: [g1][g2][g3][g4]  (size = 4 params)    │
70
        │  Bucket 1: [g5][g6][g7][g8]  (size = 4 params)    │
71
        │  ...                                                │
72
        └─────────────────────────────────────────────────────┘
73
                      ↓ 异步 AllReduce
74
        ┌─────────────────────────────────────────────────────┐
75
        │  Bucket 0: AllReduce 完成                          │
76
        │  Bucket 1: AllReduce 进行中...                      │
77
        │  ...                                                │
78
        └─────────────────────────────────────────────────────┘
79

80
        效果:
81
        → 减少通信次数 (从 N 次到 N/bucket_size 次)
82
        → 重叠计算和通信
83
        → 带宽利用率更高
84
        """
85

86
    def broadcast_buffers(self):
87
        """
88
        Buffer 广播
89
        """
90
        return """
91
        DDP 的 Buffer 处理:
92

93
        1. 同步 Buffers
94
           - BatchNorm 的 running_mean/std
95
           - 使用 AllReduce 同步
96
           - 每个 GPU 广播自己的 buffers
97

98
        2. 非同步 Buffers
99
           - Dropout 的随机状态
100
           - 不需要同步
101

102
        3. 注册 Hook
103
           - 可以在梯度同步时执行自定义操作
104
           - 用途: gradient clipping, monitoring
105
        """

1
class ZeROOptimization:
2
    """
3
    ZeRO (Zero Redundancy Optimizer)
4

5
    切分优化器状态以减少显存冗余
6
    """
7

8
    def stages(self):
9
        """
10
        ZeRO 三个阶段
11
        """
12
        return {
13
            "ZeRO-1": {
14
                "name": "分片优化器状态",
15
                "description": "将 Adam optimizer states (m, v) 分片到各 GPU",
16
                "memory_reduction": "4x (75% reduction)",
17
                "communication": "与 DDP 相同",
18
            },
19
            "ZeRO-2": {
20
                "name": "分片优化器状态 + 梯度",
21
                "description": "ZeRO-1 + 梯度也分片",
22
                "memory_reduction": "8x (~87.5% reduction)",
23
                "communication": "略增 (allreduce 变成reduce-scatter + allgather)",
24
            },
25
            "ZeRO-3": {
26
                "name": "分片所有状态",
27
                "description": "ZeRO-2 + 模型参数也分片",
28
                "memory_reduction": "线性减少 (N 倍)",
29
                "communication": "显著增加 (需要动态 gather 参数)",
30
            },
31
        }
32

33
    def zero_stage1(self):
34
        """
35
        ZeRO-1 实现
36
        """
37
        return '''
38
        # ZeRO-1: 只分片优化器状态
39

40
        # 优化器状态 (Adam):
41
        # m (first moment) - FP32
42
        # v (second moment) - FP32
43

44
        # 每个参数 p 的优化器状态:
45
        # p_fp32, m, v = 4 + 4 + 4 = 12 bytes per param
46

47
        # 原来 (DDP): 每个 GPU 持有完整的 m, v
48
        # ZeRO-1: 每个 GPU 只持有 1/N 的 m, v
49

50
        class ZeRO1Optimizer:
51
            def __init__(self, model, world_size, rank):
52
                self.world_size = world_size
53
                self.rank = rank
54
                self.param_partitions = {}  # rank → params
55

56
                # 将参数均匀分配到各 GPU
57
                all_params = list(model.parameters())
58
                for i, param in enumerate(all_params):
59
                    target_rank = i % world_size
60
                    self.param_partitions[target_rank].append(param)
61

62
            def step(self):
63
                # 只更新本地 partition 的参数
64
                for param in self.param_partitions[self.rank]:
65
                    # 更新 m, v (只存储本地部分)
66
                    self._update_adam_state(param)
67
        '''
68

69
    def zero_vs_ddp_memory(self):
70
        """
71
        显存对比
72
        """
73
        return """
74
        70B 参数模型，FP16 权重，Adam 优化器:
75

76
        DDP (完全副本):
77
        ┌──────────────────────────────────────────────────┐
78
        │  权重 (FP16):        140 GB                      │
79
        │  梯度 (FP32):        280 GB                      │
80
        │  优化器状态 (FP32):  280 GB                      │
81
        │  激活值:              ~50 GB                      │
82
        │  ─────────────────────────────────────────────  │
83
        │  总计:               ~750 GB                     │
84
        │  需要: A100 80GB × 10                          │
85
        └──────────────────────────────────────────────────┘
86

87
        ZeRO-3 + 8 GPUs:
88
        ┌──────────────────────────────────────────────────┐
89
        │  权重分片:          140 / 8 = 17.5 GB          │
90
        │  梯度分片:          280 / 8 = 35 GB            │
91
        │  优化器状态分片:    280 / 8 = 35 GB            │
92
        │  激活值:              ~50 GB                    │
93
        │  ─────────────────────────────────────────────  │
94
        │  总计:               ~137.5 GB per GPU          │
95
        │  需要: A100 80GB × 8 + Offload                  │
96
        └──────────────────────────────────────────────────┘
97
        """

1
class TensorParallelism:
2
    """
3
    Tensor Parallelism
4

5
    将单层权重矩阵沿某个维度切分到多个 GPU
6
    """
7

8
    def why_tensor_parallel(self):
9
        """
10
        为什么需要 Tensor Parallelism
11

12
        对于 LLaMA-2 70B:
13
        """
14
        return {
15
            "embedding": "4096 × 32000 = 131M",
16
            "attention": {
17
                "q_proj": "4096 × 4096 = 16.8M × 3 = 50.4M",
18
                "k_proj": "4096 × 4096",
19
                "v_proj": "4096 × 4096",
20
                "o_proj": "4096 × 4096 = 16.8M",
21
                "total_attention": "67.2M per layer × 80 layers",
22
            },
23
            "ffn": {
24
                "gate_proj": "11008 × 4096 = 45M",
25
                "up_proj": "4096 × 11008 = 45M",
26
                "down_proj": "11008 × 4096 = 45M",
27
                "total_ffn": "135M per layer × 80 layers",
28
            },
29

30
            "largest_layer": "FFN up_proj: 45M params = 90 GB FP16",
31
            "problem": "单层 FFN 比单卡显存还大！",
32
            "solution": "切分到多卡",
33
        }
34

35
    def matrix_partitioning(self):
36
        """
37
        矩阵分块策略
38
        """
39
        return """
40
        矩阵乘法 Y = XW
41

42
        W: (M, N)
43
        X: (B, M)
44
        Y: (B, N)
45

46
        ┌─────────────────────────────────────────────────┐
47
        │              W (M×N)                             │
48
        │  ┌─────────┐ ┌─────────┐ ┌─────────┐            │
49
        │  │   W0    │ │   W1    │ │   W2    │            │
50
        │  │(M, N/3) │ │(M, N/3) │ │(M, N/3) │            │
51
        │  └─────────┘ └─────────┘ └─────────┘            │
52
        │     GPU 0      GPU 1      GPU 2                │
53
        └─────────────────────────────────────────────────┘
54

55
        X: (B, M) 复制到所有 GPU
56

57
        Y0 = X @ W0 (B, N/3) → GPU 0
58
        Y1 = X @ W1 (B, N/3) → GPU 1
59
        Y2 = X @ W2 (B, N/3) → GPU 2
60

61
        输出拼接: Y = [Y0, Y1, Y2]
62
        """

1
class ParallelismTypes:
2
    """
3
    列并行与行并行
4
    """
5

6
    def column_parallel(self):
7
        """
8
        Column Parallel
9

10
        切分权重矩阵的列维度
11
        """
12
        return """
13
        用于 Linear 层的 QKV 投影:
14

15
        W_qkv: (d_model, 3 * d_hidden)
16

17
        Column Parallel (切分 3*d_hidden):
18

19
        GPU 0: W_q0 (d_model, d_hidden)
20
        GPU 1: W_k0 (d_model, d_hidden)
21
        GPU 2: W_v0 (d_model, d_hidden)
22
        ...
23

24
        前向:
25
        X: (B, d_model) 复制
26
        ↓
27
        X @ W_q0 = Q0 (B, d_hidden)
28
        X @ W_k0 = K0 (B, d_hidden)
29
        X @ W_v0 = V0 (B, d_hidden)
30

31
        通信: 无 (输出已经是分片的)
32

33
        注意: 需要 AllGather 合并 Q/K/V 以进行 Attention
34
        """
35

36
    def row_parallel(self):
37
        """
38
        Row Parallel
39

40
        切分权重矩阵的行维度
41
        """
42
        return """
43
        用于 Linear 层的输出投影:
44

45
        W_o: (d_hidden, d_model)
46

47
        Row Parallel (切分 d_hidden):
48

49
        GPU 0: W_o0 (d_hidden/2, d_model)
50
        GPU 1: W_o1 (d_hidden/2, d_model)
51

52
        前向:
53
        X0: (B, d_hidden/2) ← 从 Attention 来
54
        X1: (B, d_hidden/2) ← 从 Attention 来
55
        ↓
56
        X0 @ W_o0 = Y0 (B, d_model)
57
        X1 @ W_o1 = Y1 (B, d_model)
58

59
        Y = Y0 + Y1  ← AllReduce (求和)
60

61
        通信: AllReduce
62
        """
63

64
    def attention_parallel(self):
65
        """
66
        Attention 中的 Tensor Parallelism
67
        """
68
        return """
69
        Attention 层并行化:
70

71
        ┌─────────────────────────────────────────────────────────┐
72
        │                    Attention Layer                        │
73
        ├─────────────────────────────────────────────────────────┤
74
        │                                                           │
75
        │  Input X                                                  │
76
        │      ↓                                                    │
77
        │  ┌─────────────────┐                                    │
78
        │  │  QKV Projection  │  Column Parallel (每个 GPU 3 heads) │
79
        │  └─────────────────┘                                    │
80
        │      ↓                                                    │
81
        │  ┌─────────────────────────────────────────────────┐     │
82
        │  │  Q 分布到所有 GPU                                │     │
83
        │  │  K, V 需要 AllGather 收集                      │     │
84
        │  │  各自计算部分 attention                          │     │
85
        │  │  Output 经过 AllReduce                          │     │
86
        │  └─────────────────────────────────────────────────┘     │
87
        │      ↓                                                    │
88
        │  ┌─────────────────┐                                    │
89
        │  │  Output Proj    │  Row Parallel                       │
90
        │  └─────────────────┘                                    │
91
        │                                                           │
92
        └─────────────────────────────────────────────────────────┘
93

94
        Megatron-LM 风格: 将 heads 分组，每组分配给不同 GPU
95
        """

1
class TPCommunication:
2
    """
3
    Tensor Parallelism 通信分析
4
    """
5

6
    def communication_pattern(self):
7
        """
8
        通信模式
9
        """
10
        return """
11
        Transformer 层的 TP 通信:
12

13
        ┌──────────────────────────────────────────────────────────┐
14
        │  Layer N                                                   │
15
        ├──────────────────────────────────────────────────────────┤
16
        │                                                           │
17
        │  Input X                                                   │
18
        │     ↓                                                     │
19
        │  QKV: Column Parallel → 无通信                            │
20
        │     ↓                                                     │
21
        │  Q AllGather (all TP ranks) → N-1 次广播                  │
22
        │  K AllGather (all TP ranks) → N-1 次广播                  │
23
        │  V AllGather (all TP ranks) → N-1 次广播                  │
24
        │     ↓                                                     │
25
        │  Attention Score (本地)                                   │
26
        │     ↓                                                     │
27
        │  Softmax → 无通信                                         │
28
        │     ↓                                                     │
29
        │  Attention × V → AllReduce (求和输出)                     │
30
        │     ↓                                                     │
31
        │  Output Proj: Row Parallel → AllReduce                     │
32
        │     ↓                                                     │
33
        │  FFN Gate: Column Parallel → 无通信                        │
34
        │     ↓                                                     │
35
        │  FFN Up: Column Parallel → 无通信                          │
36
        │     ↓                                                     │
37
        │  FFN AllGather → AllGather                                │
38
        │     ↓                                                     │
39
        │  FFN Down: Row Parallel → AllReduce                       │
40
        │     ↓                                                     │
41
        │  Output X                                                 │
42
        │                                                           │
43
        └──────────────────────────────────────────────────────────┘
44

45
        每层通信: 4 AllGather + 2 AllReduce
46
        """
47

48
    def bandwidth_requirement(self):
49
        """
50
        带宽需求
51
        """
52
        return {
53
            "per_layer_per_gpu": "2 × seq_len × hidden_size × tp_size (bytes)",
54
            "for_70b_seq_4096": "~128 MB per layer per forward",
55
            "forward_per_epoch": "~128 MB × 80 layers × 1M steps",
56
            "backward": "2x forward",
57

58
            "bandwidth_analysis": {
59
                "A100_NVLink": "900 GB/s bidirectional",
60
                "Gen1_InfiniBand": "25 GB/s",
61
                "Gen2_InfiniBand": "50 GB/s",
62
                "Gen4_InfiniBand": "200 GB/s",
63
            },
64

65
            "recommendation": "TP 需要 NVLink 或高速互联",
66
        }

1
class PipelineParallelism:
2
    """
3
    Pipeline Parallelism
4

5
    将模型按层切分到多个 GPU
6
    """
7

8
    def layer_distribution(self):
9
        """
10
        层分布
11
        """
12
        return """
13
        32 层模型，4-GPU Pipeline Parallel:
14

15
        GPU 0: Layers 0-7     (embedding + first 8 layers)
16
        GPU 1: Layers 8-15    (middle layers)
17
        GPU 2: Layers 16-23   (middle layers)
18
        GPU 3: Layers 24-31   (last 8 layers + output)
19

20
        Forward:
21
        ┌─────────────────────────────────────────────────────────┐
22
        │  GPU 0: [L0][L1][L2][L3][L4][L5][L6][L7] → hidden_7   │
23
        │      ↓                                                 │
24
        │  GPU 1: [L8][L9][L10][L11][L12][L13][L14][L15]→hidd_15│
25
        │      ↓                                                 │
26
        │  GPU 2: [L16]...[L23] → hidden_23                     │
27
        │      ↓                                                 │
28
        │  GPU 3: [L24]...[L31] → output                        │
29
        └─────────────────────────────────────────────────────────┘
30
        """
31

32
    def pipeline_bubble(self):
33
        """
34
        Pipeline Bubble 问题
35
        """
36
        return """
37
        Naive Pipeline (无 interleaving):
38

39
        时间 →→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→→
40

41
        Microbatch 1: [GPU0][GPU1][GPU2][GPU3]
42
        Microbatch 2:      [GPU0][GPU1][GPU2][GPU3]
43
        Microbatch 3:           [GPU0][GPU1][GPU2][GPU3]
44
        Microbatch 4:                [GPU0][GPU1][GPU2][GPU3]
45

46
        Bubble (空闲):
47
        [  ][  ][  ][GPU0][  ][  ][  ][GPU1][  ][  ][  ][GPU2][  ][  ][  ][GPU3]
48

49
        问题:
50
        → Bubble 数量 ≈ (P-1) / num_microbatches
51
        → P=4, microbatches=4 → 75% bubble!
52
        → GPU 利用率极低
53
        """
54

55
    def interleaving(self):
56
        """
57
        Interleaving 减少 Bubble
58
        """
59
        return """
60
        Interleaving (degree=4):
61

62
        每个 GPU 处理多个 chunks:
63

64
        GPU 0: [c0_0][c1_0][c2_0][c3_0][c4_0][c5_0]...
65
        GPU 1: [c0_1][c1_1][c2_1][c3_1][c4_1][c5_1]...
66
        GPU 2: [c0_2][c1_2][c2_2][c3_2][c4_2][c5_2]...
67
        GPU 3: [c0_3][c1_3][c2_3][c3_3][c4_3][c5_3]...
68

69
        每个 microbatch 在不同 GPU 上轮流执行
70
        Bubble 大幅减少
71

72
        Bubble 比例 ≈ 1 / (num_microbatches × degree)
73
        """

1
class ScheduleStrategies:
2
    """
3
    调度策略
4
    """
5

6
    def forward_backward(self):
7
        """
8
        Forward-Backward-Pipeline
9
        """
10
        return """
11
        GPipe 调度 (Channeh et al., 2019):
12

13
        1F1B 模式 (One Forward, One Backward):
14

15
        Forward:  [F0][F1][F2][F3]
16
        Backward: [B3][B2][B1][B0]
17

18
        完整调度 ( microbatches=8 ):
19

20
        Stage 0: [F0][F1][F2][F3][B3][B2][B1][B0]
21
        Stage 1:    [F0][F1][F2][F3][B3][B2][B1][B0]
22
        Stage 2:       [F0][F1][F2][F3][B3][B2][B1][B0]
23
        Stage 3:          [F0][F1][F2][F3][B3][B2][B1][B0]
24

25
        特点:
26
        → Forward 完成后立即开始 Backward
27
        → 稳定状态时每个 GPU 只做一个操作
28
        → Bubble 只出现在开始和结束
29
        """
30

31
    def interleaved_schedule(self):
32
        """
33
        Interleaved Schedule
34
        """
35
        return """
36
        Interleaved 1F1B:
37

38
        每个 stage 被分成多个 sub-stages:
39

40
        GPU 0: [S0_0][S0_1]...[S0_k]  (layers 0..k, k+1..2k)
41
        GPU 1: [S1_0][S1_1]...[S1_k]  (layers 2k..3k, ...)
42

43
        调度:
44
        F0_0, F0_1, F1_0, F0_2, F1_1, F1_2, ...
45

46
        优势:
47
        → 减少通信等待
48
        → 更好利用计算和通信重叠
49
        """
50

51
    def async_pipeline(self):
52
        """
53
        异步 Pipeline (实验性)
54
        """
55
        return """
56
        Chimera / PipeDream 风格:
57

58
        思想: Forward 和 Backward 可以异步
59

60
        GPU 0: [F0][F1][F2][F3]...[F0][F1][F2][F3]...  (连续 forward)
61
        GPU 1:    [F0][F1][F2][F3]...[B0][B1][B2][B3]...  (backward 滞后)
62

63
        权衡:
64
        → 更高 GPU 利用率
65
        → 收敛性可能受影响
66
        → 需要 careful weight versioning
67
        """

1
class PipeDreamImplementation:
2
    """
3
    PipeDream 风格实现
4
    """
5

6
    def weight_stashing(self):
7
        """
8
        Weight Stashing
9

10
        保存每个 microbatch 对应的权重版本
11
        """
12
        return '''
13
        class PipeDreamModel:
14
            """PipeDream 实现"""
15

16
            def __init__(self, stages, num_microbatches):
17
                self.stages = stages
18
                self.num_microbatches = num_microbatches
19
                self.weights = {}  # microb_id → weights
20
                self.optimizer = Adam(self.stages.parameters())
21

22
            def forward_backward(self, input_batch, microb_ids):
23
                """
24
                1F1B 调度
25
                """
26
                # Forward
27
                for mb_id in microb_ids:
28
                    self.stages[mb_id % len(self.stages)].forward(
29
                        self.weights[mb_id]  # 使用对应版本的权重
30
                    )
31

32
                # Backward
33
                for mb_id in reversed(microb_ids):
34
                    self.stages[mb_id % len(self.stages)].backward(
35
                        self.weights[mb_id]  # 梯度累积到正确版本
36
                    )
37

38
                # Update (每隔一段时间)
39
                self.optimizer.step()
40
                self.optimizer.zero_grad()
41
        '''
42

43
    def microbatch_vs_pipeline(self):
44
        """
45
        Microbatch 与 Pipeline 对比
46
        """
47
        return {
48
            "small_microbatch": {
49
                "pros": "更好的负载均衡，减少 bubble",
50
                "cons": "更多通信开销，更多 activation 显存",
51
            },
52
            "large_microbatch": {
53
                "pros": "减少通信，简化调度",
54
                "cons": "bubble 比例增加",
55
            },
56

57
            "recommendation": {
58
                "global_batch_size": "根据 GPU 数量调整",
59
                "microbatch_size": "尽量让 pipeline 填满",
60
                "num_microbatches": "P × 4 ~ P × 10",
61
            },
62
        }

1
class SequenceParallelism:
2
    """
3
    Sequence Parallelism
4

5
    沿序列维度切分
6
    """
7

8
    def why_sequence_parallel(self):
9
        """
10
        为什么需要 Sequence Parallelism
11

12
        当 TP 处理超长序列时:
13
        """
14
        return """
15
        场景: 100K tokens 序列, TP=8
16

17
        问题:
18
        1. Attention 计算: O(N²) 显存
19
           → 即使 TP 切分 Q/K/V，单个 head 仍需处理 N/8 tokens
20

21
        2. LayerNorm / Dropout
22
           → 需要看到完整序列
23
           → 成为瓶颈
24

25
        3. 显存分布不均
26
           → embedding / output 层: 参数量大
27
           → attention / ffn: 激活值大
28

29
        解决: Sequence Parallelism
30

31
        沿序列维度切分:
32
        seq_len = 100K
33
        SP = 8
34
        每个 GPU 处理 12.5K tokens
35
        """
36

37
    def layerwise_operations(self):
38
        """
39
        需要跨序列通信的操作
40
        """
41
        return {
42
            "attention": "需要 AllGather Q/K/V，然后 AllReduce 输出",
43
            "layernorm": "需要 AllReduce 计算 mean/variance",
44
            "dropout": "可以在本地处理",
45
            "embedding": "需要 AllGather (输入) 和 AllReduce (输出)",
46
            "loss": "需要 AllReduce",
47
        }

1
class RingCommunication:
2
    """
3
    环形通信实现
4
    """
5

6
    def ring_allreduce(self):
7
        """
8
        Ring AllReduce
9
        """
10
        return """
11
        Ring AllReduce 用于 Sequence Parallelism:
12

13
        4 GPUs, 每个持有 N/4 tokens 的部分和
14

15
        Step 1: Reduce Scatter
16
        ┌─────────────────────────────────────────────────────┐
17
        │  GPU 0: [a0] → [a0+b0+c0+d0]                        │
18
        │  GPU 1: [b0] → [b0+c0+d0+a0]                        │
19
        │  GPU 2: [c0] → [c0+d0+a0+b0]                        │
20
        │  GPU 3: [d0] → [d0+a0+b0+c0]                        │
21
        └─────────────────────────────────────────────────────┘
22

23
        Step 2: AllGather
24
        ┌─────────────────────────────────────────────────────┐
25
        │  GPU 0: [sum] → [sum, sum, sum, sum]                │
26
        │  GPU 1: [sum] → [sum, sum, sum, sum]                │
27
        │  GPU 2: [sum] → [sum, sum, sum, sum]                │
28
        │  GPU 3: [sum] → [sum, sum, sum, sum]                │
29
        └─────────────────────────────────────────────────────┘
30

31
        通信量: 2 × (N-1) / num_gpus
32
        """
33

34
    def ring_attention(self):
35
        """
36
        Ring Attention
37
        """
38
        return """
39
        Ring Attention 实现:
40

41
        4 GPUs, 序列长度 N, 每个 GPU 处理 N/4 tokens
42

43
        Q[i]: 本地查询
44
        K[i], V[i]: 本地键值
45

46
        1. Q 保持在本地
47

48
        2. K, V 通过环形传递:
49
           GPU 0: Q0 + [K0,V0] → 计算 [K0,V0] 的 attention
50
           GPU 0: Q0 + [K1,V1] → 计算 [K1,V1] 的 attention (接收自 GPU 1)
51
           GPU 0: Q0 + [K2,V2] → 计算 [K2,V2] 的 attention (接收自 GPU 2)
52
           GPU 0: Q0 + [K3,V3] → 计算 [K3,V3] 的 attention (接收自 GPU 3)
53

54
        3. 累加所有 attention 结果
55

56
        通信: 每个 GPU 接收 K,V N-1 次
57
        """

1
class HybridParallelism:
2
    """
3
    混合并行策略
4
    """
5

6
    def tensor_pipeline_2d(self):
7
        """
8
        TP × PP (2D Parallelism)
9
        """
10
        return """
11
        2D Parallelism: TP × PP
12

13
        例子: 70B 模型, 16 GPUs
14

15
        TP = 4, PP = 4:
16

17
        GPU Layout:
18
        ┌─────────┬─────────┬─────────┬─────────┐
19
        │  Stage 0 │  Stage 1 │  Stage 2 │  Stage 3 │
20
        │ (TP=4)  │ (TP=4)  │ (TP=4)  │ (TP=4)  │
21
        └─────────┴─────────┴─────────┴─────────┘
22

23
        每个 Stage 有 4 GPUs (TP group)
24
        4 个 Stages (Pipeline)
25

26
        通信:
27
        - TP 内部: AllReduce (NVLink)
28
        - PP 跨 stage: P2P (IB)
29
        """
30

31
    def dp_tp_pp_3d(self):
32
        """
33
        DP × TP × PP (3D Parallelism)
34
        """
35
        return """
36
        3D Parallelism: DP × TP × PP
37

38
        例子: 175B 模型, 64 GPUs
39

40
        DP = 8, TP = 4, PP = 2:
41
        总计: 8 × 4 × 2 = 64 GPUs
42

43
        结构:
44
        ┌─────────────────────────────────────────────────────┐
45
        │  DP Groups (8)                                       │
46
        │  ┌──────────┬──────────┬──────────┬──────────┐      │
47
        │  │  PP 0    │  PP 1    │  PP 0    │  PP 1    │ ...  │
48
        │  │ (TP=4)  │ (TP=4)  │ (TP=4)  │ (TP=4)  │      │
49
        │  └──────────┴──────────┴──────────┴──────────┘      │
50
        │       ↑                                               │
51
        │    Pipeline Communication                            │
52
        └─────────────────────────────────────────────────────┘
53

54
        内存估算 (175B FP16):
55
        - TP=4: 每个 GPU 持有 1/4 的单层权重
56
        - PP=2: 每 GPU 持有 1/2 的层数
57
        - DP=8: 每个 DP replica 独立
58

59
        每 GPU: 175B / (4 × 2) × 2 bytes = ~44 GB
60
        + activations + optimizer
61
        """
62

63
    def communication_topology(self):
64
        """
65
        通信拓扑感知
66
        """
67
        return """
68
        现代 GPU 集群的拓扑:
69

70
        Node 0: [GPU 0] [GPU 1] [GPU 2] [GPU 3]
71
        Node 1: [GPU 4] [GPU 5] [GPU 6] [GPU 7]
72
        ...
73

74
        带宽层次:
75
        ┌─────────────────────────────────────────────────────┐
76
        │  NVLink (within node): ~900 GB/s                    │
77
        │  PCIe (within node): ~64 GB/s                        │
78
        │  InfiniBand (cross node): ~50 GB/s                  │
79
        └─────────────────────────────────────────────────────┘
80

81
        最佳并行策略:
82
        ┌─────────────────────────────────────────────────────┐
83
        │  TP: within node (NVLink)                           │
84
        │  PP: across nodes (IB)                             │
85
        │  DP: across all (梯度同步相对较小)                    │
86
        └─────────────────────────────────────────────────────┘
87
        """

1
class MegatronArchitecture:
2
    """
3
    Megatron-LM 架构
4
    """
5

6
    def parallel_layout(self):
7
        """
8
        并行布局
9
        """
10
        return '''
11
        # Megatron-LM 3D Parallel
12

13
        # 配置:
14
        # tensor_model_parallel_size = 4  (TP)
15
        # pipeline_model_parallel_size = 4  (PP)
16
        # data_parallel_size = N / (TP × PP)
17

18
        # 模型构建
19
        model = TransformerBlock(
20
            tensor_parallel=TensorParallel,
21
            pipeline_parallel=PipelineParallel,
22
        )
23

24
        # Forward 流程:
25
        # 1. TP group 内广播输入
26
        # 2. Pipeline 调度
27
        # 3. 各 GPU 计算本地层
28
        # 4. PP 跨 GPU 传递 hidden states
29
        '''
30

31
    def memory_estimation(self):
32
        """
33
        显存估算
34
        """
35
        return """
36
        LLaMA-2 70B, 4096 sequence, batch=1, BF16:
37

38
        单 GPU (无并行):
39
        - 模型: 140 GB
40
        - 激活值: ~50 GB
41
        - 优化器: 280 GB
42
        - 总计: ~470 GB
43

44
        TP=8, PP=4:
45
        - TP: 每层权重 / 8
46
        - PP: 层数 / 4
47
        - 每 GPU 权重: 140 / 32 = 4.4 GB
48
        - 每 GPU 激活: 50 / 8 = 6.25 GB (SP 配合)
49
        - 每 GPU 优化器: 280 / 32 = 8.75 GB
50
        - 总计: ~19.4 GB per GPU
51

52
        结论: 可以在 A100 80GB 上训练
53
        """

1
class FSDPImplementation:
2
    """
3
    FSDP (Fully Sharded Data Parallel)
4
    """
5

6
    def basic_usage(self):
7
        """
8
        FSDP 基本用法
9
        """
10
        return '''
11
        import torch
12
        from torch.distributed.fsdp import (
13
            FullyShardedDataParallel as FSDP,
14
            ShardingStrategy,
15
            MixedPrecision,
16
            BackwardPrefetch,
17
        )
18
        from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
19

20
        # 混合精度
21
        fp16_policy = MixedPrecision(
22
            param_dtype=torch.float16,
23
            reduce_dtype=torch.float16,
24
            buffer_dtype=torch.float16,
25
        )
26

27
        # 定义模型
28
        model = Transformer()
29

30
        # FSDP 包装
31
        fsdp_model = FSDP(
32
            model,
33
            sharding_strategy=ShardingStrategy.FULL_SHARD,  # ZeRO-3
34
            mixed_precision=fp16_policy,
35
            auto_wrap_policy=transformer_auto_wrap_policy,
36
            backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
37
            device_id=torch.cuda.current_device(),
38
        )
39

40
        # 训练
41
        for batch in dataloader:
42
            loss = fsdp_model(batch)
43
            loss.backward()  # FSDP 自动分片梯度
44
            fsdp_model.optimize_step()
45
        '''
46

47
    def sharding_strategies(self):
48
        """
49
        分片策略
50
        """
51
        return {
52
            "FULL_SHARD": "参数 + 梯度 + 优化器状态全分片 (ZeRO-3)",
53
            "SHARD_GRAD_OP": "梯度 + 优化器状态分片 (ZeRO-2)",
54
            "NO_SHARD": "不分片，类似 DDP",
55
            "HYBRID_SHARD": "本地分片 + 跨节点通信",
56

57
            "recommendation": {
58
                "small_model": "SHARD_GRAD_OP (低通信开销)",
59
                "large_model": "FULL_SHARD (最低显存)",
60
                "multi_node": "HYBRID_SHARD (减少跨节点通信)",
61
            },
62
        }
63

64
    def cpu_offload(self):
65
        """
66
        CPU Offload
67
        """
68
        return '''
69
        # CPU Offload 配置
70

71
        from torch.distributed.fsdp import CPUOffload
72

73
        fsdp_model = FSDP(
74
            model,
75
            sharding_strategy=ShardingStrategy.FULL_SHARD,
76
            cpu_offload=CPUOffload(
77
                offload_params=True,  # 将参数卸载到 CPU
78
                pin_memory=True,
79
            ),
80
        )
81

82
        # 效果:
83
        # → 进一步减少 GPU 显存
84
        # → 但增加 CPU-GPU 传输开销
85
        # → 适合超大型模型
86
        '''