Im vorheriger Artikel dieser Serie, Operation in allen Bereichen der Informatik: Matrixmultiplikation. Es wird häufig in neuronalen Netzen verwendet, um die Aktivierung linearer Schichten zu berechnen. Aktivierungen allein sind jedoch schwer zu interpretieren, da ihre Werte und Statistiken (Mittelwert, Varianz, Min-Max-Amplitude) von Schicht zu Schicht stark variieren können. Dies ist einer der Gründe, warum wir Aktivierungsfunktionen verwenden, zum Beispiel die Logistikfunktion (auch bekannt als Sigmoid), die jede reelle Zahl in projiziert (0; 1) Reichweite.

Die Softmax-Funktion, auch normalisierte Exponentialfunktion genannt, ist eine mehrdimensionale Verallgemeinerung des Sigmoids. Es wandelt einen Vektor von Rohwerten (Logits) in einen um Wahrscheinlichkeitsverteilung über M Klassen. Wir können es als a interpretieren gewichteter Durchschnitt das verhält sich wie ein reibungslose Funktion und kann bequem sein differenziert. Es ist eine entscheidende Komponente der Skalarproduktaufmerksamkeit, der Sprachmodellierung und der multinomialen logistischen Regression.

In diesem Artikel behandeln wir Folgendes:

1. Implementierung eines effizienten Softmax-Kernels in Triton.
2. Implementierung des Rückwärtsdurchlaufs (autograd).
3. Optimierung: Cache-Modifikatoren und Auto-Tuning.

Wenn Sie Triton noch nicht kennen, lesen Sie die vorherigen Artikel!

Haftungsausschluss: Sofern nicht anders angegeben, stammen alle Illustrationen und Animationen vom Autor.

Definition

Der Softmax ist wie folgt definiert:

Die Normalisierung stellt sicher, dass sich der Vektor summiert 1sodass sie als gültige Wahrscheinlichkeitsverteilung interpretiert werden kann.

Beachten Sie, dass diese Formulierung des Softmax sehr empfindlich reagiert numerischer Überlauf. Denken Sie daran, dass der Maximalwert ein Commonplace ist float16 darstellen kann ist 65 504was ungefähr ist exp(11). Dies bedeutet, dass jeder Eingabewert größer als ~11 zu führt exp(z_i) Überschreiten des darstellbaren Bereichs, was dazu führt Überlauf.

Ein gängiger Trick, um dieses Drawback zu entschärfen, besteht darin, den Maximalwert des Eingabevektors von jedem Factor zu subtrahieren, sodass sich das neue Most ergibt 0 vor Potenzierung und 1 nach.

Naive Implementierung

Wie Sie sehen, ist die Berechnung des Softmax erforderlich zwei ReduktionsoperationenA max und a Summe. Ein naiver Algorithmus erfordert drei separate Durchgänge über den Eingabevektor. Zuerst wird das Most berechnet, dann die Summe und schließlich die normalisierten Ausgaben.

So sieht eine naive Numpy-Implementierung aus:

Ein wiederkehrendes Thema in dieser Triton-Serie ist die Minimierung hoher Latenzzeiten globaler Speicherzugriff. Unsere aktuelle Numpy-Implementierung erfordert drei separate Speicherlesevorgänge des vollständigen Eingabevektors, was äußerst ineffizient ist.

On-line-Softmax

Glücklicherweise können wir einen cleveren Trick anwenden, der sogenannte On-line-Softmaxum die zu verschmelzen max Und sum Schritte, wodurch die Anzahl der Speicherlesevorgänge reduziert wird 2.

Zunächst definieren wir die Summe der Exponentialfunktionen rekursiv. In der folgenden Reihe von Gleichungen gilt: m_i bezieht sich auf das Most über x bis zum ich-ter Index.

Diese Gleichheit ermöglicht es uns, die Summe der Exponentialfunktionen zu berechnen iterativ unter Verwendung des Maximalwerts bis jetzt. Wir können es nutzen, um die erste und zweite Schleife in der naiven Implementierung zu verschmelzen und das Most und die Summe der Exponentialfunktionen iterativ zu berechnen.

Unser Algorithmus wird:

Dies lässt sich leicht in Numpy übersetzen:

Nachdem wir nun die Hauptprinzipien hinter dem Softmax verstanden haben, werden wir es in Triton implementieren, angefangen bei der einfachen Einzelblockversion bis hin zur On-line-Mehrblockformulierung. Letztendlich möchten wir, dass sich unser Kernel wie ein PyTorch-Modul verhält und mit diesem kompatibel ist autograd.

Leider verhalten sich Triton-Kernel aus Sicht von PyTorch wie Black Packing containers: Die von ihnen ausgeführten Operationen werden nicht nachverfolgt autograd. Dies erfordert, dass wir den Rückwärtsdurchlauf selbst implementieren und explizit angeben, wie Gradienten berechnet werden sollen. Lassen Sie uns unsere geliebte Kettenregel auffrischen und den Softmax-Gradienten ableiten.

Gradient

Da die Ausgaben des Softmax streng positiv sind, können wir das verwenden logarithmische Ableitung um die Ableitung des Gradienten zu erleichtern. Hier nehmen wir die Ableitung von Protokoll der Ausgabe und wenden Sie die Kettenregel an:

Von dort aus ordnen wir die Begriffe neu und folgen diesen Schritten:

Nehmen wir nun an, dass wir einen Upstream-Gradienten haben, der beispielsweise durch eine Verlustfunktion erzeugt wird L (zB ein Kreuzentropieverlust). Wir erhalten den folgenden Ausdruck des Gradienten:

Die Vereinfachung des linken Begriffs in (9) liegt daran, dass δ_ij wird nur gleich sein 1 für die ich-tes Factor, wodurch die Summe zusammenfällt J auf einen einzigen Begriff.

