SpikingBrain-7B su due vecchi PC: inferenza distribuita CPU-only
In questo esperimento ho provato a far funzionare
SpikingBrain-7B-W8ASpike utilizzando solamente CPU,
senza GPU, distribuendo il modello tra due vecchi computer.
Il progetto nasce soprattutto come esperimento tecnico:
capire se un modello di circa 7,7 miliardi di parametri,
con alcune caratteristiche differenti rispetto a un Transformer tradizionale,
potesse essere eseguito su hardware x86 ormai piuttosto datato.
Il modello utilizzato
Il checkpoint utilizzato è:
Abel2076/SpikingBrain-7B-W8ASpike
basato sul progetto:
BICLab/SpikingBrain-7B
Una delle parti interessanti del modello è la presenza di
dynamic spike activation.
In termini semplici, l’attivazione interna non viene trattata esattamente
come in un normale percorso completamente denso.
Il modello utilizza un meccanismo spike-inspired che introduce
sparsità nelle attivazioni.
Questo non significa che stiamo simulando veri neuroni biologici,
né che si tratti necessariamente di una classica rete neurale spiking
nel senso tradizionale del termine.
È però una caratteristica importante dell’architettura,
e nel progetto CPU ho cercato di mantenere anche queste operazioni
specifiche invece di sostituirle con normali attivazioni dense.
L’hardware
Per l’esperimento ho utilizzato due vecchi PC collegati direttamente
tramite rete Gigabit Ethernet.
-
atlas5:
Intel Core i5-4590, 4 core / 4 thread, circa 31 GB di RAM. -
argo3:
Intel Core i3-3240, 2 core / 4 thread, circa 31 GB di RAM.
In totale sono quindi disponibili circa 62 GB di RAM,
ma il modello non viene semplicemente caricato interamente su entrambi i PC.
Divisione del modello
SpikingBrain-7B contiene 28 layer principali.
Dopo alcuni benchmark ho scelto una divisione molto semplice:
- atlas5: embedding + layer da 0 a 13;
- argo3: layer da 14 a 27 + normalizzazione finale + lm_head.
In pratica il primo computer elabora la prima metà del modello
e invia lo stato nascosto al secondo computer.
Il secondo completa il calcolo e determina il token successivo.
Prompt / Tokenizer
|
v
atlas5
Embedding
Layer 0 ... 13
|
| FP32 hidden state
v
Gigabit Ethernet / Gloo
|
v
argo3
Layer 14 ... 27
Norm
LM Head
Argmax
|
v
Token successivo
La comunicazione tra i due processi utilizza
PyTorch Gloo.
Il collegamento di rete, nei test, non è risultato essere il principale
collo di bottiglia: la maggior parte del tempo viene spesa nei calcoli CPU.
Niente CUDA, Triton o FlashAttention
Una delle condizioni dell’esperimento era utilizzare un percorso
completamente CPU.
Per questo sono state realizzate implementazioni di riferimento
in PyTorch per diverse operazioni necessarie al modello, tra cui:
- RMSNorm;
- SwiGLU;
- Grouped Query Attention con sliding window;
- GLA, Gated Linear Attention;
- dynamic spike activation;
- cache GLA e KV per la generazione autoregressiva;
- QuantLinear con fake quantization.
Il progetto non utilizza CUDA, Triton, FlashAttention o kernel GPU.
L’obiettivo iniziale era prima di tutto verificare la correttezza
del percorso CPU, non ottenere alte prestazioni.
Il problema più curioso: SIGILL e AVX-512
Durante i primi test reali è successo qualcosa di interessante.
La prima generazione distribuita ha funzionato,
ma una successiva esecuzione indipendente è terminata con
SIGILL, cioè una istruzione CPU non valida.
Dopo diverse prove è stato possibile localizzare il problema
in una routine matematica utilizzata dal runtime PyTorch/MKL.
Sul vecchio i3 Ivy Bridge veniva raggiunta una istruzione AVX-512,
che quella CPU naturalmente non supporta.
La soluzione utilizzata per questo esperimento è stata costruire
una versione di PyTorch orientata alla compatibilità con Ivy Bridge,
utilizzando:
USE_CUDA=0 USE_MKL=0 USE_MKLDNN=0 BLAS=OpenBLAS USE_DISTRIBUTED=1 USE_GLOO=1 -march=ivybridge -mtune=ivybridge
Con questo runtime i test successivi non hanno più mostrato
nuovi SIGILL.
Questo risultato riguarda naturalmente la configurazione utilizzata
nel mio laboratorio e non vuole essere una dichiarazione generale
su MKL o PyTorch.
La prima generazione stabile
Dopo aver stabilizzato il runtime ho eseguito cinque sessioni
completamente indipendenti utilizzando il prompt:
Hello
Tutte e cinque hanno prodotto gli stessi token greedy:
Hello, and welcome
Durante questi test:
- nessun SIGILL;
- nessun utilizzo di swap;
- cache coerenti sui 28 layer;
- chiusura corretta dei processi;
- generazione deterministica nelle sessioni provate.
Alcuni prompt reali
Successivamente ho provato anche alcuni prompt normali.
Tra le brevi continuazioni ottenute:
The capital of France is Paris. The capital of Il cielo è un'immagine di Hello, how are you? I am fine, thank
È stato verificato anche il vero chat template del tokenizer,
compresa la gestione del token EOS.
Questi esempi sono solamente smoke test tecnici:
non rappresentano un benchmark sulla qualità del modello.
Quanto è veloce?
Non molto.
Ed era abbastanza prevedibile considerando le CPU utilizzate.
Nella configurazione stabile la generazione richiede indicativamente
circa 13-14 secondi per token.
Il principale collo di bottiglia attuale è
QuantLinear.
Nonostante il nome W8ASpike, il checkpoint che ho analizzato contiene
pesi memorizzati in FP32.
L’implementazione CPU attuale esegue quindi una fake quantization
dei pesi durante il forward, invece di utilizzare direttamente
pesi INT8 e una vera GEMM INT8 ottimizzata per CPU.
Questa è probabilmente una delle aree più interessanti
per eventuali esperimenti futuri.
Il progetto è pubblico su GitHub
Ho deciso di pubblicare codice, test, documentazione
e cronologia degli esperimenti su GitHub.
Sono documentati anche i tentativi falliti e il problema SIGILL,
non solamente i test riusciti.
Repository:
https://github.com/Uraroga/spikingbrain-cpu-cluster
Issue aperta agli sviluppatori di SpikingBrain
Ho inoltre aperto una Issue direttamente nel repository ufficiale
di BICLab/SpikingBrain-7B per presentare l’esperimento
e chiedere agli sviluppatori se un percorso CPU/reference
possa essere interessante per il progetto.
Ho chiesto anche se esistano già lavori simili per W8ASpike su CPU
e quale possa essere una buona direzione per implementare
una vera quantizzazione INT8 con GEMM CPU.
La discussione è disponibile qui:
https://github.com/BICLab/SpikingBrain-7B/issues/43
Un progetto sperimentale
Questo lavoro non vuole essere presentato come una implementazione
ufficiale di SpikingBrain e non vuole rivendicare alcun primato.
È semplicemente un esperimento indipendente:
prendere due vecchi computer, molta RAM, Linux, Docker,
PyTorch e una rete Gigabit e vedere fino a dove è possibile arrivare
senza utilizzare una GPU.
Il risultato, almeno per me, è già interessante:
SpikingBrain-7B-W8ASpike genera realmente testo
attraversando entrambi i computer e mantenendo le sue operazioni
specifiche anche nel percorso CPU.
Il codice e tutti i dettagli tecnici sono disponibili nel repository GitHub.
Progetto sviluppato da Sergio (Uraroga) con assistenza sostanziale
di ChatGPT e OpenAI Codex per progettazione, implementazione,
debug, test e documentazione.
Il progetto è indipendente e non è affiliato,
sponsorizzato o approvato da OpenAI o dagli autori di SpikingBrain.