Add new tpu backend for torch

[siyuanfoundation]

Description

Add a new torch trainer backend for torch_tpu

Related to this announcement https://developers.googleblog.com/torchtpu-running-pytorch-natively-on-tpus-at-google-scale/

Related issues

Additional information

[gemini-code-assist]

Code Review

This pull request introduces TPU support for Ray Train and Ray AIR, including a new TPUTorchDeviceManager and logic to inject environment variables for distributed TPU training. Feedback highlights a missing use_tpu field in ScalingConfig that will cause runtime errors and an empty registration method for the tpu_dist backend. It is also recommended to use math.prod instead of functools.reduce for more idiomatic calculations of TPU chip products.

[ryanaoleary]

This is only called from ray.util.set_visible_accelerator_ids - so it won't be set automatically which I think is the intended goal. I would prefer adding something like set_accelerator_env_vars to AcceleratorManager and then the path for TPU is:

1. On Ray node init we call ResourceAndLabelSpec to pass required resources/labels to the Raylet.
2. Call set_accelerator_env_vars if the AcceleratorManager is not None
3. For the TPU implementation, we check for the expected Torch and/or JAX env vars and set them in the raylet process if missing
4. We may want to consider gating the above on if the user passes some flag or env var that indicates they're going to use Torch TPU.

This would keep it extensible to other accelerator types.

[ryanaoleary]

Avoid silent pass and add some debug logging if we expect set_device to be called but want to force it to be a no-op