# export_nanogpt.py
import torch
from executorch.exir import EdgeCompileConfig, to_edge
from torch.nn.attention import sdpa_kernel, SDPBackend
from torch._export import capture_pre_autograd_graph
from torch.export import export
import copy
from executorch.backends.apple.mps.partition.mps_partitioner import MPSPartitioner
from executorch.exir.backend.backend_details import CompileSpec
from executorch.exir.backend.backend_api import to_backend
from model import GPT
# Load the model.
model = GPT.from_pretrained('gpt2')
# Create example inputs. This is used in the export process to provide
# hints on the expected shape of the model input.
example_inputs = (torch.randint(0, 100, (1, model.config.block_size), dtype=torch.long), )
# Trace the model, converting it to a portable intermediate representation.
# The torch.no_grad() call tells PyTorch to exclude training-specific logic.
with torch.nn.attention.sdpa_kernel([SDPBackend.MATH]), torch.no_grad():
m = capture_pre_autograd_graph(model, example_inputs)
traced_model = export(m, example_inputs)
# Convert the model into a runnable ExecuTorch program.
edge_config = EdgeCompileConfig(_check_ir_validity=False)
edge_manager = to_edge(traced_model, compile_config=edge_config)
edge_copy = copy.deepcopy(edge_manager)
compiler_specs = [CompileSpec("use_fp16", bytes([True]))]
et_delegate = edge_copy.to_backend(MPSPartitioner(compiler_specs))
Hi @gkteco, if you are still running into this issue, can you please share the content of /var/folders/fn/gg7mvyh5507g5d3gqz98r6l80000gn/T/tmpvo6v9l1i/schema.json from the error message? Also it would be good to share the ExecuTorch commit you are on.
By default, the exir code will delete the temp json file, but it should keep a copy if you export ET_EXIR_SAVE_FLATC_INPUTS_ON_FAILURE=1 before exporting.
Also, please post more logs if you have them. There should be log messages before the CalledProcessError with more details about the problem.