SandAI-org/MagiCompiler

A plug-and-play compiler that delivers free-lunch optimizations for both inference and training.

What it solves

MagiCompiler는 대규모 Transformer 유사 아키텍처의 성능을 제한하는 메모리 장벽과 연산자 오버헤드를 해결합니다. 단순한 로컬 연산자 최적화를 넘어, 학습과 멀티모달 추론 모두를 위한 시스템 수준 최적화를 제공합니다.

How it works

torch.compile 위에 구축된 MagiCompiler는 실행 라이프사이클의 전역 관리자로 작동합니다. 서브그래프 디스패치 관리, 데이터 흐름(오프로드 및 프리패치 포함) 오케스트레이션, 메모리 할당 제어를 수행합니다. 주요 기술 메커니즘은 다음과 같습니다.

  • Full-graph capture: 추론 시 Transformer 경계를 넘어 커널 융합 범위를 최대화합니다.
  • FSDP-aware layer-wise compilation: 학습 시 공격적인 크로스 연산자 융합을 가능하게 하면서 분산 파라미터 샤딩을 투명하게 유지합니다.
  • Selective offloading: 호스트‑to‑디바이스(H2D) 전송과 계산을 겹쳐 메모리 제한 환경에서 파이프라인 버블을 제거합니다.
  • Heuristic activation recomputation: 연산이 계산 집약적인지 메모리 집약적인지에 따라 저장할 연산과 재계산할 연산을 자동으로 판단해 피크 메모리를 감소시킵니다.

Who it’s for

대규모 Transformer 모델을 다루는 개발자와 연구자를 위해 설계되었으며, 특히 멀티모달 프레임워크를 구현하거나 FSDP와 같은 분산 전략으로 대형 모델을 학습하는 경우에 적합합니다.

Highlights

  • Plug-and-play integration: 간단한 데코레이터(@magi_compile@magi_register_custom_op)만으로 복잡한 모델 리팩토링 없이도 속도 향상을 얻을 수 있습니다.
  • Significant acceleration: NVIDIA H100 GPU에서 비디오 생성 모델에 대해 기존 SOTA 솔루션 대비 9%~26%의 성능 향상을 입증했습니다.
  • Introspection toolkit: magi_depyf를 포함하여 컴파일 타임라인, 바이트코드 흐름, 서브그래프 분할을 디버깅할 수 있습니다.
  • Hardware-aware: JIT 오프로드를 통해 고성능 GPU(H100)와 소비자용 GPU(RTX 5090)를 모두 지원합니다.