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 clientRemote 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 的真实边界。