分散式推理:張量平行、流水線平行與專家平行

從張量平行、流水線平行到專家平行,拆解當模型太大、單張 GPU 放不下時,如何用多張 GPU 協同推理。

分散式推理:張量平行、流水線平行與專家平行

上一篇我們談了推測解碼,理解了如何用小模型換大速度,在不損失品質的前提下把生成速度提升 2 到 3 倍。

但推測解碼解決的是「單一請求為什麼這麼慢」的問題,還沒有解決「模型太大,單張 GPU 放不下」的問題。

一個 70B 模型用 FP16 儲存需要 140 GB,再加上 KV Cache,單張 80 GB 的 A100 根本不夠。
即使量化到 INT4,模型仍需 35 GB,加上 KV Cache 和運行開銷,也很難在單卡上服務多個請求。

這一篇,我們要談當模型太大時,如何用多張 GPU 協同推理:分散式推理
我們會深入三種核心平行策略:張量平行、流水線平行、專家平行,並討論如何選擇與組合。

Tensor Parallel + Pipeline Parallel + Expert Parallel

一、為什麼需要分散式推理?

問題一:模型太大,單卡放不下

模型FP16 大小INT4 大小單張 A100 80GB
7B14 GB3.5 GB輕鬆
13B26 GB6.5 GB可以
70B140 GB35 GBFP16 不行,INT4 勉強
175B350 GB87.5 GB都不行
405B810 GB202 GB需要多卡

問題二:即使放得下,吞吐也不夠

單張 GPU 的算力和記憶體頻寬有限。
當你需要服務大量請求時,單卡無法提供足夠的吞吐。

問題三:KV Cache 佔用大量記憶體

如我們在 KV Cache 那篇談過的,70B 模型在序列長度 4096、批次 32 時,KV Cache 需要 320 GB
這遠超過單張 GPU 的容量。

解決方案:分散式推理

把模型和計算分散到多張 GPU 上,協同完成推理。

單卡:  [GPU 0] 全部模型
多卡:  [GPU 0] 部分模型  [GPU 1] 部分模型  [GPU 2] 部分模型  [GPU 3] 部分模型

分散式推理的核心挑戰是:如何切分模型,如何通訊,如何保持效率。

二、平行策略的分類

分散式推理有多種平行策略,每種策略切分不同的維度。

策略切分維度通訊模式適用場景
資料平行(Data Parallel)批次梯度 All-Reduce訓練
張量平行(Tensor Parallel)權重矩陣All-Reduce單節點多卡
流水線平行(Pipeline Parallel)點對點多節點
專家平行(Expert Parallel)專家All-to-AllMoE 模型
序列平行(Sequence Parallel)序列長度Ring All-Reduce長上下文

對推理來說,最常用的是張量平行、流水線平行、專家平行
這三種策略可以組合,形成 3D 平行

Data | Tensor | Pipeline | Expert | Sequence

三、張量平行(Tensor Parallelism)

核心思想

張量平行把單一層的權重矩陣切分到多張 GPU 上。
每張 GPU 計算部分結果,然後透過通訊合併。

切分方式

以一個線性層 Y = XW 為例:

列平行(Column Parallel):

W 按列切分成 [W1, W2]
GPU 0 計算 Y1 = X W1
GPU 1 計算 Y2 = X W2
輸出 Y = [Y1, Y2](拼接)

行平行(Row Parallel):

W 按行切分成 [W1; W2]
GPU 0 計算 Y1 = X1 W1
GPU 1 計算 Y2 = X2 W2
輸出 Y = Y1 + Y2(All-Reduce)

在 Transformer 中的應用

以多頭注意力為例:

注意力頭 1-8  → GPU 0
注意力頭 9-16 → GPU 1

每個 GPU 負責一部分注意力頭,獨立計算,最後拼接。

以 FFN 為例:

第一個線性層:列平行(升維)
第二個線性層:行平行(還原)

這樣只需要在兩個地方通訊:

  1. 注意力輸出後的 All-Reduce。
  2. FFN 輸出後的 All-Reduce。

通訊開銷

張量平行每層需要兩次 All-Reduce。
All-Reduce 的成本與隱藏層大小和 GPU 數量成正比。

通訊量 ≈ 2 × batch_size × seq_len × hidden_size × 2 bytes

這對頻寬要求很高。
因此,張量平行通常只在單節點內使用,因為節點內有 NVLink(高頻寬、低延遲)。

實作範例:vLLM 的張量平行

