[Train] Add TPU multi-slice support to JaxTrainer

[ryanaoleary]

Description

This PR adds support in the JaxTrainer to schedule across multiple TPU slices using the ray.util.tpu public utilities.

To support this, this PR adds new AcceleratorConfigs to the V2 scaling config, which consolidate the accelerator related fields for TPU and GPU. When TPUAcceleratorConfig is specified, the JaxTrainer utilizes a SlicePlacementGroup to atomically reserve num_slices TPU slices of the desired topology, auto-detecting the required values for num_workers and resources_per_worker when unspecified.

TODO: I'll add some manual testing and usage examples in the comments.

Related issues

#55162

Additional information

Optional: Add implementation details, API changes, usage examples, screenshots, etc.

[ryanaoleary]

cc: @andrewsykim @chiayi @matthewdeng @liulehui

[gemini-code-assist]

Code Review

This pull request introduces multi-slice TPU support for JaxTrainer by refactoring the accelerator configuration and leveraging ray.util.tpu.SlicePlacementGroup. The changes include a new AcceleratorConfig API, which provides a cleaner way to specify GPU and TPU resources, and deprecates older fields like use_gpu and use_tpu. The logic for reserving TPU slices and determining worker configurations is now encapsulated within SlicePlacementGroup, which correctly handles multi-slice reservations and auto-detects num_workers and resources_per_worker. The test suite has been significantly improved to cover various single-host, multi-host, and multi-slice scenarios. Overall, this is a well-structured and comprehensive update that greatly enhances TPU support in Ray Train. I have a couple of suggestions to address a potential runtime error and a syntax issue.

[liulehui]

thank you!!

[liulehui]

🫶

[github-actions]

This pull request has been automatically marked as stale because it has not had
any activity for 14 days. It will be closed in another 14 days if no further activity occurs.
Thank you for your contributions.

You can always ask for help on our discussion forum or Ray's public slack channel.

If you'd like to keep this open, just leave any comment, and the stale label will be removed.

[ryanaoleary]

cc: @matthewdeng @liulehui I've resolved the outstanding comments and this PR is ready for another review. I'm fixing the comments on #59136 now so that it can be merged first, and then this PR will just contain the Ray Train changes.

[ryanaoleary]

@liulehui Thank you for the reviews!! I think the CI failure for test_jax_distributed_shutdown_timeout should be fixed with 1756d33.

[matthewdeng]

Thanks! Few minor remaining comments.

[dayshah]

have a nit, but non-blocking can address in a follow-up if it makes sense @ryanaoleary

[dayshah]

Why 1 on failed calculation or invalid inputs? I feel like it makes more sense to raise, but maybe missing something here

[ryanaoleary]

@dayshah I had it default to 1 because in the Ray Train code we call validate_tpu_config and get_tpu_worker_resources for the TPU inputs, which should avoid this ever raising since we validate that the num_workers and topology / accelerator type are compatible with valid inputs.

My thought was that if for some reason this call did fail in the controller (which is the only place it's called), that it'd be better to return a default slice of 1 rather than crash the controller. I think since I don't expect this case to happen regardless due to the validation, that I can update this util to Raise like you suggest since yeah that does make more sense.