Skip to content

Commit 474f5b1

Browse files
justinchubyCopilot
andauthored
Add ONNX export and quantization skill (#238)
Document the Gemma4 ONNX export and INT4 quantization workflow as a reusable skill. ## What this adds New skill at `.agents/skills/onnx-export-quantization/SKILL.md` covering: - **mobius CLI export** — `mobius build` flags, EP variants (default/cuda/onnx-standard), multi-model outputs - **Olive quantization** — Q4_K_M (k-quant) and NF4 (4-bit NormalFloat) with code examples - **GPU acceleration** — cupy for 19-51x kquant speedup - **HuggingFace upload** — standard directory layout (dtype × EP matrix), upload verification - **Common issues** — MoE weight mapping, hybrid attention v_proj, BF16 type mismatches, incomplete uploads - **Testing** — L4 golden data generation, L5 end-to-end smoke test, expected tolerances This captures the workflow used for the Gemma4 ONNX exports so future model exports follow the same pattern. --------- Signed-off-by: Justin Chu <justinchu@microsoft.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
1 parent cefc55e commit 474f5b1

1 file changed

Lines changed: 396 additions & 0 deletions

File tree

  • .agents/skills/onnx-export-quantization
Lines changed: 396 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,396 @@
1+
---
2+
name: onnx-export-quantization
3+
description: >
4+
Use this skill when exporting ONNX models with mobius and quantizing
5+
them with Olive for deployment. Covers the mobius CLI, EP options,
6+
INT4 quantization (Q4_K_M and NF4), HuggingFace upload structure,
7+
GPU-accelerated quantization, common issues, and testing quantized
8+
models.
9+
---
10+
11+
# Skill: ONNX Export and Quantization
12+
13+
## When to use
14+
15+
Use this skill when:
16+
- Exporting a model from HuggingFace to ONNX format using `mobius build`
17+
- Quantizing an ONNX model to INT4 (Q4_K_M or NF4) with Olive
18+
- Uploading ONNX models to HuggingFace Hub in the standard directory layout
19+
- Debugging export or quantization failures
20+
- Choosing between execution provider (EP) variants
21+
22+
## Exporting models with `mobius build`
23+
24+
### Basic command
25+
26+
```bash
27+
mobius build \
28+
--model <hf-model-id> \
29+
--dtype <f16|bf16> \
30+
--ep <default|cuda|onnx-standard> \
31+
--runtime ort-genai \
32+
--external-data safetensors \
33+
--max-shard-size 5GB \
34+
<output-directory>/
35+
```
36+
37+
### Flag reference
38+
39+
| Flag | Description |
40+
|------|-------------|
41+
| `--model <id>` | HuggingFace model ID (e.g. `google/gemma-4-27b-it`) |
42+
| `--dtype <f16\|bf16>` | Model precision — `f16` (float16) or `bf16` (bfloat16) |
43+
| `--optimize [RULES]` | Apply mobius rewrite rules after building (e.g. `group_query_attention`, `packed_attention`, `skip_norm`). Use without value for all rules, or specify comma-separated names. Not needed for basic exports. |
44+
| `--ep <variant>` | Execution provider variant (see below) |
45+
| `--runtime ort-genai` | Generate `genai_config.json` and copy tokenizer files for ORT GenAI runtime |
46+
| `--external-data safetensors` | Store weights externally in safetensors format |
47+
| `--max-shard-size 5GB` | Split external data into shards ≤ 5GB |
48+
49+
### Execution provider (EP) variants
50+
51+
Build separate ONNX models per EP because each applies different graph
52+
rewrites and fused ops:
53+
54+
| EP | Flag | When to use |
55+
|----|------|-------------|
56+
| `default` | `--ep default` | Portable ONNX — no vendor-specific fusions. Compatible with all execution providers and runtimes. This is the default if `--ep` is omitted. |
57+
| `cuda` | `--ep cuda` | NVIDIA GPU inference. Emits `com.microsoft` fused ops (GroupQueryAttention, MoE, etc.) for maximum CUDA performance. |
58+
| `onnx-standard` | `--ep onnx-standard` | Strict ONNX-only — inlines all custom-domain functions into standard ONNX ops. Use when targeting runtimes that don't support `com.microsoft` ops. |
59+
60+
Other EPs are available (`cpu`, `dml`, `webgpu`, `trt-rtx`). Run
61+
`mobius list eps` to see all options.
62+
63+
**Typical export matrix:** Build each dtype × EP combination:
64+
65+
```bash
66+
for dtype in f16 bf16; do
67+
for ep in default cuda onnx-standard; do
68+
mobius build --model google/gemma-4-12b-it \
69+
--dtype $dtype --ep $ep \
70+
--runtime ort-genai \
71+
--external-data safetensors --max-shard-size 5GB \
72+
output/${dtype}/${ep}/
73+
done
74+
done
75+
```
76+
77+
### Multi-model outputs
78+
79+
For multimodal models (VLMs, audio-language), `mobius build` produces
80+
multiple sub-models:
81+
82+
```
83+
output/
84+
├── decoder/ # Text decoder
85+
│ ├── model.onnx
86+
│ └── model.onnx.data.safetensors
87+
├── embedding/ # Embedding model
88+
│ ├── model.onnx
89+
│ └── model.onnx.data.safetensors
90+
├── vision_encoder/ # Vision encoder (VLMs)
91+
│ ├── model.onnx
92+
│ └── model.onnx.data.safetensors
93+
├── audio_encoder/ # Audio encoder (ALMs)
94+
│ ├── model.onnx
95+
│ └── model.onnx.data.safetensors
96+
└── genai_config.json
97+
```
98+
99+
## Quantization with Olive
100+
101+
### Installation
102+
103+
Olive with ONNX quantization support (install from PR if needed for
104+
latest features):
105+
106+
```bash
107+
pip install olive-ai
108+
# Or from a specific PR for bleeding-edge features:
109+
pip install git+https://github.com/microsoft/Olive.git@refs/pull/2406/head
110+
```
111+
112+
For GPU-accelerated quantization (highly recommended for large models):
113+
114+
```bash
115+
pip install cupy-cuda12x
116+
```
117+
118+
### Q4_K_M quantization (k-quant)
119+
120+
K-quant quantization uses mixed block sizes with importance-based bit
121+
allocation. Q4_K_M is a good balance of quality and size.
122+
123+
The repo uses Olive's config-driven `olive.run()` pattern (see
124+
`examples/olive/` for working examples). A typical Olive config for
125+
k-quant quantization:
126+
127+
```json
128+
{
129+
"input_model": { "type": "OnnxModel", "model_path": "decoder/model.onnx" },
130+
"passes": {
131+
"kquant": {
132+
"type": "OnnxKQuantQuantization",
133+
"bits": 4,
134+
"block_size": 32
135+
}
136+
},
137+
"output_dir": "output/Q4_K_M/default/decoder"
138+
}
139+
```
140+
141+
```bash
142+
olive run --config kquant_config.json
143+
```
144+
145+
### NF4 quantization (4-bit NormalFloat)
146+
147+
NF4 uses a normal-distribution-optimized 4-bit format. Fast native C++
148+
implementation — no GPU needed.
149+
150+
```json
151+
{
152+
"input_model": { "type": "OnnxModel", "model_path": "decoder/model.onnx" },
153+
"passes": {
154+
"nf4": {
155+
"type": "OnnxBnb4Quantization",
156+
"precision": "nf4"
157+
}
158+
},
159+
"output_dir": "output/NF4/default/decoder"
160+
}
161+
```
162+
163+
> See `examples/olive/ministral-3-3b-vlm/` for a complete working
164+
> example that combines mobius export with Olive quantization.
165+
166+
### GPU acceleration with cupy
167+
168+
Installing `cupy-cuda12x` gives a **19–51x speedup** for k-quant
169+
quantization:
170+
171+
| Method | CPU time per matrix | GPU time per matrix | Speedup |
172+
|--------|--------------------|--------------------|---------|
173+
| K-quant (Q4_K_M) | 3–27s | 0.17–0.52s | 19–51x |
174+
| NF4 | 42ms for 67M params | N/A (C++ native) | Already fast |
175+
176+
```bash
177+
# Install cupy for CUDA 12.x
178+
pip install cupy-cuda12x
179+
180+
# Olive auto-detects cupy and uses GPU when available
181+
```
182+
183+
### Quantizing multi-model exports
184+
185+
Quantize each sub-model independently. Typically only the decoder is
186+
quantized (it has the most parameters). Copy all other files needed
187+
for a complete ORT GenAI package:
188+
189+
```bash
190+
# Quantize decoder only (largest model)
191+
olive run --config kquant_decoder.json
192+
193+
# Copy other sub-models as-is (already small)
194+
cp -r output/f16/default/embedding/ output/Q4_K_M/default/embedding/
195+
cp -r output/f16/default/vision_encoder/ output/Q4_K_M/default/vision_encoder/
196+
197+
# IMPORTANT: Copy config, tokenizer, and processor files too
198+
cp output/f16/default/genai_config.json output/Q4_K_M/default/
199+
cp output/f16/default/tokenizer* output/Q4_K_M/default/
200+
cp output/f16/default/image_processor.json output/Q4_K_M/default/ 2>/dev/null
201+
cp output/f16/default/audio_processor.json output/Q4_K_M/default/ 2>/dev/null
202+
```
203+
204+
Without the tokenizer and processor config files, ORT GenAI will fail
205+
to load the model.
206+
207+
## HuggingFace upload structure
208+
209+
### Standard directory layout
210+
211+
```
212+
<org>/<model>-onnx/
213+
├── f16/
214+
│ ├── default/ # Portable ONNX (no vendor fusions)
215+
│ ├── cuda/ # CUDA EP (fused ops)
216+
│ └── onnx-standard/ # Strict ONNX-only (inlined functions)
217+
├── bf16/
218+
│ ├── default/
219+
│ ├── cuda/
220+
│ └── onnx-standard/
221+
├── Q4_K_M/
222+
│ └── default/ # Quantized models typically CPU-only
223+
└── NF4/
224+
└── default/
225+
```
226+
227+
Each EP directory contains the full model structure (decoder/,
228+
embedding/, vision_encoder/, audio_encoder/ as applicable) plus
229+
`genai_config.json`.
230+
231+
### Upload with huggingface_hub
232+
233+
```python
234+
from huggingface_hub import HfApi
235+
236+
api = HfApi()
237+
api.upload_folder(
238+
folder_path="output/f16/default",
239+
path_in_repo="f16/default",
240+
repo_id="org/model-onnx",
241+
repo_type="model",
242+
)
243+
```
244+
245+
### Verify uploads
246+
247+
After uploading, verify all shards are present. Incomplete uploads are
248+
a common issue with large models:
249+
250+
```python
251+
from huggingface_hub import HfApi
252+
253+
api = HfApi()
254+
files = api.list_repo_files("org/model-onnx")
255+
# Check that all expected .safetensors shards exist
256+
for variant in ["f16/default", "f16/cuda", "bf16/default"]:
257+
shards = [f for f in files if f.startswith(variant) and f.endswith(".safetensors")]
258+
print(f"{variant}: {len(shards)} shards")
259+
```
260+
261+
## Common issues and fixes
262+
263+
### 1. MoE expert weight mapping
264+
265+
Models with Mixture-of-Experts (e.g. Gemma4 26b-a4b) may need expert
266+
weight remapping in `preprocess_weights()`. HuggingFace stores experts
267+
as 3D tensors (`experts.gate_up_proj [E, 2*inter, H]`) that must be
268+
mapped to the fused MoE op's parameter names (`fc1_experts_weights`,
269+
`fc2_experts_weights`).
270+
271+
**Symptom:** Weight loading errors or incorrect MoE outputs.
272+
273+
**Fix:** Check the model's `preprocess_weights()` maps HF expert weight
274+
names to the ONNX parameter names. See the `moe-models` skill for
275+
the pattern.
276+
277+
### 2. Hybrid attention v_proj shape mismatches
278+
279+
Models with hybrid attention (e.g. Gemma4 31b with different `head_dim`
280+
for local vs global attention layers) may have shape mismatches in
281+
value projections.
282+
283+
**Symptom:** Shape errors during weight loading or forward pass.
284+
285+
**Fix:** Ensure `v_proj` dimensions account for per-layer head
286+
configurations. Check `num_global_key_value_heads` vs
287+
`num_key_value_heads` in the config.
288+
289+
### 3. CUDA GQA head_dim limitations
290+
291+
Older versions of ORT had a limitation where `head_dim > 256` would fail
292+
with the CUDA GroupQueryAttention kernel.
293+
294+
**Symptom:** CUDA runtime error during inference with large head
295+
dimensions.
296+
297+
**Status:** This limitation has been removed in recent ORT versions.
298+
If using an older ORT build, fall back to `--ep default` or
299+
`--ep onnx-standard`.
300+
301+
### 4. Incomplete uploads
302+
303+
Large models with many shards can have incomplete uploads to HuggingFace
304+
Hub, especially on unstable connections.
305+
306+
**Symptom:** Model fails to load with file-not-found errors for
307+
specific shard files.
308+
309+
**Fix:** Verify all shards are present after upload (see the verify
310+
script above). Re-upload missing shards with `api.upload_file()`.
311+
312+
### 5. BF16 type mismatches
313+
314+
Some components may produce FP32 outputs when the model is built in
315+
BF16, causing type mismatch errors in ORT.
316+
317+
**Symptom:** `Type Error: Type parameter (T) bound to different types
318+
(tensor(bfloat16) and tensor(float))`.
319+
320+
**Fix:** Check for constants, initializers, or norm layers that stay
321+
FP32 when the model is BF16. Add `op.CastLike(result, input)` to
322+
ensure dtype consistency. See the `reusable-components` skill's
323+
section on precision behaviour.
324+
325+
## Testing quantized models
326+
327+
### L4: Golden data generation
328+
329+
Generate reference outputs from the full-precision HuggingFace model
330+
using the golden data generation script:
331+
332+
```bash
333+
# Generate golden files for all test cases
334+
python scripts/generate_golden.py
335+
336+
# Generate for a specific task type
337+
python scripts/generate_golden.py --task-type causal-lm
338+
339+
# Generate for a single test case
340+
python scripts/generate_golden.py --case testdata/cases/causal-lm/gpt2.yaml
341+
342+
# Use GPU for large models
343+
python scripts/generate_golden.py --device cuda
344+
```
345+
346+
Golden reference files are stored in `testdata/golden/` as JSON. Use
347+
`compare_golden()` from `mobius._testing.parity` to compare model
348+
outputs against the reference:
349+
350+
```python
351+
from mobius._testing.parity import compare_golden
352+
353+
compare_golden(
354+
model_output=output_logits,
355+
golden_path="testdata/golden/causal-lm/my_model.json",
356+
)
357+
```
358+
359+
### L5: End-to-end smoke test
360+
361+
Run inference with the quantized model through ORT GenAI:
362+
363+
```python
364+
import onnxruntime_genai as og
365+
366+
model = og.Model("output/Q4_K_M/default/")
367+
tokenizer = og.Tokenizer(model)
368+
params = og.GeneratorParams(model)
369+
params.set_search_options(max_length=50, do_sample=False)
370+
params.input_ids = tokenizer.encode("Hello, world!")
371+
372+
output_ids = model.generate(params)
373+
print(tokenizer.decode(output_ids[0]))
374+
```
375+
376+
### Numerical parity verification
377+
378+
Quantized models will have some numerical divergence from the
379+
full-precision model. Expected tolerances:
380+
381+
| Quantization | Typical divergence | Notes |
382+
|-------------|-------------------|-------|
383+
| Q4_K_M | Moderate | Top-1 token agreement ~95%+ for coherent text |
384+
| NF4 | Moderate | Similar to Q4_K_M |
385+
| F16 (no quant) | Minimal | Should match BF16 closely |
386+
387+
Verify that generated text is coherent and semantically correct rather
388+
than requiring exact numerical matches.
389+
390+
## Cross-references
391+
392+
- **Adding models:** `.agents/skills/adding-a-new-model/SKILL.md`
393+
- **MoE weights:** `.agents/skills/moe-models/SKILL.md`
394+
- **Component precision:** `.agents/skills/reusable-components/SKILL.md`
395+
- **ORT GenAI config:** `.agents/skills/ort-genai-config/SKILL.md`
396+
- **Quality checklist:** `.agents/skills/quality-checklist/SKILL.md`

0 commit comments

Comments
 (0)