Ray Train & Ray Tune: Distributed Training and Hyperparameter Tuning

Most machine learning projects start on a single machine, and at some point the training loop either runs out of memory, takes too long on one GPU, or needs hundreds of hyperparameter trials that a single process cannot finish in time. Ray Train and Ray Tune address those two problems directly: the first scales an existing training loop across multiple workers and GPUs, and the second runs many training runs in parallel to search a hyperparameter space. This tutorial walks through both libraries with a coherent PyTorch example, then shows how to combine them so you can tune the hyperparameters of a distributed training job without rewriting your code.

A Short Introduction to Ray and the AI Libraries

Ray is a framework for distributed Python. At its core it provides tasks and actors that run across a cluster, plus a scheduler that places work on available CPUs and GPUs. On top of that core, the Ray AI Libraries provide higher-level building blocks for common machine learning needs:

• Ray Data for distributed data loading and preprocessing.
• Ray Train for distributed model training.
• Ray Tune for hyperparameter search and experiment orchestration.
• Ray Serve for online model serving (covered in a separate tutorial, so we only mention it here for context).

The important property is that the same code runs on your laptop and on a multi-node cluster. You develop locally, then point the same script at a cluster when you need more compute. This tutorial focuses on Ray Train and Ray Tune, which together cover the training and tuning stages of a typical workflow.

Installation

Install Ray with the Train and Tune extras, plus PyTorch:

pip install "ray[train,tune]" torch torchvision

Verify the installation:

python -c "import ray; from ray import train, tune; print(ray.version)"

The examples below were written against Ray 2.x. The API for Ray Train was reworked in the 2.x series, so make sure you are not on a pre-2.7 release where the imports differ.

Ray Train: From Single-Node to Distributed

The starting point: a plain PyTorch loop

Consider a typical single-process training loop for an image classifier on FashionMNIST. There is nothing distributed about it yet.

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms


def buildmodel():

    return nn.Sequential(
        nn.Flatten(),
        nn.Linear(28  28, 256),

        nn.ReLU(),
        nn.Linear(256, 128),
        nn.ReLU(),
        nn.Linear(128, 10),
    )


def trainsinglenode(epochs=5, lr=1e-3, batchsize=64):
    transform = transforms.ToTensor()
    traindata = datasets.FashionMNIST(

        root="./data", train=True, download=True, transform=transform
    )
    loader = DataLoader(traindata, batchsize=batchsize, shuffle=True)

    device = "cuda" if torch.cuda.isavailable() else "cpu"

    model = buildmodel().to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    lossfn = nn.CrossEntropyLoss()


    for epoch in range(epochs):
        model.train()
        for images, labels in loader:
            images, labels = images.to(device), labels.to(device)
            optimizer.zerograd()
            loss = lossfn(model(images), labels)

            loss.backward()
            optimizer.step()
        print(f"epoch {epoch}: last batch loss {loss.item():.4f}")

This works, but it uses one process and one device. To train across several GPUs or several nodes with data-parallelism, we adapt the loop for Ray Train.

Converting the loop with TorchTrainer

Ray Train wraps your training function and runs one copy of it per worker. Each worker handles a shard of the data and synchronizes gradients automatically. The changes are small and localized:

Wrap the model with ray.train.torch.preparemodel, which moves it to the correct device and wraps it in DistributedDataParallel.

Wrap the data loader with ray.train.torch.preparedataloader, which inserts a DistributedSampler so each worker sees a different shard.

Report metrics and checkpoints with ray.train.report instead of plain print.

import os
import tempfile

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

import ray.train.torch
from ray.train import Checkpoint, ScalingConfig
from ray.train.torch import TorchTrainer


def buildmodel():

    return nn.Sequential(
        nn.Flatten(),
        nn.Linear(28  28, 256),
        nn.ReLU(),
        nn.Linear(256, 128),
        nn.ReLU(),
        nn.Linear(128, 10),
    )


