nemo_automodel.checkpoint._backports.filesystem#

Module Contents#

Classes#

SerializationFormat

Enumeration of supported on-disk checkpoint formats.

_StorageInfo

This is the per entry storage info.

_StoragePrefix

_TensorLoader

_SerialCpuLoader

_OverlappingCpuLoader

_StorageWriterTransforms

This is experimental, and will likely move elsewhere in the future. It lives here to minimize changes while we are still learning and gathering feedback.

FileSystemBase

FileSystem

_FileSystemWriter

Basic implementation of StorageWriter using file IO.

_StorageReaderTransforms

This is experimental, and will likely move elsewhere in the future. It lives here to minimize changes while we are still learning and gathering feedback.

FileSystemReader

FileSystemWriter

Basic implementation of StorageWriter using file IO.

Functions#

_generate_uuid

_item_size

_split_by_size_and_type

_write_item

_write_files_from_queue

Data#

__all__

_metadata_fn

DEFAULT_SUFFIX

API#

nemo_automodel.checkpoint._backports.filesystem.__all__#

[‘FileSystemWriter’, ‘FileSystemReader’, ‘FileSystem’, ‘FileSystemBase’]

nemo_automodel.checkpoint._backports.filesystem._metadata_fn: str#

‘.metadata’

class nemo_automodel.checkpoint._backports.filesystem.SerializationFormat(*args, **kwds)[source]#

Bases: enum.Enum

Enumeration of supported on-disk checkpoint formats.

Initialization

TORCH_SAVE#

‘torch_save’

SAFETENSORS#

‘safetensors’

class nemo_automodel.checkpoint._backports.filesystem._StorageInfo[source]#

This is the per entry storage info.

relative_path: str#

None

offset: int#

None

length: int#

None

transform_descriptors: Optional[collections.abc.Sequence[str]]#

None

__getstate__()[source]#
class nemo_automodel.checkpoint._backports.filesystem._StoragePrefix[source]#
prefix: str#

None

nemo_automodel.checkpoint._backports.filesystem.DEFAULT_SUFFIX#

‘.distcp’

nemo_automodel.checkpoint._backports.filesystem._generate_uuid() → str[source]#
class nemo_automodel.checkpoint._backports.filesystem._TensorLoader[source]#

Bases: abc.ABC

abstractmethod add(size: int, obj: object) → None[source]#
abstractmethod start_loading() → None[source]#
abstractmethod values() → collections.abc.Iterator[tuple[torch.Tensor, object]][source]#
class nemo_automodel.checkpoint._backports.filesystem._SerialCpuLoader(resolve_fun: Callable)[source]#

Bases: nemo_automodel.checkpoint._backports.filesystem._TensorLoader

add(size: int, obj: object) → None[source]#
start_loading() → None[source]#
values() → collections.abc.Iterator[tuple[torch.Tensor, object]][source]#
class nemo_automodel.checkpoint._backports.filesystem._OverlappingCpuLoader(
resolve_fun: Callable,
stream: Optional[torch.Stream] = None,
inflight_threshhold: int = 1000000,
)[source]#

Bases: nemo_automodel.checkpoint._backports.filesystem._TensorLoader

property _done: bool#
_drain() → list[tuple[torch.Tensor, object]][source]#
_refill() → None[source]#
_finish() → collections.abc.Iterable[tuple[torch.Tensor, object]][source]#
add(size: int, obj: object) → None[source]#
start_loading() → None[source]#
values() → collections.abc.Iterator[tuple[torch.Tensor, object]][source]#
class nemo_automodel.checkpoint._backports.filesystem._StorageWriterTransforms(
extensions: Optional[collections.abc.Sequence[torch.distributed.checkpoint._extension.StreamTransformExtension]] = None,
)[source]#

This is experimental, and will likely move elsewhere in the future. It lives here to minimize changes while we are still learning and gathering feedback.

Initialization

If the extensions arg is None, this means the implementation should provide whatever defaults it chooses. An empty sequence indicates no extensions should be used. At this time, the default extensions sequence is empty.

transform_save_stream(
write_item: torch.distributed.checkpoint.planner.WriteItem,
raw_stream: io.IOBase,
) → tuple[IO[bytes], list[str]][source]#
nemo_automodel.checkpoint._backports.filesystem._item_size(
item: torch.distributed.checkpoint.planner.WriteItem,
) → int[source]#
nemo_automodel.checkpoint._backports.filesystem._split_by_size_and_type(
bins: int,
items: list[torch.distributed.checkpoint.planner.WriteItem],
) → list[list[torch.distributed.checkpoint.planner.WriteItem]][source]#
nemo_automodel.checkpoint._backports.filesystem._write_item(
transforms: nemo_automodel.checkpoint._backports.filesystem._StorageWriterTransforms,
stream: io.IOBase,
data: Union[io.BytesIO, torch.Tensor],
write_item: torch.distributed.checkpoint.planner.WriteItem,
storage_key: str,
serialization_format: nemo_automodel.checkpoint._backports.filesystem.SerializationFormat,
) → torch.distributed.checkpoint.storage.WriteResult[source]#
nemo_automodel.checkpoint._backports.filesystem._write_files_from_queue(
create_stream: Callable,
file_queue: queue.Queue,
result_queue: queue.Queue,
planner: torch.distributed.checkpoint.planner.SavePlanner,
transforms: nemo_automodel.checkpoint._backports.filesystem._StorageWriterTransforms,
inflight_threshhold: int,
use_fsync: bool,
thread_count: int,
serialization_format: nemo_automodel.checkpoint._backports.filesystem.SerializationFormat,
) → None[source]#
class nemo_automodel.checkpoint._backports.filesystem.FileSystemBase[source]#

