Model Deployment and Inference Optimization Questions

Situation: I need a repeatable CI test plan that verifies a TensorFlow SavedModel -> ONNX -> ONNX Runtime roundtrip produces numerically equivalent outputs within acceptable tolerances.

Test plan (high-level steps)

1. Export & convert pipeline

• Scripted steps: (a) export SavedModel, (b) convert with tf2onnx/onnx-tf, (c) run ONNX Runtime inference. Fix RNG seeds and TF/ONNX runtime versions.

2. Test data selection

• Unit tests: small hand-crafted vectors that exercise edge cases (zeros, ones, large/small magnitudes, negative, inf/nan).
• Functional tests: random inputs with fixed seeds across distributions (uniform, normal, skewed).
• Coverage tests: inputs that trigger different ops, dynamic shapes, batch sizes, and quantized/dtype variants.
• Real-data smoke test: 50–200 real samples from production-ish dataset.

3. Tolerance rules & metrics

• Per-output checks:
  • Exact equality for integer outputs.
  • For floating types, use combined metrics:
    • max_abs = max(|y_tf - y_onnx|)
    • rms = sqrt(mean((y_tf - y_onnx)^2))
    • cosine_sim for embeddings/vectors.
  • Default thresholds (float32): rtol=1e-5, atol=1e-6; practical thresholds: max_abs < 1e-4 or rms < 1e-6. For fp16 or quantized: relax (rtol=1e-2, atol=1e-3).
• Special checks: NaN/Inf parity (fail if TF has finite and ONNX has NaN/Inf or vice versa).
• Relative per-output scaling: normalize by max(|y_tf|, epsilon) when outputs span orders of magnitude.

4. Pass/fail rules

• Per-test: pass if metrics under thresholds and NaN/Inf parity holds.
• Aggregate: allow up to a tiny percentage (e.g., 1–2%) of samples to exceed soft thresholds for flaky ops; failing tests trigger investigation.

5. Automation & CI hooks

• Integrate into CI pipeline (GitHub Actions / Jenkins):
  • Matrix run across runtime versions, hardware (CPU/GPU), and dtypes.
  • Store artifacts: SavedModel, ONNX model, test inputs/outputs, diff reports, and serialized failure cases.
  • Auto-generate human-readable report with metric summaries and example failing cases (show inputs, TF vs ONNX outputs, diffs).
  • Alerting: fail the PR on hard failures; for soft failures, open a ticket with attached artifacts.
  • Regression baselines: keep golden outputs and only allow changes via approved updates.

6. Additional practices

• Add randomized fuzzing tests periodically (nightly).
• Maintain converter-version compatibility tests.
• Add model-level unit tests for deterministic ops and stochastic ops (ensure seeds or compare distributions).

This plan provides deterministic, reproducible checks, clear numeric criteria per dtype, and CI automation to catch regressions early while producing helpful artifacts for debugging.

Requirements & constraints:

• Functional: stateless classification, 1,000,000 RPS peak, p95 latency ≤ 10 ms globally.
• Non-functional: high availability, cost-effective, regional traffic patterns, model updates with low churn.

High-level architecture:

1. Global CDN (content delivery network) + Edge Cache for responses and lightweight models
2. Regional serving clusters (multiple zones per region) with provisioned pools + autoscaling
3. Load balancer + regional gateway that routes to warm inference nodes
4. Optional serverless endpoints for spiky, low-throughput regions
5. Central model registry and CI/CD + canary rollout

Key components & responsibilities:

• CDN/Edge (CloudFront/Cloudflare Workers, Fastly): cache identical requests/responses for TTL; host tiny quantized model variants (e.g., 8-bit) for ultra-low-latency on edge where feasible.
• Regional Gateways (anycast DNS + regional LB): terminate TLS, do auth, rate-limit, and VPC-proxy to regional pools.
• Regional Serving Pools: provisioned instances (k8s or VM scale sets) with warmed containers running optimized inference runtimes (TorchScript/TF-TRT, ONNX Runtime). Use CPU for small models, GPU/TPU or inference accelerators for heavy models.
• Sharding/Partitioning: shard by model version + request type; for extremely large models use model-parallel inference (Tensor Parallel) or offload to specialized accelerators.
• Autoscaling: maintain a baseline of provisioned warm capacity to meet p95 SLA; scale horizontally with predictive scaling using traffic forecasting and scale-in cooldown to avoid cold starts.
• Serverless: use for <5% of unpredictable bursts; accept higher cold-start latency and cost.

Data flow:
Client -> Anycast DNS -> Edge CDN (cache hit? return) -> Regional LB -> Inference pool -> Response
If model update: Atomic swap from model registry + health checks + slow rollout to prevent tail-risk.

Performance & meeting p95=10ms:

• Minimize network hops: use anycast + regional endpoints so RTT is <5ms in most regions.
• Keep inference time ≤ 5ms target: quantize, prune, use batch=1 optimized kernels, use pinned threads and CPU vectorization or small GPUs with low queuing.
• Warm pools sized via capacity planning: e.g., if single optimized instance can handle 500 rps at p95, need 2000 instances globally; distribute regionally by traffic.

