MiCoPro: End-to-End Mixed Precision HW/SW Co-design with HW-aware Proxy Model

Technische Samenvatting: MiCoPro: End-to-End Mixed Precision HW/SW Co-design met een HW-bewust Proxy Model

Probleemstelling
Gekwantiseerde Neurale Netwerken (QNN's) zijn essentieel voor de implementatie van AI op hulpbron-beperkte edge-apparaten. Hoewel laag-specifieke mixed-precision kwantisatie (MPQ) een weg biedt om efficiëntie te maximaliseren zonder aan nauwkeurigheid in te boeten, kampen bestaande exploratie-algoritmen met aanzienlijke beperkingen. Conventionele methoden worstelen met de hoge dimensionaliteit en complexe correlaties van de MPQ-zoekruimte, en vertrouwen vaak op inefficiënte heuristische zoekmethoden of machine learning-technieken die een excessief evaluatiebudget vereisen. Bovendien falen standaard hardware-kostenmetrieken, zoals Bit Operations (BOPs), er vaak in om de werkelijke latentie op specifieke hardware-doelen nauwkeurig te voorspellen omdat ze architecturale nuances zoals operand packing, geheugenbandbreedte en SIMD-utilisatie negeren. Tot slot is er een gebrek aan een uniform, end-to-end framework dat efficiënt MPQ-schema's kan verkennen op basis van hardwarekenmerken en de resulterende modellen direct van PyTorch naar bare-metal C-code kan implementeren over diverse platforms.

Methodologie
De auteurs stellen MiCo voor, een holistisch framework voor MPQ-exploratie, en de uitbreiding ervan MiCoPro, die een Hardware-Aware Proxy (HAP) model integreert.

1. MPQ Exploratie-algoritme (MiCo):

   • Formulering: De zoektocht wordt behandeld als een budget-beperkt black-box optimalisatieprobleem.
   • Surrogaatmodellering: Het framework maakt gebruik van Random Forest (RF) ensemble learning om de modelnauwkeurigheid te voorspellen op basis van kwantisatieconfiguraties. RF is gekozen vanwege het vermogen om niet-lineaire interacties in discrete zoekruimtes te modelleren met beperkte trainingsdata.
   • Sampling Strategie:
     • Orthogonale Initiële Sampling: In plaats van uniforme random sampling genereert de methode een initiële set samples die extreme bitbreedtes bevat (bijv. W8A8, W1A1) en orthogonale combinaties waarbij specifieke lagen worden gekwantiseerd terwijl andere op volle precisie blijven. Dit zorgt ervoor dat de voorspeller de grenzen en interacties van de zoekruimte efficiënt leert.
     • Near-Constraint Sampling (NCS): Een genetisch algoritme genereert kandidaat-schema's binnen een Region of Interest (ROI) nabij de hardware-restrictie (bijv. de latentie- of BOPs-limiet). De ROI begint smal en breidt zich uit als er geen haalbare oplossingen worden gevonden, waarbij prioriteit wordt gegeven aan schema's die de nauwkeurigheid maximaliseren nabij de grens van de restrictie.
   • Evaluatiestrategie: Het framework maakt gebruik van een tweeledige evaluatiebenadering. Voor bitbreedtes $\ge$ 4 bits wordt Post-Training Quantization (PTQ) gebruikt voor snelle schatting. Voor sub-4-bit configuraties wordt een korte Quantization-Aware Training (QAT) re-training fase toegepast om de nauwkeurigheid te herstellen, wat een balans vindt tussen zoekkosten en precisie.

2. Hardware-Aware Proxy (MiCoPro):

   • Kernel Profiling: Het systeem verzamelt cyclus-nauwkeurige latentiegegevens van mixed-precision kernels (MatMul, Conv2D) over diverse vormen en precisiecombinaties op de doelhardware.
   • Feature Engineering: Om de beperkingen van BOPs te overwinnen, introduceren de auteurs Composite Bit Operations (CBOPs). Deze features breken de hardwarekosten af in:
     • BMACs (Bit Multiply-Accumulate): $\max(b_w, b_a) \times MACs$, wat de bottleneck van mixed-width operand extensie vastlegt.
     • ALoads/WLoads: Afzonderlijke termen voor activatie- en gewichtsgeheugentrafiek, wat de asymmetrische impact op bandbreedte vastlegt.
   • Kalibratie & Transfer: Een netwerk-afhankelijke min-max kalibratie stemt de proxy-voorspellingen af op de werkelijke hardware-latenties bij precisie-extremen. Daarnaast maakt een transfer learning mechanisme het mogelijk om de proxy aan te passen aan nieuwe hardwarevarianten (bijv. verschillende RISC-V core configuraties) met minimale nieuwe profiling data door de residuele fout van een vooraf getraind model te leren.

3. End-to-End Implementatie:

   • Het framework bevat een deployment pipeline die het geoptimaliseerde PyTorch-model converteert naar een PyTorch FX graph.
   • Het ondersteunt graafconversie voor specifieke simulators (bijv. BitFusion) en genereert bare-metal C-code voor CPU's.
   • Het maakt gebruik van een aangepaste mixed-precision kernel library (uitbreiding van baremetal-NN) met SIMD-geaccelereerde varianten voor 8/4/2/1-bit operaties, wat directe executie op RISC-V processoren en accelerators mogelijk maakt.

Belangrijkste Bijdragen

• Nieuw Exploratie-algoritme: Een Random Forest-gestuurde Near-Constraint Sampling (RF+NCS) strategie die MPQ-schema's efficiënt verkent en in vergelijking met heuristische en reinforcement learning-gebaseerde methoden een lagere nauwkeurigheidsdaling bereikt onder strikte restricties.
• Hardware-Aware Proxy Model: Een latentie-voorspellingsmethode met behulp van CBOPs en transfer learning die de traditionele BOPs metrieken aanzienlijk overtreft in het voorspellen van werkelijke hardware-latentie, wat echte hardware-bewuste optimalisatie mogelijk maakt.
• Holistisch Framework: Een open-source, end-to-end pipeline die de brug slaat tussen MPQ-exploratie en implementatie, met ondersteuning voor diverse targets inclus_ief toegewijde accelerators (BitFusion) en SIMD-uitgebreide RISC-V processoren.

Experimentele Resultaten
Het framework werd geëvalueerd op een divers scala aan modellen (CNNs, ResNets, ViT, TinyLLaMa) en hardware-targets (BitFusion accelerator, VexiiMiCo RISC-V varianten).

• Zoekefficiëntie: In PTQ-zoektochten presteerde MiCo consistent beter dan state-of-the-art methoden (HAWQ-V3, HAQ, BOMP, w-based) over diverse modellen en restricties heen. Bijvoorbeeld, op SqueezeNet bereikte MiCo 68,22% nauwkeurigheid in 32 iteraties, terwijl HAQ 91 iteraties nodig had om hetzelfde niveau te bereiken.
• QAT Prestaties: In low-bitwidth (1/2/4/8-bit) QAT-zoektochten behield MiCo competitieve nauwkeurigheid, met name op kleinere modellen waar andere methoden moeite hadden met hoog-dimensionale zoekruimtes.
• Proxy Nauwkeurigheid: Het HAP-model met CBOPs en volledige features behaalde een aanzienlijk hogere voorspellingsnauwkeurigheid (bijv. een $R^2$ van 0,94 versus 0,50 voor BOPs op BitFusion) en een lagere Mean Absolute Percentage Error (MAPE) vergeleken met BOPs-gebaseerde baselines.
• Implementatie Versnelling: End-to-end implementatie toonde aan dat HAP-gestuurde exploratie superieure real-world versnellingen oplevert. Op BitFusion reduceerde HAP-gestuurde MPQ de latentie tot 60% van de baseline (versus 72% voor BOPs-gestuurde), met verwaarloosbaar nauwkeurigheidsverlies. Op RISC-V CPU's bereikte HAP-gestuurde exploratie tot 40% latentiereductie vergeleken met uniforme kwantisatie, terwijl BOPs-gestuurde exploratie aanzienlijk lagere versnellingen opleverde door niet-gemodelleerde hardware-overhead.

Significantie
Het artikel stelt dat MiCoPro de kritieke kloof tussen theoretische kwantisatie-efficiëntie en praktische hardware-implementatie aanpakt. Door verder te gaan dan eenvoudige BOPs metrieken naar een hardware-bewust proxy model en het bieden van een verenigde deployment stack, stelt het framework de snelle exploratie en directe implementatie van MPQ-modellen mogelijk die echte acceleratie op edge-apparaten realiseren. De resultaten demonstreren dat hardware-bewuste optimalisatie essentieel is om het volledige potentieel van mixed-precision kwantisatie te benutten, waarbij tot 40% latentiereductie wordt bereikt met minder dan 3% nauwkeurigheidsverlies over verschillende hardware-architecturen.