Checkpointing

This module implements methods for checkpointing and resuming training from a checkpoint.

class torchft.checkpointing.CheckpointTransport

Bases: Generic[T], ABC

disallow_checkpoint() → None

Called after send_checkpoint to wait for the checkpoint to be sent.

Once this returns, the state_dict may be mutated so no further data should be sent.

abstract metadata() → str

Returns a string that will be used by the remote CheckpointTransport to fetch the checkpoint.

abstract recv_checkpoint(src_rank: int, metadata: str, step: int, timeout: timedelta) → T

Receives the checkpoint from the given rank.

Parameters

• src_rank – the rank to receive the checkpoint from

• metadata – the metadata returned by the remote CheckpointTransport

• step – the step number to receive

• timeout – the timeout to wait for the checkpoint

abstract send_checkpoint(dst_ranks: List[int], step: int, state_dict: T, timeout: timedelta) → None

Sends the checkpoint, only called when there is a rank that is behind.

This may be async.

Parameters

• dst_ranks – the ranks to send to

• step – the step number to send

• state_dict – the state dict to send

• timeout – the timeout to wait for the checkpoint to be sent

shutdown(wait: bool = True) → None

Called to shutdown the checkpoint transport.

Parameters

wait – whether to wait for the transport to shutdown

class torchft.checkpointing.HTTPTransport(timeout: timedelta, num_chunks: int)

Bases: CheckpointTransport[T]

This is an HTTP server that can be used to transfer checkpoints between workers.

This allows for fast recovery of workers by fetching the current weights from an existing worker.

Parameters

• timeout – the timeout for HTTP requests

• num_chunks – the number of chunks to split the checkpoint into (0 for no chunking)

address() → str

Returns the HTTP address to fetch a checkpoint from this server. Step must be appended to the end of the address.

Format: http://host:port/checkpoint/1234

Returns

an HTTP address

allow_checkpoint(step: int) → None

Allows serving the checkpoint with the specified step number.

Parameters

step – the step number to serve

disallow_checkpoint() → None

Disallows serving the checkpoint.

All requests will block until allow_checkpoint is called.

metadata() → str

Returns a string that will be used by the remote CheckpointTransport to fetch the checkpoint.

recv_checkpoint(src_rank: int, metadata: str, step: int, timeout: timedelta) → T

Receives the checkpoint from the given rank.

Parameters

• src_rank – the rank to receive the checkpoint from

• metadata – the metadata returned by the remote CheckpointTransport

• step – the step number to receive

• timeout – the timeout to wait for the checkpoint

send_checkpoint(dst_ranks: List[int], step: int, state_dict: T, timeout: timedelta) → None

Sends the checkpoint, only called when there is a rank that is behind.

This may be async.

Parameters

• dst_ranks – the ranks to send to

• step – the step number to send

• state_dict – the state dict to send

• timeout – the timeout to wait for the checkpoint to be sent

shutdown(wait: bool = True) → None

Shutdown the server.

Process group transport

PGTransport uses its constructor timeout when a checkpoint send or receive omits timeout or passes None. An explicit per-call timeout takes precedence, including when Manager passes its configured timeout.

For example, with a configured process group pg:

from datetime import timedelta

import torch
from torchft.checkpointing.pg_transport import PGTransport

transport = PGTransport(pg, timeout=timedelta(seconds=30), device=torch.device("cpu"))
transport.send_checkpoint([1], step=1, state_dict=state_dict)
transport.send_checkpoint([1], step=2, state_dict=state_dict, timeout=timedelta(seconds=60))

The timeout applies to each communication wait, not the total transfer duration. The underlying process group’s timeout is configured separately.

class torchft.checkpointing.pg_transport.PGTransport(pg: ProcessGroup, timeout: timedelta, device: device, state_dict: Optional[Callable[[], object]] = None)

This is a checkpoint transport that uses the process group to transfer checkpoints. This allows for fast recovery of workers by fetching the current weights from an existing worker.

Parameters

• pg – the process group to use for communication

• timeout – the default timeout for each communication wait. An explicit timeout passed to send_checkpoint or recv_checkpoint overrides this value. This is not a deadline for the entire checkpoint transfer and does not change the process group’s own timeout.

• device – the device to use for tensors

• state_dict – if specified this function will be called to do an inplace receive into the returned state_dict. This is much faster than having to allocate new tensors and transferring them to the CPU.

disallow_checkpoint() → None

Called after send_checkpoint to wait for the checkpoint to be sent.

Once this returns, the state_dict may be mutated so no further data should be sent.

metadata() → str

Returns a string that will be used by the remote CheckpointTransport to fetch the checkpoint.

recv_checkpoint(src_rank: int, metadata: str, step: int, timeout: Optional[timedelta] = None) → T

Receive a checkpoint, using the constructor timeout if timeout is None.

send_checkpoint(dst_ranks: list[int], step: int, state_dict: T, timeout: Optional[timedelta] = None) → None

Send a checkpoint, using the constructor timeout if timeout is None.