Distributed training across machines#

You run PyTorch DistributedDataParallel with torchrun on 2 machines × 8 GPUs. The run first measures all-reduce bandwidth across the 16 GPUs, which tells you whether NCCL is really using the InfiniBand fabric, then trains a ResNet-50 on synthetic images and reports its throughput. Swap the training part for your own once the numbers look right.

What you need:

Two machines with 8 GPUs each in your workspace's pools, on the same InfiniBand fabric. Machines with RDMA run their workers on the host network, which is what NCCL over InfiniBand needs.
The editor role in the workspace, an API token, curl and jq (How the recipes are written).

How Astraeus runs it#

A run asking for 16 GPUs on 2 machines gets two workers, ddp-resnet50-0 and ddp-resnet50-1, one per machine, 8 GPUs each. They start as a gang: both are placed together or not at all, and neither starts its container until both are ready. Each worker runs one torchrun, which starts 8 processes, one per GPU.

flowchart LR
  subgraph m1["gpu-01 · worker ddp-resnet50-0 · RANK=0 (leader)"]
    t0["torchrun<br/>8 processes"]
  end
  subgraph m2["gpu-02 · worker ddp-resnet50-1 · RANK=1"]
    t1["torchrun<br/>8 processes"]
  end
  t1 -- "rendezvous MASTER_ADDR:29500" --> t0
  t0 <-- "NCCL over InfiniBand" --> t1

Every worker gets these variables (a variable you set in env wins):

Variable	Value here	Meaning
`RANK`	`0`, `1`	The worker's index in the run: torchrun's `--node-rank`.
`WORLD_SIZE`	`2`	How many workers: torchrun's `--nnodes`. Not the number of GPU processes.
`MASTER_ADDR`	`10.20.0.11`	The address of the machine the leader (worker 0) runs on.
`MASTER_PORT`	`29500`	The rendezvous port.
`ASTRAEUS_GPU_COUNT`	`8`	GPUs this worker holds: torchrun's `--nproc-per-node`.
`ASTRAEUS_WORLD_SIZE`, `ASTRAEUS_REPLICA_INDEX`	`2`, `0`/`1`	The same as `WORLD_SIZE` and `RANK`.
`ASTRAEUS_LEADER_ADDRESS`, `ASTRAEUS_LEADER_TASK_NAME`		The same as `MASTER_ADDR`, and the leader worker's name.
`ASTRAEUS_TASK_NAME`, `ASTRAEUS_NODE_NAME`, `ASTRAEUS_JOB_NAME`		This worker, its machine, the run.
`SLURM_PROCID`, `SLURM_NNODES`, `SLURM_NTASKS`, `SLURM_JOB_NAME`, …		Slurm's names for the same, so scripts and libraries written for Slurm work unchanged.

torchrun sets LOCAL_RANK, RANK and WORLD_SIZE for each of the processes it starts; inside your training script they are per process (0–15 across the run), as usual. See Inside a worker for the complete list.

On a GPU worker the machine also sets:

Setting	Value
`NCCL_DEBUG`	`WARN`
`NCCL_ASYNC_ERROR_HANDLING`	`1`
`NCCL_SOCKET_IFNAME`	`^lo,docker,virbr,veth,cni,wg`
`NCCL_IB_HCA`	The machine's RDMA ports that are up and fast enough (`=mlx5_0:1,=mlx5_1:1,…`), when the run was placed on InfiniBand or RDMA or asked for it.
RDMA devices	`/dev/infiniband/*` and the `IPC_LOCK` capability, in the same case.
`NCCL_IB_DISABLE`	`1` when the run was placed on Ethernet.
`/dev/shm`	Half the worker's memory limit, between 1 GiB and 64 GiB.
Locked memory	Unlimited (`memlock`), for pinned buffers and RDMA registration.

1. Write the script#

train_ddp.py

import os
import time

import torch
import torch.distributed as dist
import torch.nn as nn
import torchvision


