AOTInductor External Weight Storage and Weight Streaming Update

A new GitHub example demonstrates how to store AOTInductor model weights outside the shared library and update them at inference runtime in a thread-safe manner, addressing the lack of documentation for advanced AOTInductor features. The example, created by Lei Mao, sets aot_inductor.package_constants_in_so to False and freezing to False to keep weights updatable, and is available at https://github.com/leimao/AOTInductor-External-Weight-Storage-Weight-Streaming-Update-Example.

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/