Export Times

fx_mode

symbolic

<<<

import time
import warnings
import numpy as np
from transformers import LlamaConfig
from transformers.models.llama.modeling_llama import LlamaModel
import onnx
import onnxruntime
import torch
import torch._dynamo
import torch.export
import onnxscript
import torch.onnx
import experimental_experiment
import experimental_experiment.torch_interpreter
import experimental_experiment.torch_interpreter.aten_functions
from experimental_experiment.torch_models.llama_helper import get_llama_model

begin = time.perf_counter()
print("creating model")
model, example_args_collection = get_llama_model(
    input_dims=[(2, 1024)],
    hidden_size=4096,
    num_hidden_layers=2,
    vocab_size=32000,
    intermediate_size=11008,
    max_position_embeddings=2048,
    num_attention_heads=32,
    _attn_implementation="eager",
)

torch._dynamo.reset()
begin = time.perf_counter()
torch._dynamo.export(model, tracing_mode="symbolic")(*example_args_collection[0])
print(f"time to export symbolic --- {time.perf_counter() - begin}")

>>>

    creating model
    time to export symbolic --- 1.4666611019999891

fake

<<<

import time
import warnings
import numpy as np
from transformers import LlamaConfig
from transformers.models.llama.modeling_llama import LlamaModel
import onnx
import onnxruntime
import torch
import torch._dynamo
import torch.export
import onnxscript
import torch.onnx
import experimental_experiment
import experimental_experiment.torch_interpreter
import experimental_experiment.torch_interpreter.aten_functions
from experimental_experiment.torch_models.llama_helper import get_llama_model

begin = time.perf_counter()
print("creating model")
model, example_args_collection = get_llama_model(
    input_dims=[(2, 1024)],
    hidden_size=4096,
    num_hidden_layers=2,
    vocab_size=32000,
    intermediate_size=11008,
    max_position_embeddings=2048,
    num_attention_heads=32,
    _attn_implementation="eager",
)

torch._dynamo.reset()
begin = time.perf_counter()
torch._dynamo.export(model, tracing_mode="fake")(*example_args_collection[0])
print(f"time to export fake --- {time.perf_counter() - begin}")

>>>

    creating model
    time to export fake --- 0.656119611999884

Custom Exporter

With a very simple model:

<<<

import time
from experimental_experiment.checks import print_import_time

print_import_time()

import torch
import experimental_experiment.torch_interpreter


class Neuron(torch.nn.Module):
    def __init__(self, n_dims: int, n_targets: int):
        super(Neuron, self).__init__()
        self.linear = torch.nn.Linear(n_dims, n_targets)

    def forward(self, x):
        return torch.sigmoid(self.linear(x))


model = Neuron(3, 1)
x = torch.rand(5, 3)

begin = time.perf_counter()
onx = experimental_experiment.torch_interpreter.to_onnx(model, (x,))
print(f"time to export 1x --- {time.perf_counter() - begin}")

begin = time.perf_counter()
onx = experimental_experiment.torch_interpreter.to_onnx(model, (x,))
print(f"time to export 2x --- {time.perf_counter() - begin}")

>>>

    time to import onnx --- 0.8925491659999807
    time to import onnx_array_api --- 0.00020604899987120007
    time to import torch --- 2.1970381749999888
    'torch.export' already imported
    time to import torch.export --- 4.470000021683518e-06
    time to import onnxscript --- 0.1432250909999766
    time to import onnxruntime --- 0.04443645499986815
    time to import torch.onnx --- 0.04214493900008165
    time to import torch._dynamo --- 1.2265451249998023
    time to import experimental_experiment.torch_interpreter --- 3.803971805999936
    time to import experimental_experiment.torch_interpreter.aten_functions --- 0.010961083999973198
    time to export 1x --- 0.29894498699991345
    time to export 2x --- 0.016950073999851156

With a bigger model:

<<<

import time
import warnings
import numpy as np
from transformers import LlamaConfig
from transformers.models.llama.modeling_llama import LlamaModel
import onnx
import onnxruntime
import torch
import torch._dynamo
import torch.export
import onnxscript
import torch.onnx
import experimental_experiment
import experimental_experiment.torch_interpreter
import experimental_experiment.torch_interpreter.aten_functions
from experimental_experiment.torch_models.llama_helper import get_llama_model

model, example_args_collection = get_llama_model(
    input_dims=[(2, 1024)],
    hidden_size=4096,
    num_hidden_layers=1,
    vocab_size=32000,
    intermediate_size=11008,
    max_position_embeddings=2048,
    num_attention_heads=32,
    _attn_implementation="eager",
)

