Le TPU sono circuiti integrati specifici per le applicazioni (ASIC) progettati su misura e sviluppati da Google. Sono progettate appositamente da zero per accelerare i carichi di lavoro di machine learning e intelligenza artificiale. Le TPU sono ottimizzate specificamente per la moltiplicazione di matrici pesanti e le operazioni sui tensori che fungono da elementi costitutivi fondamentali delle reti neurali, alimentando tutto, dai modelli linguistici di grandi dimensioni agli agenti di ragionamento complessi.

Il cuore di una TPU è un'unità di moltiplicazione a matrice (MXU) specializzata che elabora enormi quantità di operazioni con matrici contemporaneamente. Le TPU utilizzano un'architettura "systolic array" che legge i dati una sola volta e li fa fluire continuamente attraverso migliaia di unità logiche aritmetiche (ALU), accumulando risultati senza la necessità di leggere e scrivere costantemente nella memoria. Supportano anche l'aritmetica a precisione ridotta (come la rappresentazione in virgola mobile a 16 bit o 8 bit), consentendo milioni di calcoli al secondo senza sacrificare l'accuratezza richiesta per i modelli di AI.

Le Cloud TPU offrono un supporto nativo ad alte prestazioni per i principali framework di machine learning, principalmente TensorFlow, PyTorch e JAX. Il codice scritto in questi framework viene compilato dal compilatore Accelerated Linear Algebra (XLA), che ottimizza automaticamente il grafico computazionale per essere eseguito in modo efficiente sull'hardware TPU sottostante.

Le TPU eccellono nelle attività di deep learning che richiedono operazioni parallele su matrici di grandi dimensioni. Sono altamente consigliati per:

• Modelli dominati da calcoli matriciali
• Esecuzioni di addestramento che richiedono settimane o mesi per convergere
• Modelli di base di grandi dimensioni e modelli linguistici di grandi dimensioni (LLM)
• Workload con incorporamenti di dimensioni molto grandi, come i motori di consigli avanzati
• Apprendimento per rinforzo continuo e inferenza ad alto volume e bassa latenza

Un pod di TPU è un enorme cluster fisico di chip TPU collegati tra loro tramite una rete specializzata ad alta velocità (come la rete Optical Circuit Switching o Virgo di Google). Un singolo superpod può contenere migliaia di chip interconnessi. Una "slice" è un sottoinsieme più piccolo e dedicato di un pod che puoi affittare. Questa architettura di rete consente agli sviluppatori di scalare i propri carichi di lavoro di AI su enormi quantità di calcolo con poche o nessuna modifica al codice, gestendo in modo efficace l'intera sezione come una singola macchina unificata. 

Il Goodput si riferisce al tempo effettivo e produttivo che l'infrastruttura dedica all'addestramento attivo del modello, anziché all'inattività. L'addestramento di modelli di base di grandi dimensioni comporta un'orchestrazione complessa su migliaia di chip, il che significa che i cicli di calcolo possono essere facilmente sprecati a causa di ritardi nel trasferimento dei dati, ripristini dell'hardware o creazione di checkpoint. Le Cloud TPU sono progettate per massimizzare il goodput attraverso interconnessioni a larghezza di banda elevata e memorizzazione nella cache ottimizzata, garantendo che i cicli di calcolo a pagamento contribuiscano direttamente all'avanzamento del modello.

Gli utenti possono supportare l'affidabilità, il monitoraggio completo e la gestione utilizzando Google Kubernetes Engine (GKE) e Cluster Director. Sfruttando GKE per TPU, i clienti operano in un ecosistema cloud-native dotato di Inference Gateway per il bilanciamento del carico e della capacità di scalare i deployment fino a 130.000 nodi GKE.

Le Cloud TPU supportano l'erogazione di LLM ad alto throughput e bassa latenza principalmente attraverso i più diffusi motori di erogazione open source: vLLM e SGLang:

• vLLM: si integra con Cloud TPU tramite il pacchetto vllm-tpu per gestire i modelli JAX e vllm-torchtpu per i modelli Torch.
• SGLang: offre un servizio TPU ad alte prestazioni utilizzando il suo backend JAX dedicato (sglang-jax) e uno separato per Torch (sglang-torch).

Entrambi i progetti forniscono librerie ottimizzate e specifiche per i pacchetti, su misura per l'hardware TPU, al fine di massimizzare l'efficienza di pubblicazione. Per le guide alla configurazione, consulta la documentazione sull'inferenza di Cloud TPU.