def main():
    dist.init_process_group("nccl")
    rank, world = dist.get_rank(), dist.get_world_size()
    local_rank = int(os.environ["LOCAL_RANK"])  # set by torchrun
    torch.cuda.set_device(local_rank)
    dev = torch.device("cuda", local_rank)

    # 1. All-reduce bandwidth: is NCCL using the fabric?
    x = torch.ones(256 * 1024 * 1024, dtype=torch.float32, device=dev)  # 1 GiB
    for _ in range(5):
        dist.all_reduce(x)
    torch.cuda.synchronize()
    iters = 20
    t0 = time.perf_counter()
    for _ in range(iters):
        dist.all_reduce(x)
    torch.cuda.synchronize()
    dt = (time.perf_counter() - t0) / iters
    busbw = x.numel() * 4 * 2 * (world - 1) / world / dt / 1e9
    if rank == 0:
        print(f"all_reduce 1 GiB over {world} GPUs: {dt * 1e3:.1f} ms, bus bandwidth {busbw:.0f} GB/s", flush=True)
    del x

    # 2. ResNet-50 on synthetic images: throughput.
    model = torchvision.models.resnet50().to(dev, memory_format=torch.channels_last)
    model = nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])
    opt = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
    loss_fn = nn.CrossEntropyLoss()
    batch = int(os.environ.get("BATCH_PER_GPU", "256"))
    steps = int(os.environ.get("STEPS", "300"))
    warmup = 20
    images = torch.randn(batch, 3, 224, 224, device=dev).to(memory_format=torch.channels_last)
    labels = torch.randint(0, 1000, (batch,), device=dev)
    for step in range(1, steps + 1):
        if step == warmup + 1:
            torch.cuda.synchronize()
            t0 = time.perf_counter()
        with torch.autocast("cuda", dtype=torch.bfloat16):
            loss = loss_fn(model(images), labels)
        opt.zero_grad(set_to_none=True)
        loss.backward()
        opt.step()
        if rank == 0 and step % 50 == 0:
            print(f"step {step}/{steps} loss {loss.item():.3f}", flush=True)
    torch.cuda.synchronize()
    ips = batch * world * (steps - warmup) / (time.perf_counter() - t0)
    if rank == 0:
        print(f"throughput: {ips:,.0f} images/s on {world} GPUs", flush=True)
    dist.destroy_process_group()


if __name__ == "__main__":
    main()

2. Write the run specification#

run.json

{
  "metadata": {"name": "ddp-resnet50", "labels": {"project": "fabric-check"}},
  "spec": {
    "start": "Gang",
    "on_failure": "RestartJob",
    "task_template": {
      "image": "nvcr.io/nvidia/pytorch:24.08-py3",
      "command": "bash",
      "args": [
        "-c",
        "echo \"worker $RANK of $WORLD_SIZE on $ASTRAEUS_NODE_NAME, leader $MASTER_ADDR:$MASTER_PORT, $ASTRAEUS_GPU_COUNT GPUs\"; exec torchrun --nnodes=$WORLD_SIZE --nproc-per-node=$ASTRAEUS_GPU_COUNT --node-rank=$RANK --master-addr=$MASTER_ADDR --master-port=$MASTER_PORT /app/train_ddp.py"
      ],
      "env": {"NCCL_DEBUG": "INFO", "NCCL_DEBUG_SUBSYS": "INIT,NET", "STEPS": "300", "BATCH_PER_GPU": "256"},
      "time_limit_seconds": 3600,
      "requested_resources": {
        "gpu_requests": {"count": 16, "models": ["H100"], "healthy_only": true},
        "machines": 2,
        "per_gpu": {"cpu_cores": 12, "memory_bytes": 107374182400},
        "network": {"interconnect": "infiniband", "min_gbps": 400},
        "topology": {"keep_within": "fabric", "patience_seconds": 600},
        "node_selection": {"mode": "Any"}
      },
      "configs": []
    }
  }
}

$ jq --rawfile src train_ddp.py \
    '.spec.task_template.configs = [{"mounts": ["/app/train_ddp.py"], "value": $src}]' \
    run.json > run.full.json

What each part does:

