Repository: microsoft/onnxscript

Issue:
TorchScript Graph: Support model > 2GB
Dependent on `torch.Graph._export_onnx`, large model export requires specifying a target path to store the model and weights proto. Right now this path is fixed as "".

https://github.com/microsoft/onnx-script/blob/4986a4be9f8ff1b0e2b7bdde738836574f90e900/onnxscript/function_libs/torch_aten/graph_building.py#L475


Discussion:
2 models (at least, some larger ones are failing for other reasons) [from bench](https://github.com/microsoft/onnx-converters-private/issues/166) are affected. New IR and updates for exporter serialization should take care of this.

2 like: The serialized model is larger than the 2GiB limit imposed by the protobuf library. Therefore the output file must be a file path, so that the ONNX external data can be written to the same directory. Please specify the output file name.
```
        AlbertForMaskedLM
        AlbertForQuestionAnswering
```
I think we do support external data with torchscript right? Or was I hallucinating 
we do but we need a str path provided. A tedious restriction that we want to eliminate.
Removing low priority now that it is the only blocker left for `torch.onnx.dynamo_export` w/ llama.
Will we be able to provide an option in `to_model_proto` to leverage what we already have in torch for this?
I believe we have many options, need to pick one that delivers most natural export api and UX.
Would you like to share some of the options?
I'm thinking of model / weight separation by default. 
Can someone please assign this to @wschin? I have a version for demo.
One thing I think we can do is to save the model out to disk with torchscript, then load it back and pass the model proto onto the exporter for it to decide how to serialize. When we move away from using torchscript this will not be a problem anymore 

Resolving patch:
diff --git a/onnxscript/function_libs/torch_lib/graph_building.py b/onnxscript/function_libs/torch_lib/graph_building.py
--- a/onnxscript/function_libs/torch_lib/graph_building.py
+++ b/onnxscript/function_libs/torch_lib/graph_building.py
@@ -3,6 +3,7 @@
 
 import logging
 import os
+import tempfile
 import typing
 import warnings
 from typing import Any, Dict, Final, List, Mapping, Optional, Sequence, Tuple, Union
@@ -63,6 +64,9 @@
     None,
 ]
 
+# Be sure to leave ample room for the rest of the proto fields.
+_LARGE_MODEL_SIZE_THRESHOLD = int(2**30 * 1.8)  # 1.8GB
+
 # TODO(justinchuby): Build a context manager to handle source information.
 
 
