Un approfondimento su SparseCore per i modelli di incorporamento di grandi dimensioni (LEM)

SparseCore è un processore a blocchi specializzato progettato per l'accelerazione ad alte prestazioni di carichi di lavoro che comportano accesso e calcolo irregolari e sparsi della memoria, in particolare su grandi set di dati archiviati nella memoria ad alta larghezza di banda (HBM). Sebbene eccella in attività come l'incorporamento di ricerche, le sue funzionalità si estendono all'accelerazione di una serie di altri workload dinamici e sparsi.

1. Introduzione a SparseCore

Caratteristiche architettoniche principali:

• Architettura a riquadri: comprende più riquadri di calcolo (ogni riquadro è un'unità Dataflow completa con memoria locale e unità di elaborazione proprie) che consentono l'elaborazione parallela.
• Esecuzione dinamica: supporta in modo nativo il flusso di controllo e gli accessi alla memoria dipendenti dai dati, fondamentali per i dati sparsi.
• Elaborazione vettoriale: utilizza attività con vettori piccoli (8 o 16 elementi, a seconda della versione hardware) per un calcolo efficiente.
• Controllo centralizzato: un unico sequencer SparseCore orchestra le attività in tutti i riquadri, garantendo operazioni sincronizzate.
• Supporto del riepilogo dei dati: include operazioni cross-lane specializzate utili per attività come l'ordinamento, il filtraggio e le somme dei prefissi.
• Gerarchia di memoria: sfrutta strategicamente la HBM per l'archiviazione di grandi set di dati e la memoria scratchpad locale (SPMEM) per l'organizzazione temporanea dei dati a cui si accede di frequente, riducendo significativamente la latenza della HBM.

Specifiche in sintesi:

2. Preelaborazione host SparseCore

Una preparazione efficace dei dati è fondamentale per le prestazioni di SparseCore ed è qui che il pre-elaborazione dell'host svolge un ruolo fondamentale. Comprende diverse funzionalità chiave:

• Trasformazione dei dati:
  • Applica le trasformazioni necessarie ai dati di input non elaborati.
  • Gestisci le trasformazioni degli ID, il che è particolarmente importante quando si tratta di impilare funzionalità o tabelle.
  • Converti i dati di input nel formato sparso Coordinate (COO), descritto nella sezione seguente.
  • Partiziona i dati per una distribuzione efficiente tra i diversi SparseCore disponibili sul chip.
• Convalida del limite:
  • Assicurati che le caratteristiche dei dati di input (ad esempio, il numero di ID) siano conformi ai limiti operativi predefiniti di SparseCore, ad esempio max_ids_per_partition e max_unique_ids_per_partition.
  • Se i dati di input superano questi limiti, il livello di preelaborazione dell'host può tentare di segmentare i dati in mini-batch più piccoli che rientrano nei vincoli.
• Data Transfer:
  • Copia in modo efficiente i dati elaborati e convalidati nella memoria ad alta larghezza di banda (HBM) della TPU, rendendoli pronti per l'esecuzione di SparseCore.

Informazioni sull'impilamento delle tabelle:

L'impilamento delle tabelle è una tecnica di ottimizzazione significativa in cui più tabelle di incorporamento vengono combinate logicamente per migliorare l'efficienza della ricerca degli incorporamenti. Questo processo viene in genere gestito automaticamente dal framework ML sottostante.

• Stacking di funzionalità: si verifica quando più funzionalità distinte condividono la stessa tabella di incorporamento sottostante. Un esempio comune è l'utilizzo di un unico dizionario di incorporamento per varie caratteristiche categoriche come i codici postali di contesti diversi.
• Impilamento delle tabelle: in questo scenario, più tabelle di incorporamento distinte vengono impilate insieme. Le tabelle che condividono la stessa dimensione di embedding e la stessa configurazione dell'ottimizzatore vengono spesso raggruppate.

Il vantaggio principale dell'impilamento delle tabelle è la creazione di una dimensione batch effettiva più grande per le operazioni su queste tabelle impilate. Ciò riduce il sovraccarico di calcolo e può essere efficace per nascondere le latenze di comunicazione tra chip (ICI). Per prestazioni ottimali, è consigliabile un numero moderato di tabelle in pila (generalmente compreso tra 5 e 100).

3. Conversione in tensori COO

Prima che i dati possano essere elaborati da SparseCore, vengono comunemente convertiti in un formato di tensore sparso Coordinate (COO). Il formato COO è un modo per rappresentare in modo efficiente le matrici sparse, in genere utilizzando tre array:

• row_ids: un array contenente gli indici di riga per ogni elemento diverso da zero. Nel contesto dell'elaborazione in batch, questo spesso corrisponde alla dimensione batch.
• col_ids: un array contenente gli indici delle colonne per ogni elemento diverso da zero. Per gli incorporamenti, spesso si tratta dei valori delle caratteristiche o degli ID.
• values (facoltativo): un array contenente i valori effettivi degli elementi diversi da zero alle coordinate (row, col) corrispondenti. Per i calcoli dei limiti (descritti in seguito) relativi ai conteggi degli ID, questi valori (incrementi) spesso non vengono presi in considerazione.

Esempio illustrativo:

Considera una matrice sparsa di input che rappresenta batch di ID:

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

Dopo la conversione al formato COO (e potenzialmente dopo la deduplicazione degli ID all'interno dello stesso campione):

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

Questa conversione è fondamentale per il modo in cui SparseCore elabora e distribuisce il lavoro. In particolare, col_ids sono fondamentali per determinare a quale partizione SparseCore specifica appartiene un ID, consentendo uno sharding e una ricerca efficienti.

4. SparsecoreConfig: l'API di alto livello

API di incorporamento specifiche per framework:

• 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 o meccanismi equivalenti come i flag XLA, funge da interfaccia di alto livello per controllare un'ampia gamma di comportamenti di SparseCore. Una comprensione approfondita di questi parametri è fondamentale per un'ottimizzazione efficace del rendimento e per garantire il corretto funzionamento dei modelli.

• disable_table_stacking: bool = False
  • Spiegazione: questo flag controlla se l'impilamento automatico delle tabelle impedisce al framework di impilare le tabelle, il che potrebbe comportare una riduzione delle prestazioni a causa dell'aumento degli overhead e di una minore capacità di nascondere la latenza dell'interconnessione inter-chip (ICI).
  • Predefinito: False (il che implica che l'impilamento delle tabelle è generalmente abilitato per impostazione predefinita laddove il framework lo supporta).
• max_ids_per_chip_per_sample: int = 64
  • Spiegazione: questo parametro stabilisce un limite superiore globale per il numero totale di ID incorporamento che un singolo chip può elaborare da un campione nel batch di input, aggregato in tutte le tabelle. È un meccanismo per gestire le risorse a livello di chip, prima che vengano presi in considerazione limiti più granulari per tabella o per partizione. La regolazione di questo valore dipende in genere dalle caratteristiche specifiche del modello e dalla capacità complessiva del sistema.
  • Predefinito: 64.
• max_ids_per_table: Optional[Dict[str, int]] = None
  • Spiegazione: questo parametro specifica il numero massimo di ID incorporamento (che possono includere duplicati) che possono essere elaborati per ogni tabella logica, considerando tutte le sue partizioni in tutti i SparseCore. Si tratta di un limite più ampio rispetto a max_ids_per_partition. Se una tabella T è suddivisa in P partizioni, questo limite si applica alla somma degli ID indirizzati a tutte le partizioni P. Spesso è correlato a max_ids_per_partition_per_sample e alla dimensione complessiva del batch.
  • Impostazione: in genere configurata utilizzando un file dei limiti (ad esempio, utilizzando il flag xla_sparse_core_max_ids_file), in cui è definito max_ids_per_partition. Questo concetto a livello di tabella è un metodo per impostare i limiti a livello di partizione (max_ids e max_uniques).
  • Predefinito: None (il valore può essere dedotto dai limiti per partizione o da altre configurazioni se non fornito esplicitamente).
• max_unique_ids_per_table: Optional[Dict[str, int]] = None
  • Spiegazione: analogo a max_ids_per_table, ma questo parametro specifica il numero massimo di ID univoci per ogni tabella logica. Si tratta di un'impostazione fondamentale per dimensionare correttamente i buffer sul dispositivo utilizzati nell'elaborazione degli ID univoci e nelle successive operazioni sui vettori.
  • Impostazione: definita comunemente anche in un file dei limiti o derivata da max_unique_ids_per_partition_per_sample.
  • Predefinito: None.
• allow_id_dropping: bool = False
  • Spiegazione: questo flag booleano controlla l'eliminazione degli ID quando il numero di ID riscontrati nei dati di input (limiti osservati) supera i limiti impostati durante la compilazione (ad esempio, max_ids_per_partition).
    • Se True: gli ID che causerebbero il superamento dei limiti vengono eliminati automaticamente. In genere, gli ID all'interno di una partizione vengono elaborati in ordine ordinato e qualsiasi ID che superi il limite del conteggio corrente per il mini-batch designato viene eliminato. Ciò consente al programma di continuare l'esecuzione, ma potrebbe avere un impatto negativo sull'accuratezza del modello.
    • Se False: viene attivato un errore e il processo probabilmente terminerà se i limiti osservati superano i limiti compilati. Questo approccio garantisce l'elaborazione di tutti i dati, ma richiede una configurazione più conservativa dei limiti.
  • Predefinito: False (causando un errore in caso di overflow anziché l'eliminazione silenziosa dei dati).

• initialize_tables_on_host: bool = True

  • Spiegazione: questo flag determina se le tabelle di incorporamento vengono inizializzate sulla CPU host prima di essere trasferite successivamente alla memoria ad alta larghezza di banda (HBM) della TPU. La prassi standard prevede l'inizializzazione delle tabelle sull'host. Se imposti questo valore su True, segui questa convenzione. Se fosse impostato su False, implicherebbe un meccanismo di inizializzazione sul dispositivo, che potrebbe avere implicazioni diverse in termini di prestazioni o prerequisiti di inizializzazione specifici.

• enable_fast_table_initialization: bool = False

  • Spiegazione: inizializza le tabelle direttamente sulla TPU. In questo modo è possibile ridurre i tempi di avvio del modello.

5. Pipelining per le prestazioni

Il pipelining è una tecnica di ottimizzazione delle prestazioni che consente l'esecuzione simultanea di operazioni su TensorCore (TC) e SparseCore (SC). Sovrapponendo questi calcoli, la velocità effettiva complessiva può essere migliorata in modo significativo.

• Meccanismo: in un passaggio di addestramento standard che prevede ricerche di incorporamento sparse (gestite da SC) e calcoli di livelli densi (gestiti da TC), il pipelining consente a SC di lavorare alla sua parte del passaggio i (ad esempio, passaggi in avanti o indietro) mentre TC elabora contemporaneamente una parte diversa dello stesso passaggio i o anche parti di passaggi adiacenti come i-1 o i+1.
• Impatto sui gradienti: SparseCore potrebbe operare su gradienti "obsoleti". Ad esempio, i gradienti calcolati durante la fase di backpropagation del passaggio i potrebbero non essere completamente aggiornati e visibili al SC fino al passaggio i+2.
• Compromesso tra prestazioni e valori numerici: questa esecuzione sovrapposta può portare ad accelerazioni sostanziali, potenzialmente fino a un miglioramento di 2 volte del tempo di esecuzione del passaggio del dispositivo. Tuttavia, le lievi modifiche ai valori numerici (embedding_weights) derivanti dall'utilizzo di gradienti obsoleti potrebbero influenzare il comportamento di convergenza del modello o l'accuratezza finale raggiunta. L'accettabilità di questo compromesso dipende molto dal modello e spesso richiede una convalida empirica.
• Control flag: il pipelining può essere controllato da tf_xla_disable_full_embedding_pipelining. Se imposti questo flag su true, viene disattivato il pipelining completo (sovrapposizione del calcolo di TensorCore e SparseCore), mentre se lo imposti su false (o se la semantica del flag implica l'attivazione quando è false), viene attivato.

Flusso concettuale del pipelining:

• Senza pipelining (flusso sequenziale semplificato):

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

• Con il pipelining (flusso sovrapposto semplificato):

  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
  

  Nota: le fasi di pipelining effettive implementate nell'hardware e nel compilatore possono essere più complesse e spesso includono pre-loop, loop di esecuzione principali e post-loop per gestire le dipendenze dei dati e garantire la correttezza.

6. Il ruolo di XLA

XLA (Accelerated Linear Algebra) è il compilatore specifico per il dominio che traduce grafici computazionali di alto livello, in genere da framework come TensorFlow, in codice macchina altamente ottimizzato e personalizzato per le TPU. Ciò include la generazione delle istruzioni per le operazioni destinate a SparseCore.

Funzioni chiave nel contesto di SparseCore:

• Compilazione di operazioni sparse: XLA è responsabile della compilazione delle operazioni di ricerca di incorporamento (come SparseDenseMatmulOp) e di altri calcoli sparsi in programmi SparseCore eseguibili di basso livello.
• Integrazione dei limiti: utilizza i limiti operativi configurati (ad esempio max_ids_per_partition, max_unique_ids_per_partition, spesso forniti tramite un file dei limiti specificato da flag come xla_sparse_core_max_ids_file) per determinare staticamente le dimensioni e allocare i buffer di memoria sul dispositivo, in particolare all'interno di SPMEM.
• Ottimizzazioni mirate: XLA esegue una serie di ottimizzazioni progettate specificamente per l'architettura SparseCore. Questi possono includere la pianificazione delle istruzioni, le trasformazioni del layout della memoria e la fusione delle operazioni per massimizzare l'efficienza.
• Controllo tramite flag: molti aspetti del comportamento di SparseCore, dei parametri di ottimizzazione e delle strategie di ottimizzazione sono esposti e controllati tramite i flag XLA (ad esempio xla_sparse_core_estimate_max_ids per la stima dei limiti o xla_sc_detect_nan per il debug).

Stato open source:

Al momento l'implementazione di Sparsecore è interna e viene eseguita utilizzando libtpu.so.

Segnalazione e diagnostica degli errori:

Gli errori di compilazione correlati alle configurazioni SparseCore o ai vincoli delle risorse spesso si manifestano come errori di compilazione XLA:TPU. Questi messaggi di errore possono fornire informazioni preziose su problemi come limiti impostati troppo alti per la SPMEM disponibile o l'utilizzo di configurazioni non supportate.

7. Come si traducono i limiti in tabelle su SparseCore

Su SparseCore, i "limiti" sono parametri di configurazione fondamentali che si riferiscono principalmente a due impostazioni per partizione per ogni tabella partizionata (distribuita) tra gli SparseCore disponibili:

• max_ids_per_partition: definisce il numero massimo di ID totali (inclusi i duplicati) che un singolo SparseCore deve inviare o elaborare per una partizione specifica di una determinata tabella in un singolo passaggio di calcolo.
• max_unique_ids_per_partition: definisce il numero massimo di ID univoci che un singolo SparseCore deve inviare o elaborare per un

Traduzione nel layout e nell'elaborazione della tabella fisica:

• Strategia di partizionamento delle tabelle: le tabelle di incorporamento vengono in genere "partizionate in base al modulo" in tutti gli SparseCore del sistema. Ciò significa che ogni SparseCore diventa responsabile di un sottoinsieme distinto del vocabolario (righe) di ogni tabella. Un ID j viene generalmente assegnato a SparseCore_k in base a una formula come k = j % num_total_sparse_cores.
• Definizione di "partizione": in questo contesto, una "partizione" si riferisce al segmento specifico di una tabella di incorporamento per cui una singola SparseCore gestisce le ricerche.
• Allocazione del buffer SPMEM: questi limiti vengono utilizzati dal compilatore XLA per dimensionare e allocare staticamente i buffer all'interno della memoria scratchpad sul dispositivo (SPMEM). I buffer sono dimensionati in modo che tutti i dati necessari relativi agli ID per una determinata partizione (fino ai limiti specificati di max_ids e max_unique_ids) possano essere caricati in SPMEM per l'elaborazione. Ciò è particolarmente importante per i calcoli non elementwise, ad esempio la riduzione di ID duplicati all'interno di una partizione (ad esempio, quando si crea una rappresentazione CSR), in cui l'intero set di dati pertinente per gli ID della partizione deve essere facilmente disponibile nella memoria veloce.

• Limiti compilati e limiti osservati:

  • Limiti osservati: il numero effettivo di ID rilevati per ogni partizione durante l'esecuzione, in base ai dati di input in fase di elaborazione.
  • Se i limiti osservati superano i limiti compilati, possono verificarsi errori o l'eliminazione degli ID (se allow_id_dropping è abilitato).

• Calcolo dei limiti: il processo di determinazione dei limiti appropriati prevede un'attenta analisi della distribuzione dei dati di input. Per una determinata tabella (che chiameremo T1, che potrebbe a sua volta far parte di una tabella in pila più grande T):

  1. Il batch di input (ad esempio, un SparseTensor 2D di forma [BatchSize, MaxSequenceLength]) viene inizialmente suddiviso tra gli SparseCore disponibili. Ad esempio, se un TensorCore è accoppiato a due SparseCore, ogni SparseCore potrebbe ricevere un sotto-batch di forma [BatchSize/2, MaxSequenceLength].
  2. Questo sottoinsieme viene quindi convertito nel formato COO, generando row_ids e col_ids.
  3. Gli ID duplicati all'interno dello stesso campione (ovvero le voci con lo stesso row_id e col_id) vengono rimossi.
  4. Per ogni col_id univoco rimanente (all'interno di un campione), il SparseCore di destinazione responsabile di questo ID viene determinato utilizzando la regola di sharding mod: target_sc_id = col_id % num_total_sparse_cores.
  5. Viene mantenuto un conteggio del numero totale di ID (ids_per_sparse_core[target_sc_id]++) e del numero di ID univoci (unique_ids_per_sparse_core[target_sc_id]++, dopo aver verificato l'univocità per quel target_sc_id specifico) destinati a ogni target_sc_id.
  6. Il max_ids_per_partition per la tabella T1 viene quindi impostato su max(ids_per_sparse_core_array).
  7. Analogamente, il max_unique_ids_per_partition per la tabella T1 è impostato su max(unique_ids_per_sparse_core_array).
  8. Se la tabella T1 è un componente di una tabella in pila, alle distribuzioni degli ID potrebbero essere applicate trasformazioni aggiuntive, come rotazioni o spostamenti, prima di sommare le statistiche di tutte le tabelle costituenti. Ciò contribuisce a bilanciare il carico tra i chip.

L'impostazione corretta di questi limiti è un compromesso: limiti inferiori possono potenzialmente portare a prestazioni migliori (in quanto è necessario elaborare meno dati per passaggio e la pressione SPMEM è ridotta), ma se impostati troppo bassi, possono comportare un mini-batching eccessivo o l'eliminazione indesiderata degli ID.

8. Come comunica ogni SparseCore

La comunicazione SparseCore, in particolare nel contesto dell'elaborazione di un elenco di ID per le ricerche di incorporamento, si basa su diversi meccanismi coordinati:

• Mod sharding and implicit routing:
  • Le tabelle di incorporamento sono mod-shard in tutti gli SparseCore del sistema.
  • Quando l'host fornisce un batch di dati di input (che viene successivamente preelaborato in formato COO, incluso col_ids), il valore col_id viene utilizzato per determinare quale SparseCore è responsabile di quell'ID specifico: target_sc_id = col_id % num_total_sparse_cores.
  • Ogni SparseCore riceve ed elabora solo il sottoinsieme di ID mappati alle partizioni del vocabolario assegnate. La fase di preelaborazione dell'host è fondamentale per preparare i dati in modo che ogni SparseCore possa identificare e utilizzare facilmente gli ID pertinenti.
• Distribuzione dei dati per host:
  • La logica di preelaborazione dell'host partiziona il batch di input complessivo e distribuisce le parti pertinenti di row_ids e col_ids (insieme a eventuali funzionalità o pesi associati, se applicabile) alla memoria (HBM) direttamente accessibile a ogni SparseCore o a una HBM condivisa da cui gli SparseCore recupereranno i dati richiesti.
• Elaborazione intra-SparseCore:
  • Una volta ricevuto il set di ID designato per una determinata partizione della tabella, SparseCore esegue operazioni come la deduplicazione di questi ID e la raccolta dei vettori di incorporamento corrispondenti. Si tratta principalmente di calcoli locali eseguiti all'interno dei riquadri di SparseCore e che utilizzano la relativa SPMEM locale.
• Comunicazione Inter-SparseCore (All-to-All):
  • Dopo la fase di elaborazione iniziale (ad esempio le ricerche di incorporamento), è possibile utilizzare un pattern di comunicazione "all-to-all" per combinare o ridistribuire i risultati tra gli SparseCore (ad esempio, prima di inserire le attivazioni in un livello TensorCore che prevede un input corrispondente a tutte le posizioni di campionamento originali). Questo è fondamentale per ricostruire l'insieme completo di attivazioni se il batch di input originale è stato distribuito per l'elaborazione parallela.
• Comunicazione con i Tensor Core:
  • Gli SparseCore comunicano con i TensorCore per inviare attivazioni di incorporamento (durante il calcolo in avanti) e per ricevere gradienti (durante il calcolo all'indietro). Questa interazione è orchestrata dal programma compilato XLA e spesso coinvolge HBM come buffer intermedio. La strategia di pipelining (descritta in precedenza) influenza notevolmente la tempistica e la sincronizzazione di questa comunicazione SC-TC.

In sostanza, la "distribuzione" iniziale degli ID ai SparseCore appropriati viene gestita in gran parte dallo schema di sharding e dai passaggi di preelaborazione dell'host. La comunicazione successiva prevede che SparseCores operi sui propri dati locali, potenzialmente seguita da operazioni di comunicazione collettiva come all-to-all se i dati devono essere scambiati o riordinati a livello globale tra SparseCores prima di essere ulteriormente elaborati dai TensorCore.

9. Gestione della memoria SparseCore

Ogni SparseCore gestisce in modo efficiente diversi tipi di memoria per eseguire i calcoli:

• Memoria scratchpad (SPMEM):
  • Natura: una SRAM locale relativamente piccola, ma molto veloce, disponibile esclusivamente per ogni SparseCore. È importante notare che SPMEM non è una cache; il suo utilizzo è gestito e orchestrato in modo esplicito dal compilatore XLA.
  • Scopo: SPMEM viene utilizzato per "organizzare i dati in modo opportunistico". Sono inclusi input, output e risultati intermedi necessari per i calcoli SC in corso. L'organizzazione temporanea dei dati in SPMEM riduce significativamente l'elevata latenza in genere associata all'accesso a HBM.
  • Dimensionamento: come descritto nella sezione "Limiti", i buffer SPMEM vengono dimensionati staticamente in fase di compilazione. Il dimensionamento si basa su parametri come max_ids_per_partition e max_unique_ids_per_partition. Questa allocazione statica garantisce che per qualsiasi operazione su una partizione di tabella (ad esempio la riduzione del CSR), tutti i dati necessari per gli ID della partizione (fino ai limiti definiti) possano essere inseriti in SPMEM.
  • Ottimizzazioni del compilatore: il compilatore XLA incorpora ottimizzazioni sofisticate per determinare con precisione la quantità di dati e gli elementi di dati specifici che devono essere organizzati in SPMEM per nascondere efficacemente la latenza HBM e massimizzare il rendimento.
  • Vincolo di allocazione dinamica: il compilatore SparseCore al momento non supporta l'allocazione dinamica dello spazio di lavoro. Ciò evidenzia l'importanza fondamentale del dimensionamento statico attraverso la configurazione accurata dei limiti.
• Memoria a larghezza di banda elevata (HBM):
  • Natura: una risorsa di memoria condivisa di grandi dimensioni accessibile a tutti gli SparseCore, i TensorCore e il sistema host. Le tabelle di incorporamento principali sono archiviate in HBM.
  • Utilizzo dello stack: le operazioni SparseCore spesso richiedono spazio di archiviazione temporaneo in HBM per i risultati intermedi che non rientrano nella SPMEM limitata o devono essere passati tra le fasi più grandi della pipeline di elaborazione. L'utilizzo dello stack HBM durante i passaggi in avanti e indietro può essere stimato come segue:
    • Stack HBM di propagazione in avanti (singola tabella) ≈ (2 * feature_width + 1) * max_unique_nz_per_row * logical_replica_count * 4 byte
    • Stack HBM di propagazione all'indietro (singola tabella) ≈ 3 * feature_width * max_unique_nz_per_row * logical_replica_count * 4 byte
  • Utilizzo dell'heap: l'HBM ospita anche l'heap, che viene gestito dall'host. L'heap memorizza dati come i pesi dei livelli densi, le costanti utilizzate dal modello e i dati di input precaricati. L'utilizzo dell'heap tende ad aumentare con il numero di passaggi per i quali l'host precarica i dati (controllato dal flag maximum_parallel_iterations). Sebbene un maggior numero di prefetching possa migliorare il rendimento sovrapponendo i trasferimenti da host a dispositivo al calcolo del dispositivo, consuma anche più HBM.
  • Serializzazione per l'ottimizzazione HBM: il flag xla_sc_num_serialized_tables_to_optimize_hbm fornisce un meccanismo per controllare quanti dati delle tabelle vengono mantenuti "live" nella memoria dello stack HBM in un determinato momento. L'aumento di questo numero serializza effettivamente l'elaborazione per più tabelle, il che può ridurre l'utilizzo massimo dello stack HBM, ma potrebbe influire sulle prestazioni a causa del parallelismo ridotto.
• Memoria vettoriale (VMEM):
  • La VMEM è una memoria scratchpad locale utilizzata esclusivamente dal TensorCore (TC). Sebbene la VMEM non sia gestita direttamente da SparseCore, è parte integrante dell'ecosistema di memoria con cui interagisce SC, principalmente tramite TensorCore.

Strategia generale di gestione della memoria:

La strategia principale di gestione della memoria per SparseCore ruota attorno all'utilizzo della piccola e veloce SPMEM per i dati "caldi" che vengono elaborati attivamente da un riquadro SparseCore, riducendo al minimo gli accessi alla HBM più lenta. I limiti configurati sono il meccanismo principale per garantire che SPMEM non vada in overflow. La HBM viene utilizzata per archiviare tabelle di incorporamento di grandi dimensioni e dati temporanei che superano la capacità della SPMEM o devono essere condivisi tra diverse unità di elaborazione o fasi della pipeline. Il compilatore XLA è responsabile dell'orchestrazione di tutti gli spostamenti di dati e dell'allocazione dei buffer in base a questi principi architetturali e ai limiti configurati dall'utente.

10. Colli di bottiglia delle prestazioni e della memoria

Per ottenere prestazioni ottimali con SparseCore è necessario comprendere chiaramente i potenziali colli di bottiglia e come risolverli. Questi possono verificarsi sull'host, all'interno di SparseCore stesso o nella sua interazione con i TensorCore.

Colli di bottiglia delle prestazioni comuni:

• Collo di bottiglia dell'host:
  • Problema: la CPU host potrebbe non riuscire a preelaborare i dati e a inviarli alla TPU abbastanza rapidamente, il che comporta un sottoutilizzo di SparseCore e TensorCore. Si tratta di un limite alle prestazioni frequente.
  • Mitigazione: monitora l'utilizzo della CPU host e le metriche della pipeline di input. Ottimizza le routine di caricamento e preelaborazione dei dati lato host (consulta i suggerimenti per la conversione COO). Modifica il flag maximum_parallel_iterations per ottimizzare il precaricamento dei dati.
• Sincronizzazione TC/SC non ottimale (mancanza di pipeline):
  • Problema: se il pipelining tra TensorCore e SparseCore è disattivato o non funziona in modo efficiente, un'unità potrebbe trascorrere molto tempo in attesa dell'altra, riducendo così la velocità effettiva complessiva del sistema.
  • Mitigazione: assicurati che il pipelining sia abilitato (ad esempio, tf_xla_disable_full_embedding_pipelining = false o il suo equivalente).
• Colli di bottiglia indotti dai limiti:
  • Problema:
    • Limiti troppo bassi: possono attivare un mini-batching eccessivo (suddivisione dei batch di input in numerosi batch secondari più piccoli per rispettare i limiti ristretti). Sebbene ciò mantenga la correttezza, ogni mini-batch introduce un sovraccarico di elaborazione, rallentando potenzialmente l'esecuzione complessiva. Se allow_id_dropping è true, limiti eccessivamente bassi possono anche portare all'eliminazione degli ID, il che influisce sull'accuratezza del modello.
    • Limiti troppo elevati (ma comunque adatti): anche se limiti molto elevati potrebbero impedire il mini-batching, potrebbero aumentare inutilmente la pressione SPMEM se le caratteristiche effettive dei dati raramente si avvicinano a questi valori di picco. Potrebbero anche portare a un utilizzo maggiore dello stack HBM rispetto a quanto strettamente necessario.
    • Errori di compilazione: se i limiti configurati richiedono più stack SPMEM o HBM della memoria fisica disponibile, la compilazione non andrà a buon fine.
  • Mitigazione: assicurati che i limiti siano impostati correttamente.
• Asimmetria della distribuzione dei dati:
  • Problema: se alcune partizioni SparseCore ricevono costantemente un numero di ID sproporzionatamente maggiore rispetto ad altre (a indicare una distribuzione degli ID non ottimale), questi SparseCore sovraccarichi diventeranno colli di bottiglia delle prestazioni.
  • Mitigazione: il rimescolamento degli ID durante il processo di mini-batching può contribuire ad alleviare questo problema per le tabelle in pila, in particolare quelle con tabelle utente "calde". Analizza attentamente le distribuzioni degli ID per impostare limiti per tabella appropriati ed equilibrati.
• Problemi di impilamento delle tabelle:
  • Problema:
    • Troppe poche tabelle impilate: potrebbero non essere sufficienti per nascondere efficacemente la latenza ICI o ridurre adeguatamente i sovraccarichi di elaborazione.
    • Troppe tabelle impilate: potrebbe comportare la creazione di tabelle logiche molto grandi che diventano difficili da gestire o potrebbero superare i limiti delle risorse disponibili.
  • Mitigazione:
    • Garantisci un numero ottimale di tabelle per l'impilamento. Una linea guida generale suggerisce un "punto ottimale" di 5-100 tabelle per l'impilamento.
• Numeri/quantizzazione inefficienti:
  • Problema: l'utilizzo della precisione FP32 completa quando formati a precisione inferiore come BF16 o numeri interi quantizzati sarebbero sufficienti (e offrirebbero un calcolo più rapido) può rappresentare un collo di bottiglia delle prestazioni.
  • Mitigazione: esplora le opzioni di precisione inferiore. Tuttavia, tieni presente che la quantizzazione stessa ha un sovraccarico e potrebbe richiedere un'attenta ottimizzazione dei parametri di quantizzazione per mantenere l'accuratezza del modello.
• Saturazione della larghezza di banda HBM:
  • Problema: un movimento eccessivo di dati da e verso HBM, potenzialmente causato da larghezze delle funzionalità molto piccole (che comportano un overhead di padding elevato), pattern di accesso alla memoria inefficienti o un numero estremamente elevato di ricerche, può saturare la larghezza di banda HBM disponibile.
  • Mitigazione: scalare il numero di TPU può contribuire a risolvere il problema della saturazione della larghezza di banda HBM.

Colli di bottiglia della memoria comuni:

• Overflow SPMEM (errore di compilazione):
  • Problema: se max_ids_per_partition e max_unique_ids_per_partition sono impostati su valori troppo elevati, il compilatore XLA potrebbe non essere in grado di allocare SPMEM sufficiente, causando errori di compilazione come: "Fixed size allocations (...) do not fit in TileSpmem (...)". Inoltre, se il termine (sample_count * feature_width) / kNumTiles (dove kNumTiles è il numero di riquadri per SC) è troppo grande per gli operandi di raccolta temporanea all'interno di SPMEM del riquadro, possono verificarsi errori come "Gather operand too large...".
  • Mitigazione: riduci le dimensioni del batch o aumenta il numero di chip utilizzati per l'elaborazione.
• Overflow dello stack HBM (runtime o compilazione):
  • Problema: se la combinazione di feature_width, max_unique_nz_per_row e logical_replica_count comporta requisiti di memoria dello stack HBM che superano l'HBM disponibile, possono verificarsi errori di memoria insufficiente (OOM) in fase di runtime o durante la compilazione.
  • Mitigazione: regola il flag xla_sc_num_serialized_tables_to_optimize_hbm per ridurre l'utilizzo dello stack HBM serializzando l'elaborazione delle tabelle (in genere a scapito delle prestazioni).
• Esaurimento dell'heap HBM:
  • Problema: causato principalmente da pesi di layer densi molto grandi, da numerose costanti memorizzate o dal prefetching degli input eccessivamente aggressivo (maximum_parallel_iterations elevato).
  • Mitigazione: monitora l'utilizzo dell'heap utilizzando strumenti come XProf Memory Viewer.
• Overhead di padding:
  • Problema: le tabelle di incorporamento vengono riempite in modo da essere allineate a 32 byte (equivalenti a 8 numeri in virgola mobile) nella dimensione della funzionalità. Di conseguenza, le larghezze delle funzionalità ridotte (ad esempio, 1 float) comportano un overhead di padding significativo (ad esempio, 7/8 dello spazio buffer allocato è padding), con conseguente spreco di HBM. Anche la dimensione del vocabolario delle tabelle viene riempita in modo da essere un multiplo del numero di SparseCore nel sistema; tuttavia, questo impatto è in genere trascurabile per le tabelle con una dimensione del vocabolario sufficientemente elevata.

Fattori generali che influiscono sulle prestazioni e sulla memoria:

• Topologia: il numero di chip disponibili e la loro architettura di interconnessione.
• Dimensione batch: influisce direttamente su sample_count per SparseCore, che a sua volta influenza il consumo di memoria e il carico di calcolo.
• Formattazione dei dati: garantire un layout efficiente dei dati sul dispositivo è fondamentale per prestazioni ottimali.

11. Analizzare un profilo SparseCore

L'analisi di un profilo di rendimento è un passaggio fondamentale per identificare i colli di bottiglia e scoprire opportunità di ottimizzazione all'interno dei tuoi workload SparseCore.

1. Ottenere una traccia:
   • Utilizza strumenti di profilazione, come XProf, per acquisire una traccia di esecuzione dettagliata durante l'addestramento o l'inferenza del modello. Questa traccia fornirà una sequenza temporale delle operazioni che si verificano sull'host, sui TensorCore e sugli SparseCore.
2. Esamina Trace Viewer (ad esempio, in XProf o TensorBoard):
   • Attività dell'organizzatore: esamina attentamente l'attività dell'organizzatore. Esistono lacune significative nell'attività della TPU? Questi intervalli potrebbero indicare che l'host è un collo di bottiglia e non riesce a fornire i dati abbastanza rapidamente. Analizza il rendimento della pipeline di input.
   • Attività di TensorCore (TC) e SparseCore (SC):
     • Esamina le sequenze temporali di esecuzione sia per TC che per SC. Funzionano in parallelo, indicando un pipelining efficace? Oppure ci sono periodi prolungati in cui un'unità è inattiva, in attesa dell'altra?
     • Identifica le operazioni che richiedono più tempo (operazioni con esecuzione più lunga) sia su SC che su TC.
     • Gli output di traccia visiva (che spesso mostrano blocchi colorati che rappresentano diverse operazioni nel tempo, come TPU:0 SparseCore 1 (pid 1005)) sono preziosi per identificare visivamente le operazioni dominanti e i periodi di inattività.
   • Analisi del tempo di esecuzione dei passaggi: osserva il tempo di esecuzione dei passaggi complessivo e scopri come viene distribuito o suddiviso tra l'elaborazione dell'host, il calcolo SC e il calcolo TC.
3. Analisi della memoria (XProf Memory Viewer):
   • Utilizzo dell'heap: utilizza strumenti come la scheda "Visualizzatore memoria" di XProf per esaminare l'utilizzo dell'heap HBM. Questo può aiutarti a determinare se pesi, costanti o prefetching degli input eccessivamente aggressivi del modello di grandi dimensioni stanno consumando una quantità eccessiva di HBM. L'attivazione di flag come --vmodule=best_fit_allocator=1 potrebbe fornire log dell'utilizzo massimo dell'heap.
   • Utilizzo dello stack (indiretto): sebbene la profilazione diretta dello stack HBM possa essere complessa, se si verificano errori di esaurimento della memoria e l'utilizzo dell'heap sembra ragionevole, l'esaurimento dello stack HBM (spesso dovuto a limiti o larghezze delle funzionalità eccessivamente grandi) è un forte sospetto. Le formule fornite per l'utilizzo dello stack HBM possono aiutarti a stimare questo valore.
4. Cerca pattern specifici:
   • Mini-batching: se i limiti vengono superati di frequente, potresti notare prove di mini-batching nella traccia (ad esempio, un numero maggiore di operazioni SC più piccole del previsto per la dimensione del batch globale). Spesso può essere dedotto dai log o osservando i conteggi delle chiamate di determinate operazioni.
   • Eliminazione degli ID: se l'eliminazione degli ID è attivata e in corso, i log di sistema potrebbero fornire indicazioni in merito. Questo sarebbe anche un chiaro segnale che i limiti configurati sono troppo restrittivi per i dati di input.
   • Tempi di compilazione: tempi di ricompilazione prolungati, soprattutto se l'ottimizzazione basata sul feedback (FDO) è abilitata e i limiti vengono modificati di frequente, possono aumentare notevolmente il tempo di addestramento complessivo.
5. Correlare con flag e configurazione:
   • Collega il comportamento osservato nel profilo alle configurazioni di SparseCore (impostazioni nei file dei limiti, flag XLA). Ad esempio, se xla_sc_num_serialized_tables_to_optimize_hbm è impostato su un valore elevato, potresti aspettarti prestazioni SC più lente, ma un consumo inferiore dello stack HBM.
6. Processo iterativo:
   • La profilazione è spesso un processo di perfezionamento iterativo. Apporta una modifica specifica (regola un limite, attiva o disattiva una funzionalità), acquisisci un nuovo profilo e confrontalo con il profilo precedente per vedere l'impatto della modifica.

12. Flag di debug generali

È possibile attivare diversi flag per facilitare il debug dei problemi relativi all'esecuzione di SparseCore. È importante notare che l'attivazione di questi controlli spesso comporta una penalità di rendimento e, pertanto, in genere devono essere disattivati per le esecuzioni di produzione.

• Controlli ID (fuori intervallo):
  • Bandiera: xla_sparse_core_enable_id_bound_check = true
  • Scopo: consente di eseguire controlli sul sistema host per rilevare se gli ID di incorporamento nei dati di input non rientrano nell'intervallo di vocabolario valido definito per una determinata tabella di incorporamento. In questo modo è possibile rilevare problemi relativi a dati di input errati o danneggiati.
• Controllo NaN:
  • Bandiera: xla_sc_detect_nan = true
  • Scopo: consente il rilevamento di valori NaN (Not a Number) all'interno dei dati in virgola mobile elaborati su SparseCore. Se viene rilevato un NaN negli input o negli output di vari passaggi del compilatore, questo flag genera un errore. Questi errori in genere forniscono informazioni su dove è stato rilevato il NaN.
• Controllo dei limiti (accesso alla memoria):
  • Bandiera: xla_sc_assert_level=bounds
  • Scopo: questo flag attiva uno strumento in stile ASAN (AddressSanitizer) che riscrive le istruzioni di accesso alla memoria (come i carichi/archiviazione VMEM e le operazioni DMA) per includere controlli dinamici. Questi controlli verificano se l'accesso alla memoria rientra nei limiti allocati della regione di memoria di destinazione.
  • Comportamento: se viene rilevato un accesso alla memoria fuori dai limiti, l'esecuzione non andrà a buon fine.
  • Attenzione: è possibile che questo controllo produca falsi positivi, ad esempio a causa di pattern di accesso complessi con stride che non sono completamente compresi dal controllo. Questa trasformazione viene applicata in una fase avanzata del processo di compilazione del backend.
• Buffer checker (corruzione della memoria):
  • Flag:
    • 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
  • Scopo: questi flag contribuiscono a garantire che i buffer di memoria non vengano danneggiati o sovrascritti inavvertitamente da operazioni non correlate. Buffer Sanitizer controlla i contenuti dei buffer per verificare che non cambino in modo imprevisto.

13. Supporto della quantizzazione

SparseDenseMatmulOp di SparseCore è progettato per supportare le operazioni sulle tabelle di incorporamento utilizzando sia tipi di dati interi che a virgola mobile a 32 bit (FP32). Sebbene l'addestramento del modello venga in genere eseguito utilizzando la precisione FP32 per le tabelle di incorporamento, è possibile applicare la quantizzazione post-addestramento (PTQ). PTQ consente l'utilizzo di tipi di dati a precisione inferiore (come gli interi a 8 bit) per l'inferenza, il che può potenzialmente portare a prestazioni migliori e a un footprint di memoria ridotto.

Quantizzazione simulata:

SparseDenseMatmulOp può essere configurato per eseguire la "quantizzazione simulata". In questa modalità operativa, i vettori di incorporamento vengono prima quantizzati a una precisione inferiore e poi dequantizzati a una precisione superiore (ad esempio FP32) prima di essere utilizzati nei calcoli successivi. Questa tecnica consente di addestrare i modelli tenendo conto degli effetti del rumore di quantizzazione. L'addestramento con la quantizzazione simulata può migliorare l'accuratezza del modello finale quando viene quantizzato completamente per l'inferenza.

Attributi di configurazione per SparseDenseMatmulOp (per la quantizzazione):

• quantization_config_num_buckets = 256
  • Questo attributo specifica il numero di bucket o livelli discreti in cui verrà quantizzato un numero in virgola mobile a 32 bit. Ad esempio, quando si esegue la quantizzazione a numeri interi a 8 bit, in genere si specificano 2^8 =256 bucket.
• quantization_config_low = -X.X
  • Questo attributo definisce il valore minimo in virgola mobile nell'intervallo di quantizzazione. Tutti i valori di input inferiori a questo minimo specificato verranno troncati a questo valore minimo durante la quantizzazione.
• quantization_config_high = Y.Y
  • Questo attributo definisce il valore massimo in virgola mobile nell'intervallo di quantizzazione. Qualsiasi valore di input superiore a questo massimo specificato verrà troncato a questo valore massimo durante la quantizzazione.

Interazione con numeri e pipeline:

Il comportamento numerico del modello può variare a seconda che sia attivato il pipelining tra TensorCore e SparseCore. Se il pipelining è attivo, i gradienti elaborati da SparseCore potrebbero essere "obsoleti" (da un'iterazione precedente). Questo può interagire con il processo di quantizzazione e potenzialmente influire sulla dinamica di addestramento del modello o sull'accuratezza finale.

14. Funzionalità in arrivo e miglioramenti recenti

L'ecosistema SparseCore è soggetto a sviluppo e miglioramento continui.

Roadmap:

• Mini-batching della dimensione del campione:
  • Questa funzionalità è pensata per integrare le funzionalità di mini-batching della dimensione del vocabolario esistenti.
  • Ciò consentirebbe un'ulteriore partizione degli input di incorporamento lungo la dimensione del campione. Ciò si otterrebbe introducendo loop sul dispositivo in grado di filtrare ed elaborare le ricerche da un sottoinsieme di campioni alla volta. Una funzionalità di questo tipo potrebbe essere utile per gestire conteggi di ID per campione molto grandi o per migliorare il bilanciamento del carico tra le unità di elaborazione.
• Supporto migliorato per l'incorporamento con meno di 8 numeri interi per riga (larghezza ridotta delle funzionalità):
  • Il design attuale spesso utilizza un padding significativo per incorporare larghezze delle funzionalità inferiori a 8 float (che corrispondono a 32 byte). Questo padding può comportare uno spreco di HBM e risorse di calcolo potenzialmente sottoutilizzate. I miglioramenti futuri mirano a ridurre questa inefficienza per le tabelle con dimensioni delle funzionalità ridotte.

Miglioramenti recenti:

• Staging degli operandi di raccolta in HBM:
  • Questa ottimizzazione contribuisce a ridurre la pressione sulla memoria scratchpad condivisa (SPMEM) consentendo di eseguire lo staging di alcuni input o output dell'operazione di raccolta nella HBM più grande.
• Riduzione dell'utilizzo della memoria dello stack:
  • Sono stati implementati miglioramenti per ridurre il consumo di memoria dello stack HBM durante le operazioni SparseCore, idealmente senza influire negativamente sulle prestazioni o sul throughput complessivi.

Questi miglioramenti sono incentrati sul miglioramento delle prestazioni, dell'efficienza della memoria e della flessibilità operativa di SparseCore per una gamma ancora più ampia di carichi di lavoro di tipo sparse.