GPU-Training auf Unity Catalog-Tabellen mit UCDataset beschleunigen

GPUs sind in der Regel der teuerste Teil jedes Trainingslaufs, und die meisten Teams schaffen es nicht, sie konstant auszulasten. Der Engpass liegt meistens nicht beim Modell. Es geht darum, Daten von ihrem Speicherort schnell genug zur GPU zu bringen – und jede Sekunde, in der diese Chips untätig warten, verschwenden Sie sowohl Zeit als auch Geld.

Heute veröffentlichen wir das Unity Catalog Dataset (UCDataset) in der Beta-Version auf AI Runtime. UCDataset ist ein nativer PyTorch-Dateneinleser, der Unity Catalog-Tabellen direkt an jede GPU in Ihrem Trainingsauftrag streamt, ganz ohne Exporte, Kopien und eigenen Code für den Einleser.

Bei simulierten Produktionsworkloads hat UCDataset die Trainingszeit im Vergleich zum Lesen einer exportierten Parquet-Kopie derselben Daten um das 10-Fache verkürzt.

Das Problem: Ihre GPUs warten auf Ihre Daten

Die meisten Trainingsdaten für Ihre Modelle befinden sich bereits in Unity Catalog-Tabellen. Um sie in einen verteilten Trainingsauftrag zu bringen, müssen Sie einen von drei Workarounds wählen, von denen jeder Sie Zeit und Geld kostet:

• Langsame, mühsame Exporte vor jedem Lauf: Das gängigste Muster besteht darin, die Tabelle vor jedem Lauf als Parquet (oder ein anderes Shard-Format) in ein Volume zu exportieren und dann ein Dataset auf diese Dateien zu verweisen. Sie bezahlen für die Speicherung einer zweiten Kopie, sie veraltet, sobald neue Daten eintreffen, und sie fällt aus der Unity Catalog-Herkunft heraus.
• Das Laden der gesamten Tabelle in den Arbeitsspeicher skaliert nicht: Das Einlesen der Tabelle in Pandas mit toPandas() funktioniert bei kleinen Tabellen und schlägt bei großen fehl. Außerdem landet alles auf einem einzigen Node, was einem Auftrag, der auf 32 GPUs trainiert, nicht hilft.
• Das Streaming über ein Warehouse ist langsam: Warehouses wurden nicht von Haus aus dafür entwickelt, GPUs zu füttern, und werden ohne Optimierung gedrosselt, wenn viele GPUs gleichzeitig lesen.

Dann ist da noch der Einleser selbst. Jede GPU und jeder Worker benötigt einen eigenen, ausbalancierten Ausschnitt der Daten. Teams verbringen Zeit damit, dies manuell aufzubauen, und ein versteckter Fehler kann dazu führen, dass das Modell auf duplizierten oder übersprungenen Zeilen trainiert, ohne dass es jemand merkt.

Das Ergebnis: Teure GPUs bleiben untätig, das Training dauert länger als nötig, und ML-Engineers warten Daten-Pipelines, anstatt Modelle zu verbessern.

Wir stellen vor: UCDataset

UCDataset macht Unity Catalog-Tabellen zu einer erstklassigen Datenquelle für das verteilte PyTorch-Training. Sie beschreiben, was gelesen werden soll (eine Tabelle, die benötigten Spalten und einen Zeilenfilter), und UCDataset streamt trainingsfertige Batches an jede GPU. Zu den Vorteilen gehören:

• Auf Geschwindigkeit ausgelegt: UCDataset liest die Dateien der Tabelle direkt aus dem Cloud-Speicher, sodass sich nichts zwischen Ihre Daten und Ihre GPUs stellt. Es teilt die passenden Zeilen gleichmäßig auf alle GPUs und Daten-Loader-Worker auf.
• Trainieren Sie auf dem Lake, nicht auf einer Kopie: Verweisen Sie auf die Live-Tabelle. Filter und Löschvektoren werden beim Lesen der Daten angewendet, und der nächste Lauf übernimmt neu hinzugefügte Zeilen automatisch.
• Durchgehend abgesichert: Trainingsaufträge lesen Daten als die Identität, die sie ausführt, und Unity Catalog-Berechtigungen gelten genau wie überall sonst auf Databricks.

Benchmark: Schnelleres Training bei der Workload eines Empfehlungsmodells

Wir haben UCDataset anhand einer simulierten Produktionsworkload gebenchmarkt. Dabei wurde eine 300 GB große Feature-Tabelle über 8 NVIDIA H100-GPUs gestreamt, um ein Empfehlungsmodell zu trainieren. Im Vergleich zum direkten Laden der Daten aus Parquet sahen wir eine durchschnittliche Verzehnfachung der Trainingsgeschwindigkeit.

Anwendung:

Nehmen wir an, ein Team speichert Benutzer- und Artikel-Features mit einem Klick-Label in einer Tabelle namens main.recsys.ctr_features. Jeden Tag treffen neue Impressionen in der Tabelle ein. Sie möchten ein Click-Through-Rate-Modell auf 16 GPUs über zwei Nodes hinweg direkt aus der Tabelle trainieren.

1. Beschreiben Sie, was gelesen werden soll, und erstellen Sie den Loader. Wählen Sie in train.py die Tabelle aus, behalten Sie die gelabelten Trainingszeilen bei und wählen Sie nur die Spalten aus, die das Modell benötigt. Der Filter wird beim Lesen der Dateien der Tabelle angewendet, sodass Zeilen, die nicht übereinstimmen, niemals Ihren Trainingscode erreichen. UCDataset liefert bereits ganze Batches, verpacken Sie es also in einen Standard-PyTorch-DataLoader mit batch_size=None.

2. Trainieren und starten auf zwei Nodes. Trainieren Sie mit einer standardmäßigen PyTorch-DDP-Schleife. UCDataset weist jeder der 16 GPUs ihren eigenen Anteil an den passenden Zeilen zu, sodass kein Sampler- oder Sharding-Code geschrieben werden muss. Beschreiben Sie den Auftrag in workload.yaml, wodurch train.py mit torchrun auf jedem Node gestartet wird, und führen Sie ihn mit der AI Runtime CLI aus:

Um auf mehr GPUs zu trainieren, ändern Sie num_accelerators in workload.yaml. Der Trainingscode bleibt unverändert. Das vollständige train.py und workload.yaml mit einer Mustertabelle zum Ausprobieren finden Sie im AI Runtime-Kochbuch: Click-Through-Rate-Modell auf einer Unity Catalog-Tabelle.

Erste Schritte

Installieren Sie das Paket von PyPI. Führen Sie in einem Notebook %pip install databricks-uc-dataset aus.

UCDataset ist in der Beta-Version auf AI Runtime verfügbar. So probieren Sie es aus:

• Lesen Sie die Dokumentation: Unity Catalog-Tabellen mit UCDataset in PyTorch laden
• Starten Sie mit einem funktionierenden Beispiel: Kochbuch-Rezept
• Sprechen Sie mit Ihrem Account-Team, wir freuen uns auf Ihr Feedback.

Für unstrukturierte Daten, die in Unity Catalog-Volumes gespeichert sind, verwenden Sie UCVolumeDataset, um Dateien mit lokalem Caching zu streamen. Während UCDataset strukturierte Tabellen verarbeitet, erweitern wir es auch zur Unterstützung des multimodalen Trainings über den FILE-Typ – so können Sie direkt aus einer Tabelle heraus mit Bildern, Audio und Dokumenten trainieren. Die Beta-Version liest Delta-Tabellen im Unity Catalog.

FAQs

Was ist UCDataset?

UCDataset (Unity Catalog Dataset) ist ein nativer Dateneinleser, der Unity Catalog-Tabellen direkt an die GPUs in einem Trainingsauftrag streamt. Es ist in der Beta-Version auf Databricks AI Runtime verfügbar. Sie wählen eine Tabelle und die benötigten Spalten aus, und UCDataset liefert trainingsfertige Batches ohne Exportvorgang oder eigenen Loader-Code.

Wie beschleunigt UCDataset das GPU-Training?

UCDataset hält GPUs ausgelastet, indem es die Datendateien einer Tabelle direkt aus dem Cloud-Speicher liest, ohne dass eine exportierte Kopie dazwischengeschaltet ist. Es teilt die passenden Zeilen gleichmäßig auf jede GPU und jeden Daten-Loader-Worker auf, sodass jede GPU nur ihren eigenen Anteil liest.

Muss ich Unity Catalog-Tabellen nach Parquet exportieren, um darauf zu trainieren?

Nein. UCDataset liest die Live-Tabelle im Unity Catalog, sodass kein Exportschritt und keine zweite Kopie erforderlich ist, die gespeichert oder synchron gehalten werden muss. Es wendet Ihren Filter und die Löschvektoren der Tabelle beim Lesen an, und jeder neue Lauf übernimmt automatisch die Zeilen, die seit dem letzten Lauf hinzugefügt wurden.

Setzt UCDataset Unity Catalog-Berechtigungen durch?

Ja. Auf AI Runtime lesen Trainingsaufträge Daten als die Identität, die sie ausführt, und die Unity Catalog-Berechtigungen für die Tabelle gelten. Da der Auftrag die geschützte Tabelle direkt liest, gibt es keine exportierte Kopie, die separat gesichert werden müsste. In der Beta-Version kann UCDataset keine Tabellen lesen, die Zeilenfilter oder Spaltenmasken haben.

Welche Tabellen kann UCDataset lesen?

UCDataset liest im Unity Catalog registrierte Delta-Tabellen, einschließlich Tabellen mit Löschvektoren.

(Dieser Blogbeitrag wurde mit KI-gestützten Tools übersetzt.) Originalbeitrag