SparseCore für Large Embedding Models (LEM) im Detail

SparseCore ist ein spezieller gekachelter Prozessor, der für die leistungsstarke Beschleunigung von Arbeitslasten entwickelt wurde, die unregelmäßigen, spärlichen Speicherzugriff und ‑berechnungen umfassen, insbesondere bei großen Datasets, die in High Bandwidth Memory (HBM) gespeichert sind. Sie eignet sich hervorragend für Aufgaben wie das Einbetten von Lookups, kann aber auch eine Vielzahl anderer dynamischer und spärlicher Arbeitslasten beschleunigen.

1. Einführung in SparseCore

Wichtige architektonische Merkmale:

• Gekachelte Architektur: Besteht aus mehreren Rechenkacheln (jede Kachel ist eine vollständige Dataflow-Einheit mit eigenem lokalen Speicher und eigener Verarbeitungseinheit), die eine parallele Verarbeitung ermöglichen.
• Dynamische Ausführung: Unterstützt nativ datenabhängige Kontrollflüsse und Speicherzugriffe, was für spärliche Daten entscheidend ist.
• Vektorverarbeitung: Nutzt Aufgaben mit kleinen Vektoren (8 oder 16 Elemente, je nach Hardwareversion) für effiziente Berechnungen.
• Zentrale Steuerung: Ein einzelner SparseCore-Sequenzer orchestriert Aufgaben auf allen Kacheln und sorgt so für synchronisierte Vorgänge.
• Unterstützung für die Datenzusammenfassung: Enthält spezielle lane-übergreifende Vorgänge, die für Aufgaben wie Sortieren, Filtern und Präfixsummen nützlich sind.
• Speicherhierarchie: HBM wird strategisch zum Speichern großer Datensätze und lokaler Scratchpad-Speicher (SPMEM) zum Staging von Daten genutzt, auf die häufig zugegriffen wird. Dadurch wird die HBM-Latenz erheblich reduziert.

Spezifikationen auf einen Blick:

2. Vorverarbeitung von SparseCore-Hosts

Eine effektive Datenvorbereitung ist für die Leistung von SparseCore von entscheidender Bedeutung. Hier spielt die Host-Vorverarbeitung eine wichtige Rolle. Sie umfasst mehrere wichtige Funktionen:

• Datentransformation:
  • Wenden Sie die erforderlichen Transformationen auf die Rohdaten an.
  • ID-Transformationen verwalten, was besonders wichtig ist, wenn Sie Features oder Tabellen stapeln.
  • Konvertieren Sie die Eingabedaten in das COO-Format (Coordinate), das im folgenden Abschnitt beschrieben wird.
  • Partitionieren Sie die Daten, um sie effizient auf die verschiedenen SparseCores auf dem Chip zu verteilen.
• Limit-Validierung:
  • Achten Sie darauf, dass die Merkmale der Eingabedaten (z. B. die Anzahl der IDs) den vordefinierten Betriebsgrenzen des SparseCore entsprechen, z. B. max_ids_per_partition und max_unique_ids_per_partition.
  • Wenn die Eingabedaten diese Grenzwerte überschreiten, kann die Vorverarbeitungsebene des Hosts versuchen, die Daten in kleinere Minibatches zu segmentieren, die den Einschränkungen entsprechen.
• Datenübertragung:
  • Die verarbeiteten und validierten Daten werden effizient in den High Bandwidth Memory (HBM) der TPU kopiert, sodass sie für die SparseCore-Ausführung bereit sind.

Tabellen stapeln:

Das Stapeln von Tabellen ist eine wichtige Optimierungstechnik, bei der mehrere Einbettungstabellen logisch kombiniert werden, um die Effizienz der Einbettungssuche zu verbessern. Dieser Prozess wird in der Regel automatisch vom zugrunde liegenden ML-Framework ausgeführt.

• Stapelung von Features: Dies tritt auf, wenn mehrere unterschiedliche Features dieselbe zugrunde liegende Einbettungstabelle verwenden. Ein häufiges Beispiel ist die Verwendung eines einzelnen Einbettungswörterbuchs für verschiedene kategoriale Merkmale wie Postleitzahlen aus unterschiedlichen Kontexten.
• Stapeln von Tabellen: In diesem Szenario werden mehrere unterschiedliche Einbettungstabellen gestapelt. Tabellen mit derselben Einbettungsdimension und Optimiererkonfiguration werden häufig gruppiert.

Der Hauptvorteil des Stapelns von Tabellen ist die Erstellung einer größeren effektiven Batchgröße für die Vorgänge in diesen gestapelten Tabellen. Dadurch wird der Rechenaufwand reduziert und es können Latenzen bei der Inter-Chip-Kommunikation (ICI) verborgen werden. Für eine optimale Leistung wird eine moderate Anzahl von gestapelten Tabellen (in der Regel zwischen 5 und 100) empfohlen.

3. Konvertierung in COO-Tensoren

Bevor Daten von SparseCore verarbeitet werden können, werden sie in der Regel in ein COO-Format (Coordinate) für dünnbesetzte Tensoren konvertiert. Das COO-Format ist eine Möglichkeit, dünnbesetzte Matrizen effizient darzustellen, in der Regel mit drei Arrays:

• row_ids: Ein Array mit den Zeilenindizes für jedes Element ungleich null. Im Kontext der Batchverarbeitung entspricht dies häufig der Batchdimension.
• col_ids: Ein Array mit den Spaltenindizes für jedes Element ungleich null. Bei Einbettungen sind das häufig die Feature- oder ID-Werte.
• values (optional): Ein Array mit den tatsächlichen Werten der Elemente ungleich null an den entsprechenden (row, col)-Koordinaten. Bei Grenzwertberechnungen (siehe unten) in Bezug auf ID-Anzahlen werden diese Werte (Steigerungen) häufig nicht berücksichtigt.

Anschauliches Beispiel:

Angenommen, Sie haben eine Eingabe-Matrix, die Batches von IDs darstellt:

[
    [id_A],                 // Sample 0
    [id_A, id_B, id_C],     // Sample 1
    [id_B, id_B, id_D],     // Sample 2 (note duplicate id_B)
]

Nach der Konvertierung in das COO-Format (und möglicherweise nach dem Entfernen von Duplikaten von IDs innerhalb derselben Stichprobe):

row_ids = [0, 1, 1, 1, 2, 2]
col_ids = [id_A, id_A, id_B, id_C, id_B, id_D]

Diese Konvertierung ist von grundlegender Bedeutung für die Verarbeitung und Verteilung von Aufgaben durch SparseCore. Insbesondere die col_ids sind entscheidend dafür, zu welcher bestimmten SparseCore-Partition eine ID gehört. So können Daten effizient partitioniert und nachgeschlagen werden.

4. SparsecoreConfig: Die API auf hoher Ebene

Framework-spezifische Einbettungs-APIs:

