Get started with distributed training using PyTorch Lightning#

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

You learn how to do the following:

Configure the Lightning Trainer to run distributed with Ray on the correct CPU or GPU device.
Configure the training function to save checkpoints so that training can recover from failures.
Configure scaling and CPU or GPU resource requirements for a training job.
Launch a distributed training job with a TorchTrainer.

Quickstart#

Here’s the final code for reference:

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

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

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

The code has three parts:

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 Lightning training script with and without Ray Train.

PyTorch Lightning with Ray Train

import os
import tempfile

import torch
from torch.utils.data import DataLoader
from torchvision.models import resnet18
from torchvision.datasets import FashionMNIST
from torchvision.transforms import ToTensor, Normalize, Compose
import lightning.pytorch as pl

import ray.train.lightning
from ray.train.torch import TorchTrainer

# Model, Loss, Optimizer
class ImageClassifier(pl.LightningModule):
    def __init__(self):
        super(ImageClassifier, self).__init__()
        self.model = resnet18(num_classes=10)
        self.model.conv1 = torch.nn.Conv2d(
            1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False
        )
        self.criterion = torch.nn.CrossEntropyLoss()

    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        x, y = batch
        outputs = self.forward(x)
        loss = self.criterion(outputs, y)
        self.log("loss", loss, on_step=True, prog_bar=True)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.model.parameters(), lr=0.001)


def train_func():
    # 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_dataloader = DataLoader(train_data, batch_size=128, shuffle=True)

    # Training
    model = ImageClassifier()
    # [1] Configure PyTorch Lightning Trainer.
    trainer = pl.Trainer(
        max_epochs=10,
        devices="auto",
        accelerator="auto",
        strategy=ray.train.lightning.RayDDPStrategy(),
        plugins=[ray.train.lightning.RayLightningEnvironment()],
        callbacks=[ray.train.lightning.RayTrainReportCallback()],
        # [1a] Optionally, disable the default checkpointing behavior
        # in favor of the `RayTrainReportCallback` above.
        enable_checkpointing=False,
    )
    trainer = ray.train.lightning.prepare_trainer(trainer)
    trainer.fit(model, train_dataloaders=train_dataloader)

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

# [3] Launch distributed training job.
trainer = TorchTrainer(
    train_func,
    scaling_config=scaling_config,
    # [3a] 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: ray.train.Result = trainer.fit()

# [4] Load the trained model.
with result.checkpoint.as_directory() as checkpoint_dir:
    model = ImageClassifier.load_from_checkpoint(
        os.path.join(
            checkpoint_dir,
            ray.train.lightning.RayTrainReportCallback.CHECKPOINT_NAME,
        ),
    )

PyTorch Lightning

import torch
from torchvision.models import resnet18
from torchvision.datasets import FashionMNIST
from torchvision.transforms import ToTensor, Normalize, Compose
from torch.utils.data import DataLoader
import lightning.pytorch as pl

# Model, Loss, Optimizer
class ImageClassifier(pl.LightningModule):
    def __init__(self):
        super(ImageClassifier, self).__init__()
        self.model = resnet18(num_classes=10)
        self.model.conv1 = torch.nn.Conv2d(
            1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False
        )
        self.criterion = torch.nn.CrossEntropyLoss()

    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        x, y = batch
        outputs = self.forward(x)
        loss = self.criterion(outputs, y)
        self.log("loss", loss, on_step=True, prog_bar=True)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.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_dataloader = DataLoader(train_data, batch_size=128, shuffle=True)

# Training
model = ImageClassifier()
trainer = pl.Trainer(max_epochs=10)
trainer.fit(model, train_dataloaders=train_dataloader)

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, ...)

Ray Train sets up your distributed process group on each worker. You only need to make a few changes to your Lightning Trainer definition.

 import lightning.pytorch as pl
-from pl.strategies import DDPStrategy
-from pl.plugins.environments import LightningEnvironment
+import ray.train.lightning

 def train_func():
     ...
     model = MyLightningModule(...)
     datamodule = MyLightningDataModule(...)

     trainer = pl.Trainer(
-        devices=[0, 1, 2, 3],
-        strategy=DDPStrategy(),
-        plugins=[LightningEnvironment()],
+        devices="auto",
+        accelerator="auto",
+        strategy=ray.train.lightning.RayDDPStrategy(),
+        plugins=[ray.train.lightning.RayLightningEnvironment()]
     )
