[Train] Implement a JaxTrainer to support SPMD with TPUs

[ryanaoleary]

Why are these changes needed?

This PR builds off previous efforts to add a JaxTrainer and the ray-tpu package to implement support for a JaxTrainer in RayTrain that supports SPMD workloads with TPUs. Support for more types of workloads (i.e. better support for CPU and GPU) can be added incrementally.

In order to support SPMD locality-aware scheduling at the TPU slice level, we alter the WorkerGroup construction in V2 Ray Train to optionally accept multiple placement groups specs to apply to a range of workers. This enables us to reserve the "TPU head" using a placement group with label selectors, retrieve its unique ray.io/tpu-slice-name, and then schedule the remaining workers on that slice in a separate placement group.

TODO: I need to add good tests and my manual testing method with a real workload in the comments.

Related issue number

#55162

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[gemini-code-assist]

Summary of Changes

Hello @ryanaoleary, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

I've implemented a new JaxTrainer in Ray Train to enable distributed JAX training, specifically targeting SPMD workloads on Google TPUs. This involved significant enhancements to Ray's TPU accelerator management, including new APIs for detecting TPU topology and worker IDs, and integrating these details into Ray's node labeling system. I've also updated the ScalingConfig to allow users to specify TPU usage and topology, and introduced specialized placement group logic to correctly schedule multi-host TPU slices. Additionally, a mechanism to handle TPU lock files has been added to improve robustness.

Highlights

• New JaxTrainer for TPU Workloads: I've introduced a new JaxTrainer class within ray.train.jax, specifically designed to facilitate distributed JAX training on TPUs. This trainer leverages Ray's distributed capabilities to manage JAX workloads across multiple TPU devices.
• Expanded Scaling Configuration for TPUs: I've enhanced the ScalingConfig in both ray.air.config and ray.train.v2.api.config by adding use_tpu and topology parameters. These allow users to explicitly request TPU resources and specify the desired TPU topology, which is crucial for proper resource allocation and worker placement in SPMD scenarios.
• Improved TPU Accelerator Management and Metadata Discovery: Significant updates have been made to the TPUAcceleratorManager in python/ray/_private/accelerators/tpu.py. This includes adding a SPMDHost dataclass, functions to infer TPU pod types from topology, and a mechanism to fetch TPU metadata from remote hosts. Crucially, I've exposed previously private methods and added get_current_node_tpu_topology and get_current_node_accelerator_labels to provide detailed TPU-specific node information for scheduling.
• Integration of TPU-Specific Ray Node Labels: To support advanced TPU scheduling, I've integrated new TPU-specific node labels into Ray's core. New constants like kLabelKeyTpuTopology, kLabelKeyTpuSliceName, kLabelKeyTpuWorkerId, and kLabelKeyTpuPodType have been added to src/ray/common/constants.h and exposed in Python. These labels enable precise placement of TPU workers for SPMD execution.
• Specialized Placement Group Creation for SPMD on TPUs: I've implemented a specialized placement group creation logic, create_placement_group_with_spmd, within _backend_executor.py. This method handles the complex task of reserving a head node in a multi-host TPU slice and then atomically scheduling the remaining workers to that specific slice, ensuring proper SPMD setup.
• TPU Lock File Management Utility: A new utility, release_tpu_lock, has been added to python/ray/train/jax/config.py. This function helps manage and release potential TPU lock files, which can prevent JAX processes from utilizing TPUs, especially after job restarts or failures.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments or fill out our survey to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a JaxTrainer to support SPMD workloads on TPUs, which is a significant feature addition. The changes include new configurations in ScalingConfig, TPU-specific logic in the TPUAcceleratorManager, and a new JAX backend for Ray Train. The implementation looks mostly correct, but there are several critical issues that need to be addressed. These include a TypeError due to a function call with an incorrect signature in the JAX backend, another TypeError from incorrect list creation for placement group selectors, and a missing method call that would lead to an AttributeError. Additionally, there's a high-severity issue with the use of sudo in the code, which poses a security risk. I've also found an incorrect test case and some minor issues. Addressing these points will be crucial for the stability and security of this new feature.

