A friendly introduction to machine learning compilers and optimizers 요약
최근에 내 분야와 정말 다르지만 ml comiler에 대한 관심이 정말 많아졌다. 가능하다면 다음 직업으로 삼아보고 싶은…
Chip Huyen의 A friendly introduction to machine learning compilers and optimizers를 읽고 정리해 보았다.
아마도 시기가 좀 된 글이여서 업데이트가 필요할 수 있겠지만 전체적으로는 여전히 좋은 글일거라고 생각한다.
1. 컴파일러의 중간자, IR(Intermediate Representation)이란?
컴파일러의 핵심에는 ‘중간자’ 역할을 하는 **IR(Intermediate Representation, 중간 표현)**이 있습니다.
이 IR은 소스 코드(고수준)의 연산 그래프에서 하드웨어가 이해할 수 있는 코드(저수준)로 번역될 때 거치는 형태입니다.
머신러닝 모델에선 주로 연산(옵레이션) 그래프가 고수준 IR로 사용됩니다.
2. 컴파일러의 주요 기능
컴파일러가 하는 일은 크게 두 가지로 요약할 수 있습니다.
- Lowering : 여러분의 모델을 특정 하드웨어에서 실행할 수 있도록 하드웨어 네이티브 코드로 변환(codegen)합니다.
- Optimizing : 특정 하드웨어에서 더 빠르고 효율적으로 동작하도록 모델을 최적화합니다.
3. 최적화의 두 가지 방법
- 로컬(Local) 최적화: 단일 연산자 또는 소수 연산자 집합 수준에서 최적화
- 글로벌(Global) 최적화: 전체 연산 그래프(모델 전체)에 대해 엔드-투-엔드로 최적화
대표적인 로컬 최적화 기법
벡터화(Vectorization)
반복문(loop)이나 중첩 루프를 하나씩 처리하는 대신, 메모리에 연속된 여러 요소를 한 번에 처리하도록 하드웨어의 벡터 연산 기능을 활용하는 방법입니다.병렬화(Parallelization)
입력 배열(또는 다차원 배열)을 여러 독립적인 청크로 나누고, 각 청크를 동시에 처리합니다.루프 타일링(Loop tiling)
루프에서 데이터 접근 순서를 변경해, 하드웨어의 메모리 레이아웃과 캐시 사용을 극대화합니다. 어떤 접근 패턴이 최적인지는 CPU와 GPU 등 하드웨어마다 다릅니다.연산자 퓨전(Operator fusion)
여러 연산자를 하나로 합쳐 쓸데없는 메모리 접근이나 반복을 줄입니다. 예를 들어 같은 배열에 대한 두 연산을 따로따로 순회하는 대신, 둘을 결합해 한 번에 처리합니다.
4. 글로벌 최적화: 구조적 융합
- 원본 그래프: 모델을 개발할 때의 논리적이고 명확한 연산 흐름 그래프
- 수직 융합 (Vertical Fusion): 순차적으로 이루어지는 연산을 하나로 묶어 중간 메모리 접근 횟수를 줄입니다. 파이프라이닝 효과를 기대할 수 있다.
- 수평 융합 (Horizontal Fusion): 병렬로 수행이 가능한 동일 연산을 묶어 전체 실행 및 데이터 처리 효율을 높입니다.
이처럼 연산자 수준을 넘어 그래프 전체 구조를 활용하는 최적화가 더 큰 성능 향상을 가져올 수 있습니다.
5. 수작업 규칙의 한계
컴파일러가 사용하는 핸드메이드(Hand-designed) 최적화 규칙에는 다음과 같은 한계가 있습니다.
- 최적(global optimal)이 아닐 수 있다.
- 변화(새로운 하드웨어나 모델 구조)에 유연하게 잘 적응하지 못한다.
6. 머신러닝 기반 컴파일러의 등장
이런 한계를 극복하기 위해, 최근에는 머신러닝 기반 컴파일러가 등장했습니다.
- 최적화 경로를 일일이 탐색하지 않고 효율적으로 탐색 공간을 줄임
- 각각의 경로가 얼마나 걸릴지 실행 시간을 예측하여, 효율적인 탐색 경로만 선택
하지만 전체 연산 그래프 전체에 대한 실행시간을 정확히 예측하기는 매우 어렵고, 현재 기술로는 부분 그래프, 즉 ‘서브그래프(subgraph)’ 수준에서 예측하는 것이 현실적입니다.
7. autoTVM: 실제 예시
autoTVM은 아래와 같은 방식으로 최적화를 진행합니다.
- 전체 연산 그래프를 여러 개의 서브그래프로 분해
- 각 서브그래프의 크기를 예측
- 각 서브그래프별로 최적 경로 탐색에 할당할 시간을 배분
- 각각의 서브그래프에 대해 최적 실행 경로를 찾아 전체 그래프 실행 계획을 만듦
autoTVM은 각 경로에서 실제로 실행에 걸린 시간을 측정하여, 이 데이터를 바탕으로 **코스트 모델(cost model)**을 학습시킵니다. 이 접근법의 장점은, 런타임에 수집된 데이터로 학습하기 때문에 어떤 하드웨어에서든 빠르게 적응할 수 있다는 점입니다. 단점은, 코스트 모델이 유의미하게 개선되기까지 시간이 조금 더 걸린다는 것입니다.
8. 오토 튜닝의 현실
자동 튜닝(오토튜닝) 결과는 매우 인상적이지만, 한 가지 함정이 있습니다.
예를 들어, TVM의 탐색 과정은 모든 가능한 경로를 시도하며 최적을 찾기 때문에, 복잡한 대형 모델에선 이 과정만 수 시간, 심지어 수 일이 걸릴 수 있습니다.
하지만,
- 이 작업은 한 번만 하면 되고
- 이전에 저장한 최적화 결과를 캐시하여 나중에도 사용할 수 있습니다.
- 한번 특정 하드웨어 백엔드에 최적화하면, 그 백엔드의 여러 기기에서 바로 쓸 수 있는 이점이 있습니다.
따라서 대규모 모델을 실제 배포(프로덕션)에 적용할 때, 그리고 목표 하드웨어가 분명할 때 가장 이상적이라는 평가를 받고 있습니다.