Eseguire Unlimited-OCR su Apple Silicon con C e Metal

In questo esperimento creo un engine ottimizzata per il modello baidu/Unlimited-OCR specifico su Apple Silicon, costruendo un runtime di inferenza nativo in C e Metal, con una piccola API Python per l'uso quotidiano. L'obiettivo è capire quanto lavoro serva per eseguire localmente un modello vision-language reale, mantenendo sotto controllo la memoria, la latenza e il formato dei pesi.

Lo rifarò? Non credo :)

In sintesi

• Task: leggere il contenuto di immagini e documenti e produrre testo direttamente sul dispositivo.
• Architettura: frontend Python, ABI C stabile, core nativo e backend Metal; backend CPU di riferimento per la validazione.
• Visione: encoder SAM con 12 blocchi, encoder CLIP con 24 blocchi e projector 2048 → 1280.
• Decoder: Transformer in stile DeepSeek V2 da 12 layer, hidden size 1280, un layer denso e 11 layer MoE.
• Input: vista globale 1024x1024, crop locali da 640x640, immagini multiple e supporto per PDF.
• Quantizzazione: profili fp16, mixed Q8_0 e dynamic mixed Q4_0 con dequantizzazione fusa nei kernel Metal.
• Risultato: il progetto espone un'API Python installabile con uv, un'ABI C riutilizzabile e un modello mixed-q4 misurato a circa 2,24 GB.

Perché riscrivere l'inferenza di un modello OCR?

UnlimitedOCR è un modello molto interessante perchè ha un utilizzo ~costante di memoria grazie al suo ring buffer, ciò lo rendo un ottimo candidato per l'inferenza locale.

Un modello OCR moderno non è solo un Transformer che restituisce testo. Prima del decoder deve elaborare l'immagine, trasformarla in una sequenza di feature visive e inserirla nel prompt nel punto corretto. Deve poi gestire il prefill, la KV-cache, il routing degli esperti MoE e la generazione token per token.

L'implementazione Python rimane molto comoda per caricare immagini e gestire tokenizer e prompt, ma non è il livello ideale per movimentare grandi quantità di dati ed eseguire moltiplicazioni di matrici a ogni token generato. Da qui nasce unlimitedocr.c: una separazione netta tra l'ergonomia del frontend e il lavoro numericamente pesante del runtime.

Questa impostazione consente di:

• eseguire l'inferenza localmente, senza dipendere da un servizio OCR remoto;
• usare la GPU Apple attraverso Metal e Metal Performance Shaders;
• controllare ogni categoria di memoria, invece di scoprire un out-of-memory a metà richiesta;
• mantenere un backend CPU di riferimento per test, golden fixture e confronti di parity;
• sperimentare con quantizzazione e kernel specializzati senza cambiare l'API pubblica.

Il modello: una pipeline vision-language completa

Unlimited-OCR combina due encoder visivi con un decoder linguistico. La pipeline nativa può essere riassunta così:

immagine
   │
   ▼
SAM: feature locali e multi-scala
   │
   ▼
CLIP: sequenza di token visivi
   │
   ▼
Projector 2048 → 1280
   │
   ▼
prompt + token immagine
   │
   ▼
decoder DeepSeek V2-style + KV cache
   │
   ▼
testo generato

Il primo encoder è SAM, con 12 blocchi, hidden size 768 e windowed attention. Quattro blocchi usano attenzione globale; gli altri lavorano su finestre 14x14. Il secondo è il vision Transformer di CLIP, con 24 blocchi, hidden size 1024 e MLP intermedio da 4096.

Le feature prodotte dai due encoder passano per un projector che le porta nello spazio nascosto del decoder, da 2048 a 1280 dimensioni. Da quel momento immagine e testo confluiscono nella stessa sequenza di embedding: i token visivi occupano lo span corrispondente al placeholder <image>.

Il decoder e il Mixture of Experts

Il decoder ha 12 layer e 10 attention heads da 128 unità. Il primo layer usa una MLP densa; gli 11 successivi usano un Mixture of Experts con 64 esperti instradati, top-6 routing e due esperti condivisi.

Questa scelta è importante per l'inferenza: il modello contiene molti pesi, ma per ogni token non vengono attivati tutti gli esperti. Il runtime deve quindi calcolare il routing, raggruppare i token per esperto, eseguire le proiezioni e ricomporre i risultati senza introdurre copie inutili.

Il decode loop mantiene le attivazioni in fp16, aggiunge i token alla KV cache e termina con un argmax sul language head. Per i token generati viene usata una ring buffer da 128 token, così la memoria del decode resta delimitata anche quando l'output cresce.

Preprocessing: immagini globali, crop e pagine multiple

Il frontend Python prepara le immagini prima di passare la richiesta all'ABI C:

• converte l'immagine in RGB e corregge l'orientamento EXIF;
• applica padding e normalizzazione in [-1, 1];
• dispone i pixel in formato NCHW, fp16 o fp32;
• tokenizza il prompt e inserisce i placeholder visivi;
• costruisce una PreparedRequest con token, mask, viste e limiti per la generazione.

Il preset base usa una vista globale da 1024x1024 e produce 273 token visivi. Il preset gundam seleziona invece una griglia di crop da 640x640 quando l'immagine è abbastanza grande, poi aggiunge la vista globale. I crop vengono elaborati in chunk: la memoria temporanea dipende dal chunk più grande, non dal numero totale di viste.

