AOTInductor External Weight Storage and Weight Streaming Update

Introduction

AOTInductor is a PyTorch compiler backend that compiles PyTorch models into optimized shared libraries for efficient inference, which is similar to NVIDIA TensorRT. However, compared to NVIDIA TensorRT, AOTInductor is somewhat less well documented. It is unclear how to use some advanced features with AOTInductor, such as model weight updates at inference runtime.

In this blog post, I would like to share an example of how to store the model weights outside the AOTInductor shared library so file, and how to update the model weights at inference runtime in a thread-safe fashion.

AOTInductor Weight Storage and Update

The example is based on a simple PyTorch model with a single linear layer. The AOTInductor engine and the model weights are compiled and packaged using a Python script. The AOTInductor engine inference and runtime weight updates are performed using a C++ program. The example is available on GitHub.

In some computer platforms, there might be some problems if the shared library so file is too large. Consequently, we would like to store the model weights outside the AOTInductor shared library so file, which is archived in the pt2 file. To enable this, we have to set the aot_inductor.package_constants_in_so configuration to False when compiling the AOTInductor model. The always_keep_tensor_constants configuration is also set to True to ensure that the model constants are always kept as updatable constants, even if they are small. In our case, if we did not set it, the bias term of the linear layer would be folded into the graph and would not be exposed as an updatable constant.

The AOTInductor configuration freezing is also set to False to allow model constants to be updated. Otherwise, the model constants will be frozen as model attributes and its layouts might be changed to a layout that is more efficient for inference, and they can no longer be updated according to the configuration comment. For example, in convolution layers, the weight layout might be changed from NCHW in the PyTorch model to NHWC in the AOTInductor model, because CUDA kernels favor the NHWC layout for better memory coalescing. Of course, since we disabled the freezing configuration, some performance optimization opportunities might be lost.

export_aoti_artifacts.py

import os

import torch
import torch._inductor

DEVICE = "cuda"
assert torch.cuda.is_available(), "CUDA is required to run this script."


# Define a PyTorch Module
class SampleModel(torch.nn.Module):

    def __init__(self):
        super().__init__()
        self.fc = torch.nn.Linear(4, 2)

    def forward(self, x):
        return self.fc(x)


OUTPUT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                          "aoti_package")


# torch::pickle_load in C++ hardcodes the archive's internal folder name to
# "data", which only happens when the saved file itself is named "data.<ext>".
#
# Weight dicts are saved on CPU: the C++ side updates the engine's inactive
# GPU buffer with allow_h2d_copy=True, copying straight from host memory. If
# we instead handed it GPU-resident tensors, the engine would still make its
# own internal copy into the inactive buffer, leaving 3 GPU-resident copies
# alive at once (our tensor + the new inactive buffer + the old active
# buffer) instead of 2.
def save_state_dict(state_dict, subdir, to_cpu=False):
    directory = os.path.join(OUTPUT_DIR, subdir)
    os.makedirs(directory, exist_ok=True)
    if to_cpu:
        state_dict = {k: v.cpu() for k, v in state_dict.items()}
    torch.save(state_dict, os.path.join(directory, "data.pt"))


