Repository: pytorch/PiPPy

Issue:
Check if remap_qualname still works after refactorization



Discussion:
Can confirm it does not work after the re-factorization. 

Resolving patch:
diff --git a/pippy/IR.py b/pippy/IR.py
--- a/pippy/IR.py
+++ b/pippy/IR.py
@@ -492,20 +492,25 @@ class QualnameMapMixin:
 
     def __init__(
         self,
-        qualname_mapping: Dict[str, str] = None,
+        splitter_qualname_map: Dict[str, str] = None,
+        tracer_qualname_map: Dict[str, str] = None,
     ):
         self.new_to_old_qualname_mapping: Dict[str, str] = (
-            qualname_mapping or {}
+            splitter_qualname_map or {}
         )
+        self.tracer_qualname_map = tracer_qualname_map
 
     def remap_qualname(self, qualname: str):
         # TODO: annoying
         if qualname.startswith("split_gm."):
             qualname = qualname[len("split_gm.") :]
 
-        # The qualname map does not store recursive items, thus,
-        # when passed a qualname with leaves, we need to perform longest prefix match
-        if qualname not in self.new_to_old_qualname_mapping:
+        name_before_split = None
+        if qualname in self.new_to_old_qualname_mapping:
+            name_before_split = self.new_to_old_qualname_mapping[qualname]
+        else:
+            # The qualname map does not store recursive items, thus,
+            # when passed a qualname with leaves, we need to perform longest prefix match
             # Split from the right, one each time
             split_names = qualname.rsplit(".", 1)
             leaf = split_names[-1]
@@ -513,25 +518,34 @@ def remap_qualname(self, qualname: str):
                 prefix = split_names[0]
                 if prefix in self.new_to_old_qualname_mapping:
                     old_prefix = self.new_to_old_qualname_mapping[prefix]
-                    return ".".join([old_prefix, leaf])
+                    name_before_split = ".".join([old_prefix, leaf])
+                    break
                 split_names = prefix.rsplit(".", 1)
                 leaf = ".".join([split_names[-1], leaf])
 
-        # Either full name match, or key not found
-        return self.new_to_old_qualname_mapping[qualname]
+        if name_before_split is None:
+            raise RuntimeError(f"Could not find mapping for {qualname}")
+
+        if self.tracer_qualname_map is not None:
+            return self.tracer_qualname_map[name_before_split]
+        else:
+            return name_before_split
 
 
 class Pipe(QualnameMapMixin, torch.nn.Module):
     def __init__(
         self,
         split_gm: fx.GraphModule,
-        qualname_mapping: Dict[str, str],
+        splitter_qualname_map: Dict[str, str],
         num_stages: int,
         has_loss_and_backward: bool,
         loss_spec,
+        tracer_qualname_map: Optional[Dict[str, str]] = None,
     ):
         # TODO: is there a way not to hard wire init?
-        QualnameMapMixin.__init__(self, qualname_mapping)
+        QualnameMapMixin.__init__(
+            self, splitter_qualname_map, tracer_qualname_map
+        )
         torch.nn.Module.__init__(self)
         self.split_gm: fx.GraphModule = split_gm
         self.executor: DetachExecutor = DetachExecutor(self.split_gm)
@@ -728,10 +742,10 @@ def split_callback(n: fx.Node):
             return part_idx
 
         # Ask split_module to return mapping from new qualname to old qualname
-        qualname_map: Dict[str, str] = {}
+        splitter_qualname_map: Dict[str, str] = {}
         # TODO: what does split do with module invocations? does it move the modules
         # into the submodules?
-        split = split_module(traced, mod, split_callback, qualname_map)
+        split = split_module(traced, mod, split_callback, splitter_qualname_map)
         # a (custom) tracer can produce dead code like orphan get_attr nodes
         split.graph.eliminate_dead_code()
 
