Fix bug in graph partitioner and update graph signature after partitioning. #125133

tugsbayasgalan · 2024-04-29T02:57:32Z

Summary: This fix does two things:

When we add inputs from partioner to the top level graph module, we insert in the order of partioner which is not guaranteed to be same as original graph inputs. This PR fixes that.
When we replace autograd ops with HOP, we create new submodules and access their outputs via getitem calls. As a result, previous node names associated with getitem gets updated, resulting in the graph being different from produced graph signature. So I just update the graph signature accordingly.

Test Plan: CI

Differential Revision: D56688411

pytorch-bot · 2024-04-29T02:57:35Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/125133

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❌ 27 New Failures, 1 Unrelated Failure

As of commit cddf9d6 with merge base faf0015 ():

NEW FAILURES - The following jobs have failed:

inductor / cuda12.1-py3.10-gcc9-sm86 / test (aot_inductor_huggingface, 1, 1, linux.g5.4xlarge.nvidia.gpu) (gh)
##[error]The operation was canceled.
Lint / lintrunner-noclang / linux-job (gh)
>>> Lint for torch/fx/passes/runtime_assert.py:
pull / linux-focal-cuda12.1-py3.10-gcc9 / test (default, 2, 5, linux.4xlarge.nvidia.gpu) (gh)
export/test_export.py::TestExport::test_unbacked_deferred_runtime_retrace
pull / linux-focal-cuda12.1-py3.10-gcc9 / test (default, 3, 5, linux.4xlarge.nvidia.gpu) (gh)
dynamo/test_cpp_guard_manager.py::CppGuardManagerMiscTests::test_validate_outputs_unbacked_cpp_guard_manager
pull / linux-focal-cuda12.1-py3.10-gcc9 / test (default, 5, 5, linux.4xlarge.nvidia.gpu) (gh)
dynamo/test_dynamic_shapes.py::DynamicShapesMiscTests::test_validate_outputs_unbacked_dynamic_shapes
pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 2, 5, linux.g5.4xlarge.nvidia.gpu) (gh)
export/test_export.py::TestExport::test_unbacked_deferred_runtime_retrace
pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 3, 5, linux.g5.4xlarge.nvidia.gpu) (gh)
export/test_retraceability.py::RetraceExportTestExport::test_unbacked_deferred_runtime_retrace_retraceability
pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 4, 5, linux.g5.4xlarge.nvidia.gpu) (gh)
dynamo/test_cpp_guard_manager.py::CppGuardManagerMiscTests::test_validate_outputs_unbacked_cpp_guard_manager
pull / linux-focal-py3.11-clang10 / test (crossref, 2, 2, linux.2xlarge) (gh)
export/test_export.py::TestExport::test_unbacked_deferred_runtime_retrace
pull / linux-focal-py3.11-clang10 / test (default, 2, 3, linux.2xlarge) (gh)
export/test_export.py::TestExport::test_unbacked_deferred_runtime_retrace
pull / linux-focal-py3.11-clang10 / test (default, 3, 3, linux.2xlarge) (gh)
dynamo/test_cpp_guard_manager.py::CppGuardManagerMiscTests::test_validate_outputs_unbacked_cpp_guard_manager
pull / linux-focal-py3.12-clang10 / test (default, 2, 3, linux.2xlarge) (gh)
export/test_export.py::TestExport::test_unbacked_deferred_runtime_retrace
pull / linux-focal-py3.12-clang10 / test (default, 3, 3, linux.2xlarge) (gh)
dynamo/test_cpp_guard_manager.py::CppGuardManagerMiscTests::test_validate_outputs_unbacked_cpp_guard_manager
pull / linux-focal-py3.8-clang10 / test (crossref, 2, 2, linux.2xlarge) (gh)
export/test_export.py::TestExport::test_unbacked_deferred_runtime_retrace
pull / linux-focal-py3.8-clang10 / test (default, 2, 3, linux.2xlarge) (gh)
export/test_export.py::TestExport::test_unbacked_deferred_runtime_retrace
pull / linux-focal-py3.8-clang10 / test (default, 3, 3, linux.2xlarge) (gh)
dynamo/test_cpp_guard_manager.py::CppGuardManagerMiscTests::test_validate_outputs_unbacked_cpp_guard_manager
pull / linux-jammy-py3.10-clang15-asan / test (default, 2, 6, linux.4xlarge) (gh)
export/test_export.py::TestExport::test_unbacked_deferred_runtime_retrace
pull / linux-jammy-py3.10-clang15-asan / test (default, 3, 6, linux.4xlarge) (gh)
dynamo/test_cpp_guard_manager.py::CppGuardManagerMiscTests::test_validate_outputs_unbacked_cpp_guard_manager
pull / linux-jammy-py3.10-clang15-asan / test (default, 4, 6, linux.4xlarge) (gh)
export/test_serdes.py::SerDesExportTestExport::test_unbacked_deferred_runtime_retrace_serdes
pull / linux-jammy-py3.10-clang15-asan / test (default, 5, 6, linux.4xlarge) (gh)
dynamo/test_dynamic_shapes.py::DynamicShapesMiscTests::test_validate_outputs_unbacked_dynamic_shapes
pull / linux-jammy-py3.10-clang15-asan / test (default, 6, 6, linux.4xlarge) (gh)
export/test_retraceability.py::RetraceExportTestExport::test_unbacked_deferred_runtime_retrace_retraceability
pull / linux-jammy-py3.8-gcc11 / test (default, 2, 3, linux.2xlarge) (gh)
export/test_export.py::TestExport::test_unbacked_deferred_runtime_retrace
pull / linux-jammy-py3.8-gcc11 / test (default, 3, 3, linux.2xlarge) (gh)
dynamo/test_cpp_guard_manager.py::CppGuardManagerMiscTests::test_validate_outputs_unbacked_cpp_guard_manager
trunk / linux-focal-cuda12.1-py3.10-gcc9 / test (nogpu_AVX512, 1, 1, linux.2xlarge) (gh)
export/test_passes.py::TestPasses::test_runtime_assert_inline_constraints_for_cond
trunk / linux-focal-cuda12.1-py3.10-gcc9 / test (nogpu_NO_AVX2, 1, 1, linux.2xlarge) (gh)
export/test_export.py::TestExport::test_unbacked_deferred_runtime_retrace
trunk / macos-12-py3-arm64 / test (default, 1, 3, macos-m1-stable) (gh)
dynamo/test_cpp_guard_manager.py::CppGuardManagerMiscTests::test_validate_outputs_unbacked_cpp_guard_manager
trunk / macos-12-py3-arm64 / test (default, 2, 3, macos-m1-stable) (gh)
export/test_export.py::TestExport::test_unbacked_deferred_runtime_retrace

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

