Accelera l'addestramento GPU sulle tabelle di Unity Catalog con UCDataset

Le GPU sono di solito la parte più costosa di qualsiasi processo di addestramento e la maggior parte dei team non riesce a sfruttarle a pieno regime. Il collo di bottiglia di solito non è il modello. Si tratta di trasferire i dati da dove si trovano alla GPU abbastanza velocemente e, ogni secondo in cui quei chip rimangono inattivi, stai sprecando tempo e denaro.

Oggi rilasciamo la versione Beta di Unity Catalog Dataset (UCDataset) su AI Runtime. UCDataset è un data loader nativo di PyTorch che esegue lo streaming delle tabelle di Unity Catalog direttamente su ogni GPU nel tuo job di addestramento, senza esportazioni, senza copie e senza codice di caricamento personalizzato.

Su carichi di lavoro di produzione simulati, UCDataset ha ridotto il tempo di addestramento di 10 volte rispetto alla lettura di una copia Parquet esportata degli stessi dati.

Il problema: le tue GPU sono in attesa dei tuoi dati

La maggior parte dei dati di addestramento per i tuoi modelli si trova già nelle tabelle di Unity Catalog. Portarli in un job di addestramento distribuito significa scegliere una delle tre soluzioni alternative, e ognuna ti costa tempo e denaro:

• Esportazioni lente e noiose prima di ogni esecuzione: Il pattern più comune consiste nell'esportare la tabella in Parquet (o un altro formato di shard) in un volume prima di ogni esecuzione, quindi puntare un dataset verso i file. Paghi per archiviare una seconda copia, diventa obsoleta nel momento in cui arrivano nuovi dati e non rientra nel lineage di Unity Catalog.
• Estrarre l'intera tabella in memoria non è scalabile: Leggere la tabella in pandas con toPandas() funziona per tabelle piccole e fallisce su quelle grandi. Mette inoltre tutto su un unico nodo, il che non aiuta un job che si addestra su 32 GPU.
• Lo streaming tramite un warehouse è lento: I warehouse non sono stati creati nativamente per alimentare le GPU e, senza ottimizzazione, verranno limitati quando molte GPU leggono contemporaneamente.

Poi c'è il loader stesso. Ogni GPU e ogni worker ha bisogno della propria porzione bilanciata di dati. I team perdono tempo a costruire questo sistema manualmente e un piccolo bug può far sì che il modello si addestri su righe duplicate o saltate senza che nessuno se ne accorga.

Il risultato: GPU costose rimangono inattive, l'addestramento richiede più tempo del dovuto e gli ingegneri ML gestiscono infrastrutture dati invece di migliorare i modelli.

Introduzione a UCDataset

UCDataset rende le tabelle di Unity Catalog una fonte di dati di primo livello per l'addestramento distribuito con PyTorch. Descrivi cosa leggere (una tabella, le colonne di cui hai bisogno e un filtro per le righe) e UCDataset esegue lo streaming di batch pronti per l'addestramento su ogni GPU. I vantaggi includono:

• Creato per la velocità: UCDataset legge i file della tabella direttamente dallo storage cloud, quindi nulla si frappone tra i tuoi dati e le tue GPU. Suddivide le righe corrispondenti in modo uniforme tra ogni GPU e ogni worker del data loader.
• Addestrati sul lake, non su una copia: Punta alla tabella live. I filtri e i vettori di eliminazione vengono applicati man mano che i dati vengono letti e l'esecuzione successiva acquisisce automaticamente le nuove righe aggiunte.
• Governance end-to-end: I job di addestramento leggono i dati con l'identità che li esegue e le autorizzazioni di Unity Catalog si applicano esattamente come altrove su Databricks.

Benchmark: addestramento più rapido su un carico di lavoro per modelli di raccomandazione

Abbiamo eseguito un benchmark di UCDataset su un carico di lavoro di produzione simulato, effettuando lo streaming di una tabella di feature da 300 GB su 8 GPU NVIDIA H100 per addestrare un modello di raccomandazione e abbiamo riscontrato in media un miglioramento di 10 volte nel tempo di addestramento rispetto al caricamento dei dati direttamente da Parquet.

Come utilizzarlo:

