SR-Forge

Distributed & Multi-GPU Training¶

SR-Forge trains on one GPU, several GPUs in one machine, or several machines — with the same training script and the same config. Distribution is not a mode you program against; it is an environment your script runs in. The framework detects that environment once, at setup_device, and every component that must care (data loading, gradient sync, metric aggregation, logging, checkpointing) consults the same detected strategy internally.

The practical consequence: if you have a working single-GPU setup, you already have a working distributed setup. What changes is the launch command, not your code.

At a glance¶

One script — the standard scaffold from Writing Scripts — four ways to launch it:

Single GPUMulti-GPU, one machine (DDP)Multi-GPU, one machine (DataParallel)Multi-node (Linux, GPU)

python scripts/train.py --config-name my-config

Nothing to configure. system.device: cuda:0 (or cpu) as always.

torchrun --standalone --nproc_per_node=2 scripts/train.py --config-name my-config

Recommended. One process per GPU, gradients all-reduced by DistributedDataParallel. --standalone handles rendezvous on a single node — no addresses to configure.

system:
  device: [0, 1]

python scripts/train.py --config-name my-config

The historical zero-launcher path: a device list in the config selects torch.nn.DataParallel (single process, replicated forward). Kept for backward compatibility; prefer DDP for anything serious — it is faster and identical to the multi-node path.

# On node 0 (rank 0, e.g. 10.0.0.1):
torchrun --nnodes=2 --node_rank=0 --nproc_per_node=1 \
         --master_addr=10.0.0.1 --master_port=29500 \
         scripts/train.py --config-name my-config

# On node 1:
torchrun --nnodes=2 --node_rank=1 --nproc_per_node=1 \
         --master_addr=10.0.0.1 --master_port=29500 \
         scripts/train.py --config-name my-config

NCCL backend, one process per GPU per node. Your job dispatcher (Kubernetes, SkyPilot, SLURM, a shell loop over SSH) is responsible for running this command on each node — see Launchers and dispatchers.

training.batch_size is PER RANK

Under DDP every rank steps on its own batch, so the effective batch is world_size × batch_size. Launching your single-GPU config unchanged on 2 nodes silently doubles the batch the optimizer sees. To keep the training recipe identical, divide: a designed batch of 32 on 2 ranks means training.batch_size: 16.

How the strategy is chosen¶

setup_device(model, cfg.system.device, dataset) resolves a TrainingStrategy exactly once per process and caches it; every later consumer (DataLoaderFactory, runners, trainer, tracker, hooks) fetches the same instance via get_active_strategy().

flowchart TD
    A["setup_device(model, device, dataset)"] --> B{"WORLD_SIZE > 1<br/>in the environment?"}
    B -- "yes (torchrun set it)" --> C["DDPStrategy<br/>one process per GPU,<br/>gradient all-reduce"]
    B -- no --> D{"device is a list?<br/>e.g. [0, 1]"}
    D -- yes --> E["DataParallelStrategy<br/>single process,<br/>replicated forward"]
    D -- no --> F["SingleDeviceStrategy<br/>the classic path"]

Strategy	Selected when	Processes	Model wrapping	Data
`SingleDeviceStrategy`	default	1	none	full dataset
`DataParallelStrategy`	`system.device` is a list	1	`nn.DataParallel`	full dataset, batch split across GPUs per step
`DDPStrategy`	launched by `torchrun` (`WORLD_SIZE > 1`)	one per GPU	`DistributedDataParallel`	sharded — each rank sees `1/world_size` of every epoch

Topology under DDP comes exclusively from the standard environment variables (RANK, LOCAL_RANK, WORLD_SIZE, MASTER_ADDR, MASTER_PORT) that torchrun sets — the framework never guesses it. The backend is NCCL when a CUDA-capable NCCL build is available, otherwise gloo (CPU); the device follows the backend.

What actually changes under DDP¶

You don't have to act on any of these — this is what the framework does on your behalf, and what you should expect to observe.