Bases: abc.ABC

abstractmethod create_stream(
path: Union[str, os.PathLike],
mode: str,
) → collections.abc.Generator[io.IOBase, None, None][source]#
abstractmethod concat_path(
path: Union[str, os.PathLike],
suffix: str,
) → Union[str, os.PathLike][source]#
abstractmethod rename(
path: Union[str, os.PathLike],
new_path: Union[str, os.PathLike],
) → None[source]#
abstractmethod init_path(
path: Union[str, os.PathLike],
) → Union[str, os.PathLike][source]#
abstractmethod mkdir(path: Union[str, os.PathLike]) → None[source]#
abstractmethod classmethod validate_checkpoint_id(
checkpoint_id: Union[str, os.PathLike],
) → bool[source]#
abstractmethod exists(path: Union[str, os.PathLike]) → bool[source]#
abstractmethod rm_file(path: Union[str, os.PathLike]) → None[source]#
class nemo_automodel.checkpoint._backports.filesystem.FileSystem[source]#

Bases: nemo_automodel.checkpoint._backports.filesystem.FileSystemBase

create_stream(
path: Union[str, os.PathLike],
mode: str,
) → collections.abc.Generator[io.IOBase, None, None][source]#
concat_path(
path: Union[str, os.PathLike],
suffix: str,
) → Union[str, os.PathLike][source]#
init_path(
path: Union[str, os.PathLike],
) → Union[str, os.PathLike][source]#
rename(
path: Union[str, os.PathLike],
new_path: Union[str, os.PathLike],
) → None[source]#
mkdir(path: Union[str, os.PathLike]) → None[source]#
classmethod validate_checkpoint_id(
checkpoint_id: Union[str, os.PathLike],
) → bool[source]#
exists(path: Union[str, os.PathLike]) → bool[source]#
rm_file(path: Union[str, os.PathLike]) → None[source]#
ls(path: Union[str, os.PathLike]) → list[str][source]#
class nemo_automodel.checkpoint._backports.filesystem._FileSystemWriter(
path: Union[str, os.PathLike],
single_file_per_rank: bool = True,
sync_files: bool = True,
thread_count: int = 1,
per_thread_copy_ahead: int = 10000000,
overwrite: bool = True,
_extensions: Optional[collections.abc.Sequence[torch.distributed.checkpoint._extension.StreamTransformExtension]] = None,
serialization_format: nemo_automodel.checkpoint._backports.filesystem.SerializationFormat = SerializationFormat.TORCH_SAVE,
*args: Any,
**kwargs: Any,
)[source]#

Bases: torch.distributed.checkpoint.storage.StorageWriter

Basic implementation of StorageWriter using file IO.

This implementation makes the following assumptions and simplifications:

  • The checkpoint path is an empty or non-existing directory.

  • File creation is atomic

The checkpoint consist of one file per write request plus a .metadata file with the serialized metadata.

Initialization

Initialize the writer pointing to path.

Parameters:

  • path – directory where the checkpoint will be written to.

  • single_file_per_rank – Produce one file per rank instead of one file per tensor/blob. Default to True.

  • sync_files – force files to be synced to permanent storage. Default to True.

  • thread_count – Number of IO threads to use to write. Default to 1.

  • per_thread_copy_ahead – How many bytes to copy from the GPU ahead of saving then. Default 10Mb.

  • overwrite – Whether to allow overwriting existing checkpoints. Defaults to True.

  • _extensions – Extensions to apply to output streams (EXPERIMENTAL)

N. B. If sync_files is disabled, there’s no guarantee that the checkpoint will be consistent in the case of a failure.

reset(
checkpoint_id: Union[str, os.PathLike, None] = None,
) → None[source]#
set_up_storage_writer(is_coordinator: bool) → None[source]#
prepare_local_plan(
plan: torch.distributed.checkpoint.planner.SavePlan,
) → torch.distributed.checkpoint.planner.SavePlan[source]#
prepare_global_plan(
plans: list[torch.distributed.checkpoint.planner.SavePlan],
) → list[torch.distributed.checkpoint.planner.SavePlan][source]#
write_data(
plan: torch.distributed.checkpoint.planner.SavePlan,
planner: torch.distributed.checkpoint.planner.SavePlanner,
) → torch.futures.Future[list[torch.distributed.checkpoint.storage.WriteResult]][source]#
_write_data(
planner: torch.distributed.checkpoint.planner.SavePlanner,
file_queue: queue.Queue,
) → torch.futures.Future[list[torch.distributed.checkpoint.storage.WriteResult]][source]#
finish(
metadata: torch.distributed.checkpoint.metadata.Metadata,
results: list[list[torch.distributed.checkpoint.storage.WriteResult]],
) → None[source]#
storage_meta() → Optional[torch.distributed.checkpoint.metadata.StorageMeta][source]#
property metadata_path: Union[str, os.PathLike]#
property checkpoint_id: Union[str, os.PathLike]#

