휴리스틱·솔버 접근법 혼용...온칩 메모리 할당 문제 해결
구글이 기계 학습(ML) 가속기의 오프칩(off-chip) 메모리 전송을 피하고 온칩(on-chip) 메모리의 사용 효율을 극대화하는 효과적인 메모리 할당 방식을 개발했다.
구글은 스마트폰 픽셀6용 ‘텐서 SoC(Tensor System-on-Chip)’ 추론 가속기와 데이터 센터용 ‘TPUv4(Tensor Processing Unit v4)’ ML 가속기의 온칩 SRAM에 대한 새로운 메모리 버퍼 할당 방식인 ‘텔라말록(TelaMalloc)’을 개발했다고 마크테크포스트가 28일(현지시간)가 전했다.
ML은 데이터 센터나 모바일 장치의 핵심 워크로드다. ML 워크로드는 기존 애플리케이션과 구조가 다르기 때문에 많은 시스템 문제를 재평가해야 한다.
그런 영역 중 하나가 메모리 할당이다. ML 가속기는 온칩 SRAM에 물리적으로 주소가 지정된 대형 스크래치패드 메모리를 포함하는 경우가 많다. 스크래치패드 메모리는 소프트웨어적으로 제어되는 메모리로, 캐시처럼 속도가 빠르고 주소 지정이 가능하며 데이터 관리를 위한 명시적인 메모리 관리 기법을 필요로 한다.
비용이 많이 드는 오프칩 메모리 전송을 피하기 위해 ML 프레임워크는 효과적인 메모리 할당 기법을 통해 온칩 SRAM의 사용을 최대화해야 한다.
C++과 같은 언어에서 메모리 할당은 조각화(fragmentation)를 최소화하면서 동적으로 할당된 버퍼 스트림에 대한 주소를 할당한다. 반면에 ML 워크로드는 버퍼의 할당 및 할당 해제 시간, 버퍼 크기 등의 매개변수를 미리 알고 있고 사용된 총 메모리 양이 장치에서 물리적으로 사용 가능한 총 메모리 양을 초과하지 않도록 메모리를 할당한다.
기존 ML 컴파일러 프레임워크는 이러한 온칩 메모리 할당 문제를 휴리스틱 기법이나 솔버(solver) 기반 접근법 중 하나로 처리한다.
ML 메모리 버퍼 할당 문제에 대한 가장 간단한 접근 방식은 버퍼의 종료 시간을 무시하고 시작 시간 순서대로 버퍼를 할당하는 것이다. 이 접근 방식은 메모리가 풍부한 경우에는 잘 작동하지만 메모리가 부족하면 실패한다. 이 문제를 해결하는 보다 효과적인 방법은 이전에 할당된 버퍼와 겹치지 않도록 하면서 한 번에 하나의 버퍼 위치를 선택하기 위해 종료 시간을 고려하는 휴리스틱을 사용하는 것이다. 이러한 접근 방식은 성능이 훨씬 뛰어나지만 여전히 복잡한 경우를 해결할 수는 없다.
솔버 기반 접근법은 버퍼의 시작 및 종료 시간이 고정돼 있으므로 시간 차원에서 어떤 버퍼가 겹치는지 결정할 수 있는 점을 이용한다.이를 위해 각 버퍼의 위치를 변수로 인코딩하고 두 버퍼가 겹치지 않도록 제약 조건을 추가한다. 솔버 기반 접근법은 제약 조건을 계산해 서로 겹치지 않는 버퍼를 할당하는 방식이다. 이 방법은 더 복잡한 경우를 효과적으로 처리할 수 있지만, 비용이 많이 들고 메모리 할당이 중요한 상황에는 적합하지 않다.
텔라말록은 휴리스틱과 솔버 기반 접근 방식을 절충하는 접근 방식을 사용한다. 할당 문제를 검색 공간으로 모델링한다. 버퍼를 할당하는 모든 단계에서 다음에 할당할 버퍼 하나를 휴리스틱에 의해 선택하고 위치를 선택한 다음 솔버 기반 접근법으로 제약조건을 만족하는지 확인한다.
이를 통해 휴리스틱이 실패하기 전에 많은 수의 추가 블록을 할당할 수 있는 결정을 할 수 있고 문제를 해결할 수 없게 된 시기를 조기에 식별할 수 있다.
예를 들어 첫번째 그림 'ML 워크로드에 대한 온칩 메모리 할당 방식의 예'에서 텔라말록이 블록 (8), (2), (7) 및 (4)를 할당했다고 가정했을 때, 오른쪽 그림의 올바른 결과와 유사하다. 솔버 기반 할당이 없으면 블록 (5)와 (6)을 더 이상 할당할 수 없다. 이 최적화를 통해 (5)와 (6)은 이미 배치된 블록(7) 아래에 올바르게 할당될 수 있다.
결론적으로 구글 고유의 메모리 할당 기법인 텔라말록은 휴리스틱 및 솔버 기반 접근법을 사용해 ML 가속기 메모리 할당 문제를 해결한다. 그들의 방법은 특정 주요 모델에서 최대 2배의 속도 향상이라는 큰 개선을 가져왔다.
텔라말록의 중요한 결과 중 하나는 즉석에서 모델을 컴파일할 수 있다는 것이다. 이렇게 하면 다른 방법으로는 불가능했던 모델도 지원할 수 있다.
박찬 위원 cpark@aitimes.com
무단전재 및 재배포 금지
기사제보 및 보도자료 news@aitimes.com