SpenseGPT: Practical One-shot Pruning Enabling Sparse and Dense GEMMs for LLM Inference

Riepilogo Tecnico: SpenseGPT: Pruning One-shot Pratico per GEMM Sparse e Dense nell'Inferenza di LLM

1. Problema

La sparsità semi-strutturata 2:4 (dove due pesi ogni quattro pesi consecutivi sono stati rimossi/pruned verso lo zero) è ampiamente supportata dagli acceleratori moderni (ad es., NVIDIA B200, GPU AMD, Meta MTIA) e offre un incremento teorico della velocità di 2× consentendo all'hardware di saltare le computazioni dei pesi nulli. Tuttavia, l'adozione pratica per l'inferenza di Large Language Model (LLM) affronta due ostacoli primari:

1. Degradazione Rigida dell'Accuratezza: Il vincolo rigido di sparsità 2:4 porta spesso a una significativa perdita di accuratezza nel post-training pruning, in particolare su compiti di ragionamento e di seguire istruzioni (instruction-following), a meno di non eseguire un esteso retraining.
2. Impracticalità delle Esistenti Relaxations: Recenti lavori che tentano di rilassare il vincolo 2:4 (ad es., PATCH, SlideSparse) soffrono di limitazioni critiche. PATCH richiede il supporto di compilatori specializzati (Triton) non ancora ampiamente disponibili, mentre SlideSparse necessita dell'espansione delle attivazioni di input, introducendo overhead a runtime che negano i miglioramenti di velocità end-to-end, specialmente per lunghezze di sequenza ($M$) elevate.

2. Metodologia

Il paper propone Spense, un formato ibrido sparse-dense pratico, e SpenseGPT, un metodo di pruning one-shot post-training basato su di esso.

A. Il Formato Spense
Invece di imporre la sparsità 2:4 su l'intera matrice dei pesi, Spense divide ogni matrice dei pesi in due regioni contigue:

• Una regione densa (che mantiene il $p\%$ dei pesi).
• Una regione 2:4 sparse (pruned con la standard sparsità 2:4).

Questo design sfrutta le librerie GEMM esistenti e altamente ottimizzate (ad es., cuSPARSELt per lo sparse e cuBLASLt per lo dense su GPU NVIDIA) senza richiedere compilatori personalizzati o l'espansione delle attivazioni di input. Gli autori implementano due varianti:

• Output-split: Divide il tensore di output, utilizzando sparse GEMM per una parte e dense GEMM per l'altra.
• Input-split: Divide il tensore di input, elaborando la prima parte con sparse GEMM e la seconda con dense GEMM (con accumulazione).
  Per le varianti GLU degli MLP, gli autori utilizzano l'output-split per $W_{up}$ e $W_{gate}$, e l'input-split per $W_{down}$.

B. Strategie di Selezione degli Indici Densi
Un insight critico è che quali indici intermedi vengono mantenuti densi influisce significativamente sulla qualità del modello. Il paper identifica che gli indici intermedi negli MLP di tipo GLU possono essere permutati senza cambiare l'output matematico, permettendo la selezione di specifici indici da mantenere densi prima di riorganizzarli in un blocco contiguo per l'esecuzione. Vengono proposte due strategie:

1. SpenseGPT (Strategia Semplice): Seleziona l'ultimo $p\%$ degli indici intermedi affinché rimangano densi. Questo si allinea con l'ordine di elaborazione per colonne di SparseGPT, dove le colonne elaborate vengono congelate e gli errori di ricostruzione sono compensati aggiornando le colonne rimanenti a destra. Mantenere le colonne dense a destra fornisce un "buffer di compensazione".
2. SpenseGPT+ (Strategia Consapevole della Perdita/Loss-Aware): Seleziona gli indici densi basandosi su una stima della perdita di ricostruzione. Ispirandosi a Wanda, calcola un punteggio di importanza per ogni indice intermedio aggregando i contributi attraverso le proiezioni gate, up e down. Il punteggio combina la magnitudo del peso e la magnitudo dell'attivazione di input ($S_{ij} = |W_{ij}|\sqrt{H_{j,j}}$). Gli indici con il punteggio più alto vengono selezionati per rimanere densi, poi permutati nella regione contigua prima di applicare SparseGPT alle restanti regioni sparse.

