Multi-machine runs

A run that needs more GPUs than one machine has runs as several workers, one per machine, started together. Astraeus places the whole group at once (a gang), on the tightest network there is room on, gives every worker its rank and the leader's address, and restarts the group together when one member fails. This page shows how to start such a run with PyTorch, DeepSpeed or JAX, what each worker receives, and how failures are handled.

How a multi-machine run starts

With start: Gang, the scheduler places every worker in one decision or none, so a run never holds half its GPUs while it waits for the rest.

1. The run is Pending until every worker fits at once.
2. All workers are bound together. Each machine prepares its worker and pulls the image; the worker then waits at the gang barrier (ReadyToStart, Waiting at the gang barrier). The run is Starting (Starting together: 3 of 4 ready).
3. When every worker is at the barrier, all of them start. The run is Running once every worker runs.
4. If some worker has waited at the barrier for 15 minutes while another has not arrived, every worker is requeued: Gang barrier timed out waiting for all members.

The other start policies:

start	Behaviour	Use for
Gang	All placed and started together, behind the barrier.	Synchronous training (all-reduce).
MinAvailable	The first min_available workers are placed together and start without a barrier; the rest join as room appears.	Elastic training, a Spark driver and executors.
Independent	Each worker when it fits.	Workers that do not talk to each other.

The console chooses Gang for a run that may take several machines. The API's default, and what astra astraeus run sends, is Independent: set start: Gang yourself.

Ranks and the rendezvous

Workers are named <run>-<rank>: llama-0, llama-1. Rank 0 is the leader (with worker groups, rank 0 of the leader group). One worker runs on each machine and holds all of that machine's GPUs for the run.

Every worker receives:

Variable	Value
RANK	This worker's rank, from 0.
WORLD_SIZE	How many workers the run has (machines, not GPUs).
MASTER_ADDR	The address of the machine the leader runs on. Set once the leader is placed.
MASTER_PORT	29500.
ASTRAEUS_GPU_COUNT	GPUs this worker has: use it as the processes per node.
ASTRAEUS_REPLICA_INDEX, ASTRAEUS_WORLD_SIZE	The same as RANK and WORLD_SIZE.
ASTRAEUS_LEADER_TASK_NAME, ASTRAEUS_LEADER_ADDRESS	The leader's full name, and its machine's address.
SLURM_PROCID, SLURM_NODEID, SLURM_NTASKS, SLURM_NNODES, SLURM_LOCALID=0, SLURM_NTASKS_PER_NODE=1, …	Slurm's names for the same; see Slurm compatibility.

Your own env overrides any of these. The complete list is in Inside a worker.

MASTER_ADDR is the leader's machine address

It reaches the leader when workers use their machine's network, which is how machines with RDMA (InfiniBand, RoCE) run them, and machines installed with --network host. On machines where each worker has an address of its own (the node network, the installer's default without RDMA), nothing listens on the machine's address: point the rendezvous at the leader's name instead, which the cluster DNS resolves once the leader runs:

LEADER="${ASTRAEUS_LEADER_TASK_NAME#*.}"          # llama-0
until getent hosts "$LEADER" >/dev/null; do sleep 1; done
export MASTER_ADDR="$LEADER"

See Networking for which network a machine's workers use.

Train with torchrun on two machines

16 GPUs on 2 machines (8 each), the run kept within one InfiniBand fabric, with a 24-hour limit and checkpoints on a drive.

ConsoleCLIAPI

1. New run: Name llama-sft, Image nvcr.io/nvidia/pytorch:24.08-py3.
2. Command bash. Arguments cannot hold a quoted script (they are split on spaces): select Edit as JSON and set args as in the API tab.
3. GPUs 16, Machines → Exactly 2 (the hint says 16 GPUs on 2 machines: 8 on each.), CPU cores per GPU 12, Memory per GPU 120Gi.
4. Drives: datasets:/data,checkpoints:/ckpt.
5. Network and placement → Network between workers: InfiniBand at least.
6. More options → Start: Automatic (gang, since it takes 2 machines); Time limit 24h.
7. Start run.

astra astraeus run cannot start a gang or mount drives. Save the API tab's spec as a workspace template once (console: Save as template, or the templates API), then:

$ astra astraeus run --template llama-sft --name llama-sft-oct
llama-sft-oct submitted to gpu-east: https://console.astralyx.cloud/o/acme/w/nlp/jobs/gpu-east/llama-sft-oct

llama-sft.json