from vllm import LLM, SamplingParams

# 用 4 張 GPU 做張量平行
llm = LLM(
    model="meta-llama/Llama-2-70b-hf",
    tensor_parallel_size=4,  # 4 張 GPU
    gpu_memory_utilization=0.9,
)

sampling_params = SamplingParams(temperature=0.7, max_tokens=200)
outputs = llm.generate(["今天天氣真好,適合"], sampling_params)

vLLM 會自動把模型切分到 4 張 GPU 上,並處理所有通訊。

優點與缺點

優點:

  • 記憶體和計算平均分攤。
  • 延遲低(單節點 NVLink)。
  • 實作相對簡單。

缺點:

  • 通訊開銷大,需要高頻寬互連。
  • 不適合跨節點(網路太慢)。
  • GPU 數量增加時,通訊成本線性增長。

四、流水線平行(Pipeline Parallelism)

核心思想

流水線平行把模型的不同層放到不同的 GPU 上。

GPU 0:第 1-10 層
GPU 1:第 11-20 層
GPU 2:第 21-30 層
GPU 3:第 31-40 層

資料像流水線一樣,從 GPU 0 傳到 GPU 3。

微批次(Micro-batching)

如果只有一個批次,流水線會有很多空閒時間(氣泡)。

時間軸:
GPU 0: [批次 1]
GPU 1:        [批次 1]
GPU 2:               [批次 1]
GPU 3:                      [批次 1]

為了解決這個問題,我們把批次拆成多個微批次:

GPU 0: [微批次 1][微批次 2][微批次 3][微批次 4]
GPU 1:           [微批次 1][微批次 2][微批次 3][微批次 4]
GPU 2:                     [微批次 1][微批次 2][微批次 3][微批次 4]
GPU 3:                               [微批次 1][微批次 2][微批次 3][微批次 4]

這樣 GPU 可以同時處理不同的微批次,利用率大幅提升。

氣泡比例

流水線的氣泡比例可以用以下公式估算:

氣泡比例=p1m+p1\text{氣泡比例} = \frac{p - 1}{m + p - 1}

其中:

  • p\text{p}:流水線階段數(GPU 數量)
  • m\text{m}:微批次數量
階段數 p微批次 m氣泡比例
4443%
4827%
41616%
4329%
83218%

微批次越多,氣泡越小。

通訊模式

流水線平行使用點對點通訊:每個 GPU 把輸出傳給下一個 GPU。

GPU 0 → GPU 1 → GPU 2 → GPU 3

通訊量小,適合跨節點(網路較慢)。

實作範例:vLLM 的流水線平行

from vllm import LLM

# 用 4 張 GPU,2 張做張量平行,2 張做流水線平行
llm = LLM(
    model="meta-llama/Llama-2-70b-hf",
    tensor_parallel_size=2,
    pipeline_parallel_size=2,
    # 總共 4 張 GPU
)

優點與缺點

優點:

  • 通訊量小,適合跨節點。
  • 可以擴展到很多 GPU。
  • 記憶體分攤效果好。

缺點:

  • 有氣泡,利用率不是 100%。
  • 延遲較高(需要等流水線填充)。
  • 實作較複雜。

五、專家平行(Expert Parallelism)

核心思想

專家平行專門用於 MoE(Mixture of Experts) 模型。
MoE 模型包含多個「專家」,每個 token 只會被路由到少數幾個專家。

輸入 token → 路由器 → 選擇 top-k 專家 → 專家計算 → 合併輸出

專家平行把不同的專家放到不同的 GPU 上。

GPU 0:專家 1、專家 2
GPU 1:專家 3、專家 4
GPU 2:專家 5、專家 6
GPU 3:專家 7、專家 8

通訊模式:All-to-All

專家平行需要 All-to-All 通訊:

  1. 每個 GPU 把 token 發送給負責該 token 專家的 GPU。
  2. 專家計算完成後,把結果發回原來的 GPU。
GPU 0 的 token A → 專家 3(在 GPU 1)
GPU 1 的 token B → 專家 1(在 GPU 0)
GPU 2 的 token C → 專家 7(在 GPU 3)
...

All-to-All 的通訊量與批次大小和隱藏層大小成正比。

挑戰

挑戰說明
負載均衡某些專家可能被過度使用
通訊開銷All-to-All 成本高
路由不穩定訓練和推理的路由可能不同
記憶體每個 GPU 要放多個專家

實作範例:DeepSpeed-MoE

