How do I set up a distributed training stack for hyperparameter sweeps across multiple GPUs?

A good distributed training stack for hyperparameter sweeps usually has four layers:

1. Training code that can run on 1 GPU or many GPUs
2. Distributed launch to use multiple GPUs per trial
3. Sweeper/orchestrator to run many trials in parallel
4. Storage + logging so trials are tracked and recoverable

Below is a practical setup that works well for PyTorch, but the same structure applies to JAX/TF too.

---

1) Choose the right architecture

There are two common ways to use multiple GPUs for sweeps:

Option A: One trial = one multi-GPU distributed job

• Each hyperparameter trial uses several GPUs
• Best when each model is large or training is slow
• Example: each sweep run uses 4 GPUs via DDP

Option B: One GPU per trial, many trials in parallel

• Each trial uses a single GPU
• Best when the model fits easily on one GPU
• Higher sweep throughput

Hybrid approach

• Run multiple trials in parallel, and each trial uses multiple GPUs
• This is common for large models and large sweeps, but needs careful resource scheduling

---

2) Make your training script distributed-ready

If you're using PyTorch, make the script compatible with DistributedDataParallel (DDP).

Key requirements

• Initialize distributed backend
• Set device based on local rank
• Use DistributedSampler for the dataset
• Only log/save from rank 0
• Save checkpoints in a way that supports resuming

Minimal DDP skeleton

import os
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def setup_distributed():
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    return local_rank

def cleanup():
    dist.destroy_process_group()

def main():
    local_rank = setup_distributed()
    device = torch.device(f"cuda:{local_rank}")

    model = MyModel().to(device)
    model = DDP(model, device_ids=[local_rank])

    train_loader = get_loader(distributed=True)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

    for epoch in range(num_epochs):
        train_one_epoch(model, train_loader, optimizer, device)
        if dist.get_rank() == 0:
            save_checkpoint(model, optimizer, epoch)

    cleanup()

if __name__ == "__main__":
    main()

Distributed dataloader

Use a DistributedSampler:

from torch.utils.data.distributed import DistributedSampler

sampler = DistributedSampler(dataset)
loader = DataLoader(dataset, sampler=sampler, batch_size=batch_size)

for epoch in range(num_epochs):
    sampler.set_epoch(epoch)

---

3) Launch distributed training correctly

The standard launcher in PyTorch is torchrun.

Single-node, 4 GPUs

torchrun --nproc_per_node=4 train.py --config config.yaml

Multi-node example

torchrun \
  --nnodes=2 \
  --nproc_per_node=8 \
  --node_rank=0 \
  --master_addr=10.0.0.1 \
  --master_port=29500 \
  train.py

If you’re sweeping, each trial can be launched by the scheduler with its own config file or CLI args.

---

4) Use a sweep manager

You need something to generate trial configs and schedule jobs.

Good choices

• Weights & Biases Sweeps: easy, good UI, integrates with PyTorch
• Ray Tune: strong for distributed sweeps and parallel scheduling
• Optuna: good for optimization, can be paired with Ray or your own executor
• Hydra + your own launcher: flexible and lightweight
• Slurm + Bash/Python: best for HPC clusters

If you want simple and robust

A common stack is:

• Hydra for config management
• PyTorch DDP for distributed training
• Ray Tune or W&B Sweeps for hyperparameter search

---

5) Recommended stack patterns

Pattern 1: Local workstation or single node

Best for 1 machine with multiple GPUs.

Stack

• PyTorch DDP
• Hydra configs
• Ray Tune or W&B sweeps

Flow

• Sweeper generates a config
• One trial launches with torchrun --nproc_per_node=N
• Each trial uses all GPUs, or subset if you manage allocation yourself

---

Pattern 2: Cluster/HPC with Slurm

Best for compute cluster environments.

Stack

• Slurm
• PyTorch DDP
• Hydra
• Optuna/Ray/W&B

Flow

• Each sweep trial submitted as a Slurm job
• Job allocates --gres=gpu:4
• Trial runs torchrun --nproc_per_node=4
• Experiment tracking goes to W&B/MLflow

