AOTInductor: компиляция моделей Torch.Export-ed до выполнения

Компиляция модели

Чтобы скомпилировать модель с помощью AOTInductor, сначала необходимо использовать torch.export.export(), чтобы преобразовать заданную модель PyTorch в вычислительный граф. torch.export гарантирует корректность и строго специфицирует захваченный IR, на который опирается AOTInductor.

Затем мы используем torch._inductor.aoti_compile_and_package(), чтобы скомпилировать экспортированную программу с помощью TorchInductor и сохранить скомпилированные артефакты в одном пакете. Пакет имеет формат спецификации архива PT2.

import os
import torch

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = torch.nn.Linear(10, 16)
        self.relu = torch.nn.ReLU()
        self.fc2 = torch.nn.Linear(16, 1)
        self.sigmoid = torch.nn.Sigmoid()

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        x = self.sigmoid(x)
        return x

with torch.no_grad():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = Model().to(device=device)
    example_inputs=(torch.randn(8, 10, device=device),)
    batch_dim = torch.export.Dim("batch", min=1, max=1024)
    # [Optional] Specify the first dimension of the input x as dynamic.
    exported = torch.export.export(model, example_inputs, dynamic_shapes={"x": {0: batch_dim}})
    # [Note] In this example we directly feed the exported module to aoti_compile_and_package.
    # Depending on your use case, e.g. if your training platform and inference platform
    # are different, you may choose to save the exported model using torch.export.save and
    # then load it back using torch.export.load on your inference platform to run AOT compilation.
    output_path = torch._inductor.aoti_compile_and_package(
        exported,
        # [Optional] Specify the generated shared library path. If not specified,
        # the generated artifact is stored in your system temp directory.
        package_path=os.path.join(os.getcwd(), "model.pt2"),
        # [Optional] Specify Inductor configs
        # This specific max_autotune option will turn on more extensive kernel autotuning for
        # better performance.
        inductor_configs={"max_autotune": True,},
    )

В этом примере параметр Dim используется, чтобы указать первое измерение входной переменной «x» как динамическое. Обратите внимание: путь и имя скомпилированной библиотеки не указаны, поэтому общая библиотека сохраняется во временном каталоге. Чтобы получить этот путь со стороны C++, мы сохраняем его в файл, чтобы позднее прочитать его в коде C++.

Инференс в Python

Существует несколько способов развернуть скомпилированный артефакт для инференса; один из них — использовать Python. В Python доступен удобный API-метод torch._inductor.aoti_load_package() для загрузки и запуска артефакта, как показано в следующем примере:

import os
import torch

device = "cuda" if torch.cuda.is_available() else "cpu"
model = torch._inductor.aoti_load_package(os.path.join(os.getcwd(), "model.pt2"))
print(model(torch.randn(8, 10, device=device)))

Во время инференса размер, тип данных и шаг входных данных должны совпадать с соответствующими характеристиками входных данных во время экспорта.

Инференс в C++

Далее мы используем следующий пример файла C++ inference.cpp для загрузки скомпилированного артефакта, что позволяет выполнять предсказания модели непосредственно в среде C++.

#include <iostream>
#include <vector>

#include <torch/torch.h>
#include <torch/csrc/inductor/aoti_package/model_package_loader.h>

int main() {
    c10::InferenceMode mode;

    torch::inductor::AOTIModelPackageLoader loader("model.pt2");
    // Assume running on CUDA
    std::vector<torch::Tensor> inputs = {torch::randn({8, 10}, at::kCUDA)};
    std::vector<torch::Tensor> outputs = loader.run(inputs);
    std::cout << "Result from the first inference:"<< std::endl;
    std::cout << outputs[0] << std::endl;

    // The second inference uses a different batch size and it works because we
    // specified that dimension as dynamic when compiling model.pt2.
    std::cout << "Result from the second inference:"<< std::endl;
    // Assume running on CUDA
    std::cout << loader.run({torch::randn({1, 10}, at::kCUDA)})[0] << std::endl;

    return 0;
}

Для сборки файла C++ можно воспользоваться предоставленным файлом CMakeLists.txt, который автоматизирует вызов python model.py для AOT-компиляции модели и компиляцию inference.cpp в исполняемый бинарный файл с именем aoti_example.

cmake_minimum_required(VERSION 3.18 FATAL_ERROR)
project(aoti_example)

find_package(Torch REQUIRED)

add_executable(aoti_example inference.cpp model.pt2)