• JAX: https://github.com/jax-ml/jax-tpu-embedding
• TensorFlow: https://www.tensorflow.org/recommenders/api_docs/python/tfrs/layers/embedding/TPUEmbedding
• Keras: https://keras.io/keras_rs/api/embedding_layers/distributed_embedding

SparsecoreConfig oder entsprechende Mechanismen wie XLA-Flags dienen als Schnittstelle auf hoher Ebene, um eine Vielzahl von SparseCore-Verhaltensweisen zu steuern. Ein gründliches Verständnis dieser Parameter ist für die effektive Leistungsoptimierung und den korrekten Betrieb Ihrer Modelle unerlässlich.

• disable_table_stacking: bool = False
  • Erklärung: Mit diesem Flag wird gesteuert, ob das automatische Stapeln von Tabellen verhindert wird. Dies kann zu einer geringeren Leistung führen, da der Overhead erhöht wird und die Möglichkeit, die Inter-Chip Interconnect (ICI)-Latenz zu verbergen, eingeschränkt wird.
  • Standard: False (d. h., das Stapeln von Tabellen ist standardmäßig aktiviert, sofern das Framework es unterstützt).
• max_ids_per_chip_per_sample: int = 64
  • Erklärung: Mit diesem Parameter wird eine globale Obergrenze für die Gesamtzahl der Einbettungs-IDs festgelegt, die ein einzelner Chip aus einer Stichprobe im Eingabe-Batch verarbeiten kann. Die Anzahl wird über alle Tabellen hinweg aggregiert. Es ist ein Mechanismus zur Verwaltung von Ressourcen auf Chipebene, bevor detailliertere Grenzwerte pro Tabelle oder Partition berücksichtigt werden. Die Feinabstimmung dieses Werts hängt in der Regel von den spezifischen Modelleigenschaften und der Gesamtsystemkapazität ab.
  • Standard: 64.
• max_ids_per_table: Optional[Dict[str, int]] = None
  • Erklärung: Mit diesem Parameter wird die maximale Anzahl von Einbettungs-IDs (die Duplikate enthalten können) angegeben, die für jede logische Tabelle verarbeitet werden können. Dabei werden alle Partitionen über alle SparseCores hinweg berücksichtigt. Dieses Limit ist höher als max_ids_per_partition. Wenn eine Tabelle T in P Partitionen unterteilt ist, gilt dieses Limit für die Summe der IDs, die an alle P Partitionen weitergeleitet werden. Das hängt oft mit max_ids_per_partition_per_sample und der Gesamtbatchgröße zusammen.
  • Einstellung: Wird in der Regel mit einer Datei mit Grenzwerten konfiguriert, z. B. mit dem Flag xla_sparse_core_max_ids_file, wobei max_ids_per_partition definiert ist. Dieses Konzept auf Tabellenebene ist eine Methode zum Festlegen dieser Grenzwerte auf Partitionsebene (max_ids und max_uniques).
  • Standard: None (der Wert kann aus den Limits pro Partition oder anderen Konfigurationen abgeleitet werden, wenn er nicht explizit angegeben wird).
• max_unique_ids_per_table: Optional[Dict[str, int]] = None
  • Erklärung: Analog zu max_ids_per_table, aber mit diesem Parameter wird die maximale Anzahl eindeutiger IDs für jede logische Tabelle angegeben. Dies ist eine wichtige Einstellung für die richtige Dimensionierung der On-Device-Puffer, die bei der Verarbeitung eindeutiger IDs und nachfolgenden Vektoroperationen verwendet werden.
  • Einstellung: Wird in der Regel auch in einer Datei mit Grenzwerten definiert oder aus max_unique_ids_per_partition_per_sample abgeleitet.
  • Standard: None.
• allow_id_dropping: bool = False
  • Erklärung: Mit diesem booleschen Flag wird das Löschen von IDs gesteuert, wenn die Anzahl der in den Eingabedaten gefundenen IDs (beobachtete Grenzwerte) die während der Kompilierung festgelegten Grenzwerte (z. B. max_ids_per_partition) überschreitet.
    • Wenn True: IDs, die zu einer Überschreitung der Limits führen würden, werden ohne Benachrichtigung verworfen. Normalerweise werden IDs innerhalb einer Partition in sortierter Reihenfolge verarbeitet. Jede ID, die die laufende Anzahl über das Limit für den zugehörigen Mini-Batch hinaus erhöhen würde, wird verworfen. So kann das Programm weiter ausgeführt werden, was sich jedoch negativ auf die Modellgenauigkeit auswirken kann.
    • Wenn False: Es wird ein Fehler ausgelöst und der Prozess wird wahrscheinlich beendet, wenn die beobachteten Grenzwerte die kompilierten Grenzwerte überschreiten. So werden alle Daten verarbeitet, aber die Grenzwerte müssen konservativer konfiguriert werden.
  • Standard: False (führt bei einem Überlauf zu einem Fehler und nicht zu einem stillen Datenverlust).

• initialize_tables_on_host: bool = True

  • Erklärung: Mit diesem Flag wird festgelegt, ob Embedding-Tabellen auf der Host-CPU initialisiert werden, bevor sie in den High Bandwidth Memory (HBM) der TPU übertragen werden. Tabellen werden in der Regel auf dem Host initialisiert. Wenn Sie dies auf True festlegen, wird diese Konvention eingehalten. Wenn es auf False festgelegt wäre, würde dies einen Initialisierungsmechanismus auf dem Gerät implizieren, der unterschiedliche Auswirkungen auf die Leistung oder spezifische Initialisierungsvoraussetzungen haben könnte.

• enable_fast_table_initialization: bool = False

  • Erklärung: Initialisiert die Tabellen direkt auf der TPU. Das kann dazu beitragen, die Startzeiten von Modellen zu verkürzen.

5. Pipelining zur Leistungssteigerung

Pipelining ist eine Technik zur Leistungsoptimierung, die die gleichzeitige Ausführung von Vorgängen auf dem TensorCore (TC) und dem SparseCore (SC) ermöglicht. Durch die Überlappung dieser Berechnungen kann der Gesamtdurchsatz deutlich verbessert werden.

