Verwenden von Gerätetensoren in ONNX Runtime

Die Verwendung von Gerätetensoren kann ein entscheidender Teil beim Aufbau effizienter KI-Pipelines sein, insbesondere auf heterogenen Speichersystemen. Ein typisches Beispiel für solche Systeme ist jeder PC mit einer dedizierten GPU. Während eine aktuelle GPU selbst eine Speicherbandbreite von etwa 1 TB/s aufweist, kann die Verbindung PCI 4.0 x16 zur CPU oft der limitierende Faktor mit nur ca. 32 GB/s sein. Daher ist es oft am besten, Daten so weit wie möglich lokal auf der GPU zu halten oder langsamen Speicherverkehr hinter Berechnungen zu verstecken, da die GPU gleichzeitig Berechnungen und PCI-Speicherverkehr ausführen kann.

Ein typischer Anwendungsfall für diese Szenarien, bei denen der Speicher bereits lokal auf dem Inferenzgerät vorhanden ist, wäre die GPU-beschleunigte Videoverarbeitung eines kodierten Videostreams, der mit GPU-Dekodern dekodiert werden kann. Ein weiterer häufiger Fall sind iterative Netzwerke wie Diffusionsnetzwerke oder große Sprachmodelle, bei denen Zwischen-Tensoren nicht zurück an die CPU kopiert werden müssen. Kachelbasiertes Inferenz für hochauflösende Bilder ist ein weiterer Anwendungsfall, bei dem ein benutzerdefiniertes Speichermanagement wichtig ist, um GPU-Leerlaufzeiten während PCI-Kopien zu reduzieren. Anstatt jeden Kachel sequenziell zu verarbeiten, ist es möglich, PCI-Kopien und die Verarbeitung auf der GPU zu überlappen und die Arbeit auf diese Weise zu pipelinen.

Image of sequential PCI->Processing->PCI and another image of it being interleaved.

CUDA

CUDA in ONNX Runtime verfügt über zwei benutzerdefinierte Speichertypen. "CudaPinned" und "Cuda" Speicher, wobei CUDA-gepinnter Speicher tatsächlich CPU-Speicher ist, auf den die GPU direkt zugreifen kann, was vollständig asynchrone Up- und Downloads von Speicher mit cudaMemcpyAsync ermöglicht. Normale CPU-Tensoren erlauben nur synchrone Downloads von GPU zu CPU, während Kopien von CPU zu GPU immer asynchron ausgeführt werden können.

Das Allokieren eines Tensors mit dem Allokator von Ort::Sessions ist mit der C++ API sehr einfach, die direkt auf die C-API abgebildet wird.

Ort::Session session(ort_env, model_path_cstr, session_options);
Ort::MemoryInfo memory_info_cuda("Cuda", OrtArenaAllocator, /*device_id*/0,
                                 OrtMemTypeDefault);
Ort::Allocator gpu_allocator(session, memory_info_cuda);
auto ort_value = Ort::Value::CreateTensor(
        gpu_allocator, shape.data(), shape.size(),
        ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16);

Extern allokierte Daten können auch ohne Kopieren in ein Ort::Value verpackt werden

Ort::MemoryInfo memory_info_cuda("Cuda", OrtArenaAllocator, device_id,
                                 OrtMemTypeDefault);
std::array<int64_t, 4> shape{1, 4, 64, 64};
size_t cuda_buffer_size = 4 * 64 * 64 * sizeof(float);
void *cuda_resource;
CUDA_CHECK(cudaMalloc(&cuda_resource, cuda_buffer_size));
auto ort_value = Ort::Value::CreateTensor(
    memory_info_cuda, cuda_resource, cuda_buffer_size,
    shape.data(), shape.size(),
    ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT);

Diese allokierten Tensoren können dann als I/O-Bindings verwendet werden, um Kopiervorgänge im Netzwerk zu eliminieren und die Verantwortung auf den Benutzer zu übertragen. Mit solchen I/O-Bindings sind weitere Leistungsoptimierungen möglich

  • Aufgrund der festen Tensoradresse kann ein CUDA-Graph erfasst werden, um die Latenz beim CUDA-Start auf der CPU zu reduzieren
  • Durch entweder vollständig asynchrone Downloads in gepinnten Speicher oder die Eliminierung von Speicherkopien durch die Verwendung von geräte lokalem Tensor kann CUDA vollständig asynchron über eine Ausführungsoption auf seinem gegebenen Stream ausgeführt werden

