Triton Model Navigator

A Navigator Package is a serialized artifact used to facilitate model deployment and optimization portability. Because a model graph or checkpoint alone is insufficient for deployment (which requires input/output definitions, maximal batch sizes, etc.), the package bundles the following:

• Serialized Models: Includes base formats (first available serialization), the max throughput format, and the min latency format.
• Metadata & Details: Optimization status files and execution logs.
• Reproducibility: A reproduction script for each model format.
• Data Samples: Input and output data samples provided as .npy files.

This package is the recommended way to share optimized models for deployment on PyTriton or Triton Inference Server, or to re-run optimization on different hardware without needing the original model source code.

Inplace Optimize is a PyTorch-specific feature that allows you to substitute nn.Module objects within your Python code with optimized versions (e.g., TensorRT or Torch-TensorRT) using a single API.

This is achieved using the nav.Module wrapper, which decorates your pipeline models. When you call nav.optimize with a dataloader, Model Navigator audits the decorated modules, gathers metadata, performs export/conversion, and replaces the original modules with optimized versions directly within your codebase.

Key components:

• callable: A Python object, function, or callable containing one or more wrapped models to optimize.
• dataloader: A method or class that generates input data. This data is used for export, conversion, determining input/output shapes, and creating output samples for optimization.

import model_navigator as nav

# Wrap a module
wrapped_module = nav.Module(original_module, name="my_module")

# Run optimization
nav.optimize(wrapped_module, dataloader)

The Triton Model Navigator optimization process is a multi-step pipeline designed to convert deep learning models (TensorFlow 2, PyTorch, ONNX, or JAX) into high-performance formats like TensorRT.

To start the process, you must provide a model and a dataloader. You can further customize the process using custom_configs.

The pipeline follows these five stages:

1. Model export: The source model is exported to an intermediary format (TorchScript, SavedModel, or ONNX).
2. Model conversion: The model is converted to a target representation (TorchTensorRT, TensorFlowTensorRT, ONNX, or TensorRT) to maximize performance.
3. Correctness test: The Navigator calculates absolute and relative tolerance values to ensure the converted model matches the source model's output.
4. Model profiling: Using Navigator Runners, the system performs inference to find maximum throughput and latency for each configuration.
5. Verification: The system validates the performance metrics against a user-provided verify_func (if specified).

Inplace Optimization allows you to optimize model components (like converting them to TensorRT) within a Python pipeline without modifying the original pipeline code.

To use this, you must wrap specific modules of your pipeline (e.g., text encoder, UNet, VAE) using nav.Module. You can also provide an output_mapping to ensure the optimized module's output is compatible with the original pipeline's expected types (e.g., HuggingFace classes).

Workflow:

1. Wrap pipeline modules with nav.Module.
2. Prepare a dataloader where each element is a tuple: (batch_size, data_sample).
3. Call nav.optimize(pipeline, dataloader).
4. Call nav.load_optimized() to switch the pipeline to use the optimized modules directly in Python.

import model_navigator as nav
from transformers.modeling_outputs import BaseModelOutputWithPooling
from diffusers import DPMSolverMultistepScheduler, StableDiffusionPipeline

def get_pipeline():
    pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1")
    pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
    pipe = pipe.to("cuda")
    
    # Wrap modules for optimization
    pipe.text_encoder = nav.Module(
        pipe.text_encoder,
        name="clip",
        output_mapping=lambda output: BaseModelOutputWithPooling(**output),
    )
    pipe.unet = nav.Module(pipe.unet, name="unet")
    pipe.vae.decoder = nav.Module(pipe.vae.decoder, name="vae")
    return pipe

# Dataloader: first element in tuple must be batch size
def get_dataloader():
    return [(1, "a photo of an astronaut riding a horse on mars")]

pipe = get_pipeline()
dataloader = get_dataloader()

# Execute optimization
nav.optimize(pipe, dataloader)

# Load optimized modules for direct Python execution
nav.load_optimized()

The Inplace Optimize feature allows you to seamlessly substitute nn.Module objects in a PyTorch pipeline with their optimized versions. This is achieved by wrapping specific modules with nav.Module and then calling nav.optimize on the entire pipeline.

Key concepts:

• pipeline: A Python object or callable containing one or more wrapped models.
• dataloader: A method or class that generates input data. The dataloader is used to perform export/conversion, determine input shape ranges, and create output samples for optimization.

To use this, wrap your modules with nav.Module, providing a name and an optional output_mapping (useful for converting raw output data back into specific HuggingFace classes). After calling nav.optimize(pipeline, dataloader), use nav.load_optimized() to load the most performant versions of the modules into your existing pipeline object.

import model_navigator as nav
from diffusers import StableDiffusionPipeline

def get_pipeline():
    pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1")
    pipe = pipe.to("cuda")

    # Wrap modules with nav.Module
    pipe.text_encoder = nav.Module(
        pipe.text_encoder,
        name="clip",
        output_mapping=lambda output: BaseModelOutputWithPooling(**output),
    )
    pipe.unet = nav.Module(pipe.unet, name="unet")
    pipe.vae.decoder = nav.Module(pipe.vae.decoder, name="vae")

    return pipe

def get_dataloader():
    # The first element in the tuple must be the batch size
    return [(1, "a photo of an astronaut riding a horse on mars")]

pipe = get_pipeline()
dataloader = get_dataloader()

# Execute optimization
nav.optimize(pipe, dataloader)

# Load the optimized modules into the pipeline
nav.load_optimized()

# Use the pipeline normally
images = pipe(["a photo of an astronaut riding a horse on mars"])