Simple PyTorch cluster

This example shows how you can perform distributed computation with PyTorch. It is a kind of ‘hello world’ example for distributed ML training: setting up a cluster and executing a broadcast operation to share a single tensor.

Basic setup: Imports, dependencies, and a script 

Let’s get the imports out of the way first.

import os
from pathlib import Path

import modal

Communicating between nodes in a cluster requires communication libraries. We’ll use torch, so we add it to our container’s Image here.

image = modal.Image.debian_slim(python_version="3.12").uv_pip_install(
    "torch~=2.5.1", "numpy~=2.2.1"
)

The approach we’re going to take is to use a Modal Function to launch the underlying script we want to distribute over the cluster nodes. The script is located in another file in the same directory of our examples repo. In order to use it in our remote Modal Function, we need to duplicate it remotely, which we do with add_local_file.

this_directory = Path(__file__).parent

image = image.add_local_file(
    this_directory / "simple_torch_cluster_script.py",
    remote_path="/root/script.py",
)

app = modal.App("example-simple-torch-cluster", image=image)

Configuring a test cluster 

First, we set the size of the cluster in containers/nodes. This can be between 1 and 8. This is part of our Modal configuration, since Modal is responsible for spinning up our cluster.

n_nodes = 2

Next, we set the number of processes we run per node. The usual practice is to run one process per GPU, so we set those two values to be equal. Note that N_GPU is Modal configuration (“how many GPUs should we spin up for you?”) while nproc_per_node is torch.distributed configuration (“how many processes should we spawn for you?”).

n_proc_per_node = N_GPU = 4
GPU_CONFIG = f"A10G:{N_GPU}"

Lastly, we need to select our communications library: the software that will handle sending messages between nodes in our cluster. Since we are running on GPUs, we use the NVIDIA Collective Communications Library (nccl, pronounced “nickle”).

This is part of torch.distributed configuration — Modal handles the networking infrastructure but not the communication protocol.

backend = "nccl"  # or "gloo" on CPU, see https://pytorch.org/docs/stable/distributed.html#which-backend-to-use

This cluster configurations is nice for testing, but typically you’ll want to run a cluster with the maximum number of GPUs per container — 8 if you’re running on H100s, the beefiest GPUs we offer on Modal.

Launching the script 

Our Modal Function is merely a ‘launcher’ that sets up the distributed cluster environment and then calls torch.distributed.run, the underlying Python code exposed by the torchrun command line tool.

So executing this distributed job is easy! Just run

modal run simple_torch_cluster.py

in your terminal.

In addition to the values set in code above, you can pass additional arguments to torch.distributed.run via the command line:

modal run simple_torch_cluster.py --max-restarts=1

@app.function(gpu=GPU_CONFIG)
@modal.clustered(size=n_nodes)
def dist_run_script(*args):
    from torch.distributed.run import parse_args, run

    cluster = modal.Cluster.from_context()  # we populate this data for you
    # which container am I?
    container_rank = cluster.container_rank()
    # what are the IP addresses of the containers in this cluster?
    container_ips = cluster.container_ips()
    # how many containers are in this cluster?
    world_size = len(container_ips)
    # what's the leader/master/main container's address?
    main_addr = container_ips[0]
    # what's the identifier of this cluster task in Modal?
    task_id = os.environ["MODAL_TASK_ID"]
    print(f"hello from {container_rank=}")
    if container_rank == 0:
        print(
            f"reporting cluster state from rank0/main: {main_addr=}, {world_size=}, {task_id=}"
        )

    run(
        parse_args(
            [
                f"--nnodes={n_nodes}",
                f"--node_rank={container_rank}",
                f"--master_addr={main_addr}",
                f"--nproc-per-node={n_proc_per_node}",
                "--master_port=1234",
            ]
            + list(args)
            + ["/root/script.py", "--backend", backend]
        )
    )