modelzoo.common.pytorch.model_utils.checkpoint_converters.streaming_checkpoints#

Functions

convert_file_size_to_int

Converts a size expressed as a string with digits and unit (like "5MB") to an integer (in bytes).

dtype_byte_size

Returns the size (in bytes) occupied by one parameter of type dtype.

Classes

StreamingCSLeaf

Marks checkpoint keys that can be directly loaded from/saved to the H5 checkpoint.

StreamingCSWriter

Writes a Cerebras H5 checkpoint in a streaming (incremental) manner rather than accumulating the full checkpoint into memory and then writing all weights at the end.

StreamingCSWriterView

StreamingCSWriterView allows for checkpoints with arbitrarily nested dictionaries/lists to be written in a streaming (incremental) manner by offering a "view" into a StreamingCSWriter.

StreamingShardedHFReader

Allows sharded HuggingFace checkpoints to be read in a streaming manner rather than loading all shards into memory all at once.

StreamingShardedHFWriter

Writes a HuggingFace sharded checkpoint in a streaming manner rather than accumulating the full checkpoint into memory and then writing all shards at the end.