diff --git a/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/_index.md b/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/_index.md new file mode 100644 index 0000000000..7ae65d2891 --- /dev/null +++ b/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/_index.md @@ -0,0 +1,65 @@ +--- +title: Convert SmolVLA to ONNX for Arm CPUs + +draft: true +cascade: + draft: true + +description: Export SmolVLA from PyTorch to ONNX, create a packed INT4 weight-only model with TorchAO quantization, and run both models with ONNX Runtime on an Arm CPU. +minutes_to_complete: 180 +who_is_this_for: Machine learning developers who want to deploy a vision-language-action policy with ONNX Runtime on an Arm CPU. +learning_objectives: + - Export SmolVLA from PyTorch as an ONNX model. + - Run and validate the FP32 ONNX model with ONNX Runtime on an Arm CPU. + - Quantize eligible linear weights to INT4 and store them in a packed ONNX model. + - Run the FP32 and INT4 models with identical inputs and compare their action outputs and ONNX Runtime latency. +prerequisites: + - An aarch64 Linux system, such as the Radxa Orion O6. + - Enough free storage for the model data, checkpoint weights, Python environment, and generated ONNX files. + - Git and Python 3.12. + - Familiarity with Python, PyTorch, and Linux command-line tools. +author: "" +generate_summary_faq: true +rerun_summary: false +rerun_faqs: false +skilllevels: Advanced +subjects: ML +armips: + - Cortex-A +operatingsystems: + - Linux +tools_software_languages: + - Python + - PyTorch + - TorchAO + - ONNX + - ONNX Runtime + - LeRobot +further_reading: + - resource: + title: PyTorch ONNX exporter documentation + link: https://docs.pytorch.org/docs/stable/onnx.html + type: documentation + - resource: + title: TorchAO documentation + link: https://docs.pytorch.org/ao/stable/ + type: documentation + - resource: + title: ONNX Runtime 4-bit quantization + link: https://onnxruntime.ai/docs/performance/model-optimizations/quantization.html#quantize-to-int4uint4 + type: documentation + - resource: + title: ONNX Runtime Execution Providers + link: https://onnxruntime.ai/docs/execution-providers/ + type: documentation + - resource: + title: KleidiAI optimized micro-kernels for Arm CPUs + link: https://github.com/ARM-software/kleidiai + type: GitHub Repository + +### FIXED, DO NOT MODIFY +# ================================================================================ +weight: 1 +layout: "learningpathall" +learning_path_main_page: "yes" +--- diff --git a/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/_next-steps.md b/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/_next-steps.md new file mode 100644 index 0000000000..727b395ddd --- /dev/null +++ b/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/_next-steps.md @@ -0,0 +1,8 @@ +--- +# ================================================================================ +# FIXED, DO NOT MODIFY THIS FILE +# ================================================================================ +weight: 21 # The weight controls the order of the pages. _index.md always has weight 1. +title: "Next Steps" # Always the same, html page title. +layout: "learningpathall" # All files under learning paths have this same wrapper for Hugo processing. +--- diff --git a/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/convert-and-validate-onnx.md b/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/convert-and-validate-onnx.md new file mode 100644 index 0000000000..d2d249d107 --- /dev/null +++ b/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/convert-and-validate-onnx.md @@ -0,0 +1,88 @@ +--- +title: Export and validate the SmolVLA ONNX model +description: Export SmolVLA from PyTorch as an FP32 ONNX model and validate its action output on an Arm CPU. +weight: 3 +layout: learningpathall +--- + +## Understand SmolVLA + +SmolVLA is a compact vision-language-action model. It combines camera images, +a task instruction, and the robot state to generate a sequence of robot +actions. + +![SmolVLA combines camera images, a task instruction, and robot state in a vision-language model that conditions an action expert to generate an action sequence.#center](smolvla.png "SmolVLA architecture") + +Image source: [SmolVLA paper](https://arxiv.org/pdf/2506.01844). + +## Export SmolVLA to ONNX + +The checkpoint includes the SmolVLA policy and the LeRobot processors used +before and after inference. The exporter writes the policy to an ONNX graph; +the processors remain outside it. + +Run the exporter: + +```bash +work/venv/bin/python scripts/export_onnx.py \ + --checkpoint work/artifacts/smolvla_libero \ + --output work/onnx/fp32/model.onnx \ + --reference-dir work/onnx/fp32/reference +``` + +The exporter creates a fixed-shape model and a deterministic reference batch. +The batch includes an explicit flow-matching noise tensor, so the PyTorch and +ONNX Runtime paths receive the same inputs. + +The expected output ends with: + +```output +PASS: ONNX Runtime matches PyTorch within atol=0.001 and rtol=0.001 +``` + +The FP32 graph contains the exported policy computation. Large weights can be +stored in external data files beside `model.onnx`, so keep the complete +`work/onnx/fp32` directory together. + +## Review the model interface + +The graph accepts six preprocessed inputs: + +| Input | Type | Shape | Description | +| --- | --- | --- | --- | +| `camera1` | `float32` | `[1, 3, 512, 512]` | Primary RGB image scaled to `[0, 1]` | +| `camera2` | `float32` | `[1, 3, 512, 512]` | Wrist RGB image scaled to `[0, 1]` | +| `lang_tokens` | `int64` | `[1, 48]` | Tokenized task instruction | +| `lang_attention_mask` | `int64` | `[1, 48]` | Valid language-token mask | +| `state` | `float32` | `[1, 8]` | Normalized robot state | +| `noise` | `float32` | `[1, 50, 32]` | Flow-matching noise | + +The model returns an action chunk with shape `[1, 50, 7]`: 50 steps with seven +control values each. + +The ONNX boundary deliberately starts after preprocessing and ends before +postprocessing: + +```text +observation -> LeRobot preprocessor -> ONNX model -> LeRobot postprocessor -> robot action +``` + +Use the processors from the same checkpoint. They preserve the tokenization, +normalization, and action conversion expected by the policy. + +## Inspect the validation report + +View the report written by the exporter: + +```bash +work/venv/bin/python -m json.tool work/onnx/fp32/validation.json +``` + +Confirm that the report lists `CPUExecutionProvider`, reports an output shape +of `[1, 50, 7]`, and passes validation. + +## What you've accomplished and what's next + +You have exported SmolVLA to FP32 ONNX and validated its action output +with ONNX Runtime on an Arm CPU. Next, you will quantize the eligible linear +weights to packed INT4 and run the resulting model through the same interface. diff --git a/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/quantize-and-run-int4-onnx.md b/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/quantize-and-run-int4-onnx.md new file mode 100644 index 0000000000..0cbcea4b65 --- /dev/null +++ b/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/quantize-and-run-int4-onnx.md @@ -0,0 +1,66 @@ +--- +title: Quantize SmolVLA to INT4 and compare it with FP32 +description: Quantize eligible SmolVLA linear weights with TorchAO, then compare FP32 and INT4 outputs and ONNX Runtime latency on Arm. +weight: 4 +layout: learningpathall +--- + +## INT4 quantization scope + +TorchAO quantizes eligible constant linear weights across the exported SmolVLA +model. The converter replaces supported ONNX `MatMul` and `Gemm` operations +with packed `com.microsoft::MatMulNBits` operations. + +This is weight-only quantization, not an entirely INT4 graph. + +On supported Arm CPUs, ONNX Runtime can use optimized kernels such as KleidiAI. + +## Create the packed INT4 model + +Run the TorchAO converter: + +```bash +work/venv/bin/python scripts/quantize_onnx_torchao.py \ + --input work/onnx/fp32/model.onnx \ + --output work/onnx/int4/smolvla-int4.onnx +``` + +The command creates a packed ONNX file and reports how many eligible linear +operations were converted. Dynamic or unsupported matrix multiplications +remain in floating point. + +## Compare FP32 and INT4 + +Run both models with the deterministic reference batch created during export: + +```bash +work/venv/bin/python scripts/compare_onnx_outputs.py \ + --fp32-model work/onnx/fp32/model.onnx \ + --int4-model work/onnx/int4/smolvla-int4.onnx \ + --reference-dir work/onnx/fp32/reference \ + --output work/comparison/smolvla-action-comparison.png +``` + +The script runs both models with ONNX Runtime `CPUExecutionProvider` and +creates: + +```text +work/comparison/smolvla-action-comparison.png +work/comparison/smolvla-action-comparison.json +``` + +The figure compares all seven normalized output channels and median latency. +The JSON file records the latency and overall output error. + +## Review the O6 result + +![Seven plots compare FP32 and TorchAO INT4 normalized SmolVLA outputs across all 50 predicted steps for each of seven channels. A latency panel compares median ONNX Runtime latency on a Radxa Orion O6.#center](smolvla-action-comparison.png "SmolVLA Action Comparison") + +On the O6, INT4 reduced median ONNX Runtime latency from 3.33 seconds to 2.06 +seconds, a 1.61x speedup. The normalized outputs had an MAE of 0.153. + +## What you've accomplished + +You have converted eligible SmolVLA linear weights to packed INT4 in an ONNX +model, run FP32 and INT4 with identical inputs on an Arm CPU, and compared all +seven normalized output channels and ONNX Runtime latency. diff --git a/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/scripts/check_assets.py b/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/scripts/check_assets.py new file mode 100644 index 0000000000..15c4afb7ad --- /dev/null +++ b/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/scripts/check_assets.py @@ -0,0 +1,114 @@ +#!/usr/bin/env python3 +"""Verify the pinned model, source, and environment used by the Learning Path.""" + +from __future__ import annotations + +import argparse +from importlib import metadata +import json +from pathlib import Path +import subprocess +import sys + +from workspace import configure_workspace + + +EXPECTED = { + "lerobot": "30da8e687a6dfc617fcd94afc367ac7071c376ce", + "model_repo": "HuggingFaceVLA/smolvla_libero", + "model_revision": "6721902bc4d61e50a3bfdb11dfb4cb626f05d102", + "base_model_repo": "HuggingFaceTB/SmolVLM2-500M-Instruct", + "base_model_revision": "7b375e1b73b11138ff12fe22c8f2822d8fe03467", +} + +EXPECTED_PACKAGES = { + "torch": "2.11.0+cpu", + "torchvision": "0.26.0+cpu", + "torchao": "0.18.0", + "transformers": "5.5.4", + "onnx": "1.22.0", + "onnxruntime": "1.29.0", +} + +POLICY_FILES = ( + "config.json", + "model.safetensors", + "policy_preprocessor.json", + "policy_preprocessor_step_5_normalizer_processor.safetensors", + "policy_postprocessor.json", + "policy_postprocessor_step_1_unnormalizer_processor.safetensors", +) +BASE_MODEL_FILES = ("config.json", "model.safetensors", "tokenizer.json") + + +def verify_snapshot(root: Path, files: tuple[str, ...], revision: str) -> None: + """Verify required snapshot files and their Hugging Face revisions.""" + + for relative in files: + asset = root / relative + metadata_file = root / ".cache/huggingface/download" / f"{relative}.metadata" + if not asset.is_file(): + raise FileNotFoundError(f"Missing public asset: {asset}") + if not metadata_file.is_file(): + raise FileNotFoundError(f"Missing Hugging Face metadata: {metadata_file}") + lines = metadata_file.read_text(encoding="utf-8").splitlines() + if not lines or lines[0] != revision: + raise RuntimeError(f"Unexpected Hugging Face revision for {asset}") + + +def verify_environment(work_root: Path) -> None: + """Verify the Python version, pinned packages, and environment manifest.""" + + if sys.version_info[:2] != (3, 12): + raise RuntimeError("Python 3.12 is required") + environment = work_root / "environment.freeze.txt" + if not environment.is_file(): + raise FileNotFoundError(f"Missing environment manifest: {environment}") + + for package, expected in EXPECTED_PACKAGES.items(): + actual = metadata.version(package) + if actual != expected: + raise RuntimeError(f"Unexpected {package} version: {actual}") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--work-root", type=Path) + args = parser.parse_args() + work_root = configure_workspace(args.work_root) + + revisions_path = work_root / "revisions.json" + revisions = json.loads(revisions_path.read_text(encoding="utf-8")) + if revisions != EXPECTED: + raise RuntimeError(f"Revision manifest does not match the Learning Path: {revisions_path}") + + actual_lerobot = subprocess.check_output( + ["git", "-C", str(work_root / "lerobot"), "rev-parse", "HEAD"], + text=True, + ).strip() + if actual_lerobot != EXPECTED["lerobot"]: + raise RuntimeError(f"Unexpected LeRobot revision: {actual_lerobot}") + lerobot_status = subprocess.check_output( + ["git", "-C", str(work_root / "lerobot"), "status", "--porcelain"], + text=True, + ).strip() + if lerobot_status: + raise RuntimeError("LeRobot worktree has local changes") + + verify_snapshot( + work_root / "artifacts/smolvla_libero", + POLICY_FILES, + EXPECTED["model_revision"], + ) + verify_snapshot( + work_root / "artifacts/smolvlm_base", + BASE_MODEL_FILES, + EXPECTED["base_model_revision"], + ) + verify_environment(work_root) + + print("PASS: public policy, base model, LeRobot source, and environment are ready") + + +if __name__ == "__main__": + main() diff --git a/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/scripts/compare_onnx_outputs.py b/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/scripts/compare_onnx_outputs.py new file mode 100644 index 0000000000..d32009b776 --- /dev/null +++ b/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/scripts/compare_onnx_outputs.py @@ -0,0 +1,234 @@ +#!/usr/bin/env python3 +"""Compare FP32 and INT4 SmolVLA ONNX outputs and CPU latency.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +import platform +import time + +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +import onnxruntime as ort + +INPUT_INTERFACE = ( + ("camera1", "tensor(float)", (1, 3, 512, 512)), + ("camera2", "tensor(float)", (1, 3, 512, 512)), + ("lang_tokens", "tensor(int64)", (1, 48)), + ("lang_attention_mask", "tensor(int64)", (1, 48)), + ("state", "tensor(float)", (1, 8)), + ("noise", "tensor(float)", (1, 50, 32)), +) +DTYPES = {"tensor(float)": np.dtype("float32"), "tensor(int64)": np.dtype("int64")} +ACTION_SHAPE = (1, 50, 7) +COLORS = {"FP32": "#2563EB", "INT4": "#F97316"} + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--fp32-model", type=Path, required=True) + parser.add_argument("--int4-model", type=Path, required=True) + parser.add_argument("--reference-dir", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True, help="Output PNG path") + parser.add_argument("--threads", type=int, default=8) + parser.add_argument("--warmups", type=int, default=1) + parser.add_argument("--runs", type=int, default=5) + args = parser.parse_args() + if args.threads < 1 or args.runs < 1 or args.warmups < 0: + parser.error("--threads and --runs must be positive; --warmups cannot be negative") + if args.output.suffix.lower() != ".png": + parser.error("--output must end in .png") + return args + +def make_session(path: Path, threads: int) -> ort.InferenceSession: + options = ort.SessionOptions() + options.intra_op_num_threads = threads + options.inter_op_num_threads = 1 + options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL + options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL + options.log_severity_level = 3 + return ort.InferenceSession( + str(path.resolve(strict=True)), options, providers=["CPUExecutionProvider"] + ) + +def interface(session: ort.InferenceSession) -> tuple[tuple[tuple[object, ...], ...], ...]: + describe = lambda values: tuple( + (value.name, value.type, tuple(value.shape)) for value in values + ) + return describe(session.get_inputs()), describe(session.get_outputs()) + +def validate_models(fp32: ort.InferenceSession, int4: ort.InferenceSession) -> None: + description = interface(fp32) + if description != interface(int4): + raise ValueError("FP32 and INT4 model interfaces differ") + inputs, outputs = description + if inputs != INPUT_INTERFACE: + raise ValueError(f"Unexpected model inputs: {inputs}") + if outputs != (("actions", "tensor(float)", ACTION_SHAPE),): + raise ValueError(f"Expected one float actions output with shape {ACTION_SHAPE}") + +def load_reference( + directory: Path, session: ort.InferenceSession +) -> dict[str, np.ndarray]: + directory = directory.resolve(strict=True) + if not directory.is_dir(): + raise NotADirectoryError(directory) + feeds = {} + for model_input in session.get_inputs(): + path = directory / f"{model_input.name}.npy" + value = np.load(path, allow_pickle=False) + expected_dtype = DTYPES.get(model_input.type) + if expected_dtype is None or value.dtype != expected_dtype: + raise TypeError(f"{path.name} has dtype {value.dtype}; expected {expected_dtype}") + if value.shape != tuple(model_input.shape): + raise ValueError( + f"{path.name} has shape {value.shape}; expected {tuple(model_input.shape)}" + ) + if np.issubdtype(value.dtype, np.floating) and not np.isfinite(value).all(): + raise ValueError(f"{path.name} contains non-finite values") + feeds[model_input.name] = np.ascontiguousarray(value) + return feeds + +def infer(session: ort.InferenceSession, feeds: dict[str, np.ndarray]) -> tuple[np.ndarray, float]: + start = time.perf_counter_ns() + result = session.run(None, feeds) + elapsed_ms = (time.perf_counter_ns() - start) / 1_000_000 + if len(result) != 1: + raise ValueError(f"Expected one model output, found {len(result)}") + actions = np.asarray(result[0]) + if actions.shape != ACTION_SHAPE or actions.dtype != np.float32: + raise ValueError(f"Expected float32 actions with shape {ACTION_SHAPE}") + if not np.isfinite(actions).all(): + raise ValueError("A model produced non-finite actions") + return actions, elapsed_ms + +def run_pairs( + sessions: dict[str, ort.InferenceSession], + feeds: dict[str, np.ndarray], + iterations: int, + timed: bool, +) -> tuple[dict[str, np.ndarray], dict[str, list[float]]]: + outputs: dict[str, np.ndarray] = {} + timings = {"FP32": [], "INT4": []} + for index in range(iterations): + order = ("FP32", "INT4") if index % 2 == 0 else ("INT4", "FP32") + for name in order: + outputs[name], elapsed = infer(sessions[name], feeds) + if timed: + timings[name].append(elapsed) + return outputs, timings + +def render(output: Path, actions: dict[str, np.ndarray], medians: dict[str, float], speedup: float, mae: float) -> None: + background, ink, muted, grid = "#F4F7FB", "#142033", "#617086", "#DCE4EE" + plt.rcParams.update({"font.family": "DejaVu Sans", "text.color": ink}) + figure = plt.figure(figsize=(12.8, 7.2), dpi=100, facecolor=background) + outer = figure.add_gridspec( + 1, 2, left=0.055, right=0.965, bottom=0.08, top=0.82, + width_ratios=(4.4, 1.25), wspace=0.18 + ) + plots = outer[0, 0].subgridspec(4, 2, hspace=0.5, wspace=0.24) + figure.text(0.055, 0.945, "SmolVLA Action Comparison", fontsize=22, weight="bold") + figure.text( + 0.055, 0.9, + "FP32 vs INT4 · identical reference inputs · 50 predicted steps · normalized model outputs", + fontsize=10.5, color=muted, + ) + figure.text(0.055, 0.85, "SEVEN ACTION CHANNELS", fontsize=8.5, weight="bold", color=muted) + figure.text(0.055, 0.025, "Each channel uses its own vertical scale.", fontsize=7.5, color=muted) + + steps = np.arange(ACTION_SHAPE[1]) + fp32, int4 = actions["FP32"][0], actions["INT4"][0] + axes = [figure.add_subplot(plots[index // 2, index % 2]) for index in range(7)] + for channel, axis in enumerate(axes): + axis.plot(steps, fp32[:, channel], color=COLORS["FP32"], linewidth=1.7) + axis.plot(steps, int4[:, channel], color=COLORS["INT4"], linewidth=1.5, linestyle="--") + axis.set_title(f"Action {channel}", loc="left", fontsize=9, weight="bold", pad=3) + axis.text( + 1, 1.03, f"MAE {np.mean(np.abs(int4[:, channel] - fp32[:, channel])):.3g}", + transform=axis.transAxes, ha="right", va="bottom", fontsize=7, color=muted, + ) + axis.set_xlim(0, 49) + axis.set_xticks((0, 25, 49)) + axis.tick_params(labelsize=7, colors=muted, length=2) + axis.grid(axis="y", color=grid, linewidth=0.7) + axis.spines[["top", "right"]].set_visible(False) + axis.spines[["left", "bottom"]].set_color("#C9D5E3") + for axis in (axes[5], axes[6]): + axis.set_xlabel("Predicted step", fontsize=7.5, color=muted) + + legend = figure.add_subplot(plots[3, 1]) + legend.set_axis_off() + legend.plot([], [], color=COLORS["FP32"], linewidth=2, label="FP32") + legend.plot([], [], color=COLORS["INT4"], linewidth=2, linestyle="--", label="INT4") + legend.legend(loc="upper left", frameon=False, fontsize=9, ncol=2) + + latency = figure.add_subplot(outer[0, 1], facecolor="white") + names = ("FP32", "INT4") + values = [medians[name] for name in names] + bars = latency.barh(names, values, color=[COLORS[name] for name in names], height=0.48) + latency.invert_yaxis() + latency.bar_label(bars, labels=[f"{value:,.1f} ms" for value in values], padding=4, fontsize=9) + latency.set_xlim(0, max(values) * 1.28) + latency.set_title("CPU latency", loc="left", fontsize=10, weight="bold", pad=14) + latency.tick_params(axis="y", labelsize=9, length=0) + latency.tick_params(axis="x", bottom=False, labelbottom=False) + latency.spines[:].set_visible(False) + latency.text(0, -0.55, "Median session.run · lower is better", fontsize=8, color=muted) + latency.set_ylim(2.8, -0.75) + latency.text(0, 2.05, f"{speedup:.2f}×", fontsize=25, weight="bold") + latency.text(0, 2.35, "FP32 latency ÷ INT4 latency", fontsize=8, color=muted) + latency.text(0, 2.62, f"Normalized MAE {mae:.4g}", fontsize=9, weight="bold") + + output.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(output, dpi=100, facecolor=background) + plt.close(figure) + +def main() -> None: + args = parse_args() + sessions = { + "FP32": make_session(args.fp32_model, args.threads), + "INT4": make_session(args.int4_model, args.threads), + } + validate_models(sessions["FP32"], sessions["INT4"]) + feeds = load_reference(args.reference_dir, sessions["FP32"]) + if args.warmups: + run_pairs(sessions, feeds, args.warmups, timed=False) + actions, timings = run_pairs(sessions, feeds, args.runs, timed=True) + medians = {name: float(np.median(values)) for name, values in timings.items()} + speedup = medians["FP32"] / medians["INT4"] + difference = actions["INT4"].astype(np.float64) - actions["FP32"] + mae = float(np.mean(np.abs(difference))) + report = { + "timing_scope": "ONNX Runtime session.run on CPUExecutionProvider", + "platform_machine": platform.machine(), "onnxruntime_version": ort.__version__, + "threads": args.threads, "warmups": args.warmups, "runs": args.runs, + "model_interface": { + "inputs": [ + {"name": name, "type": dtype, "shape": shape} + for name, dtype, shape in INPUT_INTERFACE + ], + "output": {"name": "actions", "type": "tensor(float)", "shape": ACTION_SHAPE}, + }, + "median_latency_ms": medians, "speedup": speedup, + "normalized_output_error": { + "mae": mae, + "rmse": float(np.sqrt(np.mean(np.square(difference)))), + "max_abs": float(np.max(np.abs(difference))), + }, + } + output = args.output.resolve() + sidecar = output.with_suffix(".json") + render(output, actions, medians, speedup, mae) + sidecar.write_text( + json.dumps(report, indent=2, allow_nan=False) + "\n", encoding="utf-8" + ) + print(f"FP32 median: {medians['FP32']:.3f} ms") + print(f"INT4 median: {medians['INT4']:.3f} ms") + print(f"Speedup: {speedup:.3f}x") + print(f"Saved {args.output} and {args.output.with_suffix('.json')}") + +if __name__ == "__main__": + main() diff --git a/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/scripts/export_onnx.py b/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/scripts/export_onnx.py new file mode 100644 index 0000000000..b3d9bc8d1a --- /dev/null +++ b/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/scripts/export_onnx.py @@ -0,0 +1,525 @@ +#!/usr/bin/env python3 +"""Export a public two-camera SmolVLA checkpoint and validate ONNX Runtime.""" + +from __future__ import annotations + +import argparse +from contextlib import contextmanager +import gc +import json +import math +import os +from pathlib import Path +import platform +import shutil +import tempfile +from typing import Sequence + +import numpy as np +import torch +from torch import Tensor, nn +from torch.nn import functional as F + +from workspace import configure_workspace, load_smolvla_policy + + +WORK_ROOT = configure_workspace() + + +INPUT_NAMES = ( + "camera1", + "camera2", + "lang_tokens", + "lang_attention_mask", + "state", + "noise", +) +EXPECTED_ACTION_DIM = 7 +EXPECTED_DENOISING_STEPS = 10 + +_ORIGINAL_TORCH_CUMSUM = torch.cumsum +_ORIGINAL_TORCH_FULL = torch.full + + +@contextmanager +def fresh_directory(destination: Path): + """Build a directory beside its destination, then publish it in one rename.""" + + destination = destination.resolve() + if os.path.lexists(destination): + raise FileExistsError(f"Refusing to overwrite {destination}") + destination.parent.mkdir(parents=True, exist_ok=True) + staging = Path( + tempfile.mkdtemp( + prefix=f".{destination.name}.tmp-", dir=str(destination.parent) + ) + ) + try: + yield staging + if os.path.lexists(destination): + raise FileExistsError(f"Destination appeared during the run: {destination}") + staging.rename(destination) + except BaseException: + if os.path.lexists(staging): + shutil.rmtree(staging) + raise + + +def write_json(path: Path, payload: object) -> None: + path.write_text( + json.dumps(payload, indent=2, sort_keys=True, allow_nan=False) + "\n", + encoding="utf-8", + ) + + +def validate_array( + name: str, + value: np.ndarray, + *, + shape: Sequence[int] | None = None, + dtype: np.dtype | type | None = None, +) -> None: + if shape is not None and tuple(value.shape) != tuple(shape): + raise ValueError(f"{name} has shape {value.shape}; expected {tuple(shape)}") + if dtype is not None and value.dtype != np.dtype(dtype): + raise TypeError(f"{name} has dtype {value.dtype}; expected {np.dtype(dtype)}") + if np.issubdtype(value.dtype, np.number) and not np.isfinite(value).all(): + raise ValueError(f"{name} contains a non-finite value") + + +def policy_dimensions(policy: nn.Module) -> dict[str, int]: + state_feature = policy.config.robot_state_feature + action_feature = policy.config.action_feature + if state_feature is None or len(state_feature.shape) != 1: + raise ValueError("The checkpoint must define one robot-state feature") + if action_feature is None or len(action_feature.shape) != 1: + raise ValueError("The checkpoint must define one action feature") + dimensions = { + "state_dim": int(state_feature.shape[0]), + "action_dim": int(action_feature.shape[0]), + "action_chunk_size": int(policy.model.config.chunk_size), + "latent_action_dim": int(policy.model.config.max_action_dim), + "denoising_steps": int(policy.model.config.num_steps), + } + if dimensions["action_dim"] != EXPECTED_ACTION_DIM: + raise ValueError( + f"This Learning Path requires {EXPECTED_ACTION_DIM} action channels" + ) + if dimensions["denoising_steps"] != EXPECTED_DENOISING_STEPS: + raise ValueError( + f"This Learning Path requires {EXPECTED_DENOISING_STEPS} denoising steps" + ) + return dimensions + + +def install_exportable_rope() -> None: + """Replace SmolVLA's in-place RoPE implementation with an ONNX-safe form.""" + + import lerobot.policies.smolvla.smolvlm_with_expert as expert_module + + def apply_rope(x: Tensor, positions: Tensor, max_wavelength: float = 10_000) -> Tensor: + half = x.shape[-1] // 2 + source_dtype = x.dtype + source = x.to(torch.float32) + exponents = (2.0 / x.shape[-1]) * torch.arange( + half, dtype=torch.float32, device=x.device + ) + timescale = max_wavelength**exponents + radians = positions[..., None].to(torch.float32) / timescale[None, None, :] + radians = radians[..., None, :] + first, second = source.split(half, dim=-1) + result = torch.cat( + [ + first * torch.cos(radians) - second * torch.sin(radians), + second * torch.cos(radians) + first * torch.sin(radians), + ], + dim=-1, + ) + return result.to(source_dtype) + + expert_module.apply_rope = apply_rope + + +def install_exportable_attention(policy: nn.Module) -> None: + """Use exportable eager attention in the vision encoder.""" + + for module in policy.modules(): + config = getattr(module, "config", None) + if ( + config is not None + and getattr(config, "model_type", None) == "smolvlm_vision" + and hasattr(config, "_attn_implementation") + ): + config._attn_implementation = "eager" + + +def static_trace_length(value: object) -> object: + """Convert a traced scalar shape to a constant for this fixed-shape graph.""" + + if isinstance(value, Tensor) and value.ndim == 0: + return int(value.detach().cpu()) + return value + + +def install_exportable_masking() -> None: + """Handle Transformers 5.5 scalar shape tensors in the legacy exporter.""" + + import transformers.masking_utils as masking_utils + + original_sdpa_mask = masking_utils.sdpa_mask + + def exportable_sdpa_mask(*args, **kwargs): + positional = list(args) + if len(positional) > 1: + positional[1] = static_trace_length(positional[1]) + elif "q_length" in kwargs: + kwargs["q_length"] = static_trace_length(kwargs["q_length"]) + if len(positional) > 2: + positional[2] = static_trace_length(positional[2]) + elif "kv_length" in kwargs: + kwargs["kv_length"] = static_trace_length(kwargs["kv_length"]) + return original_sdpa_mask(*positional, **kwargs) + + masking_utils.sdpa_mask = exportable_sdpa_mask + + +def boolean_safe_cumsum(input_tensor: Tensor, *args, **kwargs) -> Tensor: + """Preserve PyTorch bool cumsum semantics with an ONNX-valid integer input.""" + + if input_tensor.dtype == torch.bool: + input_tensor = input_tensor.to(torch.int64) + return _ORIGINAL_TORCH_CUMSUM(input_tensor, *args, **kwargs) + + +def install_exportable_cumsum() -> None: + """Make boolean cumulative sums legal for ONNX Runtime.""" + + torch.cumsum = boolean_safe_cumsum + + +def dtype_stable_full(size, fill_value, *args, **kwargs) -> Tensor: + """Make PyTorch's integer fill-value dtype inference explicit for ONNX.""" + + if "dtype" not in kwargs and type(fill_value) is int: + kwargs["dtype"] = torch.int64 + return _ORIGINAL_TORCH_FULL(size, fill_value, *args, **kwargs) + + +def install_exportable_full() -> None: + """Prevent legacy ONNX ScatterND from mixing float and integer tensors.""" + + torch.full = dtype_stable_full + + +def exportable_sinusoidal_embedding( + time: Tensor, + dimension: int, + min_period: float, + max_period: float, + device: torch.device | str = "cpu", +) -> Tensor: + """Compute SmolVLA timestep features in ONNX Runtime-supported FP32.""" + + if dimension % 2 != 0: + raise ValueError(f"dimension ({dimension}) must be divisible by 2") + if time.ndim != 1: + raise ValueError("The time tensor is expected to have shape (batch_size,)") + fraction = torch.linspace( + 0.0, + 1.0, + dimension // 2, + dtype=torch.float32, + device=device, + ) + period = min_period * (max_period / min_period) ** fraction + sin_input = (1.0 / period * 2 * math.pi)[None, :] * time.to(torch.float32)[:, None] + return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1) + + +def install_exportable_timestep_embedding() -> None: + """Avoid unsupported double-precision Sin/Cos kernels in ONNX Runtime CPU.""" + + import lerobot.policies.smolvla.modeling_smolvla as smolvla_module + + smolvla_module.create_sinusoidal_pos_embedding = exportable_sinusoidal_embedding + + +class ExportableSmolVLA(nn.Module): + """Expose the LIBERO policy core with explicit images, state, and noise.""" + + def __init__(self, policy: nn.Module, dimensions: dict[str, int]): + super().__init__() + self.model = policy.model + self.state_dim = dimensions["state_dim"] + self.action_dim = dimensions["action_dim"] + if len(policy.config.image_features) != 2: + raise ValueError( + "This Learning Path expects the two-camera SmolVLA-LIBERO checkpoint" + ) + if self.model.config.max_state_dim < self.state_dim: + raise ValueError("The requested state dimension exceeds the model maximum") + if self.model.config.max_action_dim < self.action_dim: + raise ValueError("The requested action dimension exceeds the model maximum") + + def forward( + self, + camera1: Tensor, + camera2: Tensor, + lang_tokens: Tensor, + lang_attention_mask: Tensor, + state: Tensor, + noise: Tensor, + ) -> Tensor: + batch_size = state.shape[0] + image_mask = torch.ones(batch_size, dtype=torch.bool, device=state.device) + images = [camera1 * 2.0 - 1.0, camera2 * 2.0 - 1.0] + image_masks = [image_mask, image_mask] + padded_state = F.pad(state, (0, self.model.config.max_state_dim - self.state_dim)) + actions = self.model.sample_actions( + images, + image_masks, + lang_tokens, + lang_attention_mask.to(torch.bool), + padded_state, + noise=noise, + ) + return actions[:, :, : self.action_dim] + + +def make_inputs( + policy: nn.Module, dimensions: dict[str, int], seed: int +) -> tuple[Tensor, ...]: + generator = torch.Generator(device="cpu").manual_seed(seed) + tokenizer = policy.model.vlm_with_expert.processor.tokenizer + encoded = tokenizer( + "pick up the black bowl and place it on the plate", + return_tensors="pt", + padding="max_length", + truncation=True, + max_length=48, + ) + return ( + torch.rand(1, 3, 512, 512, generator=generator), + torch.rand(1, 3, 512, 512, generator=generator), + encoded["input_ids"].to(torch.int64), + encoded["attention_mask"].to(torch.int64), + torch.randn( + 1, + dimensions["state_dim"], + generator=generator, + ), + torch.randn( + 1, + dimensions["action_chunk_size"], + dimensions["latent_action_dim"], + generator=generator, + ), + ) + + +def save_reference(directory: Path, inputs: Sequence[Tensor]) -> None: + directory.mkdir(parents=True, exist_ok=False) + for name, value in zip(INPUT_NAMES, inputs, strict=True): + np.save( + directory / f"{name}.npy", + value.detach().cpu().numpy(), + allow_pickle=False, + ) + + +def session_interface(session) -> dict[str, list[dict[str, object]]]: + def describe(value) -> dict[str, object]: + return {"name": value.name, "shape": list(value.shape), "type": value.type} + + return { + "inputs": [describe(value) for value in session.get_inputs()], + "outputs": [describe(value) for value in session.get_outputs()], + } + + +def export_bundle( + args: argparse.Namespace, + checkpoint: Path, + output: Path, + reference_dir: Path, +) -> None: + import onnx + import onnxruntime as ort + + torch.manual_seed(args.seed) + policy = load_smolvla_policy(checkpoint).to(device="cpu", dtype=torch.float32) + policy.eval() + dimensions = policy_dimensions(policy) + wrapper = ExportableSmolVLA(policy, dimensions).eval() + inputs = make_inputs(policy, dimensions, args.seed) + input_arrays = { + name: np.ascontiguousarray(value.detach().cpu().numpy()) + for name, value in zip(INPUT_NAMES, inputs, strict=True) + } + expected_inputs = { + "camera1": ((1, 3, 512, 512), np.float32), + "camera2": ((1, 3, 512, 512), np.float32), + "lang_tokens": ((1, 48), np.int64), + "lang_attention_mask": ((1, 48), np.int64), + "state": ((1, dimensions["state_dim"]), np.float32), + "noise": ( + ( + 1, + dimensions["action_chunk_size"], + dimensions["latent_action_dim"], + ), + np.float32, + ), + } + for name, value in input_arrays.items(): + shape, dtype = expected_inputs[name] + validate_array(name, value, shape=shape, dtype=dtype) + + action_shape = (1, dimensions["action_chunk_size"], dimensions["action_dim"]) + with torch.inference_mode(): + baseline_output = wrapper(*inputs).detach().cpu().numpy() + validate_array( + "unmodified PyTorch actions", + baseline_output, + shape=action_shape, + dtype=np.float32, + ) + + install_exportable_rope() + install_exportable_masking() + install_exportable_cumsum() + install_exportable_full() + install_exportable_timestep_embedding() + install_exportable_attention(policy) + with torch.inference_mode(): + pytorch_output = wrapper(*inputs).detach().cpu().numpy() + validate_array( + "export-safe PyTorch actions", + pytorch_output, + shape=action_shape, + dtype=np.float32, + ) + compatibility_difference = pytorch_output.astype(np.float64) - baseline_output.astype( + np.float64 + ) + compatibility_passed = bool( + np.allclose(pytorch_output, baseline_output, atol=args.atol, rtol=args.rtol) + ) + if not compatibility_passed: + raise RuntimeError("The export-safe PyTorch path changed checkpoint output") + + with torch.inference_mode(): + torch.onnx.export( + wrapper, + inputs, + str(output), + input_names=list(INPUT_NAMES), + output_names=["actions"], + opset_version=args.opset, + export_params=True, + keep_initializers_as_inputs=False, + external_data=True, + dynamo=False, + ) + onnx.checker.check_model(str(output), full_check=False) + del wrapper, policy + gc.collect() + + options = ort.SessionOptions() + options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_BASIC + session = ort.InferenceSession( + str(output), + sess_options=options, + providers=["CPUExecutionProvider"], + ) + feeds = input_arrays + ort_output = session.run(["actions"], feeds)[0] + validate_array( + "ONNX Runtime actions", ort_output, shape=action_shape, dtype=np.float32 + ) + difference = ort_output.astype(np.float64) - pytorch_output.astype(np.float64) + passed = bool(np.allclose(ort_output, pytorch_output, atol=args.atol, rtol=args.rtol)) + save_reference(reference_dir, inputs) + + report = { + "format": "smolvla-onnx-validation-v3", + "architecture": platform.machine(), + "seed": args.seed, + "opset": args.opset, + "execution_providers": session.get_providers(), + "derived_configuration": dimensions, + "interface": session_interface(session), + "export_compatibility": { + "max_absolute_error": float(np.abs(compatibility_difference).max()), + "mean_absolute_error": float(np.abs(compatibility_difference).mean()), + "passed": compatibility_passed, + }, + "max_absolute_error": float(np.abs(difference).max()), + "mean_absolute_error": float(np.abs(difference).mean()), + "atol": args.atol, + "rtol": args.rtol, + "passed": passed, + "versions": { + "python": platform.python_version(), + "torch": torch.__version__, + "onnx": onnx.__version__, + "onnxruntime": ort.__version__, + }, + } + report_path = output.parent / "validation.json" + write_json(report_path, report) + if not passed: + raise RuntimeError("ONNX Runtime output did not meet the export tolerances") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--reference-dir", type=Path, required=True) + parser.add_argument("--seed", type=int, default=20260813) + parser.add_argument("--opset", type=int, default=17) + parser.add_argument("--atol", type=float, default=1.0e-3) + parser.add_argument("--rtol", type=float, default=1.0e-3) + args = parser.parse_args() + + if args.opset < 1: + raise ValueError("--opset must be positive") + if not np.isfinite(args.atol) or args.atol < 0: + raise ValueError("--atol must be finite and non-negative") + if not np.isfinite(args.rtol) or args.rtol < 0: + raise ValueError("--rtol must be finite and non-negative") + + checkpoint = args.checkpoint.resolve(strict=True) + expected_checkpoint = (WORK_ROOT / "artifacts/smolvla_libero").resolve(strict=True) + if checkpoint != expected_checkpoint: + raise ValueError("--checkpoint must be the prepared workspace artifact") + final_output = args.output.resolve() + final_reference = args.reference_dir.resolve() + final_bundle = final_output.parent + if final_output.suffix.lower() != ".onnx": + raise ValueError("--output must use the .onnx suffix") + if final_output.name == "validation.json": + raise ValueError("--output collides with validation.json") + if final_reference.parent != final_bundle: + raise ValueError("--reference-dir must be inside the ONNX output bundle") + if final_reference == final_output: + raise ValueError("--output and --reference-dir must be different paths") + if final_reference.name == "validation.json": + raise ValueError("--reference-dir collides with validation.json") + + with fresh_directory(final_bundle) as staging: + export_bundle( + args, + checkpoint, + staging / final_output.name, + staging / final_reference.name, + ) + print( + f"PASS: ONNX Runtime matches PyTorch within atol={args.atol:g} " + f"and rtol={args.rtol:g}" + ) + + +if __name__ == "__main__": + main() diff --git a/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/scripts/quantize_onnx_torchao.py b/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/scripts/quantize_onnx_torchao.py new file mode 100644 index 0000000000..db57236dee --- /dev/null +++ b/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/scripts/quantize_onnx_torchao.py @@ -0,0 +1,340 @@ +#!/usr/bin/env python3 +"""Quantize constant ONNX linear weights with TorchAO INT4 for Arm CPUs.""" + +from __future__ import annotations + +import argparse +import os +from pathlib import Path +import tempfile + +import numpy as np +import onnx +import torch +import torchao +from onnx import AttributeProto, TensorProto, helper, numpy_helper +from torchao.quantization import ( + IntxWeightOnlyConfig, + MappingType, + PerGroup, + quantize_, +) + + +GROUP_SIZE = 32 +ACCURACY_LEVEL = 4 + + +class Names: + def __init__(self, graph: onnx.GraphProto) -> None: + self.used = {item.name for item in graph.initializer} + self.used.update(item.name for item in (*graph.input, *graph.output, *graph.value_info)) + self.used.update( + value + for node in graph.node + for value in (*node.input, *node.output) + if value + ) + + def make(self, base: str) -> str: + candidate = base + index = 1 + while candidate in self.used: + candidate = f"{base}_{index}" + index += 1 + self.used.add(candidate) + return candidate + + +def resolve_initializer( + name: str, + initializers: dict[str, onnx.TensorProto], + producers: dict[str, onnx.NodeProto], +) -> str | None: + seen: set[str] = set() + while name not in initializers: + if name in seen: + raise ValueError("Identity cycle found while resolving a weight") + seen.add(name) + producer = producers.get(name) + if producer is None or producer.op_type != "Identity" or len(producer.input) != 1: + return None + name = producer.input[0] + return name + + +def gemm_attributes(node: onnx.NodeProto) -> dict[str, object]: + return {item.name: helper.get_attribute_value(item) for item in node.attribute} + + +def eligible_weight( + node: onnx.NodeProto, + initializers: dict[str, onnx.TensorProto], + producers: dict[str, onnx.NodeProto], +) -> tuple[str, bool] | None: + transpose = False + if node.op_type == "MatMul" and len(node.input) >= 2: + pass + elif node.op_type == "Gemm" and len(node.input) >= 2: + attrs = gemm_attributes(node) + if ( + int(attrs.get("transA", 0)) != 0 + or int(attrs.get("transB", 0)) != 1 + or float(attrs.get("alpha", 1.0)) != 1.0 + or float(attrs.get("beta", 1.0)) != 1.0 + ): + return None + transpose = True + else: + return None + + name = resolve_initializer(node.input[1], initializers, producers) + if name is None: + return None + tensor = initializers[name] + if len(tensor.dims) != 2 or tensor.data_type != TensorProto.FLOAT: + return None + return name, transpose + + +def pack_nibbles(values: np.ndarray, pad_value: int) -> np.ndarray: + if values.shape[-1] % 2: + padding = [(0, 0)] * (values.ndim - 1) + [(0, 1)] + values = np.pad(values, padding, constant_values=pad_value) + values = np.asarray(values, dtype=np.uint8) + return values[..., 0::2] | (values[..., 1::2] << np.uint8(4)) + + +def torchao_qparams(weight: np.ndarray, config: IntxWeightOnlyConfig): + """Return ORT-packed qdata, scales, zero points, and K.""" + k, n = weight.shape + linear = torch.nn.Linear(k, n, bias=False, device="meta") + linear_weight = torch.from_numpy(np.ascontiguousarray(weight.T)).clone() + linear.weight = torch.nn.Parameter(linear_weight, requires_grad=False) + quantize_(linear, config) + + qweight = linear.weight + qdata = qweight.qdata.detach().cpu().numpy().astype(np.int16) + scales = qweight.scale.detach().cpu().numpy().astype(np.float32) + zero_points = qweight.zero_point.detach().cpu().numpy().astype(np.int16) + if qdata.min() < -8 or qdata.max() > 7 or zero_points.min() < -8 or zero_points.max() > 7: + raise ValueError("TorchAO produced values outside the signed INT4 range") + + blocks = k // GROUP_SIZE + packed = pack_nibbles((qdata + 8).astype(np.uint8), pad_value=0).reshape( + n, blocks, GROUP_SIZE // 2 + ) + packed_zp = pack_nibbles( + (zero_points + 8).astype(np.uint8), pad_value=8 + ) + if scales.shape != (n, blocks) or packed_zp.shape != (n, (blocks + 1) // 2): + raise ValueError("Unexpected TorchAO group-wise parameter shape") + return packed, np.ascontiguousarray(scales), np.ascontiguousarray(packed_zp), k + + +def nested_uses(graph: onnx.GraphProto) -> set[str]: + used = {item.name for item in graph.output} + for node in graph.node: + used.update(value for value in node.input if value) + for attr in node.attribute: + if attr.type == AttributeProto.GRAPH: + used.update(nested_uses(attr.g)) + elif attr.type == AttributeProto.GRAPHS: + for child in attr.graphs: + used.update(nested_uses(child)) + return used + + +def clean_graph(graph: onnx.GraphProto, original_initializers: set[str]) -> None: + while True: + used = nested_uses(graph) + kept = [ + node + for node in graph.node + if not ( + node.op_type == "Identity" + and all(output not in used for output in node.output) + ) + ] + if len(kept) == len(graph.node): + break + del graph.node[:] + graph.node.extend(kept) + + used = nested_uses(graph) + kept_initializers = [item for item in graph.initializer if item.name in used] + del graph.initializer[:] + graph.initializer.extend(kept_initializers) + kept_inputs = [ + item + for item in graph.input + if item.name not in original_initializers or item.name in used + ] + del graph.input[:] + graph.input.extend(kept_inputs) + + +def convert(model: onnx.ModelProto, config: IntxWeightOnlyConfig) -> dict[str, int]: + graph = model.graph + names = Names(graph) + initializers = {item.name: item for item in graph.initializer} + original_initializers = set(initializers) + producers = {output: node for node in graph.node for output in node.output if output} + cache: dict[tuple[str, bool], tuple[str, str, str, int, int]] = {} + new_initializers: list[onnx.TensorProto] = [] + new_nodes: list[onnx.NodeProto] = [] + converted_matmul = converted_gemm = 0 + + def quantized_weight(name: str, transpose: bool): + key = (name, transpose) + if key in cache: + return cache[key] + source = np.asarray(numpy_helper.to_array(initializers[name]), dtype=np.float32) + weight = np.ascontiguousarray(source.T if transpose else source) + k = weight.shape[0] + if k % GROUP_SIZE: + raise ValueError( + f"Eligible weight {name!r} has K={k}; K must be divisible by {GROUP_SIZE}" + ) + packed, scales, zero_points, k = torchao_qparams(weight, config) + suffix = "__transposed" if transpose else "" + packed_name = names.make(name + suffix + "__torchao_int4") + scales_name = names.make(name + suffix + "__torchao_scales") + zp_name = names.make(name + suffix + "__torchao_zero_points") + new_initializers.extend( + ( + numpy_helper.from_array(packed, packed_name), + numpy_helper.from_array(scales, scales_name), + numpy_helper.from_array(zero_points, zp_name), + ) + ) + cache[key] = (packed_name, scales_name, zp_name, k, weight.shape[1]) + return cache[key] + + for node in graph.node: + spec = eligible_weight(node, initializers, producers) + if spec is None: + new_nodes.append(node) + continue + weight_name, transpose = spec + packed, scales, zero_points, k, n = quantized_weight(weight_name, transpose) + output = node.output[0] + has_bias = node.op_type == "Gemm" and len(node.input) >= 3 and bool(node.input[2]) + matmul_output = names.make(output + "__torchao_matmul") if has_bias else output + new_nodes.append( + helper.make_node( + "MatMulNBits", + [node.input[0], packed, scales, zero_points], + [matmul_output], + name=names.make((node.name or output) + "__TorchAO_INT4"), + domain="com.microsoft", + K=k, + N=n, + bits=4, + block_size=GROUP_SIZE, + accuracy_level=ACCURACY_LEVEL, + ) + ) + if has_bias: + new_nodes.append( + helper.make_node( + "Add", + [matmul_output, node.input[2]], + [output], + name=names.make((node.name or output) + "/BiasAdd"), + ) + ) + if node.op_type == "MatMul": + converted_matmul += 1 + else: + converted_gemm += 1 + + if not converted_matmul and not converted_gemm: + raise ValueError("The model has no eligible FP32 MatMul or standard Gemm weights") + del graph.node[:] + graph.node.extend(new_nodes) + graph.initializer.extend(new_initializers) + clean_graph(graph, original_initializers) + + final_initializers = {item.name: item for item in graph.initializer} + final_producers = {output: node for node in graph.node for output in node.output if output} + missed = sum( + eligible_weight(node, final_initializers, final_producers) is not None + for node in graph.node + ) + if missed: + raise RuntimeError(f"Conversion left {missed} eligible FP32 linear nodes") + return { + "matmul": converted_matmul, + "gemm": converted_gemm, + "weights": len(cache), + "fp32_matmul": sum(node.op_type == "MatMul" for node in graph.node), + "fp32_gemm": sum(node.op_type == "Gemm" for node in graph.node), + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--input", type=Path, required=True, help="FP32 ONNX model") + parser.add_argument( + "--output", type=Path, required=True, help="new single-file INT4 ONNX model" + ) + args = parser.parse_args() + + source = args.input.expanduser().resolve() + output = args.output.expanduser().resolve() + if not source.is_file(): + raise FileNotFoundError(source) + if output.exists(): + raise FileExistsError(f"Refusing to overwrite {output.name}") + if torchao.__version__.split("+")[0].split(".")[:2] != ["0", "18"]: + raise RuntimeError(f"This converter requires TorchAO 0.18, found {torchao.__version__}") + + config = IntxWeightOnlyConfig( + weight_dtype=torch.int4, + granularity=PerGroup(GROUP_SIZE), + mapping_type=MappingType.ASYMMETRIC, + ) + model = onnx.load(source, load_external_data=True) + stats = convert(model, config) + microsoft_imports = [ + item for item in model.opset_import if item.domain == "com.microsoft" + ] + if microsoft_imports: + microsoft_imports[0].version = 1 + for duplicate in microsoft_imports[1:]: + model.opset_import.remove(duplicate) + else: + model.opset_import.append(helper.make_opsetid("com.microsoft", 1)) + onnx.external_data_helper.convert_model_from_external_data(model) + if any(onnx.external_data_helper.uses_external_data(item) for item in model.graph.initializer): + raise RuntimeError("Failed to embed all output tensors") + + output.parent.mkdir(parents=True, exist_ok=True) + descriptor, temporary_name = tempfile.mkstemp( + dir=output.parent, prefix=f".{output.name}.", suffix=".onnx" + ) + os.close(descriptor) + temporary = Path(temporary_name) + try: + onnx.save_model(model, temporary, save_as_external_data=False) + onnx.checker.check_model(str(temporary), full_check=False) + if output.exists(): + raise FileExistsError(f"Refusing to overwrite {output.name}") + temporary.chmod(0o644) + os.replace(temporary, output) + finally: + temporary.unlink(missing_ok=True) + print( + f"Converted {stats['matmul']} MatMul and {stats['gemm']} Gemm nodes " + f"using {stats['weights']} TorchAO INT4 weights." + ) + print( + f"Kept {stats['fp32_matmul']} dynamic/unsupported MatMul and " + f"{stats['fp32_gemm']} Gemm nodes in FP32." + ) + print(f"Wrote one ONNX file ({output.stat().st_size / (1024 * 1024):.1f} MiB): {output.name}") + + +if __name__ == "__main__": + main() diff --git a/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/scripts/setup.sh b/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/scripts/setup.sh new file mode 100755 index 0000000000..aaf53b07f5 --- /dev/null +++ b/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/scripts/setup.sh @@ -0,0 +1,92 @@ +#!/usr/bin/env bash +set -euo pipefail + +if (( $# != 0 )); then + echo "Usage: setup.sh" >&2 + exit 2 +fi + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +LP_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)" +WORK_ROOT="${SMOLVLA_LP_WORK_ROOT:-$LP_ROOT/work}" +PYTHON_BIN="${SMOLVLA_LP_PYTHON:-python3}" +LEROBOT_ROOT="$WORK_ROOT/lerobot" +VENV_ROOT="$WORK_ROOT/venv" +ARTIFACT_ROOT="$WORK_ROOT/artifacts" + +LEROBOT_REVISION="30da8e687a6dfc617fcd94afc367ac7071c376ce" +MODEL_REPO="HuggingFaceVLA/smolvla_libero" +MODEL_REVISION="6721902bc4d61e50a3bfdb11dfb4cb626f05d102" +BASE_MODEL_REPO="HuggingFaceTB/SmolVLM2-500M-Instruct" +BASE_MODEL_REVISION="7b375e1b73b11138ff12fe22c8f2822d8fe03467" + +if ! "$PYTHON_BIN" -c 'import sys; raise SystemExit(sys.version_info[:2] != (3, 12))'; then + echo "Python 3.12 is required." >&2 + exit 1 +fi + +mkdir -p "$ARTIFACT_ROOT" "$WORK_ROOT/cache/huggingface" \ + "$WORK_ROOT/cache/pip" "$WORK_ROOT/tmp" + +export HF_HOME="${HF_HOME:-$WORK_ROOT/cache/huggingface}" +export PIP_CACHE_DIR="${PIP_CACHE_DIR:-$WORK_ROOT/cache/pip}" +export TMPDIR="${TMPDIR:-$WORK_ROOT/tmp}" + +if [[ ! -d "$LEROBOT_ROOT/.git" ]]; then + git clone https://github.com/huggingface/lerobot.git "$LEROBOT_ROOT" +fi +git -C "$LEROBOT_ROOT" fetch origin "$LEROBOT_REVISION" +git -C "$LEROBOT_ROOT" checkout --detach "$LEROBOT_REVISION" +if [[ -n "$(git -C "$LEROBOT_ROOT" status --porcelain)" ]]; then + echo "LeRobot worktree has local changes: $LEROBOT_ROOT" >&2 + exit 1 +fi + +if [[ ! -x "$VENV_ROOT/bin/python" ]]; then + "$PYTHON_BIN" -m venv "$VENV_ROOT" +fi +if ! "$VENV_ROOT/bin/python" -c 'import sys; raise SystemExit(sys.version_info[:2] != (3, 12))'; then + echo "Python 3.12 is required in the virtual environment: $VENV_ROOT" >&2 + exit 1 +fi + +"$VENV_ROOT/bin/python" -m pip install --upgrade \ + pip wheel 'setuptools>=71.0.0,<81.0.0' +"$VENV_ROOT/bin/python" -m pip install \ + --index-url https://download.pytorch.org/whl/cpu \ + 'torch==2.11.0+cpu' 'torchvision==0.26.0+cpu' +"$VENV_ROOT/bin/python" -m pip install \ + 'torchao==0.18.0' 'transformers==5.5.4' \ + 'onnx==1.22.0' 'onnxruntime==1.29.0' \ + 'matplotlib>=3.10.3,<4.0.0' +"$VENV_ROOT/bin/python" -m pip install -e "$LEROBOT_ROOT[smolvla]" + +"$VENV_ROOT/bin/hf" download "$MODEL_REPO" \ + --revision "$MODEL_REVISION" \ + --exclude 'onnx/**' \ + --local-dir "$ARTIFACT_ROOT/smolvla_libero" +"$VENV_ROOT/bin/hf" download "$BASE_MODEL_REPO" \ + --revision "$BASE_MODEL_REVISION" \ + --exclude 'onnx/**' \ + --local-dir "$ARTIFACT_ROOT/smolvlm_base" + +"$VENV_ROOT/bin/python" -m pip list --format=freeze > "$WORK_ROOT/environment.freeze.txt" +"$VENV_ROOT/bin/python" - "$WORK_ROOT/revisions.json" <<'PY' +import json +from pathlib import Path +import sys + +payload = { + "lerobot": "30da8e687a6dfc617fcd94afc367ac7071c376ce", + "model_repo": "HuggingFaceVLA/smolvla_libero", + "model_revision": "6721902bc4d61e50a3bfdb11dfb4cb626f05d102", + "base_model_repo": "HuggingFaceTB/SmolVLM2-500M-Instruct", + "base_model_revision": "7b375e1b73b11138ff12fe22c8f2822d8fe03467", +} +Path(sys.argv[1]).write_text( + json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8" +) +PY + +echo "Environment: $VENV_ROOT" +echo "Model: $ARTIFACT_ROOT/smolvla_libero" diff --git a/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/scripts/workspace.py b/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/scripts/workspace.py new file mode 100644 index 0000000000..6fb6066d4e --- /dev/null +++ b/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/scripts/workspace.py @@ -0,0 +1,47 @@ +#!/usr/bin/env python3 +"""Configure caches owned by this standalone Learning Path.""" + +from __future__ import annotations + +import os +from pathlib import Path + + +def configure_workspace(work_root: Path | None = None) -> Path: + """Set reproducible cache defaults and return the resolved work root.""" + + learning_path_root = Path(__file__).resolve().parent.parent + if work_root is None: + work_root = Path( + os.environ.get("SMOLVLA_LP_WORK_ROOT", learning_path_root / "work") + ) + work_root = work_root.resolve() + os.environ.setdefault("HF_HOME", str(work_root / "cache/huggingface")) + os.environ.setdefault("XDG_CACHE_HOME", str(work_root / "cache/xdg")) + os.environ.setdefault("TMPDIR", str(work_root / "tmp")) + return work_root + + +def base_model_path(work_root: Path | None = None) -> Path: + """Return the verified local base-model path created by setup.sh.""" + + root = configure_workspace(work_root) + base = root / "artifacts/smolvlm_base" + required = (base / "config.json", base / "model.safetensors", base / "tokenizer.json") + missing = [str(path) for path in required if not path.is_file()] + if missing: + raise FileNotFoundError("Missing pinned SmolVLM2 base files:\n" + "\n".join(missing)) + return base + + +def load_smolvla_policy(checkpoint: Path): + """Load SmolVLA while binding its implicit backbone to the pinned local snapshot.""" + + from lerobot.configs.policies import PreTrainedConfig + from lerobot.policies.smolvla.modeling_smolvla import SmolVLAPolicy + + checkpoint = checkpoint.resolve() + config = PreTrainedConfig.from_pretrained(checkpoint) + config.device = "cpu" + config.vlm_model_name = str(base_model_path()) + return SmolVLAPolicy.from_pretrained(checkpoint, config=config) diff --git a/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/set-up-environment.md b/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/set-up-environment.md new file mode 100644 index 0000000000..734e4dcbec --- /dev/null +++ b/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/set-up-environment.md @@ -0,0 +1,64 @@ +--- +title: Set up the SmolVLA environment +description: Create a Python environment and download the pinned SmolVLA model and LeRobot source. +weight: 2 +layout: learningpathall +--- + +## Get the companion files + +Clone the Arm Learning Paths repository and open this Learning Path directory: + +```bash +git clone https://github.com/ArmDeveloperEcosystem/arm-learning-paths.git +cd arm-learning-paths/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion +``` + +## Check the system requirements + +This Learning Path runs the exported models on an aarch64 Linux CPU. Review the +processor, Python version, and available space on the filesystem where you will +keep the project: + +```bash +lscpu +python3 --version +df -h . +``` + +Use Python 3.12. Make sure the filesystem can hold the model data, checkpoint +weights, Python environment, and generated FP32 and INT4 ONNX models. + +## Create the environment + +Run the setup script: + +```bash +bash scripts/setup.sh +``` + +The script creates `work/venv`, checks out the pinned LeRobot source, installs +the conversion and runtime dependencies, and downloads the SmolVLA policy and +its SmolVLM2 dependency. It records the installed Python packages in +`work/environment.freeze.txt` and the source and model revisions in +`work/revisions.json`. + +## Verify the downloaded assets + +Check the downloaded files and revisions: + +```bash +work/venv/bin/python scripts/check_assets.py +``` + +The expected output ends with: + +```output +PASS: public policy, base model, LeRobot source, and environment are ready +``` + +## What you've accomplished and what's next + +You have prepared an Arm Linux environment with the pinned SmolVLA checkpoint, +source, and Python dependencies. Next, you will export SmolVLA as an FP32 ONNX +model and validate it with ONNX Runtime. diff --git a/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/smolvla-action-comparison.png b/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/smolvla-action-comparison.png new file mode 100644 index 0000000000..48610bb262 Binary files /dev/null and b/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/smolvla-action-comparison.png differ diff --git a/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/smolvla.png b/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/smolvla.png new file mode 100644 index 0000000000..6e3bcb8193 Binary files /dev/null and b/content/learning-paths/embedded-and-microcontrollers/smolvla-onnx-conversion/smolvla.png differ