Knowledge 03 — Model Conversion & Compilers

The JD bullet: "Model conversion workflows," "Runtime integration and optimization," "Integrate ML models … from PyTorch, TensorFlow, and ONNX," "Contribute to the Cloud AI GitHub repo." This module is about getting a model out of a training framework and onto an accelerator — the part of the job where things break, and where a principal is the person who knows why and how to fix it. It's the least glamorous and most billable skill in the role.

Table of Contents

• 1. Why conversion exists: the framework↔hardware gap
• 2. The graph: what a model really is to a compiler
• 3. ONNX: the interchange format
• 4. The compilation pipeline, stage by stage
• 5. Operator fusion (the big optimization)
• 6. The three things that always break
• 7. The compilers you must know
• 8. torch.compile and the PyTorch-native path
• 9. A debugging methodology for a failed port
• 10. Validating a port: numerical correctness
• 11. References

1. Why conversion exists: the framework↔hardware gap

PyTorch and TensorFlow are training frameworks: flexible, eager, Python-in-the-loop, full of dynamic control flow. That flexibility is murder for inference performance — every op is a separate kernel launch (Knowledge 00 §8), Python sits in the hot path, and nothing is specialized to your specific shapes or hardware.

An inference compiler (TensorRT, OpenVINO, XLA, Qualcomm AIC, AWS Neuron) takes the model's computation, freezes it into a static graph, and aggressively optimizes it for one specific target: fusing ops, picking the best kernel for each shape, choosing memory layouts, applying quantization, and emitting a compact "engine" or "executable" that runs with no Python and minimal overhead.

PyTorch model (eager, flexible, slow)
      │  export (trace/script/dynamo)
ONNX graph  ── or ──  framework-native IR (torch.export, StableHLO)
      │  compile (fuse, quantize, select kernels, plan memory)
Optimized engine for target  (TensorRT plan / qaic binary / Neuron NEFF / OpenVINO IR / TPU XLA exe)
      │  load + run
Accelerator runtime  →  fast, fixed, hardware-specific inference

The core tension you manage: training frameworks optimize for expressiveness and iteration; inference compilers optimize for one fixed graph on one chip. Conversion is crossing that chasm, and every place the two disagree is a bug you'll triage.

2. The graph: what a model really is to a compiler

To a compiler, a model is a computational graph (DAG): nodes are operators (MatMul, Add, Softmax, LayerNorm, Conv, Attention), edges are tensors. Each node has attributes (shapes, dtypes, parameters). The compiler's job is to transform this graph into something that runs fast on the target while computing the same result.

Two ways frameworks produce the graph:

• Tracing (torch.jit.trace, torch.onnx.export default): run the model with example inputs, record the ops that fire. Fast and simple, but loses data-dependent control flow — an if that didn't fire on the trace input is gone, and the trace bakes in the example's shapes.
• Scripting / graph capture (torch.jit.script, torch.export/Dynamo, TF @tf.function): statically analyze the code to capture control flow. Handles dynamic behavior better but is pickier about what Python it accepts.

First debugging instinct when a port misbehaves on some inputs but not others: a trace captured only one control-flow path. Re-export with proper graph capture or refactor the dynamic part out of the model.

3. ONNX: the interchange format

ONNX (Open Neural Network Exchange) is the lingua franca — a vendor-neutral, serialized graph format (protobuf) with a standardized operator set ("opset"). It's how a PyTorch model reaches a non-PyTorch runtime (TensorRT, OpenVINO, Qualcomm AIC, ONNX Runtime). The JD names ONNX explicitly because it's the on-ramp to most accelerators.

What you must know:

• Export: torch.onnx.export(model, sample_input, "model.onnx", opset_version=17, dynamic_axes=...). The newer torch.onnx.dynamo_export uses graph capture and handles more models.
• Opset version: the ONNX standard evolves; newer ops (e.g., fused attention, newer normalizations) need newer opsets, but a given runtime supports up to some opset. Opset mismatch is a top-3 failure (§6).
• ONNX Runtime (ORT): a reference runtime that itself has hardware execution providers (CUDA, TensorRT, OpenVINO, QNN/Qualcomm, CoreML). Useful as a portable baseline and a debugging oracle.
• Inspection: Netron (visualize the graph), onnx.checker (validate), onnxsim/onnx-simplifier (fold constants, clean up), onnxruntime (run and compare outputs).
• ONNX optimizations: constant folding, shape inference, redundant-node elimination — often run before handing to the vendor compiler.