This is often the most reliable for large sweeps.

---

Pattern 3: Ray for scheduling + training

Best if you want automatic parallelism management.

Stack

• Ray Tune
• PyTorch Lightning or PyTorch DDP
• Object store for results
• W&B/MLflow for tracking

Flow

• Ray schedules trials across GPUs
• You specify GPU per trial
• Ray handles placement and parallel execution

Example with Ray Tune:

from ray import tune

tune.run(
    train_fn,
    config={
        "lr": tune.loguniform(1e-5, 1e-2),
        "batch_size": tune.choice([32, 64, 128]),
    },
    resources_per_trial={"cpu": 4, "gpu": 2},
)

If each trial needs multiple GPUs, Ray can allocate them directly.

---

6) Add experiment tracking

Use a tracker so each trial is reproducible and comparable.

Popular options

• Weights & Biases
• MLflow
• TensorBoard
• Neptune

Log:

• Hyperparameters
• Metrics
• GPU utilization
• Training time
• Checkpoint path
• Seed
• Git commit hash

Best practice

Only let rank 0 write logs and save checkpoints to avoid duplication.

---

7) Handle checkpoints and resume logic

For sweeps, failed jobs will happen.

Save:

• Model state dict
• Optimizer state
• Scheduler state
• Epoch / step
• Random seeds
• Config

Resume:

• Detect latest checkpoint
• Restart from saved state
• Ensure the same config is used

This is especially important when jobs are preempted on shared clusters.

---

8) Be careful with resource allocation

The most common failure mode in sweeps is oversubscribing GPUs.

Rules of thumb

• One process per GPU for DDP
• Don’t start multiple trials on the same GPU unless you intentionally want that
• Keep CPU and dataloader workers balanced
• Use CUDA_VISIBLE_DEVICES or scheduler allocations to isolate resources

Example

If you have 8 GPUs and each trial uses 2 GPUs:

• Run 4 trials in parallel
• Each trial launches 2 DDP processes

---

9) Suggested practical setup

If you want something that “just works”:

For a single machine

• Hydra for configs
• PyTorch DDP with torchrun
• W&B Sweeps or Ray Tune
• MLflow or W&B for tracking

For a cluster

• Slurm
• Hydra
• PyTorch DDP
• W&B or MLflow
• Optional: Ray Tune if you want smarter scheduling

---

10) Example workflow

Training code

• train.py accepts CLI config
• Uses DDP when WORLD_SIZE > 1
• Logs only on rank 0
• Saves checkpoints to a trial-specific folder

Sweep config

• Defines search space for lr, weight_decay, dropout, etc.

Launcher

• Submits one job per trial
• Each job runs torchrun
• Job writes results to tracker

---

11) A very simple example with Hydra + torchrun

config.yaml

train:
  lr: 1e-4
  batch_size: 64
  epochs: 10

launch

torchrun --nproc_per_node=4 train.py train.lr=3e-4 train.batch_size=128

You can generate many such runs with a script, W&B sweep agent, or Ray Tune.

---

12) What I’d recommend by use case

If you’re on one workstation

Use:

• torchrun
• Hydra
• W&B Sweeps

If you’re on a cluster

Use:

• Slurm
• torchrun
• Hydra
• W&B/MLflow
• Optionally Optuna for search logic

If you want automatic distributed scheduling

Use:

• Ray Tune
• PyTorch DDP
• W&B or MLflow

---

13) Common pitfalls

• Forgetting DistributedSampler.set_epoch(epoch)
• Logging from every rank
• Saving checkpoints from every rank
• Using wrong batch size per GPU
• Not scaling learning rate appropriately
• Oversubscribing GPUs across sweep trials
• Non-determinism due to inconsistent seeding

---

If you want, I can give you one of these next:

1. A concrete PyTorch + Hydra + DDP + W&B template
2. A Ray Tune distributed sweep example
3. A Slurm job script for multi-GPU sweeps
4. A PyTorch Lightning version