MCPcopy Create free account
hub / github.com/zai-org/CodeGeeX / check_checkpoint_args

Function check_checkpoint_args

codegeex/megatron/checkpointing.py:43–72  ·  view source on GitHub ↗

Ensure fixed arguments for a model are the same for the input arguments and the one retrieved from checkpoint.

(checkpoint_args)

Source from the content-addressed store, hash-verified

41
42
43def check_checkpoint_args(checkpoint_args):
44 """Ensure fixed arguments for a model are the same for the input
45 arguments and the one retrieved from checkpoint."""
46 args = get_args()
47
48 def _compare(arg_name, old_arg_name=None):
49 if old_arg_name is not None:
50 checkpoint_value = getattr(checkpoint_args, old_arg_name)
51 else:
52 checkpoint_value = getattr(checkpoint_args, arg_name)
53 args_value = getattr(args, arg_name)
54 error_message = (
55 "{} value from checkpoint ({}) is not equal to the "
56 "input argument value ({}).".format(arg_name, checkpoint_value, args_value)
57 )
58 assert checkpoint_value == args_value, error_message
59
60 _compare("num_layers")
61 _compare("hidden_size")
62 _compare("num_attention_heads")
63 _compare("max_position_embeddings")
64 if args.vocab_file:
65 _compare("make_vocab_size_divisible_by")
66 _compare("padded_vocab_size")
67 _compare("tokenizer_type")
68 if get_checkpoint_version() < 3.0:
69 _compare("tensor_model_parallel_size", old_arg_name="model_parallel_size")
70 if get_checkpoint_version() >= 3.0:
71 _compare("tensor_model_parallel_size")
72 _compare("pipeline_model_parallel_size")
73
74
75def ensure_directory_exists(filename):

Callers 1

load_checkpointFunction · 0.85

Calls 3

get_argsFunction · 0.90
_compareFunction · 0.85
get_checkpoint_versionFunction · 0.85

Tested by

no test coverage detected