• Mechanismus: In einem standardmäßigen Trainingsschritt, der Sparse-Embedding-Lookups (von SC verarbeitet) und Berechnungen für dichte Ebenen (von TC verarbeitet) umfasst, ermöglicht die Pipelining-Funktion, dass der SC seinen Teil von Schritt i (z. B. Forward- oder Backward-Pass) verarbeitet, während der TC gleichzeitig einen anderen Teil desselben Schritts i oder sogar Teile angrenzender Schritte wie i-1 oder i+1 verarbeitet.
• Auswirkungen auf Gradienten: Der SparseCore-Vorgang kann mit „alten“ Gradienten ausgeführt werden. So sind die Gradienten, die während der Backpropagation-Phase von Schritt i berechnet werden, möglicherweise erst in Schritt i+2 vollständig aktualisiert und für den SC sichtbar.
• Kompromiss zwischen Leistung und numerischen Werten: Diese überlappende Ausführung kann zu erheblichen Geschwindigkeitssteigerungen führen, möglicherweise bis zu einer Verdopplung der Geräte-Schrittzeit. Die geringfügigen Änderungen der numerischen Werte (embedding_weights), die sich aus der Verwendung veralteter Gradienten ergeben, können sich jedoch auf das Konvergenzverhalten des Modells oder die endgültige Genauigkeit auswirken. Die Akzeptabilität dieses Kompromisses hängt stark vom Modell ab und erfordert oft eine empirische Validierung.
• Steuerungsflag: Das Pipelining kann mit tf_xla_disable_full_embedding_pipelining gesteuert werden. Wenn Sie dieses Flag auf true setzen, wird das vollständige Pipelining (überlappende TensorCore- und SparseCore-Berechnung) deaktiviert. Wenn Sie es auf false setzen (oder wenn die Semantik des Flags die Aktivierung impliziert, wenn es auf „false“ gesetzt ist), wird es aktiviert.

Konzeptioneller Pipelining-Ablauf:

• Ohne Pipelining (vereinfachter sequenzieller Ablauf):

  Loop: SC/F_i -> TC/F_i -> TC/B_i -> SC/B_i

• Mit Pipelining (vereinfachter überlappender Ablauf):

  Time ->
  Step i:   SC/F_i | TC/F_i | TC/B_i | SC/B_i
  Step i+1:          SC/F_i+1| TC/F_i+1| TC/B_i+1| SC/B_i+1
  

  Hinweis: Die tatsächlichen Pipelining-Phasen, die in der Hardware und im Compiler implementiert sind, können komplexer sein und umfassen oft Vor-, Haupt- und Nachschleifen, um Datenabhängigkeiten zu verwalten und die Richtigkeit zu gewährleisten.

6. Die Rolle von XLA

XLA (Accelerated Linear Algebra) ist der domainspezifische Compiler, der rechenintensive Graphen auf hoher Ebene, in der Regel aus Frameworks wie TensorFlow, in hochoptimierten Maschinencode für TPUs übersetzt. Dazu gehört auch das Generieren der Anweisungen für Vorgänge, die für den SparseCore bestimmt sind.

Wichtige Funktionen im SparseCore-Kontext:

• Kompilierung von Sparse-Operationen: XLA ist für die Kompilierung von Embedding-Suchvorgängen (z. B. SparseDenseMatmulOp) und anderen Sparse-Berechnungen in ausführbare SparseCore-Programme auf niedriger Ebene verantwortlich.
• Integration von Limits: Es werden die konfigurierten Betriebslimits verwendet (z. B. max_ids_per_partition, max_unique_ids_per_partition, die häufig über eine Limitdatei bereitgestellt werden, die durch Flags wie xla_sparse_core_max_ids_file angegeben wird), um die Größen von Gerätespeicherpuffern statisch zu bestimmen und zuzuweisen, insbesondere im SPMEM.
• Gezielte Optimierungen: XLA führt eine Reihe von Optimierungen durch, die speziell für die SparseCore-Architektur entwickelt wurden. Dazu gehören die Planung von Anweisungen, Transformationen des Speicherlayouts und das Zusammenführen von Vorgängen, um die Effizienz zu maximieren.
• Steuerung über Flags: Viele Aspekte des SparseCore-Verhaltens, der Optimierungsparameter und der Optimierungsstrategien werden über XLA-Flags verfügbar gemacht und gesteuert (z. B. xla_sparse_core_estimate_max_ids für die Grenzwertschätzung oder xla_sc_detect_nan für das Debugging).

Open-Source-Status:

Derzeit ist die Sparsecore-Implementierung intern und wird über libtpu.so bereitgestellt.

Fehlerberichte und Diagnosen:

Kompilierungsfehler im Zusammenhang mit SparseCore-Konfigurationen oder Ressourcenbeschränkungen treten häufig als XLA:TPU-Kompilierungsfehler auf. Diese Fehlermeldungen können wertvolle Informationen zu Problemen wie zu hoch angesetzten Limits für den verfügbaren SPMEM oder die Verwendung nicht unterstützter Konfigurationen liefern.

7. So werden Limits in Tabellen auf SparseCore übertragen

In SparseCore sind „Limits“ grundlegende Konfigurationsparameter, die sich hauptsächlich auf zwei Einstellungen pro Partition für jede Tabelle beziehen, die auf die verfügbaren SparseCores verteilt wird:

• max_ids_per_partition: Hiermit wird die maximale Anzahl der Gesamt-IDs (einschließlich Duplikaten) definiert, die von einem einzelnen SparseCore in einem einzelnen Rechenschritt an eine bestimmte Partition einer bestimmten Tabelle gesendet oder für diese verarbeitet werden können.
• max_unique_ids_per_partition: Hier wird die maximale Anzahl eindeutiger IDs definiert, die von einem einzelnen SparseCore an einen

Übersetzung in physisches Tabellenlayout und Verarbeitung:

• Strategie für das Sharding von Tabellen: Das Sharding von Embedding-Tabellen erfolgt in der Regel modulo über alle SparseCores im System. Das bedeutet, dass jeder SparseCore für eine bestimmte Teilmenge des Vokabulars (Zeilen) jeder Tabelle verantwortlich ist. SparseCore_k wird in der Regel eine ID j auf Grundlage einer Formel wie k = j % num_total_sparse_cores zugewiesen.
• Definition einer Partition: In diesem Kontext bezieht sich eine Partition auf das spezifische Segment einer Einbettungstabelle, für das ein einzelner SparseCore-Vorgang Lookups verarbeitet.
• SPMEM-Pufferzuweisung: Diese Grenzwerte werden vom XLA-Compiler verwendet, um Puffer im Scratchpad-Speicher (SPMEM) auf dem Gerät statisch zu dimensionieren und zuzuweisen. Die Puffer sind so dimensioniert, dass alle erforderlichen Daten zu den IDs für eine bestimmte Partition (bis zu den angegebenen max_ids- und max_unique_ids-Grenzwerten) zur Verarbeitung in SPMEM geladen werden können. Dies ist besonders wichtig für Berechnungen, die nicht elementweise erfolgen, z. B. beim Reduzieren doppelter IDs innerhalb einer Partition (z. B. beim Erstellen einer Compressed Sparse Row-Darstellung (CSR)), bei denen der gesamte relevante Datensatz für die IDs dieser Partition im schnellen Speicher verfügbar sein muss.

