TensorFlow 源码拆解

7 · 分布式执行:Strategy 如何改变 step 的边界

从 strategy、replica context 到 collective、remote runtime 的分层

分布式不是“把 batch 切成几份”

tf.distribute 同时改变变量创建、输入分片、函数调用、梯度聚合、状态同步和错误处理。高层 Strategy 将用户的 train_step 包装成 per-replica function,再由 strategy runtime 决定每个 replica 如何运行和如何 reduce。

Model.fit / user step
        ↓
Strategy.run(fn, args)
        ↓
replica context: fn on each replica
        ↓
collective all-reduce / all-gather / broadcast
        ↓
cross-replica context: merged result + variable update

变量的语义先于通信

MirroredVariable、SyncOnReadVariable、分布式 optimizer slot 都不是普通 Variable 的别名。读取变量时可能得到 replica-local value 或同步后的 value;写入时可能在每个 replica 执行、只在 cross-replica 执行,或通过 aggregation 合并。

输入分布

Strategy 将 Dataset 包装成分布式 iterator,并根据 input context 计算 pipeline id、num replicas、global batch size 和 per-replica batch。若 dataset 已经自行 sharding,再叠加 auto-sharding 可能改变样本分布;所以读 distribute 代码时必须把 input pipeline 和 collective 语义一起看。

Collective 与 rendezvous

Collective op 需要 group key、instance key、设备列表、通信实现和 rendezvous。不同 replica 必须以一致的顺序进入同一个 collective,否则会死锁或超时。TensorFlow 在 C++ runtime 中维护 collective executor、通信器、协调服务和远端 context;Python API 只是配置入口。

local EagerContext
  ├─ device manager
  ├─ collective executor manager
  ├─ rendezvous / coordination service
  └─ remote context + RPC client

Remote Eager 与 Graph 分布式

tensorflow/core/common_runtime/eager/context.h 中的 EagerContext 负责函数注册、设备发现、远端 master/worker 初始化和 context view id。Graph 模式通常经过 graph partition,把跨设备边转换为 send/recv;Eager 模式则可以提交 remote execute 请求。二者都要解决设备命名、函数在远端可见、Tensor 编码与错误传播。

错误防线

分布式失败包括单 replica kernel error、collective 超时、worker 消失、coordination service error 和 preemption。可靠实现需要把本地 Status 转成跨进程可理解的 error code,并让所有参与者尽快退出当前 collective;否则一个 worker 已失败,其他 worker 可能一直等待。

源码锚点

  • tensorflow/python/distribute/distribute_lib.py:Strategy / replica / cross-replica 公共抽象。
  • tensorflow/python/distribute/:Mirrored、MultiWorker、TPU、parameter server 等策略实现。
  • tensorflow/core/common_runtime/eager/context.h:远端 context、设备与函数同步。
  • tensorflow/core/distributed_runtime/:master、worker、session、RPC、partition。
  • tensorflow/core/common_runtime/collective/:collective executor、通信与 rendezvous。
  • tensorflow/python/data/experimental/service/:分布式数据服务客户端。

源码阅读实验

先在 MirroredStrategy 下打印 strategy.run 的返回结构,再把 tf.config.set_soft_device_placement(False) 打开。前者能看到 per-replica value,后者能暴露 device placement 和 kernel availability 的真实边界。

On this page