sophgo/tpu-mlir

Machine learning compiler based on MLIR for Sophgo TPU.

해결하는 문제

TPU-MLIR은 다양한 주류 프레임워크의 사전 학습된 신경망을 TPU 하드웨어에 배포하는 문제를 해결합니다. 고수준 모델 정의(PyTorch 또는 ONNX 등)와 TPU 칩에서 효율적인 실행에 필요한 특정 바이너리 형식(bmodel) 사이의 간극을 메웁니다.

작동 방식

이 프로젝트는 MLIR(Multi-Level Intermediate Representation) 기반의 컴파일 파이프라인을 사용하여 모델을 여러 단계로 변환합니다:

  1. 프론트엔드 임포트: PyTorch, ONNX, TFLite, Caffe 또는 HuggingFace의 모델을 가져와 통일된 MLIR 표현으로 변환합니다.
  2. Lowering: 컴파일러는 패턴 재작성 및 메모리 계획을 사용하여 고수준 표현을 TPU 전용 다이얼렉트(Top 및 Tpu)로 변환합니다.
  3. 양자화: 다양한 양자화 기술(F32, BF16, F16, INT8) 및 캘리브레이션을 적용하여 하드웨어 효율성을 위해 모델을 최적화합니다.
  4. 배포: 최종 출력물은 대상 TPU 프로세서에서 실행할 수 있는 bmodel 파일입니다.

대상 사용자

딥러닝 모델(대규모 언어 모델 및 비전 모델 포함)을 Sophgo TPU 하드웨어에 배포해야 하는 AI 엔지니어 및 개발자를 위해 설계되었습니다.

주요 특징

  • 폭넓은 프레임워크 지원: PyTorch, ONNX, TFLite, Caffe 및 HuggingFace와 호환됩니다.
  • LLM 최적화: HuggingFace LLM을 위한 전용 원샷 변환을 제공하며, 히스토리 KV 캐시 및 청크 추론을 지원합니다.
  • 포괄적인 양자화: 대칭/비대칭 INT8, F16, BF16을 지원하며 AWQ, GPTQ, AutoRound를 위한 패스스루를 지원합니다.
  • 프로덕션 툴링: 범용 추론 러너(model_runner), bmodel 검사 도구(model_tool), 그리고 정확도 검증 유틸리티가 포함되어 있습니다.