Field	Why
`gpu_requests.count: 16`, `machines: 2`	16 GPUs for the whole run, split evenly: 8 on each of 2 machines, one worker per machine. The count must divide by `machines`.
`per_gpu`	12 cores and 100 GiB of memory with each GPU: 96 cores and 800 GiB per worker. Give either `per_gpu` or `cpu_cores`/`memory_bytes`, not both.
`start: Gang`	All or nothing, started together. If a worker is still not ready 15 minutes after the first one reached the start barrier, every worker goes back to the queue (Gang barrier timed out waiting for all members).
`on_failure: RestartJob`	Synchronous training does not survive the loss of a member: any failure restarts both workers.
`network.interconnect: infiniband`, `min_gbps: 400`	Only machines with InfiniBand ports of at least 400 Gb/s; the worker gets only those ports in `NCCL_IB_HCA`. Never RoCE, never Ethernet.
`topology.keep_within: fabric`	Both machines on one InfiniBand fabric, or the run waits.
`topology.patience_seconds: 600`	Wait up to 10 minutes for a tighter placement than the one free now (one rack within the fabric, say). Other work runs meanwhile.
`healthy_only`	Never a GPU showing signs of a coming fault: a fault mid-run costs more than the wait.
`NCCL_DEBUG: INFO`	Overrides the default `WARN` so the log shows which transport NCCL chose. Set it back once it works.
`exec torchrun …`	`bash -c` expands the variables; `exec` makes `torchrun` the container's main process so it receives the stop signal and passes it to the 8 processes.

Fabrics are found by themselves: machines whose InfiniBand ports report the same subnet manager are one fabric. An administrator can set the topology.astraeus.io/fabric and topology.astraeus.io/rack labels on Machines → machine → Where it is (Pools, labels and topology).

3. Submit the run#

ConsoleCLIAPI

Open Runs → New run and press Edit as JSON (worker groups, drives, probes…).
Paste the contents of run.full.json and press Start run.

The form has the same settings without JSON: GPUs 16, Machines exactly 2, GPU model, and under the advanced sections Network between workers, RDMA port speed at least, Never leave, Wait for a better network and Only healthy GPUs. The script then has to be in the image.

The CLI cannot attach config files, so bake the script into an image first:

Dockerfile

FROM nvcr.io/nvidia/pytorch:24.08-py3
COPY train_ddp.py /app/train_ddp.py

$ docker build -t registry.example.com/ml/ddp-check:1 . && docker push registry.example.com/ml/ddp-check:1
$ astra astraeus run --name ddp-resnet50 --image registry.example.com/ml/ddp-check:1 \
    --gpus 16 --machines 2 --cpus 12 --mem 100G --time 1h \
    -e NCCL_DEBUG=INFO -e NCCL_DEBUG_SUBSYS=INIT,NET \
    -- bash -c 'exec torchrun --nnodes=$WORLD_SIZE --nproc-per-node=$ASTRAEUS_GPU_COUNT --node-rank=$RANK --master-addr=$MASTER_ADDR --master-port=$MASTER_PORT /app/train_ddp.py'
ddp-resnet50 submitted to main: https://console.astralyx.cloud/o/acme/w/vision/jobs/main/ddp-resnet50

With GPUs, --cpus and --mem are per GPU. Several machines start as a gang. The CLI has no flags for the GPU model, healthy GPUs, the network, topology or on_failure: for those, submit run.full.json from the console or the API. Keep the single quotes around the bash -c command: the variables must be expanded in the worker, not by your shell.

$ curl -fsS -X POST "$API/runs" -H "Authorization: Bearer $ASTRAEUS_TOKEN" \
    -H 'content-type: application/json' -d @run.full.json | jq -r .metadata.name
ddp-resnet50

4. Watch it start#

While the two machines are found, the run is Pending; once both are placed it is Starting until both containers are ready, then Running.

$ astra astraeus runs
NAME                          CLUSTER       STATE       REASON
ddp-resnet50                  main          Starting    Starting together: 1 of 2 ready
$ astra astraeus runs
NAME                          CLUSTER       STATE       REASON
ddp-resnet50                  main          Running     2 of 2 workers running

The run's page says which network the run got, in words ("one InfiniBand fabric"), and the workers' machines. If it waits, the page and each worker's Why it is not running yet say what is missing, machine by machine.

5. Check NCCL and the numbers#

Worker 0 holds global rank 0, so its log has the results:

