Get started with distributed training using Horovod

Warning

HorovodTrainer is deprecated and scheduled for removal. Ray Train V2, the default, doesn’t support it, and constructing a HorovodTrainer raises a DeprecationWarning. Use TorchTrainer or TensorflowTrainer instead. To keep running existing Horovod code, set RAY_TRAIN_V2_ENABLED=0 to fall back to the previous implementation.

Ray Train configures the Horovod environment and rendezvous server for you, so you can run your DistributedOptimizer training script. For more information, see the Horovod documentation.

Quickstart

# TODO: [V2] Deprecated doc code to delete.
import os

os.environ["RAY_TRAIN_V2_ENABLED"] = "0"

import tempfile

import horovod.torch as hvd
import ray
from ray import train
from ray.train import Checkpoint, ScalingConfig
import ray.train.torch  # Need this to use `train.torch.get_device()`
from ray.train.horovod import HorovodTrainer
import torch
import torch.nn as nn

# If using GPUs, set this to True.
use_gpu = False


input_size = 1
layer_size = 15
output_size = 1
num_epochs = 3


class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.layer1 = nn.Linear(input_size, layer_size)
        self.relu = nn.ReLU()
        self.layer2 = nn.Linear(layer_size, output_size)

    def forward(self, input):
        return self.layer2(self.relu(self.layer1(input)))


def train_loop_per_worker():
    hvd.init()
    dataset_shard = train.get_dataset_shard("train")
    model = NeuralNetwork()
    device = train.torch.get_device()
    model.to(device)
    loss_fn = nn.MSELoss()
    lr_scaler = 1
    optimizer = torch.optim.SGD(model.parameters(), lr=0.1 * lr_scaler)
    # Horovod: wrap optimizer with DistributedOptimizer.
    optimizer = hvd.DistributedOptimizer(
        optimizer,
        named_parameters=model.named_parameters(),
        op=hvd.Average,
    )
    for epoch in range(num_epochs):
        model.train()
        for batch in dataset_shard.iter_torch_batches(
            batch_size=32, dtypes=torch.float
        ):
            inputs, labels = torch.unsqueeze(batch["x"], 1), batch["y"]
            outputs = model(inputs)
            loss = loss_fn(outputs, labels)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            print(f"epoch: {epoch}, loss: {loss.item()}")

        with tempfile.TemporaryDirectory() as tmpdir:
            torch.save(model.state_dict(), os.path.join(tmpdir, "model.pt"))
            train.report(
                {"loss": loss.item()}, checkpoint=Checkpoint.from_directory(tmpdir)
            )


train_dataset = ray.data.from_items([{"x": x, "y": x + 1} for x in range(32)])
scaling_config = ScalingConfig(num_workers=3, use_gpu=use_gpu)
trainer = HorovodTrainer(
    train_loop_per_worker=train_loop_per_worker,
    scaling_config=scaling_config,
    datasets={"train": train_dataset},
)
result = trainer.fit()

Update your training function

First, update your training function to support distributed training.

If your training function already runs with the Horovod Ray Executor, you shouldn’t need to change it.

If you’re new to Horovod, see the Horovod guide.

Create a HorovodTrainer

A trainer is the primary Ray Train class for managing state and running training. For Horovod, set up a HorovodTrainer as in the following example:

from ray.train import ScalingConfig
from ray.train.horovod import HorovodTrainer
# For GPU Training, set `use_gpu` to True.
use_gpu = False
trainer = HorovodTrainer(
    train_func,
    scaling_config=ScalingConfig(use_gpu=use_gpu, num_workers=2)
)

Always use a HorovodTrainer when you train with Horovod, regardless of the training framework, such as PyTorch or TensorFlow.

To customize the backend setup, pass a HorovodConfig:

from ray.train import ScalingConfig
from ray.train.horovod import HorovodTrainer, HorovodConfig

trainer = HorovodTrainer(
    train_func,
    tensorflow_backend=HorovodConfig(...),
    scaling_config=ScalingConfig(num_workers=2),
)

For more configuration options, see the DataParallelTrainer API.

Run a training function

After you have a distributed training function and a trainer, call trainer.fit() to start training:

trainer.fit()

Further reading

The Ray Train HorovodTrainer replaces the distributed communication backend of the native libraries with its own implementation, so the remaining integration points stay the same. If you use Horovod with PyTorch or TensorFlow, see the corresponding guide for configuration details.

If you implement your own Horovod-based training routine without any of the training libraries, read the Ray Train user guides. You can adapt much of their content to generic use cases.