Asynchronous Pytorch Distributed Checkpoint save with optimized FileSystemWriter

State dict saver for PyT Distributed format allowing asynchronous save.

class nvidia_resiliency_ext.checkpointing.async_ckpt.state_dict_saver.CheckpointMetadataCache

Bases: object

Global metadata and plans carried from one save to the next, to reuse the metadata.

Building the global metadata gathers every rank’s local plan on the coordinator. With enable_cache, save_state_dict_async_plan avoids that:

• The first save in a process reuses the metadata of a loaded checkpoint (see set_cached_global_metadata) if all ranks’ local plans write exactly the chunks it lists; verify_global_metadata_reuse checks that with one all_reduce.

• Later saves reuse the plans and metadata of the previous save without communicating: enable_cache promises that the checkpoint structure does not change. Each rank still creates its local plan and raises if it differs from the previous one.

The coordinator writes a shallow copy of the cached metadata on every save, so the cache never holds a save’s storage_data and never modifies the metadata it was given.

global_metadata

The global metadata to reuse: the loaded checkpoint’s, without its storage_data, on every rank until the first save, the previous save’s on the coordinator afterwards.

Type:

Metadata

local_plan

This rank’s local plan of the previous save, as the planner created it; None before the first save in this process.

Type:

SavePlan

central_plan

This rank’s global plan of the previous save.

Type:

SavePlan

set_cached_global_metadata(cached_global_metadata)

Sets the global metadata of a loaded checkpoint, to be reused if it still applies.

Every rank must set it. The next save checks whether all ranks’ local plans write exactly the chunks this metadata lists, and if so reuses it instead of gathering the plans and building the global metadata again. The metadata is not modified.

Parameters:

cached_global_metadata (Metadata) – The global metadata from a previous checkpoint.

nvidia_resiliency_ext.checkpointing.async_ckpt.state_dict_saver.init_checkpoint_metadata_cache(cached_global_metadata=None)

Initializes the checkpoint metadata cache.

This function creates a new CheckpointMetadataCache instance and sets the cached global metadata from the previous checkpoint

Parameters:

cached_global_metadata (Metadata)

nvidia_resiliency_ext.checkpointing.async_ckpt.state_dict_saver.save_state_dict_async_finalize(storage_writer, global_metadata, dist_wrapper)

Finalization of save_state_dict_async_plan.

The input arguments are the same as the save_state_dict_async_plan output, the write_results are retrieved from the storage_writer. Each rank sends them to the coordinator as a table (see _metadata_pickler.table) if it can, else pickled.

Parameters:

• storage_writer (FileSystemWriterAsync) – storage writer used for planning

• global_metadata (Metadata) – metadata created during planning

• dist_wrapper (_DistWrapper) – distributed wrapper created during planning

Return type:

None

Returns: None

nvidia_resiliency_ext.checkpointing.async_ckpt.state_dict_saver.save_state_dict_async_plan(state_dict, storage_writer, process_group=None, coordinator_rank=0, planner=None, enable_cache=False, metadata_cache=None)

First stage of saving a state dict to storage.

This is an async adjustment of torch.distributed.checkpoint.state_dict_saver. In order to support async save, saving should be split into three parts:

1. Planning

2. Actual saving

3. Finalization

Out of these, step (2) must happen asynchronously. The first step is realized with this function.

The planning part consists of several steps, described here: https://pytorch.org/docs/stable/distributed.checkpoint.html#torch.distributed.checkpoint.SavePlanner

Parameters:

• state_dict (STATE_DICT_TYPE) – state dict to save

• storage_writer (FileSystemWriterAsync) – in current version only an instance of FileSystemWriterAsync

• process_group (dist.ProcessGroup, optional) – process group used for save planning

• coordinator_rank (int) – coordinator rank for planning. Defaults to 0.

• planner (SavePlanner, optional) – save planner for torch.distributed.checkpoint format

• enable_cache (bool) – Reuse global metadata instead of building it: the first save reuses a loaded checkpoint’s metadata if it still applies, later saves reuse the previous save’s plans and metadata, as the checkpoint structure must not change between them (a rank whose local plan changed raises). See CheckpointMetadataCache. Defaults to False.

• metadata_cache (CheckpointMetadataCache, optional) – Custom metadata cache instance to use for storing and retrieving checkpoint metadata. If not provided, the global cache will be used.

Returns:

Contains:

• storage writer (the one passed as input)

• global metadata to write (on the coordinator; None on other ranks)

• distributed wrapper used for planning

Return type:

tuple

The return value of this function should be passed as an input to save_state_dict_async_finalize.

nvidia_resiliency_ext.checkpointing.async_ckpt.state_dict_saver.verify_global_metadata_reuse(loaded_metadata, local_plan, dist_wrapper)

Verifies that the metadata of a loaded checkpoint can be reused as this save’s global metadata.

It can if all ranks’ local plans write exactly the chunks the metadata lists, with the same types, global sizes and properties; see _metadata_reuse. Collective: every rank of the planning group must call it, also ranks without loaded metadata, which vote against reuse.

Parameters:

• loaded_metadata (Metadata, optional) – The global metadata of the loaded checkpoint.

• local_plan (SavePlan) – This rank’s local save plan.

• dist_wrapper (_DistWrapper) – distributed wrapper created during planning

Return type:

bool

Returns: True iff the global metadata reuse is possible.