3. Contributi Chiave

• Formato Spense: Un formato ibrido sparse-dense che rilassa il rigido vincolo 2:4 pur mantenendo la compatibilità con le esistenti librerie GEMM ad alte prestazioni, evitando il supporto di compilatori personalizzati o l'overhead di espansione delle attivazioni.
• Importanza della Selezione dei Densi: Dimostrazione che la selezione degli indici densi è un fattore critico per mantenere la qualità del modello, poiché una selezione ingenua porta a gravi cali di performance.
• Strategie di Pruning: Introduzione di SpenseGPT (selezione del tail contiguo) e SpenseGPT+ (selezione basata sulla perdita di ricostruzione) per il pruning one-shot.
• Velocità nel Mondo Reale: La prima dimostrazione di un pruning one-shot che ottiene miglioramenti reali della velocità di decoding end-to-end di LLM su GPU recenti (B200) utilizzando tensor core sparse semi-strutturati preservando al contempo la qualità del modello.

4. Risultati Sperimentali

I metodi sono stati valutati su Qwen3-32B e Seed-OSS-36B utilizzando precisione FP8 su GPU B200 con il framework vLLM. I benchmark includevano AIME (matematica), GPQA (conoscenza), IFEval (seguire istruzioni) e LiveCodeBench (coding).

• Preservazione dell'Accuratezza:
  • Il pruning 2:4 rigido (sparsità 50%) ha causato cali significativi di accuratezza (ad es., il punteggio medio di Seed-OSS-36B è sceso da 73.86 a 65.18 con SparseGPT).
  • SpenseGPT+ al 25% di sparsità ha raggiunto un punteggio medio di 73.12 su Seed-OSS-36B, quasi eguagliando il baseline denso (73.86).
  • Su Qwen3-32B, SpenseGPT+ al 25% di sparsità ha raggiunto 70.76, vicino al baseline denso di 71.68.
  • I confronti hanno mostrato che SpenseGPT+ supera significativamente SpenseGPT e altri baseline (HyperPrune, Wanda) selezionando correttamente gli indici densi.
• Velocità End-to-End:
  • Spense ha ottenuto fino a 1.2× di velocità di decoding end-to-end rispetto al baseline denso su GPU B200.
  • Al contrario, SlideSparse non ha fornito una velocità significativa a causa degli overhead non-GEMM derivanti dall'espansione delle attivazioni di input.
• Efficienza:
  • Il pruning di modelli da 32B–36B ha richiesto circa 1.33–2.15 ore su una singola GPU B200, paragonabile a SparseGPT e molto più veloce di HyperPrune. ProxSparse è andato in esaurimento della memoria (out of memory) su questi modelli.

5. Significato e Claim

Il paper sostiene di presentare la prima dimostrazione di pruning one-shot di velocità di decoding end-to-end reale di LLM derivante da tensor core sparse semi-strutturati su hardware recente (specificamente GPU B200) mantenendo al contempo la qualità del modello.

Gli autori sottolineano che il loro approccio colma il divario tra l'accelerazione hardware teorica e l'implementazione pratica. Evitando la necessità di compilatori specializzati o l'espansione degli input, Spense consente l'utilizzo immediato di librerie ottimizzate esistenti. Il lavoro evidenzia che il semplice rilassamento del vincolo di sparsità non è sufficiente; la strategia per selezionare quali pesi rimangono densi è fondamentale per preservare le capacità di ragionamento e di seguire le istruzioni del modello.

Limitazioni Riconosciute:

• La velocità dipende dall'efficienza delle librerie GEMM sottostanti e varia con la dimensione del batch, la lunghezza della sequenza e l'hardware.
• SpenseGPT+ si basa su dati di calibrazione; se la distribuzione di deployment differisce significativamente dal set di calibrazione, gli indici densi selezionati potrebbero non essere ottimali.