begin = time.perf_counter()
onx = experimental_experiment.torch_interpreter.to_onnx(
    model, example_args_collection[0]
)
print(f"time to export 1x --- {time.perf_counter() - begin}")

begin = time.perf_counter()
onx = experimental_experiment.torch_interpreter.to_onnx(
    model, example_args_collection[0]
)
print(f"time to export 2x --- {time.perf_counter() - begin}")

>>>

    time to export 1x --- 4.350439972999993
    time to export 2x --- 2.4030480129999887

Dynamo Exporter

<<<

import time
import warnings

from experimental_experiment.checks import print_import_time

print_import_time()

import torch
import experimental_experiment.torch_interpreter


class Neuron(torch.nn.Module):
    def __init__(self, n_dims: int, n_targets: int):
        super(Neuron, self).__init__()
        self.linear = torch.nn.Linear(n_dims, n_targets)

    def forward(self, x):
        return torch.sigmoid(self.linear(x))


model = Neuron(3, 1)
x = torch.rand(5, 3)

with warnings.catch_warnings():
    warnings.simplefilter("ignore")

    begin = time.perf_counter()
    onx = torch.onnx.export(model, x, dynamo=True)
    print(f"time to export 1x --- {time.perf_counter() - begin}")

    begin = time.perf_counter()
    onx = torch.onnx.export(model, x, dynamo=True)
    print(f"time to export 2x --- {time.perf_counter() - begin}")

>>>

    time to import onnx --- 0.92466512399983
    time to import onnx_array_api --- 0.00011090499992860714
    time to import torch --- 2.1064651730000605
    'torch.export' already imported
    time to import torch.export --- 5.9920000694546616e-06
    time to import onnxscript --- 0.11986815900013426
    time to import onnxruntime --- 0.028219053999919197
    time to import torch.onnx --- 0.0405488050000713
    time to import torch._dynamo --- 1.3904928970000583
    time to import experimental_experiment.torch_interpreter --- 4.2399748549999
    time to import experimental_experiment.torch_interpreter.aten_functions --- 0.007423783999911393
    [torch.onnx] Obtain model graph for `Neuron([...]` with `torch.export.export(..., strict=False)`...
    [torch.onnx] Obtain model graph for `Neuron([...]` with `torch.export.export(..., strict=False)`... ✅
    [torch.onnx] Run decomposition...
    [torch.onnx] Run decomposition... ✅
    [torch.onnx] Translate the graph into ONNX...
    [torch.onnx] Translate the graph into ONNX... ✅
    time to export 1x --- 1.4448936220001087
    [torch.onnx] Obtain model graph for `Neuron([...]` with `torch.export.export(..., strict=False)`...
    [torch.onnx] Obtain model graph for `Neuron([...]` with `torch.export.export(..., strict=False)`... ✅
    [torch.onnx] Run decomposition...
    [torch.onnx] Run decomposition... ✅
    [torch.onnx] Translate the graph into ONNX...
    [torch.onnx] Translate the graph into ONNX... ✅
    time to export 2x --- 0.5141066000001047

With a bigger model:

<<<

import time
import warnings
import numpy as np
from transformers import LlamaConfig
from transformers.models.llama.modeling_llama import LlamaModel
import onnx
import onnxruntime
import torch
import torch._dynamo
import torch.export
import onnxscript
import torch.onnx
import experimental_experiment
import experimental_experiment.torch_interpreter
import experimental_experiment.torch_interpreter.aten_functions
from experimental_experiment.torch_models.llama_helper import get_llama_model

model, example_args_collection = get_llama_model(
    input_dims=[(2, 1024)],
    hidden_size=4096,
    num_hidden_layers=1,
    vocab_size=32000,
    intermediate_size=11008,
    max_position_embeddings=2048,
    num_attention_heads=32,
    _attn_implementation="eager",
)

with warnings.catch_warnings():
    warnings.simplefilter("ignore")

    begin = time.perf_counter()
    onx = torch.onnx.export(model, *example_args_collection[0], dynamo=True)
    print(f"time to export 1x --- {time.perf_counter() - begin}")

    begin = time.perf_counter()
    onx = torch.onnx.export(model, *example_args_collection[0], dynamo=True)
    print(f"time to export 2x --- {time.perf_counter() - begin}")