• Kompilierte Limits im Vergleich zu beobachteten Limits:

  • Beobachtete Grenzwerte: Die tatsächliche Anzahl der IDs, die für jede Partition zur Laufzeit auf Grundlage der verarbeiteten Eingabedaten ermittelt wurden.
  • Wenn die beobachteten Grenzwerte die kompilierten Grenzwerte überschreiten, kann dies zum Löschen von IDs (wenn allow_id_dropping aktiviert ist) oder zu Fehlern führen.

• Grenzwerte berechnen: Bei der Bestimmung geeigneter Grenzwerte ist eine sorgfältige Analyse der Verteilung der Eingabedaten erforderlich. Für eine beliebige Tabelle (T1), die selbst Teil einer größeren gestapelten Tabelle T sein kann:

  1. Der Eingabe-Batch (z. B. ein 2D-SparseTensor mit der Form [BatchSize, MaxSequenceLength]) wird zuerst auf die verfügbaren SparseCores aufgeteilt. Wenn ein TensorCore beispielsweise mit zwei SparseCores gekoppelt ist, erhält jeder SparseCore möglicherweise einen Sub-Batch mit der Form [BatchSize/2, MaxSequenceLength].
  2. Dieser Unter-Batch wird dann in das COO-Format konvertiert, was zu row_ids und col_ids führt.
  3. Doppelte IDs innerhalb derselben Stichprobe (d.h. Einträge mit demselben row_id und col_id) werden entfernt.
  4. Für jede verbleibende eindeutige col_id (in einer Stichprobe) wird das Ziel-SparseCore, das für diese ID verantwortlich ist, anhand der Mod-Sharding-Regel target_sc_id = col_id % num_total_sparse_cores bestimmt.
  5. Es wird eine Zählung der Gesamtzahl der IDs (ids_per_sparse_core[target_sc_id]++) und der Anzahl der eindeutigen IDs (unique_ids_per_sparse_core[target_sc_id]++, nachdem die Eindeutigkeit für das jeweilige target_sc_id sichergestellt wurde) geführt, die für jedes target_sc_id bestimmt sind.
  6. Der max_ids_per_partition für Tabelle T1 wird dann auf max(ids_per_sparse_core_array) gesetzt.
  7. Analog dazu ist max_unique_ids_per_partition für die Tabelle T1 auf max(unique_ids_per_sparse_core_array) gesetzt.
  8. Wenn die Tabelle T1 eine Komponente einer gestapelten Tabelle ist, werden möglicherweise zusätzliche Transformationen wie Rotationen oder Verschiebungen auf die ID-Verteilungen angewendet, bevor Statistiken aus allen Bestandstabellen summiert werden. So lässt sich die Last auf die Chips verteilen.

Das richtige Festlegen dieser Grenzwerte ist ein Balanceakt: Niedrigere Grenzwerte können potenziell zu einer höheren Leistung führen, da pro Schritt weniger Daten verarbeitet werden müssen und der SPMEM-Druck verringert wird. Wenn sie jedoch zu niedrig festgelegt werden, kann dies zu übermäßigem Mini-Batching oder unerwünschtem ID-Dropping führen.

8. So kommunizieren die einzelnen SparseCores

Die SparseCore-Kommunikation, insbesondere im Zusammenhang mit der Verarbeitung einer Liste von IDs für das Nachschlagen von Einbettungen, basiert auf mehreren koordinierten Mechanismen:

• Mod-Sharding und implizites Routing:
  • Die Einbettungstabellen werden über alle SparseCores im System hinweg mod-sharded.
  • Wenn der Host einen Batch mit Eingabedaten bereitstellt, der anschließend in das COO-Format vorverarbeitet wird (einschließlich col_ids), wird der col_id-Wert verwendet, um zu bestimmen, welcher SparseCore für die jeweilige ID verantwortlich ist: target_sc_id = col_id % num_total_sparse_cores.
  • Jeder SparseCore empfängt und verarbeitet nur die Teilmenge der IDs, die den zugewiesenen Vokabelpartitionen zugeordnet sind. Die Vorverarbeitungsphase auf dem Host ist entscheidend, um die Daten so vorzubereiten, dass jeder SparseCore seine relevanten IDs leicht identifizieren und verarbeiten kann.
• Datenverteilung nach Host:
  • Die Vorverarbeitungslogik des Hosts partitioniert den gesamten Eingabe-Batch und verteilt die relevanten Teile von row_ids und col_ids (sowie alle zugehörigen Features oder Gewichte, falls zutreffend) entweder an den Speicher (HBM), auf den jeder SparseCore direkt zugreifen kann, oder an einen gemeinsamen HBM, aus dem die SparseCores die erforderlichen Daten abrufen.
• Intra-SparseCore-Verarbeitung:
  • Sobald ein SparseCore die zugewiesene Gruppe von IDs für eine bestimmte Tabellenpartition erhalten hat, werden Vorgänge wie die Deduplizierung dieser IDs und das Erfassen der entsprechenden Einbettungsvektoren ausgeführt. Dabei handelt es sich hauptsächlich um lokale Berechnungen, die in den eigenen Kacheln des SparseCore ausgeführt werden und den lokalen SPMEM nutzen.
• Inter-SparseCore communication (All-to-All) (Kommunikation zwischen SparseCores, All-to-All):
  • Nach der ersten Verarbeitungsphase (z. B. das Einbetten von Lookups) kann ein „All-to-All“-Kommunikationsmuster verwendet werden, um Ergebnisse über SparseCores hinweg zu kombinieren oder neu zu verteilen. Dies kann beispielsweise erfolgen, bevor Aktivierungen in eine TensorCore-Schicht eingegeben werden, die Eingaben entsprechend allen ursprünglichen Stichprobenpositionen erwartet. Dies ist wichtig, um den vollständigen Satz von Aktivierungen zu rekonstruieren, wenn der ursprüngliche Eingabe-Batch für die parallele Verarbeitung verteilt wurde.
• Kommunikation mit Tensor-Cores:
  • SparseCores kommunizieren mit TensorCores, um Einbettungsaktivierungen (während des Vorwärtsdurchlaufs) zu senden und Gradienten (während des Rückwärtsdurchlaufs) zu empfangen. Diese Interaktion wird vom XLA-kompilierten Programm orchestriert und umfasst häufig HBM als Zwischenpuffer. Die oben beschriebene Pipelining-Strategie hat großen Einfluss auf das Timing und die Synchronisierung dieser SC-TC-Kommunikation.

Die anfängliche „Verteilung“ von IDs auf die entsprechenden SparseCores wird weitgehend durch das Sharding-Schema und die Host-Vorverarbeitungsschritte übernommen. Bei der nachfolgenden Kommunikation werden SparseCores auf ihren lokalen Daten ausgeführt, möglicherweise gefolgt von kollektiven Kommunikationsvorgängen wie All-to-All, wenn Daten global ausgetauscht oder zwischen SparseCores neu angeordnet werden müssen, bevor sie von den TensorCores weiterverarbeitet werden.

9. Speicherverwaltung in SparseCore