Triton-Implementierung

Einzelblock Softmax

Nachdem wir nun die Ableitung des Gradienten durchgearbeitet haben, können wir die Vorwärts- und Rückwärts-Softmax-Kernel schreiben. Konzentrieren wir uns zunächst auf den PyTorch-Wrapper, um zu verstehen, wie die Einzelblockimplementierung auf hoher Ebene funktioniert. Bei einem 2D-Eingabetensor verarbeiten die Vorwärts- und Rückwärtskerne alle Zeilen parallel.

Der Einfachheit halber definieren wir die BLOCK_SIZE groß genug sein, um alle Spalten gleichzeitig zu verarbeiten. Konkret legen wir es als die nächste Potenz von 2 fest, die über der Anzahl der Spalten liegt, wie von Triton gefordert.

Dann definieren wir unser „Raster“ als die Anzahl der Zeilen (es könnte möglicherweise auch eine Batch-Dimension verarbeiten).

Der PyTorch-Wrapper für unsere SoftmaxSingleBlock ist eine Klasse, von der erbt torch.autograd.Perform das umsetzt ahead Und backward. Beide Methoden dauern eine ctx Argument, das wir verwenden werden, um die Softmax-Ausgaben während des Vorwärtsdurchlaufs zwischenzuspeichern und sie während des Rückwärtsdurchlaufs wiederzuverwenden.

Beide Kernel sind ziemlich unkompliziert. Wir beginnen mit dem Laden der Zeileneingaben mit derselben Syntax wie in meinem vorherigen Vektoraddition Artikel. Beachten Sie das BLOCK_SIZE Und num_warps werden mit a berechnet calculate_settings Funktion. Diese Funktion stammt aus dem Unfault Bibliothek und wurde in anderen Kernel-Bibliotheken wiederverwendet, z LigerKernel (auf dem die Kernel in diesem Artikel lose basieren) bietet es Heuristiken zum Optimieren beider Variablen:

def calculate_settings(n: int) -> tuple(int, int):
 MAX_FUSED_SIZE = 65536 # most grid dimension on Nvidia GPUs
    BLOCK_SIZE = next_power_of_2(n)
    if BLOCK_SIZE > MAX_FUSED_SIZE:
        # we take away this assertion on this article
        increase RuntimeError(
            f"Can not launch Triton kernel since n = {n} exceeds "
            f"the utmost CUDA blocksize = {MAX_FUSED_SIZE}."
        )
    num_warps = 4
    if BLOCK_SIZE >= 32768:
        num_warps = 32
    elif BLOCK_SIZE >= 8192:
        num_warps = 16
    elif BLOCK_SIZE >= 2048:
        num_warps = 8
    return BLOCK_SIZE, num_warps

Dann implementieren wir den regulären Softmax für den Vorwärtsdurchlauf und die Gleichung (10) für den Rückwärtspass. Die einzige Neuerung im Vergleich zu früheren Artikeln ist die Verwendung von Cache-Modifikatoren, die dem Compiler mitteilen, wie Daten zwischengespeichert und entfernt werden sollen. Im Second konzentrieren wir uns nur auf drei Cache-Modifikatoren:

• .ca (Cache auf allen Ebenen): Weist den Compiler an, die Daten sowohl in den L1- als auch in den L2-Cache zu laden, was darauf hindeutet, dass sie bald wiederverwendet werden könnten. Dieser Modifikator sollte verwendet werden, wenn die Daten klein genug sind, um in L1 zu passen (~128–192 KB professional SM auf einem A100) und wahrscheinlich wiederholt darauf zugegriffen wird.
• .cs (Streaming): Daten behandeln als Streamingwird es einmal verwendet und dann verworfen, um Platz in L1 freizugeben.
• .wb (Zurückschreiben): Normales zwischengespeichertes Schreiben, die Daten bleiben in der Cache-Hierarchie, intestine, wenn die Ausgabe wiederverwendet werden kann.

In den folgenden Kerneln verwenden wir die .ca Modifikator für Ladungen, da wir mehrere Operationen an den geladenen Daten durchführen. Zum Speichern verwenden wir .cs im Vorwärtsdurchlauf, da die Ausgaben nicht sofort wiederverwendet werden und .wb im Rückwärtsgang da im Kontext von autograd (dh die Kettenregel), Gradientenausgaben werden von nachgeschalteten Kerneln verbraucht.

Multiblock-Softmax

Werfen wir nun einen Blick auf die On-line-Formulierung des Softmax. In diesem Abschnitt implementieren wir eine Multiblock-Variante des vorherigen Kernels. Diese Model wird verwendet BLOCK_SIZE < n_colsmit anderen Worten, wir laden nur eine Kachel mit BLOCK_SIZE Elemente nacheinander, ähnlich wie wir mit gekacheltem GEMM in der umgegangen sind letztes Tutorial. Nun fragen Sie sich vielleicht: „Wie wählen wir die Blockgröße aus?“

Dies ist eine großartige Gelegenheit, Triton vorzustellen autotune Dienstprogramm. Ausgestattet mit einer Liste der Konfigurationen, autotune führt eine Rastersuche durch, um die beste Konfiguration für eine bestimmte Eingabeform zu ermitteln und zwischenzuspeichern. Dieser Vorgang wird jedes Mal wiederholt, wenn eine neue Eingabeform an den Kernel übergeben wird.

Hier führen wir mit der folgenden Hilfsfunktion eine Rastersuche über die Blockgröße und die Anzahl der Warps durch:

from itertools import product

# --- Multi Block Tuning ---
BLOCK_SIZES = (256, 512, 1024, 2048, 4096, 8192)
NUM_WARPS = (2, 4, 8, 16)