>>>

    [torch.onnx] Obtain model graph for `LlamaModelWrapper([...]` with `torch.export.export(..., strict=False)`...
    [torch.onnx] Obtain model graph for `LlamaModelWrapper([...]` with `torch.export.export(..., strict=False)`... ❌
    [torch.onnx] Obtain model graph for `LlamaModelWrapper([...]` with `torch.export.export(..., strict=True)`...
    [torch.onnx] Obtain model graph for `LlamaModelWrapper([...]` with `torch.export.export(..., strict=True)`... ❌
    [torch.onnx] Obtain model graph for `LlamaModelWrapper([...]` with `torch.export draft_export`...
    [torch.onnx] Obtain model graph for `LlamaModelWrapper([...]` with `torch.export draft_export`... ❌
    [runpythonerror]
    Traceback (most recent call last):
      File "~/vv/this312/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_capture_strategies.py", line 118, in __call__
        exported_program = self._capture(model, args, kwargs, dynamic_shapes)
                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
      File "~/vv/this312/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_capture_strategies.py", line 202, in _capture
        return torch.export.export(
               ^^^^^^^^^^^^^^^^^^^^
      File "~/vv/this312/lib/python3.12/site-packages/torch/export/__init__.py", line 319, in export
        raise e
      File "~/vv/this312/lib/python3.12/site-packages/torch/export/__init__.py", line 286, in export
        return _export(
               ^^^^^^^^
      File "~/vv/this312/lib/python3.12/site-packages/torch/export/_trace.py", line 1159, in wrapper
        raise e
      File "~/vv/this312/lib/python3.12/site-packages/torch/export/_trace.py", line 1125, in wrapper
        ep = fn(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^
      File "~/vv/this312/lib/python3.12/site-packages/torch/export/exported_program.py", line 123, in wrapper
        return fn(*args, **kwargs)
               ^^^^^^^^^^^^^^^^^^^
      File "~/vv/this312/lib/python3.12/site-packages/torch/export/_trace.py", line 2172, in _export
        ep = _export_for_training(
             ^^^^^^^^^^^^^^^^^^^^^
      File "~/vv/this312/lib/python3.12/site-packages/torch/export/_trace.py", line 1159, in wrapper
        raise e
      File "~/vv/this312/lib/python3.12/site-packages/torch/export/_trace.py", line 1125, in wrapper
        ep = fn(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^
      File "~/vv/this312/lib/python3.12/site-packages/torch/export/exported_program.py", line 123, in wrapper
        return fn(*args, **kwargs)
               ^^^^^^^^^^^^^^^^^^^
      File "~/vv/this312/lib/python3.12/site-packages/torch/export/_trace.py", line 2033, in _export_for_training
        export_artifact = export_func(
                          ^^^^^^^^^^^^
      File "~/vv/this312/lib/python3.12/site-packages/torch/export/_trace.py", line 1933, in _non_strict_export
        ) = make_fake_inputs(
            ^^^^^^^^^^^^^^^^^
      File "~/vv/this312/lib/python3.12/site-packages/torch/_export/non_strict_utils.py", line 284, in make_fake_inputs
        combined_args = _combine_args(nn_module, args, kwargs)
                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
      File "~/vv/this312/lib/python3.12/site-packages/torch/export/dynamic_shapes.py", line 654, in _combine_args
        return signature.bind(*args, **kwargs).arguments
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
      File "/usr/lib/python3.12/inspect.py", line 3277, in bind
        return self._bind(args, kwargs)
               ^^^^^^^^^^^^^^^^^^^^^^^^
      File "/usr/lib/python3.12/inspect.py", line 3190, in _bind
        raise TypeError(msg) from None
    TypeError: missing a required argument: 'attention_mask'
    
    The above exception was the direct cause of the following exception:
    
    Traceback (most recent call last):
      File "<stdin>", line 40, in <module>
      File "~/vv/this312/lib/python3.12/site-packages/torch/onnx/__init__.py", line 367, in export
        return _compat.export_compat(
               ^^^^^^^^^^^^^^^^^^^^^^
      File "~/vv/this312/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_compat.py", line 119, in export_compat
        onnx_program = _core.export(
                       ^^^^^^^^^^^^^
      File "~/vv/this312/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_flags.py", line 20, in wrapper
        return func(*args, **kwargs)
               ^^^^^^^^^^^^^^^^^^^^^
      File "~/vv/this312/lib/python3.12/site-packages/torch/onnx/_internal/exporter/_core.py", line 1332, in export
        raise _errors.TorchExportError(
    torch.onnx._internal.exporter._errors.TorchExportError: Failed to export the model with torch.export. This is step 1/3 of exporting the model to ONNX. Next steps:
    - Modify the model code for `torch.export.export` to succeed. Refer to https://pytorch.org/docs/stable/generated/exportdb/index.html for more information.
    - Debug `torch.export.export` and summit a PR to PyTorch.
    - Create an issue in the PyTorch GitHub repository against the *torch.export* component and attach the full error stack as well as reproduction scripts.
    
    ## Exception summary
    
    <class 'TypeError'>: missing a required argument: 'attention_mask'
    
    (Refer to the full stack trace above for more information.)