Workflow reality: PyTorch → ONNX → vendor compiler is the dominant path because vendors invest in ONNX front-ends rather than parsing PyTorch directly. Mastering ONNX export + Netron inspection + ORT comparison is the portable conversion skill across every accelerator.

4. The compilation pipeline, stage by stage

Whatever the vendor, the compiler does roughly these stages. Knowing them lets you reason about any toolchain (TensorRT, AIC, Neuron, XLA):

1. Import / parse the graph (from ONNX or framework IR).
2. Shape inference & freezing: resolve tensor shapes; decide static vs dynamic dimensions; fold constants.
3. Graph rewriting / canonicalization: normalize equivalent patterns, remove no-ops, fold BatchNorm into preceding Conv.
4. Operator fusion (§5): merge chains of ops into single kernels.
5. Quantization / precision assignment (Knowledge 02): place quant/dequant nodes, assign per-layer precision, calibrate.
6. Kernel selection / autotuning: for each op + shape, pick the fastest kernel from a library or autotune by timing candidates (TensorRT does this — it's why building an engine is slow and hardware-specific).
7. Memory planning: assign tensor buffers, reuse memory across non-overlapping lifetimes, plan the workspace.
8. Code/engine emission: produce the target binary (TensorRT .plan, Neuron NEFF, qaic binary, XLA executable).

Two consequences for the job:

• A compiled engine is specialized to the GPU SKU, precision, and (often) batch/shape it was built for. Build on an H100, it may not load on an A100. Build on the deployment hardware. This catches teams constantly.
• Build time vs run time: autotuning makes builds slow (minutes) but runs fast. Cache engines; don't rebuild per request.

5. Operator fusion (the big optimization)

The headline compiler optimization, and a direct application of the memory wall. A sequence like y = dropout(gelu(x @ W + b)) naively launches 4 kernels, each reading its input from HBM and writing its output back to HBM — 4 round-trips for data that could have stayed on-chip.

Fusion merges them into one kernel: read x and W once, do matmul → add → gelu → dropout entirely in registers/SRAM, write y once. The intermediate tensors never touch HBM. For memory-bound elementwise chains this is often a 2–5× speedup, and it also removes kernel-launch overhead.

Classic fusions you'll see named in profiles:

• MatMul + bias + activation (the MLP).
• LayerNorm/RMSNorm fusion (the reduction + normalize + scale in one pass).
• Attention fusion → FlashAttention (Knowledge 01 §6) is the ultimate fused attention kernel.
• Fused dequant + matmul for quantized models (dequantize weights inside the matmul kernel so you never store FP16 weights in HBM).

Whiteboard explanation of fusion: "Each unfused op is a separate trip to HBM. Fusion keeps intermediates in fast on-chip memory so we touch HBM once for the whole chain. It's the same arithmetic, far less data movement — the recurring theme of inference optimization." (Same lesson as FlashAttention, generalized.)

6. The three things that always break

When you port a real model, expect these. Being the calm person who recognizes them on sight is the JD's "issue triage and technical escalation" value.

(a) Unsupported / custom operators

The model uses an op the target compiler doesn't implement (a novel activation, a custom CUDA kernel, a new attention variant, nonzero, complex indexing). Symptoms: export fails, or compile fails with "unsupported op X." Fixes, in order of preference: upgrade opset/compiler; rewrite the op using supported primitives; decompose it; mark it to run on CPU/fallback (ONNX Runtime can split execution); or write a custom plugin/kernel (TensorRT plugin, custom op) — the last resort, real C++/CUDA work.

(b) Dynamic shapes

Training is happy with arbitrary batch sizes and sequence lengths; compilers love static shapes (so they can autotune and plan memory). Symptoms: engine built for shape [1, 512] fails or runs slow on [8, 2048]. Fixes: declare dynamic_axes in ONNX export; use optimization profiles (TensorRT lets you specify min/opt/max shapes and builds for the range); bucket inputs into a few fixed shapes and pad; for LLMs, this is the hard part — variable sequence length and KV cache growth are why LLM-specific compilers (TensorRT-LLM) exist rather than plain TensorRT.

(c) Numerical divergence

The ported model runs but gives slightly (or wildly) different outputs. Causes: FP16/FP8 precision differences, a fused kernel computing in a different order, a quantization scale, a subtly different op semantics (e.g., a different LayerNorm epsilon or rounding mode), or layout/transpose bugs. Fixes: compare layer-by-layer against the FP16 reference (§10); identify where divergence starts; decide if it's acceptable (tiny FP drift) or a real bug (wrong op).

The principal's calm: "Ports break in exactly three ways — an op the compiler doesn't know, a shape it didn't expect, or numbers that drifted. Let's find which one." That framing turns a panicked war room into a checklist.

7. The compilers you must know

JD2 reality: a Qualcomm/accelerator solutions role lives mostly in ONNX → vendor compiler (AIC/Neuron/TensorRT) + vendor quantizer (AIMET/Model Optimizer). The generic skill — export clean ONNX, inspect it, feed the vendor compiler, debug the three failure modes, validate numerics — transfers across all of them. Learn the pipeline, not one vendor's flags.

8. torch.compile and the PyTorch-native path

Since PyTorch 2.0, torch.compile(model) captures the graph with TorchDynamo, optimizes it with TorchInductor (which generates fused Triton kernels for GPU), and falls back to eager for un-capturable code ("graph breaks"). It's the lowest-friction way to get fusion + reduced overhead without leaving PyTorch — no ONNX export, no separate engine.

When to use it vs a full compiler:

• torch.compile: fast iteration, research-to-prod, when you want most of the win with minimal porting risk. Watch for graph breaks (Python in the hot path) that kill the speedup — TORCH_LOGS=graph_breaks to find them.
• TensorRT/vendor compiler: when you need maximum performance, INT8/FP8 with autotuning, deployment without a Python runtime, or a non-NVIDIA target.

Many production stacks use both: torch.compile for flexible parts, TensorRT-LLM/vLLM (which internally uses compiled kernels) for the LLM hot path.

9. A debugging methodology for a failed port

A repeatable procedure (this is the "issue triage and escalation" the JD asks for):

1. Reproduce minimally: smallest model/op that fails. Isolate one layer if possible.
2. Read the actual error, not the stack trace summary — compilers name the offending op/shape.
3. Classify into the three failure modes: unsupported op? dynamic shape? numerical drift?
4. Inspect the graph in Netron — is the export even correct? Often the bug is upstream in the ONNX, not the compiler.
5. Bisect precision: does it work in FP32? FP16? Fails only in INT8? → it's a quantization/calibration issue, not a conversion issue.
6. Compare against ORT as an oracle: if ONNX Runtime gives the right answer but TensorRT doesn't, the bug is in the TensorRT path, not your ONNX.
7. Escalate with a clean repro: when you file a vendor GitHub issue (a JD duty), attach the minimal model, the exact versions, and the diff vs the reference. This is what "contribute to the developer docs and repo" looks like in practice.

10. Validating a port: numerical correctness

Never declare a port done because it "runs." Prove it computes the right thing:

• End-to-end output comparison: same inputs through FP16 reference and the ported engine; compare logits/outputs with a tolerance (np.allclose(rtol=1e-2, atol=1e-2) for FP16; looser for FP8/INT8). Report max abs error and cosine similarity of logits.
• Layer-by-layer comparison when end-to-end diverges: dump intermediate activations from both, find the first layer where they part — that localizes the bug.
• Task-level eval (Knowledge 02 §11): the port plus its quantization must hold task accuracy, not just logit closeness. A 1e-2 logit error can still shift argmax on hard tokens.
• Throughput/latency validation: confirm the port is actually faster (sometimes a bad fusion or fallback-to-CPU op makes it slower — measure, don't assume).

Deliverable framing: a finished port comes with a one-page report — "ported to TensorRT FP8, max logit error 7e-3, cosine 0.9998, MMLU 78.1 vs 78.4 FP16, 2.3× throughput, 0.51× memory." That report is the customer-ready documentation the JD wants.

11. References

• ONNX specification and ONNX Runtime docs (operators, opsets, execution providers).
• NVIDIA TensorRT Developer Guide and TensorRT-LLM docs (optimization profiles, plugins, FP8).
• Qualcomm Cloud AI SDK / AIC compiler and AIMET docs; AWS Neuron SDK docs; OpenVINO docs; XLA/StableHLO docs.
• TorchDynamo / TorchInductor / torch.export documentation (PyTorch 2.x).
• Netron (graph viewer), onnx-simplifier, Triton (OpenAI) kernel language.
• Chen et al., "TVM: An Automated End-to-End Optimizing Compiler for Deep Learning" (OSDI 2018) — the compiler-stack mental model.

➡ Next: Knowledge 04 — Serving Frameworks — wrapping the optimized model in a high-throughput server.