trunk / win-vs2019-cpu-py3 / test (default, 1, 3, windows.4xlarge.nonephemeral) (gh) (trunk failure)
[ FAILED ] LoggingTest.ExceptionWhat

This comment was automatically generated by Dr. CI and updates every 15 minutes.

SherlockNoMad · 2024-04-29T15:49:26Z

torch/export/_trace.py

+ if node.op == "output":
+ updated_output_names = node.args
+
+ # It could be updated because replace_set_grad_with_hop_pass adds new


addd a test case for this?

The internal diff has a torchrec unittest.

SherlockNoMad · 2024-04-29T15:49:34Z

torch/fx/passes/split_module.py

@@ -477,6 +477,16 @@ def instantiate_node_partition_mapping(node):
 orig_mod_attr_nodes: List[Node] = [
 orig_mod_env[key] for key in partition.inputs
 ]
+
+ # We actually need to insert the placeholder nodes in the original order


The internal diff has a torchrec unittest.

ydwu4 · 2024-04-29T15:58:29Z

torch/export/_trace.py

+ updated_output_names = node.args
+
+ # It could be updated because replace_set_grad_with_hop_pass adds new
+ # getitem nodes into the graph.


Should we move this to replace_set_grad_with_hop_pass or there are other cases that also require this pass? Or we can organize the change into a pass or a function?

We should probably have a simplified version of the internal unittest in OSS so that we could get an early signal when we add new features in OSS in the future.

Right now it is only replace_set_grad_with_hop_pass, but in theory, it could be useful for any pass that adds submodules in the graph. I can make it a free function.

