fix several misrecorded rnn functions

heheda12345 · Mar 14, 2024 · 66ff6af · 66ff6af
1 parent 9362a1f
commit 66ff6af
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 8 deletions.
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -35,5 +35,5 @@ jobs:
           spack load cuda@11.8.0 /jb4mlxg
           spack load python@3.9.12%gcc@=11.3.0
           source ~/venv/frontend-env/bin/activate
-          srun --gres=gpu:v100:1 --exclusive ./scripts/pytest_with_preload.sh -vs test
-          FORCE_RUN_SKIPPED_TEST=1 srun --gres=gpu:v100:1 --exclusive ./scripts/pytest_with_preload.sh -vs test/test_model_blockdrop.py -k test_blockdrop_dyn
+          srun -p ja --gres=gpu:v100:1 --exclusive ./scripts/pytest_with_preload.sh -vs test
+          FORCE_RUN_SKIPPED_TEST=1 srun -p ja --gres=gpu:v100:1 --exclusive ./scripts/pytest_with_preload.sh -vs test/test_model_blockdrop.py -k test_blockdrop_dyn
diff --git a/frontend/guard_tracker.py b/frontend/guard_tracker.py
@@ -202,9 +202,10 @@ def as_fx_node(arg: Any) -> NodeArgs:
                     if common_device is not None and common_device != torch.device(
                             'cpu'):
                         cpu_node = var.as_fx_node()
-                        return self.fx_graph.create_node(
-                            "call_method", "to", (cpu_node,),
-                            {"device": common_device})
+                        # return self.fx_graph.create_node(
+                        #     "call_method", "to", (cpu_node,),
+                        #     {"device": common_device})
+                        return cpu_node
                 else:
                     # TODO: record all operation in SymInt or SymFloat
                     pass
@@ -352,6 +353,8 @@ def record_function(self,
 
             fx_node = self.fx_graph.create_node("call_method", func.__name__,
                                                 pargs, pkwargs)
+            if func.__name__ == 'tolist':
+                add_partial_var = False
             if add_partial_var:
                 self.partial_var = {
                     -1: [
@@ -1737,9 +1740,11 @@ def set_if_inplace_return() -> None:
                     ]
                 })
                 return
-            if hasattr(func,
-                       "__name__") and func.__name__ in ("flatten_parameters",
-                                                         "numel", "children"):
+            if hasattr(func, "__name__") and func.__name__ in (
+                    "flatten_parameters", "numel", "children",
+                    "named_parameters", "_weights_have_changed",
+                    "check_forward_args", "permute_hidden", "_check_input_dim",
+                    "parameters"):
                 return
             if hasattr(func, "__module__"
                       ) and func.__module__ == 'torch.autograd.profiler':