cerebras.modelzoo.tools.checkpoint_converters.base_converter.BaseCheckpointConverter_HF_CS#

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

Bases: cerebras.modelzoo.tools.checkpoint_converters.base_converter.BaseCheckpointConverter_CS_CS

HF checkpoints contain model only while CS checkpoints package model, optimizer, and lr_scheduler into a single checkpoint. This class overrides the post_checkpoint_convert to automatically extract/package the state_dict correctly.

Methods

`attempt_mup_to_sp`
`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`
`file_formats`
`formats`
`get_config_converter_class`
`get_converter_indices`
`get_mup_converter`	Allows models to override the default muP converters with their own
`init_output_checkpoint`
`load`
`post_checkpoint_convert`
`post_model_convert`
`pre_checkpoint_convert`
`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`
`supports_mup_conversion`

get_mup_converter()[source]#

Allows models to override the default muP converters with their own

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

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.

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.

pre_model_convert(old_state_dict, new_state_dict, configs, converter_indices, drop_unmatched_keys)#

Hook executes right before model conversion.

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.

cerebras.modelzoo.tools.checkpoint_converters.base_converter.BaseCheckpointConverter_CS_CS

cerebras.modelzoo.tools.checkpoint_converters.base_converter.BaseCheckpointConverter_UnpackedHF_PackedCS