`core.dist_checkpointing.strategies.checkpointable`#

Module Contents#

Classes#

`CheckpointableShardedTensor`	ShardedTensor extension compatible with PyTorch DCP checkpointing library.
`LocalShardsContainer`	DCP compatible container for local shards.

API#

class core.dist_checkpointing.strategies.checkpointable.CheckpointableShardedTensor( data: torch.Tensor, sh_ten: core.dist_checkpointing.mapping.ShardedTensor, )#

Bases: torch.Tensor

ShardedTensor extension compatible with PyTorch DCP checkpointing library.

Implements the torch.distributed._checkpointable._Checkpointable protocol.

Initialization

__new__( data: torch.Tensor, sh_ten: core.dist_checkpointing.mapping.ShardedTensor, )#

__create_write_items__( fqn: str, sh_ten: core.dist_checkpointing.strategies.checkpointable.CheckpointableShardedTensor, index: int = None, ) → list[torch.distributed.checkpoint.planner.WriteItem]#

Simple translation from ShardedTensor offsets into DCP offsets.

Parameters:

fqn (str) – tensor FQN.
sh_ten (CheckpointableShardedTensor) – same as self
index (int) – specifies index within the LocalShardsContainer. This is an optimization hint used in DCP.

Returns:

list of DCP WriteItem metadata objects.

Return type:

List[WriteItem]

__create_chunk_list__() → list[torch.distributed.checkpoint.metadata.ChunkStorageMetadata]#

Simple translation from ShardedTensor offsets into DCP offsets.

Returns:: list of DCP ChunkStorageMetadata metadata objects.
Return type:: List[ChunkStorageMetadata]

__get_tensor_shard__( index: torch.distributed.checkpoint.metadata.MetadataIndex, ) → torch.Tensor#

Trivial implementation which simply yields the underlying tensor.

Parameters:: index (MetadataIndex) – unused
Returns:: the underlying data tensor
Return type:: Tensor

_local_chunk() → torch.distributed.checkpoint.metadata.ChunkStorageMetadata#

Describes the local shard within the global tensor.

Prepended axes (ShardedTensor.prepend_axis_num) reflect the global tensor shape on top of the local data, but a local shard always covers exactly one index along each of them. They are therefore reported with extent one in the chunk metadata, consistently with the data view returned by _shard_data_view.

_shard_data_view() → torch.Tensor#

Local data with the prepended axes explicitly represented as singleton dimensions.

This is a view of the local data, no payload is copied. DCP indexes shards with offsets expressed in the global tensor coordinates (including the prepended axes), so the local data has to be exposed with matching dimensionality.

classmethod from_sh_ten( sh_ten: core.dist_checkpointing.mapping.ShardedTensor, ) → core.dist_checkpointing.strategies.checkpointable.CheckpointableShardedTensor#

Constructor which turns a ShardedTensor into CheckpointableShardedTensor

Parameters:: sh_ten (ShardedTensor) – a sharded tensor to wrap
Returns:: wrapped ShardedTensor
Return type:: CheckpointableShardedTensor

abstractmethod classmethod __torch_dispatch__(func, types, args, kwargs=None)#: Placeholder implementation.

__repr__()#

class core.dist_checkpointing.strategies.checkpointable.LocalShardsContainer(local_shards: list[torch.Tensor])#

Bases: torch.Tensor

DCP compatible container for local shards.

PyTorch DCP requires a single tensor per rank for a given global tensor FQN. This class acts as a container allowing multiple checkpointable shards per rank.