Distribution strategy

The tf.distribute.* API and its underlying mechanics. Lives in tensorflow/python/distribute/. DTensor lives nearby in tensorflow/dtensor/ and tensorflow/python/dtensor/.

Purpose

• Provide a high-level API to distribute training across multiple GPUs on one host, multiple workers, or TPU pods, without rewriting model code.
• Implement the strategies: MirroredStrategy, MultiWorkerMirroredStrategy, ParameterServerStrategy, TPUStrategy, OneDeviceStrategy, CentralStorageStrategy.
• Drive the underlying distributed runtime and collective primitives.

Directory layout

tensorflow/python/distribute/
├── distribute_lib.py              # Strategy / StrategyExtended base classes
├── mirrored_strategy.py
├── mirrored_run.py
├── multi_worker_strategy.py        # MultiWorkerMirroredStrategy
├── parameter_server_strategy.py
├── parameter_server_strategy_v2.py
├── tpu_strategy.py
├── one_device_strategy.py
├── central_storage_strategy.py
├── values.py                       # PerReplica, Mirrored, DistributedValues
├── reduce_util.py                  # ReduceOp.SUM/MEAN/etc
├── input_lib.py                    # DistributedDataset
├── cross_device_ops.py             # AllReduceCrossDeviceOps, NcclAllReduce, ReductionToOneDevice
├── coordinator/                    # ClusterCoordinator for ParameterServerStrategy
├── failure_handling/               # Preemption-tolerant training helpers
├── combinations.py                 # Test combinations
└── ...

tensorflow/dtensor/
└── (DTensor — typed distributed tensor with explicit layouts)

Strategy taxonomy

Strategy	Use case
MirroredStrategy	Single host, multiple GPUs. Variables mirrored on every GPU; gradients all-reduced.
MultiWorkerMirroredStrategy	Multiple hosts, each with one or more GPUs. NCCL or ring all-reduce across hosts.
ParameterServerStrategy	Async or sync training with PS workers holding variables and worker tasks running steps.
TPUStrategy	Run on a TPU pod; all-reduce uses TPU collective.
OneDeviceStrategy	Convenience: forces all ops to one named device. Useful for tests.
CentralStorageStrategy	Variables on CPU, computation on GPU.

Strategy is the user-facing interface; StrategyExtended is the per-strategy implementation surface (distribute_lib.py).

How model.fit uses a strategy

graph TD
    User[user code]
    Scope[strategy.scope]
    Model[Keras Model]
    Data[strategy.experimental_distribute_dataset]
    Step[strategy.run train_step]
    Reduce[strategy.reduce ALL_REDUCE]

    User --> Scope
    Scope --> Model
    User --> Data
    Data --> Step
    Step --> Reduce
    Reduce --> Step

1. The user enters with strategy.scope(): and creates the model. Variables created inside the scope are wrapped in MirroredVariable/PSStrategyVariable/etc., living on multiple devices.
2. strategy.experimental_distribute_dataset(ds) returns a DistributedDataset that yields per-replica batches.
3. Inside model.fit (or a custom loop), strategy.run(train_step, args=(batch,)) runs train_step once per replica.
4. The replica's gradients are reduced (typically ReduceOp.SUM or MEAN) using cross-device ops in cross_device_ops.py — NCCL on GPU, CollectiveReduce on CPU/multi-host.
5. Variables are updated identically on every replica thanks to the all-reduce.

Cross-device ops

tensorflow/python/distribute/cross_device_ops.py is where the per-strategy reduction logic lives:

• NcclAllReduce — uses tensorflow/core/nccl/ for GPU.
• HierarchicalCopyAllReduce — designed for DGX-style topologies.
• ReductionToOneDevice — copies per-replica values to a single device, reduces, broadcasts back.
• CollectiveAllReduce — uses CollectiveReduce op which routes to the collectives in the distributed runtime.

Coordinator (ParameterServerStrategy)

ClusterCoordinator (tensorflow/python/distribute/coordinator/) is a higher-level abstraction over ParameterServerStrategy. The user submits async closures (schedule(fn)); the coordinator distributes them across worker tasks. The PSs hold variables; workers fetch, compute, push updates.

DTensor

DTensor (tensorflow/dtensor/, tensorflow/python/dtensor/) is a different model: tensors carry an explicit Layout that says how they're sharded across a Mesh of devices. Operations are SPMD-compiled across the mesh. DTensor is younger than tf.distribute and is the basis for parts of newer TPU-pod and large-scale training stacks.

Failure handling

tensorflow/python/distribute/failure_handling/ provides preemption-tolerant training: workers checkpoint regularly, detect preemption via a SIGTERM-like signal, and restart from the latest checkpoint. Used heavily on cloud TPU/GPU for spot instances.

Integration points

• Variables: tensorflow/python/distribute/distributed_variable.py and friends create per-strategy variable subclasses.
• tf.function: traced functions are aware of strategies via tf.distribute.has_strategy() and the _distribution_strategy attribute.
• Distributed runtime: collectives and rendezvous come from systems/distributed-runtime.
• Coordination service: tf.distribute calls into the coordination service for barriers and peer health.

Entry points for modification

• New strategy: subclass tf.distribute.Strategy (and StrategyExtended), implement _distribute_dataset, _run, variable creator, gradient reduction.
• New cross-device op: extend cross_device_ops.py with a new CrossDeviceOps subclass.
• Coordinator/scheduler tweaks: tensorflow/python/distribute/coordinator/.

Related

• distributed-runtime — collectives, gRPC, coordination service.
• features/saved-model — how distributed models are exported.
• apps/keras — model.fit integration.