La stessa interfaccia permette di preparare più pagine come un'unica richiesta. La parte Python include anche una funzione opzionale per rasterizzare un PDF con PyMuPDF; per il core nativo, una pagina PDF è semplicemente un'altra vista globale.

Memoria: il vincolo principale

In questo progetto la gestione della memoria non è un dettaglio dell'allocatore. Il runtime assegna ogni buffer a una categoria e ne traccia l'uso:

• model views;
• KV cache;
• prompt embeddings;
• workspace e feature finali della visione;
• scratch del decoder e degli esperti MoE;
• logits readback e buffer transitori.

Prima di eseguire una richiesta, il core stima il working set per ciascuna categoria e lo confronta con il budget configurato. Se il budget non è sufficiente, la richiesta viene rifiutata con un errore dettagliato.

La pipeline visiva usa scratch buffer riutilizzabili e assembla le feature finali una vista alla volta. Le arene del decoder vengono allocate o ampliate dopo aver conosciuto la dimensione effettiva della parte visiva. In questo modo il picco della visione non si somma inutilmente a quello del decoder.

Quantizzazione: Q8 prima, Q4 dove conviene

Il checkpoint fp16 originale occupa circa 6,7 GB. Per renderlo più pratico su un portatile il progetto supporta profili misti:

La quantizzazione opera su gruppi di 64 valori (group size = 64). La dequantizzazione non crea una copia fp16 completa dei pesi: viene eseguita direttamente nei kernel che fanno embedding lookup, GEMV, GEMM, SwiGLU, combine degli esperti o argmax del language head.

Il profilo Q8 copre le matrici del decoder e, dove sono disponibili kernel fusi, anche i moduli della visione e il projector. Il profilo Q4 estende l'int4 agli esperti MoE e ad altri moduli lineari già validati; l'attenzione resta in Q8 perché è una delle parti più sensibili alla perdita di precisione, mentre normalizzazioni, bias, convoluzioni e attivazioni rimangono in fp16.

Il passaggio a Q4 ha richiesto un quality gate prima del rollout: il packing dei nibble è stato confrontato con alternative più semplici e poi sottoposto a test QA end-to-end contro le baseline fp16 e Q8. Sul percorso di decode degli esperti instradati, la misura riportata dal progetto è di circa 2,7x rispetto a Q8 su M1 Pro; l'argmax del language head è circa 1,2x più veloce.

Dall'ABI C all'API Python

L'interfaccia pubblica C espone un numero volutamente ridotto di operazioni: apertura e chiusura dell'engine, generazione di richieste preparate, risultati tokenizzati, report di memoria e profiling. Il frontend Python costruisce le strutture ctypes, mantiene vivi i buffer durante la chiamata e decodifica i token restituiti dal runtime nativo.

L'utilizzo base resta molto semplice:

from unlimitedocr_c import UnlimitedOCR

ocr = UnlimitedOCR(quant="q4")
text = ocr.generate("page.png")
print(text)
ocr.close()

generate() accetta percorsi locali, URL, bytes, file-like object, immagini PIL e stringhe base64/data URI. Al primo utilizzo il wrapper può scaricare il checkpoint da Hugging Face e convertirlo in un file .uocr memorizzato nella cache; le chiamate successive riutilizzano il modello convertito.

Una singola istanza è sicura da condividere tra thread, ma le chiamate native vengono serializzate da un lock per engine. Se servono realmente più inferenze contemporanee, l'applicazione può creare un'istanza per ogni execution lane: i pesi memory-mapped possono essere condivisi dal sistema operativo, mentre KV cache, scratch e altre strutture runtime restano separati per ogni istanza.

Risultati

Il risultato principale di questo esperimento è un runtime utilizzabile, non soltanto un benchmark isolato:

• il modello gira attraverso una pipeline Metal nativa con un'API Python installabile;
• fp16, mixed Q8_0 e mixed Q4_0 condividono lo stesso formato .uocr e lo stesso contratto di generazione;
• l'engine espone report di memoria e profiling per analizzare i colli di bottiglia;
• il backend CPU di riferimento permette di testare validazione, formati, sequenze e risultati senza dipendere dalla GPU;
• la gestione della visione è chunked e il budget di memoria viene controllato prima del forward;
• la variante mixed-q4 ha superato i test QA end-to-end sul percorso OCR reale.

La versione corrente della libreria Python è la 0.4.1; il progetto è pensato soprattutto per macOS su Apple Silicon. L'ABI C è però separata dal backend: in futuro la stessa interfaccia potrà ospitare altri runtime, anche se CUDA rimane al momento una milestone futura.

Direzioni future: più backend e più osservabilità

Il passo successivo più naturale è estendere il modello dei backend oltre Metal: Vulkan, CUDA e Apple Neural Engine sono già contemplati dall'architettura come direzioni possibili, ma richiedono kernel, gestione della memoria e QA specifici.

In futuro, credo che spingerò per ottimizzare kernel di modelli interessanti su framework già esistenti (e.g., mlx, pytorch), piuttosto che creare formati custom da 0.

Riepilogo: architettura

Citazioni e risorse

• Modello upstream: baidu/Unlimited-OCR su Hugging Face
• Codice: unlimitedocr.c su GitHub
• Repo di riferimento: llama.cpp e ds4