def trainfunc(config):
    lr = config["lr"]
    epochs = config["epochs"]
    batchsize = config["batchsize"]

    transform = transforms.ToTensor()
    traindata = datasets.FashionMNIST(

        root="./data", train=True, download=True, transform=transform
    )
    loader = DataLoader(traindata, batchsize=batchsize, shuffle=True)
    # Inserts a DistributedSampler and moves batches to the worker device.
    loader = ray.train.torch.preparedataloader(loader)

    model = buildmodel()

    # Moves the model to the device and wraps it in DistributedDataParallel.
    model = ray.train.torch.preparemodel(model)

    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    lossfn = nn.CrossEntropyLoss()


    for epoch in range(epochs):
        # With a DistributedSampler the epoch must be set so shuffling differs.
        if ray.train.getcontext().getworldsize() > 1:
            loader.sampler.setepoch(epoch)


        model.train()
        totalloss, nbatches = 0.0, 0

        for images, labels in loader:
            optimizer.zerograd()
            loss = lossfn(model(images), labels)

            loss.backward()
            optimizer.step()
            totalloss += loss.item()
            nbatches += 1


        avgloss = totalloss / max(nbatches, 1)

        # Save a checkpoint and report metrics back to Ray Train.
        with tempfile.TemporaryDirectory() as tmpdir:

            torch.save(
                model.module.statedict(),
                os.path.join(tmpdir, "model.pt"),

            )
            checkpoint = Checkpoint.fromdirectory(tmpdir)

            ray.train.report(
                {"loss": avgloss, "epoch": epoch},
                checkpoint=checkpoint,
            )

A few details worth noting. We call model.module.statedict() because preparemodel wraps the model in DistributedDataParallel, and the real weights live under .module. Reporting and checkpointing happen on every worker, but Ray Train coordinates so only one rank persists the checkpoint by default.

Configuring scale with ScalingConfig

ScalingConfig declares how many workers to launch and whether each one needs a GPU. The training function does not change when you scale up; only this configuration does.

scalingconfig = ScalingConfig(
    numworkers=4,     # four data-parallel workers
    usegpu=True,      # one GPU per worker

    # resourcesperworker={"CPU": 2},  # optional: extra CPU per worker

)

trainer = TorchTrainer(
    trainloopperworker=trainfunc,

    trainloopconfig={"lr": 1e-3, "epochs": 5, "batchsize": 64},
    scalingconfig=scalingconfig,
)

result = trainer.fit()

To run on CPU only (for local development), set usegpu=False and pick a small numworkers such as 2. The effective batch size is batchsize * numworkers because each worker processes its own shard, so you may want to scale the learning rate accordingly.

Accessing results and the best checkpoint

trainer.fit() returns a Result object that holds the final metrics and the checkpoint history.

print("Final reported metrics:", result.metrics)

The latest checkpoint.
bestcheckpoint = result.checkpoint


Or select by a metric across reported checkpoints.
bestcheckpoint = result.getbestcheckpoint(metric="loss", mode="min")

Load weights back into a fresh model.
with bestcheckpoint.asdirectory() as ckptdir:

    statedict = torch.load(os.path.join(ckptdir, "model.pt"))

    model = buildmodel()
    model.loadstatedict(statedict)

The same pattern for other frameworks

The data-parallel pattern is not specific to raw PyTorch. Ray Train ships trainers that follow the same structure for other ecosystems, so the surrounding ScalingConfig, reporting, and result handling stay identical:

• PyTorch Lightning through ray.train.lightning utilities that integrate with the Trainer.
• Hugging Face Transformers by wrapping the transformers.Trainer with the Ray Train reporting callback.
• XGBoost and LightGBM through XGBoostTrainer and LightGBMTrainer, which run distributed gradient boosting.

In each case you keep the framework's own training code and let Ray Train handle worker orchestration and resource placement.

Ray Tune: Scalable Hyperparameter Search

Ray Tune runs many training configurations in parallel and tracks their results. A single configuration is called a trial. You give Tune a search space, a trainable, and a budget, and it manages scheduling, early stopping, and bookkeeping.

Defining a search space

The search space is a dictionary of distributions from ray.tune:

from ray import tune

searchspace = {
    "lr": tune.loguniform(1e-4, 1e-1),
    "batchsize": tune.choice([32, 64, 128]),

    "epochs": 5,  # a fixed value is fine; not everything must be searched
}

Common space functions include tune.uniform, tune.loguniform, tune.choice, tune.randint, and tune.gridsearch([...]) for exhaustive grids.

A simple trainable

A trainable is a function that takes a config, trains, and reports a metric. For Tune-only experiments (no distributed training) the function reports with ray.train.report just like before:

import ray.train


def tunetrainable(config):

    # ... build model, loader, optimizer using config["lr"], etc. ...
    for epoch in range(config["epochs"]):
        avgloss = trainoneepoch(config)  # your loop
        accuracy = evaluate(config)
        ray.train.report({"loss": avgloss, "accuracy": accuracy})

Running the Tuner

tune.Tuner ties the trainable, the search space, and a TuneConfig together. TuneConfig sets the number of trials (numsamples), the target metric, and the optimization direction (mode).

from ray import tune
from ray.tune import TuneConfig

