Machine Learning: Programmiersprache Triton von OpenAI ist für GPUs optimiert

Die Open-Source-Programmiersprache zum Erstellen von Deep-Learning-Anwendungen setzt auf Python und automatisiert die Optimierung der GPU-Architektur.

Lesezeit: 3 Min.
In Pocket speichern
vorlesen Druckansicht Kommentare lesen 5 Beiträge
Künstliche Intelligenz als Hilfswissenschaftler: KI revolutioniert die Forschung

(Bild: Phonlamai Photo/Shutterstock.com)

Von
  • Rainald Menge-Sonnentag

OpenAI hat mit Triton eine neue Programmiersprache vorgestellt, die speziell für Machine-Learning-Anwendungen (ML) ausgelegt ist. Die nun in Version 1.0 veröffentlichte Open-Source-Sprache baut auf das im ML-Umfeld verbreitete Python. Unter der Haube zielt sie auf Optimierungen für die Architektur der GPU – derzeit ausschließlich aus dem Hause Nvidia.

Motivation zur Entwicklung der 2019 vorgestellten Sprache ist laut dem OpenAI-Blog, dass ML-Frameworks auf der einen Seite nicht performant genug für umfangreiche künstliche neuronale Netze (KNN) seien. Die hardwarenahe GPU-Programmierung beispielsweise mit Nvidias CUDA hat auf der anderen Seite eine hohe Einstiegshürde und stellt hohe Anforderungen an die manuelle Optimierung der Prozesse.

Um Software möglichst effizient auf GPUs auszuführen, muss der Code auf die Architektur zugeschnitten sein. So sorgt das Memory Coalescing dafür, dass Speichertransfers vom DRAM in große Transaktionen zusammengefasst werden. Außerdem gilt es das Shared-Memory-Management für die im SRAM gespeicherten Daten zu optimieren. Schließlich müssen die Berechnungen sowohl innerhalb einzelner Streaming Multiprocessors (SMs) als auch SM-übergreifend partitioniert und disponiert werden.

Die GPU-Architektur erfordert spezielle Codeanpassungen zum Optimieren von Anwendungen.

Während Entwicklerinnen und Entwickler unter CUDA die zugehörigen Compiler-Optimierungen manuell vornehmen müssen, kümmert sich der Triton-Compiler offensichtlich automatisch um das Memory Coalescing, das Shared-Memory-Management und das Scheduling innerhalb der SMs. Lediglich für das übergreifende Scheduling sind manuelle Anpassungen erforderlich.

Von der Syntax her setzt Triton auf Python. Beim Umsetzen der spezifischen GPU-Funktionen ähnelt es wohl dem Softwarepaket Numba, das numerische Funktionen mit einem JIT-Compiler (Just-in Time) in Maschinencode übersetzt. Wie Numba verwendet Triton dekorierte Python-Funktionen. Darüber lassen sich Kernels definieren, die parallel mit unterschiedlichen program_ids in einem Grid aus Instanzen laufen.

Im Gegensatz zu Numba kennt Triton die Zeigerarithmetik. Folgendes Codebeispiel zeigt die einfachste Umsetzung eines Kernels, der alle Blockoperationen in einem einzelnen Thread durchführt und Zeigerarithmetik verwendet:

BLOCK = 512

# This is a GPU kernel in Triton.
# Different instances of this
# function may run in parallel.
@jit
def add(X, Y, Z, N):
   # In Triton, each kernel instance
   # executes block operations on a
   # single thread: there is no construct
   # analogous to threadIdx
   pid = program_id(0)
   # block of indices
   idx = pid * BLOCK + arange(BLOCK)
   mask = idx < N
   # Triton uses pointer arithmetics  
   # rather than indexing operators
   x = load(X + idx, mask=mask)
   y = load(Y + idx, mask=mask)
   store(Z + idx, x + y, mask=mask)


...
grid = (ceil_div(N, BLOCK),)
# no thread-block
add[grid](x, y, z, x.shape[0])

Aus Performancegründen setzt Triton auf eine modulare Systemarchitektur, in deren Mittelpunkt mit Triton-IR ein Zwischencode steht (Intermediate Representation, IR). In der IR spielen mehrdimensionale Werteblöcken eine zentrale Rolle. Das Compiler-Backend von Triton vereinfacht und optimiert den Zwischencode, um ihn für die Compilerarchitektur LLVM aufzubereiten und letztlich auf PTX (Parallel Thread Execution) zuzuschneiden. In Version 1.0 zielen alle Optimierungen auf die GPUs von Nvidia, und das Kompilieren für CPUs oder AMD-GPUs ist derzeit nicht vorgesehen.

(Bild: OpenAI)

Weitere Details lassen sich einem Beitrag auf dem OpenAI-Blog entnehmen, den Philippe Tillet verfasst hat. Er hat die Sprache maßgeblich erfunden und 2019 noch vor seiner Zeit als OpenAI-Mitarbeiter vorgestellt. OpenAI entstand 2015 als Forschungsprojekt, und zu den Gründern gehörte Elon Musk. Anfangs entstanden viele Projekte im Bereich Reinforcement Learning. Große Bekanntheit erlangte das Unternehmen mit dem Sprachmodell Generative Pre-trained Transformer 3 (GPT-3). 2019 verabschiedete sich OpenAI vom reinen Non-Profit-Business und investiert inzwischen selbst in Start-ups im Bereich Machine Learning – zuletzt mit 100 Millionen US-Dollar.

Das Triton-Repository findet sich auf GitHub, und Tillet lädt ausdrücklich zum Forken ein. Außerdem bittet er um Mitarbeit an der Umsetzung für Plattformen jenseits der Nvidia-GPUs.

(rme)