{
  "metadata": { "name": "llama-sft" },
  "spec": {
    "start": "Gang",
    "on_failure": "RestartJob",
    "task_template": {
      "image": "nvcr.io/nvidia/pytorch:24.08-py3",
      "command": "bash",
      "args": ["-c", "torchrun --nnodes=$WORLD_SIZE --node-rank=$RANK --nproc-per-node=$ASTRAEUS_GPU_COUNT --master-addr=$MASTER_ADDR --master-port=$MASTER_PORT /data/code/sft.py --resume-from /ckpt/llama-sft"],
      "time_limit_seconds": 86400,
      "datavolume_refs": [
        { "name": "datasets", "mount_path": "/data", "mode": "ReadOnly" },
        { "name": "checkpoints", "mount_path": "/ckpt", "mode": "ReadWrite" }
      ],
      "requested_resources": {
        "gpu_requests": { "count": 16 },
        "machines": 2,
        "per_gpu": { "cpu_cores": 12, "memory_bytes": 128849018880 },
        "network": { "interconnect": "infiniband" }
      }
    }
  }
}

$ curl -sS -X POST "$API/jobs" -H "Authorization: Bearer $ASTRA_TOKEN" \
    -H 'Content-Type: application/json' --data-binary @llama-sft.json | jq -r .metadata.name
llama-sft

Inside, each machine's torchrun starts 8 processes and sets LOCAL_RANK for each; your script calls torch.distributed.init_process_group("nccl") as usual.

Leave out machines to let the scheduler choose the shape: 16 GPUs may run as 2×8, or 4×4 inside one fabric when that is what is free. Add "gpu_requests": {"count": 16, "min_per_machine": 8} to never go below 8 per machine.

DeepSpeed

Launch DeepSpeed through torchrun: DeepSpeed reads the same variables, and Astraeus does not provide the SSH between workers that the deepspeed launcher's hostfile mode needs.

"command": "bash",
"args": ["-c", "torchrun --nnodes=$WORLD_SIZE --node-rank=$RANK --nproc-per-node=$ASTRAEUS_GPU_COUNT --master-addr=$MASTER_ADDR --master-port=$MASTER_PORT train.py --deepspeed --deepspeed_config /data/code/ds_zero3.json"]

In train.py, call deepspeed.init_distributed() (or let deepspeed.initialize do it).

JAX

One process per machine, with the GPUs of that machine:

train_jax.py

import os
import jax

jax.distributed.initialize(
    coordinator_address=f"{os.environ['MASTER_ADDR']}:{os.environ['MASTER_PORT']}",
    num_processes=int(os.environ["WORLD_SIZE"]),
    process_id=int(os.environ["RANK"]),
)
print(jax.process_index(), jax.device_count(), jax.local_device_count())

MPI

Astraeus starts one container per worker and provides no mpirun integration, PMI or SSH between workers. Programs launched with mpirun across machines do not work as they are. Use a framework that bootstraps from environment variables (PyTorch, DeepSpeed through torchrun, JAX, Horovod with Gloo), or bring your own launcher in the image.

Check the network first

Before a long run, start a short one with the same shape that does an all-reduce across every GPU:

allreduce.py

import os, time, torch, torch.distributed as dist