Um den benutzerdefinierten Compute-Stream für CUDA festzulegen, beziehen Sie sich auf die V2-Option-API, die den undurchsichtigen Strukturzeiger Ort[CUDA|TensorRT]ProviderOptionsV2* und die Funktion Update[CUDA|TensorRT]ProviderOptionsWithValue(options, "user_compute_stream", cuda_stream); zur Festlegung seines Stream-Members verfügbar macht. Weitere Details finden Sie in der Dokumentation jedes Ausführungsanbieters.

Wenn Sie Ihre Optimierungen überprüfen möchten, hilft Nsight System bei der Korrelation von CPU-API und GPU-Ausführung von CUDA-Operationen. Dies ermöglicht auch die Überprüfung, ob die gewünschten Synchronisationen durchgeführt wurden und keine asynchrone Operation auf synchrone Ausführung zurückfällt. Es wird auch in diesem GTC-Vortrag verwendet, der die optimale Nutzung von Gerätetensoren erklärt.

Python API

Die Python-API unterstützt die gleichen Leistungsmöglichkeiten wie die oben genannte C++-API. Gerätetensoren können wie hier gezeigt allokiert werden. Darüber hinaus kann der user_compute_stream über diese API festgelegt werden

sess = onnxruntime.InferenceSession("model.onnx", providers=["TensorrtExecutionProvider"])
option = {}
s = torch.cuda.Stream()
option["user_compute_stream"] = str(s.cuda_stream)                    
sess.set_providers(["TensorrtExecutionProvider"], [option])

Das Aktivieren der asynchronen Ausführung in Python ist über die gleiche Ausführungsoption wie bei der C++-API möglich.

DirectML

Das Erreichen des gleichen Verhaltens ist über DirectX-Ressourcen möglich. Um eine asynchrone Verarbeitung auszuführen, ist es entscheidend, das gleiche Management von Ausführungs-Streams wie bei CUDA durchzuführen. Für DirectX bedeutet dies die Verwaltung des Geräts und seiner Befehlswarteschlange, was über die C-API möglich ist. Details zur Festlegung der Befehlswarteschlange für Berechnungen sind mit der Verwendung von SessionOptionsAppendExecutionProvider_DML1 dokumentiert.

Wenn separate Befehlswarteschlangen für Kopie und Berechnung verwendet werden, ist es möglich, PCI-Kopien und die Ausführung zu überlappen sowie die Ausführung asynchron zu gestalten.

#include <onnxruntime/dml_provider_factory.h>
Ort::MemoryInfo memory_info_dml("DML", OrtDeviceAllocator, device_id,
                                OrtMemTypeDefault);

std::array<int64_t, 4> shape{1, 4, 64, 64};
void *dml_resource;
size_t d3d_buffer_size = 4 * 64 * 64 * sizeof(float);
const OrtDmlApi *ort_dml_api;
Ort::ThrowOnError(Ort::GetApi().GetExecutionProviderApi(
                  "DML", ORT_API_VERSION, reinterpret_cast<const void **>(&ort_dml_api)));

// Create d3d_buffer using D3D12 APIs
Microsoft::WRL::ComPtr<ID3D12Resource> d3d_buffer = ...;

// Create the dml resource from the D3D resource.
ort_dml_api->CreateGPUAllocationFromD3DResource(d3d_buffer.Get(), &dml_resource);


Ort::Value ort_value(Ort::Value::CreateTensor(memory_info_dml, dml_resource,
                     d3d_buffer_size, shape.data(), shape.size(),
                     ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT));

Ein Einzeldatei-Beispiel ist auf GitHub zu finden, das zeigt, wie Befehlswarteschlangen für Kopie und Ausführung verwaltet und erstellt werden.

Python API

Obwohl die Zuweisung von DirectX-Eingaben aus Python möglicherweise kein Hauptanwendungsfall ist, ist die API verfügbar. Dies kann sich als sehr vorteilhaft erweisen, insbesondere für Zwischen-Netzwerk-Caches, wie z. B. Key-Value-Caching in großen Sprachmodellen (LLMs).

import onnxruntime as ort
import numpy as np

session = ort.InferenceSession("model.onnx",
                               providers=["DmlExecutionProvider"])

cpu_array = np.zeros((1, 4, 512, 512), dtype=np.float32)
dml_array = ort.OrtValue.ortvalue_from_numpy(cpu_array, "dml")

binding = session.io_binding()
binding.bind_ortvalue_input("data", dml_array)
binding.bind_output("out", "dml")
# if the output dims are known we can also bind a preallocated value
# binding.bind_ortvalue_output("out", dml_array_out)

session.run_with_iobinding(binding)