import deepspeed

# DeepSpeed 的 MoE 配置
config = {
    "train_micro_batch_size_per_gpu": 1,
    "zero_optimization": {"stage": 3},
    "moe": {
        "enabled": True,
        "ep_size": 4,  # 專家平行度
        "num_experts": 8,
        "top_k": 2,
    },
}

model = deepspeed.init_inference(
    model,
    config=config,
)

優點與缺點

優點:

  • 模型容量大,但計算量不變。
  • 每個 token 只激活部分專家。
  • 適合超大模型。

缺點:

  • All-to-All 通訊開銷大。
  • 負載均衡困難。
  • 實作複雜。

六、序列平行(Sequence Parallelism)

核心思想

序列平行把序列長度切分到多張 GPU 上。
適合處理超長上下文(如 128K、1M token)。

Ring Attention

每張 GPU 負責一部分序列,用環狀通訊交換 KV。

GPU 0:token 1-4096
GPU 1:token 4097-8192
GPU 2:token 8193-12288
GPU 3:token 12289-16384

計算注意力時,每張 GPU 需要其他 GPU 的 KV。
透過環狀通訊,逐步交換。

適用場景

  • 長上下文推理。
  • 長文檔問答。
  • 超長序列生成。

實作範例:Ring Attention

# 概念性程式碼
def ring_attention(query, key, value, rank, world_size):
    """環狀注意力"""
    output = None
    for step in range(world_size):
        # 計算當前 KV 的注意力
        partial_output = attention(query, key, value)

        # 把 KV 傳給下一個 GPU
        key = send_recv(key, next_rank)
        value = send_recv(value, next_rank)

        # 合併輸出
        output = merge(output, partial_output)

    return output

七、混合平行:3D 平行

實務上,單一平行策略往往不夠。
我們會組合多種策略,形成 3D 平行

總 GPU 數 = 資料平行 × 張量平行 × 流水線平行

例如,64 張 GPU:

資料平行:4
張量平行:4
流水線平行:4
總計:4 × 4 × 4 = 64

分配策略

平行策略適合的網路建議大小
張量平行NVLink(節點內)≤ 8
流水線平行節點間網路任意
資料平行節點間網路任意

實作範例:vLLM 的 3D 平行

from vllm import LLM

llm = LLM(
    model="meta-llama/Llama-2-405b-hf",
    tensor_parallel_size=8,      # 單節點 8 張 GPU
    pipeline_parallel_size=4,    # 跨 4 個節點
    # 總共 32 張 GPU
)

DeepSpeed 的 3D 平行配置

{
    "train_batch_size": 32,
    "train_micro_batch_size_per_gpu": 1,
    "tensor_parallel": {
        "enabled": true,
        "tp_size": 4
    },
    "pipeline_parallel": {
        "enabled": true,
        "pp_size": 2
    },
    "zero_optimization": {
        "stage": 3
    }
}

八、如何選擇平行策略?

決策樹

模型放得下單卡嗎?
├── 是 → 需要更高吞吐嗎?
│         ├── 是 → 資料平行(多副本)
│         └── 否 → 單卡即可
└── 否 → 單節點內有高頻寬互連嗎?
          ├── 是 → 張量平行(≤ 8)
          └── 否 → 流水線平行(跨節點)

              是 MoE 模型嗎?
              ├── 是 → 加入專家平行
              └── 否 → 完成

選擇原則

場景推薦策略
單卡放得下,吞吐不足資料平行
單卡放不下,單節點多卡張量平行
跨節點,模型很大流水線平行
MoE 模型專家平行 + 張量平行
超長上下文序列平行
超大模型(405B+)3D 平行

網路頻寬的影響

網路頻寬適合的平行
NVLink600-900 GB/s張量平行
InfiniBand200-400 GB/s流水線平行、資料平行
Ethernet10-100 GB/s資料平行

張量平行對頻寬最敏感,所以只適合 NVLink。

九、最佳實踐

1. 先量化,再平行

# 先量化減少模型大小
llm = LLM(
    model="meta-llama/Llama-2-70b-hf",
    quantization="awq",           # 先量化
    tensor_parallel_size=2,       # 再平行
)

2. 張量平行不超過 8

# 不好的例子:張量平行 16,跨節點,通訊太慢
tensor_parallel_size=16

# 好的例子:張量平行 8(單節點 NVLink 上限)
tensor_parallel_size=8

3. 流水線平行用足夠的微批次

