[논문 리뷰] Runtime Prediction of AI Model Operations Using a GRU-Based Neural Network

Mallapragada, Suma|arXiv (Cornell University)|2023. 08. 25.

Parallel Computing and Optimization Techniques인용 수 5

한 줄 요약

이 논문은 TpuGraphs를 소개한다. TpuGraphs는 TPU에서 전체 기계학습 계산 그래프의 런타임을 예측하기 위한 대규모 데이터셋으로, 레이아웃 및 타일 컴파일 구성에 따라 성능을 모델링하기 위해 GRU 기반 신경망을 사용한다. 이 데이터셋은 대규모 모델에서 그래프 수준의 성능 예측을 가능하게 하며, 실제 하드웨어 평가 대비 최대 2400배 빠른 추론 속도를 제공한다. 기준 모델은 ResNet, BERT 및 Transformers와 같은 복잡하고 다양한 워크로드에서 중간 정도의 정확도를 보인다.

ABSTRACT

Modern AI models can be represented as computational graphs, where eachnode corresponds to a tensor operation (e.g., matrix multiplication, convolution),and edges represent tensor data flows. Optimizing the executionof these graphs on hardware accelerators such as Tensor Processing Units(TPUs) requires careful selection of compiler configurations that control layoutand tiling strategies.The compilation configuration involves two key types of optimizations:• Layout Configuration: Controls how tensors are arranged in physicalmemory by specifying the dimension order for inputs and outputs ofeach operation node.• Tile Configuration: Controls the tile size of each fused subgraph,impacting data locality and parallelism. Accurately predicting the runtime of AI model graphs under various configurationscan automate and improve the selection of optimal compiler settings,reducing execution time and resource consumption. The Kaggle competition dataset “Google - Fast or Slow? Predict AIModel Runtime” provides runtime data for XLA High Level Optimizer (HLO)graphs running on TPU v3 hardware. This dataset, called TPUGraphs,comprises multiple collections with diverse layouts and tiling configurations,posing a challenging performance prediction task. This work proposes a GRU-based runtime prediction pipeline leveragingopcode runtime features, graph structure dependencies, and configurablenode embeddings. The method consolidates node features, integrates configurationconvolutions, and trains a neural network to predict runtime withmean squared error loss. The approach captures both static graph propertiesand dynamic configuration effects, enabling enhanced runtime estimation toguide compiler heuristics.

연구 동기 및 목표

• TPU와 같은 가속기에서 실행되는 전체 기계학습 프로그램을 대상으로 한 대규모 그래프 수준 성능 예측 데이터셋의 부족을 해결하기 위해.
• 레이아웃 및 타일 컴파일 구성과 같은 다양한 설정에서 계산 그래프의 실행 시간을 효율적이고 정확하게 예측하기 위해.
• 비용이 많이 드는 실제 하드웨어 측정을 대체하여 컴파일러 자동튜닝을 가속화할 수 있는 기계학습 모델의 개발을 지원하기 위해.
• ResNet, BERT 및 Transformers와 같은 다양한 생산 규모의 모델을 포함한 총 3100만 개의 레이아웃 구성과 1300만 개의 타일 구성이 포함된 벤치마크 데이터셋을 제공하기 위해.
• 고도로 재현되는 구조와 메모리 제약 조건이 존재하는 대규모이고 다양한 그래프에서 스케일링되고 일반화 가능한 모델을 훈련시키는 데 직면한 과제를 탐구하기 위해.

제안 방법

