Export NanoGPT Example to MPS Delegate

Following the guide here: Getting Started with LLMs via ExecuTorch — ExecuTorch 0.3 documentation shows how to export nanoGPT to the XNNPACK delegate but does not provide other backend examples. The guide here Building and Running ExecuTorch with MPS Backend — ExecuTorch 0.3 documentation is not really what I’m looking for. I’d like a guide on how to support my own models (such as in the nanoGPT example).

Is there any advice on where to find this information? Or better yet, how to do this?

Thank you in advanced for your time.

UPDATE: I’m running into the error:

CalledProcessError: Command '['flatc', '--binary', '-o', '/var/folders/fn/gg7mvyh5507g5d3gqz98r6l80000gn/T/tmpvo6v9l1i', '/var/folders/fn/gg7mvyh5507g5d3gqz98r6l80000gn/T/tmpvo6v9l1i/schema.fbs', '/var/folders/fn/gg7mvyh5507g5d3gqz98r6l80000gn/T/tmpvo6v9l1i/schema.json']' returned non-zero exit status 1.

when trying to export GitHub - karpathy/nanoGPT: The simplest, fastest repository for training/finetuning medium-sized GPTs. with Executorch to the MPS backend.

Here is my code:

# 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))

Does anyone have experience with this?

This sounds like an issues with flatcc installed in your enviornment. Maybe @dbort or @larryliu0820 might be able to answer this more quickly.

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.