AI-Hypercomputer / AI-Hypercomputer/maxtext

[bug] Mid-Training Evaluation Loop Spawns Subprocess Claiming /dev/vfio/0 ('Device or resource busy')

オープン
#4,612 コメント 0 件 リアクション 0 件 担当者 0 名 GitHub で見る
bug
主要言語
Python
スター
2.4k
フォーク
607
平均マージ
2日 19時間
マージ済み PR(30日)
158

説明

### Bug report

1. Deploy single-host GRPO RL training job on an 8-chip Cloud TPU v6e slice (`v6e-2x4`).
2. Set `eval_interval=10` and `num_test_batches=25` in training parameters.
3. Upon hitting Step 10, `train_rl.py` lines 747–755 invoke `evaluate(...)`, which calls `create_models_and_meshes(...)` in `model_creation_utils.py`.
4. `create_models_and_meshes` executes `jax.devices()`, attempting to open `/dev/vfio/0` while the main RL trainer process holds the hardware lock.

**Expected:** In-flight evaluation reuses active TPU device handles without opening new hardware locks.
**Actual:** Evaluation fails with `RuntimeError: Unable to initialize backend 'tpu': FAILED_PRECONDITION: open(/dev/vfio/0): Device or resource busy`.

### Logs/Output

FAILED_PRECONDITION: TPU initialization failed: open(/dev/vfio/0): Device or resource busy
Traceback (most recent call last):
File "//train.py", line 140, in
rl_train(config_argv, {})
File "/usr/local/lib/python3.12/site-packages/maxtext/trainers/post_train/rl/train_rl.py", line 546, in rl_train
reference_model, reference_mesh, actor_model, actor_mesh, rollout_mesh = model_creation_utils.create_models_and_meshes(...)
File "/usr/local/lib/python3.12/site-packages/maxtext/utils/model_creation_utils.py", line 650, in create_models_and_meshes
reference_model, reference_mesh = from_pretrained(trainer_config, devices=trainer_devices, wrap_with_tunix_adapter=True)
File "/usr/local/lib/python3.12/site-packages/jax/_src/xla_bridge.py", line 482, in devices
return get_backend(backend).devices()
RuntimeError: Unable to initialize backend 'tpu': FAILED_PRECONDITION: TPU initialization failed: open(/dev/vfio/0): Device or resource busy

### Environment Information

- **Framework:** MaxText (GRPO RL / JAX Flax NNX)
- **Hardware:** 8x Cloud TPU v6e (`ct6e-standard-8t`, topology `v6e-2x4`)
- **OS:** Linux (GKE Standard Container Image)
- **Python:** 3.12
- **JAX:** 0.4.35+
- **Inference Engine:** vLLM V1 TPU (`v0.20.1rc1.dev136`)

### Additional Context

Org: Google Cloud GTM: Global Solutions Team
**Suggested Fix:**
In `src/maxtext/trainers/post_train/rl/train_rl.py`, update `evaluate(...)` to accept the existing active `actor_mesh` and `trainer_devices` directly, avoiding
`create_models_and_meshes(...)` re-execution on single-host TPU hardware.

コントリビューションガイド

コントリビューションガイドを開く

評価

この issue はまだ評価されていません。

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。