fix the missing argument in test and typo (#5730)

This PR fixes the issue mentioned in [PR5722](#5722) that causes the hangs in the nv-torch-latest-v100 tests. Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com>
microsoft · Jul 8, 2024 · 774b897 · 774b897
1 parent 8411816
commit 774b897
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/deepspeed/checkpoint/ds_to_universal.py b/deepspeed/checkpoint/ds_to_universal.py
@@ -69,7 +69,7 @@ def parse_arguments():
                         dest='strict',
                         action='store_false',
                         help='Do not perform validity checks on converted checkpoint.')
-    parser.add_argument('--inject-missing-state',
+    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()

diff --git a/tests/unit/checkpoint/test_universal_checkpoint.py b/tests/unit/checkpoint/test_universal_checkpoint.py
@@ -110,7 +110,8 @@ def train_save_convert(ds_config, hidden_dim, load_optim, use_torch_adam, dtype,
                            num_extract_workers=1,
                            num_merge_workers=1,
                            keep_temp_folder=False,
-                           strict=True)
+                           strict=True,
+                           inject_missing_state=False)
 
     dist.barrier()
     if dist.get_rank() == 0: