Distribution strategy

Cross-cutting view of how multi-device, multi-host, and TPU training is organised. The Python API is documented in systems/distribution-strategy; this page describes the feature end-to-end with pointers to the Python and runtime layers.

What "distributed" means here

TensorFlow has at least four different distribution stories:

1. Single-host, multiple GPUs — MirroredStrategy, NCCL all-reduce.
2. Multi-host, multiple GPUs — MultiWorkerMirroredStrategy, gRPC + NCCL.
3. Parameter server — ParameterServerStrategy, async or sync, used for very large embedding models.
4. TPU pods — TPUStrategy, XLA-compiled, TPU-specific collectives.

A fifth story is DTensor (tensorflow/dtensor/), which differs in that the user describes layouts of tensors over a Mesh and the SPMD compiler partitions the program; this is closer to JAX's pjit model.

End-to-end flow

graph TD
    Code[user code with strategy.scope]
    Strategy[tf.distribute.Strategy]
    Vars[per-replica Variables]
    Datasets[DistributedDataset]
    Step[strategy.run train_step]
    XLA[Optional XLA compilation]
    Reduce[Cross-device reduce]
    Coord[Coordination service]
    RT[Distributed runtime gRPC]
    Devices[Per-replica devices]

    Code --> Strategy
    Strategy --> Vars
    Code --> Datasets
    Strategy --> Step
    Step --> XLA
    Step --> Reduce
    Reduce --> Devices
    Strategy --> Coord
    Coord --> RT
    RT --> Devices

Where the pieces live

Concern	Code	Wiki
Strategy classes	tensorflow/python/distribute/	systems/distribution-strategy
Cross-device reductions	tensorflow/python/distribute/cross_device_ops.py	"
Collectives in C++	tensorflow/core/common_runtime/collective_*, nccl/	systems/distributed-runtime
gRPC master/worker	tensorflow/core/distributed_runtime/	"
Coordination service (heartbeats)	tensorflow/core/distributed_runtime/coordination/	"
TPU strategy	tensorflow/python/distribute/tpu_strategy.py	systems/tpu
DTensor	tensorflow/dtensor/, tensorflow/python/dtensor/	systems/distribution-strategy
Failure handling / preemption	tensorflow/python/distribute/failure_handling/	"
Data sharding	tensorflow/python/distribute/input_lib.py	features/tf-data

What model.fit does on multiple GPUs

1. The user creates a MirroredStrategy().
2. Inside with strategy.scope(): they build the model. Variables are created as MirroredVariables (one copy per GPU).
3. model.fit(ds, ...) distributes ds via strategy.experimental_distribute_dataset(ds).
4. Inside fit's training loop, it calls strategy.run(train_step, args=(per_replica_batch,)). Each replica runs train_step on its local GPU.
5. Gradients are reduced across replicas (NCCL all-reduce).
6. The optimizer applies the reduced gradients to each MirroredVariable's local copy — they stay in sync.

For MultiWorkerMirroredStrategy, steps 4–6 are the same but the all-reduce is across hosts, gated by the coordination service for liveness checks and barrier sync.

Saving and loading distributed models

A distributed model saves like any other. Variables save their primary (replica-0) values; the strategy ensures all replicas are in sync at the time of save. Restoring is single-process; the restored model can then be re-distributed inside a new strategy.scope(). See features/saved-model.

Related

• systems/distribution-strategy — Python implementation details.
• systems/distributed-runtime — gRPC and collectives.
• systems/tpu — TPU-specific.
• features/tf-data — distributed input.