[DRAFT] INitial version of sticky export (pytorch#151047)

Summary: This is to make torchnative demos and benchmarking real models more simple by not requiring ppl to find example inputs first. Test Plan: CI Differential Revision: D72815584 Pull Request resolved: pytorch#151047 Approved by: https://github.com/zhxchen17

Author: tugsbayasgalan

test/export/test_experimental.py

@@ -1,5 +1,6 @@
 # Owner(s): ["oncall: export"]
 # flake8: noqa
+import types
 import unittest
 from typing import Dict, List, Tuple
 
@@ -9,7 +10,7 @@
 from torch._functorch.aot_autograd import aot_export_module
 from torch.export import export, export_for_training
 from torch.export._trace import _convert_ts_to_export_experimental
-from torch.export.experimental import _export_forward_backward
+from torch.export.experimental import _export_forward_backward, _sticky_export
 from torch.export.graph_signature import OutputKind
 from torch.testing import FileCheck
 
@@ -333,6 +334,111 @@ def forward(self, x, label):
  OutputKind.LOSS_OUTPUT,
  )
 
+ def test_sticky_export(self):
+ class Model(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.linear = torch.nn.Linear(4, 4)
+
+ def forward(self, x):
+ return self.linear(x)
+
+ class Pipeline:
+ def __init__(self, model):
+ self.model = model
+
+ def generate(self, *args, **kwargs):
+ return self.model(*args, **kwargs)
+
+ inp = torch.randn(4, 4)
+
+ p = Pipeline(Model())
+ orig_forward = p.model.forward
+ p.model.forward = _sticky_export(p.model.forward)
+ res = p.generate(inp)
+
+ p.model.forward = orig_forward
+ res2 = p.generate(inp)
+ self.assertTrue(torch.allclose(res, res2))
+
+ def test_sticky_export_dynamic(self):
+ class Model(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.linear = torch.nn.Linear(4, 4)
+
+ def forward(self, x):
+ if x.shape[0] < 5:
+ return self.linear(x)
+ return x.sin()
+
+ class Pipeline:
+ def __init__(self, model):
+ self.model = model
+
+ def generate(self, *args, **kwargs):
+ return self.model(*args, **kwargs)
+
+ inp = torch.randn(4, 4)
+
+ def callback(*args, **kwargs):
+ # I think it is bit weird to use the forward arg name here, so
+ # lets just use ShapeCollections
+
+ flat_args, _ = torch.utils._pytree.tree_flatten((args, kwargs))
+ collections = torch.export.ShapesCollection()
+ for arg in flat_args:
+ if isinstance(arg, torch.Tensor):
+ collections[arg] = {
+ i: torch.export.Dim.AUTO for i in range(len(arg.shape))
+ }
+ return collections
+
+ p = Pipeline(Model())
+ p.model.forward = _sticky_export(
+ p.model.forward, dynamic_shapes_callback=callback
+ )
+ _ = p.generate(inp)
+ self.assertExpectedInline(
+ str(p.model.forward._exported_artifact.code).strip(),
+ """\
+def forward(self, x):
+ x, = fx_pytree.tree_flatten_spec(([x], {}), self._in_spec)
+ linear_weight = self.linear.weight
+ linear_bias = self.linear.bias
+ linear = torch.ops.aten.linear.default(x, linear_weight, linear_bias); x = linear_weight = linear_bias = None
+ return pytree.tree_unflatten((linear,), self._out_spec)""",
+ )
+
+ def test_sticky_export_nested_inp(self):
+ class Model(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.linear = torch.nn.Linear(4, 4)
+
+ def forward(self, *, inputs):
+ return self.linear(inputs[0]) + self.linear(inputs[1])
+
+ class Pipeline:
+ def __init__(self, model):
+ self.model = model
+
+ def generate(self, *, input_tensor, input_tensor2):
+ inputs = [input_tensor, input_tensor2]
+ return self.model(inputs=inputs)
+
+ inp = torch.randn(4, 4)
+ inp2 = torch.randn(4, 4)
+
+ p = Pipeline(Model())
+ orig_forward = p.model.forward
+ p.model.forward = _sticky_export(p.model.forward)
+ res = p.generate(input_tensor=inp, input_tensor2=inp2)
+
+ p.model.forward = orig_forward
+ res2 = p.generate(input_tensor=inp, input_tensor2=inp2)
+ self.assertTrue(torch.allclose(res, res2))
+
 
 if __name__ == "__main__":
  run_tests()

torch/export/experimental/__init__.py

@@ -1,4 +1,6 @@
 import copy
+import functools
+import types
 import typing
 
 import torch
@@ -67,3 +69,39 @@ def _export_forward_backward(
  _remove_detach_pass(gm, new_graph_signature)
 
  return ep._update(gm, new_graph_signature)
+
+
+@typing.no_type_check
+def _sticky_export(forward_func, dynamic_shapes_callback=None):
+ """
+ Lazily export the model on first forward call.
+ Usage:
+ model.forward = _sticky_export(model.forward, dynamic_shapes_callback=callback)
+ """
+ model = forward_func.__self__
+ original_forward = forward_func.__func__
+
+ @functools.wraps(forward_func)
+ def wrapper(*args, **kwargs):
+ # Unpatch forward to avoid recursion during export
+ model.forward = types.MethodType(original_forward, model)
+
+ dynamic_shapes_spec = None
+ if dynamic_shapes_callback:
+ dynamic_shapes_spec = dynamic_shapes_callback(*args, **kwargs)
+
+ try:
+ exported = torch.export.export(
+ model,
+ args,
+ kwargs,
+ dynamic_shapes=dynamic_shapes_spec,
+ ).module()
+ wrapper._exported_artifact = exported
+ finally:
+ # Restore the wrapper after export
+ model.forward = wrapper
+
+ return exported(*args, **kwargs)
+
+ return wrapper