AI-Hypercomputer / AI-Hypercomputer/maxtext

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

Đang mở
#4,933 0 bình luận 0 reaction 0 người được giao Xem trên GitHub
Ngôn ngữ chính
Python
Star
2.4k
Fork
607
Merge trung bình
2 ngày 19 giờ
Pull request đã merge (30 ngày)
158

Mô tả

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

Hướng dẫn đóng góp

Mở hướng dẫn đóng góp

Đánh giá

Issue này chưa được đánh giá.

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.