GPT-2 preparation and export

You’ll compare a PyTorch eager baseline against ExecuTorch with no XNNPACK backend and ExecuTorch with the XNNPACK delegate.

First, create a clean exported program that ExecuTorch can lower in the next step.

The prepare_gpt2.py script loads the tokenizer and GPT-2 model from Hugging Face, wraps the model, creates one fixed-shape example input, exports the model with torch.export, and writes a small model_spec.json.

model_spec.json is a course-specific manifest used by these scripts. ExecuTorch does not require this file. It keeps the model name, exported program path, .pte artifact names, input shapes, tokenizer details, and example inputs in one place so the preparation, lowering, benchmarking, and demo scripts all use the same settings.

Why wrap the model? Transformer language models usually produce logits (prediction scores) with shape [B, T, V]:

B: batch size
T: sequence length
V: vocabulary size

During training, logits for every token position are useful because loss is computed across the sequence. During generation, earlier tokens are already fixed. Only the logits at the final real token position are needed to choose the next token.

That is why the wrapper returns [B, V] instead of [B, T, V]. It keeps the output focused on what the text generation loop actually uses. The wrapper uses the attention mask to find the final non-padding token, so fixed-shape padded inputs still return logits for the prompt rather than logits for a padding position.

The core of the script looks like this:

    

        
        
tokenizer = AutoTokenizer.from_pretrained(args.model_id)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

model = load_hf_model(AutoModelForCausalLM, args.model_id).to(args.device).eval()
wrapped = LastTokenLogits(model).eval()

encoded = tokenizer(
    args.example_prompt,
    return_tensors="pt",
    padding="max_length",
    truncation=True,
    max_length=args.seq_len,
)

example_inputs = (
    encoded["input_ids"].to(args.device),
    encoded["attention_mask"].to(args.device),
)

with torch.inference_mode():
    exported_program = export(wrapped, example_inputs)

save(exported_program, exported_path)

This step writes gpt2_exported.pt2 plus model_spec.json. Keep both files: the lowering script reads them next.

Run on the development machine.

    

        
        
module-1-transformer-inference-with-xnnpack/scripts/aot/prepare_gpt2.py \
  --out module-1-transformer-inference-with-xnnpack/artifacts/gpt2

The exported program expects two inputs:

input_ids: integer token IDs with shape [1, 128]
attention_mask: a matching integer mask with shape [1, 128]

Here 1 is the batch size. The sequence is padded or truncated to 128 tokens. In the attention mask, 1 marks real input tokens and 0 marks padding.

GPT-2 preparation and export

Nice work