Add an argument to enable the injection of missing state during the conversion of universal checkpoints

deepspeed/checkpoint/ds_to_universal.py

@@ -32,6 +32,8 @@
     SUB_PARAM_SHAPE,
     VOCAB_TENSOR,
     UNIVERSAL_CHECKPOINT_INFO,
+    UNIVERSAL_CHECKPOINT_VERSION_KEY,
+    UNIVERSAL_CHECKPOINT_VERSION_VALUE,
     VOCABULARY_PARAMETER_PATTERNS,
     PIPELINE_REPLICATED_PARAMETER_PATTERNS,
     TP_REPLICATED_PARAMETER_PATTERNS,
@@ -65,6 +67,9 @@ def parse_arguments():
                         dest='strict',
                         action='store_false',
                         help='Do not perform validity checks on converted checkpoint.')
+    parser.add_argument('--inject-missing-state',
+                        action='store_true',
+                        help='Inject missing checkpoint state into the checkpoint if it is absent.')
     args = parser.parse_args()
     print(f'args = {args}')
     return args
@@ -344,6 +349,15 @@ def _save_optimizer_state(args, ds_checkpoint):
     _save_checkpoint(output_file_path, output_sd)
 
 
+def _inject_missing_state(ds_checkpoint):
+    if UNIVERSAL_CHECKPOINT_INFO not in ds_checkpoint.global_state:
+        sd = torch.load(ds_checkpoint.mp_rank_files[0], map_location=torch.device('cpu'))
+        if UNIVERSAL_CHECKPOINT_INFO not in sd:
+            ds_checkpoint.global_state[UNIVERSAL_CHECKPOINT_INFO] = {}
+            ds_checkpoint.global_state[UNIVERSAL_CHECKPOINT_INFO][
+                UNIVERSAL_CHECKPOINT_VERSION_KEY] = UNIVERSAL_CHECKPOINT_VERSION_VALUE
+
+
 def _check_for_required_state(ds_checkpoint):
     universal_checkpoint_info = ds_checkpoint.get_checkpoint_info(UNIVERSAL_CHECKPOINT_INFO)
     assert universal_checkpoint_info is not None, f'Required {UNIVERSAL_CHECKPOINT_INFO} state is missing in checkpoint. Verify that client creates this state.'
@@ -355,7 +369,10 @@ def main(args):
     print(f'Converting DeepSpeed checkpoint in {args.input_folder} to Universal checkpoint in {args.output_folder}')
 
     ds_checkpoint = DeepSpeedCheckpoint(args.input_folder)
-    _check_for_required_state(ds_checkpoint)
+    if args.inject_missing_state:
+        _inject_missing_state(ds_checkpoint)
+    else:
+        _check_for_required_state(ds_checkpoint)
 
     iteration = ds_checkpoint.get_iteration()
     #_create_latest_file(args.output_folder, iteration)