[ryanaoleary]

cc: @matthewdeng @andrewsykim I moved all the code under V2 since this change adds new API fields to the ScalingConfig and changes the logic in the Train controller to accept multiple placement group specs. I'll update this PR with my manual tests with Jax today.

[matthewdeng]

This logic is a bit specific to TPUs so I don't know if it belongs in this generic Controller code.

I think we should aim to find a way to modularize this more, perhaps as part of a Callback that gets created/injected whenever use_tpu=True.

[ryanaoleary]

Done in a341722, I added a new callback TPUReservationCallback to ControllerCallback that gets called in _start_worker_group before creating the immutable WorkerGroupContext. This new callback returns a bundle_label_selector to influence the scheduling of the WorkerGroup. Currently the only callback that ends up getting called is TPUReservationCallback, but this way the logic in the controller does not directly reference TPUs and is extensible to other backends.

[matthewdeng]

Not sure if it's a good idea to pass in a created PlacementGroup. The divergence in PlacementGroup creation logic may lead to unexpected behavior depending on which path it goes down.

[ryanaoleary]

I was worried about the head PlacementGroup being released before the full slice was scheduled, causing a race condition. Thinking about it more I think that the PG shouldn't be removed until the JaxTrainer is cleaned up, so this shouldn't be a concern. I'll update the implementation to just return the slice name and pass a bundle_label_selector to the WorkerGroup.

[ryanaoleary]

Done in a341722.

[matthewdeng]

Is topology a first-class Ray Core concept? We'd want to make sure it's easy to understand from the API what inputs this takes in and how it'll be used.

[matthewdeng]

Also for TPU users how familiar are topology/accelerator? Would it be easier for the user to just specify the pod type directly?

[ryanaoleary]

I don't see topology used in Ray core at all, except to configure TPU env vars and node labels - but any users of multi-host TPUs should be familiar with the concept. The concept is also already introduced in KubeRay through the numOfHosts field: https://docs.ray.io/en/latest/cluster/kubernetes/user-guides/tpu.html.

I think topology and accelerator type are the best top-level variables for users to specify, since currently in GKE these are the two values users configure when creating their GKE nodepool and when scheduling pods to it using the cloud.google.com/gke-tpu-accelerator and cloud.google.com/gke-tpu-topology nodeSelectors: https://cloud.google.com/kubernetes-engine/docs/how-to/tpus.

[ultrons]

Topology is quite standard TPU concept. TPU type / Pod Type is in some cases not uniquely mapped to a topology.

[matthewdeng]

Awesome, is it safe to say that this API would then be super intuitive for a TPU user? Is there any other grouping/organization that might be more natural to how a user thinks about setting up their workload?

        scaling_config=ScalingConfig(
            use_tpu=True,
            num_workers=4,
            topology="2x2x4",
            accelerator_type="TPU-V4",
            resources_per_worker={"TPU": 4},
            placement_strategy="SPREAD",
        ),

[ryanaoleary]

Yeah I think this top level API should be clear for TPU users - the only thing I can think of is that we could have num_workers, resources_per_worker and placement_strategy be auto-set based on the topology if not provided. For example, if we have a multi-host topology of 4x4 v6e we could automatically detect that num_workers should be 4, resources_per_worker should be TPU: 4 since that's the number of chips on each host, and placement_strategy should be SPREAD.

[ryanaoleary]

Validated the JaxTrainer manually using the following MaxText example:

1. Setup a GKE cluster with a "2x2x2" v4 TPU nodepool, Workload Identity enabled, and GCS Fuse enabled.
2. Enable the Ray operator addon to install KubeRay and the KubeRay TPU webhook.
3. Setup kubernetes service account and GS bucket for checkpoints, set the names in your env under KSA and GSBUCKET respectively.
4. Create a RayCluster with the following spec:

apiVersion: ray.io/v1
kind: RayCluster
metadata:
  name: ray-tpu-cluster
