dump compile trace to CI output for debugging

[ghstack-poisoned]
pytorch · Dec 16, 2024 · 376f8f5 · 376f8f5
1 parent 86c8e55
commit 376f8f5
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/test_runner.py b/test_runner.py
@@ -413,6 +413,8 @@ def run_test(test_flavor: OverrideDefinitions, full_path: str, output_dir: str):
 
     for override_arg in test_flavor.override_args:
         cmd = f"CONFIG_FILE={full_path} NGPU={test_flavor.ngpu} LOG_RANK={all_ranks} ./run_llama_train.sh"
+        # dump compile trace for debugging purpose
+        cmd = f"TORCH_TRACE=\"{output_dir}/{test_name}/compile_trace\" " + cmd
         if test_name == "fsdp2_memory_estimation":
             cmd = (
                 f"CONFIG_FILE={full_path} NGPU={test_flavor.ngpu} LOG_RANK={all_ranks} "