5.2 分布式策略与性能优化

5.2 分布式策略与性能优化

本节摘要：单卡跑不动或跑不快时，出路分两条：横向扩容用分布式策略，把批次切开、副本并行、梯度聚合；纵向提速用性能剖析定位瓶颈，再按瓶颈段对症下药——数据管道、算子效率、通信开销各有专属药方。本节讲透数据并行的调度时序（每步发生了什么、全局批次怎么切、梯度在哪一刻被聚合），给出 MirroredStrategy 与 MultiWorkerMirroredStrategy 的实战代码，并整理一张"症状到药方"的优化速查表。

本节能力清单

阅读完本节，你应当能够：

1. 说出数据并行一步内的四个阶段：切批、副本前向、梯度聚合、参数同步；
2. 用 MirroredStrategy 把既有模型扩到单机多卡，改动控制在三行内；
3. 用 TensorBoard 剖析器读出瓶颈段并归因；
4. 按症状速查表选用混合精度、XLA 编译或管道调整。

数据并行的调度时序

数据并行的思想朴素：模型不变，数据切开。每个设备持有一份完整的模型副本（镜像），一个全局批次被均分成 N 份，各设备同时对自己那份做前向与梯度，然后梯度被聚合（默认 all-reduce 求平均），每个设备用聚合后的梯度同步更新自己的参数副本——镜像保证参数始终一致。排程时序一步四拍：

理解这四拍，两个实践推论立刻成立：全局批次大小是单卡批次乘卡数，学习率通常要随之放大；聚合阶段是新增的通信开销，算子太碎或模型太小时，通信可能吃掉并行的收益。

图 23 数据并行一步的设备视角

实战：三行扩到多卡

MirroredStrategy 的接入点极其克制：策略对象包住模型构建与 compile，批次大小按卡数放大，其余不动。fit 与自定义循环都能直接跑在策略作用域里：

import tensorflow as tf

strategy = tf.distribute.MirroredStrategy()   # 单机多卡，默认 NVLink 或 PCIe
print("devices:", strategy.num_replicas_in_sync)
# 单卡机器输出：devices: 1
# 多卡机器输出：devices: 2（或更多）

with strategy.scope():                        # 变量在各副本上镜像创建
    model = tf.keras.Sequential([tf.keras.layers.Input(shape=(8,)),
                                 tf.keras.layers.Dense(64, activation="relu"),
                                 tf.keras.layers.Dense(1)])
    model.compile(optimizer=tf.optimizers.Adam(1e-3), loss="mse")

import numpy as np
X = np.random.rand(8000, 8).astype("float32")
y = np.random.rand(8000).astype("float32")
model.fit(X, y, epochs=2, batch_size=128, verbose=0)   # 全局批次按卡数放大
print("distributed fit done")
# 输出：distributed fit done
# fit 内部自动把 128 切成每卡 64，梯度聚合对用户透明
​

多机训练换 MultiWorkerMirroredStrategy，额外要求每台机器配 TF_CONFIG 环境变量（声明集群节点角色与地址），代码主体不变。参数服务器策略（ParameterServerStrategy）适合超大嵌入表场景——模型大到单卡装不下参数时才需要，普通任务别碰。

# 多机配置示例（每台机器启动前设置，示意结构）
# import json, os
# os.environ["TF_CONFIG"] = json.dumps({
#     "cluster": {"worker": ["host1:12345", "host2:12345"]},
#     "task": {"type": "worker", "index": 0},     # 0 号机填 0，1 号机填 1
# })
strategy_multi = tf.distribute.MultiWorkerMirroredStrategy()
print("multi-worker strategy ready")
# 输出：multi-worker strategy ready
# 与单机版同名 API，模型代码零改动迁移
​

性能剖析：先定位再下药

优化的大忌是无的放矢。第一步永远是测量：TensorBoard 剖析器给出各段时间线——设备空闲段（数据饥饿）、算子间隙（调度开销）、通信段（聚合开销）一目了然。轻量级的手测先用总吞吐：固定步数计时，改一处测一次：

import time

def throughput(ds, steps=100):
    it = iter(ds)
    next(it)
    t0 = time.time()
    for _ in range(steps):
        next(it)
    dt = time.time() - t0
    return steps / dt

# 瓶颈实验：同一数据，管道有无预取的吞吐对比
raw = tf.data.Dataset.from_tensor_slices(np.random.rand(20000, 32, 32, 3).astype("float32"))
slow = raw.map(lambda img: tf.image.resize(img, [64, 64])).batch(64)
fast = (raw.map(lambda img: tf.image.resize(img, [64, 64]),
                num_parallel_calls=tf.data.AUTOTUNE)
        .batch(64).prefetch(tf.data.AUTOTUNE))
print(f"no-tuning: {throughput(slow):.0f} batches/s")
print(f"tuned    : {throughput(fast):.0f} batches/s")
# 参考输出：no-tuning: 310 batches/s / tuned: 1450 batches/s
# 四倍多差距全来自 2.5 节的三件套——管道几乎总是第一嫌疑
​

症状到药方：优化速查表

# 药方一：混合精度——计算切 16 位、关键处保 32 位
tf.keras.mixed_precision.set_global_policy("mixed_float16")
mp = tf.keras.Sequential([tf.keras.layers.Input(shape=(32, 32, 3)),
                          tf.keras.layers.Conv2D(32, 3, padding="same", activation="relu"),
                          tf.keras.layers.GlobalAveragePooling2D(),
                          tf.keras.layers.Dense(10, dtype="float32")])  # 末层保 float32
print("mixed precision on:", tf.keras.mixed_precision.global_policy().name)
# 输出：mixed precision on: mixed_float16
# 显存降三成上下、吞吐提升，末层 float32 防数值下溢

# 药方二：XLA——整图编译融合算子
@tf.function(jit_compile=True)
def fused_add(a, b, c):
    return (a + b) * c
r = fused_add(tf.constant(2.0), tf.constant(3.0), tf.constant(4.0))
print(float(r))
# 输出：20.0
# 小算子密集的计算图收益明显，大算子收益有限
​

⚠️ 常见坑：分布式下全局批次与学习率联动被忽略。卡数翻倍后批次翻倍而学习率不动，等效步长骤减，收敛变慢——经验法则是批次线性放大时学习率按同比例或平方根比例上调。

💡 关键直觉：性能优化是"测量、归因、对症"的循环，不是参数玄学。每改一处测一次吞吐，数字不涨就回退——手指比直觉可靠。

本节要点回顾

• 四拍时序：切批、副本前向、梯度聚合、参数同步，聚合是通信瓶颈。
• 三行扩容：策略作用域包住建模与 compile，全局批次按卡数放大。
• 多机迁移：MultiWorker 加 TF_CONFIG，代码主体不变。
• 先测后改：剖析器看时间线，吞吐计时做前后对照。
• 药方速查：管道三件套、混合精度、XLA，按症状对号入座。

下节解决最后一公里：把训练成果交付上线。