tuner = tune.Tuner(
    tunetrainable,

    paramspace=searchspace,

    tuneconfig=TuneConfig(
        numsamples=20,        # twenty trials

        metric="accuracy",
        mode="max",
    ),
)

results = tuner.fit()

Search algorithms

By default Tune samples the space randomly (or runs a grid if you used tune.gridsearch). For smarter sampling you can plug in a search algorithm that learns from completed trials. Two common choices:

from ray.tune.search.optuna import OptunaSearch
from ray.tune.search.hyperopt import HyperOptSearch

Tree-structured Parzen estimator via Optuna.
optunasearch = OptunaSearch(metric="accuracy", mode="max")


Or HyperOpt's TPE implementation.
hyperoptsearch = HyperOptSearch(metric="accuracy", mode="max")

tuner = tune.Tuner(
    tunetrainable,

    paramspace=searchspace,

    tuneconfig=TuneConfig(
        numsamples=20,

        searchalg=optunasearch,

    ),
)

Note that this is how Ray Tune relates to Optuna: Optuna provides the sampler, while Tune provides parallel execution, scheduling, and resource management across a cluster. If you only need single-process tuning, plain Optuna is simpler; Tune adds value once you want trials to run in parallel.

Schedulers and early stopping

A scheduler decides which trials to continue, pause, or stop. The most widely used is ASHA (Asynchronous Successive Halving), which stops underperforming trials early so the budget concentrates on promising configurations.

from ray.tune.schedulers import ASHAScheduler

scheduler = ASHAScheduler(
    metric="accuracy",
    mode="max",
    maxt=10,            # max epochs (or training iterations) per trial
    graceperiod=2,      # run at least two iterations before stopping

    reductionfactor=2,
)

tuner = tune.Tuner(
    tunetrainable,

    paramspace=searchspace,

    tuneconfig=TuneConfig(
        numsamples=40,

        scheduler=scheduler,
        metric="accuracy",
        mode="max",
    ),
)

Population Based Training (PopulationBasedTraining) is another scheduler. Instead of only stopping trials, it periodically copies the weights and hyperparameters of strong trials into weak ones and perturbs them, which lets the hyperparameters change during a single run. It is well suited to schedules like learning-rate decay but is heavier to set up, so reach for ASHA first unless you specifically need evolving hyperparameters.

Analyzing results

tuner.fit() returns a ResultGrid. You can pull out the best trial and inspect every trial as a dataframe.

results = tuner.fit()

best = results.getbestresult(metric="accuracy", mode="max")

print("Best config:", best.config)
print("Best accuracy:", best.metrics["accuracy"])

The best trial's checkpoint, if the trainable saved one.
bestcheckpoint = best.checkpoint

All trials as a pandas DataFrame for plotting or export.
df = results.getdataframe()

print(df[["accuracy", "loss", "config/lr", "config/batchsize"]].head())

Combining Ray Train and Ray Tune

The two libraries compose cleanly: you can tune the hyperparameters of a distributed TorchTrainer by wrapping it in a Tuner. Each trial then launches its own group of data-parallel workers. The key is to pass the search space under the trainloopconfig key so Tune knows which part of the trainer to vary.

from ray import tune
from ray.tune import TuneConfig, Tuner
from ray.tune.schedulers import ASHAScheduler
from ray.train import ScalingConfig
from ray.train.torch import TorchTrainer

Reuse the distributed trainfunc from the Ray Train section.

trainer = TorchTrainer(
    trainloopperworker=trainfunc,

    scalingconfig=ScalingConfig(numworkers=2, usegpu=True),
)

paramspace = {

    "trainloopconfig": {

        "lr": tune.loguniform(1e-4, 1e-1),
        "batchsize": tune.choice([32, 64, 128]),
        "epochs": 5,
    }
}

tuner = Tuner(
    trainer,
    paramspace=paramspace,
    tuneconfig=TuneConfig(

        numsamples=16,
        metric="loss",
        mode="min",
        scheduler=ASHAScheduler(maxt=5, graceperiod=1),
    ),
)

results = tuner.fit()
best = results.getbestresult(metric="loss", mode="min")
print("Best distributed config:", best.config["trainloopconfig"])

Be mindful of total resource demand. With numsamples=16 and numworkers=2 and a GPU per worker, a fully parallel run wants 32 GPUs. On a smaller cluster Tune simply queues trials until resources free up, but you should size numsamples and numworkers against the hardware you actually have.

Local Machine vs a Ray Cluster

Everything above runs unchanged on a single machine. Ray starts a local instance automatically the first time you call trainer.fit() or tuner.fit(). To use more than one node, start a cluster and point your driver at it; the training and tuning code does not change.