Cost-performance tradeoffs:

• Provisioned instances increase cost but ensure strict latency; serverless reduces ops but higher per-inference cost and cold starts.
• Edge-hosted tiny model reduces network + latency and CDN egress cost, but increases complexity and model fragmentation.
• GPU/accelerator instances increase throughput per node (fewer machines) but higher fixed cost; use mixed fleet: CPU for predictable baseline, GPUs for peak and heavy models.

Reliability, monitoring & operations:

• Observability: per-region SLO (service-level objective) telemetry, p95/p99 latency, queue depth, CPU/GPU utilization, error rates.
• Circuit breakers, retry with jitter, backpressure, and graceful degradation (fall back to lightweight model or cached response).
• Security and compliance: TLS, model signing, audit logs.

Trade-offs & rationale:

• Prioritize provisioned warm capacity to hit strict p95 SLA; add edge caching & tiny models to reduce load/cost.
• Use autoscaling + predictive forecasting to avoid overprovisioning.
• Favor simpler sharding (by region/version) unless model size forces model-parallel setups.

This design balances latency, throughput, and cost by combining CDN edge caching, regional warm pools for strict SLAs, and selective use of serverless/GPU where cost-effective.

Performance & numerical considerations

• BLAS/LAPACK: use vendor-optimized BLAS (Intel MKL, OpenBLAS) on CPU and cuBLAS/cuSOLVER on GPU. These provide tuned kernels and stable algorithms for GEMM, SVD, etc.

Memory layout

• Match library conventions: BLAS/LAPACK expect column-major. If your C++ uses row-major (std::vector/row-major Eigen), either transpose views or use libraries that support row-major. Mismatched layouts force copies and hurt performance.

Cache blocking

• Implement/block GEMM in tiles sized to L1/L2 cache (e.g., blocking M,N,K) or rely on BLAS which already does this. For custom inner loops, ensure contiguous memory accesses and avoid stride >1 in hot loops.

Numeric precision choices

• Use float32 for most NN workloads; use float64 for reductions/accumulators when summing many small terms. On GPU, use tensor cores (mixed-precision) with FP32 accumulation to benefit throughput without catastrophic numerical loss. Test sensitivity to precision and fallback to FP32/FP64 where needed.

Interoperability with Python

• Expose C++ via pybind11 or C-API. Use NumPy arrays with proper dtype and memory-order flags; avoid implicit copies by accepting contiguous arrays or documenting needed layout. For GPU work, use CUDA/cuBLAS and provide PyTorch/TensorFlow custom ops or DLPack to transfer tensors without copies.

Guidelines

• Measure performance and numerical error across precisions. Keep accumulators in higher precision. Align memory to 64 bytes, pad to avoid false sharing. Use unit tests against reference BLAS results and property tests (associativity bounds) to catch regressions.

These practices yield high-performance, numerically predictable inner loops compatible with Python ML stacks.

Also covers (folded from merged near-duplicates): a1ed479c folds x86 SIMD/AVX and MKL-DNN/oneDNN specifics into the same low-level CPU story. Also folds baff70b6 (memory-bandwidth-limited inference optimization: fusion, tiling, precision reduction).

Below are practical techniques to reduce inference memory footprint, with what they do, typical memory vs accuracy trade-offs, and where they’re most applicable.

1. Weight quantization

• What: Reduce numerical precision (e.g., FP32 → FP16/INT8/INT4) for weights and activations; can be post-training or quant-aware training.
• Trade-off: Big memory (and compute) savings (×2–×8). INT8 typically has negligible accuracy loss for many conv/CNNs; INT4/ternary risks larger accuracy drop unless retrained or calibrated.
• Applies: Edge (highly valuable) and server (for throughput/latency); use server for large-batch fast inference.

2. Weight sharing (codebook / k-means / product quantization)

• What: Replace many unique weights with indices into a small codebook.
• Trade-off: Good compression (often ×4–×16) with modest accuracy loss if codebook well-designed; overhead for lookup tables.
• Applies: Edge and bandwidth-constrained deployment. More complex on-device decoding than simple quantization.

3. Pruning (unstructured and structured)

• What: Remove parameters - unstructured zeroing vs structured (filter/channel) removal.
• Trade-off: Unstructured pruning yields high sparsity but needs sparse kernels to realize memory/compute gains; structured pruning gives direct memory/latency benefits with usually greater accuracy loss at high sparsity.
• Applies: Edge when target supports sparse/structured kernels; server when you can exploit sparse-matrix libraries or fine-grained acceleration.

4. Operator streaming (layer-by-layer or chunked execution)

• What: Stream inputs and intermediate tensors, recompute cheaply when needed, or process slices to avoid holding entire activations.
• Trade-off: Lowers peak activation memory at cost of extra compute or increased latency due to recomputation/IO.
• Applies: Edge devices with tiny RAM or constrained GPU memory; also useful on servers to run huge models without extra GPU memory.