def get_autotune_config(
    block_sizes: listing(int), num_warps: listing(int)
) -> listing(triton.Config):
    return (
        triton.Config(kwargs={"BLOCK_SIZE": bs}, num_warps=nw)
        for (bs, nw) in listing(product(block_sizes, num_warps))
    )

Wir können jetzt unsere Multi-Block-Kernel damit dekorieren autotune und übergeben Sie die Liste der Konfigurationen, key=”n_cols” gibt an, dass die optimale Konfiguration von der Anzahl der Spalten der Eingabe abhängt.

Die Implementierung dieser Kernel ähnelt konzeptionell sehr dem zuvor behandelten On-line-Softmax. Der Hauptunterschied besteht darin, dass wir über Kacheln iterieren (nicht über einzelne Elemente wie in Numpy), was einige Anpassungen erfordert. Zum Beispiel addieren wir eine Summe über die Kachel im d Replace und der Rückwärtskernel erfordern nun ebenfalls zwei Iterationen.

Hinweis: Der PyTorch-Wrapper ist genau derselbe, außer dass wir die Zeile wo löschen BLOCK_SIZE Und num_warps werden deklariert (da sie ausgewählt werden von autotune).

Testen und Benchmarking

Wir können nun einen Vorwärts- und Rückwärtsdurchlauf mit beiden Kerneln ausführen und sicherstellen, dass sie mit den PyTorch-Baselines übereinstimmen:

def validate_kernel(kernel_fn: callable) -> None:
    system = "cuda:0" if torch.cuda.is_available() else "cpu"
    torch.random.manual_seed(0)

    # Generate inputs
    x = torch.randn((256, 512), system=system) # triton enter
    x.requires_grad = True
    xt = deepcopy(x) # torch enter

    triton_output = kernel_fn(x)
    torch_output = torch.softmax(xt, dim=1)
    torch.testing.assert_close(triton_output, torch_output) # take a look at fwd kernel

    # Setup pretend labels
    y = torch.zeros_like(x)
    inds = (torch.arange(0, y.form(0)), torch.randint(0, 3, (y.form(0),)))
    y(inds) = 1

    # Outline loss and run backward go
    loss_fn = torch.nn.CrossEntropyLoss()
    loss = loss_fn(torch_output, y)
    loss.backward()

    # Save gradient tensor for later
    torch_xgrad = xt.grad.detach().clone()
    triton_loss = loss_fn(triton_output, y)
    triton_loss.backward()
    torch.testing.assert_close(x.grad, torch_xgrad) # take a look at grad outputs

validate_kernel(softmax_sb)
validate_kernel(softmax_mb)

Abschließend vergleichen wir unsere Implementierung mit der PyTorch-Basislinie mithilfe des folgenden Snippets:

# --- Supply: Triton softmax tutorial ---
@triton.testing.perf_report(
    triton.testing.Benchmark(
        x_names=("N"),  # argument names to make use of as an x-axis for the plot
        x_vals=(
            128 * i for i in vary(2, 100)
        ),  # totally different doable values for `x_name`
        line_arg="supplier",  # argument title whose worth corresponds to a unique line within the plot
        line_vals=(
            "triton_single_block",
            "triton_multi_block",
            "torch",
        ),  # doable values for `line_arg``
        line_names=(
            "Triton_single_block",
            "Triton_multi_block",
            "Torch",
        ),  # label title for the traces
        kinds=(("blue", "-"), ("inexperienced", "-"), ("crimson", "-")),
        ylabel="GB/s",  # label title for the y-axis
        plot_name="softmax-performance",  # title for the plot. Used additionally as a file title for saving the plot.
        args={"M": 4096},  # values for perform arguments not in `x_names` and `y_name`
    )
)
def benchmark(M, N, supplier):
    x = torch.randn(M, N, system=DEVICE, dtype=torch.float32)
    stream = getattr(torch, DEVICE.sort).Stream()
    getattr(torch, DEVICE.sort).set_stream(stream)
    if supplier == "torch":
        ms = triton.testing.do_bench(lambda: torch.softmax(x, axis=-1))
    if supplier == "triton_single_block":
        torch.cuda.synchronize()
        ms = triton.testing.do_bench(lambda: softmax_sb(x))
        torch.cuda.synchronize()
    if supplier == "triton_multi_block":
        torch.cuda.synchronize()
        ms = triton.testing.do_bench(lambda: softmax_mb(x))
        torch.cuda.synchronize()
    gbps = lambda ms: 2 * x.numel() * x.element_size() * 1e-9 / (ms * 1e-3)
    return gbps(ms)

benchmark.run(show_plots=True, print_data=True)

Gute Nachrichten! Unser Single-Block-Kernel übertrifft durchweg die PyTorch-Basislinie, während die Multi-Block-Variante bei Eingaben mit mehr als 6.000 Spalten abfällt:

Betrachtet man größere Eingaben, können wir mehrere Beobachtungen machen:

1. Der Multiblock-Kernel stabilisiert schließlich einen Durchsatz von etwa 900 GB/s und übertrifft damit die PyTorch-Basislinie für Eingaben mit mehr als 30.000 Spalten.
2. Interessanterweise scheint die Multiblock-Variante bei Eingaben mit mehr als 60.000 Spalten dominant zu sein.
3. Auch wenn wir mit der Single-Block-Variante die maximale Blockgröße überschreiten, läuft der Kernel aus irgendeinem Grund immer noch reibungslos. Tatsächlich verwaltet Triton die Blockgröße automatisch unter der Haube.
   Wann n_cols größer als das Hardwarelimit ist, zerlegt Triton die Eingabe und iteriert darüber. Dies scheint jedoch langsamer zu sein als der Multiblock-Ansatz.