Jeder SparseCore verwaltet effizient mehrere verschiedene Speichertypen, um seine Berechnungen durchzuführen:

• Scratchpad-Speicher (SPMEM):
  • Nature: Ein relativ kleiner, aber sehr schneller lokaler SRAM, der ausschließlich für jeden SparseCore verfügbar ist. Wichtig: SPMEM ist kein Cache. Die Verwendung wird explizit vom XLA-Compiler verwaltet und orchestriert.
  • Zweck: SPMEM wird verwendet, um Daten „opportunistisch zu stagen“. Dazu gehören Ein- und Ausgaben sowie Zwischenergebnisse, die für laufende SC-Berechnungen erforderlich sind. Durch das Bereitstellen von Daten in SPMEM wird die hohe Latenz, die normalerweise mit dem Zugriff auf HBM verbunden ist, erheblich reduziert.
  • Größe: Wie im Abschnitt „Limits“ beschrieben, werden SPMEM-Puffer statisch zur Kompilierzeit dimensioniert. Die Größe basiert auf Parametern wie max_ids_per_partition und max_unique_ids_per_partition. Durch diese statische Zuweisung wird sichergestellt, dass für jeden Vorgang für eine Tabellenpartition (z. B. CSR-Reduzierung) alle erforderlichen Daten für die IDs dieser Partition (bis zu den definierten Limits) in SPMEM passen.
  • Compiler-Optimierungen: Der XLA-Compiler enthält ausgefeilte Optimierungen, um genau zu bestimmen, wie viele Daten und welche spezifischen Datenelemente im SPMEM bereitgestellt werden müssen, um die HBM-Latenz effektiv zu verbergen und die Leistung zu maximieren.
  • Einschränkung für die dynamische Zuweisung: Der SparseCore-Compiler unterstützt derzeit keine dynamische Zuweisung von Scratchpads. Das unterstreicht die entscheidende Bedeutung der statischen Größenanpassung durch die sorgfältige Konfiguration von Limits.
• High Bandwidth Memory (HBM):
  • Nature: Eine große, gemeinsam genutzte Speicherressource, auf die alle SparseCores, TensorCores und das Hostsystem zugreifen können. Primäre Einbettungstabellen werden in HBM gespeichert.
  • Stapelspeicher: Für SparseCore-Vorgänge ist häufig temporärer Speicher in HBM für Zwischenergebnisse erforderlich, die entweder nicht in den begrenzten SPMEM passen oder zwischen größeren Phasen der Verarbeitungspipeline übergeben werden müssen. Die HBM-Stapelauslastung während der Vorwärts- und Rückwärts-Durchläufe kann so geschätzt werden:
    • Forward Pass HBM Stack (single table) ≈ (2 * feature_width + 1) * max_unique_nz_per_row * logical_replica_count * 4 bytes
    • Backward Pass HBM Stack (single table) ≈ 3 * feature_width * max_unique_nz_per_row * logical_replica_count * 4 bytes
  • Heap-Nutzung: HBM bietet auch Platz für den Heap, der vom Host verwaltet wird. Im Heap werden Daten wie Gewichte für dichte Layer, vom Modell verwendete Konstanten und vorab abgerufene Eingabedaten gespeichert. Die Heap-Nutzung steigt in der Regel mit der Anzahl der Schritte, für die der Host Daten vorab abruft (gesteuert durch das Flag maximum_parallel_iterations). Durch mehr Prefetching kann die Leistung verbessert werden, da Host-zu-Gerät-Übertragungen mit Geräteberechnungen überlappen. Es wird jedoch auch mehr HBM verbraucht.
  • Serialisierung zur HBM-Optimierung: Das Flag xla_sc_num_serialized_tables_to_optimize_hbm bietet einen Mechanismus, um zu steuern, wie viele Tabellendaten zu einem bestimmten Zeitpunkt im HBM-Stapelspeicher „live“ gehalten werden. Wenn Sie diese Zahl erhöhen, wird die Verarbeitung für mehr Tabellen effektiv serialisiert. Dadurch kann die maximale HBM-Stapelauslastung reduziert werden, was jedoch aufgrund der geringeren Parallelität auf Kosten der Leistung gehen kann.
• Vektorspeicher (VMEM):
  • VMEM ist ein lokaler Scratchpad-Speicher, der ausschließlich vom TC (TensorCore) verwendet wird. VMEM wird zwar nicht direkt vom SparseCore verwaltet, ist aber ein integraler Bestandteil des Arbeitsspeichersystems, mit dem der SC hauptsächlich über den TensorCore interagiert.

Gesamtstrategie für die Speicherverwaltung:

Die zentrale Speicherverwaltungsstrategie für SparseCore basiert auf der Verwendung des kleinen, schnellen SPMEM für die „heißen“ Daten, die aktiv von einer SparseCore-Kachel verarbeitet werden. Dadurch werden Zugriffe auf den langsameren HBM minimiert. Die konfigurierten Grenzwerte sind der primäre Mechanismus, um sicherzustellen, dass der SPMEM nicht überläuft. HBM wird zum Speichern großer Einbettungstabellen und temporärer Daten verwendet, die entweder die SPMEM-Kapazität überschreiten oder auf verschiedene Verarbeitungseinheiten oder Pipelinephasen verteilt werden müssen. Der XLA-Compiler ist für die Orchestrierung aller Datenbewegungen und Pufferzuweisungen auf Grundlage dieser Architekturprinzipien und der vom Nutzer konfigurierten Grenzwerte verantwortlich.

10. Leistungs- und Speicherengpässe

Um mit SparseCore eine optimale Leistung zu erzielen, müssen Sie potenzielle Engpässe und Möglichkeiten zur Behebung dieser Engpässe kennen. Diese können auf dem Host, im SparseCore selbst oder in der Interaktion mit den TensorCores auftreten.

Häufige Leistungsengpässe:

• Host-Engpass:
  • Problem: Die Host-CPU kann Daten möglicherweise nicht schnell genug vorverarbeiten und an die TPU übertragen, was zu einer Unterauslastung von SparseCores und TensorCores führt. Das ist ein häufiger Leistungsbegrenzer.
  • Abhilfe: Überwachen Sie die Host-CPU-Auslastung und die Messwerte der Eingabepipeline. Optimieren Sie das Laden und Vorverarbeiten von Daten auf der Hostseite (siehe Tipps zur COO-Conversion). Passen Sie das Flag maximum_parallel_iterations an, um das Vorabrufen von Daten zu optimieren.
• Suboptimale TC-/SC-Synchronisierung (kein Pipelining):
  • Problem: Wenn das Pipelining zwischen dem TensorCore und dem SparseCore deaktiviert ist oder nicht effizient funktioniert, muss eine Einheit möglicherweise lange auf die andere warten, was den Gesamtdurchsatz des Systems verringert.
  • Risikominderung: Achten Sie darauf, dass Pipelining aktiviert ist (z. B. tf_xla_disable_full_embedding_pipelining = false oder das entsprechende Äquivalent).
