Multi-GPU training tests on GitHub Actions

Test changes to your distributed training loop before starting a long GPU run. A change to gradient accumulation can still produce a falling loss while updating the model by the wrong amount. We can catch that in CI with one optimizer step and a small, fixed batch.

This example compares a two-GPU PyTorch DistributedDataParallel (DDP) update with a single-process reference using the same samples. Each GPU processes half the batch in two microbatches. The test fails if the resulting weights differ beyond a small floating-point tolerance.

What belongs in this test

Use this when changing gradient accumulation or loss scaling in your training loop. Two independent jobs in a matrix won’t exercise the gradient exchange between those workers.

The fixture is deliberately small. We’re checking whether a training step is correct before paying for a full run. The test imports the model and training function from your repository, so a PR changing that code changes what CI runs.

Put the training code under test

Here is a small training module in src/training.py. If your project already has a model and training step, import those instead and adapt the fixture’s input shapes. Keep the production training function outside the regression test.

from contextlib import nullcontext

import torch.nn.functional as F
from torch import nn


def build_model():
    return nn.Sequential(nn.Linear(4, 8), nn.Tanh(), nn.Linear(8, 1))


def train_step(model, inputs, targets, optimizer):
    optimizer.zero_grad(set_to_none=True)
    microbatches = list(zip(inputs.chunk(2), targets.chunk(2)))
    for index, (x, y) in enumerate(microbatches):
        last = index == len(microbatches) - 1
        context = nullcontext() if last else model.no_sync()
        with context:
            loss = F.mse_loss(model(x), y) / len(microbatches)
            loss.backward()
    optimizer.step()

Compare it with an independent update

Save this as tests/ddp_regression.py. It imports the checked-out training code; the reference calculation does not call train_step. Both workers start with the same weights and get different samples, so skipping gradient synchronisation changes the result.

import copy
import os
from datetime import timedelta

import torch
import torch.distributed as dist
import torch.nn.functional as F
from torch.nn.parallel import DistributedDataParallel as DDP

from training import build_model, train_step


def main():
    local_rank = int(os.environ['LOCAL_RANK'])
    print(f'Worker {local_rank}: starting training regression', flush=True)
    torch.cuda.set_device(local_rank)
    device = torch.device('cuda', local_rank)
    dist.init_process_group(
        'nccl', device_id=device, timeout=timedelta(seconds=90)
    )
    try:
        print(f'Worker {local_rank}: NCCL ready', flush=True)
        assert dist.get_world_size() == 2, 'Run this test with two workers'
        torch.backends.cuda.matmul.allow_tf32 = False
        torch.manual_seed(7)
        model = build_model()
        reference = copy.deepcopy(model)
        inputs = torch.arange(32, dtype=torch.float32).reshape(8, 4) / 16 - 1
        targets = 0.7 * inputs[:, :1] - 0.2 * inputs[:, 1:2] + 0.1

        # Reference: one mean loss over all eight samples, one SGD update.
        reference_optimizer = torch.optim.SGD(reference.parameters(), lr=0.1)
        reference_optimizer.zero_grad(set_to_none=True)
        F.mse_loss(reference(inputs), targets).backward()
        reference_optimizer.step()

        distributed = DDP(model.to(device), device_ids=[local_rank])
        optimizer = torch.optim.SGD(distributed.parameters(), lr=0.1)
        rank = dist.get_rank()
        train_step(
            distributed,
            inputs[rank::2].to(device),
            targets[rank::2].to(device),
            optimizer,
        )

        expected = dict(reference.named_parameters())
        for name, actual in distributed.module.named_parameters():
            torch.testing.assert_close(
                actual.detach().cpu(), expected[name].detach(),
                rtol=1e-5, atol=1e-6,
                msg=f'Rank {rank}: {name} differs from the full-batch update',
            )
        print(f'Rank {rank}: distributed update matches the reference')
    finally:
        dist.destroy_process_group()


if __name__ == '__main__':
    main()

DDP averages gradients across workers. With equal local batches, mean loss and the same total samples, that should match the full-batch reference. Dividing by the number of microbatches keeps accumulation from multiplying the update. The first microbatch uses no_sync(); the final backward pass synchronises the accumulated gradients through NCCL.

Run it on two GPUs in CI

Complete the quickstart and give the Machine GitHub App access to the repository. This runner requires purchased compute credit or a paid plan. Starter credit only supports one GPU with up to 24 GB of VRAM.

Save this workflow as .github/workflows/distributed-training.yml. It runs manually and on same-repository PRs that change training code or tests. Adjust paths to match your repository. Fork PRs don’t launch these paid runners.

name: Distributed training regression
on:
  workflow_dispatch:
  pull_request:
    paths:
      - 'src/**'
      - 'tests/**'
      - '.github/workflows/distributed-training.yml'
permissions:
  contents: read
concurrency:
  group: ddp-${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
  cancel-in-progress: true
jobs:
  training-step:
    if: github.event_name == 'workflow_dispatch' || github.event.pull_request.head.repo.full_name == github.repository
    runs-on: machine/id=${{ github.run_id }}-${{ github.run_attempt }}/runner=rtx6000-12xl
    timeout-minutes: 15
    steps:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v5
        with:
          python-version: '3.11'
      - name: Install PyTorch
        run: |
          python -m venv .venv
          .venv/bin/python -m pip install torch==2.8.0 --index-url https://download.pytorch.org/whl/cu128
      - name: Test the distributed training step
        timeout-minutes: 3
        run: PYTHONPATH=src .venv/bin/torchrun --standalone --nproc-per-node=2 tests/ddp_regression.py

rtx6000-12xl gives both workers GPUs in the same VM. Check the current GPU prices before running it. The example uses the smallest multi-GPU runner in the catalog; larger runners don’t improve this test. The job has a 15-minute limit, including installation. The training test has a three-minute step limit, with a 90-second process-group timeout for collectives.

Make the assertion useful for your model

Keep the reference calculation independent of the training function under test. If both paths share the same loss-scaling bug, they can agree for the wrong reason. Check every worker’s updated parameters, rather than printing a loss and declaring success.

This fixture has equal batches, no dropout and no batch-dependent layers. It uses a fixed sample split; test your data sampler separately. BatchNorm and mixed precision need their own expected behaviour and tolerances. A passing test covers this training step, but it won’t tell you whether a model converges or whether a longer run is fast.

See PyTorch’s DDP gradient synchronisation and no_sync() reference when adapting the test to your training loop.