dist.init_process_group("nccl")
local = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local)
x = torch.ones(256 * 1024 * 1024 // 4, device="cuda")   # 256 MiB
for _ in range(5):
    dist.all_reduce(x)
torch.cuda.synchronize()
t = time.time()
for _ in range(20):
    dist.all_reduce(x)
torch.cuda.synchronize()
if dist.get_rank() == 0:
    gb = 20 * x.numel() * 4 * 2 * (dist.get_world_size() - 1) / dist.get_world_size() / 1e9
    print(f"bus bandwidth: {gb / (time.time() - t):.1f} GB/s over {dist.get_world_size()} GPUs")

Run it with the torchrun command above and "NCCL_DEBUG": "INFO" in env: the log of rank 0 shows which transport NCCL chose (NET/IB for InfiniBand, NET/Socket for TCP).

The network between workers

Where the scheduler put the run decides what each worker gets:

Placed on	Each worker gets
One machine	Nothing extra; GPUs talk over NVLink or PCIe inside it.
One NVLink domain	IMEX channel 0, so GPU memory is shared across the domain's machines.
An InfiniBand fabric, or RDMA	The machine's RDMA devices (/dev/infiniband/…), IPC_LOCK, unlimited locked memory, and NCCL_IB_HCA listing the ports that are up and at least network.min_gbps.
Ethernet across machines	NCCL_IB_DISABLE=1, so NCCL does not try InfiniBand ports that cannot reach the other machines.

Every GPU worker also gets NCCL_DEBUG=WARN, NCCL_ASYNC_ERROR_HANDLING=1 and NCCL_SOCKET_IFNAME=^lo,docker,virbr,veth,cni,wg. Any of these is replaced by a value you set in env.

To require a fast network, set network.interconnect (rdma, infiniband or nvlink) and network.min_gbps: the run waits rather than run over Ethernet. See GPUs and placement.

Workers reach each other by name through the cluster DNS when they have addresses of their own: <run>-<rank> is a worker, <run> is the leader while it runs. Names resolve within the workspace. See Names and service discovery.

Worker groups

Use task_groups when the workers are not all alike: a launcher and its workers, parameter servers, an engine your workers call.

spark-like.json

{
  "metadata": { "name": "etl" },
  "spec": {
    "start": "Independent",
    "leader_group": "driver",
    "task_groups": [
      { "name": "driver",
        "task_template": { "image": "apache/spark:3.5.1", "command": "/opt/spark/bin/driver.sh",
          "requested_resources": { "cpu_cores": 4, "memory_bytes": 17179869184 } } },
      { "name": "executor",
        "depends_on": [{ "group_name": "driver", "condition": "Running" }],
        "task_template": { "image": "apache/spark:3.5.1", "command": "/opt/spark/bin/executor.sh",
          "requested_resources": { "cpu_cores": 16, "memory_bytes": 68719476736, "machines": 4 } } }
    ]
  }
}

• Workers are named <run>-<group>-<i>: etl-driver-0, etl-executor-0 … etl-executor-3. Ranks run across groups in order: the driver is rank 0, the executors 1–4.
• depends_on holds a group until the groups it names are Running (or Healthy: running and passing their health_check). The waiting workers say Waits for worker group driver to be running. A gang cannot use depends_on.
• Each worker also gets ASTRAEUS_GROUP_NAME and ASTRAEUS_GROUP_REPLICA_INDEX (its index within the group).
• complete_with: "<group>" ends the run when that group completes and stops the others (servers that never exit on their own).

When a worker fails

on_failure	A worker fails	Its machine is lost
RestartJob (default for a gang)	Every worker is requeued and the gang starts again: Gang restart: llama-sft-1 failed.	Same: the lost worker is Down and the gang restarts.
RestartTask	Only that worker restarts; if it is the leader, every worker does (the rendezvous goes with it).	Same.
FailJob	No restart: the run fails and the other workers are stopped.	Same.

• Restarts back off from 10 s, doubling to 5 min, with a budget of 10; a worker running 10 minutes earns its budget back. When a gang member's budget is spent, the run fails: Worker llama-sft-1 failed on gpu-07: Exited with code 1; the gang does not run without it.
• A worker past its time limit fails the gang: the others are stopped with A gang member exceeded its time limit.
• A gang restart starts from the beginning of your script: write checkpoints to a drive and resume from the latest. See Priorities, preemption and checkpoints.
• NCCL_ASYNC_ERROR_HANDLING=1 makes a collective that loses a peer fail instead of hanging, so the gang restarts rather than sits idle until its time limit.

Troubleshooting

Symptom	Cause	Fix
The run stays Starting, one worker in Pulling	One machine is still pulling a large image.	Wait; after 15 minutes at the barrier the gang is requeued. Pre-pull large images, or use a smaller one.
Gang cannot be placed whole: …	Not every worker fits at once.	See Read a pending reason.
Rank 0 waits forever at init_process_group, others time out connecting	MASTER_ADDR is a machine address while workers have their own addresses.	Use the leader's name; see the warning in Ranks and the rendezvous.
The run hangs at its first collective	NCCL is trying a network the other machines cannot reach.	Ask for interconnect: rdma/infiniband, or set NCCL_IB_DISABLE=1 and NCCL_SOCKET_IFNAME in env.
NCCL WARN NET/IB : No device found	Placed over Ethernet, so no RDMA devices were passed.	Set network.interconnect to require RDMA.
Cannot plan workers: cannot evenly distribute 12 gpus across 5 machines	machines does not divide the GPUs.	Change one of them.

Related

• GPUs and placement
• Inside a worker
• Priorities, preemption and checkpoints
• Run specification
• Networking
• Slurm compatibility