Convert a random-weight model with create_model

This example shows the core workflow of mbext on a tiny, randomly initialised model. Nothing is downloaded from the Hugging Face hub: a small LlamaConfig is built with a single hidden layer and random weights, converted to ONNX with modelbuilder.builder.create_model(), and the ONNX logits are compared to the PyTorch reference.

This is exactly how the fast test suite validates every architecture, which is why the tests run offline and in seconds (see the Design page).

import os
import tempfile

import numpy as np

Build a tiny random-weight model

A minimal LlamaForCausalLM with small dimensions and a single hidden layer. The weights are randomly initialised (seeded for reproducibility).

import torch
from tokenizers import Tokenizer
from tokenizers.models import WordLevel
from transformers import AutoModelForCausalLM, LlamaConfig, PreTrainedTokenizerFast

num_hidden_layers = 1
config = LlamaConfig(
    architectures=["LlamaForCausalLM"],
    hidden_size=512,
    intermediate_size=1376,
    max_position_embeddings=2048,
    num_attention_heads=8,
    num_key_value_heads=4,
    num_hidden_layers=num_hidden_layers,
    vocab_size=1024,
    bos_token_id=1,
    eos_token_id=2,
)
head_size = config.hidden_size // config.num_attention_heads

torch.manual_seed(42)
model = AutoModelForCausalLM.from_config(config)
model.eval()

LlamaForCausalLM(
  (model): LlamaModel(
    (embed_tokens): Embedding(1024, 512)
    (layers): ModuleList(
      (0): LlamaDecoderLayer(
        (self_attn): LlamaAttention(
          (q_proj): Linear(in_features=512, out_features=512, bias=False)
          (k_proj): Linear(in_features=512, out_features=256, bias=False)
          (v_proj): Linear(in_features=512, out_features=256, bias=False)
          (o_proj): Linear(in_features=512, out_features=512, bias=False)
        )
        (mlp): LlamaMLP(
          (gate_proj): Linear(in_features=512, out_features=1376, bias=False)
          (up_proj): Linear(in_features=512, out_features=1376, bias=False)
          (down_proj): Linear(in_features=1376, out_features=512, bias=False)
          (act_fn): SiLUActivation()
        )
        (input_layernorm): LlamaRMSNorm((512,), eps=1e-06)
        (post_attention_layernorm): LlamaRMSNorm((512,), eps=1e-06)
      )
    )
    (norm): LlamaRMSNorm((512,), eps=1e-06)
    (rotary_emb): LlamaRotaryEmbedding()
  )
  (lm_head): Linear(in_features=512, out_features=1024, bias=False)
)

A minimal word-level tokenizer is enough for create_model to write the processing files next to the ONNX model.

tokenizer = PreTrainedTokenizerFast(
    tokenizer_object=Tokenizer(WordLevel(vocab={"<unk>": 0, "<s>": 1, "</s>": 2}, unk_token="<unk>")),
    bos_token="<s>",
    eos_token="</s>",
    unk_token="<unk>",
)

Convert to ONNX with create_model

The PyTorch model and tokenizer are saved to a folder, then converted to ONNX for the CPU execution provider in fp32 precision.

from modelbuilder.builder import create_model  # noqa: E402

work_dir = tempfile.mkdtemp(prefix="mbext_example_")
model_dir = os.path.join(work_dir, "model")
output_dir = os.path.join(work_dir, "onnx")
cache_dir = os.path.join(work_dir, "cache")

model.save_pretrained(model_dir)
tokenizer.save_pretrained(model_dir)

create_model(
    model_name="tiny-llama",
    input_path=model_dir,
    output_dir=output_dir,
    precision="fp32",
    execution_provider="cpu",
    cache_dir=cache_dir,
    num_hidden_layers=num_hidden_layers,
)

onnx_path = os.path.join(output_dir, "model.onnx")
print("ONNX model written to:", onnx_path)

Writing model shards:   0%|          | 0/1 [00:00<?, ?it/s]
Writing model shards: 100%|██████████| 1/1 [00:00<00:00, 221.18it/s]
GroupQueryAttention (GQA) is used in this model.

Loading weights:   0%|          | 0/12 [00:00<?, ?it/s]
Loading weights: 100%|██████████| 12/12 [00:00<00:00, 7076.01it/s]
Reading embedding layer
Reading layer 0
Reading final norm
Reading LM head
Saving ONNX model in /tmp/mbext_example_vryhgkx5/onnx

