diff --git a/examples/auto_deploy/README.md b/examples/auto_deploy/README.md
index cba226e7310f..9ece4b7ddc91 100644
--- a/examples/auto_deploy/README.md
+++ b/examples/auto_deploy/README.md
@@ -1,36 +1,12 @@
-
+# π₯πβ‘ AutoDeploy Examples
-# π₯πβ‘ AutoDeploy
-
-
Seamless Model Deployment from PyTorch to TRT-LLM
-
-
-
-AutoDeploy is a prototype feature in beta stage designed to simplify and accelerate the deployment of PyTorch models, including off-the-shelf models like those from Hugging Face, to TensorRT-LLM. It automates graph transformations to integrate inference optimizations such as tensor parallelism, KV-caching and quantization. AutoDeploy supports optimized in-framework deployment, minimizing the amount of manual modification needed.
+This folder contains runnable examples for **AutoDeploy**. For general AutoDeploy documentation, motivation, support matrix, and feature overview, please see the [official docs](https://nvidia.github.io/TensorRT-LLM/torch/auto_deploy/auto-deploy.html).
______________________________________________________________________
-## Motivation & Approach
-
-Deploying large language models (LLMs) can be challenging, especially when balancing ease of use with high performance. Teams need simple, intuitive deployment solutions that reduce engineering effort, speed up the integration of new models, and support rapid experimentation without compromising performance.
-
-AutoDeploy addresses these challenges with a streamlined, (semi-)automated pipeline that transforms in-framework PyTorch models, including Hugging Face models, into optimized inference-ready models for TRT-LLM. It simplifies deployment, optimizes models for efficient inference, and bridges the gap between simplicity and performance.
-
-### **Key Features:**
-
-- **Seamless Model Transition:** Automatically converts PyTorch/Hugging Face models to TRT-LLM without manual rewrites.
-- **Unified Model Definition:** Maintain a single source of truth with your original PyTorch/Hugging Face model.
-- **Optimized Inference:** Built-in transformations for sharding, quantization, KV-cache integration, MHA fusion, and CudaGraph optimization.
-- **Immediate Deployment:** Day-0 support for models with continuous performance enhancements.
-- **Quick Setup & Prototyping:** Lightweight pip package for easy installation with a demo environment for fast testing.
-
-______________________________________________________________________
+## Quick Start
-## Get Started
-
-1. **Install AutoDeploy:**
-
-AutoDeploy is accessible through TRT-LLM installation.
+AutoDeploy is included with the TRT-LLM installation.
```bash
sudo apt-get -y install libopenmpi-dev && pip3 install --upgrade pip setuptools && pip3 install tensorrt_llm
@@ -38,11 +14,7 @@ sudo apt-get -y install libopenmpi-dev && pip3 install --upgrade pip setuptools
You can refer to [TRT-LLM installation guide](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/installation/linux.md) for more information.
-2. **Run Llama Example:**
-
-You are ready to run an in-framework LLama Demo now.
-
-The general entrypoint to run the auto-deploy demo is the `build_and_run_ad.py` script, Checkpoints are loaded directly from Huggingface (HF) or a local HF-like directory:
+Run a simple example with a Hugging Face model:
```bash
cd examples/auto_deploy
@@ -51,93 +23,17 @@ python build_and_run_ad.py --model "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
______________________________________________________________________
-## Support Matrix
-
-AutoDeploy streamlines the model deployment process through an automated workflow designed for efficiency and performance. The workflow begins with a PyTorch model, which is exported using `torch.export` to generate a standard Torch graph. This graph contains core PyTorch ATen operations alongside custom attention operations, determined by the attention backend specified in the configuration.
-
-The exported graph then undergoes a series of automated transformations, including graph sharding, KV-cache insertion, and GEMM fusion, to optimize model performance. After these transformations, the graph is compiled using one of the supported compile backends (like `torch-opt`), followed by deploying it via the TRT-LLM runtime.
-
-### Supported Models
-
-**Bring Your Own Model**: AutoDeploy leverages `torch.export` and dynamic graph pattern matching, enabling seamless integration for a wide variety of models without relying on hard-coded architectures.
-
-Additionally, we have officially verified support for the following models:
-
-
-Click to expand supported models list
-
-| Model Series | HF Model Card | Model Factory | Precision | World Size | Runtime | Compile Backend ||| Attention Backend |||
-|--------------|----------------------|----------------|-----------|------------|---------|-----------------|--------------------|--------------------|--------------------|----------|----------|
-| | | | | | | torch-simple | torch-compile | torch-opt | triton | flashinfer | MultiHeadLatentAttention |
-| LLaMA | meta-llama/Llama-2-7b-chat-hf
meta-llama/Meta-Llama-3.1-8B-Instruct
meta-llama/Llama-3.1-70B-Instruct
codellama/CodeLlama-13b-Instruct-hf | AutoModelForCausalLM | BF16 | 1,2,4 | demollm, trtllm | β
| β
| β
| β
| β
| n/a |
-| LLaMA-4 | meta-llama/Llama-4-Scout-17B-16E-Instruct
meta-llama/Llama-4-Maverick-17B-128E-Instruct | AutoModelForImageTextToText | BF16 | 1,2,4,8 | demollm, trtllm | β
| β
| β | β
| β
| n/a |
-| Nvidia Minitron | nvidia/Llama-3_1-Nemotron-51B-Instruct
nvidia/Llama-3.1-Minitron-4B-Width-Base
nvidia/Llama-3.1-Minitron-4B-Depth-Base | AutoModelForCausalLM | BF16 | 1,2,4 | demollm, trtllm | β
| β
| β
| β
| β
| n/a |
-| Nvidia Model Optimizer | nvidia/Llama-3.1-8B-Instruct-FP8
nvidia/Llama-3.1-405B-Instruct-FP8 | AutoModelForCausalLM | FP8 | 1,2,4 | demollm, trtllm | β
| β
| β
| β
| β
| n/a |
-| DeepSeek | deepseek-ai/DeepSeek-R1-Distill-Llama-70B | AutoModelForCausalLM | BF16 | 1,2,4 | demollm, trtllm | β
| β
| β
| β
| β
| n/a |
-| Mistral | mistralai/Mixtral-8x7B-Instruct-v0.1
mistralai/Mistral-7B-Instruct-v0.3 | AutoModelForCausalLM | BF16 | 1,2,4 | demollm, trtllm | β
| β
| β
| β
| β
| n/a |
-| BigCode | bigcode/starcoder2-15b | AutoModelForCausalLM | FP32 | 1,2,4 | demollm, trtllm | β
| β
| β
| β
| β
| n/a |
-| Deepseek-V3 | deepseek-ai/DeepSeek-V3 | AutoModelForCausalLM | BF16 | 1,2,4 | demollm | β
| β | β | n/a | n/a | β
|
-| Phi4 | microsoft/phi-4
microsoft/Phi-4-reasoning
microsoft/Phi-4-reasoning-plus | AutoModelForCausalLM | BF16 | 1,2,4 | demollm, trtllm | β
| β
| β
| β
| β
| n/a |
-| Phi3/2 | microsoft/Phi-3-mini-4k-instruct
microsoft/Phi-3-mini-128k-instruct
microsoft/Phi-3-medium-4k-instruct
microsoft/Phi-3-medium-128k-instruct
microsoft/Phi-3.5-mini-instruct | AutoModelForCausalLM | BF16 | 1,2,4 | demollm, trtllm | β
| β
| β
(partly) | β
| β | n/a |
-
-
-
-### Runtime Integrations
-
-AutoDeploy runs natively with the entire `TRT-LLM` stack via the `LLM` API. In addition, we provide a light-weight wrapper of the `LLM` API for onboarding and debugging new models:
-
-| `"runtime"` | Description |
-|-------------|-------------|
-| `trtllm` | A robust, production-grade runtime optimized for high-performance inference. |
-| `demollm` | A lightweight runtime wrapper designed for development and testing, featuring a naive scheduler and KV-cache manager for simplified debugging and testing. |
-
-### Compile Backends
-
-AutoDeploy supports multiple backends for compiling the exported Torch graph:
-
-| `"compile_backend"` | Description |
-|--------------------|-------------|
-| `torch-simple` | Exports the graph without additional optimizations. |
-| `torch-compile` | Applies `torch.compile` to the graph after all AutoDeploy transformations have been completed. |
-| `torch-cudagraph` | Performs CUDA graph capture (without torch.compile). |
-| `torch-opt` | Uses `torch.compile` along with CUDA Graph capture to enhance inference performance. |
-
-### Attention backends
-
-Optimize attention operations using different attention kernel implementations:
-
-| `"attn_backend"` | Description |
-|----------------------|-------------|
-| `triton` | Custom fused multi-head attention (MHA) with KV Cache kernels for efficient attention processing. |
-| `flashinfer` | Uses off-the-shelf optimized attention kernels with KV Cache from the [`flashinfer`](https://github.com/flashinfer-ai/flashinfer.git) library. |
-
-### Precision Support
-
-AutoDeploy supports a range of precision formats to enhance model performance, including:
-
-- BF16, FP32
-- Quantization formats like FP8.
-
-______________________________________________________________________
-
-## Advanced Usage
-
-### Example Run Script ([`build_and_run_ad.py`](./build_and_run_ad.py))
-
-To build and run AutoDeploy example, use the [`build_and_run_ad.py`](./build_and_run_ad.py) script:
+## Example Run Script ([`build_and_run_ad.py`](./build_and_run_ad.py))
-```bash
-cd examples/auto_deploy
-python build_and_run_ad.py --model "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-```
+This script demonstrates end-to-end deployment of HuggingFace checkpoints using AutoDeployβs graph-transformation pipeline.
-You can arbitrarily configure your experiment. Use the `-h/--help` flag to see available options:
+You can configure your experiment with various options. Use the `-h/--help` flag to see available options:
```bash
python build_and_run_ad.py --help
```
-Below is a non-exhaustive list of common config options:
+Below is a non-exhaustive list of common configuration options:
| Configuration Key | Description |
|-------------------|-------------|
@@ -159,7 +55,7 @@ Below is a non-exhaustive list of common config options:
For default values and additional configuration options, refer to the [`ExperimentConfig`](./build_and_run_ad.py) class in [build_and_run_ad.py](./build_and_run_ad.py) file.
-Here is a more complete example of using the script:
+The following is a more complete example of using the script:
```bash
cd examples/auto_deploy
@@ -172,282 +68,16 @@ python build_and_run_ad.py \
--benchmark.enabled True
```
-### Logging Level
-
-Use the following env variable to specify the logging level of our built-in logger ordered by
-decreasing verbosity;
-
-```bash
-AUTO_DEPLOY_LOG_LEVEL=DEBUG
-AUTO_DEPLOY_LOG_LEVEL=INFO
-AUTO_DEPLOY_LOG_LEVEL=WARNING
-AUTO_DEPLOY_LOG_LEVEL=ERROR
-AUTO_DEPLOY_LOG_LEVEL=INTERNAL_ERROR
-```
-
-The default level is `INFO`.
-
-### Model Evaluation with LM Evaluation Harness
-
-lm-evaluation-harness is supported. To run the evaluation, please use the following command:
-
-```bash
-# model is defined the same as above. Other config args can also be specified in the model_args (comma separated).
-# You can specify any tasks supported with lm-evaluation-harness.
-cd examples/auto_deploy
-python lm_eval_ad.py \
---model autodeploy --model_args model=meta-llama/Meta-Llama-3.1-8B-Instruct,world_size=2 --tasks mmlu
-```
-
-### Mixed-precision Quantization using TensorRT Model Optimizer
-
-TensorRT Model Optimizer [AutoQuantize](https://nvidia.github.io/TensorRT-Model-Optimizer/reference/generated/modelopt.torch.quantization.model_quant.html#modelopt.torch.quantization.model_quant.auto_quantize) algorithm is a PTQ algorithm from ModelOpt which quantizes a model by searching for the best quantization format per-layer while meeting the performance constraint specified by the user. This way, `AutoQuantize` enables to trade-off model accuracy for performance.
-
-Currently `AutoQuantize` supports only `effective_bits` as the performance constraint (for both weight-only quantization and weight & activation quantization). See
-[AutoQuantize documentation](https://nvidia.github.io/TensorRT-Model-Optimizer/reference/generated/modelopt.torch.quantization.model_quant.html#modelopt.torch.quantization.model_quant.auto_quantize) for more details.
-
-#### 1. Quantize a model with ModelOpt
-
-Refer to [NVIDIA TensorRT Model Optimizer](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/examples/llm_autodeploy/README.md) for generating quantized model checkpoint.
-
-#### 2. Deploy the quantized model with AutoDeploy
-
-```bash
-cd examples/auto_deploy
-python build_and_run_ad.py --model "" --args.world-size 1
-```
-
-### Incorporating `auto_deploy` into your own workflow
-
-AutoDeploy can be seamlessly integrated into your existing workflows using TRT-LLM's LLM high-level API. This section provides a blueprint for configuring and invoking AutoDeploy within your custom applications.
-
-Here is an example of how you can build an LLM object with AutoDeploy integration:
-
-```
-from tensorrt_llm._torch.auto_deploy import LLM
-
-
-# Construct the LLM high-level interface object with autodeploy as backend
-llm = LLM(
- model=,
- world_size=,
- compile_backend="torch-compile",
- model_kwargs={"num_hidden_layers": 2}, # test with smaller model configuration
- attn_backend="flashinfer", # choose between "triton" and "flashinfer"
- attn_page_size=64, # page size for attention (tokens_per_block, should be == max_seq_len for triton)
- skip_loading_weights=False,
- model_factory="AutoModelForCausalLM", # choose appropriate model factory
- mla_backend="MultiHeadLatentAttention", # for models that support MLA
- free_mem_ratio=0.8, # fraction of available memory for cache
- simple_shard_only=False, # tensor parallelism sharding strategy
- max_seq_len=,
- max_batch_size=,
-)
-
-```
-
-Please consult the [AutoDeploy `LLM` API](../../tensorrt_llm/_torch/auto_deploy/llm.py) and the
-[`AutoDeployConfig` class](../../tensorrt_llm/_torch/auto_deploy/llm_args.py)
-for more detail on how AutoDeploy is configured via the `**kwargs` of the `LLM` API.
-
-### Expert Configuration of LLM API
-
-For expert TensorRT-LLM users, we also expose the full set of [`LlmArgs`](../../tensorrt_llm/_torch/auto_deploy/llm_args.py)
-*at your own risk* (the argument list diverges from TRT-LLM's argument list):
-
-
-Click to expand for more details on using LlmArgs directly
-
-- All config fields that are used by the AutoDeploy core pipeline (i.e. the `InferenceOptimizer`) are
- _exclusively_ exposed in the [`AutoDeployConfig` class](../../tensorrt_llm/_torch/auto_deploy/llm_args.py).
- Please make sure to refer to those first.
-- For expert users we expose the full set of [`LlmArgs`](../../tensorrt_llm/_torch/auto_deploy/llm_args.py)
- that can be used to configure the [AutoDeploy `LLM` API](../../tensorrt_llm/_torch/auto_deploy/llm.py) including runtime options.
-- Note that some fields in the full [`LlmArgs`](../../tensorrt_llm/_torch/auto_deploy/llm_args.py)
- object are overlapping, duplicated, and/or _ignored_ in AutoDeploy, particularly arguments
- pertaining to configuring the model itself since AutoDeploy's model ingestion+optimize pipeline
- significantly differs from the default manual workflow in TensorRT-LLM.
-- However, with the proper care the full [`LlmArgs`](../../tensorrt_llm/_torch/auto_deploy/llm_args.py)
- objects can be used to configure advanced runtime options in TensorRT-LLM.
-- Note that any valid field can be simply provided as keyword argument ("`**kwargs`") to the
- [AutoDeploy `LLM` API](../../tensorrt_llm/_torch/auto_deploy/llm.py).
-
-
-
-### Expert Configuration of `build_and_run_ad.py`
-
-For expert users, `build_and_run_ad.py` provides advanced configuration capabilities through a flexible argument parser powered by PyDantic Settings and OmegaConf. You can use dot notation for CLI arguments, provide multiple YAML configuration files, and leverage sophisticated configuration precedence rules to create complex deployment configurations.
-
-
-Click to expand for detailed configuration examples
-
-#### CLI Arguments with Dot Notation
-
-The script supports flexible CLI argument parsing using dot notation to modify nested configurations dynamically. You can target any field in both the [`ExperimentConfig`](./build_and_run_ad.py) and nested [`AutoDeployConfig`](../../tensorrt_llm/_torch/auto_deploy/llm_args.py)/[`LlmArgs`](../../tensorrt_llm/_torch/auto_deploy/llm_args.) objects:
-
-```bash
-# Configure model parameters
-# NOTE: config values like num_hidden_layers are automatically resolved into the appropriate nested
-# dict value ``{"args": {"model_kwargs": {"num_hidden_layers": 10}}}`` although not explicitly
-# specified as CLI arg
-python build_and_run_ad.py \
- --model "meta-llama/Meta-Llama-3.1-8B-Instruct" \
- --args.model-kwargs.num-hidden-layers=10 \
- --args.model-kwargs.hidden-size=2048 \
- --args.tokenizer-kwargs.padding-side=left
-
-# Configure runtime and backend settings
-python build_and_run_ad.py \
- --model "TinyLlama/TinyLlama-1.1B-Chat-v1.0" \
- --args.world-size=2 \
- --args.compile-backend=torch-opt \
- --args.attn-backend=flashinfer
-
-# Configure prompting and benchmarking
-python build_and_run_ad.py \
- --model "microsoft/phi-4" \
- --prompt.batch-size=4 \
- --prompt.sp-kwargs.max-tokens=200 \
- --prompt.sp-kwargs.temperature=0.7 \
- --benchmark.enabled=true \
- --benchmark.bs=8 \
- --benchmark.isl=1024
-```
-
-#### YAML Configuration Files
-
-Both [`ExperimentConfig`](./build_and_run_ad.py) and [`AutoDeployConfig`](../../tensorrt_llm/_torch/auto_deploy/llm_args.py)/[`LlmArgs`](../../tensorrt_llm/_torch/auto_deploy/llm_args.py) inherit from [`DynamicYamlMixInForSettings`](../../tensorrt_llm/_torch/auto_deploy/utils/_config.py), enabling you to provide multiple YAML configuration files that are automatically deep-merged at runtime.
-
-Create a YAML configuration file (e.g., `my_config.yaml`):
-
-```yaml
-# my_config.yaml
-args:
- model_kwargs:
- num_hidden_layers: 12
- hidden_size: 1024
- world_size: 4
- compile_backend: torch-compile
- attn_backend: triton
- max_seq_len: 2048
- max_batch_size: 16
- transforms:
- sharding:
- strategy: auto
- quantization:
- enabled: false
-
-prompt:
- batch_size: 8
- sp_kwargs:
- max_tokens: 150
- temperature: 0.8
- top_k: 50
-
-benchmark:
- enabled: true
- num: 20
- bs: 4
- isl: 1024
- osl: 256
-```
-
-Create an additional override file (e.g., `production.yaml`):
-
-```yaml
-# production.yaml
-args:
- world_size: 8
- compile_backend: torch-opt
- max_batch_size: 32
-
-benchmark:
- enabled: false
-```
-
-Then use these configurations:
-
-```bash
-# Using single YAML config
-python build_and_run_ad.py \
- --model "meta-llama/Meta-Llama-3.1-8B-Instruct" \
- --yaml-configs my_config.yaml
-
-# Using multiple YAML configs (deep merged in order, later files have higher priority)
-python build_and_run_ad.py \
- --model "meta-llama/Meta-Llama-3.1-8B-Instruct" \
- --yaml-configs my_config.yaml production.yaml
-
-# Targeting nested AutoDeployConfig with separate YAML
-python build_and_run_ad.py \
- --model "meta-llama/Meta-Llama-3.1-8B-Instruct" \
- --yaml-configs my_config.yaml \
- --args.yaml-configs autodeploy_overrides.yaml
-```
-
-#### Configuration Precedence and Deep Merging
-
-The configuration system follows a strict precedence order where higher priority sources override lower priority ones:
-
-1. **CLI Arguments** (highest priority) - Direct command line arguments
-1. **YAML Configs** - Files specified via `--yaml-configs` and `--args.yaml-configs`
-1. **Default Settings** (lowest priority) - Built-in defaults from the config classes
-
-**Deep Merging**: Unlike simple overwriting, deep merging intelligently combines nested dictionaries recursively. For example:
-
-```yaml
-# Base config
-args:
- model_kwargs:
- num_hidden_layers: 10
- hidden_size: 1024
- max_seq_len: 2048
-```
-
-```yaml
-# Override config
-args:
- model_kwargs:
- hidden_size: 2048 # This will override
- # num_hidden_layers: 10 remains unchanged
- world_size: 4 # This gets added
-```
-
-**Nested Config Behavior**: When using nested configurations, outer YAML configs become init settings for inner objects, giving them higher precedence:
-
-```bash
-# The outer yaml-configs affects the entire ExperimentConfig
-# The inner args.yaml-configs affects only the AutoDeployConfig
-python build_and_run_ad.py \
- --model "meta-llama/Meta-Llama-3.1-8B-Instruct" \
- --yaml-configs experiment_config.yaml \
- --args.yaml-configs autodeploy_config.yaml \
- --args.world-size=8 # CLI override beats both YAML configs
-```
-
-#### Built-in Default Configuration
-
-Both [`AutoDeployConfig`](../../tensorrt_llm/_torch/auto_deploy/llm_args.py) and [`LlmArgs`](../../tensorrt_llm/_torch/auto_deploy/llm_args.py) classes automatically load a built-in [`default.yaml`](../../tensorrt_llm/_torch/auto_deploy/config/default.yaml) configuration file that provides sensible defaults for the AutoDeploy inference optimizer pipeline. This file is specified in the [`_get_config_dict()`](../../tensorrt_llm/_torch/auto_deploy/llm_args.py) function and defines default transform configurations for graph optimization stages.
-
-The built-in defaults are automatically merged with your configurations at the lowest priority level, ensuring that your custom settings always override the defaults. You can inspect the current default configuration to understand the baseline transform pipeline:
-
-```bash
-# View the default configuration
-cat tensorrt_llm/_torch/auto_deploy/config/default.yaml
-
-# Override specific transform settings
-python build_and_run_ad.py \
- --model "TinyLlama/TinyLlama-1.1B-Chat-v1.0" \
- --args.transforms.export-to-gm.strict=true
-```
+### Advanced Configuration
-
+The script supports flexible configs:
-## Roadmap
+- CLI dot notation for nested fields
+- YAML configs with deep merge
+- Precedence: CLI > YAML > defaults
-Check out our [Github Project Board](https://github.com/orgs/NVIDIA/projects/83) to learn more about
-the current progress in AutoDeploy and where you can help.
+Please refer to [Expert Configuration of LLM API](https://nvidia.github.io/TensorRT-LLM/torch/auto_deploy/advanced/expert_configurations.html) for details.
## Disclaimer
-This project is in active development and is currently in an early (beta) stage. The code is in prototype, subject to change, and may include backward-incompatible updates. While we strive for correctness, we provide no guarantees regarding functionality, stability, or reliability. Use at your own risk.
+This project is under active development and is currently in a prototype stage. The code is experimental, subject to change, and may include backward-incompatible updates. While we strive for correctness, there are no guarantees regarding functionality, stability, or reliability.
diff --git a/examples/auto_deploy/lm_eval_ad.py b/examples/auto_deploy/lm_eval_ad.py
deleted file mode 100644
index 5cd4377848ce..000000000000
--- a/examples/auto_deploy/lm_eval_ad.py
+++ /dev/null
@@ -1,120 +0,0 @@
-from typing import Any, Dict, List, Tuple
-
-import torch.nn.functional as F
-from build_and_run_ad import ExperimentConfig, build_llm_from_config
-from lm_eval.__main__ import cli_evaluate
-from lm_eval.api.model import TemplateLM
-from lm_eval.api.registry import register_model
-from tqdm import tqdm
-
-from tensorrt_llm._torch.auto_deploy import LlmArgs
-from tensorrt_llm._torch.auto_deploy.utils.logger import ad_logger
-from tensorrt_llm.llmapi import RequestOutput
-from tensorrt_llm.sampling_params import SamplingParams
-
-
-@register_model("autodeploy")
-class AutoDeployEval(TemplateLM):
- def __init__(self, chunk_size: int = 200, **kwargs):
- super().__init__()
-
- # some lm-eval specific default values
- kwargs["max_tokens"] = max(int(kwargs.get("max_tokens", 0)), 256)
- kwargs["max_seq_len"] = max(2048 + kwargs["max_tokens"], int(kwargs.get("max_seq_len", 0)))
- if "batch_size" in kwargs:
- kwargs["batch_size"] = int(kwargs["batch_size"])
-
- self.config = ExperimentConfig(args=LlmArgs(**kwargs))
- self.chunk_size = chunk_size
-
- ad_logger.info(f"AutoDeploy config: {self.config}")
- self.llm = build_llm_from_config(self.config)
- ad_logger.info("Loaded AutoDeploy model")
-
- @property
- def eot_token_id(self) -> int:
- return self.llm.tokenizer.eos_token_id
-
- def tok_encode(self, string, add_special_tokens=False, **kwargs):
- return self.llm.tokenizer.encode(string, add_special_tokens=add_special_tokens, **kwargs)
-
- def _loglikelihood_tokens(
- self, requests: List[Any], disable_tqdm: bool = False
- ) -> List[Tuple[float, bool]]:
- """Compute the log likelihood of the continuation given the context."""
- # some book-keeping...
- num_r = len(requests)
- desc = "Processing loglikelihood requests"
- sampling_params = SamplingParams(max_tokens=1, return_context_logits=True)
-
- # process requests
- futures: Dict[int, RequestOutput] = {}
- results = []
- for i, request in tqdm(enumerate(requests), desc=desc, total=num_r, disable=disable_tqdm):
- # asynchronously submit a chunk of requests ahead of time...
- if i % self.chunk_size == 0:
- for j in range(i, min(i + self.chunk_size, num_r)):
- prompt_ids = requests[j][1] + requests[j][2]
- futures[j] = self.llm.generate_async(prompt_ids, sampling_params)
-
- # process the output of the request i
- r_out: RequestOutput = futures.pop(i).result()
-
- # check continuation portion of the prompt
- # NOTE: context_logits are offset by 1 since they predict future token
- ctxlen = len(request[1])
- token_ids_cont = request[2]
- logits_cont = r_out.context_logits[ctxlen - 1 : -1] # [sl, vocab]
- logprobs_cont = F.log_softmax(logits_cont, dim=-1) # [sl, vocab]
- top_tokens_cont = logprobs_cont.argmax(dim=-1).tolist() # [sl]
-
- # compute logprob and check for greedy
- logprob_sum = sum(logprobs_cont[list(range(len(logprobs_cont))), token_ids_cont]).item()
- is_greedy = top_tokens_cont == token_ids_cont
-
- results.append((logprob_sum, is_greedy))
-
- # clear response
- del r_out
-
- return results
-
- def loglikelihood_rolling(self, requests, disable_tqdm: bool = False):
- raise NotImplementedError
-
- def generate_until(self, requests: List[Any], disable_tqdm: bool = False) -> List[str]:
- # some book-keeping and parameters...
- num_r = len(requests)
- desc = "Processing generate requests"
-
- def _get_sp(gen_kwargs):
- k_mapping = {
- "temperature": "temperature",
- "top_p": "top_p",
- "until": "stop",
- }
- kwargs_mapped = {
- k_sp: gen_kwargs[k_gen] for k_gen, k_sp in k_mapping.items() if k_gen in gen_kwargs
- }
- return SamplingParams(max_tokens=self.config.max_tokens, **kwargs_mapped)
-
- # process requests
- futures: Dict[int, RequestOutput] = {}
- results = []
- for i, _ in tqdm(enumerate(requests), desc=desc, total=num_r, disable=disable_tqdm):
- # asynchronously submit a chunk of requests ahead of time...
- if i % self.chunk_size == 0:
- for j in range(i, min(i + self.chunk_size, num_r)):
- context, gen_kwargs = requests[j].args
- prompt_ids = self.tok_encode(context)
- futures[j] = self.llm.generate_async(prompt_ids, _get_sp(gen_kwargs))
-
- # process the output of the request i
- r_out: RequestOutput = futures.pop(i).result()
- results.append(r_out.outputs[0].text)
-
- return results
-
-
-if __name__ == "__main__":
- cli_evaluate()
diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py b/tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py
index 80f9d440c156..9ded1bea1c3a 100644
--- a/tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py
+++ b/tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py
@@ -232,7 +232,8 @@ def _get_mem_info_in_mb():
current_cache_size = cm.current_cache_size_bytes()
current_num_pages = cm.info.num_pages
ad_logger.info(
- f"Current cache size: {current_cache_size}, Current num pages: {current_num_pages}"
+ f"Current cache size (MB): {current_cache_size // 1024 // 1024}, "
+ f"Current num pages (MB): {current_num_pages}"
)
if free_mem_ratio == 0.0:
diff --git a/tests/unittest/_torch/auto_deploy/integration/test_lm_eval.py b/tests/unittest/_torch/auto_deploy/integration/test_lm_eval.py
deleted file mode 100644
index c1b7b5c26557..000000000000
--- a/tests/unittest/_torch/auto_deploy/integration/test_lm_eval.py
+++ /dev/null
@@ -1,148 +0,0 @@
-"""Testing lm-eval with AutoDeploy for basic accuracy validation.
-
-NOTE (lucaslie): this test is for local testing only. It is not registered to run as part of CI.
-"""
-
-import gc
-import glob
-import json
-import os
-from unittest import mock
-
-import lm_eval_ad # noqa: F401
-import pytest
-from _dist_test_utils import param_with_device_count
-from _model_test_utils import _hf_model_dir_or_hub_id
-from _torch_test_utils import fp4_compatible, fp8_compatible
-from lm_eval.__main__ import cli_evaluate
-from lm_eval.api.task import TaskConfig
-from utils.llm_data import llm_models_root
-
-_task_config_post_init = TaskConfig.__post_init__
-
-
-def _patched_post_init(self: TaskConfig, *args, **kwargs):
- # Call the original __post_init__
- ret = _task_config_post_init(self, *args, **kwargs)
-
- # Redirect dataset_path to the local directory if it exists.
- # This is to enable the test works in offline mode when we are running on CI.
- dataset_path_lookup = {
- "hails/mmlu_no_train": f"{llm_models_root()}/datasets/hails/mmlu_no_train",
- "gsm8k": f"{llm_models_root()}/datasets/openai/gsm8k",
- }
-
- # If the dataset key isn't in dataset_path_lookup, or the local path doesn't exist,
- # the test fallbacks to downloading dataset from huggingface hub; it will fail if
- # offline mode is enabled by HF_DATASETS_OFFLINE=1.
- if self.dataset_path in dataset_path_lookup:
- local_path = dataset_path_lookup[self.dataset_path]
- if os.path.exists(local_path):
- self.dataset_path = local_path
-
- return ret
-
-
-def _cli_evaluate_with_mocks(args):
- # Mock post-init and argparser's reference to sys.argv
- with (
- mock.patch.object(TaskConfig, "__post_init__", new=_patched_post_init),
- mock.patch("argparse._sys.argv", [""] + args),
- ):
- # set up args via argparser
- cli_evaluate()
-
-
-# LM eval limit setting has a bug: https://github.com/EleutherAI/lm-evaluation-harness/issues/2324
-# So we run each task individually.
-@pytest.mark.parametrize(
- "world_size,model_args,tasks,score_keys,scores",
- [
- param_with_device_count(
- 2,
- {
- "model": _hf_model_dir_or_hub_id(
- f"{llm_models_root()}/llama-3.1-model/Llama-3.1-8B-Instruct",
- "meta-llama/Meta-Llama-3.1-8B-Instruct",
- ),
- "attn_backend": "flashinfer",
- "max_batch_size": 32,
- },
- ["gsm8k", "mmlu"],
- ["exact_match,strict-match", "acc,none"],
- [0.75, 0.675],
- ),
- param_with_device_count(
- 2,
- {
- "model": _hf_model_dir_or_hub_id(
- f"{llm_models_root()}/llama-3.1-model/Llama-3.1-8B-Instruct-FP8",
- "nvidia/Llama-3.1-8B-Instruct-FP8",
- )
- },
- ["gsm8k", "mmlu"],
- ["exact_match,strict-match", "acc,none"],
- [0.70, 0.67],
- marks_extra=[
- pytest.mark.skipif(not fp8_compatible(), reason="Requires fp8 support"),
- ],
- ),
- param_with_device_count(
- 2,
- {
- "model": "nvidia/Llama-3.1-8B-Instruct-FP4",
- },
- ["gsm8k", "mmlu"],
- ["exact_match,strict-match", "acc,none"],
- [0.70, 0.64],
- marks_extra=[
- pytest.mark.skipif(not fp4_compatible(), reason="Requires fp4 support"),
- ],
- ),
- param_with_device_count(
- 4,
- {
- "model": _hf_model_dir_or_hub_id(
- f"{llm_models_root()}/Mixtral-8x7B-Instruct-v0.1",
- "mistralai/Mixtral-8x7B-Instruct-v0.1",
- ),
- "compile_backend": "torch-simple",
- },
- ["gsm8k", "mmlu"],
- ["exact_match,strict-match", "acc,none"],
- [0.583, 0.67],
- ),
- ],
-)
-def test_lm_eval(world_size: int, model_args, tasks, score_keys, scores, tmp_path):
- # setup base args for our testing
- args = ["--model", "autodeploy", "--batch_size", "4", "--limit", "0.1"]
-
- # set up model args with world size
- model_args["world_size"] = world_size
- args += ["--model_args", ",".join([f"{k}={v}" for k, v in model_args.items()])]
- gen_kwargs = {}
- # Greedy to avoid variance
- gen_kwargs["temperature"] = 1e-4
- gen_kwargs["top_k"] = 1
- args += ["--gen_kwargs", ",".join([f"{k}={v}" for k, v in gen_kwargs.items()])]
-
- # set up output path and tasks
- args += ["--output_path", str(tmp_path)]
- args += ["--tasks", ",".join(tasks) if len(tasks) > 1 else tasks[0]]
-
- # run the cli with CI patches
- _cli_evaluate_with_mocks(args)
-
- json_files = glob.glob(os.path.join(tmp_path, "**", "*.json"), recursive=True)
- assert json_files
- with open(json_files[0], "r") as f:
- results = json.load(f)["results"]
-
- for task, score_key, score in zip(tasks, score_keys, scores):
- assert results[task][score_key] >= score, (
- f"{task}:{score_key} blow expected threshold. Details: {results[task]}"
- )
-
- # Terminate and free the executor resource.
- gc.collect()