• Durch Limits verursachte Engpässe:
  • Problem:
    • Zu niedrige Limits: Können zu übermäßigem Mini-Batching führen, bei dem Eingabe-Batches in zahlreiche kleinere Unter-Batches aufgeteilt werden, um die strengen Limits einzuhalten. Dadurch wird zwar die Richtigkeit beibehalten, aber jeder Mini-Batch führt zu einem gewissen Verarbeitungsaufwand, der die Gesamtausführung möglicherweise verlangsamt. Wenn allow_id_dropping auf „true“ gesetzt ist, können zu niedrige Grenzwerte auch dazu führen, dass IDs entfernt werden, was sich auf die Modellgenauigkeit auswirkt.
    • Grenzwerte zu hoch (aber noch passend): Sehr hohe Grenzwerte können zwar Mini-Batching verhindern, aber den SPMEM-Druck unnötig erhöhen, wenn die tatsächlichen Datenmerkmale diese Spitzenwerte nur selten erreichen. Sie können auch zu einer höheren HBM-Stapelauslastung führen, als unbedingt erforderlich ist.
    • Kompilierungsfehler: Wenn für die konfigurierten Limits mehr SPMEM- oder HBM-Stack als der verfügbare physische Arbeitsspeicher erforderlich ist, schlägt die Kompilierung fehl.
  • Abhilfe: Prüfen Sie, ob die Limits richtig festgelegt sind.
• Verzerrung der Datenverteilung:
  • Problem: Wenn bestimmte SparseCore-Partitionen im Vergleich zu anderen durchweg eine unverhältnismäßig große Anzahl von IDs erhalten (was auf eine schlechte ID-Verteilung hinweist), werden diese überlasteten SparseCores zu Leistungsengpässen.
  • Abhilfe: Das Mischen von IDs während des Mini-Batch-Prozesses kann dieses Problem bei gestapelten Tabellen beheben, insbesondere bei Tabellen mit „heißen“ Nutzer-IDs. Analysieren Sie die ID-Verteilungen sorgfältig, um angemessene und ausgewogene tabellenspezifische Limits festzulegen.
• Probleme beim Stapeln von Tabellen:
  • Problem:
    • Zu wenige gestapelte Tabellen: Möglicherweise nicht ausreichend, um die ICI-Latenz effektiv zu verbergen oder den Verarbeitungsaufwand angemessen zu reduzieren.
    • Zu viele gestapelte Tabellen: Dies kann dazu führen, dass sehr große logische Tabellen erstellt werden, die sich nur schwer verwalten lassen oder die verfügbaren Ressourcenlimits überschreiten.
  • Abhilfe:
    • Sorgen Sie für eine optimale Anzahl von Tabellen für das Stapeln. Eine allgemeine Richtlinie empfiehlt einen „Sweetspot“ von 5 bis 100 Tabellen für das Stapeln.
• Ineffiziente numerische Werte/Quantisierung:
  • Problem: Die Verwendung der vollen FP32-Genauigkeit, wenn niedrigere Genauigkeitsformate wie BF16 oder quantisierte Ganzzahlen ausreichen (und schnellere Berechnungen ermöglichen), kann zu einem Leistungsengpass führen.
  • Abhilfe: Erwägen Sie Optionen mit geringerer Genauigkeit. Die Quantisierung selbst hat jedoch einen gewissen Overhead und erfordert möglicherweise eine sorgfältige Abstimmung der Quantisierungsparameter, um die Modellgenauigkeit beizubehalten.
• HBM-Bandbreitensättigung:
  • Problem: Wenn zu viele Daten zum und vom HBM übertragen werden, kann die verfügbare HBM-Bandbreite gesättigt werden. Dies kann durch sehr kleine Featurebreiten (die zu einem hohen Padding-Overhead führen), ineffiziente Arbeitsspeicherzugriffsmuster oder eine extrem große Anzahl von Lookups verursacht werden.
  • Abhilfe: Durch Skalieren der Anzahl der TPUs kann die HBM-Bandbreitensättigung verringert werden.

Häufige Speicherengpässe:

• SPMEM-Überlauf (Kompilierungsfehler):
  • Problem: Wenn max_ids_per_partition und max_unique_ids_per_partition zu hoch eingestellt sind, kann der XLA-Compiler möglicherweise nicht genügend SPMEM zuweisen, was zu Kompilierungsfehlern wie "Fixed size allocations (...) do not fit in TileSpmem (...)" führt. Wenn der Begriff (sample_count * feature_width) / kNumTiles (wobei kNumTiles die Anzahl der Kacheln pro SC ist) für das Staging von Operanden im SPMEM der Kachel zu groß ist, können Fehler wie "Gather operand too large..." auftreten.
  • Abhilfe: Reduzieren Sie die Batchgröße oder erhöhen Sie die Anzahl der für die Verarbeitung verwendeten Chips.
• HBM-Stacküberlauf (Laufzeit oder Kompilierung):
  • Problem: Wenn die Kombination aus feature_width, max_unique_nz_per_row und logical_replica_count zu HBM-Stapelspeicheranforderungen führt, die den verfügbaren HBM überschreiten, kann dies zu OOM-Fehlern (Out-Of-Memory) führen, entweder zur Laufzeit oder während der Kompilierung.
  • Abhilfe: Optimieren Sie das Flag xla_sc_num_serialized_tables_to_optimize_hbm, um die HBM-Stapelspeichernutzung zu reduzieren, indem Sie die Verarbeitung von Tabellen serialisieren. Dies geht in der Regel mit Leistungseinbußen einher.
• HBM-Heap-Erschöpfung:
  • Problem: Wird hauptsächlich durch sehr große Gewichte für dichte Ebenen, zahlreiche im Arbeitsspeicher gespeicherte Konstanten oder ein zu aggressives Vorabrufen von Eingaben (hoher maximum_parallel_iterations) verursacht.
  • Abhilfe: Überwachen Sie die Heap-Nutzung mit Tools wie XProf Memory Viewer.
• Padding-Overhead:
  • Problem: Die Einbettungstabellen werden in der Merkmalsdimension auf 32 Byte (entspricht 8 Gleitkommazahlen) aufgefüllt. Kleine Featurebreiten (z. B. 1 Gleitkommazahl) führen daher zu einem erheblichen Padding-Overhead (z. B. sind 7/8 des zugewiesenen Pufferbereichs Padding), was zu verschwendetem HBM führt. Die Vokabeldimension von Tabellen wird ebenfalls so aufgefüllt, dass sie ein Vielfaches der Anzahl der SparseCores im System ist. Diese Auswirkung ist jedoch in der Regel für Tabellen mit einer ausreichend hohen Vokabelgröße vernachlässigbar.

Allgemeine Faktoren, die sich auf Leistung und Arbeitsspeicher auswirken:

• Topologie: Die Anzahl der verfügbaren Chips und ihre Verbindungsarchitektur.
• Batchgröße: Wirkt sich direkt auf die sample_count pro SparseCore aus, was wiederum den Arbeitsspeicherverbrauch und die Rechenlast beeinflusst.
• Datenformatierung: Ein effizientes Datenlayout auf dem Gerät ist entscheidend für eine optimale Leistung.

11. SparseCore-Profil analysieren

Die Analyse eines Leistungsprofils ist ein wichtiger Schritt, um Engpässe zu identifizieren und Optimierungsmöglichkeiten für Ihre SparseCore-Arbeitslasten zu finden.

1. Trace abrufen:
   • Verwenden Sie Profiling-Tools wie XProf, um einen detaillierten Ausführungstrace zu erfassen, während Ihr Modell trainiert wird oder Inferenz ausgeführt wird. Dieser Trace enthält eine Zeitachse der Vorgänge, die auf dem Host, den TensorCores und den SparseCores ausgeführt werden.
2. Trace Viewer untersuchen (z. B. in XProf oder TensorBoard):
   • Aktivitäten des Hosts: Prüfen Sie die Aktivitäten des Hosts. Gibt es erhebliche Lücken bei der TPU-Aktivität? Solche Lücken können darauf hindeuten, dass der Host ein Engpass ist, da er Daten nicht schnell genug bereitstellt. Analysieren Sie die Leistung Ihrer Eingabepipeline.
   • TensorCore- (TC) und SparseCore-Aktivität (SC):
     • Sehen Sie sich die Ausführungszeitachsen für TC und SC an. Werden sie parallel ausgeführt, was auf ein effektives Pipelining hindeutet? Oder gibt es längere Zeiträume, in denen eine Einheit untätig ist und auf die andere wartet?
     • Ermitteln Sie die Vorgänge, die auf dem SC und dem TC am meisten Zeit in Anspruch nehmen (Vorgänge mit der längsten Laufzeit).
     • Visuelle Trace-Ausgaben (oft mit farbigen Blöcken, die verschiedene Vorgänge im Zeitverlauf darstellen, wie TPU:0 SparseCore 1 (pid 1005)) sind von unschätzbarem Wert, um dominante Vorgänge und Leerlaufzeiten visuell zu erkennen.
   • Schrittzeitanalyse: Sehen Sie sich die gesamte Schrittzeit an und erfahren Sie, wie sie sich auf die Hostverarbeitung, die Berechnung von SC und die Berechnung von TC verteilt oder aufgeschlüsselt wird.
3. Arbeitsspeicheranalyse (XProf Memory Viewer):
   • Heap-Nutzung: Verwenden Sie Tools wie den Tab „Memory Viewer“ von XProf, um die HBM-Heap-Nutzung zu prüfen. So lässt sich feststellen, ob große Modellgewichte, Konstanten oder ein zu aggressives Vorabrufen von Eingaben zu viel HBM beanspruchen. Wenn Sie Flags wie --vmodule=best_fit_allocator=1 aktivieren, erhalten Sie möglicherweise Logs zur maximalen Heap-Nutzung.
   • Stack-Nutzung (indirekt): Das direkte Profiling von HBM-Stacks kann komplex sein. Wenn jedoch Fehler vom Typ „Out-Of-Memory“ auftreten und die Heap-Nutzung angemessen erscheint, ist eine Erschöpfung des HBM-Stacks (oft aufgrund zu großer Grenzwerte oder Feature-Breiten) ein starker Verdacht. Die Formeln für die HBM-Stack-Nutzung können bei der Schätzung helfen.
4. Nach bestimmten Mustern suchen:
   • Mini-Batching: Wenn Limits häufig überschritten werden, kann es sein, dass im Trace Mini-Batching zu sehen ist, z. B. eine höhere Anzahl kleinerer SC-Vorgänge als für die globale Batchgröße erwartet. Dies lässt sich oft aus Logs oder durch Beobachten der Aufrufanzahl bestimmter Vorgänge ableiten.
   • Löschen von IDs: Wenn das Löschen von IDs aktiviert ist und erfolgt, können Systemprotokolle Hinweise darauf enthalten. Das wäre auch ein deutliches Zeichen dafür, dass die konfigurierten Grenzwerte für die Eingabedaten zu restriktiv sind.
   • Kompilierungszeiten: Längere Rekompilierungszeiten, insbesondere wenn FDO (Feedback Directed Optimization) aktiviert ist und Grenzwerte häufig angepasst werden, können die gesamte Trainingszeit erheblich verlängern.
5. Mit Flags und Konfiguration korrelieren:
   • Setzen Sie das beobachtete Verhalten im Profil in Beziehung zu Ihren SparseCore-Konfigurationen (Einstellungen in Limit-Dateien, XLA-Flags). Wenn xla_sc_num_serialized_tables_to_optimize_hbm beispielsweise auf einen hohen Wert gesetzt ist, ist mit einer langsameren SC-Leistung, aber einem geringeren HBM-Stack-Verbrauch zu rechnen.
6. Iterativer Prozess:
   • Das Profiling ist oft ein iterativer Prozess. Nehmen Sie eine bestimmte Änderung vor (z. B. ein Limit anpassen oder eine Funktion aktivieren oder deaktivieren), erfassen Sie ein neues Profil und vergleichen Sie es mit dem vorherigen Profil, um die Auswirkungen der Änderung zu sehen.

12. Allgemeine Debugging-Flags

Es können mehrere Flags aktiviert werden, um die Fehlerbehebung bei der Ausführung von SparseCore zu unterstützen. Das Aktivieren dieser Prüfungen führt oft zu Leistungseinbußen. Daher sollten sie in der Regel für Produktionsläufe deaktiviert werden.

• ID-Prüfungen (außerhalb des gültigen Bereichs):
  • Flag: xla_sparse_core_enable_id_bound_check = true
  • Zweck: Ermöglicht die Überprüfung des Hostsystems, um festzustellen, ob Einbettungs-IDs in den Eingabedaten außerhalb des gültigen Vokabularbereichs liegen, der für eine bestimmte Einbettungstabelle definiert ist. So lassen sich Probleme mit falschen oder beschädigten Eingabedaten erkennen.
• NaN-Prüfung:
  • Flag: xla_sc_detect_nan = true
  • Zweck: Ermöglicht die Erkennung von NaN-Werten (Not a Number) in Gleitkommadaten, die auf dem SparseCore verarbeitet werden. Wenn in den Ein- oder Ausgaben verschiedener Compilerdurchläufe ein NaN erkannt wird, wird durch dieses Flag ein Fehler ausgelöst. Diese Fehler enthalten in der Regel Informationen dazu, wo der NaN-Wert aufgetreten ist.
