Update pretrain_gpt_alcf.py

argonne-lcf · Apr 4, 2024 · 7794fc0 · 7794fc0
1 parent 8c6c91f
commit 7794fc0
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/pretrain_gpt_alcf.py b/pretrain_gpt_alcf.py
@@ -571,18 +571,21 @@ def main():
                 model_provider,
                 ModelType.encoder_or_decoder,
                 forward_step,
-                args_defaults={'tokenizer_type': 'GPT2BPETokenizer'},
+                # args_defaults={'tokenizer_type': 'GPT2BPETokenizer'},
                 data_post_process=data_post_process
             )
-
-        prof.export_chrome_trace(f"{args.tensorboard_dir}/torch-trace-{RANK}-of-{WORLD_SIZE}.json")
+        args = get_args()
+        prof.export_chrome_trace(
+            f"{args.tensorboard_dir}"
+            "/torch-trace-{RANK}-of-{WORLD_SIZE}.json"
+        )
     else:
         model = pretrain(
             train_valid_test_datasets_provider,
             model_provider,
             ModelType.encoder_or_decoder,
             forward_step,
-            args_defaults={'tokenizer_type': 'GPT2BPETokenizer'},
+            # args_defaults={'tokenizer_type': 'GPT2BPETokenizer'},
             data_post_process=data_post_process
         )
     return model