Data is sharded, not duplicated. The loader factory swaps your sampler for a DistributedSampler: each rank iterates a disjoint 1/world_size slice of the dataset each epoch, re-shuffled per epoch (the trainer calls set_epoch for you). Your shuffle: true config is honored through the sampler. Validation is sharded the same way.

Gradients are averaged every step. Each rank computes forward/backward on its shard; DDP all-reduces gradients before the optimizer step, so all replicas stay bit-identical. This is why the effective batch is world_size × batch_size.

Metrics are globally reduced. At epoch end the trainer count-weight- averages every metric across ranks before firing the epoch-finished event, so the numbers in W&B and in LossLogger output are fleet-global means, not rank 0's shard. Band-keyed metric dicts keep their structure through the reduction.

Side effects happen once. The W&B run exists on rank 0 only (other ranks run with the tracker disabled but stay API-consistent for resume checks). Checkpoints, logged images, and epoch metrics are written once per fleet — see Writing distribution-safe hooks for the contract your own hooks should follow.

Progress bars are per-rank, tagged. Every rank prints its own bar over its own shard, prefixed [r0], [r1], … so interleaved log lines stay attributable. Loss values shown per-rank are local; trust the epoch summary (reduced) for the global picture.

GANs work without special handling. setup_device wraps each trainable sub-module (generator, discriminator) separately — DDP only synchronises a module's own forward, and the GAN runner drives each part through its own step, so a single wrap of the composite would never sync gradients. The runner also freezes the discriminator during the generator step so DDP's reducer sees exactly one backward per module per iteration, and modules are wrapped with static_graph=True so parameters reused within one forward (per-band shared encoders and the like) don't trip DDP's readiness tracking. If you write a custom composite model, implement trainable_modules() returning {attribute_name: module} for each independently-stepped part — that is the whole contract.

Data loading and workers¶

With num_workers > 0 the loader factory applies a worker policy tuned for long trainings:

forkserver start method (where the platform offers it): workers are spawned from a pristine process instead of forking a thread-heavy parent, which makes an entire class of fork-after-threads deadlocks impossible by construction.
persistent_workers: true: workers survive across epochs instead of being re-created.
pin_memory is stripped automatically when training on CPU.

Anything you set explicitly in the loader config wins over the policy — pass your own multiprocessing_context or persistent_workers to override.

Datasets must pickle under forkserver

Forkserver workers receive the dataset by pickling. Datasets holding open file handles, thread locks, or huge in-memory state either fail to pickle or transfer slowly. Keep dataset objects lightweight (paths and coordinates, not decoded arrays); if a dataset genuinely can't pickle, set multiprocessing_context: fork explicitly (Linux) or num_workers: 0 — nothing else changes.

Two environment switches help in the field: DEBUG_HANG=1 makes the main process and every worker dump all thread stacks periodically (first-step hangs stop being guesswork), and SRFORGE_CV2_THREADS=1 re-enables OpenCV's internal thread pool that srforge.init() pins to zero for fork safety.

Checkpointing, resume, and reproducibility¶

Checkpoints are written by rank 0 only; every rank restores identically at startup, so a resumed fleet is consistent by construction. With a fixed tracker run_id (resume: allow), every relaunch continues the same W&B run — crash recovery never spawns a new experiment.

Resume is epoch-aligned by design: the epoch saver writes at epoch boundaries, and with per-epoch seeded sampling a resumed run replays exactly the trajectory an uninterrupted run would have produced — the run's history stays reproducible no matter how many times it was interrupted. A mid-epoch checkpoint scheme would resume faster after a kill, but it splices step-N weights onto a fresh replay of the epoch — a trajectory no clean run can produce — which is why SR-Forge deliberately does not ship one.

Writing distribution-safe hooks¶

Runner-level hook points fire on every rank (control hooks — gradient clipping, step ratios, loss injection — must run everywhere or replicas diverge). Side effects must not. The contract, in full detail in Hooks:

