Need JAX computations to span multiple hosts? Call jax.distributed.initialize() after starting the cluster and before any JAX computation. It establishes distributed execution for JAX’s multicontroller model, where the same Python program runs independently on every host. section 7.4
Keep these operational rules in mind:
- Run the program manually on each host; a single Colab notebook does not currently provide a good process-management path. section 7.4
- Initialization is required for GPU systems and optional but recommended for TPU systems. section 7.4
- Use jax.process_index() to identify a process, then distinguish its local devices from global devices used through collective operations. section 7.4
Initialize every TPU worker
Place jax.distributed.initialize() at the top of the worker program, then print jax.process_index(). In a v2-32 TPU Pod slice, four workers report 32 global devices and eight local devices per host.