Um noch weiter zu gehen, könnten wir beide Ansätze in einem einzigen Kernel kombinieren, der explizit den optimalen Kernel basierend auf der Eingabegröße auswählt. Auf diese Weise würden wir von der hohen Leistung des Single-Block-Kernels für kleine Eingaben und dem höheren Durchsatz der Multi-Block-Variante für Eingaben mit mehr als 60.000 Spalten profitieren.

Damit ist die dritte Folge dieser Triton-Serie abgeschlossen. Nochmals vielen Dank für Ihre Unterstützung!

Im nächsten Artikel werden wir die On-line-Softmax-Formulierung im Kontext von nutzen Blitzaufmerksamkeit.

Bis zum nächsten Mal!

Ressourcen:

Eine kleine Optimierung reicht aus. Modelle mögen Gpt4 kosten mehr als 100 Millionen US -Greenback zum Coachingwas macht a 1% Effizienzgewinn Wert über eine Million Greenback. Eine leistungsstarke Möglichkeit, die Effizienz von Modellen für maschinelles Lernen zu optimieren, besteht darin, einige ihrer Komponenten zu schreiben direkt auf der GPU. Wenn Sie so etwas wie ich sind, reicht die einfache Erwähnung von Cuda -Kerneln aus, um Schüttelfrost in die Wirbelsäule zu schicken, da sie notorisch komplex sind, um zu schreiben und zu debuggen.

Glücklicherweise, Openai freigegeben Triton Im Jahr 2021, eine neue Sprache und Compiler, die einen Großteil der Komplexität von CUDA abstrahiert und weniger erfahrene Praktizierende ermöglicht, Performantenkörner zu schreiben. Ein bemerkenswertes Beispiel ist Nicht slothein LLM-Trainingsdienst, der verspricht 30x schnelleres Coaching mit 60% weniger Speicherverbrauchalles dank Ersetzen von Schichten in Pytorch durch Triton -Kernel.

In dieser Tutorial-Serie lernen wir die Grundlagen der GPU-Architektur und der Implementierung von Hochleistungs-Triton-Kerneln! Alle in dieser Serie vorgestellten Code sind bei verfügbar https://github.com/rpegoud/triton-kernels.

Grundlagen der GPU -Architektur