You see on the class/method	It means
`@main_process_only` on the class	every handler runs once per fleet — the whole class is a rank-0 side effect
`@main_process_only` on a method	mixed class: this handler is rank-0, undecorated siblings run everywhere on purpose
no decorator	runs on every rank, deliberately

from srforge.distributed import main_process_only
from srforge.training.hooks import Hook, hooks_into

@main_process_only          # whole class is a rank-0 side effect
class MyMediaLogger(Hook):
    @hooks_into("on_post_step")
    def maybe_write(self, ctx):
        ...

All gates are call-time checks and always-true no-ops on single-device runs — decorated hooks cost nothing outside DDP.

Launchers and dispatchers¶

SR-Forge deliberately knows nothing about job dispatchers. Its entire distributed contract is the standard torchrun environment — any system that runs torchrun (or sets RANK / LOCAL_RANK / WORLD_SIZE / MASTER_ADDR / MASTER_PORT itself) works unmodified:

Dispatcher	Status
`torchrun` directly (any SSH-reachable machines)	works
Kubernetes + SkyPilot, Kubeflow `PyTorchJob`	works — they run/emulate torchrun
SLURM via `sbatch` + `torchrun`	works
Ray Train, SageMaker (torch launcher)	works — they set the torch env vars
SLURM `srun` direct, `mpirun`, LSF `jsrun`	needs a 5-line env shim (below)

# SLURM direct-srun shim: map scheduler vars to the torch contract
export RANK=$SLURM_PROCID
export LOCAL_RANK=$SLURM_LOCALID
export WORLD_SIZE=$SLURM_NTASKS
export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n1)
export MASTER_PORT=29500
python scripts/train.py --config-name my-config

Useful NCCL knobs for multi-node runs (set them in your dispatcher's job spec, not in SR-Forge):

NCCL_SOCKET_IFNAME=eth0   # pin rendezvous/transport to the fast NIC on multi-NIC nodes
NCCL_IB_DISABLE=1         # force TCP when InfiniBand/RoCE userspace (rdma-core) is absent
NCCL_DEBUG=INFO           # verbose transport selection when debugging connectivity

Troubleshooting¶

Symptom	Cause	Fix
Loss curves shifted after moving to 2 nodes	effective batch doubled	halve `training.batch_size` (it is per rank)
`Expected to mark a variable ready only once`	custom composite model wrapped as one module, or parameters reused across steps	implement `trainable_modules()` on the model; the framework then wraps each part with `static_graph=True`
Rank hangs at `init_process_group`	ranks can't reach `MASTER_ADDR:MASTER_PORT`, or rendezvous picked the wrong NIC	verify connectivity; set `NCCL_SOCKET_IFNAME`
`NCCL WARN Bootstrap: no socket interface found`	the named interface doesn't exist on this node	fix `NCCL_SOCKET_IFNAME` for that node's NIC naming
First training step hangs forever with workers	a library started threads before workers forked	default policy already uses `forkserver` and pins OpenCV threads; set `DEBUG_HANG=1` to see every thread's stack
Workers crash with pickling errors	dataset not picklable under `forkserver`	slim the dataset state, or set `multiprocessing_context: fork` / `num_workers: 0`
Same image/metric logged twice per step	a custom side-effect hook without a gate	add `@main_process_only` (class or method)
Windows multi-node picks CPU	no NCCL on Windows → gloo backend, and the device follows the backend	expected: multi-node GPU training requires Linux/NCCL; Windows is fully supported single-machine

Current limits¶

Honest boundaries of what ships today: no FSDP or DeepSpeed (model must fit on one GPU; strategy: fsdp fails loudly rather than degrading), no self-spawning launcher (DDP always enters via torchrun or equivalent env), no native SLURM/MPI env auto-detection (use the shim above), and multi-node on Windows runs gloo/CPU only. The W&B-based warm-start/resume paths assume network access to the W&B API; fully air-gapped clusters should keep checkpoints on a shared filesystem and pass local paths.