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

기술 요약: MiCoPro: HW-aware Proxy Model을 활용한 엔드 투 엔드 혼합 정밀도 HW/SW 공동 설계

문제 정의
양자화 신경망(QNN)은 자원이 제한된 엣지 디바이스에 AI를 배포하는 데 필수적입니다. 레이어별 혼합 정밀도 양자화(MPQ)는 정확도를 희생하지 않으면서 효율성을 극대화할 수 있는 경로를 제공하지만, 기존의 탐색 알고리즘은 상당한 한계에 직면해 있습니다. 전통적인 방법들은 MPQ 탐색 공간의 높은 차원성과 복잡한 상관관계로 인해 어려움을 겪으며, 종종 과도한 평가 예산을 요구하는 비효로적인 휴리스틱 탐색이나 머신 러닝 기술에 의존합니다. 또한, Bit Operations(BOPs)와 같은 표준 하드웨어 비용 메트릭은 연산자 패킹(operand packing), 메모리 대역폭, SIMD 활용도와 같은 아키텍처의 미묘한 차이를 무시하기 때문에 특정 하드웨어 타겟에서의 실제 레이턴시를 정확하게 예측하지 못합니다. 마지막으로, 하드웨어 특성에 기반하여 MPQ 스킴을 효율적으로 탐색하고, 결과 모델을 PyTorch에서 다양한 플랫폼을 위한 베어메탈(bare-metal) C 코드로 직접 배포할 수 있는 통합된 엔드 투 엔드 프레임워크가 부족한 실정입니다.

방법론
저자들은 MPQ 탐색을 위한 홀리스틱 프레임워크인 MiCo와, 하드웨어 인지 프록시(Hardware-Aware Proxy, HAP) 모델을 통합한 확장판인 MiCoPro를 제안합니다.

1. MPQ 탐색 알고리즘 (MiCo):

   • 정식화: 탐색은 예산 제약이 있는 블랙박스 최적화 문제로 취급됩니다.
   • 대리 모델링(Surrogate Modeling): 프레임워크는 양자화 구성에 따른 모델 정확도를 예측하기 위해 Random Forest(RF) 앙상블 학습을 사용합니다. RF는 제한된 학습 데이터로 이산적 탐색 공간에서의 비선형 상호작용을 모델링하는 능력이 뛰어나기 때문에 선택되었습니다.
   • 샘플링 전략:
     • 직교 초기 샘플링(Orthogonal Initial Sampling): 균등 랜덤 샘플링 대신, 극한의 비트 폭(예: W8A8, W1A1)과 특정 레이어는 양자화하고 다른 레이어는 풀 정밀도를 유지하는 직교 조합을 포함하는 초기 샘상 집합을 생성합니다. 이는 예측 모델이 탐색 공간의 경계와 상호작용을 효율적으로 학습하도록 보장합니다.
     • 근접 제약 샘플링(Near-Constraint Sampling, NCS): 유전 알고리즘이 하드웨어 제약(예: 레이턴시 또는 BOPs 제한)에 근접한 관심 영역(ROI) 내에서 후보 스킴을 생성합니다. ROI는 좁게 시작하여 적절한 솔루션을 찾지 못할 경우 확장되며, 제약 경계 근처에서 정확도를 극대화하는 스킴을 우선시합니다.
   • 평가 전략: 프레임워크는 이분화된 평가 접근 방식을 사용합니다. 비트 폭이 4비트 이상인 경우, 빠른 추정을 위해 사후 학습 양자화(PTQ)를 사용합니다. 4비트 미만의 구성의 경우, 정확도 회복을 위해 짧은 양자화 인식 학습(QAT) 재학습 단계를 수행하여 탐색 비용과 정밀도 사이의 균형을 맞춥니다.

2. 하드웨어 인지 프록시 (MiCoPro):

   • 커널 프로파일링: 시스템은 다양한 형태(shape)와 정밀도 조합에 대해 대상 하드웨어에서의 혼합 정밀도 커널(MatMul, Conv2D)에 대한 사이클 단위의 정확한 레이턴시 데이터를 수집합니다.
   • 특징 공학(Feature Engineering): BOPs의 한계를 극복하기 위해 저자들은 **복합 비트 연산(Composite Bit Operations, CBOPs)**을 도입했습니다. 이 특징들은 하드웨어 비용을 다음과 같이 분해합니다:
     • BMACs (Bit Multiply-Accumulate): $\max(b_w, b_a) \times MACs$. 혼합 폭 피연산자 확장(operand extension)의 병목 현상을 포착합니다.
     • ALoads/WLoads: 비대칭적인 대역폭 영향을 포착하기 위해 활성화(activation) 및 가중치(weight) 메모리 트래픽을 별도로 구분합니다.
   • 캘리브레이션 및 전이(Calibration & Transfer): 네트워크 의존적 최소-최대 캘리브레이션은 프록시 예측을 정밀도 극한에서의 실제 하드웨어 레이턴시와 정렬합니다. 또한, 전이 학습 메커니즘을 통해 사전 훈련된 모델의 잔차 오차(residual error)를 학습함으로써, 최소한의 새로운 프로파일링 데이터만으로 새로운 하드웨어 변체(예: 서로 다른 RISC-V 코어 구성)에 프록시를 적응시킬 수 있습니다.

