示例还需另装模型依赖 pip install diffusers transformers；JIT 模式每次新进程都会重新调优，要复用产物需走 AOT 流程；README 未列出具体 PyTorch/CUDA 版本对应关系，需查安装指南。
aitune
ai-dynamoApache-2.0活跃维护面向 NVIDIA GPU 的深度学习推理调优工具包，帮助把训练好的模型调优并部署到 TensorRT、ONNX Runtime 等推理后端上跑得更快。
NVIDIA AITune is an inference toolkit designed for tuning and deploying Deep Learning models with a focus on NVIDIA GPUs.
v0.6.0发布于 2026-09-16 08:54 UTC290star35fork同步于 2026-11-10 18:03 UTC
pip install --extra-index-url https://pypi.nvidia.com aitune 硬件要求按 README 说明
需要 Linux 环境、NVIDIA GPU 以及兼容的 PyTorch 与 CUDA 安装；README 未给出具体显存要求。
需要 CUDA多卡 / 张量并行
NVIDIA 出品的 PyTorch 推理调优工具，自动为各模块挑选后端并验证精度与性能
面向 NVIDIA GPU 的深度学习推理调优工具包，帮助把训练好的模型调优并部署到 TensorRT、ONNX Runtime 等推理后端上跑得更快。
核心能力
- 用统一 API 接入多种推理后端并支持自定义扩展
- 按 nn.Module 粒度评估候选后端，同一流水线可混用
- 对比原始实现做数值校验，再按吞吐或延迟策略选型
- 候选可包含 PTQ 量化、kernel 选择与 CUDA graphs
- 支持 JIT 在线调优与 AOT 生成 .ait 产物部署
容易踩的坑
内容由模型根据 README 生成，可能不完整或不准确，以官方文档为准。看原文 README
- v0.6.0AITune Release v0.6.0最新正式版2026-09-16 · 2026-09-16 08:54 UTC
AITune Release v0.6.0
AITune v0.6.0 focuses on accelerating image and video generation with multi-GPU tuning and experimental automatic kernel replacement. It brings coordinated backend selection to tensor- and context-parallel models, adds collective NVIDIA Dynamo serving, and introduces attention kernel optimization that works alongside AOT and JIT compilation.
Summary
This release extends the AITune workflow to models that an application distributes across multiple GPUs. Inspection, compilation, correctness validation, and profiling are coordinated across ranks so the model uses a consistent backend selection. Examples cover FLUX image generation and WAN 2.1 text-to-video generation with context parallelism.
The experimental kernel replacement feature profiles expensive functional calls, evaluates compatible implementations, and selects replacements that pass correctness checks and improve measured performance. Initial providers focus on attention, including PyTorch SDPA, Diffusers attention backends, SageAttention, and FlashAttention-4.
Key Dependencies
- Python
>=3.10,<3.13. - PyTorch 2.8 or newer, including supported NGC prerelease builds; version-specific extras cover PyTorch 2.8 through 2.12 with matching Torch-TensorRT and TorchAO constraints.
- TensorRT
>=10.3, TorchAO>=0.12, and NumPy>=1.26.4. - Optional NVIDIA Dynamo integration through
aitune[dynamo]. - Additional kernel providers require their corresponding runtimes. Diffusers attention providers require
diffusers>=0.35.0; Hub-backed implementations also requirekernels>=0.12.
Breaking Changes and Default Behavior
- Both AOT and JIT tuning now default to
MaxThroughputStrategy, profiling compatible candidates to select the highest-throughput backend. This can increase tuning time compared with stopping at the first successful backend. Explicit strategy and backend choices remain supported. - Custom backends now declare
BuildModeand must explicitly opt into multi-GPU execution. - The keyword argument in
enable_performance_validation(enable=...)is nowmode=.... Positional boolean calls remain supported.
Features & Improvements
Multi-GPU Tuning and Serving
- Added AOT and JIT tuning for application-managed tensor and context parallelism, preserving distributed module and tensor placement.
- Coordinated backend acceptance across ranks: a candidate must build and pass correctness checks on every rank. Profiling uses minimum throughput and maximum latency across ranks to make one consistent selection.
- Added topology-aware backend selection, including TorchInductor and Torch-TensorRT support for distributed modules.
- Isolated caches, artifacts, and telemetry by rank and synchronized inspection ordering across ranks.
- Added collective NVIDIA Dynamo serving: one endpoint receives requests, all ranks execute the distributed model, and rank 0 returns the response.
Image and Video Acceleration
- Added a WAN 2.1 example covering Ulysses and ring context parallelism, backend comparison, video generation, and collective Dynamo serving.
- Expanded multi-GPU integration guidance with FLUX image-generation and WAN video-generation examples.
- Reduced memory pressure when tuning multiple modules by storing recorded samples on disk and loading them lazily.
- Added bounded TensorRT CUDA graph caching per static optimization profile, with fallback for dynamic shapes or capture failures.
Experimental Kernel Replacement
- Added
KernelOptimizerto identify expensivetorch.nn.functionalcalls and select faster compatible implementations using representative inputs. - Added attention providers for PyTorch SDPA, Diffusers attention dispatch, SageAttention, and FlashAttention-4.
- Validate and benchmark candidates before accepting a replacement; unavailable or incompatible providers are skipped.
- Added
KernelSelectorBackendto combine kernel selection with a delegate AOT
- Python
- v0.5.0AITune Release v0.5.02026-08-04 · 2026-08-04 11:16 UTC
查看更新说明
AITune v0.5.0 adds explicit dynamic-shape configuration, latency-oriented tuning strategies, runtime performance attribution, and more robust JIT inspection and correctness validation.
Summary
This release allows AOT modules to declare explicit dynamic input shapes and validates compiled model correctness at the minimum and maximum boundaries of those shapes.
It introduces
MinLatencyStrategyfor selecting the backend with the lowest latency andLatencyBudgetStrategyfor maximizing throughput while satisfying a latency constraint.The release also improves deferred JIT tuning, enriches tuning reports with inspection details, normalizes equivalent forward-call layouts, and adds runtime attribution reports for tuned modules, untuned modules, pipeline overhead, PyTorch operators, and CUDA kernels.
Key Dependencies
- Python
>=3.10,<3.13 - PyTorch
>=2.8; version extras are provided for PyTorch 2.8 through 2.12 - Matching Torch-TensorRT versions for PyTorch 2.8 through 2.12
- TorchAO
>=0.12, with version-specific constraints provided by the PyTorch extras - TensorRT
>=10.3 - NumPy
>=1.26.4 - Optional Dynamo integration via
aitune[dynamo]
Breaking Changes
- Calls that pass the same forward parameters positionally or by keyword are now treated as the same execution graph. Existing tuned checkpoints must be regenerated.
- Forward parameter paths replace
args_*andkwargs_*identifiers in metadata reports and TensorRT optimization profiles. Custom optimization profiles must be updated to use the new paths. aitune.torch.jit.tune.deferred()now schedules tuning for the next normal forward pass instead of tuning immediately.FindMaxBatchSizeMixinConfigandPerformanceValidationMixinConfigwere replaced by a single strategy-levelProfilingConfig.
Features & Improvements
Dynamic Shapes and Correctness
- Added explicit dynamic input-shape configuration for AOT modules.
- Added correctness validation at minimum and maximum dynamic-shape boundaries.
- Normalized positional and keyword forward-call layouts into a canonical graph representation.
- Added stable forward parameter paths for nested input structures and backend optimization profiles.
Tune Strategies
- Added
MinLatencyStrategy, which profiles successful backends and selects the one with the lowest latency. - Added
LatencyBudgetStrategy, which selects the highest-throughput result that satisfies a configured latency budget. - Unified profiling, maximum-batch-size discovery, and performance-validation configuration under
ProfilingConfig.
JIT Tuning and Reporting
- Deferred-mode tuning now runs on the next normal forward pass after
aitune.torch.jit.tune.deferred()is called. - Added full JIT inspection details and configuration to deferred-mode tuning reports.
- Added basic JIT inspection information to eager-mode tuning reports.
- Improved validation of user-provided tuning configuration.
Runtime Performance Attribution
- Added the user-facing
aitune.torch.profile()API. - Added attribution for runtime spent in AITune-managed modules, untuned modules, and pipeline or framework code outside modules.
- Added PyTorch operator and CUDA kernel summaries.
- Added structured JSON-ready results, Markdown rendering, and optional Chrome trace export.
Compatibility and Documentation
- Added dependency profiles for PyTorch 2.8, 2.9, 2.10, 2.11, and 2.12.
- Refreshed installation instructions and NGC container guidance.
- Updated Dynamo examples and expanded validation across supported PyTorch versions.
Known Issues
- Existing tuned checkpoints are incompatible with the new canonical forward-parameter representation and must be regenerated.
- Custom TensorRT optimization profiles using
args_*orkwargs_*identifiers must be migrated
- Python
- v0.4.1AITune Release v0.4.12026-06-30 · 2026-06-30 18:27 UTC
查看更新说明
AITune Release v0.4.1
AITune v0.4.1 improves profiling accuracy, lowers tuning memory spikes, and cleans up performance validation configuration names.
Summary
This release focuses on profiling and validation fixes. Baseline performance now uses the recorded global batch size, throughput is computed from mean latency, and fixed-step sampling is the default after warmup.
It also reduces memory spikes in baseline profiling and Torch Inductor JIT with autocast. Optional stable-window profiling now has clearer bounds, warmer failure messages, and a less strict default CV threshold.
Key Dependencies
- Python
=3.10,3.13 - PyTorch
=2.7,2.11 - Torch-TensorRT
2,2.11 - TorchAO
=0.13,0.17 - TensorRT
=10.5 - NumPy
=2.0.0 - Optional Dynamo integration via
aitune[dynamo]
Breaking Changes
- Profiling and inspection threshold arguments were renamed to ratio-based names.
enable_validate_against_baseline(False)was replaced byenable_performance_validation(False).
Fixes & Improvements
Profiling
- Uses the recorded global batch size for baseline performance profiling.
- Computes profiling throughput from mean latency.
- Defaults profiling measurements to fixed-step sampling after warmup.
- Excludes warmup samples from optional stable-window CV checks.
- Adds max-sample bounds and clearer failure errors for stable-window profiling.
- Relaxes the default stable-window CV threshold to 10%.
- Validates ratio and positive profiling configuration values.
Memory Use
- Removes a memory spike during baseline performance measurement.
- Removes a memory spike in the Torch Inductor JIT backend when autocast is active.
Backend Behavior
- Aligns Torch-TensorRT defaults.
- Synchronizes CUDA during JIT inspection so measurements are more accurate.
Known Issues
- Hardware-specific quantization paths such as NVFP4 require compatible NVIDIA hardware and installed backend support.
- GPU functional tests require the relevant backend dependencies and hardware; unit tests remain the default local validation path.
- PyTorch 2.11 is outside the supported range for this release.
- Python
- v0.4.0AITune Release v0.4.02026-06-03 · 2026-06-03 04:13 UTC
查看更新说明
AITune Release v0.4.0
AITune v0.4.0 expands backend coverage, improves JIT tuning for complex pipelines, and adds observability needed to understand tuning decisions.
Summary
This release adds ONNXRuntime and Torch Inductor AOT backends, extends quantization support for TorchAO and TensorRT, introduces Dynamo serving support, and improves JIT tuning control for complex pipelines such as diffusion workloads.
It also adds performance validation against Torch eager baselines, richer tuning telemetry, optional hardware metrics, refreshed documentation and examples.
Key Dependencies
- Python
>=3.10,<3.13 - PyTorch
>=2.7,<2.11 - Torch-TensorRT
>2,<2.11 - TorchAO
>=0.13,<0.17 - TensorRT
>=10.5 - NumPy
>=2.0.0 - Optional Dynamo integration via
aitune[dynamo]
Breaking Changes
TorchInductorBackendwas renamed toTorchInductorJitBackend.HighestThroughputStrategywas renamed toMaxThroughputStrategy.NVTX_ENABLEwas renamed toAITUNE_NVTX_EVENTS.aitune.torch.jit_config.backendswas removed. Configure JIT backends throughaitune.torch.jit_config.strategy, for exampleFirstWinsStrategy(backends=[...]).- The previous public system-monitoring APIs were removed:
SystemMonitorsystem_resource_monitorenable_gpu_memory_logging
Features & Improvements
Backend Coverage
- Added ONNXRuntime backend support with CUDA and TensorRT execution providers.
- Added Torch Inductor AOT backend support for saved compiled artifacts.
- Added TorchAO NVFP4DQ and MXFP8DQ quantization options, including filter support and hardware validation.
- Added TensorRT NVFP4 quantization support.
- Improved dynamic-shape handling across Torch Inductor, Torch-TensorRT, and ONNX export paths.
JIT Tuning
- Added deferred JIT tuning mode for pipelines where modules run variable numbers of times before the best tuning point is known.
- Added JIT tune strategy selection through
aitune.torch.jit_config.strategy. - Added package and module-class exclusions for JIT patching.
- Added Diffusers integration hooks for pipeline compatibility.
- Improved wrapped-descendant handling so JIT tuning restores patched modules correctly.
Tune Strategies
- Backends that are correct but slower than Torch eager can now be rejected automatically.
- Added per-module speedup reporting during tuning.
- Speedup summaries are visible at the default warning log level.
Observability
- Added tuning telemetry reports covering runs, modules, graphs, backend builds, throughput, selected backends, and failures.
- Added
AITUNE_TUNING_DATA_PATHfor choosing the telemetry output path. - Added
snapshot_tuning_data()for long-running processes. - Added optional hardware metrics collection with
AITUNE_HARDWARE_METRICS=1. - Added
AITUNE_HARDWARE_METRICS_PATHfor hardware metrics output. - Added NVTX annotation support through
AITUNE_NVTX_EVENTS=1.
Deployment & Runtime
- Added Dynamo worker support for serving AITune-tuned models as Dynamo endpoints.
- Fixed checkpoint backend artifacts portability by storing relative artifact paths.
- Added disk-space checks before build and save operations.
Bug Fixes
- Fixed Torch-TensorRT AOT save failures for wrapt-decorated forwards.
- Fixed shared dynamic dimensions in Torch-TensorRT AOT.
- Fixed bounded dimension handling for Torch-TensorRT AOT export.
- Fixed Torch Inductor AOT dynamic-shape handling.
- Fixed JIT tuning state handling when wrappers are still in initial state.
- Preserved externally registered forward hooks across save and restore.
- Fixed
UserDicttraversal for inputs such astransformers.BatchEncoding. - Fixed Blackwell-specific quantization configuration handling.
- Dropped the Torch-TensorRT
enabled_precisions={float16}default so engines follow the model dtype.
Documentation
- Migrated the documentation preview to Fern.
- Added backend guides
- Python
- v0.3.0AITune Release v0.3.02026-03-15 · 2026-03-15 16:28 UTC
查看更新说明
Summary
AITune is an open-source (Apache 2.0) inference toolkit, hosted under the ai-dynamo GitHub organization and distributed via PyPI. It is designed for tuning and deploying Deep Learning models on NVIDIA GPUs, significantly improving inference speed and efficiency across various AI workloads.
Major Features & Improvements
Tuning Modes
- Just-in-Time (JIT) Tuning: Zero-code model tuning and inspection controlled through a single import or environment flag. Tunes on the very first model call using only one sample, with automatic fallback to Torch Inductor when a backend cannot compile a module.
- Ahead-of-Time (AOT) Tuning: Low-code API for explicit model inspection, backend selection, and module-level tuning. Supports forward hooks for custom pre/post-processing logic around tuned modules.
Backend Support
- TensorRT: Multi-profile engines with auto-generated and user-provided profiles, CUDA graph capture, FP16/FP8/INT8 mixed precision via TensorRT Model Optimizer, and Dynamo-based ONNX export (
torch.onnx.export(dynamo=True)) for improved graph fidelity. - TorchInductor: Added support for static and dynamic HuggingFace models, broadening model compatibility beyond TensorRT workflows.
Model Compatibility
- Complex Inputs: Support for dataclasses, user-defined objects in
module.forward()arguments, and lists/dicts within Torch module containers for more complete model analysis. - LLM Support: Added KV cache support to enable tuning of autoregressive large language models.
Performance & Observability
- Memory Optimization: Reduced CPU/GPU memory usage during tuning by offloading inactive modules to the
metadevice, with optimized input/output metadata handling. - Profiling: Extended metrics collection through NVTX annotations for Nsight Systems integration. Added configurable console output suppression with automatic log-to-file.
Documentation & Examples
- Documentation & Examples: Added comprehensive documentation and extended end-to-end examples across Computer Vision, Generative AI, Speech Recognition, and NLP workloads.
Bug Fixes
- Fixed dynamic shapes handling in TorchTensorRT AoT and TensorRT ONNX Dynamo export paths, calibration data creation for ModelOpt PTQ, bfloat16 precision in TensorRT, JIT cache directory collisions, and profiling for models without batching support.
Known Issues
- AITune currently only supports single-GPU configurations.
- Just-in-Time tuning does not support
transformers>=5due to@capture_outputsdecorator.
Full Changelog
用户点评
登录用户的真实使用记录，可按版本和硬件筛选共 0 条
还没有人点评。第一个说出真实体验的人最有价值。
NVIDIA AITune
NVIDIA AITune automates inference tuning for PyTorch models and pipelines on NVIDIA GPUs. It brings inference engines and acceleration techniques together under a single, extensible API. AITune combines backend evaluation, numerical validation, and performance measurement into one workflow, selecting implementations for individual modules according to the configured tuning strategy.
Start with a model from Hugging Face or timm, or bring your own model and checkpoint. For supported applications, enable just-in-time tuning during inference, or use the explicit tuning API to prepare artifacts for deployment.
Find your model in the recipe catalog, or follow the quick start.
Why AITune?
- Integrate through one API. Use a shared workflow across supported backends. Extend the workflow with custom backends and tuning strategies.
- Evaluate options per module. Use default candidates that account for execution needs, or configure candidates explicitly. Different modules in a pipeline can use different backends.
- Select using measurements. Check numerical outputs and measure performance on representative inputs. Choose strategies for throughput, latency, or a latency budget to guide selection.
- Prepare for deployment. Reuse tuned modules in Python, generate Triton model repositories, or serve through Dynamo workers, as supported by the recipe.
How it works
AITune combines these stages into a workflow for each model or pipeline:
- Inspect. Provide a PyTorch model and representative inputs. AITune finds tunable
nn.Modulecomponents and observes their inputs and execution. - Wrap. Customize the selected modules' tuning strategies and backend configurations.
- Tune. Evaluate configured backend candidates, check numerical outputs against the original implementation, and measure performance on representative inputs. Select a backend using the configured strategy. Candidates can incorporate post-training quantization (PTQ), kernel selection, or CUDA graphs where supported by the backend.
- Run & deploy. Run tuned modules in Python. Use the explicit workflow to save an
.aitartifact for reuse or deployment with Triton or Dynamo, as supported by the recipe.
A shared tuning workflow brings supported backends and acceleration techniques together, with selection guided by the configured strategy.
Install
Use a Linux environment with an NVIDIA GPU and a compatible PyTorch and CUDA installation:
pip install --extra-index-url https://pypi.nvidia.com aitune
See the installation guide for requirements, containers, and PyTorch/CUDA version selection.
Quick start
Both examples use Stable Diffusion. Install the model dependencies:
pip install diffusers transformers
Just-in-time (JIT): try acceleration in your existing application
Just-in-time tuning happens as your application runs, using inputs captured during inference.
Save this as your_script.py:
import torch
from diffusers import DiffusionPipeline
pipe = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5",
torch_dtype=torch.float16,
).to("cuda")
image = pipe("A landscape with mountains and a lake").images[0]
image.save("landscape.png")
Enable AITune when launching the script:
AUTOWRAPT_BOOTSTRAP=aitune_enable_jit_tuning python your_script.py
The example saves an image to landscape.png. Use representative prompts for your workload.
Eligible modules use the selected implementations in the running application.
A new process starts tuning again; use the explicit workflow below for saved artifacts.
See the JIT guide for configuration and deferred tuning
for pipelines.
Ahead-of-time (AOT): prepare a reusable tuning artifact
Ahead-of-time tuning prepares your model for deployment with a consistent flow:
Provide your model and representative inputs:
import torch
from diffusers import DiffusionPipeline
import aitune.torch as ait
pipe = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5",
torch_dtype=torch.float16,
).to("cuda")
inputs = [{"prompt": "A landscape with mountains and a lake"}]
modules = ait.inspect(pipe, inputs)
pipe = ait.wrap(pipe, modules.get_modules())
ait.tune(pipe, inputs)
image = pipe(**inputs[0]).images[0]
image.save("landscape.png")
ait.save(pipe, "tuned_pipe.ait")
This produces an image and a tuned_pipe.ait artifact that can be loaded into a compatible pipeline.
Use representative prompts or data for your workload. See the
AOT guide for tuning
configuration and the checkpoint guide
for saving and loading.
Model recipes
Explore ready-to-use model recipes for tuning, validation, benchmarking, and deployment. Choose a model and a configuration for your GPU setup, precision, and deployment target, then adapt the recipe with your own compatible checkpoint and data.
Learn more
- Documentation
- Supported backends
- Multi-GPU tuning
- Profiling and hardware metrics
- Deployment overview
- Triton deployment
- Dynamo deployment
Notice
NOTICE AND DISCLAIMER: This software automatically retrieves, accesses or interacts with external materials. Those retrieved materials are not distributed with this software and are governed solely by separate terms, conditions and licenses. You are solely responsible for finding, reviewing and complying with all applicable terms, conditions, and licenses, and for verifying the security, integrity and suitability of any retrieved materials for your specific use case. This software is provided "AS IS", without warranty of any kind. The author makes no representations or warranties regarding any retrieved materials, and assumes no liability for any losses, damages, liabilities or legal consequences from your use or inability to use this software or any retrieved materials. Use this software and the retrieved materials at your own risk.