[PT2 Inference] Prototype of Inference Runtime (pytorch#108482)

Summary:
This diff demonstrates a simplified E2E workflow for PT2 Inference stack:
1. Model author with `torch.export()`
2. Model processing with `aot_inductor.compile()`
3. Model served with a new Inference Runtime API, named `ModelRunner`

`torch.export()` and `aot_inductor.compile()` produces a zip file using `PyTorchStreamWriter`.
Runtime reads the zip file with `PyTorchStreamReader`.
The zip file contains
 {F1080328179}
More discussion on packaging can be found in https://docs.google.com/document/d/1C-4DP5yu7ZhX1aB1p9JcVZ5TultDKObM10AqEtmZ-nU/edit?usp=sharing

Runtime can now switch between two Execution modes:
1. Graph Interpreter mode, implemented based on Sigmoid's Executor
2. AOTInductor mode, implemented based on FBAOTInductorModel

Test Plan:
buck2 run  mode/dev-nosan mode/inplace -c fbcode.enable_gpu_sections=True //sigmoid/inference/test:e2e_test

Export and Lower with AOTInductor
buck2 run mode/dev-sand mode/inplace -c fbcode.enable_gpu_sections=True sigmoid/inference:export_package

Run with GraphInterpreter and AOTInducotr
buck2 run mode/dev-nosan //sigmoid/inference:main

Reviewed By: suo

Differential Revision: D47781098

Pull Request resolved: pytorch#108482
Approved by: https://github.com/zhxchen17

docs/source/export.rst

@@ -17,6 +17,10 @@ torch.export
 .. autoclass:: ExportedProgram
 
     .. automethod:: module
+    .. automethod:: buffers
+    .. automethod:: named_buffers
+    .. automethod:: parameters
+    .. automethod:: named_parameters
 
 .. autoclass:: ExportBackwardSignature
 .. autoclass:: ExportGraphSignature

torch/export/__init__.py

@@ -4,7 +4,7 @@
 import pathlib
 import typing
 from enum import auto, Enum
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union
 
 import sympy
 
@@ -286,6 +286,39 @@ def graph_signature(self):
     def state_dict(self):
         return self._state_dict
 
+    @compatibility(is_backward_compatible=False)
+    def parameters(self) -> Iterator[torch.nn.Parameter]:
+        """
+        Returns an iterator over original module's parameters.
+        """
+        for _, param in self.named_parameters():
+            yield param
+
+    @compatibility(is_backward_compatible=False)
+    def named_parameters(self) -> Iterator[Tuple[str, torch.nn.Parameter]]:
+        """
+        Returns an iterator over original module parameters, yielding
+        both the name of the parameter as well as the parameter itself.
+        """
+        for param_name in self.graph_signature.parameters:
+            yield param_name, self.state_dict[param_name]
+
+    @compatibility(is_backward_compatible=False)
+    def buffers(self) -> Iterator[torch.Tensor]:
+        """Returns an iterator over original module buffers."""
+
+        for _, buf in self.named_buffers():
+            yield buf
+
+    @compatibility(is_backward_compatible=False)
+    def named_buffers(self) -> Iterator[Tuple[str, torch.Tensor]]:
+        """
+        Returns an iterator over original module buffers, yielding
+        both the name of the buffer as well as the buffer itself.
+        """
+        for buffer_name in self.graph_signature.buffers:
+            yield buffer_name, self.state_dict[buffer_name]
+
     @property
     @compatibility(is_backward_compatible=False)
     def call_spec(self):