Get started with distributed training using PyTorch

This tutorial shows you how to convert an existing PyTorch script to use Ray Train.

Learn how to do the following:

1. Configure a model to run distributed and on the correct CPU or GPU device.

2. Configure a data loader to shard data across the workers and place data on the correct CPU or GPU device.

3. Configure a training function to save checkpoints so that training can recover from failures.

4. Configure scaling and CPU or GPU resource requirements for a training job.

5. Launch a distributed training job with TorchTrainer.

Quickstart

For reference, the final code looks similar to the following:

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

def train_func():
    # Your PyTorch training code here.
    ...

scaling_config = ScalingConfig(num_workers=2, use_gpu=True)
trainer = TorchTrainer(train_func, scaling_config=scaling_config)
result = trainer.fit()

• train_func is the Python code that runs on each distributed training worker.

• ScalingConfig defines the number of distributed training workers and whether to use GPUs.

• TorchTrainer launches the distributed training job.

Compare a PyTorch training script with and without Ray Train.

PyTorch with Ray Train

import os
import tempfile

import torch
from torch.nn import CrossEntropyLoss
from torch.optim import Adam
from torch.utils.data import DataLoader
from torchvision.models import resnet18
from torchvision.datasets import FashionMNIST
from torchvision.transforms import ToTensor, Normalize, Compose

import ray.train.torch

def train_func():
    # Model, Loss, Optimizer
    model = resnet18(num_classes=10)
    model.conv1 = torch.nn.Conv2d(
        1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False
    )
    # [1] Prepare model.
    model = ray.train.torch.prepare_model(model)
    # model.to("cuda")  # This is done by `prepare_model`
    criterion = CrossEntropyLoss()
    optimizer = Adam(model.parameters(), lr=0.001)

    # Data
    transform = Compose([ToTensor(), Normalize((0.28604,), (0.32025,))])
    data_dir = os.path.join(tempfile.gettempdir(), "data")
    train_data = FashionMNIST(root=data_dir, train=True, download=True, transform=transform)
    train_loader = DataLoader(train_data, batch_size=128, shuffle=True)
    # [2] Prepare dataloader.
    train_loader = ray.train.torch.prepare_data_loader(train_loader)

    # Training
    for epoch in range(10):
        if ray.train.get_context().get_world_size() > 1:
            train_loader.sampler.set_epoch(epoch)

        for images, labels in train_loader:
            # This is done by `prepare_data_loader`!
            # images, labels = images.to("cuda"), labels.to("cuda")
            outputs = model(images)
            loss = criterion(outputs, labels)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        # [3] Report a checkpoint so that training can resume after a failure.
        metrics = {"loss": loss.item(), "epoch": epoch}
        with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
            torch.save(
                model.module.state_dict(),
                os.path.join(temp_checkpoint_dir, "model.pt")
            )
            ray.train.report(
                metrics,
                checkpoint=ray.train.Checkpoint.from_directory(temp_checkpoint_dir),
            )
        if ray.train.get_context().get_world_rank() == 0:
            print(metrics)

# [4] Configure scaling and resource requirements.
scaling_config = ray.train.ScalingConfig(num_workers=2, use_gpu=True)

# [5] Launch distributed training job.
trainer = ray.train.torch.TorchTrainer(
    train_func,
    scaling_config=scaling_config,
    # [5a] If running in a multi-node cluster, this is where you
    # should configure the run's persistent storage that is accessible
    # across all worker nodes.
    # run_config=ray.train.RunConfig(storage_path="s3://..."),
)
result = trainer.fit()

# [6] Load the trained model.
with result.checkpoint.as_directory() as checkpoint_dir:
    model_state_dict = torch.load(os.path.join(checkpoint_dir, "model.pt"))
    model = resnet18(num_classes=10)
    model.conv1 = torch.nn.Conv2d(
        1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False
    )
    model.load_state_dict(model_state_dict)

PyTorch

import os
import tempfile

import torch
from torch.nn import CrossEntropyLoss
from torch.optim import Adam
from torch.utils.data import DataLoader
from torchvision.models import resnet18
from torchvision.datasets import FashionMNIST
from torchvision.transforms import ToTensor, Normalize, Compose

# Model, Loss, Optimizer
model = resnet18(num_classes=10)
model.conv1 = torch.nn.Conv2d(
    1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False
)
model.to("cuda")
criterion = CrossEntropyLoss()
optimizer = Adam(model.parameters(), lr=0.001)