Supponiamo che un team mantenga le feature di utenti e articoli con un'etichetta di clic memorizzata in una tabella main.recsys.ctr_features. Ogni giorno arrivano nuove impressioni nella tabella. Vogliono addestrare un modello di click-through-rate su 16 GPU in due nodi, direttamente dalla tabella.

1. Descrivi cosa leggere e crea il loader. In train.py, seleziona la tabella, mantieni le righe di addestramento etichettate e seleziona solo le colonne di cui il modello ha bisogno. Il filtro viene applicato durante la lettura dei file della tabella, quindi le righe che non corrispondono non raggiungono mai il tuo codice di addestramento. UCDataset produce già batch interi, quindi incapsulalo in un DataLoader standard di PyTorch con batch_size=None.

2. Addestra e avvia su due nodi. Addestra con un ciclo DDP standard di PyTorch. UCDataset assegna a ciascuna delle 16 GPU la propria quota di righe corrispondenti, quindi non c'è alcun codice di campionamento o sharding da scrivere. Descrivi il job in workload.yaml, che avvia train.py con torchrun su ogni nodo, ed eseguilo con la CLI di AI Runtime:

Per eseguire l'addestramento su più GPU, modifica num_accelerators in workload.yaml. Il codice di addestramento rimane lo stesso. Il file train.py completo e workload.yaml, con una tabella di esempio su cui provarli, si trovano nel cookbook di AI Runtime: Modello di click-through-rate su una tabella di Unity Catalog.

Inizia

Installa il pacchetto da PyPI. In un notebook, esegui %pip install databricks-uc-dataset.

UCDataset è disponibile in Beta su AI Runtime. Per provarlo:

• Leggi la documentazione: Carica le tabelle di Unity Catalog in PyTorch con UCDataset
• Inizia da un esempio funzionante: ricetta del cookbook
• Parla con il tuo team di account, saremmo felici di ricevere il tuo feedback.

Per i dati non strutturati memorizzati nei volumi di Unity Catalog, usa UCVolumeDataset per eseguire lo streaming dei file con caching locale. Mentre UCDataset gestisce le tabelle strutturate, lo stiamo anche espandendo per supportare l'addestramento multimodale tramite il tipo FILE, consentendoti di addestrare su immagini, audio e documenti direttamente da una tabella. La versione Beta legge le tabelle Delta in Unity Catalog.

Domande frequenti

Cos'è UCDataset?

UCDataset (Unity Catalog Dataset) è un data loader nativo che esegue lo streaming delle tabelle di Unity Catalog direttamente sulle GPU in un job di addestramento. È disponibile in Beta su Databricks AI Runtime. Scegli una tabella e le colonne di cui hai bisogno, e UCDataset fornisce batch pronti per l'addestramento senza passaggi di esportazione o codice di caricamento personalizzato.

In che modo UCDataset accelera l'addestramento su GPU?

UCDataset mantiene le GPU attive leggendo i file di dati di una tabella direttamente dallo storage cloud, senza una copia esportata nel mezzo. Suddivide le righe corrispondenti in modo uniforme tra ogni GPU e worker del data loader, in modo che ogni GPU legga solo la propria quota.

Devo esportare le tabelle di Unity Catalog in Parquet per addestrarvi sopra?

No. UCDataset legge la tabella live in Unity Catalog, quindi non c'è alcun passaggio di esportazione e nessuna seconda copia da archiviare o mantenere sincronizzata. Applica il filtro e i vettori di eliminazione della tabella durante la lettura e ogni nuova esecuzione acquisisce le righe aggiunte dall'ultima volta.

UCDataset applica le autorizzazioni di Unity Catalog?

Sì. Su AI Runtime, i job di addestramento leggono i dati con l'identità che li esegue e si applicano le autorizzazioni di Unity Catalog sulla tabella. Poiché il job legge direttamente la tabella soggetta a governance, non c'è alcuna copia esportata da proteggere separatamente. Nella versione Beta, UCDataset non può leggere tabelle che dispongono di filtri per righe o maschere per colonne.

Quali tabelle può leggere UCDataset?

UCDataset legge le tabelle Delta registrate in Unity Catalog, incluse le tabelle con vettori di eliminazione.

(Questo post sul blog è stato tradotto utilizzando strumenti basati sull'intelligenza artificiale) Post originale