Scale Multinode Launcher

This script launches SCOUT MAE training on multiple Mithril nodes using PyTorch DDP. It should be run from mithril1. The script is mainly responsible for distributed orchestration: choosing the master address, SSH-launching workers, setting NCCL variables, logging, cleanup, and automatic relaunch.

Cluster setup

HOSTS=(mithril1 mithril2 mithril3 mithril4 mithril6)
NNODES=${#HOSTS[@]}
NPROC_PER_NODE=8

This runs on 5 nodes, skipping mithril5. Since each node launches 8 local processes, the job has:

DDP ranks total. mithril1 is node rank 0 and acts as the launcher.

Config and resume

CONFIG="${1:-configs/tile1/150k/E/scale_170M.yaml}"
RESUME="${RESUME:-auto}"
REPO=/mnt/phase1-data-a/ryan/scout
PYBIN=/mnt/phase1-data-a/ryan/miniconda3/envs/scout/bin
LOG_DIR=$REPO/logs

The first CLI argument chooses the config; otherwise it uses the default scale_170M.yaml. RESUME=auto resumes the latest checkpoint, while RESUME=none starts from scratch on the first launch.

After the first exit/crash, the script forces:

RESUME_FLAG="--resume auto"

Network pinning

MASTER_ADDR=\$(ip -4 addr show enp6s0 | awk '/inet / {print \$2}' | cut -d/ -f1)
NCCL_SOCKET_IFNAME=enp6s0

The script uses the enp6s0 IP on mithril1 as the PyTorch master address, and also forces NCCL to use enp6s0. This avoids accidentally using docker0 or another wrong interface.

Preflight checks

Before launch, each remote node is checked for:

• train_mae.py exists,
• torchrun exists and is executable,
• enp6s0 exists.

This avoids a partial distributed launch where some ranks start and then hang waiting for a missing/broken node.

NCCL settings

NCCL_P2P_DISABLE=1
NCCL_IB_DISABLE=1
NCCL_SOCKET_IFNAME=enp6s0
TORCH_NCCL_ASYNC_ERROR_HANDLING=1
TORCH_NCCL_BLOCKING_WAIT=1
TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC=900

P2P and IB are disabled because this cluster has known issues with those paths. NCCL is forced onto TCP over enp6s0. The heartbeat timeout makes dead nodes fail after about 15 minutes instead of hanging forever.

Process tag and cleanup

PROC_TAG="scale_\${USER}_\${MASTER_PORT}_\$\$"

Every launched training command includes this tag via --proc-tag. Cleanup then uses:

pkill -9 -f "\${PROC_TAG}"

on every node. This matters because SSH can disconnect while remote torchrun or Python workers keep running.

Launch and logs

Each node runs a command of the form:

\$PYBIN/torchrun \$TR_ARGS --node_rank=\$i train_mae.py \
  --config \${CONFIG} \${RESUME_FLAG} --proc-tag \${PROC_TAG}

Node 0 launches locally. Other nodes launch through SSH. Each node writes to:

\${LOG_DIR}/scale_node\${i}.log

To monitor:

tail -f /mnt/phase1-data-a/ryan/scout/logs/scale_node*.log

Failure and relaunch behavior

The script waits for all node processes, records any nonzero exit, then calls kill_all. It does this manually instead of letting set -e kill the launcher on the first failed rank.

After cleanup:

RESUME_FLAG="--resume auto"
sleep 900

Then the whole job is launched again.

Stopping

Ctrl-C triggers:

kill_all

so both local and remote tagged workers are killed immediately.

For a graceful stop after the current attempt exits:

touch /mnt/phase1-data-a/ryan/scout/logs/STOP_SCALE

The sentinel is checked only between iterations, not during a currently running training attempt.

Mental model

Use this launcher when the problem is distributed survival: node failure, NCCL hang, SSH disconnect, stale workers, or relaunch behavior.

Debug train_mae.py or the YAML config when the problem is model behavior, loss curves, checkpoint semantics, data order, or optimizer state.

Cross Region Training

Coreweave SUNK gives a SLURM interface on Kubernetes, so user-facing workflow is still mostly sbatch, srun, and torchrun rather than raw Kubernetes.

Architecture:

• Slurm/Kubernetes: allocates resources
• PyTorch DDP: synchronizes GPUs inside a region
• DiLoCo: synchronizes model replicas across weakly connected regions
• Object storage: rendezvous transport (not training algorithm)

Cluster Setup

We have $16 \times \text{H}200$s on EU island and $1 \times \text{GH}200$ on US island.

We must assume /mnt/data$_\text{EU} \neq$ /mnt/data$_\text{US}$ (some data must stay in Spain). This requires each region to have its own:

• code checkout
• Python/conda env
• dataset shard
• credentials
• Slurm launcher
• local config

We must build environment on the compute node, not login node!

srun --gres=gpu:1 --pty bash
python - <<'PY'
print("arch:", platform.machine())
print("cuda:", torch.cuda.is_available())
print("gpu:", torch.cuda.get_device_name(0))
PY

Local Training (DDP)

Local DDP does:

However across a WAN the cross-region link would sit on a critical path of every optimizer step. So across regions should not have a shared DDP process group nor WAN NCCL (instead we should have periodic DiLoCo sync).

WAN

Core network is the high-capacity central backbone that moves massive volumes of data, where the edges (CNEs) link.

WAN NCCL should never be done (extremely overkill) unless we requires exact consistency of a single-cluster run.

torchrun \
    --nnodes="$NNODES" \
    --nproc_per_node="$GPUS_PER_NODE" \
    --rdzv_backend=c10d \
    --rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
    train.py \
    --config "$CONFIG" \
    --xr-enabled

c10d: C++ backend library inside PyTorch that handles distributed communication.

• torch.distributed.init_process_group() wakes up c10d
• provides collective communication APIs like all_reduce, broadcast, and all_gather.
• it manages the backends: NCCL (Nvidia GPUs), Gloo (CPUs), MPI

rdz: nodes don't automatically know each other's IP address or who's rank $0$. Rendezvous process is a roll call where all nodes show up, agree on participants, assign ranks, and pick master node.

• Intra-node: multiple GPUs on the same machine. Processes can rendezvous through a local TCP store / file store / launcher-provided environment variables. Actual gradient sync uses fast local GPU links such as NVLink or PCIe through NCCL.
• Inter-node, same cluster/region: multiple machines in one Slurm/Kubernetes cluster. Usually no bucket is needed. The launcher picks a master node/address, all workers connect to a shared TCP store such as c10d, ranks are assigned, then NCCL performs direct node-to-node collectives over the cluster network.
• Cross-region / independent clusters: jobs may be launched separately, sit behind different schedulers/networks, and may not be able to form one clean NCCL world. In this case we often need an external rendezvous or mailbox: object storage bucket, database, TCP service, Redis, etc.

For normal same-cluster DDP:

For cross-region DiLoCo:

Object Storage

Object storage acts as rendezvous mailbox:

For our setup, .pt is delta payload and .done is barrier marker.

Kueue is. Kubernetes-native queing system for batch, HPC, and AI/ML jobs (decide which jobs wait and where they run based on quotas). Helps with multi-cluster dispatch where a job is assigned to a ClusterQueue configured with a MultiKueue admission check.

def exchange(tag, region, regions, payload):
    data = serialize(payload)
    
    put(f"{tag}/{region}.pt", data)
    put(f"{tag}/{region}.done", b"1")
    
    out = {region: payload}
    
    for peer in regions:
        if peer == regions:
            continue
        wait_until_exists(f"{tag}/{peer}.done")
        out[peer] = deserialize(get(f"{tag}/{peer}.pt"))
    return out

DiLoCo Core Algorithm

At sync round $t$, all regions start with same outer parameters $\theta_t$.

Each region $i$ trains locally for $H$ steps, $\theta_t \rightarrow \theta_{t,H}^{(i)}$. We exchange the parameter delta:

For simple SGD, this becomes

We aggregate to get,

Then apply outer optimizer:

def diloco_outer_step(model, theta_outer, mom, transport, round_id):
    theta_local = flatten(model)
    delta_local = theta_outer - theta_local
    
    deltas = transport.exchange(
        tag=f"r{round_id:06d}",
        payload=delta_local,
    )
    
    delta = sum(weights[r] * deltas[r] for r in regions)
    
    mom.mul_(outer_momentum).add_(delta)
    theta_new = theta_outer - outer_lr * mom
    
    assign_model_params(model, theta_new)
    return theta_new.clone(), mom
    

Training Loop Pattern

Per-region, training is normal.

for step, batch in enumerate(loader):
    loss = model_loss(model, batch)
    
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()
    
    if outer_sync is not None:
        outer_sync.maybe_sync()

with defined class:

class OuterSync:
    def maybe_sync(self, force=False):
        self.inner += 1
        
        if not force and self.inner %= self.every != 0:
            return False
            
        self.round += 1
        self.sync_round()
        return True

def sync_round(self):
    if self.rank == 0:
        theta_local = self._gather_flat_params()
        delta_local = self.theta_outer - theta_local
        
        deltas = self._exchange(
            tag=f"r{round_id:06d}",
            payload=delta_local,
        )
        
        delta = sum(self.w[r] * deltas[r] for r in self.regions)
        
        self.mom.mul_(self.momentum).add_(delta)
        theta_new = self.theta_outer - self.outer_lr * self.mom
    else:
        theta_new = None
    
    theta_new = self._broadcast(theta_new) # rank 0 to all local ranks
    self.theta_outer = theta_new.clone()
    self._scatter_to_model(theta_new)

Transport-Agnostic Design

Separate the tranining algorithm from transport backend.

DiLoCo shouldn't care whether regions sync through S3, GCS, HTTP, rsync, Redis, etc.

So use a small rendezvous interface:

class Transport:
     def exchange(self, tag: str, payload):
        buf = io.BytesIO()
        torch.save(payload, buf)
        data = buf.getvalue()
        
        self.put(f"{tag}/{self.region}.pt", data)
        self.put(f"{tag}/{self.region}.done", b"1")
        
        out = {self.region: payload}
        
        for peer in self.regions:
            if peer == self.region:
                continue
            
            self.get(f"{tag}/{peer}.done")
            peer_bytes = self.get(f"{tag}/{peer}.pt")
            out[peer] = torch.load(io.BytesIO(peer_bytes), map_location="cpu")
        
        return out

Algorithm can just call peer_payloads = transport.exchange(tag, local_payload)

S3 is just one backend:

class S3Transport(Transport):
    def put(self, key, data):
        aws_s3_cp_stdin(data, key)
    
    def get(self, key, timeout=1800):
        return poll_until_exists_and_download(key, timeout)
    
    def rm_round(self, tag):
        rm_recursive(tag)

Use aws configure set s3.addressing_style virtual for launcher setup.

Initialization, Rounds, and Observability

We run an initial alignment step:

theta0 = flatten(model)
if rank == 0:
    theta0 = cross_region_align(thetae0)
    
theta0 = broadcast_from_rank0(theta0)
assign_model_params(model, theta0)
    
theta_outer = theta0.clone()

Otherwise, $\Delta_i = \theta_\text{outer} - \theta_i$ will be mismatched.

Use explicit round IDs: tag = f"r{round_id:06d}"

For timing, make sure to split time between up, wait, down, and down_bw.

Additional Optimizations

Despite correctness, the aws s3 cp stream from rank was poor because only one process, one object, and one node were using the WAN.

1. Persistent object-store client. Replace shelling out to aws s3 cp with a persistent boto3 client. This avoids process startup every sync round and reuses HTTP/S3 connections.
   s3 vs boto3

   AWS CLI is a command-line tool for executing quick, single tasks directly from your terminal, whereas Boto3 is the official Python Software Development Kit (SDK) used to integrate AWS services directly inside programmatic Python applications. Under the hood, both tools are actually built by Amazon on top of the exact same core translation library, botocore.
2. Send bf16 deltas. The WAN payload is a parameter delta, \[ \Delta_i = \theta_\text{outer} - \theta_i, \]

   not the full training state. Sending \(\Delta_i\) in bf16 roughly halves transfer size:

   \[ \text{fp32} = 4\text{ bytes/param}, \qquad \text{bf16} = 2\text{ bytes/param}. \]

3. Shard the delta across nodes. Instead of one EU rank uploading the whole delta, split the flat delta into shards:
   \[ \text{owner}(s) = s \bmod \text{num\_nodes}. \]

   Then each EU node leader uploads its own shards. This lets both EU nodes use their WAN paths.

4. Upload many shard objects concurrently. One object per node is still close to one stream per node. Use many shard objects and upload/download them with a thread pool. This gives object storage enough parallel requests to fill the link.
5. Avoid the PyTorch slice serialization trap. Before saving a shard, clone it:

   
   shard = delta[lo].clone()
   torch.save(shard, buffer)
   

   Otherwise torch.save(delta[lo]) can serialize the backing storage of the larger tensor, accidentally sending much more than the shard.

6. Use large payloads to saturate WAN. Small models do not generate enough delta bytes to reach steady-state bandwidth. A 1B-param model has a bf16 delta around:
   \[ 10^9 \times 2 \approx 2\text{ GB}. \]

   With two EU nodes, this is about \(1\text{ GB}\) per node, large enough to better saturate the cross-region path.

7. Keep the outer optimizer step on GPU. At 1B params, the weighted merge and Nesterov outer update become a large vector operation. Keeping \(\theta_\text{outer}\), momentum, and received deltas on GPU avoids CPU-side reduce becoming the new bottleneck.
8. Send raw shard bytes instead of torch.save. For fixed flat parameter shards, both sides already know shard offsets, dtype, and shape. So instead of serializing each shard with torch.save, materialize a contiguous bf16 shard and send its raw bytes directly:
   \[ \text{tensor slice} \rightarrow \text{contiguous bf16 bytes} \rightarrow \text{object store}. \]

   This removes torch.save/torch.load overhead and avoids Python/PyTorch serialization becoming the bottleneck. The cost is that the payload is no longer self-describing, so sender and receiver must use the exact same parameter ordering and shard metadata.

Gradient Sync

Sync gradients every step.

After local backward, each region already has its intra-region DDP-averaged gradient:

Then each region exchanges its flattened gradient through the same object-storage transport:

The averaged gradient is written back into p.grad, then normal clipping and optimizer step happen:

loss.backward()
    
if outer is not None and outer.pre_step:
    outer.sync_grads()      # cross-region grad average
    
clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
    
if outer is not None and not outer.pre_step:
    outer.maybe_sync()      # DiLoCo path

The implementation reuses the DiLoCo plumbing: same object-store mailbox, shard ownership, node leaders, round tags, and cleanup. The only real difference is the payload:

Hardware

Terminology

NIC (Network Interface Card): the hardware device that connects the server to the network. What sends/receives data between nodes, like gradients, checkpoints, S3 traffic, or NCCL communication.

MFU: how much useful model math did we achieve versus dense hardware peak?

HFU: how much hardware math did GPU execute versus dense hardware peak?

Topology table: how close each NIC is to the GPU through PCIe/NUMA.

nvidia-smi topo -m

• PCIe (Peripheral Component Interconnect Express): high-speed bus connecting GPUs/NICs/SSDs to the CPU/system.
• NUMA (Non-Uniform Memory Access): Tells us which CPU/memory island the GPU or NIC belongs to (each region has its own nearby RAM and nearby PCIe devices). Same NUMA node is faster/closer than crossing nodes.

PIX vs NODE vs SYS

• PIX means devices are very close on PCIe. Usually at most one PCIe bridge (best kind of GPU-NIC path).
• NODE means device are on the same NUMA node, but path is less direct. For example it might cross PCIe host bridges or internal interconnects inside NUMA node (better than crossing sockets).
• SYS is worse as it crosses CPU/socket interconnect between NUMA nodes.

As a result NIC0 is the GPU-local NIC. So if we can control NIC selection, we would prefer NIC0 for GPU-adjacent traffic.

• Other NICs are still reachable but bandwidth or latency would be worse
• Note that NIC8-NIC16 are mutually PIX which means they are close to each other (under same PCIe switch). However they are still NODE away from GPU0.

NUMA affinity: Assigned to a NUMA node, tells us which CPU/memory domain (which is on same node) a device is closest to.

CPU affinity: which CPU cores a process or device should use.

Improving MFU/HFU

• High MFU usually means the training step is dominated by large, efficient tensor-core kernels: matmuls, MLPs, QKV projections, output projections, and well-sized attention kernels.
• Low MFU often means the GPU is stalled or underfilled. Common causes are small batch size, too few heads, too few tokens, tiny kernels, CPU synchronization, dataloading stalls, communication, memory copies, or many unfused operations.
• Being compute-bound is usually good for MFU, but only if the compute is efficient.
• The main way to improve MFU is to increase useful parallel work per GPU launch: \[ B \times H \times N_{\text{tokens}} \]

  where \(B\) is batch size, \(H\) is number of attention heads, and \(N_{\text{tokens}}\) is the number of tokens participating in the kernel.

• Activation checkpointing often increases HFU because it recomputes forward work during backward. This can make the GPU busier, but it does not necessarily improve MFU, because MFU counts the original model FLOPs, not the recomputed FLOPs.
• Recall
  MFU is not the same as wall-clock speed.
• When diagnosing low MFU,
  • Is step time is dominated by useful model kernels or by overhead?
  • If overhead dominates, fix dataloading, communication, CPU syncs, and kernel launch gaps.
  • If useful kernels dominate but MFU is still low, fix the kernel shapes by increasing batch, sequence length, heads, width, or packing.
• The goal is not just to be compute-bound. The goal is to be compute-bound on large, efficient, useful tensor-core work.

Nvidia Architecture

A100

• A100: NVIDIA Ampere-generation datacenter GPU.
• Useful baseline because many earlier large-scale ML systems and training papers used A100 clusters.
• Common memory sizes: 40 GB or 80 GB HBM2e.
• A100 80 GB has around 2 TB/s memory bandwidth.
• Supports FP16, BF16, TF32 Tensor Cores, and structured sparsity.
• Does not have Hopper's FP8 Transformer Engine.
• Main limitation compared to H100/H200: lower Tensor Core throughput and no native FP8 transformer path.

H100

• H100: NVIDIA Hopper-generation datacenter GPU.
• Major jump over A100 for transformer training and inference.
• Adds the Transformer Engine, which enables FP8 usage in transformer layers.
• Common memory size: 80 GB HBM3.
• Much higher memory bandwidth and Tensor Core throughput than A100.
• Important for modern LLM training because FP8/BF16 Tensor Core paths are much stronger.
• For dense BF16 MFU, reference the dense BF16 peak, not the “with sparsity” peak.
• first major transformer-specialized NVIDIA GPU generation.

H200

• H200: Hopper-generation GPU with upgraded memory.
• Memory: about 141 GB HBM3e.
• Memory bandwidth: about 4.8 TB/s.
• Compute is similar to H100 for many dense BF16/FP16 workloads.
• Biggest gains come when the workload is memory-capacity-bound or memory-bandwidth-bound.
• Helps with large-context attention, large-batch inference, activation-heavy training, and memory-heavy kernels.

GH200

• GH200: Grace-Hopper superchip.
• Combines a Grace CPU with a Hopper GPU.
• CPU and GPU are connected through NVLink-C2C, which is much tighter than a normal PCIe CPU-GPU connection.
• Grace CPU has large LPDDR memory. Hopper GPU has HBM.
• Useful when CPU-GPU communication, host-memory capacity, or CPU-side preprocessing matters.
• In distributed training, GH200 still behaves like a GPU node, but the CPU-GPU memory path is stronger than a standard PCIe setup.

B200

• B200: NVIDIA Blackwell-generation datacenter GPU.
• Successor generation after Hopper.
• Uses HBM3e memory.
• Has stronger Tensor Cores and better low-precision support.
• Adds stronger FP4 support, especially relevant for inference and quantized workloads.
• Compared to H200, B200 improves compute architecture and low-precision throughput.

B300

• B300: NVIDIA Blackwell Ultra GPU.
• Uses Blackwell Ultra Tensor Cores.
• Has 288 GB HBM3e per GPU in DGX B300 systems.
• DGX B300 has 8 B300 GPUs, for about 2.3 TB total GPU memory.
• NVIDIA lists DGX B300 at 72 PFLOPS FP8 training and 144 PFLOPS FP4 inference across the 8-GPU system.
• Important improvement over B200: more memory capacity and stronger low-precision AI throughput.
• Especially relevant for large-model inference, long-context workloads, and memory-heavy serving.

GB200

• GB200: Grace-Blackwell superchip.
• Combines Grace CPU with Blackwell GPU.
• The important large-system version is GB200 NVL72 which connects 72 Blackwell GPUs in one rack-scale NVLink domain.
• Designed for giant model training and inference where GPU-GPU communication becomes a major bottleneck.
• Provides very large aggregate HBM capacity and very high GPU communication bandwidth.

GB300

• GB300: Grace-Blackwell Ultra superchip / rack-scale platform.
• Uses B300-class Blackwell Ultra GPUs with Grace CPUs.
• GB300 NVL72 contains 72 B300 GPUs and 36 Grace CPUs per rack.
• Each B300 GPU has 288 GB HBM3e, so the rack has about 20 TB of HBM.
• Compared with GB200, GB300 improves memory capacity and low-precision AI throughput.