AI-Hypercomputer / AI-Hypercomputer/maxtext

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

Ouverte
#4,933 0 commentaires 0 réactions 0 personnes assignées Voir sur GitHub
Langage dominant
Python
Étoiles
2.4k
Forks
607
Merge moyen
2 j 19 h
PR mergées (30 j)
158

Description

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

```python
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`.

## 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

Guide de contribution

Ouvrir le guide de contribution

Évaluation

Cette issue n'a pas encore été évaluée.

Recevez les nouvelles issues par e-mail

Un résumé court des issues GitHub adaptées aux débutants.