transformers
5602e09e - fix(distributed): bind TPU ranks to device index 0, not LOCAL_RANK (#49543)

Commit
6 days ago
fix(distributed): bind TPU ranks to device index 0, not LOCAL_RANK (#49543) * fix(distributed): bind TPU ranks to device index 0, not LOCAL_RANK torch_tpu gives each torchrun process a single chip, which the process always sees as device index 0. Binding the device with LOCAL_RANK made set_device_index() fail on every rank but 0, so tensor-parallel loading crashed on ranks 1..N-1 and rank 0 hung on a collective. Add a `_local_device_index` helper that returns the current device index on TPU and LOCAL_RANK elsewhere, and use it wherever the device is bound. Behaviour on other accelerators is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * refactor(distributed): factor device binding into _bind_local_device * fix(distributed): default LOCAL_RANK to 0 in initialize_tensor_parallelism * refactor(distributed): inline TPU device index into _bind_local_device --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Author
Parents
Loading