ThunderFX failure: KeyError: 'l_stack0_' #1293
Description
On the latest Thunder (dafc79d) I tried the snippet from
#1174 (comment)
and got the following error:
File ~/dev/lightning-thunder/thunder/dynamo/compiler.py:81, in ThunderCompiler.__call__(self, gm, sample_args)
77 recompile_graph(gm)
79 # The whole graph may not be supported by `thunder`, so we split it in `thunder` supported sections
80 # and unsupported sections which are passed to `torch.compile(backend='inductor')`
---> 81 split_module, subgraph_info = _splitter(gm, self._thunder_jit, self._torch_compile, sample_args)
82 self.subgraph_infos.append(subgraph_info)
83 return split_module
File ~/dev/lightning-thunder/thunder/dynamo/splitter.py:133, in _splitter(gm, thunder_jit, torch_inductor, _unused_sample_args)
130 return partition_cnt
132 # `split_module` iterates over nodes and determines the partition to place them based on the callback.
--> 133 split_gm: torch.fx.GraphModule = split_module(
134 gm, root_m=None, split_callback=callback, keep_original_order=True, keep_original_node_name=True
135 )
137 def is_thunder_supported_partition(node: torch.fx.Node) -> bool:
138 return node.name.startswith("submod") and int(node.name.replace("submod_", "")) in supported_partitions
File ~/dev/pytorch/main/torch/fx/passes/split_module.py:513, in split_module.<locals>.<lambda>(n)
510 for node in m.graph.nodes:
511 if node.op == "output":
512 base_mod_graph.output(
--> 513 torch.fx.graph.map_arg(node.args[0], lambda n: base_mod_env[n.name])
514 ) # noqa: B950
516 return torch.fx.graph_module.GraphModule(base_mod_attrs, base_mod_graph)
BackendCompilerFailed: backend='<thunder.dynamo.compiler.ThunderCompiler object at 0x7fd27cb43310>' raised:
KeyError: 'l_stack0_'My PyTorch version is '2.4.0a0+git3827810'.