Product
이번 게시물은 Optimium을 실행할 때 활용되는 runtime에 대한 글입니다. Runtime이 하는 일은 굉장히 다양한데요. 오늘은 그 중에서 memory를 효율적으로 할당하는 memory planning에 대해 소개드리고자 합니다.
Jinhwan Shin
June 17, 2024
안녕하세요? ENERZAi에서 runtime을 개발하고 있는 신진환이라고 합니다. 그 동안의 글에서는 고성능 추론 엔진 Optimium과 그 기반이 되는 프로그래밍 언어 Nadya를 설명드렸는데요. 이번 주에는 Optimium을 실행할 때 활용되는 runtime에 대한 글입니다. Runtime이 하는 일은 굉장히 다양한데요. 오늘은 그 중에서 memory를 효율적으로 할당하는 memory planning에 대해 소개드리고자 합니다.
Memory는 무한하지 않다
과거에 비해 요즘에는 메모리 사용량을 크게 신경 쓰지 않습니다. 스마트폰만 하더라도 초창기에는 128메가바이트, 512메가바이트 하던 램 크기가 8기가바이트 이상 하는 경우는 심심치 않게 볼 수 있습니다. 몇 년 전 노트북도 2기가바이트가 기본 이였지만 요즘에는 16기가바이트나 32기가바이트까지 제공하죠. 이런 환경에서 대부분의 소프트웨어는 메모리를 다 잡아먹거나 메모리가 부족해서 멈추는 현상은 거의 발생하지 않습니다.
하지만 AI에서는 이야기가 다릅니다. 최근 유행하는 LLM들을 보면 메모리를 최소 수 기가 바이트 에서 수 백 기가 바이트까지 사용합니다. 큰 메모리를 보유하고 있는 서버 급 컴퓨터가 아니라면 대부분 모델을 실행하기 어려운 크기이죠. 이 모델을 serving 하기 위해 여러 모델을 동시에 실행한다면 더욱 어려운 문제가 되죠. 이러한 문제는 edge에서 더욱 커집니다. IoT 장치 같은 edge 디바이스에서는 스마트폰 보다 더 작은 메모리를 보유할 가능성이 높고, 단순히 모델을 실행하는 것 뿐만 아니라 입력으로 사용할 데이터를 수집하거나, 사용자와의 상호작용을 하기 위해 또 다른 작업이 동시에 수행될 가능성이 높기 때문에 추론을 위해 많은 메모리를 사용 할 수 없습니다.
이처럼 예전에 비해 메모리의 크기가 많이 증가 했음에도 불구하고 AI 추론 엔진 입장에서는 메모리는 여전히 유한한 자원입니다. 따라서 메모리를 효율적으로 할당함으로써 필요한 최대 사용량을 줄이면서도 모델을 추론할 수 있는 방법인 memory planning이 필요합니다.
Memory planning이란?
조그마한 서랍장이 있다고 생각해봅시다. 여기에 그동안 작성했던 모든 서류들을 쌓아놓겠습니다. 이것이 한 달, 두 달이 지나 1년이 되면 점점 서랍장이 가득 차서 더 이상 새로운 서류들을 넣지 못할 때가 올 것입니다. 이 때 우리에게 주어진 option은 2가지입니다. 서랍장을 하나 더 사거나, 또는 기존의 서랍장을 비우는 것인데요. 기존에 있던 서류들에는 앞으로도 사용할 수 있는 것들도 있겠지만, 평생이 지나도 보지 않을 것들이 많이 있을 것입니다. 우리는 필요 없는 서류들을 버림으로써 하나의 서랍장으로 계속 서류들을 보관할 수 있습니다.
이렇듯 memory는 궁극적으로 현재의 그리고 앞으로의 어떤 연산을 위해 변수에 대한 정보를 보관해야 합니다. 그런데, memory는 썼다 지우기를 반복할 수 있기 때문에 더 이상 사용되지 않는 특정 변수들을 제거하고 해당 위치에 앞으로 필요하게 될 다른 변수들을 저장할 수 있습니다. 이렇게 어떤 정보를 저장하고 어떤 정보를 지울 것인지 설계하여 메모리 사용량을 효율적으로 절감하는 것이 바로 memory planning입니다. 이러한 memory planning은 어디서나 활용되지만 저희 Optimium이 AI에서 이를 어떻게 활용하는 지 이제부터 소개드리겠습니다.
Memory planning — Tensor 자르고 붙이기
먼저 모델 안 Layer들의 실행 순서를 결정해야 합니다. 어떤 순서로 Layer들이 실행되는 지 알아야 어떤 Tensor를 먼저 배치할지 알 수 있기 때문이죠. 예를 들어 다음과 같은 간단한 모델이 있다고 가정 해봅시다.
위 그래프에서 1 → 2 → 4 → 3 → 5 순서로 Layer가 실행된다고 가정합니다.
위 정보를 통해 Layer가 어떤 Tensor에 접근하는지 tracking합니다. 예를 들어 1번 Layer를 수행하기 위해서는 Tensor A, B에 접근해야 하고, 2번 Layer를 수행하기 위해서는 Tensor B, C에 접근해야 합니다.
위 그림과 같이 Layer가 어떤 Tensor에 접근하고, Tensor들은 몇 번 참조되는 지 분석이 완료되면 이 정보를 기반으로 Planning을 시작됩니다.
Memory planning을 진행할 때 현재 활성 상태의 Tensor와 그 Memory 범위를 추적하는 Allocation Table과 실제 Tensor가 어느 범위의 Memory에 할당할 것인지를 기록하는 Memory Plan Table을 두어 진행합니다.
설명의 명료성을 위해 모든 Tensor의 크기는 100으로 고정했습니다.
예를 들어 4번 Layer를 수행하기 위해서는 Tensor C, E에 접근해야 합니다. Tensor C는 2번 Layer에 대해 Memory Planning을 수행할 시점에 이미 할당 했으므로 그대로 이용하지만 Tensor E는 아직 할당되기 전입니다.
이 때 Allocation Table의 Tensor B는 더 이상 참조하고 있는 Layer가 없으므로 이 memory를 안전하게 재활용 할 수 있습니다. 따라서 Tensor E를 Tensor B에 대치 하여 그 영역을 재사용 합니다.
위와 같은 과정으로 Layer 실행 순서와 Tensor 접근 순서를 고려해 Memory Planning을 수행하게 됩니다.
아래 가상의 모델을 보면 원래 메모리를 600만큼 사용했지만 Memory Planning 후 메모리 사용량이 50%가 감소한 300만큼 사용하는 것을 볼 수 있습니다.
Memory 다이어트 성공!
Optimium에서는 보다 효율적인 메모리 사용을 위해 위와 같은 Memory Planning을 도입하여 제공중에 있습니다. 아래 그래프처럼 Memory Planning을 활성화 및 비활성화 시 메모리 사용량을 비교해보면 최소 76%에서 최대 92%만큼의 메모리 사용량이 감소한 것을 볼 수 있습니다. (모델별로 상이할 수 있음)
또한 다양한 memory optimization을 통해 메모리 사용량 뿐 만 아니라 성능까지 같이 도모하고 있습니다. AMD64 Platform에서는 최대 1.55 배 가속, ARM64 Platform에서는 최대 1.34배 가속을 볼 수 있습니다.
마무리
이번 게시물에서는 Optimium의 runtime이 어떻게 메모리 사용량을 줄이고 실행 속도를 가속화하는 지 간략하게 알아보았습니다. Runtime이 하는 일은 사실 이것 말고도 매우 다양합니다. 현재 Optimium 팀은 사용자 편의를 위한 remote API, 가속화를 위한 parallel execution 등 다양한 기능을 개발하기 위해 달려나가고 있습니다. 앞으로도 Optimium에 적용된 다양한 최적화 기법에 대한 게시물들을 업로드할 예정이니, 많은 관심 부탁 드립니다!