• TpuGraphs 데이터셋은 오픈소스 기계학습 프로그램의 계산 그래프로 구성되며, 각 그래프는 TPU v3에서의 실행 시간과 함께 레이아웃 및 타일 컴파일 구성과 함께 제공된다.
• 그래프 구조, 구성, 실행 시간 간의 관계를 모델링하기 위해 GRU 기반 신경망을 제안한다. 이는 그래프 연산의 순차적 의존성을 포착한다.
• 모델은 그래프 노드와 간선을 시퀀스로 처리하며, GRU 레이어를 사용해 구조적 및 운영적 특징을 인코딩하여 성능 예측을 수행한다.
• 특징 공학은 텐서 레이아웃 순서와 타일 크기를 입력 임베딩으로 인코딩하여, 모델이 구성별 성능 패턴을 학습할 수 있도록 한다.
• 훈련 파이프라인은 큰 그래프를 배치 처리하고 메모리 사용량을 최적화하여 처리하며, 동일한 그래프의 여러 구성에 대해 특징 추출을 분산 처리한다.
• 기준 모델은 그래프 신경망(GNN)과 손실 함수로 MSE 및 ListMLE를 사용한 지도 학습을 통해 구현되었으며, 순위 기반 최적화를 위해 사용된다.

Figure 1: Scale of TpuGraphs compared to other graph property prediction datasets.

실험 결과

연구 질문

• RQ1 그래프 및 구성 특성만을 사용하여 대규모 계산 그래프의 런타임을 정확하게 예측할 수 있는 학습된 성능 모델이 존재하는가?
• RQ2 ResNet, BERT 및 Transformers와 같은 다양한 실제 워크로드에서 모델 성능는 어떻게 변화하는가?
• RQ3 최대 7,700개의 노드와 높은 구성 재현성으로 인한 고려 대상 그래프를 다룰 때, 훈련 효율성, 모델 확장성 및 예측 정확도 사이의 상충 관계는 어떠한가?
• RQ4 모델 기반 예측은 실제 하드웨어 평가 대비 얼마나 더 빠른가? 이 속도 향상은 자동튜닝 파이프라인에서 활용될 수 있는가?
• RQ5 생산 환경에서의 기계학습 워크로드에서, 하나의 모델이 예상치 못한 그래프 구조와 구성에 얼마나 잘 일반화되는가?

주요 결과

• TpuGraphs 데이터셋은 총 3,100만 개의 레이아웃 구성과 1,300만 개의 타일 구성으로 이루어져 있으며, 평균적으로 7,700개의 노드를 가진 그래프를 포함한다. 이는 가장 큰 비교 가능한 그래프 특성 데이터셋보다 25배 많고, 기존의 기계학습 성능 데이터셋보다 770배 이상 크다.
• 모델 예측은 실제 하드웨어 평가 대비 최대 2,400배 빠르며, 구성당 추론 시간은 0.004초에서 0.1초 사이로, 컴파일 및 실행에 45~475초가 소요되는 것과 대비된다.
• GRU 및 GNN 아키텍처를 사용한 기준 모델은 레이아웃 컬렉션에서 각각 ListMLE 기준 10.5%의 중앙값 상위 1 오차와 MSE 기준 19.0%의 오차를 기록했으며, BERT 및 ResNet 같은 특정 모델에서 뚜렷한 향상이 있었다.
• 가장 뛰어난 성능을 보인 모델(LISTMLE + EarlySAGE)은 레이아웃 컬렉션에서 10.5%의 상위 1 오차를 기록했고, 타일 컬렉션에서도 동일한 10.5%의 오차를 기록하여 다양한 워크로드에 걸쳐 강력한 일반화 능력을 입증했다.
• 모델의 특징 추출 시간은 동일한 그래프의 여러 구성에 걸쳐 분산 처리되어, 여러 구성이 평가되는 자동튜닝 파이프라인에서 매우 효율적이다.
• 이 연구는 학습된 비용 모델이 자동튜닝을 크게 가속화할 수 있음을 확인했으며, 후보 구성 평가에 소요되는 시간을 수시간에서 수밀리초로 단축시킬 수 있다.

Figure 2: Important optimizations in ML compilers include graph-level and kernel-level optimizations. A graph-level optimization requires the context of the entire graph to make optimal decisions and transforms the entire graph accordingly. A kernel-level optimization transforms each kernel (a fused