Hyperparameter tuning with Ray Tune#

Important

This guide shows how to integrate Ray Train and Ray Tune to tune hyperparameters for distributed training runs with Ray Train V2. Ray Train V2 is the default starting in Ray 2.51. In Ray 2.43 through 2.50, enable it by setting the environment variable RAY_TRAIN_V2_ENABLED=1.

For information about the deprecation and migration, see Tuner(trainer) API deprecation.

You can combine Ray Train with Ray Tune to run hyperparameter sweeps over distributed training runs. This combination is often useful for a small sweep over critical hyperparameters before you launch a long run with the best-performing hyperparameters on all available cluster resources.

Quickstart#

The following example uses these components:

  • Tuner launches the tuning job, which runs trials of train_driver_fn with different hyperparameter configurations.

  • train_driver_fn takes in a hyperparameter configuration, instantiates a TorchTrainer or another framework’s trainer, and launches the distributed training job.

  • ScalingConfig defines the number of training workers and resources per worker for a single Ray Train run.

  • train_fn_per_worker is the Python code that executes on each distributed training worker for a trial.

import random
import tempfile
import uuid

import ray.train
import ray.train.torch
import ray.tune
from ray.tune.integration.ray_train import TuneReportCallback


# [1] Define your Ray Train worker code.
def train_fn_per_worker(train_loop_config: dict):
    # Unpack train worker hyperparameters.
    # Train feeds in the `train_loop_config` defined below.
    lr = train_loop_config["lr"]

    # training code here...
    print(
        ray.train.get_context().get_world_size(),
        ray.train.get_context().get_world_rank(),
        train_loop_config,
    )
    # model = ray.train.torch.prepare_model(...)  # Wrap model in DDP.
    with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
        ray.train.report(
            {"loss": random.random()},
            checkpoint=ray.train.Checkpoint.from_directory(temp_checkpoint_dir),
        )


# [2] Define a function that launches the Ray Train run.
def train_driver_fn(config: dict):
    # Unpack run-level hyperparameters.
    # Tune feeds in hyperparameters defined in the `param_space` below.
    num_workers = config["num_workers"]

    trainer = ray.train.torch.TorchTrainer(
        train_fn_per_worker,
        train_loop_config=config["train_loop_config"],
        scaling_config=ray.train.ScalingConfig(
            num_workers=num_workers,
            # Uncomment to use GPUs.
            # use_gpu=True,
        ),
        run_config=ray.train.RunConfig(
            # [3] Assign unique names to each run.
            # Recommendation: use the trial id as part of the run name.
            name=f"train-trial_id={ray.tune.get_context().get_trial_id()}",
            # [4] (Optional) Pass in a `TuneReportCallback` to propagate
            # reported results to the Tuner.
            callbacks=[TuneReportCallback()],
            # (If multi-node, configure S3 / NFS as the storage path.)
            # storage_path="s3://...",
        ),
    )
    trainer.fit()


# Launch a single Train run.
# Note that you can only create a TuneReportCallback in a Ray Tune session.
# train_driver_fn({"num_workers": 4, "train_loop_config": {"lr": 1e-3}})


# Launch a sweep of hyperparameters with Ray Tune.
tuner = ray.tune.Tuner(
    train_driver_fn,
    param_space={
        "num_workers": ray.tune.choice([2, 4]),
        "train_loop_config": {
            "lr": ray.tune.grid_search([1e-3, 3e-4]),
            "batch_size": ray.tune.grid_search([32, 64]),
        },
    },
    run_config=ray.tune.RunConfig(
        name=f"tune_train_example-{uuid.uuid4().hex[:6]}",
        # (If multi-node, configure S3 / NFS as the storage path.)
        # storage_path="s3://...",
    ),
    # [5] (Optional) Set the maximum number of concurrent trials
    # in order to prevent too many Train driver processes from
    # being launched at once.
    tune_config=ray.tune.TuneConfig(max_concurrent_trials=2),
)
results = tuner.fit()

print(results.get_best_result(metric="loss", mode="min"))

What does Ray Tune provide?#

Ray Tune provides utilities for the following tasks:

This guide focuses only on the integration layer between Ray Train and Ray Tune. For details on using Ray Tune, see the Ray Tune documentation.

Configure resources for multiple trials#

Ray Tune launches multiple trials that each run a user-defined function in a remote Ray actor. Each trial gets a different sampled hyperparameter configuration.

