Multi-host config validation assumes 1 device, so global_batch_size_to_train_on is derived wrong

[WandLZhang]

On multi-host, config is validated before the JAX distributed system is initialised. get_num_target_devices() in configs/types.py calls jax.devices(), which raises at that point. The handler assumes 1 device:

self.num_target_devices = 1  # Default for validation when JAX is not initialized
try:
  self.num_target_devices = get_num_target_devices()
except (RuntimeError, IndexError):
  logger.warning("JAX device system not available for config validation. Assuming 1 device.")

With 1 device assumed, global_batch_size_to_train_on becomes per_device_batch_size regardless of the real chip count. The run then fails inside pjit with an indivisibility error that names neither the batch nor the device count, so the cause isn't obvious from the message.

Reproduce

Run any multi-host job on GKE without setting global_batch_size_to_train_on explicitly. A 4-host, 16-chip slice shows the warning above and then fails in pjit.

Workaround

Set global_batch_size_to_train_on explicitly. That is what I do, and it's enough.

Notes on a real fix

I also tried inferring the chip count from TPU_CHIPS_PER_HOST_BOUNDS and TPU_WORKER_HOSTNAMES in the except branch. It works on multi-host GKE, but I don't think it belongs upstream:

• It reads environment variables that multi-controller GKE happens to set. They aren't a supported interface.
• Under Pathways those variables are absent. That case is harmless, because train.py calls pathwaysutils.initialize() before pyconfig.initialize(argv), so the proxy backend is registered and jax.devices() returns the real device list. The except branch never runs. I confirmed this on a v5e slice: TPU_CHIPS_PER_HOST_BOUNDS and TPU_WORKER_HOSTNAMES were both unset and jax.devices() still returned 8.

A better fix runs the checks that need a device count after distributed init. Or it raises an error naming the missing setting instead of quietly assuming 1. Both are bigger than something I should propose on my own.

cc @mmcsa