# 不好的例子:微批次 2,氣泡很大
micro_batches=2

# 好的例子:微批次 16 以上
micro_batches=16

4. 監控通訊開銷

def monitor_communication(llm):
    """監控通訊開銷"""
    stats = llm.llm_engine.statistics
    return {
        "all_reduce_time": stats.all_reduce_time,
        "all_to_all_time": stats.all_to_all_time,
        "compute_time": stats.compute_time,
        "communication_ratio": (
            stats.all_reduce_time + stats.all_to_all_time
        ) / stats.total_time,
    }

5. 使用高效的通信庫

# 安裝 NCCL
pip install nvidia-nccl-cu12

# 設定環境變數
export NCCL_IB_DISABLE=0
export NCCL_NET_GDR_LEVEL=2

6. 測試不同配置

def benchmark_parallel_configs(model, configs):
    """測試不同平行配置"""
    results = []

    for config in configs:
        llm = LLM(model=model, **config)
        throughput, latency = benchmark(llm)
        results.append({
            "config": config,
            "throughput": throughput,
            "latency": latency,
        })

    return sorted(results, key=lambda r: -r["throughput"])

7. 考慮成本

平行越多,通訊成本越高。找到效能與成本的平衡點。

8. 使用 vLLM 或 DeepSpeed

# 不要自己實作平行,用成熟的框架
from vllm import LLM
llm = LLM(model="...", tensor_parallel_size=4)

十、常見的陷阱

1. 張量平行跨節點

# 不好的例子:跨節點做張量平行,通訊太慢
tensor_parallel_size=16  # 跨 2 個節點

# 好的例子:張量平行只在單節點內
tensor_parallel_size=8
pipeline_parallel_size=2

2. 流水線微批次太少

# 不好的例子:微批次 2,氣泡 43%
# 好的例子:微批次 16,氣泡 16%

3. 忽略網路頻寬

# 不好的例子:用 Ethernet 做張量平行
# 好的例子:用 NVLink 做張量平行,Ethernet 做資料平行

4. 沒有監控

# 不好的例子:不知道通訊佔了多少時間
# 好的例子:監控 compute vs communication 比例

5. 平行度過高

# 不好的例子:用 64 張 GPU 跑 7B 模型
# 好的例子:用 1 張 GPU 跑 7B,或 2 張做張量平行

6. 忽略 KV Cache

# 不好的例子:只考慮模型權重,忽略 KV Cache
# 好的例子:同時考慮兩者

7. 沒有壓力測試

# 不好的例子:上線後才發現通訊瓶頸
# 好的例子:先做壓力測試

8. 選錯平行策略

# 不好的例子:MoE 模型不用專家平行
# 好的例子:MoE 用專家平行 + 張量平行

十一、總結:從單卡到多卡,從平行到協同

讓我們回顧這一篇的核心:

  • 為什麼需要分散式推理:模型太大、吞吐不夠、KV Cache 佔用大量記憶體。
  • 平行策略的分類:資料平行、張量平行、流水線平行、專家平行、序列平行。
  • 張量平行:切分權重矩陣,All-Reduce 通訊,適合單節點 NVLink。
  • 流水線平行:切分層,點對點通訊,適合跨節點,需足夠微批次。
  • 專家平行:切分專家,All-to-All 通訊,適合 MoE 模型。
  • 序列平行:切分序列,Ring Attention,適合長上下文。
  • 混合平行:3D 平行,組合多種策略。
  • 如何選擇:根據模型大小、GPU 數量、網路頻寬、場景。
  • 最佳實踐:先量化再平行、張量平行不超過 8、足夠微批次、監控通訊、用高效通信庫、測試配置、考慮成本、用成熟框架。
  • 常見陷阱:張量平行跨節點、微批次太少、忽略網路、沒有監控、平行度過高、忽略 KV Cache、沒有壓力測試、選錯策略。

分散式推理是 LLM 部署的終極挑戰。
它讓你在多張 GPU 上跑超大模型,服務大量使用者,同時保持可接受的延遲與成本。
但它也帶來了通訊、同步、調度等新的複雜性。

理解了分散式推理,你就掌握了部署超大模型的關鍵。
接下來,我們要談部署系列的最後一個主題:成本優化。


下一篇預告

《成本優化:快取、路由、降級與自動擴縮》

我們會談如何在不犧牲品質的前提下,降低 LLM 服務的成本,包括:快取策略、模型路由、降級機制、自動擴縮、以及如何計算與優化每請求成本。