Reducing High-Bandwidth Memory Bottlenecks in JAX-Based LLM Training with Host Offloading
Di cosa parla
Si spiega che, nell’addestramento di grandi modelli di linguaggio con JAX, la memoria ad alta velocità delle schede di calcolo si esaurisce prima che la potenza di calcolo sia pienamente sfruttata. Pesi del modello, gradienti e stati dell’ottimizzatore occupano molta memoria. L’idea illustrata è spostare temporaneamente parte di questi dati nella memoria principale del computer per alleviare il collo di bottiglia.
Cosa permette di osservare
Serve a esplorare quando conviene trasferire dati tra la memoria rapida della scheda e la memoria principale per addestrare modelli più grandi, e quali compromessi progettuali tra uso della memoria e velocità possono emergere usando JAX.
Dalla fonte
Large language model (LLM) training workloads increasingly run into GPU memory limits before compute is fully used. Model weights, gradients, optimizer states,...