# AOTInductor External Weight Storage and Weight Streaming Update

> Source: <https://leimao.github.io/blog/AOTInductor-External-Weight-Storage-Weight-Streaming-Update/>
> Published: 2026-08-27 04:48:11.013704+00:00

# 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](https://github.com/leimao/AOTInductor-External-Weight-Storage-Weight-Streaming-Update-Example).

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](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/_inductor/config.py#L1559). 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.

```
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137
```

 | 

``` python
import osimport torchimport torch._inductorDEVICE = "cuda"assert torch.cuda.is_available(), "CUDA is required to run this script."# Define a PyTorch Moduleclass 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.

```
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134
```

 | 

```
// 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;}} // namespaceint 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 External Weight Storage and Weight Streaming Update

[https://leimao.github.io/blog/AOTInductor-External-Weight-Storage-Weight-Streaming-Update/](https://leimao.github.io/blog/AOTInductor-External-Weight-Storage-Weight-Streaming-Update/)