checkpoint_utils_hf

Utilities for loading and saving Hugging Face-format checkpoints (AutoConfig + optional block_configs).

Functions

force_cache_dynamic_modules
load_model_config	Load model configuration from a checkpoint directory.
init_model_from_config	Build a model from config on meta/uninitialized weights (used e.g. for subblock param counts).
save_checkpoint
save_checkpoint_from_shards	Save a checkpoint when the model's weights are sharded across distributed ranks.
save_subblocks
save_model_config
load_passthrough_state_dict	Load descriptor-declared passthrough tensors from a checkpoint directory.

force_cache_dynamic_modules(config, checkpoint_dir, trust_remote_code=False)

Parameters:

• config (PreTrainedConfig)

• checkpoint_dir (Path | str)

• trust_remote_code (bool)

init_model_from_config(config, *, trust_remote_code=False, **kwargs)

Build a model from config on meta/uninitialized weights (used e.g. for subblock param counts).

trust_remote_code defaults to False (only AutoModelForCausalLM.from_config uses it). Pass True when loading configs that rely on custom modeling code from the checkpoint.

Parameters:

• config (PreTrainedConfig)

• trust_remote_code (bool)

Return type:

PreTrainedModel

load_model_config(checkpoint_dir, model_config_overrides=None, ignore_unexpected_config_keys=False, trust_remote_code=False)

Load model configuration from a checkpoint directory.

Parameters:

• checkpoint_dir (Path | str) – Path to the checkpoint directory (e.g. containing config.json).

• model_config_overrides (Mapping | None) – Optional mapping of config overrides.

• ignore_unexpected_config_keys (bool) – If True, ignore unexpected config keys.

• trust_remote_code (bool) – If True, allows execution of custom code from the model repository. This is a security risk if the model source is untrusted. Only set to True if you trust the source of the model. Defaults to False for security.

Returns:

Loaded model configuration (PretrainedConfig).

load_passthrough_state_dict(checkpoint_dir, descriptor)

Load descriptor-declared passthrough tensors from a checkpoint directory.

Parameters:

• checkpoint_dir (Path | str)

• descriptor (ModelDescriptor)

Return type:

dict[str, Tensor]

save_checkpoint(model, checkpoint_dir, descriptor)

Parameters:

• model (PreTrainedModel)

• checkpoint_dir (Path | str)

• descriptor (ModelDescriptor)

Return type:

None

save_checkpoint_from_shards(model, checkpoint_dir, descriptor, extra_state_dict=None, reference_checkpoint_dir=None, changed_keys=None)

Save a checkpoint when the model’s weights are sharded across distributed ranks.

On distributed runs, rank 0 gathers only tensor-name metadata up front and then gathers tensors one safetensors file at a time. This avoids materializing the full model from all ranks on rank 0 while still producing a single HF-compatible checkpoint/index.

Parameters:

• model (PreTrainedModel) – The model instance whose local state_dict contains this rank’s

• weights. (shard of)

• checkpoint_dir (Path | str) – Destination directory for the checkpoint files.

• descriptor (ModelDescriptor) – Descriptor used to partition weights into subblocks and build

• index. (the safetensors)

• reference_checkpoint_dir (Path | str | None) – Optional checkpoint with the same descriptor weight grouping. If provided with changed_keys, unchanged subblock files are hardlinked or symlinked from this checkpoint instead of rewritten.

• changed_keys (set[str] | None) – State-dict keys that may differ from the reference checkpoint.

• extra_state_dict (dict[str, Tensor] | None)

Return type:

None

save_model_config(model_config, checkpoint_dir)

Parameters:

• model_config (PreTrainedConfig)

• checkpoint_dir (Path | str)

Return type:

None

save_subblocks(state_dict, checkpoint_dir, weight_map=None, multi_threaded=True, max_workers=None)

Parameters:

• state_dict (dict[str, Tensor])

• checkpoint_dir (Path | str)

• weight_map (dict[str, str] | None)

• multi_threaded (bool)

• max_workers (int | None)

Return type:

None