Fault tolerance intra-job training

Fault tolerance intra-job training

During large-scale pretraining on a Slurm cluster, node failures (GPU errors, NVLink/network faults, ECC errors, etc.) are common and grow more frequent with node count and job duration, especially at node counts of 256-1024 which is our target. By default, a single node failure kills the whole job, which then has to requeue and restart from the last checkpoint which costs time. Instead, we want to allocate a small number of spare nodes alongside the training job (e.g. 2 spares). When a node fails, the job swaps in a spare and continues within the same Slurm allocation, avoiding a full requeue and minimizing lost progress. This thus works with a static world size.

Inspiration from Bloom: https://github.com/stas00/ml-engineering/tree/master/training/fault-tolerance#fault-tolerance

Tasks

• Survey existing solutions and assess fit for our stack (Megatron-LM, torchrun elastic launcher, shared sudoless HPC, Slurm). Start with NVIDIA Resiliency Extension (NVRx): https://github.com/NVIDIA/nvidia-resiliency-ext
• Check what Megatron-LM already supports natively (it integrates parts of NVRx / fault tolerance hooks)
• Evaluate torchrun elastic + rendezvous backend (e.g. c10d/etcd) for node replacement
  Note any requirements that conflict with a sudoless / shared HPC environment
• Design the spare-node mechanism: how spares are allocated in the Slurm job, held idle, and promoted on failure
• Implement the feature in the oellm-autoexp codebase
• Test at small scale and simulate a node failure (e.g. scancel a node, kill a process, or inject a fault) to verify recovery
• Measure recovery cost: wall-clock time to detect + swap + resume, and steps lost per failure
• Document usage (how many spares to request, config flags) for other users of the codebase