When you use Ray Tune by itself, trials do computation directly inside the Ray actor. For example, each trial could request 1 GPU and run single-process model training within the remote actor. When you use Ray Train inside Ray Tune functions, the Tune trial doesn’t do extensive computation inside this actor. Instead, it acts as a driver process that launches and monitors the Ray Train workers running elsewhere.

Ray Train requests its own resources through the ScalingConfig. For details, see Configure scale and accelerators.

../../_images/train_without_tune.png

A single Ray Train run. The next figure shows how Ray Tune adds a layer of hierarchy to this tree of processes.#

../../_images/train_tune_interop.png

Ray Train runs launched from within Ray Tune trials.#

Limit the number of concurrent Ray Train runs#

A Ray Train run starts only when it can acquire resources for all of its workers at once. As a result, multiple Tune trials that spawn Train runs compete for the logical resources available in the Ray cluster.

If a cluster resource such as GPUs is limited, you can’t run training for all hyperparameter configurations concurrently. Because the cluster has enough resources for only a handful of concurrent trials, set tune.TuneConfig(max_concurrent_trials) on the Tuner to limit the number of in-flight Train runs so that no trial starves for resources.

# For a fixed size cluster, calculate this based on the limiting resource (ex: GPUs).
total_cluster_gpus = 8
num_gpu_workers_per_trial = 4
max_concurrent_trials = total_cluster_gpus // num_gpu_workers_per_trial


def train_driver_fn(config: dict):
    trainer = ray.train.torch.TorchTrainer(
        train_fn_per_worker,
        scaling_config=ray.train.ScalingConfig(
            num_workers=num_gpu_workers_per_trial, use_gpu=True
        ),
    )
    trainer.fit()


tuner = ray.tune.Tuner(
    train_driver_fn,
    tune_config=ray.tune.TuneConfig(max_concurrent_trials=max_concurrent_trials),
)

For example, consider a fixed-size cluster with 128 CPUs and 8 GPUs.

  • The Tuner(param_space) runs a grid search over four hyperparameter configurations, defined by param_space={"train_loop_config": {"batch_size": tune.grid_search([8, 16, 32, 64])}}.

  • Each Ray Train run trains with four GPU workers, as set by ScalingConfig(num_workers=4, use_gpu=True). Because the cluster has only 8 GPUs, only two Train runs can acquire their full set of resources at a time.

  • However, the cluster has many CPUs, and each Ray Tune trial requests 1 CPU by default, so Ray Tune launches all four trials immediately. That launches two extra Ray Tune trial processes whose inner Ray Train runs wait for resources until one of the other trials finishes. While Train waits for resources, it emits noisy, repetitive log messages. If the total number of hyperparameter configurations is large, there might also be an excessive number of Ray Tune trial processes.

  • To fix this issue, set Tuner(tune_config=tune.TuneConfig(max_concurrent_trials=2)). Only two Ray Tune trial processes then run at a time. Calculate this number from the limiting cluster resource and the amount of that resource each trial requires.

Advanced: Set Train driver resources#

By default, the Train driver runs as a Ray Tune function with 1 CPU. Ray Tune schedules these functions on any node in the cluster that has free logical CPU resources.

Tip

If you launch longer-running training jobs or use spot instances, run the Tune functions that act as the Ray Train driver process on safe nodes, which are at lower risk of going down. For example, don’t schedule them on preemptible spot instances, and don’t run them on the same nodes as training workers. A safe node could be the head node or a dedicated CPU node in your cluster.

Placing the driver on a safe node matters because the Ray Train driver process handles fault tolerance for the worker processes, which are more likely to fail. Nodes that run Train workers can crash because of spot preemption or other errors that come from your model training code.

  • If a Train worker node dies, the Ray Train driver process on a different node is still alive and can handle the error gracefully.

  • If the driver process dies, all Ray Train workers exit ungracefully, and some of the run state might not be fully committed.

One way to place the Train driver on safe nodes is to set custom resources on certain node types and configure the Tune functions to request those resources.

# Cluster setup:
# head_node:
#     resources:
#         CPU: 16.0
# worker_node_cpu:
#     resources:
#         CPU: 32.0
#         TRAIN_DRIVER_RESOURCE: 1.0
# worker_node_gpu:
#     resources:
#         GPU: 4.0

import ray.tune


def train_driver_fn(config):
    # trainer = TorchTrainer(...)
    ...


tuner = ray.tune.Tuner(
    ray.tune.with_resources(
        train_driver_fn,
        # Note: 0.01 is an arbitrary value to schedule the actor
        # onto the `worker_node_cpu` node type.
        {"TRAIN_DRIVER_RESOURCE": 0.01},
    ),
)

