refine && add optim placeholder

clemente0731 · Jul 3, 2024 · 33e72a3 · 33e72a3
1 parent 7a1228b
commit 33e72a3
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 5 deletions.
diff --git a/pprobe/bootstrap/hook_setup.py b/pprobe/bootstrap/hook_setup.py
@@ -18,16 +18,25 @@ def __init__(self, module, module_fullname):
         ### ENABLE FIELD
         ####################
         self.pprobe_enabled = self.pprobe_toggle.get_toggle("PPROBE_ENABLE")
-        self.torch_reproduce_enabled = self.pprobe_toggle.get_toggle("TORCH_REPRODUCE")
 
+        #########################
+        ## TORCH-CONTROL
+        #########################
+        self.torch_reproduce_enabled = self.pprobe_toggle.get_toggle("TORCH_REPRODUCE")
         self.torch_catch_step_enabled = self.pprobe_toggle.get_toggle(
             "TORCH_CATCH_STEP"
         )
         self.torch_catch_loss_enabled = self.pprobe_toggle.get_toggle(
             "TORCH_CATCH_LOSS"
         )
         self.torch_catch_lr_enabled = self.pprobe_toggle.get_toggle("TORCH_CATCH_LR")
+        self.torch_catch_memory_enabled = self.pprobe_toggle.get_toggle(
+            "TORCH_CATCH_MEMORY"
+        )
 
+        #########################
+        ## TORCH-SNAPSHOT
+        #########################
         self.torch_dump_op_enabled = self.pprobe_toggle.get_toggle("TORCH_DUMP_OP")
         self.torch_dump_dist_enabled = self.pprobe_toggle.get_toggle("TORCH_DUMP_DIST")
         self.torch_dump_module_enabled = self.pprobe_toggle.get_toggle(
@@ -36,9 +45,10 @@ def __init__(self, module, module_fullname):
         self.torch_dump_optim_enabled = self.pprobe_toggle.get_toggle(
             "TORCH_DUMP_OPTIM"
         )
-        self.torch_dump_memory_enabled = self.pprobe_toggle.get_toggle(
-            "TORCH_DUMP_MEMORY"
-        )
+
+        #########################
+        ## TORCH-TESTCASE DUMP
+        #########################
         self.torch_test_dump_op_enabled = self.pprobe_toggle.get_toggle(
             "TORCH_TEST_DUMP_OP"
         )
@@ -48,12 +58,17 @@ def __init__(self, module, module_fullname):
         self.torch_test_dump_module_enabled = self.pprobe_toggle.get_toggle(
             "TORCH_TEST_DUMP_MODULE"
         )
+
+        #########################
+        ## TORCH-SCAN
+        #########################
         self.torch_perf_issue_enabled = self.pprobe_toggle.get_toggle(
             "TORCH_PERF_ISSUE"
         )
         self.torch_torch_trace_file_enabled = self.pprobe_toggle.get_toggle(
             "TORCH_TRACE_FILE"
         )
+
         ####################
         ### INIT HOOK
         ####################

diff --git a/pprobe/bootstrap/hooks/pytorch_optim.py b/pprobe/bootstrap/hooks/pytorch_optim.py
diff --git a/pprobe/toggle/hook.toggle.default b/pprobe/toggle/hook.toggle.default
@@ -10,6 +10,7 @@ TORCH_REPRODUCE=True
 TORCH_CATCH_STEP=False
 TORCH_CATCH_LOSS=False
 TORCH_CATCH_LR=False
+TORCH_CATCH_MEMORY=False
 
 #########################
 ## TORCH-SNAPSHOT
@@ -18,7 +19,6 @@ TORCH_DUMP_OP=False
 TORCH_DUMP_MODULE=False
 TORCH_DUMP_DIST=False
 TORCH_DUMP_OPTIM=False
-TORCH_DUMP_MEMORY=False
 
 #########################
 ## TORCH-TESTCASE DUMP