SandAI-org/MagiCompiler

해결하는 문제

MagiCompiler는 대규모 Transformer 유사 아키텍처에서 발생하는 메모리 병목 현상과 연산자 오버헤드를 해결합니다. 단순한 로컬 연산자 최적화를 넘어서 데이터 흐름과 메모리의 시스템 수준 조율을 제공하여, 학습 및 다중 모달 추론 모두에서 '메모리 월' 효과를 줄입니다.

작동 방식

torch.compile 위에 구축된 MagiCompiler는 실행 라이프사이클의 글로벌 매니저 역할을 합니다. 추론 시에는 전체 그래프 캡처를 사용하여 커널 융합을 극대화하고, 학습 시에는 FSDP 인식 레이어 단위 컴파일을 통해 크로스 오퍼레이션 융합을 가능하게 하면서도 분산 파라미터 샤딩을 투명하게 유지합니다. 또한 H2D 전송을 계산과 겹치는 선택적 오프로딩 정책과, 연산이 계산 중심인지 메모리 중심인지에 따라 자동으로 어떤 연산을 저장할지 재계산할지 결정하는 히ュ리스틱 기반 활성화 재계산 전략을 구현합니다.

대상 사용자

대규모 Transformer 모델을 다루는 개발자 및 연구자, 특히 다중 모달 프레임워크를 구현하거나 분산 환경에서 대규모 모델을 학습하는 사람들에게 적합합니다.

주요 특징

• 통합 프레임워크: 추론(전체 그래프 캡처)과 학습(FSDP 인식 컴파일) 모두를 지원합니다.
• 플러그 앤 플레이: 기존 모델에 최소한의 코드 변경으로 통합할 수 있도록 @magi_compile 및 @magi_register_custom_op과 같은 간단한 데코레이터를 사용합니다.
• 스마트 오프로딩: 메모리 제약 환경에서 파이프라인 버블을 제거하기 위해 내장된 비동기 오프로딩을 제공합니다.
• 자동 재계산: 수동 체크포인팅을 대체하여 히ュ리스틱 기반 활성화 재계산을 통해 최대 메모리 사용량을 낮춥니다.
• 인스펙션 도구: magi_depyf를 사용하여 컴파일 타임라인과 디컴파일된 바이트코드 흐름을 디버깅할 수 있습니다.