Update

[ghstack-poisoned]
pytorch · May 15, 2024 · c1b4c08 · c1b4c08
1 parent 76216c5
commit c1b4c08
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 22 deletions.
diff --git a/.github/workflows/unit_test_4gpu.yaml b/.github/workflows/unit_test_4gpu.yaml
@@ -30,8 +30,8 @@ jobs:
         python -m pip install --force-reinstall --pre torch --index-url https://download.pytorch.org/whl/nightly/cu121
         python -m pip install -r requirements.txt
         python -m pip install -r dev-requirements.txt
-        python ./test_runner.py
-        mv outputs artifacts-to-be-uploaded
+        mkdir artifacts-to-be-uploaded
+        python ./test_runner.py artifacts-to-be-uploaded
   # upload-coverage:
   #     - name: Upload Coverage to Codecov
   #       uses: codecov/codecov-action@v3
diff --git a/test_runner.py b/test_runner.py
@@ -4,9 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-import glob
+import argparse
 import os
-import shutil
 import subprocess
 from collections import defaultdict
 from dataclasses import dataclass
@@ -18,6 +17,11 @@
     import tomli as tomllib
 
 
+parser = argparse.ArgumentParser()
+parser.add_argument("output_dir")
+args = parser.parse_args()
+
+
 @dataclass
 class OverrideDefinitions:
     """
@@ -29,7 +33,6 @@ class OverrideDefinitions:
 
 
 CONFIG_DIR = "./train_configs"
-test_checkpoint_dir = "./test_runner_checkpoint"
 
 """
 key is the config file name and value is a list of OverrideDefinitions
@@ -40,25 +43,39 @@ class OverrideDefinitions:
 integration_tests_flavors["debug_model.toml"] = [
     OverrideDefinitions(
         [
-            ["--training.compile"],
+            [
+                f"--job.dump_folder {args.output_dir}/default/",
+            ],
+        ],
+        "Default",
+    ),
+    OverrideDefinitions(
+        [
+            [
+                "--training.compile",
+                f"--job.dump_folder {args.output_dir}/1d_compile/",
+            ],
         ],
         "1D compile",
     ),
     OverrideDefinitions(
         [
-            ["--training.tensor_parallel_degree 2 --model.norm_type=rmsnorm"],
+            [
+                "--training.tensor_parallel_degree 2 --model.norm_type=rmsnorm",
+                f"--job.dump_folder {args.output_dir}/eager_2d/",
+            ],
         ],
         "Eager mode 2DParallel",
     ),
     OverrideDefinitions(
         [
             [
                 "--checkpoint.enable_checkpoint",
-                f"--checkpoint.folder {test_checkpoint_dir}_full_checkpoint",
+                f"--job.dump_folder {args.output_dir}/full_checkpoint/",
             ],
             [
                 "--checkpoint.enable_checkpoint",
-                f"--checkpoint.folder {test_checkpoint_dir}_full_checkpoint",
+                f"--job.dump_folder {args.output_dir}/full_checkpoint/",
                 "--training.steps 20",
             ],
         ],
@@ -68,7 +85,7 @@ class OverrideDefinitions:
         [
             [
                 "--checkpoint.enable_checkpoint",
-                f"--checkpoint.folder {test_checkpoint_dir}_model_weights_only_fp32",
+                f"--job.dump_folder {args.output_dir}/model_weights_only_fp32/",
                 "--checkpoint.model_weights_only",
             ],
         ],
@@ -78,7 +95,7 @@ class OverrideDefinitions:
         [
             [
                 "--checkpoint.enable_checkpoint",
-                f"--checkpoint.folder {test_checkpoint_dir}_model_weights_only_bf16",
+                f"--job.dump_folder {args.output_dir}/model_weights_only_bf16/",
                 "--checkpoint.model_weights_only",
                 "--checkpoint.export_dtype bfloat16",
             ],
@@ -118,15 +135,5 @@ def run_test(test_flavor: OverrideDefinitions, full_path: str):
             config = tomllib.load(f)
             is_integration_test = config["job"].get("use_for_integration_test", False)
             if is_integration_test:
-                test_flavors = [OverrideDefinitions()] + integration_tests_flavors[
-                    config_file
-                ]
-
-                for test_flavor in test_flavors:
+                for test_flavor in integration_tests_flavors[config_file]:
                     run_test(test_flavor, full_path)
-
-                    # Deleting checkpoint folder from test
-                    dir_list = glob.iglob(f"{test_checkpoint_dir}_*")
-                    for path in dir_list:
-                        if os.path.exists(path):
-                            shutil.rmtree(path)