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>