3. 엔드 투 엔드 배포:

   • 프레임워크는 최적화된 PyTorch 모델을 PyTorch FX 그래프로 변환하는 배포 파이프라인을 포함합니다.
   • 특정 시뮬레이터(예: BitFusion)를 위한 그래프 변환을 지원하며, CPU를 위한 베어메탈 C 코드를 생성합니다.
   • baremetal-NN을 확장한 커스텀 혼합 정밀도 커널 라이브러리를 활용하여, 8/4/2/1-bit 연산에 대한 SIMD 가속 버전을 지원하며, 이를 통해 RISC-V 프로세서 및 가속기에서 직접 실행할 수 있도록 합니다.

주요 기여

• 새로운 탐색 알고리즘: 효율적으로 MPQ 스킴을 탐색하는 Random Forest 기반 근접 제약 샘플링(RF+NCS) 전략을 통해, 휴리스틱 및 강화 학습 기반 방법들보다 엄격한 제약 조건 하에서 더 낮은 정확도 저하를 달-성합니다.
• 하드웨어 인지 프록시 모델: CBOPs와 전이 학습을 사용하는 레이턴시 예측 방법은 전통적인 BOPs 메트릭보다 실제 하드웨어 레이턴시를 훨씬 더 잘 예측하여, 진정한 하드웨어 인지 최적화를 가능하게 합니다.
• 홀리스틱 프레임워크: MPQ 탐색과 배포를 연결하는 오픈 소스 엔드 투 엔드 파이프라인을 제공하며, BitFusion과 같은 전용 가속기 및 SIMD가 확장된 RISC-V 프로세서를 포함한 다양한 타겟을 지원합니다.

실험 결과
프레임워크는 다양한 모델(CNNs, ResNets, ViT, TinyLLaMa)과 하드웨어 타겟(BitFusion 가속기, VexiiMiCo RISC-V 변체)에 대해 평가되었습니다.

• 탐색 효율성: PTQ 탐색에서 MiCo는 다양한 모델과 제약 조건에 대해 최신 기술(HAWQ-V3, HAQ, BOMP, w-based)보다 일관되게 우수한 성능을 보였습니다. 예를 들어, SqueezeNet의 경우 MiCo는 32회의 반복(iteration) 만에 68.22%의 정확도에 도end 도달한 반면, HAQ는 동일한 수준에 도달하기 위해 91회의 반복이 필요했습니다.
• QAT 성능: 저비트 폭(1/2/4/8-bit) QAT 탐색에서 MiCo는 경쟁력 있는 정확도를 유지했으며, 특히 고차원 탐색 공간에서 다른 방법들이 어려움을 겪는 작은 모델에서 두드러졌습니다.
• 프록시 정확도: CBOPs와 전체 특징을 사용한 HAP 모델은 BOPs보다 훨씬 높은 예측 정확도(예: BitFusion에서 BOPs의 $R^2$ 0.50 대비 0.94)와 더 낮은 평균 절대 백분율 오차(MAPE)를 달성했습니다.
• 배포 속도 향상: 엔드 투 엔드 배포 결과, HAP 가이드 탐색이 우수한 실제 속도 향상을 제공함을 입증했습니다. BitFusion에서 HAP 가이드 MPQ는 (BOPs 가이드의 72% 대비) 무시할 만한 정확도 손실만으로 레이턴시를 베이스라인의 60%까지 줄였습니다. RISC-V CPU에서 HAP 가이드 탐색은 균등 양자화 대비 최대 40%의 레이턴시 감소를 달착한 반면, BOPs 가이드 탐색은 모델링되지 않은 하드웨어 오버헤드로 인해 훨씬 낮은 속도 향상을 보였습니다.

의의
본 논문은 MiCoPro가 이론적인 양자화 효율성과 실제 하드웨어 배포 사이의 중요한 간극을 해결한다고 주장합니다. 단순한 BOPs 메트릭을 넘어 하드웨어 인지 프록시 모델로 이동하고 통합된 배포 스택을 제공함으로써, 이 프레임워크는 다양한 하드웨어 아키텍처에서 진정한 가속을 달성하는 MPQ 모델의 신속한 탐색과 직접적인 배포를 가능하게 합니다. 결과는 하드웨어 인지 최적화가 엣지 디바이스에서 MPQ의 잠재력을 실현하기 위해 필수적임을 보여주며, 서로 다른 하드웨어 아키텍처 전반에서 3% 미만의 정확도 저하로 최대 40%의 레이턴시 감소를 달성했습니다.