def main():
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    model = SampleModel().to(DEVICE).eval()
    x = torch.randn(2, 4, device=DEVICE)

    with torch.no_grad():
        orig_output = model(x)

    # Export the model and compile it ahead-of-time with AOTInductor.
    exported_program = torch.export.export(model, (x, ))
    package_path = torch._inductor.aoti_compile_and_package(
        exported_program,
        package_path=os.path.join(OUTPUT_DIR, "model.pt2"),
        inductor_configs={
            "aot_inductor.package_constants_in_so": True,
            "freezing": False,
            # Without this, small constants like the bias get folded into
            # the graph and are not exposed as updatable constants.
            "always_keep_tensor_constants": True,
        },
    )

    # Dump the initial weights so the C++ side can load them via load_constants.
    weights_dict = {
        **dict(model.named_parameters()),
        **dict(model.named_buffers()),
    }
    save_state_dict(weights_dict, "weights_initial", to_cpu=True)

    # Verify the compiled AOTI package matches PyTorch before handing it off
    # to the C++ side.
    #
    # NOTE: swap_constant_buffer()/free_inactive_constant_buffer() (needed for
    # a safe double-buffered swap) are only exposed on the C++
    # AOTIModelContainerRunner, not on the Python AOTIModelPackageLoader
    # binding used by aoti_load_package(), so this Python check updates the
    # active buffer directly via load_constants(), which (unlike the
    # lower-level update_constant_buffer()) translates FQN keys (e.g.
    # "fc.weight") to the engine's internal constant names. See
    # aoti_cpp_inference.cpp for the safe swap.
    compiled_model = torch._inductor.aoti_load_package(package_path)
    print(f"AOTI engine constant FQNs: {compiled_model.get_constant_fqns()}")
    compiled_model.load_constants(weights_dict,
                                  check_full_update=True,
                                  allow_h2d_copy=True)
    aoti_output = compiled_model(x)
    assert torch.allclose(orig_output, aoti_output, rtol=1e-5, atol=1e-6)
    print("✅ Initial AOTI inference succeeded and matches PyTorch output.")

    # Update the weights dynamically and dump them too, to demonstrate
    # updating the constants of an already-loaded AOTI model from C++.
    with torch.no_grad():
        model.fc.weight.add_(1.0)
        model.fc.bias.add_(0.5)
        updated_orig_output = model(x)

    updated_weights_dict = {
        **dict(model.named_parameters()),
        **dict(model.named_buffers()),
    }
    save_state_dict(updated_weights_dict, "weights_updated", to_cpu=True)

    # Verify the updated constants also match PyTorch.
    compiled_model.load_constants(updated_weights_dict,
                                  check_full_update=True,
                                  allow_h2d_copy=True)
    updated_aoti_output = compiled_model(x)
    assert torch.allclose(updated_orig_output,
                          updated_aoti_output,
                          rtol=1e-5,
                          atol=1e-6)
    print(
        "✅ Weight update verified successfully! Updated AOTI output matches PyTorch."
    )

    # Dump the sample input and expected outputs so the C++ side can verify
    # its AOTI inference results without needing PyTorch's eager model.
    save_state_dict(
        {
            "x": x,
            "orig_output": orig_output,
            "updated_orig_output": updated_orig_output
        },
        "io_data",
    )

    print(f"AOTI package and weights saved under: {OUTPUT_DIR}")
    print(f"  - {os.path.join(OUTPUT_DIR, 'model.pt2')}")
    print(f"  - {os.path.join(OUTPUT_DIR, 'weights_initial', 'data.pt')}")
    print(f"  - {os.path.join(OUTPUT_DIR, 'weights_updated', 'data.pt')}")
    print(f"  - {os.path.join(OUTPUT_DIR, 'io_data', 'data.pt')}")


if __name__ == "__main__":
    main()

Some models being served online would require hot weight updates, and ideally we would like to stream the updated weights from the corresponding PyTorch model without additional processing to the AOTInductor engine. In addition, it is also extremely important to do this in a thread-safe fashion because in online serving, multiple threads may be performing inference simultaneously on one AOTInductor engine instance. To support this, we will use the load_constants API to load the updated weights into the inactive buffer, and then swap the inactive buffer with the active buffer. The swap_constant_buffer and free_inactive_constant_buffer APIs are used to perform this double-buffered swap safely. The load_constants API is used instead of the lower-level update_constant_buffer because it translates the fully qualified names (FQNs) of the constants (e.g., “fc.weight”) to the engine’s internal mangled constant names.

aoti_cpp_inference.cpp

// AOTI C++ inference example.
//
// This program loads an AOTInductor-compiled model package together with
// externally-dumped weights (produced by export_aoti_artifacts.py) and runs
// inference, demonstrating how to update the model's constants at runtime.
#include <torch/csrc/inductor/aoti_package/model_package_loader.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAStream.h>
#include <torch/cuda.h>
#include <torch/torch.h>

#include <fstream>
#include <iostream>
#include <string>
#include <unordered_map>
#include <vector>

namespace
{

// torch::pickle_load() expects the zip archive's internal root folder to be
// named "data", which is only the case when the file was saved by Python as
// "data.<ext>" (see save_state_dict() in export_aoti_artifacts.py).
std::unordered_map<std::string, at::Tensor>
load_tensor_dict(const std::string& path)
{
    std::ifstream file(path, std::ios::binary);
    TORCH_CHECK(file.is_open(), "Failed to open file: ", path);
    std::vector<char> data((std::istreambuf_iterator<char>(file)),
                           std::istreambuf_iterator<char>());

    torch::IValue ivalue = torch::pickle_load(data);
    std::unordered_map<std::string, at::Tensor> result;
    for (const auto& item : ivalue.toGenericDict())
    {
        result.emplace(item.key().toStringRef(), item.value().toTensor());
    }
    return result;
}

} // namespace