$ astra astraeus logs ddp-resnet50-0 | tail -n 400 | grep -E 'worker |NET/IB|all_reduce|throughput'
worker 0 of 2 on gpu-01, leader 10.20.0.11:29500, 8 GPUs
gpu-01:212:390 [0] NCCL INFO NET/IB : Using [0]mlx5_0:1/IB [1]mlx5_1:1/IB [2]mlx5_2:1/IB [3]mlx5_3:1/IB [4]mlx5_4:1/IB [5]mlx5_5:1/IB [6]mlx5_6:1/IB [7]mlx5_7:1/IB [RO]; OOB ibp24s0:10.20.0.11<0>
all_reduce 1 GiB over 16 GPUs: 5.6 ms, bus bandwidth 360 GB/s
throughput: 46,912 images/s on 16 GPUs

What to look for:

NET/IB : Using [0]mlx5_…/IB lists the InfiniBand ports. NET/Socket instead means NCCL fell back to TCP: see Troubleshooting.
The bus bandwidth should be close to what the fabric can carry per GPU. On 8 × 400 Gb/s NDR ports per machine, expect roughly 300 to 380 GB/s; a figure near 10 GB/s means TCP.

Your numbers depend on the GPUs, the fabric and the image.

6. Clean up#

ConsoleCLIAPI

On the run's page press Delete and confirm with the run's name.

$ astra astraeus delete ddp-resnet50
ddp-resnet50 deleted

$ curl -fsS -X DELETE "$API/runs/ddp-resnet50" -H "Authorization: Bearer $ASTRAEUS_TOKEN"

Variations#

Let Astraeus choose the machines. Leave out machines and say how small a slice may be: "gpu_requests": {"count": 16, "min_per_machine": 8}. Astraeus picks the fewest machines on the best network with room (two of 8, or one of 16) and may change its mind while nothing has started. With per_gpu, the cores and memory follow the shape it picks. RANK, WORLD_SIZE and ASTRAEUS_GPU_COUNT always describe the shape it chose, so the same torchrun line works.

A range of machines. "topology": {"min_nodes": 2, "max_nodes": 4} with "gpu_requests": {"count": 32} runs on 2, 3 or 4 machines, one worker on each. The GPUs must split evenly over the number chosen.

One NVLink domain (GB200 NVL72). "topology": {"keep_within": "nvlink_domain"} keeps every worker in one multi-node NVLink domain and gives each worker IMEX channel 0, so GPU memory is shared across the machines. "network": {"interconnect": "nvlink"} asks for NVLink between the workers: one machine, or one NVLink domain.

RoCE. "network": {"interconnect": "rdma"} accepts InfiniBand or RoCE but never plain Ethernet.

Spread for resilience instead. "topology": {"spread": "rack"} puts each worker in a different rack. Use it for replicas that must not fail together, not for training: it is the opposite of what NCCL wants.

Training data. Mount a dataset read-only with datavolume_refs (see A dataset on a drive, shared by runs) and write checkpoints from global rank 0 only, to a drive both machines reach (a drive on one machine or on a shared filesystem).

Troubleshooting#

Symptom	Cause	Fix
The log shows `NET/Socket` and the bus bandwidth is low	The run was placed without RDMA (no `network` rule, and no fabric had room), or the image lacks the RDMA libraries.	Add `"network": {"interconnect": "infiniband"}` so the run waits for a fabric. NVIDIA's NGC images include the libraries.
`NCCL WARN … NET/IB : No device found`	The worker got no RDMA devices: the machine's ports are down or slower than `min_gbps`.	Check the machine's RDMA ports on its page; lower `min_gbps`.
The run waits in `Pending`; workers say no fabric has room	`keep_within: fabric` with not enough free machines on one fabric.	Wait, free machines, or drop the rule. `patience_seconds` only applies to the network the run could get now, not to the rules.
Workers on machines without RDMA time out at the rendezvous	On the mesh network every worker has its own address, and `MASTER_ADDR` is the leader's machine address.	Use the run's name as the rendezvous address: `--master-addr=$SLURM_JOB_NAME`. It resolves to the leader worker in your workspace's cluster DNS.
The workers go back to the queue with Gang barrier timed out waiting for all members	A gang member did not become ready within 15 minutes of the first (an image pull failing on one machine, say).	Open the stuck worker's page: its state and History say why.
`torch.distributed.DistStoreError: … timed out` on worker 1	The leader's `torchrun` was not up yet, or port 29500 is in use on the leader's machine.	Set another port with `"env": {"MASTER_PORT": "29600"}`; your value wins.