0it [00:00, ?it/s]
Saving model.layers.0.input_layernorm.weight (f32, [512]):   8%|▊         | 1/12 [00:00<00:00, 3622.02it/s]
Saving model.layers.0.post_attention_layernorm.weight (f32, [512]):  17%|█▋        | 2/12 [00:00<00:00, 3788.89it/s]
Saving model.layers.1.final_norm_layernorm.weight (f32, [512]):  25%|██▌       | 3/12 [00:00<00:00, 4281.36it/s]
Saving cos_cache (f32, [2048,32]):  33%|███▎      | 4/12 [00:00<00:00, 4672.02it/s]
Saving sin_cache (f32, [2048,32]):  42%|████▏     | 5/12 [00:00<00:00, 4497.43it/s]
Saving model.layers.0.attn.o_proj.MatMul.weight (f32, [512,512]):  50%|█████     | 6/12 [00:00<00:00, 4547.49it/s]
Saving model.embed_tokens.weight (f32, [1024,512]):  58%|█████▊    | 7/12 [00:00<00:00, 3407.63it/s]
Saving model.layers.0.attn.qkv_proj.MatMul.weight (f32, [512,1024]):  67%|██████▋   | 8/12 [00:00<00:00, 2215.55it/s]
Saving lm_head.MatMul.weight (f32, [512,1024]):  75%|███████▌  | 9/12 [00:00<00:00, 1602.17it/s]
Saving model.layers.0.mlp.gate_proj.MatMul.weight (f32, [512,1376]):  83%|████████▎ | 10/12 [00:00<00:00, 1436.11it/s]
Saving model.layers.0.mlp.up_proj.MatMul.weight (f32, [512,1376]):  92%|█████████▏| 11/12 [00:00<00:00, 1221.34it/s]
Saving model.layers.0.mlp.down_proj.MatMul.weight (f32, [1376,512]): 100%|██████████| 12/12 [00:00<00:00, 1143.54it/s]
Saving model.layers.0.mlp.down_proj.MatMul.weight (f32, [1376,512]): 100%|██████████| 12/12 [00:00<00:00, 894.12it/s]
Saving GenAI config in /tmp/mbext_example_vryhgkx5/onnx
Saving processing files in /tmp/mbext_example_vryhgkx5/onnx for GenAI
Saving VS Code settings in /tmp/mbext_example_vryhgkx5/onnx/.vscode/settings.json
ONNX model written to: /tmp/mbext_example_vryhgkx5/onnx/model.onnx

Compare the ONNX and PyTorch logits

The exported graph expects input_ids, attention_mask, position_ids and an (empty) KV cache. We run a single prefill step through OnnxRuntime and compare the logits with the PyTorch model.

import onnxruntime  # noqa: E402

sess = onnxruntime.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
onnx_input_names = {i.name for i in sess.get_inputs()}

batch_size, seq_len = 1, 5
torch.manual_seed(0)
input_ids = torch.randint(0, config.vocab_size, (batch_size, seq_len))

feed = {
    "input_ids": input_ids.numpy().astype(np.int64),
    "attention_mask": np.ones((batch_size, seq_len), dtype=np.int64),
    "position_ids": np.arange(seq_len, dtype=np.int64).reshape(batch_size, seq_len),
}
for i in range(num_hidden_layers):
    feed[f"past_key_values.{i}.key"] = np.zeros((batch_size, config.num_key_value_heads, 0, head_size), dtype=np.float32)
    feed[f"past_key_values.{i}.value"] = np.zeros((batch_size, config.num_key_value_heads, 0, head_size), dtype=np.float32)
feed = {k: v for k, v in feed.items() if k in onnx_input_names}

onnx_logits = sess.run(None, feed)[0]

with torch.no_grad():
    torch_logits = model(input_ids).logits.numpy()

max_abs_diff = np.abs(onnx_logits - torch_logits).max()
print("logits shape:", onnx_logits.shape)
print("max absolute discrepancy (ONNX vs PyTorch):", float(max_abs_diff))

logits shape: (1, 5, 1024)
max absolute discrepancy (ONNX vs PyTorch): 1.3113021850585938e-06

The discrepancy is tiny, confirming the ONNX model matches the PyTorch reference. This is the check that runs for every supported architecture in the fast test suite.