+    trainer = ray.train.lightning.prepare_trainer(trainer)

     trainer.fit(model, datamodule=datamodule)

The following sections discuss each change.

Configure the distributed strategy#

Ray Train offers several distributed strategies for Lightning, each a subclass of a Lightning base strategy. They accept the same arguments as their base strategy classes, and internally they configure the root device and the distributed sampler arguments. RayDDPStrategy disables process_group_backend, timeout, and start_method. Set them in TorchConfig instead. Choose from the following strategies:

 import lightning.pytorch as pl
-from pl.strategies import DDPStrategy
+import ray.train.lightning

 def train_func():
     ...
     trainer = pl.Trainer(
         ...
-        strategy=DDPStrategy(),
+        strategy=ray.train.lightning.RayDDPStrategy(),
         ...
     )
     ...

Configure the Ray cluster environment plugin#

Ray Train also provides a RayLightningEnvironment class as a specification for the Ray cluster. This utility class configures the worker’s local, global, and node rank and world size.

 import lightning.pytorch as pl
-from pl.plugins.environments import LightningEnvironment
+import ray.train.lightning

 def train_func():
     ...
     trainer = pl.Trainer(
         ...
-        plugins=[LightningEnvironment()],
+        plugins=[ray.train.lightning.RayLightningEnvironment()],
         ...
     )
     ...

Configure parallel devices#

TorchTrainer already configures the correct CUDA_VISIBLE_DEVICES for you. Always use all available GPUs by setting devices="auto" and accelerator="auto".

 import lightning.pytorch as pl

 def train_func():
     ...
     trainer = pl.Trainer(
         ...
-        devices=[0,1,2,3],
+        devices="auto",
+        accelerator="auto",
         ...
     )
     ...

Save checkpoints for fault tolerance#

To persist checkpoints so that training can recover from failures, add a ray.train.lightning.RayTrainReportCallback utility callback to your Lightning Trainer.

 import lightning.pytorch as pl
 from ray.train.lightning import RayTrainReportCallback

 def train_func():
     ...
     trainer = pl.Trainer(
         ...
-        callbacks=[...],
+        callbacks=[..., RayTrainReportCallback()],
     )
     ...

Report metrics and checkpoints to Ray Train to support fault-tolerant training and hyperparameter optimization.

You can also configure asynchronous checkpointing and asynchronous validation through the callback. The checkpoint_upload_mode flag offloads checkpoint uploading to a Ray Train-managed background thread instead of blocking the Lightning training loop. The validation flag launches an asynchronous Ray task to validate the checkpoint instead of running validation_step synchronously in the training workers. This is incompatible with Lightning’s AsyncCheckpointIO plugin, because Ray Train needs to control the upload thread so it can wait for the upload to finish before committing the checkpoint.

 import lightning.pytorch as pl
 from ray.train.lightning import RayTrainReportCallback
+from ray.train import CheckpointUploadMode
+from ray.train import ValidationConfig, ValidationTaskConfig

 def train_func():
     ...
     trainer = pl.Trainer(
         ...
-        callbacks=[..., RayTrainReportCallback()],
+        callbacks=[..., RayTrainReportCallback(
+            checkpoint_upload_mode=CheckpointUploadMode.ASYNC,
+            validation=ValidationTaskConfig(fn_kwargs={}),
+        )],
     )
     ...

+def validation_fn(checkpoint):
+    # Load and validate the checkpoint, return metrics.
+    return {"val_score": ...}
+
+trainer = TorchTrainer(
+    train_func,
+    validation_config=ValidationConfig(fn=validation_fn),
+    ...
+)

The ray.train.lightning.RayTrainReportCallback class provides only a simple implementation. You can customize it further.

Prepare your Lightning `Trainer`#

Finally, pass your Lightning Trainer to prepare_trainer() to validate your configurations.

 import lightning.pytorch as pl
 import ray.train.lightning

 def train_func():
     ...
     trainer = pl.Trainer(...)
+    trainer = ray.train.lightning.prepare_trainer(trainer)
     ...

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 Lightning 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 how to use Ray Train.
See the API reference for details on the classes and methods in this tutorial.

