大模型并行的最小闭环:集合通信、TP、FSDP/ZeRO-3 到权重分片

大模型并行可以先从集合通信理解。无论是训练还是推理,多个设备之间都要交换张量,只是交换的时机和目的不同。

通信原语

常见原语包括:

  • Gather:收集。
  • Reduce:归约。
  • AllReduce:所有设备都得到归约结果。
  • Broadcast:广播。
  • Scatter:切分发送。

这些原语是 TP、DP、FSDP、ZeRO 的共同底层语言。

Tensor Parallel

Tensor Parallel 关注把单层计算拆到多张卡上。以线性层为例:

  • ColumnParallelLinear 按输出维度切权重。
  • RowParallelLinear 按输入维度切权重。

前者通常需要后续拼接或保持分片,后者通常需要对 partial output 做 AllReduce。

TP 的目标不是保存 optimizer state,而是让单层矩阵乘法跨设备执行。

FSDP / ZeRO-3

FSDP 和 ZeRO-3 更偏训练省显存。它们把参数、梯度、优化器状态分片。计算某一层时再 gather 需要的参数,算完后释放或重新分片。

因此它们和 TP 的核心区别是:TP 拆的是计算,FSDP/ZeRO-3 拆的是训练状态。

权重分片直觉

代码里常见 narrowoffset,本质是在完整权重中取出当前 rank 负责的切片。理解这些切片,比死记并行名词更重要。

大模型并行的最小闭环是:知道张量怎么切、什么时候通信、通信后每张卡手里有什么。

知识补全:训练并行和推理并行的目标不同

训练时显存压力来自参数、梯度、优化器状态和激活。FSDP / ZeRO-3 的核心价值是把这些状态分散到多张卡上,必要时再临时 gather。

推理时没有梯度和优化器状态,压力主要来自模型权重、KV cache 和单步 decode 延迟。Tensor Parallel 更常见,因为它能把单层权重和计算拆开,让单个大模型放进多张卡并并行执行。

这也是为什么同一个“并行”名词在训练和推理里含义不同。训练关心能不能省显存并保持反向传播正确;推理关心能不能降低单 token 延迟和提升吞吐。

通信直觉

通信不是免费的。每次 AllReduce、AllGather 都会带来同步等待。并行设计的关键是让计算节省大于通信成本。

判断一个并行方案时,可以问:

  1. 切分的是参数、激活、梯度,还是计算维度。
  2. 每一层需要几次通信。
  3. 通信发生在前向、反向,还是 decode 每一步。
  4. 是否会增加尾延迟。
  5. 当前瓶颈是显存、算力还是网络带宽。

这些问题比记住 TP、DP、PP、FSDP 的缩写更重要。

Source Log Coverage

The excerpts below are generated from Renyuan_Log.md and preserve the original tables, code fences, ASCII diagrams, commands, links, and explanations with source line numbers.

SourceLinesTitle
2026-05-012595-2606GPU 集合通信基础
2026-05-052647-2818张量并行与 Linear 切分
2026-05-072819-2820短记录占位
2026-05-082821-2822短记录占位
2026-05-092823-2824短记录占位
2026-05-102825-2830短记录占位
2026-05-112831-3015FSDP/ZeRO-3、TP 与权重分片

Source Log: 2026-05-01

Source lines: Renyuan_Log.md:2595-2606


2595 |# 2026-05-01
2596 |
2597 |## GPU操作
2598 |
2599 |Gather  
2600 |All-Gather  
2601 |Reduce  
2602 |All-Reduce  
2603 |Reduce-Scatter  
2604 |All-to-All  
2605 |
2606 |

Source Log: 2026-05-05

Source lines: Renyuan_Log.md:2647-2818