add_custom_command(
    OUTPUT model.pt2
    COMMAND python ${CMAKE_CURRENT_SOURCE_DIR}/model.py
    DEPENDS model.py
)

target_link_libraries(aoti_example "${TORCH_LIBRARIES}")
set_property(TARGET aoti_example PROPERTY CXX_STANDARD 17)

Если структура каталогов выглядит следующим образом, для сборки бинарного файла можно выполнить приведенные ниже команды. Обратите внимание, что переменная CMAKE_PREFIX_PATH необходима CMake для поиска библиотеки LibTorch и должна содержать абсолютный путь. Учтите, что ваш путь может отличаться от указанного в этом примере.

aoti_example/
    CMakeLists.txt
    inference.cpp
    model.py

$ mkdir build
$ cd build
$ CMAKE_PREFIX_PATH=/path/to/python/install/site-packages/torch/share/cmake cmake ..
$ cmake --build . --config Release

После создания бинарного файла aoti_example в каталоге build его запуск выведет результаты, похожие на следующие:

$ ./aoti_example
Result from the first inference:
0.4866
0.5184
0.4462
0.4611
0.4744
0.4811
0.4938
0.4193
[ CUDAFloatType{8,1} ]
Result from the second inference:
0.4883
0.4703
[ CUDAFloatType{2,1} ]

Устранение неполадок

Ниже перечислены полезные инструменты для отладки AOT Inductor.

Инструменты отладки

• torch._logging
• Минимизатор AOTInductor
• Руководство по отладке AOTInductor

Чтобы включить проверки входных данных во время выполнения, задайте для переменной окружения AOTI_RUNTIME_CHECK_INPUTS значение 1. Если размер, тип данных или шаги входных данных скомпилированной модели отличаются от использованных при экспорте, будет вызвано исключение RuntimeError.

Справочник по API

torch._inductor.aoti_compile_and_package(exported_program, _deprecated_unused_args=None, _deprecated_unused_kwargs=None, *, package_path=None, inductor_configs=None) [исходный код]

Компилирует экспортированную программу с помощью AOTInductor и упаковывает ее в артефакт .pt2, указанный в package_path. Чтобы загрузить пакет, вызовите torch._inductor.aoti_load_package(package_path).

Пример использования:

ep = torch.export.export(M(), ...)
aoti_file = torch._inductor.aoti_compile_and_package(
    ep, package_path="my_package.pt2"
)
compiled_model = torch._inductor.aoti_load_package("my_package.pt2")

Чтобы скомпилировать и сохранить несколько моделей в одном артефакте .pt2, выполните следующие действия:

ep1 = torch.export.export(M1(), ...)
aoti_file1 = torch._inductor.aot_compile(
    ep1, ..., options={"aot_inductor.package": True}
)
ep2 = torch.export.export(M2(), ...)
aoti_file2 = torch._inductor.aot_compile(
    ep2, ..., options={"aot_inductor.package": True}
)

from torch._inductor.package import package_aoti, load_package

package_aoti("my_package.pt2", {"model1": aoti_file1, "model2": aoti_file2})

compiled_model1 = load_package("my_package.pt2", "model1")
compiled_model2 = load_package("my_package.pt2", "model2")

Параметры:

• exported_program (ExportedProgram) – Экспортированная программа, созданная вызовом torch.export
• package_path (FileLike | None) – Необязательный путь к создаваемому артефакту .pt2.
• inductor_configs (dict[str, Any] | None) – Необязательный словарь конфигураций для управления inductor.

Возвращает:

Путь к созданному артефакту

Тип возвращаемого значения:

str

torch._inductor.aoti_load_package(path, run_single_threaded=False, device_index=-1) [исходный код]

Загружает модель из пакета PT2.

Если в пакет PT2 было упаковано несколько моделей, будет загружена модель по умолчанию. Чтобы загрузить определенную модель, можно напрямую вызвать API загрузки.

from torch._inductor.package import load_package

compiled_model1 = load_package("my_package.pt2", "model1")
compiled_model2 = load_package("my_package.pt2", "model2")

Параметры:

• path (FileLike) – Путь к пакету .pt2
• run_single_threaded (bool) – Нужно ли запускать модель без логики синхронизации потоков. Это полезно для предотвращения конфликтов с CUDAGraphs.
• device_index (int) – Индекс устройства, на которое следует загрузить пакет PT2. По умолчанию используется device_index=-1, соответствующий устройству cuda при использовании CUDA. Например, передача device_index=1 загрузит пакет на cuda:1.

Тип возвращаемого значения:

AOTICompiledModel