Version compatibility#

The Ray project tests Ray Train with pytorch_lightning versions 1.8.6 and 2.6.1. For full compatibility, use pytorch_lightning>=1.6.5. Earlier versions aren’t prohibited but might cause unexpected issues. If you run into compatibility issues, consider upgrading your PyTorch Lightning version or file an issue.

Note

If you’re using Lightning 2.x, use the import path lightning.pytorch.xxx instead of pytorch_lightning.xxx.

`LightningTrainer` migration guide#

Ray 2.4 introduced LightningTrainer, which exposed a LightningConfigBuilder for defining configurations for pl.LightningModule and pl.Trainer. LightningTrainer then instantiated the model and trainer objects and ran a predefined training function in a black box. This design limited your control over the training code.

Ray 2.7 introduced the unified TorchTrainer API, which aligns more closely with standard PyTorch Lightning scripts and gives you more control over your native Lightning code.

Deprecated LightningTrainer

from ray.train.lightning import LightningConfigBuilder, LightningTrainer

config_builder = LightningConfigBuilder()
# [1] Collect model configs
config_builder.module(cls=MyLightningModule, lr=1e-3, feature_dim=128)

# [2] Collect checkpointing configs
config_builder.checkpointing(monitor="val_accuracy", mode="max", save_top_k=3)

# [3] Collect pl.Trainer configs
config_builder.trainer(
    max_epochs=10,
    accelerator="gpu",
    log_every_n_steps=100,
)

# [4] Build datasets on the head node
datamodule = MyLightningDataModule(batch_size=32)
config_builder.fit_params(datamodule=datamodule)

# [5] Execute the internal training function in a black box
ray_trainer = LightningTrainer(
    lightning_config=config_builder.build(),
    scaling_config=ScalingConfig(num_workers=4, use_gpu=True),
    run_config=RunConfig(
        checkpoint_config=CheckpointConfig(
            num_to_keep=3,
            checkpoint_score_attribute="val_accuracy",
            checkpoint_score_order="max",
        ),
    )
)
result = ray_trainer.fit()

# [6] Load the trained model from an opaque Lightning-specific checkpoint.
lightning_checkpoint = result.checkpoint
model = lightning_checkpoint.get_model(MyLightningModule)

TorchTrainer

import os

import lightning.pytorch as pl

import ray.train
from ray.train.torch import TorchTrainer
from ray.train.lightning import (
    RayDDPStrategy,
    RayLightningEnvironment,
    RayTrainReportCallback,
    prepare_trainer
)

def train_func():
    # [1] Create a Lightning model
    model = MyLightningModule(lr=1e-3, feature_dim=128)

    # [2] Report Checkpoint with callback
    ckpt_report_callback = RayTrainReportCallback()

    # [3] Create a Lightning Trainer
    trainer = pl.Trainer(
        max_epochs=10,
        log_every_n_steps=100,
        # New configurations below
        devices="auto",
        accelerator="auto",
        strategy=RayDDPStrategy(),
        plugins=[RayLightningEnvironment()],
        callbacks=[ckpt_report_callback],
    )

    # Validate your Lightning trainer configuration
    trainer = prepare_trainer(trainer)

    # [4] Build your datasets on each worker
    datamodule = MyLightningDataModule(batch_size=32)
    trainer.fit(model, datamodule=datamodule)

# [5] Explicitly define and run the training function
ray_trainer = TorchTrainer(
    train_func,
    scaling_config=ray.train.ScalingConfig(num_workers=4, use_gpu=True),
    run_config=ray.train.RunConfig(
        checkpoint_config=ray.train.CheckpointConfig(
            num_to_keep=3,
            checkpoint_score_attribute="val_accuracy",
            checkpoint_score_order="max",
        ),
    )
)
result = ray_trainer.fit()

# [6] Load the trained model from a simplified checkpoint interface.
checkpoint: ray.train.Checkpoint = result.checkpoint
with checkpoint.as_directory() as checkpoint_dir:
    print("Checkpoint contents:", os.listdir(checkpoint_dir))
    checkpoint_path = os.path.join(checkpoint_dir, "checkpoint.ckpt")
    model = MyLightningModule.load_from_checkpoint(checkpoint_path)