@@ -783,13 +797,15 @@ def move_param_to_callee(
             # Update qualname mapping
             # New qualname will have submodule prefix
             new_qualname = f"{callee_name}.{new_param_name}"
-            if node.target in qualname_map:
-                # Just in case the target name is already in the qualname_map
+            if node.target in splitter_qualname_map:
+                # Just in case the target name is already in the splitter_qualname_map
                 # returned by split_module() -- we update the mapping using the
                 # new name as a new key
-                qualname_map[new_qualname] = qualname_map.pop(node.target)
+                splitter_qualname_map[new_qualname] = splitter_qualname_map.pop(
+                    node.target
+                )
             else:
-                qualname_map[new_qualname] = node.target
+                splitter_qualname_map[new_qualname] = node.target
 
             ph_counter = 0
             for sn in callee.graph.nodes:
@@ -1005,12 +1021,17 @@ def move_param_to_callee(
                 "Pipeline is in inference mode, backward pass not generated"
             )
 
+        # Tracer may modify qualname, get the qualname mapping before and after tracing.
+        # This qualname mapping is different from the mapping before and after splitting.
+        tracer_qualname_map = Pipe._get_param_buffer_mapping(mod, traced)
+
         return Pipe(
             split,
-            qualname_map,
+            splitter_qualname_map,
             num_stages,
             has_loss_and_backward,
             generated_loss_spec,
+            tracer_qualname_map,
         )
 
     @staticmethod
@@ -1117,6 +1138,52 @@ def __str__(self):
     def __repr__(self):
         return self.split_gm.__repr__()
 
+    # TODO: this util comes from pytorch/pytorch#115462, delete it from PiPPy
+    # when PyTorch 2.3 comes with support, or when PiPPy migrates from
+    # `_export_to_torch_ir` to export + unflattener.
+    @staticmethod
+    def _get_param_buffer_mapping(
+        original_module: torch.nn.Module,
+        traced_module: torch.nn.Module,
+    ) -> Dict[str, str]:
+        """
+        Returns a mapping of parameter/buffer names from the new module to the
+        original model. This is to help with restoring the FQN for parameter/buffers
+        of a traced module to what the original module contains.
+        """
+
+        param_lookup: Dict[int, List[str]] = {}
+        buffer_lookup: Dict[int, List[str]] = {}
+        for name, param in original_module.named_parameters(
+            remove_duplicate=False
+        ):
+            param_lookup.setdefault(id(param), []).append(name)
+        for name, buffer in original_module.named_buffers(
+            remove_duplicate=False
+        ):
+            buffer_lookup.setdefault(id(buffer), []).append(name)
+
+        param_buffer_table: Dict[str, str] = {}
+        for dynamo_name, dynamo_param in traced_module.named_parameters(
+            remove_duplicate=False
+        ):
+            assert dynamo_name not in param_buffer_table
+            if id(dynamo_param) in param_lookup:
+                param_buffer_table[dynamo_name] = param_lookup[
+                    id(dynamo_param)
+                ].pop()
+
+        for dynamo_name, dynamo_buffer in traced_module.named_buffers(
+            remove_duplicate=False
+        ):
+            assert dynamo_name not in param_buffer_table
+            if id(dynamo_buffer) in buffer_lookup:
+                param_buffer_table[dynamo_name] = buffer_lookup[
+                    id(dynamo_buffer)
+                ].pop()
+
+        return param_buffer_table
+
 
 class PipeSplitWrapper(torch.nn.Module):
     class SplitPoint(Enum):


Test patch:
diff --git a/test/test_ir.py b/test/test_ir.py
--- a/test/test_ir.py
+++ b/test/test_ir.py
@@ -5,7 +5,6 @@
 import unittest
 from typing import NamedTuple
 
-import pippy.fx
 import torch
 
 from pippy.IR import (
@@ -26,11 +25,6 @@
 )
 
 
-@pippy.fx.wrap
-def arange_wrapper(*args, **kwargs):
-    return torch.arange(*args, **kwargs)
-
-
 class ExampleCode(torch.nn.Module):
     def __init__(self):
         super().__init__()
@@ -74,6 +68,7 @@ def setUp(self):
         mods += [mods[0]]
         self.seq = torch.nn.Sequential(*mods)
         self.ec = ExampleCode()
+        self.example_inputs = (torch.randn(50, 512),)
 
     def test_sequential(self):
         pipe_seq = PipeSequential.from_sequential(self.seq)
@@ -337,39 +332,6 @@ def forward(self, x, target):
             v_ref = ref_grads[k_ref]
             torch.testing.assert_close(v_test, v_ref)
 
-    def test_custom_tracer_serialization(self):
-        class CustomTracer(pippy.fx.Tracer):
-            def trace(self, root, concrete_args=None):
-                rv = super().trace(root, concrete_args)
-                for node in rv.nodes:
-                    if node.target == arange_wrapper:
-                        node.target = torch.arange
-                        node.meta.clear()
-                return rv
-
-        class FooMod(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.param = torch.nn.Parameter(torch.randn(1))
-
-            def forward(self, x):
-                return x + arange_wrapper(x.shape[-1]) + 1 + torch.zeros(1)
-
-        fm = FooMod()
-
-        tracer = CustomTracer()
-        pipe = Pipe.from_tracing(fm, tracer=tracer)
-
-        with tempfile.TemporaryDirectory() as d:
-            with open(d + "tmp.pkl", "wb") as f:
-                pickle.dump(pipe.split_gm.submod_0, f)
-
-            with open(d + "tmp.pkl", "rb") as f:
-                loaded = pickle.load(f)
-
-        x = torch.randn(5, 3)
-        torch.testing.assert_close(pipe.split_gm.submod_0(x), loaded(x))
-
     def test_direct_serialization_recursion_depth(self):
         class LomgBoi(torch.nn.Module):
             def forward(self, x):
@@ -814,31 +776,8 @@ def forward(
             chunks_merged_masked["multiplied"], ref_out["multiplied"]
         )
 
-    def test_remap_qualname_transmit(self):
-        ec_pipe = Pipe.from_tracing(self.ec, MultiUseParameterConfig.TRANSMIT)
-
-        # Get the first field of all tuples, i.e. names
-        old_named_params = zip(*list(self.ec.named_parameters()))
-        old_names = list(old_named_params)[0]
-
-        # Check qualname mapping for pipe
-        for new_name, _ in ec_pipe.named_parameters():
-            old_name = ec_pipe.remap_qualname(new_name)
-            # print(f"{new_name} -> {old_name}")
-            assert (
-                old_name in old_names
-            ), f"Remapped parameter {old_name} not found in {old_names}"
-
-        # Check qualname mapping for submodule
-        for _, stage_mod in ec_pipe.split_gm.named_children():
-            for new_name, _ in stage_mod.named_parameters():
-                old_name = stage_mod.remap_qualname(new_name)
-                assert (
-                    old_name in old_names
-                ), f"Remapped parameter {old_name} not found in {old_names}"
-
-    def test_remap_qualname_replicate(self):
-        ec_pipe = Pipe.from_tracing(self.ec, MultiUseParameterConfig.REPLICATE)
+    def test_remap_qualname(self):
+        ec_pipe = Pipe.from_tracing(self.ec, 1, self.example_inputs)
 
         # Get the first field of all tuples, i.e. names
         old_named_params = zip(*list(self.ec.named_parameters()))
@@ -853,12 +792,15 @@ def test_remap_qualname_replicate(self):
             ), f"Remapped parameter {old_name} not found in {old_names}"
 
         # Check qualname mapping for submodule
+        # Not supported at the moment
+        """
         for _, stage_mod in ec_pipe.split_gm.named_children():
             for new_name, _ in stage_mod.named_parameters():
                 old_name = stage_mod.remap_qualname(new_name)
                 assert (
                     old_name in old_names
                 ), f"Remapped parameter {old_name} not found in {old_names}"
+        """
 
 
 if __name__ == "__main__":
diff --git a/test/test_pipe.py b/test/test_pipe.py
--- a/test/test_pipe.py
+++ b/test/test_pipe.py
@@ -89,6 +89,17 @@ def run_worker(args, model_class):
     torch.testing.assert_close(out, ref_out)
     print(f"equivalence test passed {torch.sum(out)} ref {torch.sum(ref_out)}")
 
+    # Check qualname
+    old_names = set(mod.state_dict().keys())
+    new_names = set(pipe.state_dict().keys())
+    for new_name in new_names:
+        old_name = pipe.remap_qualname(new_name)
+        assert (
+            old_name in old_names
+        ), f"Remapped parameter {old_name} not found in {old_names}"
+        print(f"{new_name} -> {old_name}")
+    print("Qualname check passed")
+
 
 def main(args=None):
     parser = argparse.ArgumentParser()