Skip to content

Commit 84daf79

Browse files
martinlsmmartinlsm
committed
Arm backend: Add Jupyter notebook example of MXFP
Signed-off-by: Martin Lindström <Martin.Lindstroem@arm.com> Change-Id: Ib8b77c3814a7410afff34afeb839aeb92e2d3bb3
1 parent 4f1b772 commit 84daf79

1 file changed

Lines changed: 247 additions & 0 deletions

File tree

‎examples/arm/mxfp_example.ipynb‎

Lines changed: 247 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,247 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "code",
5+
"execution_count": null,
6+
"id": "927014ab",
7+
"metadata": {},
8+
"outputs": [],
9+
"source": [
10+
"# Copyright 2026 Arm Limited and/or its affiliates.\n",
11+
"#\n",
12+
"# This source code is licensed under the BSD-style license found in the\n",
13+
"# LICENSE file in the root directory of this source tree."
14+
]
15+
},
16+
{
17+
"cell_type": "markdown",
18+
"id": "96dd9def",
19+
"metadata": {},
20+
"source": [
21+
"**NOTICE:** *MXFP is not yet fully supported in the VGF backend. This example instead uses the TOSA reference model until VGF has full support.*"
22+
]
23+
},
24+
{
25+
"cell_type": "markdown",
26+
"id": "080aefe7",
27+
"metadata": {},
28+
"source": [
29+
"# Running MXFP on Arm backend\n",
30+
"\n",
31+
"This guide demonstrates the full flow for quantizing select submodules to MXFP on the Arm backend.\n",
32+
"\n",
33+
"Before you begin:\n",
34+
"1. (In a clean virtual environment with a compatible Python version) Install executorch using `./install_executorch.sh`\n",
35+
"2. Install Arm cross-compilation toolchain and simulators using `./examples/arm/setup.sh --i-agree-to-the-contained-eula`\n",
36+
"3. Export vulkan environment variables and add MLSDK components to PATH and LD_LIBRARY_PATH using `examples/arm/arm-scratch/setup_path.sh`\n",
37+
"\n",
38+
"With all commands executed from the base `executorch` folder."
39+
]
40+
},
41+
{
42+
"cell_type": "markdown",
43+
"id": "12b6bd63",
44+
"metadata": {},
45+
"source": [
46+
"### Define the module and quantize it to MXFP\n",
47+
"\n",
48+
"The following code block shows how to use the `to_mxfp` API to quantize a PyTorch module to an MXFP representation. We start by defining the `torch.nn.Module` we want to quantize, then an `MXFPOpConfig` specifiying how to perform the quantization is created. The config selects which datatype to quantize to, and which exact layers to target. `to_mxfp` then quantizes the module in place."
49+
]
50+
},
51+
{
52+
"cell_type": "code",
53+
"execution_count": null,
54+
"id": "05a233cc",
55+
"metadata": {},
56+
"outputs": [],
57+
"source": [
58+
"import torch\n",
59+
"from torch import nn\n",
60+
"\n",
61+
"from executorch.backends.arm.ao_ext import MXFPOpConfig, to_mxfp\n",
62+
"\n",
63+
"torch.manual_seed(0)\n",
64+
"\n",
65+
"def filter_only_fc1(mod: nn.Module, name: str) -> bool:\n",
66+
" return name == \"fc1\" \n",
67+
"\n",
68+
"\n",
69+
"class Mod(nn.Module):\n",
70+
" def __init__(self):\n",
71+
" super().__init__()\n",
72+
" self.fc1 = nn.Linear(32, 64)\n",
73+
" self.fc2 = nn.Linear(64, 16)\n",
74+
"\n",
75+
" def forward(self, x):\n",
76+
" x = self.fc1(x)\n",
77+
" x = torch.relu(x)\n",
78+
" x = self.fc2(x)\n",
79+
" return x\n",
80+
"\n",
81+
"\n",
82+
"module = Mod().eval()\n",
83+
"print(f\"Initial module:\\n{module}\\n\")\n",
84+
"\n",
85+
"# F8E4M3 is used here. See backends/arm/ao_ext/mxfp.py for all supported datatypes.\n",
86+
"config = MXFPOpConfig(\n",
87+
" weight_dtype=torch.float8_e4m3fn,\n",
88+
")\n",
89+
"# To quantize only module.fc1, add back the commented out filter_fn arg.\n",
90+
"to_mxfp(\n",
91+
" module,\n",
92+
" config,\n",
93+
" # filter_fn=filter_only_fc1,\n",
94+
")\n",
95+
"print(f\"MXFP-quantized module:\\n{module}\\n\")\n",
96+
"\n",
97+
"\n",
98+
"example_inputs = (torch.randn(1, 32),)\n",
99+
"exported_program = torch.export.export(module, example_inputs)\n",
100+
"graph_module = exported_program.module(check_guards=False)\n",
101+
"print(\"Exported module:\")\n",
102+
"_ = graph_module.print_readable()"
103+
]
104+
},
105+
{
106+
"cell_type": "code",
107+
"execution_count": null,
108+
"id": "48cb1113",
109+
"metadata": {},
110+
"outputs": [],
111+
"source": [
112+
"import os\n",
113+
"from executorch.backends.arm.tosa.partitioner import TOSAPartitioner\n",
114+
"from executorch.backends.arm.tosa.compile_spec import TosaCompileSpec\n",
115+
"from executorch.exir import (\n",
116+
" EdgeCompileConfig,\n",
117+
" ExecutorchBackendConfig,\n",
118+
" to_edge_transform_and_lower,\n",
119+
")\n",
120+
"from executorch.extension.export_util.utils import save_pte_program\n",
121+
"\n",
122+
"\n",
123+
"# TODO MLETORCH-2141: MXFP is not fully supported in the VGF toolchain yet.\n",
124+
"# Use the TOSA reference model in the mean time and switch to VgfPartitioner\n",
125+
"# when full support is in place.\n",
126+
"compile_spec = TosaCompileSpec(\"TOSA-1.1+FP+mxfp\")\n",
127+
"partitioner = TOSAPartitioner(compile_spec)\n",
128+
"\n",
129+
"# Lower the exported program to the TOSA backend\n",
130+
"edge_program_manager = to_edge_transform_and_lower(\n",
131+
" exported_program,\n",
132+
" partitioner=[partitioner],\n",
133+
" compile_config=EdgeCompileConfig(\n",
134+
" _check_ir_validity=False,\n",
135+
" ),\n",
136+
")\n",
137+
"\n",
138+
"# Convert edge program to executorch\n",
139+
"executorch_program_manager = edge_program_manager.to_executorch(\n",
140+
" config=ExecutorchBackendConfig(extract_delegate_segments=False)\n",
141+
")\n",
142+
"executorch_program_manager.exported_program().module(check_guards=False).print_readable()\n",
143+
"\n",
144+
"# Save pte file\n",
145+
"cwd_dir = os.getcwd()\n",
146+
"pte_base_name = \"mxfp_example\"\n",
147+
"pte_name = pte_base_name + \".pte\"\n",
148+
"pte_path = os.path.join(cwd_dir, pte_name)\n",
149+
"save_pte_program(executorch_program_manager, pte_name)\n",
150+
"assert os.path.exists(pte_path), \"Build failed; no .pte-file found\""
151+
]
152+
},
153+
{
154+
"cell_type": "code",
155+
"execution_count": null,
156+
"id": "2c22f2c9",
157+
"metadata": {},
158+
"outputs": [],
159+
"source": [
160+
"from executorch.backends.arm.test.runner_utils import TosaReferenceModelDispatch\n",
161+
"\n",
162+
"# TODO MLETORCH-2141: Run on VGF backend instead when it's supported.\n",
163+
"with torch.no_grad():\n",
164+
" lowered_module = executorch_program_manager.exported_program().graph_module\n",
165+
" with TosaReferenceModelDispatch():\n",
166+
" tosa_ref_output = lowered_module(*example_inputs)\n",
167+
"\n",
168+
"print(\"Model output:\")\n",
169+
"print(tosa_ref_output)\n"
170+
]
171+
},
172+
{
173+
"cell_type": "markdown",
174+
"id": "89941ec7",
175+
"metadata": {},
176+
"source": [
177+
"### Porting over already quantized modules\n",
178+
"\n",
179+
"The flow presented in this notebook shows how to quantize a module to MXFP, but what if we have a module that was quantized outside of the Arm backend's `to_mxfp` flow? Then we need to port this module to the representation that is compatible with the Arm backend. Doing this requires manual work, but it is possible; note that `to_mxfp` simply replaces submodules with corresponding MXFP implementations, e.g., `torch.nn.Linear` is replaced with `executorch.backends.arm.ao_ext.ops.MXFPLinearOp`. With this observation we can replicate this procedure manually by reassigning the module's weights. The code snippet below shows how to replace the model `Mod`'s `self.fc1` with some made-up pretrained MXFP4 weights we could have gotten from somewhere else."
180+
]
181+
},
182+
{
183+
"cell_type": "code",
184+
"execution_count": null,
185+
"id": "11b72b86",
186+
"metadata": {},
187+
"outputs": [],
188+
"source": [
189+
"import torch\n",
190+
"from executorch.backends.arm.ao_ext.ops import MXFPLinearOp\n",
191+
"\n",
192+
"block_size = 32\n",
193+
"in_features = 32\n",
194+
"out_features = 64\n",
195+
"\n",
196+
"# Note: float4_e2m1fn_x2 weights are packed pairwise into uint8 arrays\n",
197+
"example_external_w_data = torch.randint(\n",
198+
" 0,\n",
199+
" 256,\n",
200+
" (1, out_features, in_features // 2),\n",
201+
" dtype=torch.uint8,\n",
202+
")\n",
203+
"example_external_w_scale = torch.randint(\n",
204+
" 0,\n",
205+
" 256,\n",
206+
" (1, out_features, in_features // block_size),\n",
207+
" dtype=torch.float8_e8m0fnu,\n",
208+
")\n",
209+
"example_external_bias = torch.randn(out_features, dtype=torch.float32)\n",
210+
"\n",
211+
"module = Mod().eval()\n",
212+
"module.fc1 = MXFPLinearOp(\n",
213+
" example_external_w_data,\n",
214+
" example_external_w_scale,\n",
215+
" example_external_bias,\n",
216+
" torch.float4_e2m1fn_x2,\n",
217+
" block_size,\n",
218+
")\n",
219+
"\n",
220+
"port_output = module(example_inputs[0])\n",
221+
"assert port_output.shape == (1, 16)\n",
222+
"print(module)\n"
223+
]
224+
}
225+
],
226+
"metadata": {
227+
"kernelspec": {
228+
"display_name": "env",
229+
"language": "python",
230+
"name": "python3"
231+
},
232+
"language_info": {
233+
"codemirror_mode": {
234+
"name": "ipython",
235+
"version": 3
236+
},
237+
"file_extension": ".py",
238+
"mimetype": "text/x-python",
239+
"name": "python",
240+
"nbconvert_exporter": "python",
241+
"pygments_lexer": "ipython3",
242+
"version": "3.10.4"
243+
}
244+
},
245+
"nbformat": 4,
246+
"nbformat_minor": 5
247+
}

0 commit comments

Comments
 (0)