return the checkpoint_id that will be used to save the checkpoint.

classmethod validate_checkpoint_id(
checkpoint_id: Union[str, os.PathLike],
) → bool[source]#
class nemo_automodel.checkpoint._backports.filesystem._StorageReaderTransforms(
extension_registry: Optional[torch.distributed.checkpoint._extension.ExtensionRegistry] = None,
)[source]#

This is experimental, and will likely move elsewhere in the future. It lives here to minimize changes while we are still learning and gathering feedback.

Initialization

transform_load_stream(
read_item: torch.distributed.checkpoint.planner.ReadItem,
transform_descriptors: collections.abc.Sequence[str],
raw_stream: IO[bytes],
) → IO[bytes][source]#
class nemo_automodel.checkpoint._backports.filesystem.FileSystemReader(
path: Union[str, os.PathLike],
_extension_registry: Optional[torch.distributed.checkpoint._extension.ExtensionRegistry] = None,
)[source]#

Bases: torch.distributed.checkpoint.storage.StorageReader

_slice_file(
file,
sinfo: nemo_automodel.checkpoint._backports.filesystem._StorageInfo,
) → IO[bytes][source]#
reset(
checkpoint_id: Union[str, os.PathLike, None] = None,
) → None[source]#
read_data(
plan: torch.distributed.checkpoint.planner.LoadPlan,
planner: torch.distributed.checkpoint.planner.LoadPlanner,
) → torch.futures.Future[None][source]#
read_metadata() → torch.distributed.checkpoint.metadata.Metadata[source]#
set_up_storage_reader(
metadata: torch.distributed.checkpoint.metadata.Metadata,
is_coordinator: bool,
) → None[source]#
prepare_local_plan(
plan: torch.distributed.checkpoint.planner.LoadPlan,
) → torch.distributed.checkpoint.planner.LoadPlan[source]#
prepare_global_plan(
plans: list[torch.distributed.checkpoint.planner.LoadPlan],
) → list[torch.distributed.checkpoint.planner.LoadPlan][source]#
property checkpoint_id: Union[str, os.PathLike]#

return the checkpoint_id that will be used to load the checkpoint.

classmethod validate_checkpoint_id(
checkpoint_id: Union[str, os.PathLike],
) → bool[source]#
class nemo_automodel.checkpoint._backports.filesystem.FileSystemWriter(
path: Union[str, os.PathLike],
single_file_per_rank: bool = True,
sync_files: bool = True,
thread_count: int = 1,
per_thread_copy_ahead: int = 10000000,
cache_staged_state_dict: bool = False,
overwrite: bool = True,
_extensions: Optional[collections.abc.Sequence[torch.distributed.checkpoint._extension.StreamTransformExtension]] = None,
serialization_format: nemo_automodel.checkpoint._backports.filesystem.SerializationFormat = SerializationFormat.TORCH_SAVE,
)[source]#

Bases: nemo_automodel.checkpoint._backports.filesystem._FileSystemWriter, torch.distributed.checkpoint.staging.BlockingAsyncStager

Basic implementation of StorageWriter using file IO.

This implementation makes the following assumptions and simplifications:

  • The checkpoint path is an empty or non-existing directory.

  • File creation is atomic

The checkpoint consist of one file per write request plus a .metadata file with the serialized metadata.

Initialization

Initialize the writer pointing to path.

Parameters:

  • path – directory where the checkpoint will be written to.

  • single_file_per_rank – Produce one file per rank instead of one file per tensor/blob. Default to True.

  • sync_files – force files to be synced to permanent storage. Default to True.

  • thread_count – Number of IO threads to use to write. Default to 1.

  • per_thread_copy_ahead – How many bytes to copy from the GPU ahead of saving then. Default 10Mb.

  • cache_staged_state_dict – Whether to cache the staged state_dict. This option decreases staging latency at the cost of increases memory usage. Additionally, if this parameter is set to True, it’s the expectation that the stager is maintained and re-used for multiple dcp.async_save calls. Default to False.

  • overwrite – Whether to allow overwriting existing checkpoints. Defaults to True.

  • _extensions – Extensions to apply to output streams (EXPERIMENTAL)

N. B. If sync_files is disabled, there’s no guarantee that the checkpoint will be consistent in the case of a failure.

stage(
state_dict: torch.distributed.checkpoint.metadata.STATE_DICT_TYPE,
) → torch.distributed.checkpoint.metadata.STATE_DICT_TYPE[source]#

Override of AsyncStager.stage