npu-torchair-infer

NPU TorchAir Inference & Migration (Main Process)

Overview

Bring an arbitrary HuggingFace model up on Ascend NPU torchair graph mode, verify it numerically against CPU/NPU-eager, and measure its speedup. The bundled scripts are model-agnostic (config-driven inputs, recursive output comparison), so the same flow applies to the next model with only a changed --model_name_or_path.

Bundled assets:

• scripts/torch_air_infer.py — minimal single-backend inference + latency stats.
• scripts/benchmark.py — 3-backend (torchair / npu_eager / cpu) accuracy + perf → JSON + Markdown.
• scripts/run_benchmark.sh — generic driver: prepare model → sanity infer → benchmark.

The one rule you cannot break

torch_npu MUST be imported before torchair, or graph mode does not engage.

import torch
import torch_npu        # before torchair
import torchair
config = torchair.CompilerConfig()
npu_backend = torchair.get_npu_backend(compiler_config=config)
model = torch.compile(model, backend=npu_backend, dynamic=False)

Details, version matrix, and CompilerConfig knobs: see references/environment.md.

Main flow (run in order)

source $ASCEND_HOME/ascend-toolkit/set_env.sh   # CANN runtime first
export HF_ENDPOINT=https://hf-mirror.com         # China-friendly mirror

# 0) Sanity on CPU (fast) — confirms inputs/outputs resolve for this model
python scripts/torch_air_infer.py --model_name_or_path <path-or-id> --device cpu --backend eager

# 1) NPU eager — confirms the model runs on NPU at all
python scripts/torch_air_infer.py --model_name_or_path <path-or-id> --device npu --backend eager --dtype float16

# 2) NPU torchair graph mode — the migration target
python scripts/torch_air_infer.py --model_name_or_path <path-or-id> --device npu --backend torchair --dtype float16

# 3) Full 3-backend accuracy + performance comparison (writes JSON + Markdown)
python scripts/benchmark.py --model_name_or_path <path-or-id> \
    --backends torchair,npu_eager,cpu --device npu \
    --dtype float16 --batch_size 1 --warmup 10 --iterations 100 \
    --output_dir benchmark_results/<model>

# strict <1e-3 parity check (fits-in-memory models): re-run step 3 with --dtype float32

Or do steps 2+3 in one shot: bash scripts/run_benchmark.sh <path-or-id> [out_subdir].

Decision tree

Step 0 (CPU) fails to build inputs?  → model needs a custom input builder (see references/methodology.md, BaseModelLoader.build_inputs).
Step 1 (NPU eager) fails?            → environment problem: re-check import order + set_env.sh (references/environment.md).
Step 2 (torchair) fails to compile?  → walk references/troubleshooting.md; fall back to npu_eager and keep going.
Step 3 accuracy not within gate?     → compare torchair vs npu_eager first (references/methodology.md): same ⇒ fp16 rounding (PASS_COSINE); differ ⇒ graph bug.
Weights gated / undownloadable?      → build a config-init random stand-in (references/methodology.md); valid for backend comparison, not task accuracy.

Accuracy & performance gates (summary)

• Golden = CPU fp32. Metrics: cosine sim, max abs diff, max rel err.
• PASS = max_abs_diff ≤ 1e-3 and cosine ≥ 0.999; PASS_COSINE = cosine passes (expected in fp16); else FAIL.
• Perf: first iteration = compile cost (reported separately); steady p50/p95/p99 over ≥100 iters with torch.npu.synchronize() before each timer stop; report throughput + peak HBM.

Full methodology, the model interface, gated-weights handling, and expected output: references/methodology.md.

References (load as needed)

• references/environment.md — import order, CANN/torch_npu/torchair/transformers compatibility matrix, pre-flight checklist, CompilerConfig.
• references/troubleshooting.md — compile-failure decision tree, known op failures (interpolate/contiguous), fallback order, bug localization.
• references/methodology.md — BaseModelLoader interface, input synthesis, accuracy/perf methodology, config-init stand-ins, expected output.

Reuse checklist for the next model

1. Step 0 on CPU — inputs/outputs resolve (fast).
2. Step 1 NPU eager works.
3. Step 2 torchair — compiles ⇒ done; else walk references/troubleshooting.md.
4. Step 3 benchmark — gate on cosine (fp16) or <1e-3 (fp32).
5. If inputs could not be inferred, add a small build_inputs override (references/methodology.md).