模型分片基础：什么是模型分片？
模型分片，说白了就是把一个大模型拆开。
你想想看，一个千亿参数的模型，单张显卡根本装不下。我最早接触这个思路是在做BERT-large部署的时候，那时候一张V100 32G放不下整个模型，急得我直挠头。后来才意识到，既然放不下，那就分开放呗。
模型分片的核心思想很简单：把模型的不同部分分配到不同的设备上。每个设备只负责模型的一部分计算，设备之间通过网络通信来交换中间结果。
模型分片 ≠ 数据并行
数据并行是每张卡都放一份完整的模型，只是把数据切分。模型分片是每张卡只放模型的一部分。这是两个完全不同的思路。
为什么要做模型分片？
原因其实就三个字：放不下。
- 显存墙：单卡显存有上限，H100也就80G。GPT-3 175B参数，光参数就要350GB，加上优化器状态、梯度，轻松上千GB。
- 计算墙：单卡算力有限，推理延迟要求高。分片后可以并行计算，降低单次推理的延迟。
- 带宽墙：模型太大，从CPU搬运到GPU太慢。分片后数据在GPU之间流动，走NVLink或InfiniBand，快得多。
我在项目中遇到过最头疼的情况：一个70B的模型，单卡推理延迟要8秒。分片到4张卡后，延迟降到了2.3秒。虽然通信有开销，但整体收益非常明显。
分片粒度：从粗到细
分片不是随便切一刀就完事。切多细？切哪里？这直接决定了你的推理效率和工程复杂度。
我个人习惯把分片粒度分成三个层次：
1. 层级别分片（Layer-wise）
这是最直观的方式。Transformer模型有几十层，把前10层放GPU0，中间10层放GPU1，最后10层放GPU2。
优点：实现简单，代码改动小。很多框架（如DeepSpeed、Megatron）原生支持。
缺点：存在严重的流水线气泡。GPU0算完才能传给GPU1，GPU1算完才能传给GPU2。大部分时间GPU在等待。
我的经验：层级别分片适合模型层数多、每层计算量大的场景。如果层数少于20层，流水线气泡会吃掉大部分并行收益。
2. 算子级别分片（Operator-wise）
把单个算子（比如矩阵乘法、注意力计算）拆开，分配到不同设备上。比如一个超大矩阵乘法，拆成多个小矩阵乘法并行计算。
优点：粒度更细，并行度更高。
缺点：实现复杂，需要手动改写算子。通信开销大，每个算子都要做一次all-reduce。
我曾经踩过的坑：算子级别分片时，通信和计算的比率要算清楚。如果通信时间 > 计算时间，分片反而更慢。我有个项目把注意力算子拆到4张卡上，结果通信开销占了60%，得不偿失。
3. 张量级别分片（Tensor-wise）
这是最细的粒度。把张量（权重矩阵）本身切分成多个分块，每个设备只存一个分块。比如一个[4096, 4096]的权重矩阵，切成4个[1024, 4096]的分块。
优点：显存利用率最高，每个设备只存自己需要的部分。
缺点：通信模式复杂，需要频繁做all-gather或reduce-scatter。对网络带宽要求极高。
| 分片粒度 | 实现难度 | 显存效率 | 通信开销 | 适用场景 |
|---|---|---|---|---|
| 层级别 | 低 | 中 | 低 | 层数多、模型深 |
| 算子级别 | 高 | 高 | 中 | 单算子计算量大 |
| 张量级别 | 极高 | 极高 | 高 | 超大模型、高带宽网络 |
分片与并行的关系
很多人搞混这两个概念。我简单说清楚：
- 分片是手段：把模型拆开，解决放不下的问题。
- 并行是目的：让多个设备同时干活，解决算不够的问题。
分片之后，天然就带来了并行。但并行不一定要靠分片——数据并行也是并行，模型并行也是并行。
这里有个关键点：分片粒度决定了并行模式。
- 层级别分片 → 流水线并行（Pipeline Parallelism）
- 算子级别分片 → 算子并行（Operator Parallelism）
- 张量级别分片 → 张量并行（Tensor Parallelism）
我在实际项目中，经常把这几种并行方式混着用。比如一个70B模型，先用张量并行把每层的权重切到4张卡上，再用流水线并行把层切到2组上，总共8张卡。这就是所谓的3D并行。
一句话总结：分片解决显存问题，并行解决算力问题。分片是并行的前提，并行是分片的价值体现。
分片的核心挑战
分片不是银弹。我踩过的坑，列出来给你参考：
- 通信瓶颈：分片越细，通信越频繁。NVLink带宽是600GB/s，但跨节点走IB只有200GB/s。分片策略要匹配硬件拓扑。
- 负载不均：不同层的计算量不一样。Attention层和FFN层的计算量差好几倍。简单按层数分片会导致有的卡忙死，有的卡闲死。
- 显存碎片：分片后每张卡的显存使用不连续，容易产生碎片。我有个项目分片后显存利用率只有65%，后来手动调整了分片边界才提到85%。
- 调试困难：分片后的模型，错误定位非常麻烦。是计算错了？还是通信丢了？还是同步没对齐？我曾经花了两天时间，才发现是一个all-reduce的buffer大小没对齐。
一个简单的分片示例
假设我们有一个两层Transformer，每层有一个Attention和一个FFN。用层级别分片到2张卡上：
# 伪代码：层级别分片
# GPU0: Layer 0
# GPU1: Layer 1
# GPU0 上的计算
def forward_gpu0(x):
h = layer0_attention(x)
h = layer0_ffn(h)
# 把结果发送给GPU1
send_to_gpu1(h)
return
# GPU1 上的计算
def forward_gpu1():
# 等待GPU0的结果
h = recv_from_gpu0()
h = layer1_attention(h)
h = layer1_ffn(h)
return h嗯，这里要注意：实际工程中不会这么裸写通信。框架会帮你封装好。但理解这个底层逻辑很重要——分片本质上就是切分 + 通信。
分片策略的选择建议
我个人的经验法则：
- 模型小于10B：不需要分片，数据并行就够了。
- 模型10B-70B：层级别分片 + 张量并行混合使用。层级别解决显存，张量并行解决计算。
- 模型大于70B：必须用3D并行。张量并行 + 流水线并行 + 数据并行，三管齐下。
一个小技巧：分片前先用profiler跑一下模型的计算图和显存分布。看看哪些层最吃显存，哪些算子最耗时。针对性地分片，比均匀分片效果好得多。
本章小结
模型分片就是把大模型拆开，放到多个设备上。分片粒度从粗到细有层级别、算子级别、张量级别。分片是手段，并行是目的。选择分片策略时，要综合考虑模型大小、硬件拓扑、通信带宽和负载均衡。
下一章我们会深入张量并行的具体实现，包括如何切分矩阵乘法、如何做all-reduce通信。到时候我会拿一个实际的LLM推理案例，手把手带你走一遍。
公众号:蓝海资料掘金营，微信deep3321