• Bounds Checker (Speicherzugriff):
  • Flag: xla_sc_assert_level=bounds
  • Zweck: Mit diesem Flag wird ein Tool im ASAN-Stil (AddressSanitizer) aktiviert, das speicherzugreifende Anweisungen (z. B. VMEM-Lade-/Speichervorgänge und DMA-Vorgänge) so umschreibt, dass dynamische Prüfungen enthalten sind. Bei diesen Prüfungen wird geprüft, ob der Speicherzugriff innerhalb der zugewiesenen Grenzen des Zielspeicherbereichs liegt.
  • Verhalten: Wenn ein Speicherzugriff außerhalb des zulässigen Bereichs erkannt wird, schlägt die Ausführung fehl.
  • Achtung: Es ist möglich, dass diese Prüfung falsch-positive Ergebnisse liefert, z. B. aufgrund komplexer Zugriffsmuster mit Schrittweite, die von der Prüfung nicht vollständig erfasst werden. Diese Transformation wird in einer späten Phase des Backend-Kompilierungsprozesses angewendet.
• Pufferprüfung (Speicherbeschädigung):
  • Flags:
    • xla_tpu_buffer_contents_sanitizer_config='cores_to_sanitize: [TC, SC_SCS, SC_TILE], sanitizer_mode: LOCAL_ONLY'
    • xla_tpu_verify_launch_id_across_cores=true
  • Zweck: Mit diesen Flags wird sichergestellt, dass Speicherpuffer nicht versehentlich durch nicht verwandte Vorgänge beschädigt oder überschrieben werden. Der Buffer Sanitizer prüft den Inhalt von Puffern, um sicherzustellen, dass sie sich nicht unerwartet ändern.

13. Unterstützung der Quantisierung

Der SparseDenseMatmulOp von SparseCore wurde entwickelt, um Vorgänge für Einbettungstabellen mit 32-Bit-Gleitkomma- (FP32) und Ganzzahldatentypen zu unterstützen. Das Modelltraining erfolgt in der Regel mit FP32-Präzision für Einbettungstabellen. Es kann jedoch eine Quantisierung nach dem Training (Post-Training Quantization, PTQ) angewendet werden. Bei der PTQ können Datentypen mit geringerer Genauigkeit (z. B. 8-Bit-Ganzzahlen) für die Inferenz verwendet werden, was potenziell zu einer besseren Leistung und einem geringeren Speicherbedarf führen kann.

Simulierte Quantisierung:

Die SparseDenseMatmulOp kann für die „simulierte Quantisierung“ konfiguriert werden. In diesem Betriebsmodus werden Einbettungsvektoren zuerst auf eine niedrigere Genauigkeit quantisiert und dann wieder auf eine höhere Genauigkeit (z. B. FP32) dequantisiert, bevor sie in nachfolgenden Berechnungen verwendet werden. Mit dieser Technik können Modelle trainiert werden, wobei die Auswirkungen von Quantisierungsrauschen berücksichtigt werden. Das Training mit simulierter Quantisierung kann die Genauigkeit des endgültigen Modells verbessern, wenn es für die Inferenz vollständig quantisiert wird.

Konfigurationsattribute für SparseDenseMatmulOp (für die Quantisierung):

• quantization_config_num_buckets = 256
  • Mit diesem Attribut geben Sie die Anzahl der diskreten Buckets oder Ebenen an, in die eine 32-Bit-Gleitkommazahl quantisiert wird. Wenn Sie beispielsweise auf 8‑Bit-Ganzzahlen quantisieren, geben Sie in der Regel 2^8 =256 Buckets an.
• quantization_config_low = -X.X
  • Mit diesem Attribut wird der minimale Gleitkommawert im Quantisierungsbereich definiert. Alle Eingabewerte unter diesem angegebenen Mindestwert werden während der Quantisierung auf diesen Mindestwert begrenzt.
• quantization_config_high = Y.Y
  • Dieses Attribut definiert den maximalen Gleitkommawert im Quantisierungsbereich. Alle Eingabewerte, die über diesem angegebenen Höchstwert liegen, werden während der Quantisierung auf diesen Höchstwert begrenzt.

Numerische Werte und Pipelining-Interaktion:

Das numerische Verhalten des Modells kann sich ändern, je nachdem, ob das Pipelining zwischen TensorCore und SparseCore aktiviert ist. Wenn Pipelining aktiv ist, sind die von SparseCore verarbeiteten Gradienten möglicherweise „veraltet“ (aus einer vorherigen Iteration). Dies kann sich auf den Quantisierungsprozess auswirken und möglicherweise die Dynamik des Modelltrainings oder die endgültige Genauigkeit beeinflussen.

14. Anstehende Funktionen und aktuelle Verbesserungen

Das SparseCore-Ökosystem wird ständig weiterentwickelt und verbessert.

Roadmap:

• Mini-Batching auf Dimensionsebene:
  • Diese Funktion ist als Ergänzung zu den vorhandenen Funktionen für Mini-Batching mit Vokabeldimensionen geplant.
  • So können Embedding-Eingaben entlang der Stichprobendimension weiter partitioniert werden. Dies wird durch die Einführung von On-Device-Schleifen erreicht, die Suchvorgänge aus einer Teilmenge von Stichproben gleichzeitig filtern und verarbeiten können. Eine solche Funktion kann hilfreich sein, wenn Sie eine sehr große Anzahl von IDs pro Stichprobe verwalten oder das Load-Balancing zwischen Verarbeitungseinheiten verbessern möchten.
• Verbesserte Unterstützung für das Einbetten mit weniger als 8 Ganzzahlen pro Zeile (geringe Featurebreiten):
  • Beim aktuellen Design wird häufig ein erhebliches Padding für die Breite von Einbettungs-Features verwendet, die weniger als 8 Gleitkommazahlen (entspricht 32 Byte) umfassen. Diese Auffüllung kann zu verschwendetem HBM und potenziell nicht ausgelasteten Rechenressourcen führen. Zukünftige Verbesserungen sollen diese Ineffizienz bei Tabellen mit kleinen Feature-Dimensionen verringern.

Letzte Verbesserungen:

• Staging von Gather-Operanden in HBM:
  • Diese Optimierung trägt dazu bei, die Belastung des gemeinsam genutzten Scratchpad-Speichers (SPMEM) zu verringern, indem einige Eingaben oder Ausgaben von Gather-Vorgängen im größeren HBM bereitgestellt werden.
• Reduzierung der Stack-Arbeitsspeichernutzung:
  • Es wurden Verbesserungen implementiert, um den HBM-Stack-Speicherverbrauch während SparseCore-Vorgängen zu reduzieren, idealerweise ohne die Gesamtleistung oder den Durchsatz negativ zu beeinträchtigen.

Diese Verbesserungen zielen darauf ab, die Leistung, Speichereffizienz und betriebliche Flexibilität von SparseCore für eine noch größere Bandbreite an spärlichen Arbeitslasten zu verbessern.