Fault tolerance and elastic training

Fault tolerance and elastic training#

Large-scale jobs run for days across thousands of GPUs, where hardware faults, NIC flaps, and node loss are routine. This guide covers the mechanisms Primus exposes to survive and recover from failures: graceful exit + checkpoint-based resume, Megatron’s fault-tolerance package and in-process restart, and TorchTitan’s torchft-based elastic training. Parameters are grounded in primus/configs/modules/megatron/trainer_base.yaml, primus_megatron_module.yaml, and primus/configs/modules/torchtitan/pre_trainer.yaml.

The foundation of all recovery is checkpointing—read Checkpoint management first.

1. The recovery model#

There are three layers, from simplest to most advanced:

Checkpoint + restart—periodically save state; on failure, relaunch the job and resume from the last checkpoint. Works on every backend; relies on the scheduler (Slurm --requeue, Kubernetes restart policy) to relaunch.
Graceful exit—detect a signal or time/iteration budget, save a final checkpoint, and exit cleanly so the restart resumes with no lost work.
In-job fault tolerance / elastic—detect a failed rank and restart in-process (Megatron) or continue with a reduced/replaced replica group (TorchTitan + torchft) without tearing down the whole job.

2. Graceful exit and auto-resume (Megatron)#

Controls in trainer_base.yaml let a run stop cleanly at a boundary so the next launch resumes seamlessly:

Parameter	Default	Purpose
`exit_signal_handler`	`false`	Install a signal handler that saves a checkpoint and exits gracefully on SIGTERM (e.g. Slurm preemption).
`exit_duration_in_mins`	`null`	Exit (after saving) once the job has run this many minutes—useful to fit scheduler time limits.
`exit_interval`	`null`	Exit after this many iterations.
`adlr_autoresume`	`false`	Enable ADLR auto-resume integration.
`adlr_autoresume_interval`	`1000`	Iterations between auto-resume checks.

Primus-level continuation (primus_megatron_module.yaml):

Parameter	Default	Purpose
`auto_continue_train`	`false`	Automatically continue from the latest checkpoint in the experiment’s save directory on relaunch.
`disable_last_saving`	`false`	Disable the final end-of-run checkpoint save (leave `false` so resume points exist).

Pattern: enable exit_signal_handler + exit_duration_in_mins (or rely on preemption signals), set a reasonable checkpoint save_interval, and turn on auto_continue_train so requeued jobs pick up where they left off.

3. Megatron fault-tolerance package and in-process restart#

Megatron integrates an optional fault-tolerance package and in-process restart (trainer_base.yaml):

Parameter	Default	Purpose
`enable_ft_package`	`false`	Enable the Megatron fault-tolerance package (rank monitoring / heartbeat).
`calc_ft_timeouts`	`false`	Auto-calculate fault-tolerance timeouts from observed step times.
`run_workload_inspector_server`	`false`	Run the workload inspector server for health/diagnostics.
`inprocess_restart`	`false`	Restart failed ranks in process to avoid a full job teardown.

In-process restart reduces recovery time by re-initializing the process group and reloading state without re-scheduling the whole allocation. Combine with frequent checkpoints so the restarted ranks have a recent resume point.

4. Numerical safety nets (Megatron)#

Detecting corruption early prevents wasted compute and divergence (trainer_base.yaml):

Parameter	Default	Purpose
`check_for_nan_in_loss_and_grad`	`true`	Abort/handle on NaN/Inf in loss or gradients.
`check_for_spiky_loss`	`false`	Detect anomalous loss spikes.
`check_for_large_grads`	`false`	Detect abnormally large gradients.
`decrease_batch_size_if_needed`	`false`	Reduce batch size when needed instead of failing.

These don’t recover from hardware faults but stop a corrupted run before it pollutes downstream checkpoints.

5. Elastic training with torchft (TorchTitan)#

TorchTitan supports semi-synchronous, replica-based fault tolerance via torchft. Configured under fault_tolerance: in primus/configs/modules/torchtitan/pre_trainer.yaml:

Parameter	Default	Purpose
`enable`	`false`	Enable torchft fault tolerance.
`process_group`	`gloo`	Process group backend for fault-tolerance coordination.
`process_group_timeout_ms`	`10000`	Coordination timeout (ms).
`replica_id`	`0`	This replica’s ID.
`group_size`	`0`	Number of replica groups (`0` = auto).
`min_replica_size`	`1`	Minimum replicas required to keep training.
`semi_sync_method`	`null`	Semi-synchronous algorithm (e.g. DiLoCo-style), `null` = standard.

With replica groups, the loss of one replica can be tolerated as long as min_replica_size is still satisfied—training continues while the failed replica recovers/rejoins, rather than crashing the whole job.

Install the optional dependencies before enabling (requirements-torchft.txt):

pip install -r requirements-torchft.txt   # torchft-nightly + OpenTelemetry exporters

TorchTitan also exposes communication timeouts under comm: (init_timeout_seconds: 300, train_timeout_seconds: 100) that govern how long collectives wait before declaring a fault.

6. Scheduler integration#

In-job mechanisms still need the scheduler to relaunch on full-job failure:

Slurm—submit with --requeue so preempted/failed jobs are re-queued; pair with exit_signal_handler to checkpoint on SIGTERM. See Deployment.
Kubernetes—use a restart policy / operator that recreates pods; mount checkpoint storage on a shared/persistent volume.
Shared checkpoint storage—all ranks must read the same checkpoint directory after relaunch (NFS, Lustre, or object storage). See Checkpoint management.

7. Recommended setup#

Always checkpoint—set a save_interval matched to your mean-time-between-failures; use async/distributed checkpointing to keep overhead low.
Exit cleanly—exit_signal_handler: true (+ exit_duration_in_mins for time-boxed allocations).
Resume automatically—auto_continue_train: true (Megatron) and --requeue (Slurm).
Reduce recovery time at scale—enable_ft_package + inprocess_restart (Megatron) or torchft replica groups (TorchTitan).
Guard numerics—keep check_for_nan_in_loss_and_grad on; consider spiky/large-grad checks for unstable configs.