AI-Hypercomputer / AI-Hypercomputer/maxtext
Multi-host config validation assumes 1 device, so global_batch_size_to_train_on is derived wrong
- Linguagem predominante
- Python
- Estrelas
- 2.4k
- Forks
- 607
- Merge médio
- 2d 19h
- PRs com merge (30d)
- 158
Descrição
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
Guia de contribuição
Avaliação
Esta issue ainda não foi avaliada.