Skip to content

← JAX Performance & Distributed step 25 of 25

Hard Primitives

Multi-Host JAX (Conceptual)

Jump to the task

Why this matters

A distributed JAX program can use devices controlled by several Python processes, often on different machines. This can provide more memory or more compute than one machine offers. A host is a machine, a process is a running program, and a device is a CPU, GPU or TPU exposed by a JAX backend. These counts are not interchangeable.

After distributed initialization, jax.device_count() reports the global number of devices for the selected backend. With no backend argument, it uses the default backend. jax.local_device_count() counts devices local to the current process. A single process may control several devices.

Worked mini-example

This inspection code works in an ordinary local session. It does not start a distributed job:

import jax

print("Processes:", jax.process_count())
print("This process:", jax.process_index())
print("Global devices:", jax.device_count())
print("Local devices:", jax.local_device_count())
assert jax.device_count() == len(jax.devices())

For example, four processes with eight devices each have 32 devices in total. Each process reports a local count of eight. This example assumes an equal device count per process; the global count is a sum, not always one process's local count multiplied by the number of hosts.

Distributed setup

In a real distributed launch, each participating process calls jax.distributed.initialize before initializing the JAX backend. Importing JAX is fine; querying devices or creating JAX arrays first is too late.

An explicit setup supplies a coordinator address reachable by all workers, a common num_processes, and a unique process_id from zero through num_processes - 1. These are process counts and IDs, even when your launcher happens to run exactly one process per host. Some cluster configurations can discover these settings automatically.

All expected processes must join. They need not start at the exact same instant, but missing workers or unreachable networking can cause a timeout. Initialization alone does not distribute your training data. You still need compatible meshes, data loading and coordinated computations; processes must issue collective operations in a consistent order.

Problem

Implement n_devices_in_mesh(dummy) to return jnp.float32(jax.device_count()). Ignore dummy; it is an unused scalar argument retained by this exercise's interface.

The function name is historical: you are counting the default backend's global devices, not constructing or inspecting a particular Mesh. Do not call jax.distributed.initialize in your submission. The runtime has already initialized JAX, and this exercise requires no distributed cluster.

Return a scalar JAX array with dtype float32. For one visible device, return 1.0; for four, return 4.0. The checker derives the expected count from your runtime rather than assuming everyone has one CPU device. Its displayed checks report whether the result has the required type and matches that count. Your function returns the count itself, not those check results.

Common pitfalls

  • Do not return a hard-coded one. Device counts depend on the runtime.
  • Local and global counts are equal in a single-process session, which does not make them interchangeable in distributed code.
  • A mesh may use only a subset of the available devices. Its size is not necessarily jax.device_count().
  • This exercise verifies an API result, not successful distributed setup, communication or performance across machines.