zhxchen17 · 2024-04-29T22:00:28Z

torch/_export/passes/replace_set_grad_with_hop_pass.py

+ # It could be updated because replace_set_grad_with_hop_pass adds new
+ # getitem nodes into the graph.
+ assert len(original_output_nodes) == len(updated_output_nodes)
+ old_output_name_to_new_output_name = {}
+ for k, v in zip(*original_output_nodes, *updated_output_nodes):
+ if isinstance(k, torch.fx.Node) and isinstance(k, torch.fx.Node):
+ old_output_name_to_new_output_name[k.name] = v.name
+ # If there is constant return in the end, it should be true for updated outputs too
+ else:
+ constant_types = (int, float, str, type(None))
+ assert k == v
+ assert isinstance(k, constant_types)
+
+ buffers_to_mutate_copy = graph_signature.buffers_to_mutate.copy()
+ user_inputs_to_mutate_copy = graph_signature.user_inputs_to_mutate.copy()
+ for k in old_output_name_to_new_output_name:
+ if k in graph_signature.buffers_to_mutate:
+ graph_signature.buffers_to_mutate[
+ old_output_name_to_new_output_name[k]
+ ] = buffers_to_mutate_copy[k]
+ if k not in old_output_name_to_new_output_name.values():
+ del graph_signature.buffers_to_mutate[k]
+
+ if k in graph_signature.user_inputs_to_mutate:
+ graph_signature.user_inputs_to_mutate[
+ old_output_name_to_new_output_name[k]
+ ] = user_inputs_to_mutate_copy[k]
+ if k not in old_output_name_to_new_output_name.values():
+ del graph_signature.user_inputs_to_mutate[k]
+
+ for i, k in enumerate(graph_signature.user_outputs):
+ if k in old_output_name_to_new_output_name:
+ new_k = old_output_name_to_new_output_name[k]
+ graph_signature.user_outputs[i] = new_k


I think you can get away these regeneration code with:

with graph_module._set_replace_hook(ep.graph_signature.get_replace_hook(): _replace_set_grad_with_hop_pass(gm)

Hmm we actually cannot use ep here because this is before we create ep. We need it before ep because ep verifier won't pass as we have non-functional ops in the graph at this point.

qq: is ep.graph_signature the same as the type as the graph signature here? If so, we could still use the method?

A bunch of loop and mapping is difficult to parse what it does at first sight lol

discussed offline with @tugsbayasgalan, we should still use something more structural here. will follow up

zhxchen17

sorry wrong button clicked lol

ydwu4 · 2024-04-30T18:10:29Z

torch/_export/passes/replace_set_grad_with_hop_pass.py

@@ -160,7 +160,58 @@ def _maybe_inline_or_replace_with_hop(node: torch.fx.Node):
 return gm


-def replace_set_grad_with_hop_pass(gm: torch.fx.GraphModule):
+def replace_set_grad_with_hop_pass(gm: torch.fx.GraphModule, graph_signature):


nit: maybe a type annotation?

yeah adding a type annotation here is better.

ydwu4

Overall looks good to me. Left some minor comments on code structure.

zhxchen17

some nits.

zhxchen17 · 2024-05-01T14:58:06Z

torch/_export/passes/replace_set_grad_with_hop_pass.py

@@ -160,7 +160,58 @@ def _maybe_inline_or_replace_with_hop(node: torch.fx.Node):
 return gm


-def replace_set_grad_with_hop_pass(gm: torch.fx.GraphModule):
+def replace_set_grad_with_hop_pass(gm: torch.fx.GraphModule, graph_signature):


yeah adding a type annotation here is better.

zhxchen17 · 2024-05-01T14:58:25Z

torch/_export/passes/replace_set_grad_with_hop_pass.py

+ def _maybe_inline_or_replace_with_hop(node: torch.fx.Node):
+ if _is_set_grad_enabled_sub_mod(node, omit_if_same_with_ambient=True):
+ _replace_with_hop(node)
+ else:
+ _remove_set_grad_and_inline(node)


we can define this function out of the context manager.

facebook-github-bot · 2024-05-01T17:22:12Z