TensorPlay
API reference
latest (dev)
latest (dev)2026-10-08
Copy
View Markdown

Latest development documentation · Updated 2026-10-08

tensorplay.distributed.checkpoint API

Functions 12

#

async_save

functionFull reference ↗
tensorplay.distributed.checkpoint.async_save(state_dict, *, checkpoint_id=None, storage_writer=None, planner=None, process_group=None, async_checkpointer_type: AsyncCheckpointerType = AsyncCheckpointerType.THREAD, async_stager: AsyncStager | None = None, no_dist=False, use_collectives=True) → Future[Any] | AsyncSaveResponse[source]

Stage the input and execute the checkpoint write asynchronously.

#

load

functionFull reference ↗
tensorplay.distributed.checkpoint.load(state_dict, *, checkpoint_id=None, storage_reader=None, planner=None, process_group=None, no_dist=False) → None[source]

Load checkpoint values into an existing state dictionary.

#

save

functionFull reference ↗
tensorplay.distributed.checkpoint.save(state_dict, *, checkpoint_id=None, storage_writer=None, planner=None, process_group=None, no_dist=False, use_collectives=True) → Any[source]

Save a state dictionary with coordinated metadata commit.

Classes 36

#

AsyncSaveResponse

classFull reference ↗
class tensorplay.distributed.checkpoint.AsyncSaveResponse(staging_completion: 'Future[None]', upload_completion: 'Future[Any]')[source]
#

ChunkStorageMetadata

classFull reference ↗
class tensorplay.distributed.checkpoint.ChunkStorageMetadata(offsets: 'tuple[int, ...]', sizes: 'tuple[int, ...]')[source]
#

DefaultLoadPlanner

classFull reference ↗
class tensorplay.distributed.checkpoint.DefaultLoadPlanner(flatten_state_dict: bool = True, flatten_sharded_tensors: bool = True, allow_partial_load: bool = False)[source]
#

DefaultSavePlanner

classFull reference ↗
class tensorplay.distributed.checkpoint.DefaultSavePlanner(flatten_state_dict: bool = True, flatten_sharded_tensors: bool = True, dedup_replicated_tensors: bool | None = None, dedup_save_to_lowest_rank: bool = False, enable_plan_caching: bool = False)[source]
#

LoadPlan

classFull reference ↗
class tensorplay.distributed.checkpoint.LoadPlan(items: 'list[ReadItem]', storage_data: 'Any' = None, planner_data: 'Any' = None)[source]
#

MegaStorageWriter

classFull reference ↗
class tensorplay.distributed.checkpoint.MegaStorageWriter(path: str | PathLike, fqn_to_index_mapping: dict[str, int] | None = None)[source]

Writes MEGA-format shards (model[-N-of-M].mega) plus model.mega.index.json; accepts plain directories and mega:// URIs.

#

Metadata

classFull reference ↗
class tensorplay.distributed.checkpoint.Metadata(state_dict_metadata: 'dict[str, TensorStorageMetadata | BytesStorageMetadata]', planner_data: 'Any' = None, storage_data: 'Any' = None, storage_meta: 'StorageMeta | None' = None, version: 'str | None' = None)[source]
#

MetadataIndex

classFull reference ↗
class tensorplay.distributed.checkpoint.MetadataIndex(fqn: 'str', offset: 'Sequence[int] | None' = None, index: 'int | None' = None) → 'None'[source]
#

QuantizedHuggingFaceStorageReader

classFull reference ↗
class tensorplay.distributed.checkpoint.QuantizedHuggingFaceStorageReader(path: str, thread_count: int = 1, target_dtype: Any = None, block_size: int = 128)[source]
#

ReadItem

classFull reference ↗
class tensorplay.distributed.checkpoint.ReadItem(type: 'LoadItemType', dest_index: 'MetadataIndex', dest_offsets: 'tuple[int, ...]', storage_index: 'MetadataIndex', storage_offsets: 'tuple[int, ...]', lengths: 'tuple[int, ...]')[source]
#

SavePlan

classFull reference ↗
class tensorplay.distributed.checkpoint.SavePlan(items: 'list[WriteItem]', storage_data: 'Any' = None, planner_data: 'Any' = None, usable: 'bool' = True)[source]
#

StateDictOptions

classFull reference ↗
class tensorplay.distributed.checkpoint.StateDictOptions(full_state_dict: 'bool' = False, cpu_offload: 'bool' = False, ignore_frozen_params: 'bool' = False, keep_submodule_prefixes: 'bool' = True, strict: 'bool' = True, broadcast_from_rank0: 'bool' = False, flatten_optimizer_state_dict: 'bool' = False, dsd_fqn_modifiers: 'str' = '_fqn_modifiers')[source]
#

StorageMeta

classFull reference ↗
class tensorplay.distributed.checkpoint.StorageMeta(checkpoint_id: 'str | os.PathLike[str] | None' = None, save_id: 'str | None' = None, load_id: 'str | None' = None, modules: 'list[str]' = <factory>)[source]
#

TensorProperties

classFull reference ↗
class tensorplay.distributed.checkpoint.TensorProperties(dtype: 'Any' = <factory>, layout: 'Any' = <factory>, requires_grad: 'bool' = False, memory_format: 'Any' = <factory>, pin_memory: 'bool' = False)[source]
#

TensorStorageMetadata

classFull reference ↗
class tensorplay.distributed.checkpoint.TensorStorageMetadata(properties: 'TensorProperties', size: 'tuple[int, ...]', chunks: 'list[ChunkStorageMetadata]')[source]
#

TensorWriteData

classFull reference ↗
class tensorplay.distributed.checkpoint.TensorWriteData(chunk: 'ChunkStorageMetadata', properties: 'TensorProperties', size: 'tuple[int, ...]')[source]
#

WriteItem

classFull reference ↗
class tensorplay.distributed.checkpoint.WriteItem(index: 'MetadataIndex', type: 'WriteItemType', bytes_io_data: 'BytesIOWriteData | None' = None, tensor_data: 'TensorWriteData | None' = None)[source]

Exceptions 1

tensorplay.distributed.autograd API

Complete API reference for tensorplay.distributed.autograd, including signatures, parameters, examples and members.

tensorplay.distributed.device_mesh API

Complete API reference for tensorplay.distributed.device_mesh, including signatures, parameters, examples and members.

On this page

Functions 12async_saveget_model_state_dictget_optimizer_state_dictget_state_dictload_sharded_optimizer_state_dictload_state_dictloadsave_state_dictsaveset_model_state_dictset_optimizer_state_dictset_state_dictClasses 36AsyncCheckpointerTypeAsyncSaveResponseBytesIOWriteDataBytesStorageMetadataCheckpointableTensorChunkStorageMetadataDefaultLoadPlannerDefaultSavePlannerFileSystemFileSystemBaseFileSystemReaderFileSystemWriterHuggingFaceStorageReaderHuggingFaceStorageWriterLoadItemTypeLoadPlanLoadPlannerMegaStorageReaderMegaStorageWriterMetadataMetadataIndexQuantizedHuggingFaceStorageReaderReadItemSavePlanSavePlannerSerializationFormatStateDictOptionsStatefulStorageMetaStorageReaderStorageWriterTensorPropertiesTensorStorageMetadataTensorWriteDataWriteItemWriteItemTypeExceptions 1CheckpointException
Ask DeepWiki