`nemo_automodel.checkpoint._backports.hf_utils`#

Module Contents#

Classes#

_HFStorageInfo

This is the per entry storage info.

Functions#

`_gen_file_name`
`_get_safetensors_file_metadata`
`_get_dtype`
`_get_dcp_custom_metadata`

Data#

`_metadata_fn`
`FILE_NAME`
`SHARDED_FILE_NAME`
`SUFFIX`
`CUSTOM_METADATA_KEY`
`DEFAULT_EXTRA_METADATA_KEY`
`SAVED_OFFSETS_KEY`
`SHAPE_KEY`
`DATA_KEY`
`DTYPE_KEY`
`DATA_OFFSETS_KEY`
`DTYPE_MAP`
`HF_DCP_VERSION`
`DCP_VERSION_KEY`
`DCP_SHARDING_INFO_KEY`

API#

nemo_automodel.checkpoint._backports.hf_utils._metadata_fn: str#: ‘model.safetensors.index.json’

nemo_automodel.checkpoint._backports.hf_utils.FILE_NAME#: ‘model-{cpt_idx}-of-{num_files}’

nemo_automodel.checkpoint._backports.hf_utils.SHARDED_FILE_NAME#: ‘shard-{shard_idx}-model-{cpt_idx}-of-{num_files}’

nemo_automodel.checkpoint._backports.hf_utils.SUFFIX#: ‘.safetensors’

nemo_automodel.checkpoint._backports.hf_utils.CUSTOM_METADATA_KEY#: ‘DCP_SHARDING_INFO’

nemo_automodel.checkpoint._backports.hf_utils.DEFAULT_EXTRA_METADATA_KEY#: ‘metadata’

nemo_automodel.checkpoint._backports.hf_utils.SAVED_OFFSETS_KEY#: ‘saved_offsets’

nemo_automodel.checkpoint._backports.hf_utils.SHAPE_KEY#: ‘shape’

nemo_automodel.checkpoint._backports.hf_utils.DATA_KEY#: ‘data’

nemo_automodel.checkpoint._backports.hf_utils.DTYPE_KEY#: ‘dtype’

nemo_automodel.checkpoint._backports.hf_utils.DATA_OFFSETS_KEY#: ‘data_offsets’

nemo_automodel.checkpoint._backports.hf_utils.DTYPE_MAP#: None

nemo_automodel.checkpoint._backports.hf_utils.HF_DCP_VERSION: float#: 1.0

nemo_automodel.checkpoint._backports.hf_utils.DCP_VERSION_KEY#: ‘DCP_VERSION’

nemo_automodel.checkpoint._backports.hf_utils.DCP_SHARDING_INFO_KEY#: ‘DCP_SHARDING_INFO’

class nemo_automodel.checkpoint._backports.hf_utils._HFStorageInfo[source]#

This is the per entry storage info.

relative_path: str#: None

offset: int#: None

length: int#: None

shape: torch.Size#: None

dtype: torch.dtype#: None

__getstate__()[source]#

nemo_automodel.checkpoint._backports.hf_utils._gen_file_name( index: int, largest_index: int, shard_index: Optional[int] = None, ) → str[source]#

nemo_automodel.checkpoint._backports.hf_utils._get_safetensors_file_metadata( file_bytes: io.IOBase, ) → tuple[Any, int][source]#

nemo_automodel.checkpoint._backports.hf_utils._get_dtype(dtype_str: str) → torch.dtype[source]#

nemo_automodel.checkpoint._backports.hf_utils._get_dcp_custom_metadata( metadata: Any, ) → Optional[Any][source]#