Deployment Guide

This guide covers the full deployment story for AITune-tuned models: saving a tuned model to a checkpoint, loading it in production, and optionally serving it as an OpenAI-compatible HTTP endpoint via NVIDIA Dynamo.

Save a Tuned Model

Basic Save

import aitune.torch as ait
# After tuning
ait.save(model, "checkpoints/model.ait")

This creates:

• checkpoints/model.ait: Compressed checkpoint with tuned modules
• checkpoints/model_sha256_sums.txt: SHA256 checksums
• checkpoints/model/: Decompressed artifacts (after first load)

With Custom Storage

from aitune.torch import LocalTorchStorage
storage = LocalTorchStorage(
    base_folder="production/models",
    remove_checkpoint_after_tune=False,
)
ait.save(model, "model_v2.ait", storage=storage)

Load in Production

Basic Load

import aitune.torch as ait
model = YourModel()
model.eval()
model.to("cuda")
ait.load(model, "checkpoints/model.ait")
output = model(input_data)

With Custom Storage

from aitune.torch import LocalTorchStorage
storage = LocalTorchStorage(base_folder="production/models")
ait.load(model, "model.ait", storage=storage)

Loading Process

1. First load — decompresses .ait file, extracts artifacts, verifies checksums, loads backend and weights. Slower due to decompression.
2. Subsequent loads — uses decompressed files from checkpoints/, skips decompression. Faster startup.

Serve with Dynamo Worker

After loading a tuned model, you can expose it as an OpenAI-compatible HTTP endpoint using AITune’s Dynamo integration. The worker registers the model with the Dynamo HTTP frontend, deserializes incoming requests, packs inference results into the Dynamo wire format, and blocks until SIGTERM/SIGINT.

Prerequisites

Install the Dynamo extra:

uv pip install "aitune[dynamo]"

The "audio" modality requires NVIDIA Dynamo 1.1 or later and serves text-to-speech requests through /v1/audio/speech. Automatic speech recognition is not currently exposed by this worker API.

For local development without etcd/NATS, set DYN_DISCOVERY_BACKEND=file before starting any Dynamo process.

Quick Start — Embedding Model

import numpy as np
import aitune.dynamo as dyn
import aitune.torch as ait  # for ait.config, ait.load, ait.save
from sentence_transformers import SentenceTransformer
model = SentenceTransformer("intfloat/e5-large-v2")
ait.config.device_after_tuning = "cpu"
ait.load(model, "checkpoints/e5large.ait")
def mapping(req) -> dict:
    sentences = req.input if isinstance(req.input, list) else [req.input]
    return {"sentences": sentences}
def embed(sentences: list[str]) -> np.ndarray:
    return model.encode(sentences, normalize_embeddings=True, device="cuda")
config = dyn.DynamoWorkerConfig(
    type="embedding",
    model_path="intfloat/e5-large-v2",
    mapping=mapping,
)
dyn.dynamo_worker(embed, config)  # blocks until shutdown

API Reference

Import from aitune.dynamo:

import aitune.dynamo as dyn

DynamoWorkerConfig

dynamo_worker(model_or_fn, config, *, setup=None, warmup=None)

Functional API. Validates config, starts the Dynamo runtime, and blocks until shutdown.

• model_or_fn: any callable, or a torch.nn.Module (requires config.mapping)
• config: DynamoWorkerConfig
• setup: optional rank-local initialization callback
• warmup: optional warmup callback run only after setup succeeds on every rank

DynamoWorker (class-based API)

For more control, subclass DynamoWorker and override setup() and serve():

import aitune.dynamo as dyn
class MyEmbeddingWorker(dyn.DynamoWorker):
    def setup(self) -> None:
        # called once at startup — load or tune your model here
        self.model = load_my_model()
    async def serve(self, request):
        sentences = request.input if isinstance(request.input, list) else [request.input]
        embeddings = self.model.encode(sentences)
        yield embeddings
MyEmbeddingWorker().run()

Override on_ready(runtime, endpoint) for post-startup work such as custom register_model calls. Override warmup() to run model warmup after setup() succeeds on every rank and before endpoint registration.

Multi-GPU Workers

Use the same DynamoWorker API for local and multi-GPU models. Initialize the application-owned process group and construct the rank-local model before calling run() or dynamo_worker() on every rank. AITune starts the Dynamo endpoint only on rank 0, serializes and broadcasts incoming requests, executes the request on every rank, and returns only rank 0’s response.

For rank-local initialization that may fail, pass it through setup= and put any collective model warmup in warmup=. AITune exchanges setup failures across the process group before allowing any rank to start warmup or register the endpoint.

The worker does not initialize or destroy the process group and does not add inference barriers or CUDA synchronization. The application remains responsible for device placement, model sharding or context parallelism, and process-group teardown after the worker returns. An initialized multi-rank process group represents one collective model worker. Deploy independent model replicas as separate worker groups or pods.

Modality Types

Audio bytes are automatically packed as WAV output. To serve another audio codec, return a fully formed Dynamo response dict. Plain dictionaries are forwarded to the runtime as-is for every modality.

Serving with run_dynamo.sh

The recommended way to start all processes locally is a run_dynamo.sh script that:

1. Starts the Dynamo HTTP frontend in the background
2. Starts the backend worker in the background
3. Polls /health until the endpoint is registered
4. Runs a smoke-test client request

#!/bin/bash
export DYN_DISCOVERY_BACKEND=file
python -m dynamo.frontend --http-port 8000 &
FRONTEND_PID=$!
python -m myapp.dynamo.backend &
BACKEND_PID=$!
trap "kill -9 $FRONTEND_PID; kill -9 $BACKEND_PID" EXIT
for i in {1..10}; do
  curl -s http://localhost:8000/health | grep -q '"dyn://aitune.backend.generate"' && break
  echo "Waiting for endpoint... (attempt $i)"
  sleep 5
done
python -m myapp.dynamo.client

See the E5Large example for a complete working version.

Next Steps

• AOT Tuning Guide — tuning a model before saving
• Tune Strategies — selecting the right optimization strategy
• Backend Guides — backend-specific deployment notes
• E5Large example — end-to-end embedding worker
• FLUX example — end-to-end image generation worker
• WAN example — context-parallel text-to-video worker