int main(int argc, char* argv[])
{
    const std::string artifacts_dir = argc > 1 ? argv[1] : "aoti_package";
    const std::string package_path = artifacts_dir + "/model.pt2";

    torch::inductor::AOTIModelPackageLoader loader(package_path,
                                                   /*model_name=*/"model",
                                                   /*run_single_threaded=*/false,
                                                   /*num_runners=*/1);
    // swap_constant_buffer()/free_inactive_constant_buffer() are only
    // exposed on the runner, not directly on AOTIModelPackageLoader.
    torch::inductor::AOTIModelContainerRunner* runner = loader.get_runner();

    // Run inference on an explicit stream (instead of the default one) and
    // hand its handle to run().
    at::cuda::CUDAStream stream = at::cuda::getStreamFromPool();
    at::cuda::CUDAStreamGuard stream_guard(stream);
    void* stream_handle = stream.stream();

    auto io_data = load_tensor_dict(artifacts_dir + "/io_data/data.pt");
    const at::Tensor& x = io_data.at("x");
    const at::Tensor& orig_output = io_data.at("orig_output");
    const at::Tensor& updated_orig_output = io_data.at("updated_orig_output");

    // 1. First-time load: nothing is being served yet, so there's no active
    // buffer worth preserving -- load the weights straight into the active
    // buffer instead of going through the inactive-buffer/swap dance.
    //
    // load_constants() (not the lower-level update_constant_buffer()) is
    // used because it translates the FQN keys (e.g. "fc.weight") to the
    // engine's internal mangled constant names; update_constant_buffer()
    // expects those internal names directly.
    //
    // The dumped weights are CPU tensors, so allow_h2d_copy=true lets the
    // engine copy them straight from host memory into its GPU-resident
    // buffer. If we instead loaded them as GPU tensors here, the engine
    // would still make its own copy into its buffer, leaving 2 GPU-resident
    // copies alive at once (ours + the engine's) instead of 1.
    auto initial_weights =
        load_tensor_dict(artifacts_dir + "/weights_initial/data.pt");
    loader.load_constants(initial_weights,
                          /*use_inactive=*/false,
                          /*check_full_update=*/true,
                          /*user_managed=*/false,
                          /*allow_h2d_copy=*/true);
    // The H2D copy above is asynchronous, and with multiple runners/streams
    // potentially consuming the shared constant buffer, a device-wide sync
    // (not just our own stream) is needed to guarantee it has completed
    // before any of them run inference against it.
    torch::cuda::synchronize();

    auto initial_outputs = loader.run({x}, stream_handle);
    // run() executed on our stream asynchronously; synchronize it before
    // reading the results on the CPU below.
    stream.synchronize();
    TORCH_CHECK(
        torch::allclose(initial_outputs[0], orig_output, /*rtol=*/1e-5,
                        /*atol=*/1e-6),
        "Initial AOTI output does not match the expected PyTorch output.");
    std::cout << "Initial AOTI inference succeeded and matches PyTorch output."
              << std::endl;

    // 2. Hot weight update: the model could already be serving inference at
    // this point, so load into the inactive buffer, atomically swap it in,
    // and free the now-inactive (old) buffer instead of overwriting the
    // active buffer in place.
    auto updated_weights =
        load_tensor_dict(artifacts_dir + "/weights_updated/data.pt");
    loader.load_constants(updated_weights,
                          /*use_inactive=*/true,
                          /*check_full_update=*/true,
                          /*user_managed=*/false,
                          /*allow_h2d_copy=*/true);
    runner->swap_constant_buffer();
    runner->free_inactive_constant_buffer();
    // Same reasoning as above: synchronize the whole device, not just our
    // stream, before running inference against the swapped-in buffer.
    torch::cuda::synchronize();

    auto updated_outputs = loader.run({x}, stream_handle);
    // Same reasoning as above: wait for our stream before reading results.
    stream.synchronize();
    TORCH_CHECK(
        torch::allclose(updated_outputs[0], updated_orig_output, /*rtol=*/1e-5,
                        /*atol=*/1e-6),
        "Updated AOTI output does not match the expected PyTorch output.");
    std::cout << "Weight update verified successfully! Updated AOTI output "
                 "matches PyTorch."
              << std::endl;

    return 0;
}

Also note that our buffer is stored on host memory initially, and the AOTInductor engine will copy it to the GPU inactive buffer when load_constants is called with allow_h2d_copy=True, because we set user_managed=False, which means the engine will manage the GPU buffer’s lifetime and allocation.

If our buffer is stored on GPU memory, we can set user_managed=True to let the engine use our buffer directly without making a copy. Otherwise, it would behave like having triple-buffering on GPU memory, and in some cases, for large models, GPU would not be able to tolerate such peak memory usage.

References

• AOTInductor Weight External Storage and Weight Streaming Update Example - GitHub
• TorchInductor Config