2647 |# 2026-05-05  
2648 |
2649 |分层 AllReduce + SHARP合在一起的真实执行路径讲清楚,并给一个具体数值例子。重点关注:哪一步在节点内做、哪一步跨节点、谁在算、谁在发。
2650 |
2651 |场景设定(例子)
2652 |2 个节点(Node A / Node B)
2653 |每个节点 4 张 GPU(共 8 张)
2654 |每张 GPU 有数据大小 B = 8 MB
2655 |网络:
2656 |节点内:NVLink / NVSwitch(很快)
2657 |节点间:InfiniBand + SHARP(较慢但可做网络内归约)
2658 |总体流程(一句话)
2659 |
2660 |先在节点内“压缩”(AllReduce),再把“压缩结果”交给交换机做全局归约,最后再分发回节点内。
2661 |
2662 |Step 1️⃣ 节点内 AllReduce(intra-node)⚡
2663 |
2664 |在 Node A 内(4 张 GPU):
2665 |
2666 |用 NCCL 做一次 AllReduce
2667 |Node B 同样做一遍
2668 |
2669 |结果:
2670 |
2671 |Node A 的每张 GPU 都拿到:A 节点的归约结果(8 MB)
2672 |Node B 同理
2673 |
2674 |👉 这里通信量确实是 B×N = 8MB × 4 = 32MB(节点内)
2675 |但因为 NVLink 很快,这一步不是瓶颈
2676 |
2677 |Step 2️⃣ 选“代表”参与跨节点 🌐
2678 |
2679 |每个节点选一个“代表”(通常是一个 GPU 或 NIC):
2680 |
2681 |Node A:选 GPU0
2682 |Node B:选 GPU0
2683 |
2684 |👉 注意:
2685 |不再是 4 张 GPU 都出去通信,而是每个节点只出 1 份数据
2686 |
2687 |Step 3️⃣ GPU → 交换机(上行)⬆️
2688 |Node A 的代表发送 8 MB
2689 |Node B 的代表发送 8 MB
2690 |
2691 |👉 跨节点总发送:
2692 |
2693 |每节点只发 B = 8MB
2694 |而不是 B×N = 32MB
2695 |Step 4️⃣ 交换机内做 Reduce(SHARP 核心)🧠
2696 |
2697 |在支持 SHARP 的交换机里(由 NVIDIA 网络实现):
2698 |
2699 |收到来自 Node A / B 的数据
2700 |直接在交换机 ASIC 里做:
2701 |result = A + B
2702 |
2703 |👉 这一步原本是在 GPU 上做的,现在搬到网络里
2704 |
2705 |Step 5️⃣ 交换机 → GPU(下行,多播)⬇️
2706 |交换机把结果(8 MB)广播回:
2707 |Node A 的代表
2708 |Node B 的代表
2709 |Step 6️⃣ 节点内广播(intra-node broadcast)🔁
2710 |
2711 |在每个节点内部:
2712 |
2713 |代表 GPU 把结果发给本节点其他 GPU(NVLink)
2714 |
2715 |最终:
2716 |
2717 |所有 8 张 GPU 都拿到全局 AllReduce 结果(8 MB)
2718 |📊 对比数据量(关键)
2719 |❌ 不用分层 + 不用 SHARP
2720 |
2721 |跨节点:
2722 |
2723 |4 张 GPU × 8MB = 32MB / 节点
2724 |✅ 分层 + SHARP
2725 |
2726 |跨节点:
2727 |
2728 |每节点只发 8MB
2729 |
2730 |👉 直接减少 4 倍(= N)
2731 |
2732 |🧠 直觉总结
2733 |
2734 |把整个过程想成“三段式”:
2735 |
2736 |① 节点内(快)
2737 |
2738 |“我们自己先算好一个结果”
2739 |
2740 |② 网络(SHARP)
2741 |
2742 |“交换机帮我们把各节点结果再合并”
2743 |
2744 |③ 节点内(快)
2745 |
2746 |“把最终答案发给自己人”
2747 |
2748 |🔥 为什么这是大模型训练的标配
2749 |
2750 |在:
2751 |
2752 |Megatron-LM
2753 |DeepSpeed
2754 |大规模 DDP
2755 |
2756 |中都会这样做,因为:
2757 |
2758 |NVLink 带宽 ≫ IB 带宽
2759 |跨节点才是瓶颈
2760 |SHARP 能进一步减少 GPU 参与通信
2761 |一句话总结
2762 |
2763 |分层 AllReduce 负责减少“谁跨节点通信”,SHARP 负责减少“通信时谁做归约 + 几轮通信”,两者结合把跨节点流量从 B×N 降到 B,并减少一轮计算/通信。
2764 |
2765 |
2766 |
2767 |4️⃣ 总体时间模型(关键)
2768 |
2769 |可以写成:
2770 |
2771 |不分层:
2772 |T = T_slow(BN)
2773 |
2774 |分层 + SHARP:
2775 |T = T_fast(BN) + T_slow(B)
2776 |
2777 |
2778 |
2779 |## 不同的并行方式
2780 |
2781 |### ColumnParallelLinear: 
2782 |**按照列维度分开。 某个GPU计算完后,结果的*某个维度*是*最终结果*,但是某GPU只有这些局部维度的信息。所以最后通过通讯来收集别的GPU结果**
2783 |
2784 |### RowParallelLinear:
2785 |**按照行维度分开。 某个GPU计算完后,有*每个*维度的信息,但是完整的维度上,都不是最终结果。所以最后*每个维度都*要再与来自其他GPU的中间信息进行计算,得到最终结果**
2786 |
2787 |
2788 |  ┌────────────────────────────┬─────────────┬──────────┬──────────────────────────┬────────────┬────────────────────┐
2789 |  │             类             │  切哪一维   │ 输入状态 │         输出状态         │    通信    │      放在哪里      │
2790 |  ├────────────────────────────┼─────────────┼──────────┼──────────────────────────┼────────────┼────────────────────┤
2791 |  │ ReplicatedLinear           │ 不切        │ 完整     │ 完整                     │ 无         │ 非并行场景或小矩阵 │
2792 |  ├────────────────────────────┼─────────────┼──────────┼──────────────────────────┼────────────┼────────────────────┤
2793 |  │ ColumnParallelLinear       │ out (dim 0) │ 完整     │ 切开                     │ 无         │ 一段计算的入口     │
2794 |  ├────────────────────────────┼─────────────┼──────────┼──────────────────────────┼────────────┼────────────────────┤
2795 |  │ RowParallelLinear          │ in (dim 1)  │ 切开     │ 完整(需 all-reduce)    │ all-reduce │ 一段计算的出口     │
2796 |  ├────────────────────────────┼─────────────┼──────────┼──────────────────────────┼────────────┼────────────────────┤
2797 |  │ MergedColumnParallelLinear │ out (dim 0) │ 完整     │ 切开(由多个子矩阵拼成) │ 无         │ gate+up 合并       │
2798 |  ├────────────────────────────┼─────────────┼──────────┼──────────────────────────┼────────────┼────────────────────┤
2799 |  │ QKVColumnParallelLinear    │ out (dim 0) │ 完整     │ 切开(Q/K/V 三段)       │ 无         │ attention 的 QKV   │
2800 |  └────────────────────────────┴─────────────┴──────────┴──────────────────────────┴────────────┴────────────────────┘
2801 |```
2802 |
2803 |```在 Megatron-LM / vLLM 中典型结构:
2804 |X
2805 | │
2806 | ├── ColumnParallel + Merged (QKV / MLP expand)
2807 | │
2808 | ▼
2809 | attention / activation
2810 | │
2811 | ├── RowParallel
2812 | ▼
2813 | Y
2814 |```
2815 |
2816 |
2817 |
2818 |

Source Log: 2026-05-07

Source lines: Renyuan_Log.md:2819-2820


2819 |# 2026-05-07
2820 |I can do this all day ...   

Source Log: 2026-05-08

Source lines: Renyuan_Log.md:2821-2822


2821 |# 2026-05-08  
2822 |赶路到广州    

Source Log: 2026-05-09

Source lines: Renyuan_Log.md:2823-2824


2823 |# 2026-05-09  
2824 |毕业照  

Source Log: 2026-05-10

Source lines: Renyuan_Log.md:2825-2830


2825 |# 2026-05-10  
2826 |赶路到深圳  
2827 |
2828 |
2829 |
2830 |

Source Log: 2026-05-11

Source lines: Renyuan_Log.md:2831-3015


2831 |# 2026-05-11
2832 |
2833 |## FSDP / ZeRO-3 和张量并行 TP 的区别
2834 |
2835 |  FSDP:切“模型状态”,主要为了省显存
2836 |  TP:切“单层计算”,主要为了让多卡一起算一个大矩阵
2837 |
2838 |  更具体地说:
2839 |
2840 |  | 对比点 | FSDP / ZeRO-3 | 张量并行 TP |
2841 |  |---|---|---|
2842 |  | 切什么 | 参数、梯度、优化器状态 | 线性层/注意力层里的大矩阵 |
2843 |  | 激活怎么切 | 通常按 batch/token 切 | 通常按 hidden dim / intermediate dim 切 |
2844 |  | 每张卡算什么 | 每张卡处理不同 batch 数据 | 多张卡一起算同一个 token 的同一层 |
2845 |  | 主要目的 | 降低显存占用 | 降低单卡计算量,并让超大层能并行计算 |
2846 |  | 通信内容 | 主要通信权重/梯度 | 主要通信激活/中间结果 |
2847 |  | 常见通信 | AllGather 参数,ReduceScatter 梯度 | AllGather 激活,ReduceScatter 或 AllReduce 输出 |
2848 |  | 对模型结构的侵入 | 相对低 | 较高,需要改 Linear/Attention 的实现 |
2849 |
2850 |  举个线性层例子:
2851 |
2852 |  Y = X @ W
2853 |
2854 |  假设:
2855 |
2856 |  X: [B, D]
2857 |  W: [D, F]
2858 |  Y: [B, F]
2859 |
2860 |  FSDP 的思路
2861 |
2862 |  每张卡存一部分 W,但真正算这一层时,会先把完整 W 收集回来:
2863 |
2864 |  平时:
2865 |  GPU0: W 第 0 片
2866 |  GPU1: W 第 1 片
2867 |  GPU2: W 第 2 片
2868 |  GPU3: W 第 3 片
2869 |
2870 |  计算前:
2871 |  AllGather -> 每张卡临时拿到完整 W
2872 |
2873 |  计算时:
2874 |  GPU0 算 batch 第 0 份
2875 |  GPU1 算 batch 第 1 份
2876 |  GPU2 算 batch 第 2 份
2877 |  GPU3 算 batch 第 3 份
2878 |
2879 |  所以 FSDP 本质上还是偏 数据并行:每张卡处理不同数据,只是参数平时不完整保存。
2880 |
2881 |  TP 的思路
2882 |
2883 |  每张卡不收集完整 W,而是真的只用自己那片 W 参与计算:
2884 |
2885 |  GPU0: W[:, 0:F/4]      -> 算 Y[:, 0:F/4]
2886 |  GPU1: W[:, F/4:F/2]    -> 算 Y[:, F/4:F/2]
2887 |  GPU2: W[:, F/2:3F/4]   -> 算 Y[:, F/2:3F/4]
2888 |  GPU3: W[:, 3F/4:F]     -> 算 Y[:, 3F/4:F]
2889 |
2890 |  最后再通过通信把结果拼起来或规约。
2891 |
2892 |  所以 TP 是真正把 一个矩阵乘法拆开计算。
2893 |
2894 |  一个简单类比:
2895 |
2896 |  FSDP:
2897 |  书太大,每个人平时只背几章。
2898 |  要读某一页时,大家临时把整本书凑出来,各自读自己的题。
2899 |
2900 |  TP:
2901 |  一道题太大,一个人算不完。
2902 |  大家每人算这道题的一部分,最后合并答案。
2903 |
2904 |
2905 |## narrow
2906 |
2907 |narrow 是 PyTorch 里用来沿某一个维度取连续切片的方法。
2908 |
2909 |  你的代码:
2910 |
2911 |  sliced_weight = loaded_weights.narrow(0, start_index, shard_size)
2912 |
2913 |  意思是:
2914 |
2915 |  从 loaded_weights 的第 0 维开始切
2916 |  起点是 start_index
2917 |  长度是 shard_size
2918 |
2919 |  等价于常见切片:
2920 |
2921 |  sliced_weight = loaded_weights[start_index : start_index + shard_size]
2922 |
2923 |  如果写完整一点:
2924 |
2925 |  loaded_weights.narrow(dim, start, length)
2926 |
2927 |  参数含义是:
2928 |
2929 |  dim: 沿哪个维度切
2930 |  start: 从这个维度的哪个下标开始
2931 |  length: 取多少个元素
2932 |
2933 |
2934 |## slided_weight = loaded_weights.narrow(0, start_index, shard_size)
2935 |
2936 |slided_weight 不是“纯指针”,而是一个 Tensor 视图(view)对象。
2937 |
2938 |  可以粗略理解成:
2939 |
2940 |  slided_weight = 指向原始数据的一张“切片窗口”
2941 |
2942 |  它本身是个 Tensor,带着:
2943 |
2944 |  - 数据在哪
2945 |  - 形状是多少
2946 |  - 从原始 loaded_weights 的哪一段开始看
2947 |  - 步长怎么走
2948 |
2949 |  但它通常不拷贝底层数据。
2950 |
2951 |  例如:
2952 |
2953 |  slided_weight = loaded_weights.narrow(0, start_index, shard_size)
2954 |
2955 |  这里 slided_weight 和 loaded_weights 共享同一块底层存储,只是 slided_weight 只看其中一部分。
2956 |
2957 |  所以你可以把它理解成:
2958 |
2959 |  - 不是新开一份完整内存
2960 |  - 也不是 C 语言里那种裸指针
2961 |  - 更像“带形状信息的视图引用”
2962 |
2963 |  真正把数据写进本地参数的是:
2964 |
2965 |  param_data.copy_(slided_weight)
2966 |
2967 |  这一步才发生拷贝。
2968 |
2969 |  一句话:
2970 |
2971 |  slided_weight 不是数据副本,而是原 tensor 上的一个切片视图。
2972 |
2973 |
2974 |## offset:
2975 |
2976 |offset = sum(self.output_sizes[:loaded_weight_id]) // self.tp_size ,// self.tp_size是什么作用
2977 |
2978 |// self.tp_size 的作用是:
2979 |
2980 |  把“完整大矩阵里的偏移量”换算成“当前 rank 本地分片里的偏移量”。
2981 |
2982 |  ———
2983 |
2984 |  假设:
2985 |
2986 |  output_sizes = [4096, 4096, 4096]
2987 |  tp_size = 4
2988 |
2989 |  完整合并矩阵是:
2990 |
2991 |  [ q: 4096 行 ][ k: 4096 行 ][ v: 4096 行 ]
2992 |
2993 |  如果现在加载 k,也就是:
2994 |
2995 |  loaded_weight_id = 1
2996 |
2997 |  完整矩阵里的偏移是:
2998 |
2999 |  sum(output_sizes[:1]) = 4096
3000 |
3001 |  也就是 k 在完整大矩阵里从第 4096 行开始。
3002 |
3003 |  但是当前 rank 本地只保存每个子矩阵的 1/4:
3004 |
3005 |  rank 本地矩阵:
3006 |  [ q shard: 1024 行 ][ k shard: 1024 行 ][ v shard: 1024 行 ]
3007 |
3008 |  所以 k 在本地矩阵里的起点不是 4096,而是:
3009 |
3010 |  4096 // 4 = 1024
3011 |
3012 |
3013 |
3014 |
3015 |