# Data
transform = Compose([ToTensor(), Normalize((0.28604,), (0.32025,))])
train_data = FashionMNIST(root='./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_data, batch_size=128, shuffle=True)

# Training
for epoch in range(10):
    for images, labels in train_loader:
        images, labels = images.to("cuda"), labels.to("cuda")
        outputs = model(images)
        loss = criterion(outputs, labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    metrics = {"loss": loss.item(), "epoch": epoch}
    checkpoint_dir = tempfile.mkdtemp()
    checkpoint_path = os.path.join(checkpoint_dir, "model.pt")
    torch.save(model.state_dict(), checkpoint_path)
    print(metrics)

Set up a training function

First, update your training code to support distributed training. Wrap your code in a training function:

def train_func():
    # Your model training code here.
    ...

Each distributed training worker executes this function.

You can also pass train_func a dictionary as its input argument through the train_loop_config parameter of TorchTrainer. The following example passes a learning rate and a number of epochs:

def train_func(config):
    lr = config["lr"]
    num_epochs = config["num_epochs"]

config = {"lr": 1e-4, "num_epochs": 10}
trainer = ray.train.torch.TorchTrainer(train_func, train_loop_config=config, ...)

Warning

To reduce serialization and deserialization overhead, avoid passing large data objects through train_loop_config. Instead, initialize large objects, such as datasets and models, directly in train_func.

 def load_dataset():
     # Return a large in-memory dataset
     ...

 def load_model():
     # Return a large in-memory model instance
     ...

-config = {"data": load_dataset(), "model": load_model()}

 def train_func(config):
-    data = config["data"]
-    model = config["model"]

+    data = load_dataset()
+    model = load_model()
     ...

 trainer = ray.train.torch.TorchTrainer(train_func, train_loop_config=config, ...)

Set up a model

Use the ray.train.torch.prepare_model() utility function to do the following:

• Move your model to the correct device.

• Wrap it in DistributedDataParallel.

-from torch.nn.parallel import DistributedDataParallel
+import ray.train.torch

 def train_func():

     ...

     # Create model.
     model = ...

     # Set up distributed training and device placement.
-    device_id = ... # Your logic to get the right device.
-    model = model.to(device_id or "cpu")
-    model = DistributedDataParallel(model, device_ids=[device_id])
+    model = ray.train.torch.prepare_model(model)

     ...

Set up a dataset

Use the ray.train.torch.prepare_data_loader() utility function, which does the following:

• Adds a DistributedSampler to your DataLoader.

• Moves the batches to the right device.

You don’t need this step if you pass a Ray Data dataset to your trainer. See Data loading and preprocessing.

 from torch.utils.data import DataLoader
+import ray.train.torch

 def train_func():

     ...

     dataset = ...

     data_loader = DataLoader(dataset, batch_size=worker_batch_size, shuffle=True)
+    data_loader = ray.train.torch.prepare_data_loader(data_loader)

     for epoch in range(10):
+        if ray.train.get_context().get_world_size() > 1:
+            data_loader.sampler.set_epoch(epoch)

         for X, y in data_loader:
-            X = X.to_device(device)
-            y = y.to_device(device)

     ...

Tip

The batch_size that DataLoader takes is the batch size for each worker. To convert between the worker batch size and the global batch size, use the following equation:

global_batch_size = worker_batch_size * ray.train.get_context().get_world_size()

Note

If you already set up your DataLoader with a DistributedSampler, prepare_data_loader() doesn’t add another one and respects the configuration of the existing sampler.

Note

DistributedSampler doesn’t work with a DataLoader that wraps IterableDataset. To work with a dataset iterator, consider using Ray Data instead of a PyTorch DataLoader. Ray Data provides performant streaming data ingestion for large-scale datasets.

For details, see Data loading and preprocessing.

Save checkpoints for fault tolerance

To recover from worker and node failures without losing progress, save checkpoints with the ray.train.report() utility function. Ray Train uploads each checkpoint to persistent storage. When you enable worker fault tolerance, Ray Train restarts failed workers, and your training function can resume from the latest checkpoint with ray.train.get_checkpoint().

Ray Train attaches the metrics that you pass to ray.train.report to the checkpoint and uses them to decide which checkpoints to keep. To track metrics over the course of training, log them to an experiment tracking tool. See Experiment tracking.

+import os
+import tempfile

+import ray.train

 def train_func():

     ...

     with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
        torch.save(
            model.state_dict(), os.path.join(temp_checkpoint_dir, "model.pt")
        )

+       metrics = {"loss": loss.item()}  # Metrics to attach to the checkpoint.

        # Build a Ray Train checkpoint from a directory
+       checkpoint = ray.train.Checkpoint.from_directory(temp_checkpoint_dir)

        # Ray Train will automatically save the checkpoint to persistent storage,
        # so the local `temp_checkpoint_dir` can be safely cleaned up after.
+       ray.train.report(metrics=metrics, checkpoint=checkpoint)

     ...

For details, see Save and load checkpoints and Handle failures and node preemption.

Configure scale and GPUs

Outside your training function, create a ScalingConfig object that sets the following two parameters:

• num_workers: The number of distributed training worker processes.

• use_gpu: Whether each worker uses a GPU or a CPU.

from ray.train import ScalingConfig
scaling_config = ScalingConfig(num_workers=2, use_gpu=True)

For details, see Configure scale and accelerators.

Configure persistent storage

Create a RunConfig object to specify the path where Ray Train saves results, including checkpoints and artifacts.

from ray.train import RunConfig

# Local path (/some/local/path/unique_run_name)
run_config = RunConfig(storage_path="/some/local/path", name="unique_run_name")

# Shared cloud storage URI (s3://bucket/unique_run_name)
run_config = RunConfig(storage_path="s3://bucket", name="unique_run_name")

# Shared NFS path (/mnt/nfs/unique_run_name)
run_config = RunConfig(storage_path="/mnt/nfs", name="unique_run_name")

Warning

A shared storage location, such as cloud storage or NFS, is optional for single-node clusters but required for multi-node clusters. On a multi-node cluster, a local path raises an error during checkpointing.

For details, see Configure persistent storage.

Launch a training job

To launch a distributed training job, pass the training function, scaling configuration, and run configuration to a TorchTrainer and call fit().

from ray.train.torch import TorchTrainer

trainer = TorchTrainer(
    train_func, scaling_config=scaling_config, run_config=run_config
)
result = trainer.fit()

Access training results

After training completes, trainer.fit() returns a Result object with information about the training run, including the metrics and checkpoints reported during training.

result.metrics     # The metrics reported during training.
result.checkpoint  # The latest checkpoint reported during training.
result.path        # The path where logs are stored.
result.error       # The exception that was raised, if training failed.

For more usage examples, see Inspect training results.

Next steps

After you convert your PyTorch training script to use Ray Train, explore the following resources:

• See the user guides to learn how to perform specific tasks.

• Browse the examples for end-to-end examples of Ray Train.

• See the API reference for details on the classes and methods that this tutorial uses.