Easy Per Step Fault Tolerance for PyTorch
| Documentation | Poster | Design Doc |
This repository implements techniques for doing a per-step fault tolerance so you can keep training if errors occur without interrupting the entire training job.
This is based on the large scale training techniques presented at PyTorch Conference 2024.
Overview
torchft is designed to provide the primitives required to implement fault tolerance in any application/train script as well as the primitives needed to implement custom fault tolerance strategies.
Out of the box, torchft provides the following algorithms:
- Fault Tolerant DDP
- Fault Tolerant HSDP: fault tolerance across the replicated dimension with any mix of FSDP/TP/etc across the other dimensions.
- LocalSGD
- DiLoCo
To implement these, torchft provides some key reusable components:
- Coordination primitives that can determine which workers are healthy via heartbeating on a per-step basis
- Fault tolerant ProcessGroup implementations that report errors sanely and be reinitialized gracefully.
- Checkpoint transports that can be used to do live recovery from a healthy peer when doing scale up operations.
The following component diagram shows the high level components and how they relate to each other:
See torchft's documentation for more details.
Examples
torchtitan (Fault Tolerant HSDP)
torchtitan provides an out of the box fault tolerant HSDP training loop built on top of torchft that can be used to train models such as Llama 3 70B.
It also serves as a good example of how you can integrate torchft into your own training script for use with HSDP.
See torchtitan's documentation for end to end usage.
Fault Tolerant DDP
We have a minimal DDP train loop that highlights all of the key components in torchft.
See train_ddp.py for more info.
DiLoCo
LocalSGD and DiLoCo are currently experimental.
See the diloco_train_loop/local_sgd_train_loop tests for an example on how to integrate these algorithms into your training loop.
Design
torchft is designed to allow for fault tolerance when using training with replicated weights such as in DDP or HSDP (FSDP with DDP).
See the design doc for the most detailed explanation.
Lighthouse
torchft implements a lighthouse server that coordinates across the different replica groups and then a per replica group manager and fault tolerance library that can be used in a standard PyTorch training loop.
This allows for membership changes at the training step granularity which can greatly improve efficiency by avoiding stopping the world training on errors.
Fault Tolerant HSDP Algorithm
torchft provides an implementation of a fault tolerant HSDP/DDP algorithm. The following diagram shows the high level operations that need to happen in the train loop to ensure everything stays consistent during a healing operation.
See the design doc linked above for more details.
Installing from PyPI
We have nighty builds available at https://pypi.org/project/torchft-nightly/
To install torchft with minimal dependencies you can run:
pip install torchft-nightly
If you want all development dependencies you can install:
pip install torchft-nightly[dev]
Installing from Source
Prerequisites
Before proceeding, ensure you have the following installed:
- Rust (with necessary dependencies)
protobuf-compilerand the corresponding development package for Protobuf.- PyTorch 2.7 RC+ or Nightly
Note that the Rust versions available in many conda environments may be outdated. To install the latest version of Rust, we recommend downloading it directly from the official website as shown in the below command:
curl --proto '=https' --tlsv1.2 https://sh.rustup.rs -sSf | sh
To install the required packages on a Debian-based system (such as Ubuntu) using apt, run:
sudo apt install protobuf-compiler libprotobuf-dev
or for a Red Hat-based system, run:
sudo dnf install protobuf-compiler protobuf-devel
Installation
pip install .
This uses pyo3+maturin to build the package, you'll need maturin installed.
If the installation command fails to invoke cargo update due to an inability to fetch the manifest, it may be caused by the proxy, proxySSLCert, and proxySSLKey settings in your .gitconfig file affecting the cargo command. To resolve this issue, try temporarily removing these fields from your .gitconfig before running the installation command.
To install in editable mode w/ the Rust extensions and development dependencies, you can use the normal pip install command:
pip install -e '.[dev]'
Usage
Lighthouse
The lighthouse is used for fault tolerance across replicated workers (DDP/FSDP) when using synchronous training.
You can start a lighthouse server by running:
RUST_BACKTRACE=1 torchft_lighthouse --min_replicas 1 --quorum_tick_ms 100 --join_timeout_ms 10000
Example Training Loop (DDP)
See train_ddp.py for the full example.
Invoke with:
TORCHFT_LIGHTHOUSE=http://localhost:29510 torchrun --master_port 29501 --nnodes 1 --nproc_per_node 1 train_ddp.py
train.py:
from torchft import Manager, DistributedDataParallel, Optimizer, ProcessGroupGloo
manager = Manager(
pg=ProcessGroupGloo(),
load_state_dict=...,
state_dict=...,
)
m = nn.Linear(2, 3)
m = DistributedDataParallel(manager, m)
optimizer = Optimizer(manager, optim.AdamW(m.parameters()))
for i in range(1000):
batch = torch.rand(2, 2, device=device)
optimizer.zero_grad()
out = m(batch)
loss = out.sum()
loss.backward()
optimizer.step()
Running DDP
After starting the lighthouse server by running:
RUST_BACKTRACE=1 torchft_lighthouse --min_replicas 1 --quorum_tick_ms 100 --join_timeout_ms 10000
A test DDP script can be launched with torchX with:
torchx run
Or Diloco with:
USE_STREAMING=True torchx run ./torchft/torchx.py:hsdp --script='train_diloco.py'
See .torchxconfig, torchx.py and the torchX documentation to understand how DDP is being ran.
torchx.py could also launch HSDP jobs when workers_per_replica is set > 1, if the training script supports it. For an example HSDP training implementation with torchFT enabled, see torchtitan.
Alternatively, to test on a node with two GPUs, you can launch two replica groups running train_ddp.py by:
On shell 1 (one replica groups starts initial training):
export REPLICA_GROUP_ID=0
export NUM_REPLICA_GROUPS=2
CUDA_VISIBLE_DEVICES=0 TORCHFT_LIGHTHOUSE=http://localhost:29510 torchrun --master_port=29600 --nnodes=1 --nproc_per_node=1 -- train_ddp.py
On shell 2 (a second replica group joins):
export REPLICA_GROUP_ID=1
export NUM_REPLICA_GROUPS=2
CUDA_VISIBLE_DEVICES=1 TORCHFT_LIGHTHOUSE=http://localhost:29510 torchrun --master_port=29601 --nnodes=1 --nproc_per_node=1 -- train_ddp.py
By observing the outputs from both shells, you should observe process group reconfiguration and live checkpoint recovery.
Example Parameter Server
torchft has a fault tolerant parameter server implementation built on it's reconfigurable ProcessGroups. This does not require/use a Lighthouse server.
See parameter_server_test.py for an example.
Contributing
We welcome PRs! See the CONTRIBUTING file.
License
torchft is BSD 3-Clause licensed. See LICENSE for more details.
1 maintainer
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distributions
Built Distributions
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file torchft_nightly-2026.8.8-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.
File metadata
- Download URL: torchft_nightly-2026.8.8-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
- Upload date:
- Size: 2.5 MB
- Tags: CPython 3.13, manylinux: glibc 2.17+ x86-64
- Uploaded using Trusted Publishing? No
- Uploaded via:
twine/7.0.0 CPython/3.10.20
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
477fd2101e2dc4c1e39252c80c96928420bd929eb099edd9e3c4c662a3456f1b
|
|
| MD5 |
4bad7ef503310bc45073f1c28e7293f9
|
|
| BLAKE2b-256 |
8ad43e3a6c67e85179a5b7456e337128b60875092687870e5c293ceb2c27f515
|
File details
Details for the file torchft_nightly-2026.8.8-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.
File metadata
- Download URL: torchft_nightly-2026.8.8-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
- Upload date:
- Size: 2.5 MB
- Tags: CPython 3.12, manylinux: glibc 2.17+ x86-64
- Uploaded using Trusted Publishing? No
- Uploaded via:
twine/7.0.0 CPython/3.10.20
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
cd52034f42fac0c481e5e7c0a60133cf2a0e258d221aca3a25e1759a75f1bb59
|
|
| MD5 |
36ce58827e33e1877c9242e6278539d7
|
|
| BLAKE2b-256 |
9328d88273efb938bb1b3b4b260d821f3618a5506f0f33c41f61975cc669d093
|
File details
Details for the file torchft_nightly-2026.8.8-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.
File metadata
- Download URL: torchft_nightly-2026.8.8-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
- Upload date:
- Size: 2.5 MB
- Tags: CPython 3.11, manylinux: glibc 2.17+ x86-64
- Uploaded using Trusted Publishing? No
- Uploaded via:
twine/7.0.0 CPython/3.10.20
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
6d3c6a5e8dfe6b50bcf5001d2d5ab1dac22687e01e1cc352aecebff0b531d0fb
|
|
| MD5 |
f3101f4bb8962247141b5b6105d8261b
|
|
| BLAKE2b-256 |
977b597cdd2c31aa4e5659f71d362a606d9ccebeb74ec5ed3f5ceee79584e8ac
|
File details
Details for the file torchft_nightly-2026.8.8-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.
File metadata
- Download URL: torchft_nightly-2026.8.8-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
- Upload date:
- Size: 2.5 MB
- Tags: CPython 3.10, manylinux: glibc 2.17+ x86-64
- Uploaded using Trusted Publishing? No
- Uploaded via:
twine/7.0.0 CPython/3.10.20
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
09a5acb61aed74378226002d6c5a13dad66c11577b98e28adcf715517405e74f
|
|
| MD5 |
87e87a7f13fd8dfdd5fc6946602b5315
|
|
| BLAKE2b-256 |
0f553ab9c54b2d8ad612d553edeec437fd606474d89b5c81f5107cc15956c480
|
File details
Details for the file torchft_nightly-2026.8.8-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.
File metadata
- Download URL: torchft_nightly-2026.8.8-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
- Upload date:
- Size: 2.5 MB
- Tags: CPython 3.9, manylinux: glibc 2.17+ x86-64
- Uploaded using Trusted Publishing? No
- Uploaded via:
twine/7.0.0 CPython/3.10.20
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
d052966da859fe57d223a772891d5d353b7a0779ecf006af393db3f4373f1a82
|
|
| MD5 |
b1580d8efbc461eb166043c1a5ab5367
|
|
| BLAKE2b-256 |
9fb0ded71203b1e8009dfb1e575e487b0301dc04153403b1b819dd35ccd09976
|