In diesem Abschnitt werden wir die Grundlagen von ((Nvidia) GPUs, um uns anzufangen und bis zum Ende dieses Artikels unseren ersten Triton -Kernel zu schreiben.

Ausgehend von der kleinsten Softwareinheit können wir die Hierarchie der Ausführungseinheiten wie folgt beschreiben:

• Themen: Der kleinste ArbeitseinheitSie führen den benutzerdefinierten Kernelcode aus.
• Verzerrt: Der kleinste PlanungseinheitSie bestehen immer aus 32 parallelen Threads mit jeweils ihren eigenen Anweisungsaditeln und Registrierungsstatus. Themen in einer Kette anfangen zusammen aber sind frei zu Zweig Und unabhängig ausführen.
• Fadenblöcke: Gruppe der Warps, wo alle Threads können kooperieren über gemeinsame Speicher und Synchronisierungsbarrieren. Es ist erforderlich, dass Threadblöcke ausgeführt werden können unabhängig und in jeder Reihenfolge parallel oder nacheinander. Diese Unabhängigkeit ermöglicht Threadblöcke in beliebiger Reihenfolge über eine beliebige Anzahl von Kernen geplantdamit GPU -Programme effizient mit der Anzahl der Kerne skalieren. Bei Bedarf können wir die Threads in einem Block an bestimmten Stellen im Kernel synchronisieren, um den Speicherzugriff zu synchronisieren.
• Streaming -Multiprozessor (SM): Eine Einheit, die verantwortlich ist viele Warps parallel ausführenEs besitzt den gemeinsamen Speicher und einen L1-Cache (enthält die neuesten globalen Reminiscence-Linien, auf die der SM zugegriffen hat). Ein SM hat eine engagierte Warp Scheduler Das Ziehen verzerrt aus den Fadenblöcken, die bereit sind, zu laufen.

Auf der Hardwareseite ist die kleinste Arbeitseinheit a Cuda -Kernder physische Arithmetische Logik UniT (Alu), der funktioniert arithmetische Operationen für einen Thread (oder Teile davon).

Um diesen Abschnitt mit einer Analogie zusammenzufassen, konnten wir sehen Cuda -Kerne als einzelne Arbeiterwährend a Kette ist a Truppe von 32 Arbeitern gleichzeitig die gleiche Anweisung. Sie können diese Aufgabe auf die gleiche Weise ausführen (verzweigen) und können sie möglicherweise zu einem anderen Zeitpunkt (Unabhängigkeit) ausführen. A Fadenblock besteht aus Mehrere Trupps teilen sich einen gemeinsamen Arbeitsbereich (dh das gemeinsame Gedächtnis) Arbeiter aus allen Trupps im Arbeitsbereich können aufeinander warten, um gleichzeitig zu Mittag zu essen. A Streaming -Multiprozessor ist a Fabrikboden mit vielen Trupps, die zusammenarbeiten und Instruments und Speicher teilen. Schließlich die GPU ist a ganze Pflanzemit vielen Böden.

Optimierungsgrundlagen

Bei der Optimierung von Deep Studying -Modellen jonglieren wir mit drei Hauptkomponenten:

1. Berechnen: Zeit, die von den GPU Computing Floating Level Operations (FLOPS) aufgewendet wird.
2. Erinnerung: Zeit für die Übertragung von Tensoren innerhalb einer GPU.
3. Überkopf: Alle anderen Operationen (Python -Dolmetscher, Pytorch Dispatch,…).

Wenn Sie diese Komponenten im Auge behalten, können Sie den richtigen Weg herausfinden, um einen Engpass zu lösen. Zum Beispiel hilft das Erhöhen des Berechnung (z. B. mit einer leistungsstärkeren GPU) nicht, wenn die meisten Zeit Speichertransfers ausgegeben werden. Im Idealfall sollte jedoch die meiste Zeit für Berechnung ausgegeben werden, genauer gesagt für Matrix -Multiplikationen, für die genaue Operation GPUs optimiert werden.

Dies impliziert, dass die Kosten für die Umschaltung von Daten entweder von der CPU in die GPU (”minimiert werden, um Daten umzusetzen (”Datenübertragungskosten”) Von einem Knoten zum anderen (”Netzwerkkosten”) Oder aus CUDA International Reminiscence (Drambillig aber langsam), um den gemeinsam genutzten Speicher (CUDA) (Sramteurer, aber schnellster Speicher des Geräts). Die später heißt es Bandbreite Kosten und wird vorerst unser Hauptaugenmerk sein. Zu den gängigen Strategien zur Reduzierung der Bandbreitenkosten gehören:

1. Wiederverwendung Daten in gemeinsamen Speicher geladen für mehrere Schritte. Ein Paradebeispiel hierfür ist die Multiplikation mit gekachelter Matrix, die wir in einem zukünftigen Beitrag behandeln werden.
2. Verschmelzen Mehrere Operationen in einem einzelnen Kernel (da jeder Kernel -Begin von DRAM zu SRAM bewegt wird), können wir beispielsweise eine Matrixmultiplikation mit einer Aktivierungsfunktion verschmelzen. Allgemein, Bedienerfusion Kann eine large Leistungserhöhung bieten, da es viele globale Speicherlesungen/Schreibvorgänge verhindert und zwei Betreiber eine Gelegenheit zur Fusion haben.

In diesem Beispiel führen wir eine Matrix -Multiplikation durch x@W und speichern das Ergebnis in einer Zwischenvariablen a. Wir wenden dann a an relu Zu a und speichern das Ergebnis in einer Variablen y. Dies erfordert, dass die GPU aus lesen kann x Und W Schreiben Sie das Ergebnis im globalen Gedächtnis in alesen aus a wieder und schließlich einschreiben y. Stattdessen würde uns Operator Fusion ermöglichen, die Anzahl der Lesevorgänge und Schreibvorgänge in International Reminiscence durch Durchführung der Matrix -Multiplikation und der Anwendung der Relu in einem einzelnen Kernel zu halbieren.

Triton

Wir werden jetzt unseren ersten Triton -Kernel schreiben, einen einfachen Vektor -Zusatz. Gehen wir zunächst durch, wie diese Operation auf einer GPU abgebaut und ausgeführt wird.

Erwägen Sie, die Einträge von zwei Vektoren zusammenzufassen wollen X Und Yjeweils mit 7 Elementen (n_elements=7).

Wir werden die GPU anweisen, dieses Downside in Stücken von 3 Elementen gleichzeitig anzugehen (BLOCK_SIZE=3). Um alle 7 Elemente der Eingangsvektoren abzudecken, startet die GPU 3 parallele „Programme“, unabhängige Instanz unseres Kernels, jeweils eine eindeutige Programm -ID, pid:

• Programm 0 wird Elemente zugewiesen 0, 1, 2.
• Programm 1 wird Elemente zugewiesen 3, 4, 5.
• Programm 2 wird Component zugewiesen 6.

Dann schreiben diese Programme die Ergebnisse in einem Vektor zurück Z in globalem Gedächtnis gespeichert.

Ein wichtiges Element ist, dass ein Kernel keinen ganzen Vektor erhält Xstattdessen erhält es a Zeiger auf die Speicheradresse des ersten ComponentsAnwesend X(0). Um auf die tatsächlichen Werte von zuzugreifen XWir müssen sie manuell aus dem globalen Speicher laden.

Wir können auf die Daten für jeden Block zugreifen, indem wir die Programm -ID verwenden: block_start = pid * BLOCK_SIZE. Von dort aus können wir die verbleibenden Elementadressen für diesen Block durch Pc erhalten offsets = block_start + vary(0, BLOCK_SIZE) und laden Sie sie in den Speicher.

Denken Sie jedoch daran, dass Programm 2 nur Component 6 zugewiesen ist, aber seine Offsets sind es (6, 7, 8). Um einen Indexierungsfehler zu vermeiden, lässt uns Triton a definieren Maske Hier zu identifizieren, um gültige Zielelemente zu identifizieren masks = offsets < n_elements.

Wir können jetzt sicher laden X Und Y und fügen Sie sie zusammen, bevor Sie das Ergebnis zu einer Ausgangsvariablen zurückschreiben Z in globalem Gedächtnis in ähnlicher Weise.

Schauen wir uns den Code genauer an, hier ist der Triton -Kernel:

import triton
import triton.language as tl

@triton.jit
def add_kernel(
	x_ptr, # pointer to the primary reminiscence entry of x
	y_ptr, # pointer to the primary reminiscence entry of y
	output_ptr, # pointer to the primary reminiscence entry of the output
	n_elements, # dimension of x and y
	BLOCK_SIZE: tl.constexpr, # dimension of a single block
):
	# --- Compute offsets and masks ---
	pid = tl.program_id(axis=0) # block index
	block_start = pid * BLOCK_SIZE # begin index for present block
	offsets = block_start + tl.arange(0, BLOCK_SIZE) # index vary
	masks = offsets < n_elements # masks out-of-bound parts
	
	# --- Load variables from world reminiscence ---
	x = tl.load(x_ptr + offsets, masks=masks)
	y = tl.load(y_ptr + offsets, masks=masks)

	# --- Operation ---
	output = x + y	
	
	# --- Save outcomes to world reminiscence ---
	tl.retailer(pointer=output_ptr + offsets, worth=output, masks=masks)

Lassen Sie uns einige der tritonspezifischen Syntax aufschlüsseln:

• Erstens ist ein Triton -Kernel immer von dekoriert von <a href="http://twitter.com/triton" goal="_blank" rel="noreferrer noopener">@triton</a>.jit.
• Zweitens müssen einige Argumente als statisch deklariert werden, was bedeutet, dass sie zur Rechenzeit bekannt sind. Dies ist erforderlich für BLOCK_SIZE und wird erreicht, indem das hinzufügen tl.constexpr Geben Sie Annotation ein. Beachten Sie auch, dass wir keine anderen Variablen kommentieren, da sie keine ordnungsgemäßen Python -Variablen sind.
• Wir verwenden tl.program_id Zugriff auf die ID des aktuellen Blocks, tl.arange verhält sich ähnlich wie Numpy’s np.arange.
• Das Laden und Speichern von Variablen wird durch Anrufe erreicht tl.load Und tl.retailer mit Anordnungen von Zeigern. Beachten Sie, dass es keine gibt return Aussage, diese Rolle wird delegiert tl.retailer.

Um unseren Kernel zu benutzen, müssen wir jetzt a schreiben Wrapper auf Pytorch-Ebene Das liefert Speicherzeiger und definiert a Kernelgitter. Im Allgemeinen ist das Kernelnetz ein 1D-, 2D- oder 3D -Tupel, das das enthält Anzahl der am Kernel zugewiesenen Fadenblöcken entlang jeder Achse. In unserem vorherigen Beispiel haben wir ein 1D -Gitter mit 3 Fadenblöcken verwendet: grid = (3, ).

Um unterschiedliche Arraygrößen zu behandeln, sind wir standardmäßig grid = (ceil(n_elements / BLOCK_SIZE), ).

def add(X: torch.Tensor, Y: torch.Tensor) -> torch.Tensor:
	"""PyTorch wrapper for `add_kernel`."""
	output = torch.zeros_like(x) # allocate reminiscence for the output
	n_elements = output.numel()  # dimension of X and Y
	
	# cdiv = ceil div, computes the variety of blocks to make use of
	grid = lambda meta: (triton.cdiv(n_elements, meta("BLOCK_SIZE")),)
	# calling the kernel will robotically retailer `BLOCK_SIZE` in `meta`
	# and replace `output`
	add_kernel(grid)(X, Y, output, n_elements, BLOCK_SIZE=1024)
	
	return output

Hier sind zwei letzte Notizen über den Wrapper:

Sie haben das vielleicht bemerkt grid wird als Lambda -Funktion definiert. Dadurch kann Triton die Anzahl der zu starteten Threadblöcke berechnen zur Startzeit. Daher berechnen wir die Gittergröße basierend auf der in gespeicherten Blockgröße metaein Wörterbuch von Kompilierungs-Zeitkonstanten, die dem Kernel ausgesetzt sind.

Wenn Sie den Kernel nennen, den Wert von output wird an Ort und Stelle geändert, sodass wir uns nicht neu zuweisen müssen output = add_kernel(…).
Wir können dieses Tutorial abschließen, indem wir überprüfen, ob unser Kernel ordnungsgemäß funktioniert:

x, y = torch.randn((2, 2048), machine="cuda")

print(add(x, y))
>> tensor(( 1.8022, 0.6780, 2.8261, ..., 1.5445, 0.2563, -0.1846), machine='cuda:0')

abs_difference = torch.abs((x + y) - add(x, y))
print(f"Max absolute distinction: {torch.max(abs_difference)}")
>> Max absolute distinction: 0.0

Das ist es für diese Einführung. In folgenden Beiträgen werden wir lernen, interessantere Kernel wie geflieste Matrix -Multiplikation zu implementieren und zu sehen autograd.

Bis zum nächsten Mal!

Referenzen und nützliche Ressourcen

Bild von Autor | Ideogramm

GPUs eignen sich hervorragend für Aufgaben, bei denen Sie den gleichen Betrieb über verschiedene Datenstücke hinweg ausführen müssen. Dies ist als die bekannt Einzelanweisung, mehrere Daten (SIMD) Ansatz. Im Gegensatz zu CPUs, die nur wenige leistungsstarke Kerne haben, haben GPUs Tausende kleinerer, die diese sich wiederholten Operationen auf einmal ausführen können. Sie sehen dieses Muster im maschinellen Lernen viel, beispielsweise beim Hinzufügen oder Multiplizieren großer Vektoren, da jede Berechnung unabhängig ist. Dies ist das ideale Szenario für die Verwendung von GPUs, um die Aufgaben parallelität zu beschleunigen.

Nvidia erstellt CUDA als eine Möglichkeit für Entwickler, Programme zu schreiben, die auf der GPU anstelle der CPU ausgeführt werden. Es basiert auf C und ermöglicht es Ihnen, spezielle Funktionen mit dem Namen Kernel zu schreiben, die viele Operationen gleichzeitig ausführen können. Das Downside ist, dass das Schreiben von CUDA in C oder C ++ nicht gerade anfängerfreundlich ist. Sie müssen sich mit Dingen wie Schaltspeicherzuweisung, Thread -Koordination und dem Verständnis der Funktionsweise der GPU auf einer niedrigen Ebene befassen. Dies kann überwältigend sein, insbesondere wenn Sie es gewohnt sind, Code in Python zu schreiben.

Hier Numba kann Ihnen helfen. Es ermöglicht das Schreiben von Cuda-Kerneln mit Python mit der Compiler-Infrastruktur von LLVM (Digital Machine), um Ihren Python-Code direkt zu CUDA-kompatiblen Kerneln zu kompilieren. Mit Simply-in-Time (JIT) -Kompilation können Sie Ihre Funktionen mit einem Dekorateur kommentieren und Numba alles andere für Sie behandelt.

In diesem Artikel werden wir ein gemeinsames Beispiel für die Addition des Vektors verwenden und einfachen CPU -Code in einen Cuda -Kernel mit Numba konvertieren. Die Addition der Vektor ist ein ideales Beispiel für die Parallelität, da die Zugabe über einen einzelnen Index unabhängig von anderen Indizes ist. Dies ist das perfekte SIMD -Szenario, sodass alle Indizes gleichzeitig hinzugefügt werden können, um die Vektor -Addition in einem Vorgang zu vervollständigen.

Beachten Sie, dass Sie eine CUDA -GPU benötigen, um diesem Artikel zu folgen. Sie können verwenden Colabs Kostenlose T4 -GPU oder eine lokale GPU mit NVIDIA -Toolkit und NVCC installiert.

# Einrichten der Umgebung und Set up von Numba

Numba ist als Python -Paket erhältlich und Sie können es mit PIP installieren. Darüber hinaus werden wir verwenden Numpy Für Vektoroperationen. Richten Sie die Python -Umgebung mit den folgenden Befehlen ein:

python3 -m venv venv
supply venv/bin/activate
pip set up numba-cuda numpy

# Vektoraddition auf der CPU

Nehmen wir ein einfaches Beispiel für die Addition von Vektor. Für zwei gegebene Vektoren fügen wir die entsprechenden Werte aus jedem Index hinzu, um den Endwert zu erhalten. Wir werden Numpy verwenden, um zufällige Float32 -Vektoren zu erzeugen und die endgültige Ausgabe mit A for Loop zu generieren.

import numpy as np 

N = 10_000_000 # 10 million components 
a = np.random.rand(N).astype(np.float32) 
b = np.random.rand(N).astype(np.float32) 
c = np.zeros_like(a) # Output array 

def vector_add_cpu(a, b, c): 
    """Add two vectors on CPU""" 
    for i in vary(len(a)): 
        c(i) = a(i) + b(i)

Hier ist eine Aufschlüsselung des Codes:

• Initialisieren Sie jeweils zwei Vektoren mit 10 Millionen zufälligen Gleitkomma-Zahlen
• Wir erstellen auch einen leeren Vektor c das Ergebnis speichern
• Der vector_add_cpu Funktion schaltet einfach jeden Index durch und fügt die Elemente aus a Und bdas Ergebnis speichern c

Dies ist a Serienbetrieb; Jede Zugabe erfolgt nacheinander. Dies funktioniert zwar einwandfrei, ist jedoch nicht der effizienteste Ansatz, insbesondere für große Datensätze. Da jede Ergänzung unabhängig von den anderen ist, ist dies ein perfekter Kandidat für die parallele Ausführung bei einer GPU.

Im nächsten Abschnitt sehen Sie, wie Sie dieselbe Operation für die GPU mit Numba konvertieren. Durch die Verteilung der einzelnen Elemente über Tausende von GPU-Threads können wir die Aufgabe erheblich schneller erledigen.

# Vektor Addition auf der GPU mit Numba

Sie werden nun Numba verwenden, um eine Python -Funktion zu definieren, die auf CUDA ausgeführt werden kann, und sie innerhalb von Python ausführen kann. Wir führen den gleichen Vektor -Addition -Betrieb durch, aber jetzt kann es für jeden Index des Numpy -Arrays parallel ausgeführt werden, was zu einer schnelleren Ausführung führt.

Hier ist der Code zum Schreiben des Kernels:

from numba import config

# Required for newer CUDA variations to allow linking instruments. 
# Prevents CUDA toolkit and NVCC model mismatches.
config.CUDA_ENABLE_PYNVJITLINK = 1

from numba import cuda, float32

@cuda.jit
def vector_add_gpu(a, b, c):
	"""Add two vectors utilizing CUDA kernel"""
	# Thread ID within the present block
	tx = cuda.threadIdx.x
	# Block ID within the grid
	bx = cuda.blockIdx.x
	# Block width (variety of threads per block)
	bw = cuda.blockDim.x

	# Calculate the distinctive thread place
	place = tx + bx * bw

	# Be certain that we do not exit of bounds
	if place < len(a):
    	    c(place) = a(place) + b(place)

def gpu_add(a, b, c):
	# Outline the grid and block dimensions
	threads_per_block = 256
	blocks_per_grid = (N + threads_per_block - 1) // threads_per_block

	# Copy knowledge to the machine
	d_a = cuda.to_device(a)
	d_b = cuda.to_device(b)
	d_c = cuda.to_device(c)

	# Launch the kernel
	vector_add_gpu(blocks_per_grid, threads_per_block)(d_a, d_b, d_c)

	# Copy the consequence again to the host
	d_c.copy_to_host(c)

def time_gpu():
	c_gpu = np.zeros_like(a)
	gpu_add(a, b, c_gpu)
	return c_gpu

Lassen Sie uns aufschlüsseln, was oben passiert.

// Verständnis der GPU -Funktion

Der @cuda.jit Der Dekorateur fordert Numba auf, die folgende Funktion als Cuda -Kernel zu behandeln. Eine spezielle Funktion, die parallel über viele Threads der GPU ausgeführt wird. Zur Laufzeit kompiliert Numba diese Funktion mit CUDA-kompatiblen Code und verarbeitet die C-API-Transpilation für Sie.

@cuda.jit
def vector_add_gpu(a, b, c):
	...

Diese Funktion wird gleichzeitig auf Tausenden von Threads ausgeführt. Wir brauchen jedoch eine Möglichkeit, herauszufinden, an welchem Teil der Daten jeder Thread funktionieren sollte. Das tun die nächsten Zeilen:

• tx ist die ID des Threads in seinem Block
• bx ist die ID des Blocks im Netz
• bw ist wie viele Fäden in einem Block sind

Wir kombinieren diese mit Berechnen Sie eine einzigartige Placewas jedem Thread mitteilt, welches Component der Arrays es hinzufügen sollte. Beachten Sie, dass die Threads und Blöcke möglicherweise nicht immer einen gültigen Index liefern, da sie in Befugnissen von 2. wirken. Dies kann zu ungültigen Indizes führen, wenn die Vektorlänge nicht der zugrunde liegenden Architektur entspricht. Daher fügen wir eine Schutzbedingung hinzu, um den Index zu validieren, bevor wir die Vektor -Addition durchführen. Dies verhindert, dass beim Zugriff auf das Array einen außergewöhnlichen Laufzeitfehler ist.

Sobald wir die einzigartige Place kennen, können wir die Werte jetzt genau wie für die CPU -Implementierung hinzufügen. Die folgende Zeile entspricht der CPU -Implementierung:

c(place) = a(place) + b(place)

// Begin des Kernels

Der gpu_add Funktion stellt Dinge auf:

• Es definiert, wie viele Themen und Blöcke verwendet werden sollen. Sie können mit unterschiedlichen Werten der Block- und Fadengrößen experimentieren und die entsprechenden Werte im GPU -Kernel drucken. Dies kann Ihnen helfen, zu verstehen, wie die zugrunde liegende GPU -Indexierung funktioniert.
• Es kopiert die Eingangsarrays (aAnwesend bUnd c) vom CPU -Speicher zum GPU -Speicher, sodass die Vektoren im GPU -RAM zugänglich sind.
• Es führt den GPU -Kernel mit vector_add_gpu(blocks_per_grid, threads_per_block).
• Schließlich kopiert es das Ergebnis von der GPU zurück in die c Array, damit wir auf die Werte der CPU zugreifen können.

# Vergleich der Implementierungen und der potenziellen Beschleunigung

Jetzt, da wir sowohl die CPU- als auch die GPU -Versionen von Vektorabzug haben, ist es Zeit zu sehen, wie sie sich vergleichen. Es ist wichtig, die Ergebnisse und den Ausführungsschub zu überprüfen, den wir mit CUDA -Parallelität erhalten können.

import timeit

c_cpu = time_cpu()
c_gpu = time_gpu()

print("Outcomes match:", np.allclose(c_cpu, c_gpu))

cpu_time = timeit.timeit("time_cpu()", globals=globals(), quantity=3) / 3
print(f"CPU implementation: {cpu_time:.6f} seconds")

gpu_time = timeit.timeit("time_gpu()", globals=globals(), quantity=3) / 3
print(f"GPU implementation: {gpu_time:.6f} seconds")

speedup = cpu_time / gpu_time
print(f"GPU speedup: {speedup:.2f}x")

Zunächst führen wir beide Implementierungen aus und überprüfen, ob ihre Ergebnisse übereinstimmen. Dies ist wichtig, um sicherzustellen, dass unser GPU -Code korrekt funktioniert und die Ausgabe mit den CPUs übereinstimmt.

Als nächstes verwenden wir Pythons integriert timeit Modul, um zu messen, wie lange jede Model dauert. Wir führen jede Funktion ein paar Mal aus und nehmen den Durchschnitt, um ein zuverlässiges Timing zu erhalten. Schließlich berechnen wir, wie oft schneller die GPU -Model mit der CPU verglichen wird. Sie sollten einen großen Unterschied sehen, da die GPU viele Operationen gleichzeitig ausführen kann, während die CPU sie jeweils in einer Schleife umgeht.

Hier ist die erwartete Ausgabe der T4 -GPU von NVIDIA auf Colab. Beachten Sie, dass sich die genaue Beschleunigung basierend auf CUDA -Versionen und der zugrunde liegenden {Hardware} unterscheiden kann.

Outcomes match: True
CPU implementation: 4.033822 seconds
GPU implementation: 0.047736 seconds
GPU speedup: 84.50x

Dieser einfache Take a look at hilft, die Leistung der GPU -Beschleunigung zu demonstrieren und warum er für Aufgaben, die große Datenmengen und parallele Arbeiten beinhalten, so nützlich ist.

# Einpacken

Und das ist es. Sie haben jetzt Ihren ersten Cuda -Kernel mit Numba geschrieben, ohne tatsächlich einen C- oder CUDA -Code zu schreiben. Numba ermöglicht eine einfache Schnittstelle zur Verwendung der GPU durch Python und macht es für Python -Ingenieure viel einfacher, mit CUDA -Programmierung zu beginnen.

Sie können jetzt dieselbe Vorlage verwenden, um erweiterte CUDA -Algorithmen zu schreiben, die im maschinellen Lernen und im tiefen Lernen häufig vorkommen. Wenn Sie nach dem SIMD -Paradigma ein Downside finden, ist es immer eine gute Idee, GPU zur Verbesserung der Ausführung zu verwenden.

Der komplette Code ist in Colab Pocket book verfügbar, auf das Sie zugreifen können Hier. Fühlen Sie sich frei, es zu testen und einfache Änderungen vorzunehmen, um ein besseres Verständnis dafür zu erhalten, wie CUDA -Indexierung und Ausführung intern funktioniert.

Kanwal Mehreen ist ein Ingenieur für maschinelles Lernen und technischer Schriftsteller mit einer tiefgreifenden Leidenschaft für die Datenwissenschaft und die Schnittstelle von KI mit Medizin. Sie hat das eBook „Produktivität mit Chatgpt maximieren“. Als Google -Era -Gelehrte 2022 für APAC setzt sie sich für Vielfalt und akademische Exzellenz ein. Sie wird auch als Teradata -Vielfalt in Tech Scholar, MITACS Globalink Analysis Scholar und Harvard Wecode Scholar anerkannt. Kanwal ist ein leidenschaftlicher Verfechter der Veränderung, nachdem er Femcodes gegründet hat, um Frauen in STEM -Bereichen zu stärken.