# Owner(s): ["oncall: pt2"] # Copyright (c) Facebook, Inc. and its affiliates. # All rights reserved. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. from contextlib import ExitStack import torch import torch.fx.traceback as fx_traceback import torch.nn as nn import torch.utils._pytree as pytree from torch._decomp import decomposition_table from torch._dynamo.testing import normalize_gm from torch._functorch._aot_autograd.descriptors import ( BufferAOTInput, ParamAOTInput, PlainAOTInput, PlainAOTOutput, ) from torch._functorch._aot_autograd.fx_utils import ( get_all_input_and_grad_nodes, get_all_output_and_tangent_nodes, get_buffer_nodes, get_named_buffer_nodes, get_named_param_nodes, get_param_and_grad_nodes, get_param_nodes, get_plain_input_and_grad_nodes, get_plain_output_and_tangent_nodes, ) from torch._functorch.aot_autograd import ( aot_compile_joint_with_descriptors, aot_export_joint_with_descriptors, ) from torch.testing._internal.common_utils import run_tests, TestCase class TestAOTJointWithDescriptors(TestCase): def test_simple_linear_module(self): """Test basic linear module with aot_export_joint_with_descriptors""" class SimpleLinear(nn.Module): def __init__(self): super().__init__() self.linear = nn.Linear(3, 2) def forward(self, x): return self.linear(x) model = SimpleLinear() inputs = (torch.randn(4, 3),) with ExitStack() as stack: # Export joint with descriptors joint_with_descriptors = aot_export_joint_with_descriptors( stack, model, inputs, decompositions=decomposition_table ) # Test the exported graph structure graph_code = joint_with_descriptors.graph_module.print_readable( print_output=False, expanded_def=True ) self.assertExpectedInline( normalize_gm(graph_code), """\ class inner_f(torch.nn.Module): def forward( self, primals, tangents, ): primals_1: "f32[2, 3]" # ParamAOTInput(target='linear.weight') primals_2: "f32[2]" # ParamAOTInput(target='linear.bias') primals_3: "f32[4, 3]" # PlainAOTInput(idx=0) tangents_1: "f32[4, 2]" # TangentAOTInput(output=PlainAOTOutput(idx=0)) primals_1, primals_2, primals_3, tangents_1, = fx_pytree.tree_flatten_spec([primals, tangents], self._in_spec) transpose: "f32[3, 2]" = torch.ops.prims.transpose.default(primals_1, [1, 0]); primals_1 = None mm: "f32[4, 2]" = torch.ops.aten.mm.default(primals_3, transpose); transpose = None mul: "f32[4, 2]" = torch.ops.prims.mul.default(mm, 1.0); mm = None mul_1: "f32[2]" = torch.ops.prims.mul.default(primals_2, 1.0); primals_2 = None broadcast_in_dim: "f32[4, 2]" = torch.ops.prims.broadcast_in_dim.default(mul_1, [4, 2], [1]); mul_1 = None add: "f32[4, 2]" = torch.ops.prims.add.default(mul, broadcast_in_dim); mul = broadcast_in_dim = None transpose_1: "f32[2, 4]" = torch.ops.prims.transpose.default(tangents_1, [1, 0]) mm_1: "f32[2, 3]" = torch.ops.aten.mm.default(transpose_1, primals_3); transpose_1 = primals_3 = None transpose_2: "f32[3, 2]" = torch.ops.prims.transpose.default(mm_1, [1, 0]); mm_1 = None sum_1: "f32[2]" = torch.ops.prims.sum.default(tangents_1, [0]); tangents_1 = None broadcast_in_dim_1: "f32[1, 2]" = torch.ops.prims.broadcast_in_dim.default(sum_1, [1, 2], [1]); sum_1 = None as_strided: "f32[2]" = torch.ops.aten.as_strided.default(broadcast_in_dim_1, [2], [1]); broadcast_in_dim_1 = None transpose_3: "f32[2, 3]" = torch.ops.prims.transpose.default(transpose_2, [1, 0]); transpose_2 = None return pytree.tree_unflatten([ add, # PlainAOTOutput(idx=0) transpose_3, # GradAOTOutput(grad_of=ParamAOTInput(target='linear.weight')) as_strided, # GradAOTOutput(grad_of=ParamAOTInput(target='linear.bias')) None, # None ], self._out_spec) """, ) # Compile the result parallel_model_fn = aot_compile_joint_with_descriptors( joint_with_descriptors ) # Test functional correctness expected_output = model(*inputs) actual_output = parallel_model_fn( *dict(model.named_parameters()).values(), *inputs ) self.assertEqual(expected_output, actual_output) def test_conv_bn_module(self): """Test convolutional + batch norm module""" class ConvBN(nn.Module): def __init__(self): super().__init__() self.conv = nn.Conv2d(1, 3, 3, padding=1) self.bn = nn.BatchNorm2d(3) def forward(self, x): x = self.conv(x) x = self.bn(x) return torch.relu(x) model = ConvBN() model.train() # Important for batch norm inputs = (torch.randn(2, 1, 4, 4),) with ExitStack() as stack: # Export joint with descriptors joint_with_descriptors = aot_export_joint_with_descriptors( stack, model, inputs, decompositions=decomposition_table ) # Test the exported graph structure graph_code = joint_with_descriptors.graph_module.print_readable( print_output=False, expanded_def=True ) # Test parameter and buffer specs self.assertIn("conv.weight", joint_with_descriptors.params_spec) self.assertIn("conv.bias", joint_with_descriptors.params_spec) self.assertIn("bn.weight", joint_with_descriptors.params_spec) self.assertIn("bn.bias", joint_with_descriptors.params_spec) self.assertIn("bn.running_mean", joint_with_descriptors.buffers_spec) self.assertIn("bn.running_var", joint_with_descriptors.buffers_spec) self.assertIn("bn.num_batches_tracked", joint_with_descriptors.buffers_spec) # Expect test on the printed graph self.assertExpectedInline( normalize_gm(graph_code), """\ class inner_f(torch.nn.Module): def forward( self, primals, tangents, ): primals_1: "f32[3, 1, 3, 3]" # ParamAOTInput(target='conv.weight') primals_2: "f32[3]" # ParamAOTInput(target='conv.bias') primals_3: "f32[3]" # ParamAOTInput(target='bn.weight') primals_4: "f32[3]" # ParamAOTInput(target='bn.bias') primals_5: "f32[3]" # BufferAOTInput(target='bn.running_mean') primals_6: "f32[3]" # BufferAOTInput(target='bn.running_var') primals_7: "i64[]" # BufferAOTInput(target='bn.num_batches_tracked') primals_8: "f32[2, 1, 4, 4]" # PlainAOTInput(idx=0) tangents_1: "f32[2, 3, 4, 4]" # TangentAOTInput(output=PlainAOTOutput(idx=0)) primals_1, primals_2, primals_3, primals_4, primals_5, primals_6, primals_7, primals_8, tangents_1, = fx_pytree.tree_flatten_spec([primals, tangents], self._in_spec) convolution: "f32[2, 3, 4, 4]" = torch.ops.aten.convolution.default(primals_8, primals_1, primals_2, [1, 1], [1, 1], [1, 1], False, [0, 0], 1); primals_2 = None add: "i64[]" = torch.ops.prims.add.default(primals_7, 1); primals_7 = None var: "f32[3]" = torch.ops.prims.var.default(convolution, [0, 2, 3], 0.0) broadcast_in_dim: "f32[1, 3, 1, 1]" = torch.ops.prims.broadcast_in_dim.default(var, [1, 3, 1, 1], [1]); var = None sum_1: "f32[3]" = torch.ops.prims.sum.default(convolution, [0, 2, 3]) broadcast_in_dim_1: "f32[1, 3, 1, 1]" = torch.ops.prims.broadcast_in_dim.default(sum_1, [1, 3, 1, 1], [1]); sum_1 = None div: "f32[1, 3, 1, 1]" = torch.ops.prims.div.default(broadcast_in_dim_1, 32.0); broadcast_in_dim_1 = None add_1: "f32[1, 3, 1, 1]" = torch.ops.prims.add.default(broadcast_in_dim, 1e-05) rsqrt: "f32[1, 3, 1, 1]" = torch.ops.prims.rsqrt.default(add_1); add_1 = None broadcast_in_dim_2: "f32[2, 3, 4, 4]" = torch.ops.prims.broadcast_in_dim.default(div, [2, 3, 4, 4], [0, 1, 2, 3]) sub: "f32[2, 3, 4, 4]" = torch.ops.prims.sub.default(convolution, broadcast_in_dim_2); broadcast_in_dim_2 = None broadcast_in_dim_3: "f32[2, 3, 4, 4]" = torch.ops.prims.broadcast_in_dim.default(rsqrt, [2, 3, 4, 4], [0, 1, 2, 3]) mul: "f32[2, 3, 4, 4]" = torch.ops.prims.mul.default(sub, broadcast_in_dim_3); sub = broadcast_in_dim_3 = None squeeze: "f32[1, 3, 1]" = torch.ops.prims.squeeze.default(div, [3]); div = None squeeze_1: "f32[1, 3]" = torch.ops.prims.squeeze.default(squeeze, [2]); squeeze = None squeeze_2: "f32[3]" = torch.ops.prims.squeeze.default(squeeze_1, [0]); squeeze_1 = None squeeze_3: "f32[1, 3, 1]" = torch.ops.prims.squeeze.default(rsqrt, [3]); rsqrt = None squeeze_4: "f32[1, 3]" = torch.ops.prims.squeeze.default(squeeze_3, [2]); squeeze_3 = None squeeze_5: "f32[3]" = torch.ops.prims.squeeze.default(squeeze_4, [0]); squeeze_4 = None mul_1: "f32[3]" = torch.ops.prims.mul.default(squeeze_2, 0.1) mul_2: "f32[3]" = torch.ops.prims.mul.default(primals_5, 0.9); primals_5 = None add_2: "f32[3]" = torch.ops.prims.add.default(mul_1, mul_2); mul_1 = mul_2 = None squeeze_6: "f32[1, 3, 1]" = torch.ops.prims.squeeze.default(broadcast_in_dim, [3]); broadcast_in_dim = None squeeze_7: "f32[1, 3]" = torch.ops.prims.squeeze.default(squeeze_6, [2]); squeeze_6 = None squeeze_8: "f32[3]" = torch.ops.prims.squeeze.default(squeeze_7, [0]); squeeze_7 = None mul_3: "f32[3]" = torch.ops.prims.mul.default(squeeze_8, 1.032258064516129); squeeze_8 = None mul_4: "f32[3]" = torch.ops.prims.mul.default(mul_3, 0.1); mul_3 = None mul_5: "f32[3]" = torch.ops.prims.mul.default(primals_6, 0.9); primals_6 = None add_3: "f32[3]" = torch.ops.prims.add.default(mul_4, mul_5); mul_4 = mul_5 = None broadcast_in_dim_4: "f32[3, 1]" = torch.ops.prims.broadcast_in_dim.default(primals_3, [3, 1], [0]) broadcast_in_dim_5: "f32[3, 1, 1]" = torch.ops.prims.broadcast_in_dim.default(broadcast_in_dim_4, [3, 1, 1], [0, 1]); broadcast_in_dim_4 = None broadcast_in_dim_6: "f32[2, 3, 4, 4]" = torch.ops.prims.broadcast_in_dim.default(broadcast_in_dim_5, [2, 3, 4, 4], [1, 2, 3]); broadcast_in_dim_5 = None mul_6: "f32[2, 3, 4, 4]" = torch.ops.prims.mul.default(mul, broadcast_in_dim_6); mul = broadcast_in_dim_6 = None broadcast_in_dim_7: "f32[3, 1]" = torch.ops.prims.broadcast_in_dim.default(primals_4, [3, 1], [0]); primals_4 = None broadcast_in_dim_8: "f32[3, 1, 1]" = torch.ops.prims.broadcast_in_dim.default(broadcast_in_dim_7, [3, 1, 1], [0, 1]); broadcast_in_dim_7 = None broadcast_in_dim_9: "f32[2, 3, 4, 4]" = torch.ops.prims.broadcast_in_dim.default(broadcast_in_dim_8, [2, 3, 4, 4], [1, 2, 3]); broadcast_in_dim_8 = None add_4: "f32[2, 3, 4, 4]" = torch.ops.prims.add.default(mul_6, broadcast_in_dim_9); mul_6 = broadcast_in_dim_9 = None le: "b8[2, 3, 4, 4]" = torch.ops.prims.le.default(add_4, 0.0) where: "f32[2, 3, 4, 4]" = torch.ops.prims.where.default(le, 0.0, add_4); le = add_4 = None view_of: "f32[2, 3, 4, 4]" = torch.ops.prims.view_of.default(where) view_of_1: "f32[2, 3, 4, 4]" = torch.ops.prims.view_of.default(view_of); view_of = None view_of_2: "f32[2, 3, 4, 4]" = torch.ops.prims.view_of.default(view_of_1); view_of_1 = None view_of_3: "f32[2, 3, 4, 4]" = torch.ops.prims.view_of.default(view_of_2); view_of_2 = None le_1: "b8[2, 3, 4, 4]" = torch.ops.prims.le.default(view_of_3, 0.0); view_of_3 = None where_1: "f32[2, 3, 4, 4]" = torch.ops.prims.where.default(le_1, 0.0, tangents_1); le_1 = tangents_1 = None broadcast_in_dim_10: "f32[1, 3]" = torch.ops.prims.broadcast_in_dim.default(squeeze_2, [1, 3], [1]); squeeze_2 = None broadcast_in_dim_11: "f32[1, 3, 1]" = torch.ops.prims.broadcast_in_dim.default(broadcast_in_dim_10, [1, 3, 1], [0, 1]); broadcast_in_dim_10 = None broadcast_in_dim_12: "f32[1, 3, 1, 1]" = torch.ops.prims.broadcast_in_dim.default(broadcast_in_dim_11, [1, 3, 1, 1], [0, 1, 2]); broadcast_in_dim_11 = None sum_2: "f32[3]" = torch.ops.prims.sum.default(where_1, [0, 2, 3]) broadcast_in_dim_13: "f32[2, 3, 4, 4]" = torch.ops.prims.broadcast_in_dim.default(broadcast_in_dim_12, [2, 3, 4, 4], [0, 1, 2, 3]) sub_1: "f32[2, 3, 4, 4]" = torch.ops.prims.sub.default(convolution, broadcast_in_dim_13); broadcast_in_dim_13 = None mul_7: "f32[2, 3, 4, 4]" = torch.ops.prims.mul.default(where_1, sub_1); sub_1 = None sum_3: "f32[3]" = torch.ops.prims.sum.default(mul_7, [0, 2, 3]); mul_7 = None mul_8: "f32[3]" = torch.ops.prims.mul.default(sum_2, 0.03125) broadcast_in_dim_14: "f32[1, 3]" = torch.ops.prims.broadcast_in_dim.default(mul_8, [1, 3], [1]); mul_8 = None broadcast_in_dim_15: "f32[1, 3, 1]" = torch.ops.prims.broadcast_in_dim.default(broadcast_in_dim_14, [1, 3, 1], [0, 1]); broadcast_in_dim_14 = None broadcast_in_dim_16: "f32[1, 3, 1, 1]" = torch.ops.prims.broadcast_in_dim.default(broadcast_in_dim_15, [1, 3, 1, 1], [0, 1, 2]); broadcast_in_dim_15 = None mul_9: "f32[3]" = torch.ops.prims.mul.default(sum_3, 0.03125) mul_10: "f32[3]" = torch.ops.prims.mul.default(squeeze_5, squeeze_5) mul_11: "f32[3]" = torch.ops.prims.mul.default(mul_9, mul_10); mul_9 = mul_10 = None broadcast_in_dim_17: "f32[1, 3]" = torch.ops.prims.broadcast_in_dim.default(mul_11, [1, 3], [1]); mul_11 = None broadcast_in_dim_18: "f32[1, 3, 1]" = torch.ops.prims.broadcast_in_dim.default(broadcast_in_dim_17, [1, 3, 1], [0, 1]); broadcast_in_dim_17 = None broadcast_in_dim_19: "f32[1, 3, 1, 1]" = torch.ops.prims.broadcast_in_dim.default(broadcast_in_dim_18, [1, 3, 1, 1], [0, 1, 2]); broadcast_in_dim_18 = None mul_12: "f32[3]" = torch.ops.prims.mul.default(squeeze_5, primals_3); primals_3 = None broadcast_in_dim_20: "f32[1, 3]" = torch.ops.prims.broadcast_in_dim.default(mul_12, [1, 3], [1]); mul_12 = None broadcast_in_dim_21: "f32[1, 3, 1]" = torch.ops.prims.broadcast_in_dim.default(broadcast_in_dim_20, [1, 3, 1], [0, 1]); broadcast_in_dim_20 = None broadcast_in_dim_22: "f32[1, 3, 1, 1]" = torch.ops.prims.broadcast_in_dim.default(broadcast_in_dim_21, [1, 3, 1, 1], [0, 1, 2]); broadcast_in_dim_21 = None broadcast_in_dim_23: "f32[2, 3, 4, 4]" = torch.ops.prims.broadcast_in_dim.default(broadcast_in_dim_12, [2, 3, 4, 4], [0, 1, 2, 3]); broadcast_in_dim_12 = None sub_2: "f32[2, 3, 4, 4]" = torch.ops.prims.sub.default(convolution, broadcast_in_dim_23); convolution = broadcast_in_dim_23 = None broadcast_in_dim_24: "f32[2, 3, 4, 4]" = torch.ops.prims.broadcast_in_dim.default(broadcast_in_dim_19, [2, 3, 4, 4], [0, 1, 2, 3]); broadcast_in_dim_19 = None mul_13: "f32[2, 3, 4, 4]" = torch.ops.prims.mul.default(sub_2, broadcast_in_dim_24); sub_2 = broadcast_in_dim_24 = None sub_3: "f32[2, 3, 4, 4]" = torch.ops.prims.sub.default(where_1, mul_13); where_1 = mul_13 = None broadcast_in_dim_25: "f32[2, 3, 4, 4]" = torch.ops.prims.broadcast_in_dim.default(broadcast_in_dim_16, [2, 3, 4, 4], [0, 1, 2, 3]); broadcast_in_dim_16 = None sub_4: "f32[2, 3, 4, 4]" = torch.ops.prims.sub.default(sub_3, broadcast_in_dim_25); sub_3 = broadcast_in_dim_25 = None broadcast_in_dim_26: "f32[2, 3, 4, 4]" = torch.ops.prims.broadcast_in_dim.default(broadcast_in_dim_22, [2, 3, 4, 4], [0, 1, 2, 3]); broadcast_in_dim_22 = None mul_14: "f32[2, 3, 4, 4]" = torch.ops.prims.mul.default(sub_4, broadcast_in_dim_26); sub_4 = broadcast_in_dim_26 = None mul_15: "f32[3]" = torch.ops.prims.mul.default(sum_3, squeeze_5); sum_3 = squeeze_5 = None convolution_backward = torch.ops.aten.convolution_backward.default(mul_14, primals_8, primals_1, [3], [1, 1], [1, 1], [1, 1], False, [0, 0], 1, [False, True, True]); mul_14 = primals_8 = primals_1 = None getitem_1: "f32[3, 1, 3, 3]" = convolution_backward[1] getitem_2: "f32[3]" = convolution_backward[2]; convolution_backward = None return pytree.tree_unflatten([ add_2, # InputMutationAOTOutput(mutated_input=BufferAOTInput(target='bn.running_mean')) add_3, # InputMutationAOTOutput(mutated_input=BufferAOTInput(target='bn.running_var')) add, # InputMutationAOTOutput(mutated_input=BufferAOTInput(target='bn.num_batches_tracked')) where, # PlainAOTOutput(idx=0) getitem_1, # GradAOTOutput(grad_of=ParamAOTInput(target='conv.weight')) getitem_2, # GradAOTOutput(grad_of=ParamAOTInput(target='conv.bias')) mul_15, # GradAOTOutput(grad_of=ParamAOTInput(target='bn.weight')) sum_2, # GradAOTOutput(grad_of=ParamAOTInput(target='bn.bias')) None, # None None, # None None, # None None, # None ], self._out_spec) """, # noqa: B950 ) # Compile the result parallel_model_fn = aot_compile_joint_with_descriptors( joint_with_descriptors ) # Test functional correctness expected_output = model(*inputs) all_params_buffers = ( *dict(model.named_parameters()).values(), *dict(model.named_buffers()).values(), ) actual_output = parallel_model_fn(*all_params_buffers, *inputs) torch.testing.assert_close(expected_output, actual_output, rtol=1e-4, atol=1e-4) def test_module_with_kwargs(self): """Test module with keyword arguments""" class ModuleWithKwargs(nn.Module): def __init__(self): super().__init__() self.linear = nn.Linear(3, 2) def forward(self, x, scale=1.0): return self.linear(x) * scale model = ModuleWithKwargs() inputs = (torch.randn(4, 3),) kwargs = {"scale": 2.0} with ExitStack() as stack: # Export joint with descriptors joint_with_descriptors = aot_export_joint_with_descriptors( stack, model, inputs, kwargs, decompositions=decomposition_table ) # Test the exported graph structure graph_code = joint_with_descriptors.graph_module.print_readable( print_output=False, expanded_def=True ) # Expect test on the printed graph self.assertExpectedInline( normalize_gm(graph_code), """\ class inner_f(torch.nn.Module): def forward( self, primals, tangents, ): primals_1: "f32[2, 3]" # ParamAOTInput(target='linear.weight') primals_2: "f32[2]" # ParamAOTInput(target='linear.bias') primals_3: "f32[4, 3]" # PlainAOTInput(idx=0) tangents_1: "f32[4, 2]" # TangentAOTInput(output=PlainAOTOutput(idx=0)) primals_1, primals_2, primals_3, primals_4 , tangents_1, = fx_pytree.tree_flatten_spec([primals, tangents], self._in_spec) transpose: "f32[3, 2]" = torch.ops.prims.transpose.default(primals_1, [1, 0]); primals_1 = None mm: "f32[4, 2]" = torch.ops.aten.mm.default(primals_3, transpose); transpose = None mul: "f32[4, 2]" = torch.ops.prims.mul.default(mm, 1.0); mm = None mul_1: "f32[2]" = torch.ops.prims.mul.default(primals_2, 1.0); primals_2 = None broadcast_in_dim: "f32[4, 2]" = torch.ops.prims.broadcast_in_dim.default(mul_1, [4, 2], [1]); mul_1 = None add: "f32[4, 2]" = torch.ops.prims.add.default(mul, broadcast_in_dim); mul = broadcast_in_dim = None mul_2: "f32[4, 2]" = torch.ops.prims.mul.default(add, 2.0); add = None mul_3: "f32[4, 2]" = torch.ops.prims.mul.default(tangents_1, 2.0); tangents_1 = None transpose_1: "f32[2, 4]" = torch.ops.prims.transpose.default(mul_3, [1, 0]) mm_1: "f32[2, 3]" = torch.ops.aten.mm.default(transpose_1, primals_3); transpose_1 = primals_3 = None transpose_2: "f32[3, 2]" = torch.ops.prims.transpose.default(mm_1, [1, 0]); mm_1 = None sum_1: "f32[2]" = torch.ops.prims.sum.default(mul_3, [0]); mul_3 = None broadcast_in_dim_1: "f32[1, 2]" = torch.ops.prims.broadcast_in_dim.default(sum_1, [1, 2], [1]); sum_1 = None as_strided: "f32[2]" = torch.ops.aten.as_strided.default(broadcast_in_dim_1, [2], [1]); broadcast_in_dim_1 = None transpose_3: "f32[2, 3]" = torch.ops.prims.transpose.default(transpose_2, [1, 0]); transpose_2 = None return pytree.tree_unflatten([ mul_2, # PlainAOTOutput(idx=0) transpose_3, # GradAOTOutput(grad_of=ParamAOTInput(target='linear.weight')) as_strided, # GradAOTOutput(grad_of=ParamAOTInput(target='linear.bias')) None, # None None, # None ], self._out_spec) """, ) # Compile the result parallel_model_fn = aot_compile_joint_with_descriptors( joint_with_descriptors ) # Test functional correctness expected_output = model(*inputs, **kwargs) actual_output = parallel_model_fn( *dict(model.named_parameters()).values(), *inputs, **kwargs ) self.assertEqual(expected_output, actual_output) def test_multiple_outputs_module(self): """Test module with multiple outputs""" class MultiOutputModule(nn.Module): def __init__(self): super().__init__() self.linear1 = nn.Linear(3, 2) self.linear2 = nn.Linear(3, 4) def forward(self, x): out1 = self.linear1(x) out2 = self.linear2(x) return out1, out2 model = MultiOutputModule() inputs = (torch.randn(4, 3),) with ExitStack() as stack: # Export joint with descriptors joint_with_descriptors = aot_export_joint_with_descriptors( stack, model, inputs, decompositions=decomposition_table ) # Test the exported graph structure graph_code = joint_with_descriptors.graph_module.print_readable( print_output=False, expanded_def=True ) # Expect test on the printed graph self.assertExpectedInline( normalize_gm(graph_code), """\ class inner_f(torch.nn.Module): def forward( self, primals, tangents, ): primals_1: "f32[2, 3]" # ParamAOTInput(target='linear1.weight') primals_2: "f32[2]" # ParamAOTInput(target='linear1.bias') primals_3: "f32[4, 3]" # ParamAOTInput(target='linear2.weight') primals_4: "f32[4]" # ParamAOTInput(target='linear2.bias') primals_5: "f32[4, 3]" # PlainAOTInput(idx=0) tangents_1: "f32[4, 2]" # TangentAOTInput(output=PlainAOTOutput(idx=0)) tangents_2: "f32[4, 4]" # TangentAOTInput(output=PlainAOTOutput(idx=1)) primals_1, primals_2, primals_3, primals_4, primals_5, tangents_1, tangents_2, = fx_pytree.tree_flatten_spec([primals, tangents], self._in_spec) transpose: "f32[3, 2]" = torch.ops.prims.transpose.default(primals_1, [1, 0]); primals_1 = None mm: "f32[4, 2]" = torch.ops.aten.mm.default(primals_5, transpose); transpose = None mul: "f32[4, 2]" = torch.ops.prims.mul.default(mm, 1.0); mm = None mul_1: "f32[2]" = torch.ops.prims.mul.default(primals_2, 1.0); primals_2 = None broadcast_in_dim: "f32[4, 2]" = torch.ops.prims.broadcast_in_dim.default(mul_1, [4, 2], [1]); mul_1 = None add: "f32[4, 2]" = torch.ops.prims.add.default(mul, broadcast_in_dim); mul = broadcast_in_dim = None transpose_1: "f32[3, 4]" = torch.ops.prims.transpose.default(primals_3, [1, 0]); primals_3 = None mm_1: "f32[4, 4]" = torch.ops.aten.mm.default(primals_5, transpose_1); transpose_1 = None mul_2: "f32[4, 4]" = torch.ops.prims.mul.default(mm_1, 1.0); mm_1 = None mul_3: "f32[4]" = torch.ops.prims.mul.default(primals_4, 1.0); primals_4 = None broadcast_in_dim_1: "f32[4, 4]" = torch.ops.prims.broadcast_in_dim.default(mul_3, [4, 4], [1]); mul_3 = None add_1: "f32[4, 4]" = torch.ops.prims.add.default(mul_2, broadcast_in_dim_1); mul_2 = broadcast_in_dim_1 = None transpose_2: "f32[4, 4]" = torch.ops.prims.transpose.default(tangents_2, [1, 0]) mm_2: "f32[4, 3]" = torch.ops.aten.mm.default(transpose_2, primals_5); transpose_2 = None transpose_3: "f32[3, 4]" = torch.ops.prims.transpose.default(mm_2, [1, 0]); mm_2 = None sum_1: "f32[4]" = torch.ops.prims.sum.default(tangents_2, [0]); tangents_2 = None broadcast_in_dim_2: "f32[1, 4]" = torch.ops.prims.broadcast_in_dim.default(sum_1, [1, 4], [1]); sum_1 = None as_strided: "f32[4]" = torch.ops.aten.as_strided.default(broadcast_in_dim_2, [4], [1]); broadcast_in_dim_2 = None transpose_4: "f32[4, 3]" = torch.ops.prims.transpose.default(transpose_3, [1, 0]); transpose_3 = None transpose_5: "f32[2, 4]" = torch.ops.prims.transpose.default(tangents_1, [1, 0]) mm_3: "f32[2, 3]" = torch.ops.aten.mm.default(transpose_5, primals_5); transpose_5 = primals_5 = None transpose_6: "f32[3, 2]" = torch.ops.prims.transpose.default(mm_3, [1, 0]); mm_3 = None sum_2: "f32[2]" = torch.ops.prims.sum.default(tangents_1, [0]); tangents_1 = None broadcast_in_dim_3: "f32[1, 2]" = torch.ops.prims.broadcast_in_dim.default(sum_2, [1, 2], [1]); sum_2 = None as_strided_1: "f32[2]" = torch.ops.aten.as_strided.default(broadcast_in_dim_3, [2], [1]); broadcast_in_dim_3 = None transpose_7: "f32[2, 3]" = torch.ops.prims.transpose.default(transpose_6, [1, 0]); transpose_6 = None return pytree.tree_unflatten([ add, # PlainAOTOutput(idx=0) add_1, # PlainAOTOutput(idx=1) transpose_7, # GradAOTOutput(grad_of=ParamAOTInput(target='linear1.weight')) as_strided_1, # GradAOTOutput(grad_of=ParamAOTInput(target='linear1.bias')) transpose_4, # GradAOTOutput(grad_of=ParamAOTInput(target='linear2.weight')) as_strided, # GradAOTOutput(grad_of=ParamAOTInput(target='linear2.bias')) None, # None ], self._out_spec) """, # noqa: B950 ) # Compile the result parallel_model_fn = aot_compile_joint_with_descriptors( joint_with_descriptors ) # Test functional correctness expected_output = model(*inputs) actual_output = parallel_model_fn( *dict(model.named_parameters()).values(), *inputs ) # Check both outputs self.assertEqual(len(expected_output), len(actual_output)) for exp, act in zip(expected_output, actual_output): self.assertEqual(exp, act) def test_in_out_specs(self): """Test that in_spec and out_spec are properly set""" class SimpleModule(nn.Module): def __init__(self): super().__init__() self.linear = nn.Linear(3, 2) def forward(self, x): return self.linear(x) model = SimpleModule() inputs = (torch.randn(4, 3),) with ExitStack() as stack: # Export joint with descriptors joint_with_descriptors = aot_export_joint_with_descriptors( stack, model, inputs, decompositions=decomposition_table ) # Test that specs are available self.assertIsNotNone(joint_with_descriptors.in_spec) self.assertIsNotNone(joint_with_descriptors.out_spec) self.assertIsNotNone(joint_with_descriptors.params_spec) self.assertIsNotNone(joint_with_descriptors.buffers_spec) # Test that they work with pytree operations flat_inputs, _ = pytree.tree_flatten((inputs, {})) self.assertTrue(len(flat_inputs) > 0) # Test parameter and buffer specs contain expected entries self.assertIn("linear.weight", joint_with_descriptors.params_spec) self.assertIn("linear.bias", joint_with_descriptors.params_spec) self.assertEqual( len(joint_with_descriptors.buffers_spec), 0 ) # No buffers in simple linear # Compile the result to ensure everything works together parallel_model_fn = aot_compile_joint_with_descriptors( joint_with_descriptors ) # Test functional correctness expected_output = model(*inputs) actual_output = parallel_model_fn( *dict(model.named_parameters()).values(), *inputs ) self.assertEqual(expected_output, actual_output) def test_fx_utils_simple_linear(self): """Test FX utilities on a simple linear module""" class SimpleLinear(nn.Module): def __init__(self): super().__init__() self.linear = nn.Linear(3, 2) def forward(self, x): return self.linear(x) model = SimpleLinear() inputs = (torch.randn(4, 3),) with ExitStack() as stack: joint_with_descriptors = aot_export_joint_with_descriptors( stack, model, inputs, decompositions=decomposition_table ) graph = joint_with_descriptors.graph_module.graph # Test get_named_param_nodes named_params = get_named_param_nodes(graph) self.assertIn("linear.weight", named_params) self.assertIn("linear.bias", named_params) self.assertEqual(len(named_params), 2) # Test get_param_nodes param_nodes = get_param_nodes(graph) self.assertEqual(len(param_nodes), 2) # Test get_named_buffer_nodes (should be empty for simple linear) named_buffers = get_named_buffer_nodes(graph) self.assertEqual(len(named_buffers), 0) # Test get_buffer_nodes buffer_nodes = get_buffer_nodes(graph) self.assertEqual(len(buffer_nodes), 0) # Test get_all_input_and_grad_nodes input_grad_nodes = get_all_input_and_grad_nodes(graph) self.assertEqual(len(input_grad_nodes), 4) # 2 params + 1 input + 1 tangent # Verify that parameters have gradients param_grads = get_param_and_grad_nodes(graph) self.assertEqual(len(param_grads), 2) for desc, (param_node, grad_node) in param_grads.items(): self.assertIsInstance(desc, ParamAOTInput) self.assertIsNotNone(param_node) self.assertIsNotNone(grad_node) # Should have gradients # Test get_plain_input_and_grad_nodes plain_input_grads = get_plain_input_and_grad_nodes(graph) self.assertEqual(len(plain_input_grads), 1) # 1 plain input for desc, (input_node, grad_node) in plain_input_grads.items(): self.assertIsInstance(desc, PlainAOTInput) self.assertIsNotNone(input_node) self.assertIsNone(grad_node) # Plain inputs don't have gradients # Test get_all_output_and_tangent_nodes output_tangent_nodes = get_all_output_and_tangent_nodes(graph) self.assertEqual(len(output_tangent_nodes), 3) # 1 output + 2 grad outputs # Test get_plain_output_and_tangent_nodes plain_output_tangents = get_plain_output_and_tangent_nodes(graph) self.assertEqual(len(plain_output_tangents), 1) for desc, (output_node, tangent_node) in plain_output_tangents.items(): self.assertIsInstance(desc, PlainAOTOutput) self.assertIsNotNone(output_node) self.assertIsNotNone( tangent_node ) # Should have tangents for backward pass def test_fx_utils_conv_bn_module(self): """Test FX utilities on a conv+batchnorm module with buffers""" class ConvBN(nn.Module): def __init__(self): super().__init__() self.conv = nn.Conv2d(1, 3, 3, padding=1) self.bn = nn.BatchNorm2d(3) def forward(self, x): x = self.conv(x) x = self.bn(x) return torch.relu(x) model = ConvBN() model.train() # Important for batch norm inputs = (torch.randn(2, 1, 4, 4),) with ExitStack() as stack: joint_with_descriptors = aot_export_joint_with_descriptors( stack, model, inputs, decompositions=decomposition_table ) graph = joint_with_descriptors.graph_module.graph # Test get_named_param_nodes named_params = get_named_param_nodes(graph) expected_params = ["conv.weight", "conv.bias", "bn.weight", "bn.bias"] for param_name in expected_params: self.assertIn(param_name, named_params) self.assertEqual(len(named_params), 4) # Test get_named_buffer_nodes named_buffers = get_named_buffer_nodes(graph) expected_buffers = [ "bn.running_mean", "bn.running_var", "bn.num_batches_tracked", ] for buffer_name in expected_buffers: self.assertIn(buffer_name, named_buffers) self.assertEqual(len(named_buffers), 3) # Test get_buffer_nodes buffer_nodes = get_buffer_nodes(graph) self.assertEqual(len(buffer_nodes), 3) # Test that all inputs include params, buffers, and plain inputs input_grad_nodes = get_all_input_and_grad_nodes(graph) self.assertEqual( len(input_grad_nodes), 9 ) # 4 params + 3 buffers + 1 input + 1 tangent # Verify buffer handling buffer_count = 0 for desc, (node, grad_node) in input_grad_nodes.items(): if isinstance(desc, BufferAOTInput): buffer_count += 1 self.assertIsNotNone(node) # Buffers typically don't have gradients unless they're trainable self.assertEqual(buffer_count, 3) def test_fx_utils_multiple_outputs(self): """Test FX utilities on a module with multiple outputs""" class MultiOutputModule(nn.Module): def __init__(self): super().__init__() self.linear1 = nn.Linear(3, 2) self.linear2 = nn.Linear(3, 4) def forward(self, x): out1 = self.linear1(x) out2 = self.linear2(x) return out1, out2 model = MultiOutputModule() inputs = (torch.randn(4, 3),) with ExitStack() as stack: joint_with_descriptors = aot_export_joint_with_descriptors( stack, model, inputs, decompositions=decomposition_table ) graph = joint_with_descriptors.graph_module.graph # Test get_all_output_and_tangent_nodes output_tangent_nodes = get_all_output_and_tangent_nodes(graph) self.assertEqual(len(output_tangent_nodes), 6) # 2 outputs + 4 grad outputs # Test get_plain_output_and_tangent_nodes plain_output_tangents = get_plain_output_and_tangent_nodes(graph) self.assertEqual(len(plain_output_tangents), 2) # Verify each output has a tangent for desc, (output_node, tangent_node) in plain_output_tangents.items(): self.assertIsInstance(desc, PlainAOTOutput) self.assertIsNotNone(output_node) self.assertIsNotNone(tangent_node) # Test parameter handling with multiple outputs param_grads = get_param_and_grad_nodes(graph) self.assertEqual(len(param_grads), 4) # 2 weights + 2 biases # All parameters should have gradients for desc, (param_node, grad_node) in param_grads.items(): self.assertIsInstance(desc, ParamAOTInput) self.assertIsNotNone(param_node) self.assertIsNotNone(grad_node) def test_fx_utils_node_consistency(self): """Test that FX utilities return consistent node references""" class SimpleModule(nn.Module): def __init__(self): super().__init__() self.linear = nn.Linear(3, 2) def forward(self, x): return self.linear(x) model = SimpleModule() inputs = (torch.randn(4, 3),) with ExitStack() as stack: joint_with_descriptors = aot_export_joint_with_descriptors( stack, model, inputs, decompositions=decomposition_table ) graph = joint_with_descriptors.graph_module.graph # Get nodes through different APIs and verify consistency named_params = get_named_param_nodes(graph) param_nodes = get_param_nodes(graph) param_grads = get_param_and_grad_nodes(graph) all_input_grads = get_all_input_and_grad_nodes(graph) # Check that get_param_nodes returns the same nodes as get_named_param_nodes self.assertEqual(len(param_nodes), len(named_params)) for node in param_nodes: self.assertIn(node, named_params.values()) # Check that param_grads contains the same parameter nodes for desc, (param_node, grad_node) in param_grads.items(): self.assertIn(param_node, param_nodes) self.assertEqual(param_node, named_params[desc.target]) # Check that all_input_grads contains the parameter nodes param_count = 0 for desc, (input_node, grad_node) in all_input_grads.items(): if isinstance(desc, ParamAOTInput): param_count += 1 self.assertIn(input_node, param_nodes) self.assertEqual(input_node, named_params[desc.target]) self.assertEqual(param_count, len(param_nodes)) def test_export_and_compile(self): class SimpleModule(nn.Module): def __init__(self): super().__init__() self.linear = nn.Linear(3, 2) def forward(self, x): return self.linear(x) model = SimpleModule() inputs = (torch.randn(4, 3),) with ExitStack() as stack: joint_with_descriptors = aot_export_joint_with_descriptors( stack, model, inputs ) model_fn = aot_compile_joint_with_descriptors(joint_with_descriptors) compiled_fn = torch.compile(fullgraph=True)(model_fn) compiled_fn(*dict(model.named_parameters()).values(), inputs).sum().backward() self.assertIsNotNone(model.linear.weight.grad) def test_preserve_annotate_simple(self): """Test basic linear module with aot_export_joint_with_descriptors""" class SimpleLinear(nn.Module): def __init__(self): super().__init__() self.linear = nn.Linear(3, 2) def forward(self, x): with fx_traceback.annotate({"pp_stage": 0}): y = self.linear(x) return y - 1 inputs = (torch.randn(4, 3),) for with_export in [True, False]: with ExitStack() as stack: model = None with fx_traceback.preserve_node_meta(): if with_export: ep = torch.export.export(SimpleLinear(), inputs) model = ep.module() else: model = SimpleLinear() joint_with_descriptors = aot_export_joint_with_descriptors( stack, model, inputs, decompositions=decomposition_table ) for node in joint_with_descriptors.graph_module.graph.nodes: if ( node.target in ( torch.ops.prims.transpose.default, torch.ops.aten.mm.default, torch.ops.prims.mul.default, torch.ops.prims.broadcast_in_dim.default, torch.ops.prims.add.default, ) # TODO: add annotation to backward graph nodes and node.meta.get("partitioner_tag") != "is_backward" ): self.assertTrue(node.meta["custom"], {"pp_stage": 0}) if node.target == torch.ops.aten.sub.default: self.assertTrue(node.meta.get("custom", {}), {}) if __name__ == "__main__": run_tests()