5. Parameter offloading (host/flash <-> accelerator)

• What: Keep most weights in CPU RAM or flash and stream to accelerator on-demand.
• Trade-off: Saves accelerator memory but increases latency and bandwidth usage; risk of stalls.
• Applies: Edge with external flash or servers with NVMe where model exceeds device memory; good when throughput tolerant of staging overhead.

6. Model partitioning (pipeline or device-split)

• What: Split model across devices/hosts - pipeline stages or split layers across CPU/GPU.
• Trade-off: Reduces per-device memory but introduces communication overhead and potential pipeline bubble latency.
• Applies: Server clusters (multi-GPU inference), edge+cloud hybrid (part on device, heavy layers on cloud).

7. Model architecture changes (smaller nets, efficient blocks, distillation)

• What: Design or replace with efficient architectures (MobileNet, EfficientNet, transformers with sparse attention) or distill large models into compact students.
• Trade-off: Typically best accuracy-vs-memory trade-offs; requires retraining and design effort.
• Applies: Primarily edge where budget tight; also server when maximizing throughput per memory budget.

General guidance:

• Combine techniques: quantization + pruning + architecture changes often yields the best practical footprint vs accuracy.
• Measure end-to-end metrics (latency, peak memory, bandwidth) on target hardware - some techniques look good on paper but need hardware support (INT8 kernels, sparse ops).
• Prefer accuracy-preserving-first: try post-training quantization and structured pruning; use QAT (quantization-aware training) or distillation if accuracy drops.

Start with the principle: make every dependency and build step explicit, immutable, and recorded so any environment can be rebuilt identically.

Practices and tools

• Dependency pinning & lockfiles:

  • Python packages: pin exact versions and hashes. Use pip-compile/pip-tools or Poetry to produce a deterministic lockfile (poetry.lock or requirements.txt with --hash). For pip, check in requirements.txt with hashes produced by pip-compile; for reproducible installs use pip install --no-deps --require-hashes.
  • Conda: use environment.yml for high-level specs but generate conda-lock (conda-lock) to produce platform-specific lockfiles (.conda-lock) with exact package builds and channels.
  • Strategy: prefer conda for system-level and compiled deps (MKL, CUDA, libsndfile) and pip/poetry for pure-Python packages; always generate lockfiles for both and commit them.

• Conda vs pip:

  • Use conda when you need binary libraries, controlled channels (conda-forge), or cross-platform compiled artifacts. Use pip/poetry for Python packaging/virtualenvs where binaries aren’t required.
  • Hybrid: create a minimal conda env (python + system libs) then pip install locked requirements inside it.

• Deterministic container builds:

  • Pin base image by digest (ubuntu@sha256:..., python:3.10-slim@sha256:...).
  • Use multi-stage builds, avoid apt-get install without version pins; pin apt packages where possible and add deterministic ordering.
  • Build with BuildKit, set --no-cache to avoid implicit layers; supply lockfiles into the image and use pip install --require-hashes or conda-lock install to ensure exact binaries.
  • Embed metadata: LABELs with git commit, lockfile hash, build timestamp.
  • Push images to an immutable registry and reference by digest in deployment manifests.

• Artifact stores & provenance:

  • Store built artifacts and models in an artifact registry (Docker registry, S3/GS with versioned prefixes, Artifactory) and record their digests.
  • Use an ML model registry (MLflow, SageMaker Model Registry, or DVC + remote storage) to store model binaries, training data hashes, and evaluation metrics.
  • Record build provenance in CI: git sha, lockfile checksum, container digest, trained model id.

• Hermetic builds:

  • Use hermetic build tools when full isolation is required: Bazel, Nix, or fully self-contained Docker images created from scratch. These ensure all inputs are declared and reproducible.
  • For reproducible ML pipelines, consider using Nix or conda-lock + pip hash inside containers to make environment builds hermetic.

Operationalize in CI/CD

• CI pipeline steps: run tests, generate lockfiles, build container with pinned base/digests, run a reproducibility check (rebuild and compare artifact digests), push image/artifacts and record metadata.
• Automate environment checks: small smoke tests that assert package versions, CUDA/cuDNN versions, GPU availability, and RNG seeds.

Extra considerations

• Seed randomness and record seed for model training. Log hardware (GPU model, driver/cuda versions).
• Monitor drift: periodically re-run build and tests to detect upstream changes.
• Document reproducibility steps in README and make a single reproducible entrypoint (Makefile or scripts/ci) to rebuild environments.

Example minimal workflow

1. Developer runs conda env + pip install from conda-lock + requirements.lock.
2. CI uses same lockfiles, builds Docker image pinned to base digest, installs from lockfiles, runs tests.
3. CI pushes immutable image digest and model artifact to registry and records them in MLflow/DVC.

This combination of pinned versions, lockfiles, hermetic container builds, artifact registries, and automated CI ensures "works on my machine" failures are eliminated.