cerebras.modelzoo.tools.checkpoint_converters.base_converter.BaseCheckpointConverter#

class cerebras.modelzoo.tools.checkpoint_converters.base_converter.BaseCheckpointConverter[source]#

Bases: cerebras.modelzoo.tools.checkpoint_converters.base_converter.BaseDictionaryConverter, abc.ABC

Converts between checkpoint state_dict formats.

Methods

convert

convert_all_keys

convert_helper

Converts all keys in a checkpoint from converter_indices.direction format to the other format.

convert_key

Attempts to convert the old key by matching against the list of conversion rules.

extract_model_dict

Hook to extract model state dicts out of the input/output checkpoint.

file_formats

formats

get_config_converter_class

get_converter_indices

init_output_checkpoint

(Pre)Initializes the output checkpoint at a supplied path.

load

post_checkpoint_convert

Hook executes after checkpoint conversion.

post_model_convert

Hook executes right after model conversion.

pre_checkpoint_convert

Hook executes before checkpoint conversion.

pre_model_convert

Hook executes right before model conversion.

replaceKey

Copies value that exists at old_state_dict's old_key to new_state_dict's new_key.

save

supports_conversion

abstract classmethod init_output_checkpoint(file_without_ext, converter_indices, **kwargs)[source]#

(Pre)Initializes the output checkpoint at a supplied path. This is used in streaming conversion when the checkpoint is written to file as conversion is performed rather than accumulating the full checkpoint in memory and saving to file at the very end.

convert_helper(input_checkpoint, configs, converter_indices, output_checkpoint={}, drop_unmatched_keys=False, no_progress_bar=True, debug=False)[source]#

Converts all keys in a checkpoint from converter_indices.direction format to the other format. Conversion will fail if at least one of the keys did not match on any conversion rules and drop_unmatched_keys is not enabled. Returns the newly converted checkpoint.

pre_model_convert(old_state_dict, new_state_dict, configs, converter_indices, drop_unmatched_keys)[source]#

Hook executes right before model conversion.

post_model_convert(old_state_dict, new_state_dict, configs, converter_indices, drop_unmatched_keys, key_prefix='')[source]#

Hook executes right after model conversion.

pre_checkpoint_convert(input_checkpoint, output_checkpoint, configs, converter_indices)[source]#

Hook executes before checkpoint conversion.

extract_model_dict(input_checkpoint, output_checkpoint, configs, converter_indices)[source]#

Hook to extract model state dicts out of the input/output checkpoint.

post_checkpoint_convert(input_checkpoint, output_checkpoint, configs, converter_indices)[source]#

Hook executes after checkpoint conversion.

convert_key(old_key, old_state_dict, new_state_dict, from_index, match_start=0, prefix='', action_fn_args=None, debug=False)#

Attempts to convert the old key by matching against the list of conversion rules. The first rule to match is used for conversion (i.e. even if multiple rules would match, the latter ones are never used). Returns True if a conversion occurred.

static replaceKey(old_key, new_key, old_state_dict, new_state_dict, from_index, action_fn_args=None)#

Copies value that exists at old_state_dict’s old_key to new_state_dict’s new_key.