Report metrics and checkpoints#

Ray Train and Ray Tune both provide utilities to upload and track checkpoints through the ray.train.report and ray.tune.report APIs. For details, see the Save and load checkpoints user guide.

If the Ray Train workers report checkpoints, you don’t need to save another Ray Tune checkpoint at the Train driver level, because the driver doesn’t hold any extra training state. The Ray Train driver process already snapshots its status periodically to the configured storage_path. For details, see Job driver fault tolerance.

To access the checkpoints from the Tuner output, append the checkpoint path as a metric. The provided TuneReportCallback does this. It propagates reported Ray Train results to Ray Tune, where the checkpoint path is attached as a separate metric.

Advanced: Fault tolerance#

To recover when the Ray Tune trials that run the Ray Train driver process crash, enable trial fault tolerance on the Ray Tune side with ray.tune.Tuner(run_config=ray.tune.RunConfig(failure_config)).

Fault tolerance on the Ray Train side is configured and handled separately. For details, see the Handle failures and node preemption user guide.

import tempfile

import ray.tune
import ray.train
import ray.train.torch


def train_fn_per_worker(train_loop_config: dict):
    # [1] Train worker restoration logic.
    checkpoint = ray.train.get_checkpoint()
    if checkpoint:
        with checkpoint.as_directory() as temp_checkpoint_dir:
            # model.load_state_dict(torch.load(...))
            ...

    with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
        # torch.save(...)
        ray.train.report(
            {"loss": 0.1},
            checkpoint=ray.train.Checkpoint.from_directory(temp_checkpoint_dir),
        )


def train_fn_driver(config: dict):
    trainer = ray.train.torch.TorchTrainer(
        train_fn_per_worker,
        run_config=ray.train.RunConfig(
            # [2] Train driver restoration is automatic, as long as
            # the (storage_path, name) remains the same across trial restarts.
            # The easiest way to do this is to attach the trial ID in the name.
            # **Do not include any timestamps or random values in the name.**
            name=f"train-trial_id={ray.tune.get_context().get_trial_id()}",
            # [3] Enable worker-level fault tolerance to gracefully handle
            # Train worker failures.
            failure_config=ray.train.FailureConfig(max_failures=3),
            # (If multi-node, configure S3 / NFS as the storage path.)
            # storage_path="s3://...",
        ),
    )
    trainer.fit()


tuner = ray.tune.Tuner(
    train_fn_driver,
    run_config=ray.tune.RunConfig(
        # [4] Enable trial-level fault tolerance to gracefully handle
        # Train driver process failures.
        failure_config=ray.tune.FailureConfig(max_failures=3)
    ),
)
tuner.fit()

Advanced: Use Ray Tune callbacks#

Pass Ray Tune callbacks into ray.tune.RunConfig(callbacks) at the Tuner level.

If you depend on the behavior of built-in or custom Ray Tune callbacks, run Ray Train as a single-trial Tune run and pass the callbacks to the Tuner.

If any callback depends on reported metrics, pass ray.tune.integration.ray_train.TuneReportCallback to the trainer callbacks. This callback propagates results to the Tuner.

import ray.tune
from ray.tune.integration.ray_train import TuneReportCallback
from ray.tune.logger import TBXLoggerCallback


def train_driver_fn(config: dict):
    trainer = TorchTrainer(
        ...,
        run_config=ray.train.RunConfig(..., callbacks=[TuneReportCallback()])
    )
    trainer.fit()


tuner = ray.tune.Tuner(
    train_driver_fn,
    run_config=ray.tune.RunConfig(callbacks=[TBXLoggerCallback()])
)

Tuner(trainer) API deprecation#

The Tuner(trainer) API takes a Ray Train trainer instance directly. It’s deprecated as of Ray 2.43 and will be removed in a future release.

Motivation#

This API change decouples the responsibilities of Ray Train and Ray Tune, and it makes hyperparameter and run configuration more explicit and flexible.

Migration steps#

To migrate from the Tuner(trainer) API to the function-based pattern, do the following:

  1. If you use Ray 2.50 or earlier, enable the environment variable RAY_TRAIN_V2_ENABLED=1.

  2. Replace Tuner(trainer) with a function-based approach that launches Ray Train inside a Tune trial.

  3. Move your training logic into a driver function that Tune calls with different hyperparameters.

Additional resources#