A module that uses torch.fx.Interpreter to execute instead of the usual codegen that GraphModule uses. This provides better stack trace information and makes it easier to debug execution.
| 45 | |
| 46 | |
| 47 | class InterpreterModule(torch.nn.Module): |
| 48 | """A module that uses torch.fx.Interpreter to execute instead of the usual |
| 49 | codegen that GraphModule uses. This provides better stack trace information |
| 50 | and makes it easier to debug execution. |
| 51 | """ |
| 52 | |
| 53 | def __init__( |
| 54 | self, |
| 55 | graph: torch.fx.Graph, |
| 56 | module_call_signature: Optional[ModuleCallSignature], |
| 57 | ): |
| 58 | super().__init__() |
| 59 | self.graph = graph |
| 60 | self.graph.owning_module = self |
| 61 | self.module_call_signature = module_call_signature |
| 62 | |
| 63 | def forward(self, *args, **kwargs): |
| 64 | assert self.graph_module is not None, "Didn't finalize this InterpreterModule" |
| 65 | if torch._dynamo.is_compiling(): |
| 66 | # Dynamo cannot trace through torch.fx.Interpreter, so fall back to |
| 67 | # GraphModule codegen in this instance. |
| 68 | return self.graph_module(*args, **kwargs) |
| 69 | else: |
| 70 | if kwargs: |
| 71 | # Handle **kwargs. FX only natively supports positional |
| 72 | # arguments (through placeholders). So in order to pass in |
| 73 | # kwargs, we must correspond the names of the placeholders with |
| 74 | # the keys in the kwarg dict. |
| 75 | arg_list = list(args) |
| 76 | kwarg_names = self.arg_names[len(arg_list) :] |
| 77 | for kwarg_name in kwarg_names: |
| 78 | if kwarg_name in kwargs: |
| 79 | arg_list.append(kwargs[kwarg_name]) |
| 80 | |
| 81 | # Assert that the kwargs passed in exactly match the positional |
| 82 | # arguments specified by the GraphModule. This should be |
| 83 | # guaranteed by the unflattening process. |
| 84 | assert len(kwarg_names) == len(kwargs) |
| 85 | assert len(arg_list) == len(self.arg_names) |
| 86 | args = tuple(arg_list) |
| 87 | |
| 88 | return torch.fx.Interpreter(self, graph=self.graph).run( |
| 89 | *args, enable_io_processing=False |
| 90 | ) |
| 91 | |
| 92 | def finalize(self): |
| 93 | # We need to "finalize" because GraphModule populates its own state_dict |
| 94 | # based on the get_attrs observed in the graph. So we need to fully |
| 95 | # construct the graph and call _sink_params before generating this |
| 96 | # GraphModule. |
| 97 | |
| 98 | # need to set `graph_module` directly on the dict to avoid it getting |
| 99 | # registered as a submodule. |
| 100 | self.__dict__["graph_module"] = torch.fx.GraphModule(self, self.graph) |
| 101 | self.graph.lint() |
| 102 | |
| 103 | # Cache arg names for kwarg handling (see forward()) |
| 104 | self.arg_names = [] |
no outgoing calls
no test coverage detected
searching dependent graphs…