spec:
  rayVersion: "2.48.0"
  headGroupSpec:
    rayStartParams: {}
    template:
      metadata:
        annotations:
          gke-gcsfuse/volumes: "true"
          gke-gcsfuse/cpu-limit: "0"
          gke-gcsfuse/memory-limit: "0"
          gke-gcsfuse/ephemeral-storage-limit: "0"
      spec:
        serviceAccountName: $KSA_NAME
        containers:
          - name: ray-head
            image: us-central1-docker.pkg.dev/ryanaoleary-gke-dev/ryanaoleary-ray/ray:ray-maxtext-image
            imagePullPolicy: Always
            ports:
            - containerPort: 6379
              name: gcs-server
            - containerPort: 8265
              name: dashboard
            - containerPort: 10001
              name: client
            resources:
              limits:
                memory: "16Gi"
              requests:
                cpu: "8"
                memory: "16Gi"
            volumeMounts:
            - name: gcs-fuse-csi-ephemeral
              mountPath: /data
            - name: dshm
              mountPath: /dev/shm
        volumes:
        - name: gke-gcsfuse-cache
          emptyDir:
            medium: Memory
        - name: dshm
          emptyDir:
            medium: Memory
        - name: gcs-fuse-csi-ephemeral
          csi:
            driver: gcsfuse.csi.storage.gke.io
            volumeAttributes:
              bucketName: $GSBUCKET
              mountOptions: "implicit-dirs"
  workerGroupSpecs:
    - replicas: 1
      numOfHosts: 4
      groupName: tpu-group
      rayStartParams: {}
      template:
        metadata:
          annotations:
            gke-gcsfuse/volumes: "true"
            gke-gcsfuse/cpu-limit: "0"
            gke-gcsfuse/memory-limit: "0"
            gke-gcsfuse/ephemeral-storage-limit: "0"
        spec:
          serviceAccountName: $KSA_NAME
          containers:
            - name: ray-worker
              image: us-central1-docker.pkg.dev/ryanaoleary-gke-dev/ryanaoleary-ray/ray:ray-maxtext-image
              imagePullPolicy: Always
              resources:
                limits:
                  memory: 200G
                  google.com/tpu: "4"
                requests:
                  cpu: "8"
                  memory: 200G
                  google.com/tpu: "4"
              env:
                - name: JAX_PLATFORMS
                  value: tpu,cpu
                - name: ENABLE_PJRT_COMPATIBILITY
                  value: "true"
              volumeMounts:
              - name: gcs-fuse-csi-ephemeral
                mountPath: /data
              - name: dshm
                mountPath: /dev/shm
          volumes:
          - name: gke-gcsfuse-cache
            emptyDir:
              medium: Memory
          - name: dshm
            emptyDir:
              medium: Memory
          - name: gcs-fuse-csi-ephemeral
            csi:
              driver: gcsfuse.csi.storage.gke.io
              volumeAttributes:
                bucketName: $GSBUCKET
                mountOptions: "implicit-dirs"
          nodeSelector:
            cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice
            cloud.google.com/gke-tpu-topology: 2x2x4

5. Apply the spec with envsubst to replace env var names with those you've set in your environment:

envsubst < ray-cluster-tpu-maxtext.yaml | kubectl apply -f -

6. Connect to the RayCluster:

kubectl port-forward svc/ray-tpu-cluster-head-svc 8265:8265 2>&1 >/dev/null &

Verified that TPU specific Ray node labels are set:

7. Create the following training script with MaxText and the new JaxTrainer:

import os
from absl import app
import logging
from typing import Sequence

import ray
from ray.train.v2.api.config import ScalingConfig, RunConfig
from ray.train.v2.jax import JaxTrainer
from MaxText.train import main as maxtext_main

def train_loop_per_worker(config):
    argv = config["argv"]
    maxtext_main(argv)

