AI-Hypercomputer / AI-Hypercomputer/maxtext

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

未關閉
#4,933 0 則留言 0 個 reaction 已指派 0 人 在 GitHub 檢視
主要語言
Python
星號
2.4k
分支
607
平均合併
2 天 19 小時
30 天內合併 PR
158

描述

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

貢獻指南

開啟貢獻指南

評估

這個 Issue 還沒有評估資料。

把新 issue 寄到你的電子郵件信箱

精選適合新手參與的 GitHub issue 摘要。