@@ -342,6 +346,18 @@ def _create_op_call_in_torch_graph(
     return node_ouputs
 
 
+def _tensor_rawdata_size(tensor: torch.Tensor) -> int:
+    """Estimate the size of a tensor in bytes.
+
+    Args:
+        tensor: The tensor to estimate the size of.
+
+    Returns:
+        The estimated size of the tensor in bytes.
+    """
+    return tensor.numel() * tensor.element_size()
+
+
 class TorchScriptGraph:
     _LOCAL_FUNCTION_DOMAIN_NAME: Final[str] = "torch_export"
     """The domain name for local functions."""
@@ -683,15 +699,15 @@ def to_model_proto(
             # TODO(BowenBao): All local function domain versions are hardcoded as 1.
             unique_custom_domains[function_proto.domain] = 1
 
-        (
-            proto,
-            _,
-            _,
-            _,
-        ) = self._torch_graph._export_onnx(  # type: ignore[attr-defined] # pylint: disable=protected-access
+        initializers_size = sum(
+            _tensor_rawdata_size(tensor) for tensor in self.initializers.values()
+        )
+
+        large_model = initializers_size > _LARGE_MODEL_SIZE_THRESHOLD
+
+        export_kwargs: dict[str, Any] = dict(
             initializers=self.initializers if include_initializers else {},
             onnx_opset_version=opset_version,
-            # TODO(justinchuby): Figure out how to get the dynamic axes from the inputs
             dynamic_axes={},
             defer_weight_export=False,
             operator_export_type=torch.onnx.OperatorExportTypes.ONNX,
@@ -699,15 +715,42 @@ def to_model_proto(
             keep_initializers_as_inputs=False,
             custom_opsets={},
             add_node_names=True,
-            # TODO(#493): Passing in this instead of reading from env.
-            # User must put the exported model file in the same folder to launch ORT.
-            onnx_file_path=os.path.join(
-                os.getenv("EXTERNAL_ONNX_INITIALIZER_FOLDER", ""), "dummy_model_path.onnx"
-            ),
             node_attr_to_name={},
         )
 
-        onnx_model = onnx.load_from_string(proto)
+        # We decided to cache the model to disk when the model is large.
+        # Alternatively, we could build the ONNX `TensorProto`s in memory
+        # and append them to the model proto.
+        # We did not do it because it is harder to get right (vs. PyTorch's battle-tested
+        # implementation) and creating the `TensorProto`s naively (by converting to numpy)
+        # is slow.
+        cache_model_to_disk = include_initializers and large_model
+
+        if cache_model_to_disk:
+            with tempfile.TemporaryDirectory() as temp_dir:
+                onnx_file_path = os.path.join(temp_dir, "exported_model.onnx")
+                export_kwargs["onnx_file_path"] = onnx_file_path
+                (
+                    proto,
+                    _,
+                    _,
+                    _,
+                ) = self._torch_graph._export_onnx(  # type: ignore[attr-defined] # pylint: disable=protected-access
+                    **export_kwargs
+                )
+                onnx_model = onnx.load_from_string(proto)
+                onnx.load_external_data_for_model(onnx_model, temp_dir)
+        else:
+            (
+                proto,
+                _,
+                _,
+                _,
+            ) = self._torch_graph._export_onnx(  # type: ignore[attr-defined] # pylint: disable=protected-access
+                **export_kwargs
+            )
+            onnx_model = onnx.load_from_string(proto)
+
         onnx_model.functions.extend(function_proto_dict.values())
 
         # `_export_onnx` only exports opset_imports that is visible to it. It does not
@@ -725,10 +768,14 @@ def to_model_proto(
         )
 
         try:
-            onnx_model = onnx.shape_inference.infer_shapes(
-                onnx_model, check_type=True, strict_mode=False, data_prop=True
-            )
-            onnx.checker.check_model(onnx_model, full_check=True)
+            if not cache_model_to_disk:
+                # Only check the model if it is in memory.
+                # Otherwise the checker and shape_inference will fail because
+                # we cannot serialize the model.
+                onnx_model = onnx.shape_inference.infer_shapes(
+                    onnx_model, check_type=True, strict_mode=False, data_prop=True
+                )
+                onnx.checker.check_model(onnx_model, full_check=True)
         except (onnx.checker.ValidationError, onnx.shape_inference.InferenceError) as e:
             warnings.warn(f"ONNX model is invalid: {e}", stacklevel=1)
             logging.debug(


Test patch:
diff --git a/onnxscript/function_libs/torch_lib/graph_building_test.py b/onnxscript/function_libs/torch_lib/graph_building_test.py
--- a/onnxscript/function_libs/torch_lib/graph_building_test.py
+++ b/onnxscript/function_libs/torch_lib/graph_building_test.py
@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import os
-import tempfile
 import unittest
 
 import torch
@@ -11,7 +10,7 @@
 import onnxscript
 import onnxscript.testing
 from onnxscript import FLOAT, evaluator
-from onnxscript import opset17 as op
+from onnxscript import opset18 as op
 from onnxscript._internal import version_utils
 from onnxscript.function_libs.torch_lib import graph_building, ops
 
@@ -165,14 +164,9 @@ def forward(self, x):
         model = MLP(input_size, hidden_size, output_size)
         x = torch.randn(batch_size, input_size)
 
-        with tempfile.TemporaryDirectory() as temp_dir:
-            os.environ["EXTERNAL_ONNX_INITIALIZER_FOLDER"] = temp_dir
-            torch.onnx.dynamo_export(
-                model,
-                x,
-            )
-            # 3 initializers are saved to files as external data.
-            self.assertEqual(len(os.listdir(temp_dir)), 3)
+        model_proto = torch.onnx.dynamo_export(model, x).model_proto
+        # Assert model is larger than 2GB (~=3GB)
+        self.assertGreater(model_proto.ByteSize(), 2**31)
 
 
 if __name__ == "__main__":