def main(argv: Sequence[str]):
    logging.basicConfig(level=logging.INFO)
    ray.init()

    trainer = JaxTrainer(
        train_loop_per_worker=train_loop_per_worker,
        train_loop_config={"argv": absolute_argv},
        scaling_config=ScalingConfig(
            use_tpu=True,
            num_workers=4,
            topology="2x2x4",
            accelerator_type="TPU-V4",
            resources_per_worker={"TPU": 4},
            placement_strategy="SPREAD",
        ),
        run_config=RunConfig(
            name="maxtext_jaxtrainer",
            worker_runtime_env={
                "env_vars": {
                    "JAX_PLATFORMS": "tpu",
                    "ENABLE_PJRT_COMPATIBILITY": "true",
                    "TPU_SLICE_BUILDER_DUMP_CHIP_FORCE": "true",
                    "TPU_SLICE_BUILDER_DUMP_ICI": "true",
                    "XLA_FLAGS": "--xla_dump_to=/tmp/xla_dump_file --xla_dump_hlo_as_proto",
                }
            },
        ),
    )

    result = trainer.fit()
    logging.info("Training complete!")

    ray.shutdown()

if __name__ == "__main__":
    app.run(main)

8. Submit the RayJob:

ray job submit \
  --address http://localhost:8265 \
  -- python maxtext_ray_trainer.py \
       MaxText/configs/base.yml \
       base_output_directory=gs://ray-maxtext-checkpoints/tmp/jax-trainer \
       dataset_type=synthetic \
       per_device_batch_size=2 \
       max_target_length=8192 \
       model_name=default \
       steps=100 \
       run_name=rayjob-with-maxtext

9. Training succeds:
   jax-trainer-maxtext-logs.txt

[matthewdeng]

Since we are only using this for finding the head host, we can simplify the logic to schedule an Actor rather than a Placement Group.

[ryanaoleary]

Do we know when the Ray Actor will get cleaned up? My concern is that the Ray Actor will go out of scope and release the head resource before the full slice schedules, causing a race condition. Conversely, if this does work is it okay to create a bunch of Ray Actors that are never cleaned up, running on each TPU head? I guess I wasn't sure whether that case or placement groups had a lower overhead.

[matthewdeng]

@MengjinYan do you know what is recommended from Core side?

I believe the lifecycle of Actors will be easier to track due to reference counting. We can create a local attribute in this Callback which will hold the head resource actor reference. If this Callback goes out of scope or if another job is launched, that reference would go away and get cleaned up.

[MengjinYan]

Ray will be automatically cleaned up when all the handle of it goes out of scope. So as long as the actor handle exists, the actor will be cleaned up.

I might misunderstood this but from the code, looks like after we create the placement group for the TPU head, we won't schedule anything to it other than getting the slice id from the head and when the function returns, we also didn't return the placement group for further training task scheduling. Wondering is it the expected behavior? If that's the case, then I think if using an actor, the actor will actually be out-of-scope and be released.

[matthewdeng]

For my understanding, are these always needed? Should they be set up automatically as part of the JaxBackend?

[ryanaoleary]

I believe the requirements differ across TPU generations, I think they're only currently used on TPU v6e. We could start passing the accelerator type into the JaxConfig and set the pjrt vars automatically for v6e? Or we could just leave it to the user to configure, my initial thought was to leave the JaxConfig as user specified as possible, and then add logic to autodetect the values for fields like num_workers, placement_strategy, reosurces_per_worker, and env vars to automatically set later based on the topology and accelerator type.

[matthewdeng]

I see. We can keep it manual here for now but over time as we discover patterns/boiler plate code that we can abstract away from the user by including into the default behavior, we can add this directly into the JaxBackend logic to reduce the barrier of entry for users.

[matthewdeng]

Is this specific to TPUs? Should it be tpu_topology instead?

[ryanaoleary]

I'm not super familiar with GPUs, but I think the field can probably be extended to set fields automatically in the Config (when left out) for GPUs too - so leaving it as topology might be fine. I don't have much of a preference either way though.

[matthewdeng]

Very nice!

[matthewdeng]

Going back on what I said before (oops), we can just check num_workers here.

Suggested change

	if scaling_config.use_tpu and (
	num_workers > 1 or scaling_config.num_workers > 1
	):
	if scaling_config.use_tpu and num_workers > 1 :

[andrewsykim]

done

[matthewdeng]

We can remove this.

Suggested change

	from ray.tune.search.sample import Domain
	SampleRange = Union["Domain", Dict[str, List]]

[andrewsykim]

done

[edoakes]

@ryanaoleary don't forget to add go label for CI