Starting a cluster manually:

# On the head node.
ray start --head --port=6379

On each worker node, pointing at the head address.
ray start --address='HEADNODEIP:6379'

Then in your script, connect to the running cluster instead of starting a local one:

import ray

ray.init(address="auto")  # joins the cluster started above

For cloud environments, Ray's cluster launcher reads a YAML file describing the head and worker node types and autoscaling limits, and brings the cluster up for you:

ray up cluster.yaml
ray submit cluster.yaml trainscript.py

The point is that ScalingConfig and your training function stay the same. You move from a laptop to a 50-node cluster by changing where the script connects and how many workers you request, not by rewriting the loop.

Experiment Tracking Integration

Ray Tune integrates with experiment trackers through callbacks, so trial metrics flow into your usual dashboards. Two common options:

from ray.air.integrations.wandb import WandbLoggerCallback
from ray.air.integrations.mlflow import MLflowLoggerCallback
from ray.tune import RunConfig

runconfig = RunConfig(
    callbacks=[
        WandbLoggerCallback(project="ray-train-demo"),
        # or:
        MLflowLoggerCallback(
            trackinguri="http://localhost:5000",

            experimentname="ray-train-demo",
        ),
    ]
)

tuner = tune.Tuner(
    tunetrainable,

    paramspace=searchspace,

    tuneconfig=TuneConfig(numsamples=20, metric="accuracy", mode="max"),

    runconfig=runconfig,

)

Every value you pass to ray.train.report is logged automatically, and each trial appears as a separate run in Weights & Biases or MLflow.

Fault Tolerance and Resuming

Long jobs fail for ordinary reasons: a spot instance is reclaimed, a node runs out of memory, or the driver is interrupted. Ray Train and Ray Tune persist checkpoints and experiment state so you can resume rather than restart.

Configure where state is stored and how many failures to retry through RunConfig and FailureConfig:

from ray.train import RunConfig, FailureConfig

runconfig = RunConfig(
    storagepath="s3://my-bucket/ray-experiments",  # or a shared local path

    name="fashionmnist-tune",
    failureconfig=FailureConfig(maxfailures=2),   # retry a failed trial twice

)

If the whole experiment is interrupted, restore it from the same storage path and continue:

from ray.tune import Tuner

tuner = Tuner.restore(
    path="s3://my-bucket/ray-experiments/fashionmnist-tune",
    trainable=tunetrainable,
    resumeerrored=True,

)
results = tuner.fit()

For this to work the storage path must be reachable from every node, which is why a shared filesystem or object store is recommended on multi-node clusters rather than a local directory on the head node.

Best Practices

• Develop on CPU first. Set usegpu=False and numworkers=2 to validate the loop quickly, then scale up. The distributed code path runs the same way at small scale.
• Report at a sensible cadence. Call ray.train.report once per epoch or every N steps, not every batch, to keep coordination overhead low.
• Account for the effective batch size. Data-parallel training multiplies the per-worker batch size by the number of workers. Scale the learning rate or warmup accordingly.
• Let ASHA cut the budget. Early stopping with ASHAScheduler usually finds good configurations with far less compute than running every trial to completion.
• Search the right things. Fix parameters that do not matter and search only the ones that do; smaller spaces converge faster.
• Use a shared storage path. On a cluster, point storagepath at an object store or shared filesystem so checkpoints and restore work across nodes.
• Size trials against hardware. Multiply numsamples by per-trial workers to estimate peak resource demand before launching.

Conclusion and Key Takeaways

Ray Train and Ray Tune cover the two scaling problems that single-machine training eventually hits. Ray Train turns a standard PyTorch loop into distributed data-parallel training with three small edits (preparemodel, preparedataloader, and ray.train.report), while ScalingConfig controls the number of workers and GPUs without touching the loop. Ray Tune runs many configurations in parallel, supports smarter search through OptunaSearch or HyperOptSearch, and saves compute with schedulers like ASHA. Because the two compose, you can tune a distributed trainer simply by wrapping it in a Tuner.

Key takeaways:

• Convert a training loop with preparemodel, preparedataloader, and ray.train.report; scale it with ScalingConfig.
• Use TorchTrainer.fit() results to retrieve metrics and the best checkpoint.
• Define a search space, run tune.Tuner with TuneConfig, and analyze the ResultGrid with getbestresult.
• Add OptunaSearch/HyperOptSearch for sample efficiency and ASHAScheduler for early stopping.
• Wrap a TorchTrainer in a Tuner to tune distributed training; the same code runs locally or on a cluster.
• Configure storagepath and FailureConfig so long runs can resume after failures.