TpuGraphs: 대형 텐서 계산 그래프의 성능 예측 데이터셋

정확한 하드웨어 성능 모델은 코드 최적화에서 중요한 역할을 합니다. 이러한 모델들은 컴파일러가 휴리스틱 결정을 내리는 데 도움을 주거나, 자동 튜너가 주어진 프로그램에 대한 최적의 구성 설정을 식별하는 데 활용될 수 있습니다. 예를 들어, 머신 러닝 컴파일러인 XLA의 자동 튜너는 Google에서 대규모 생산 트래픽을 처리하는 최신 모델들의 성능을 10-20% 향상시키는 것을 발견했습니다. 프로그램 성능 예측을 위한 몇 가지 데이터셋이 존재하지만, 이들은 기본 블록이나 커널과 같은 작은 서브 프로그램들을 대상으로 하고 있습니다. 본 논문에서는 Tensor Processing Units (TPUs)에서 실행되는 전체 텐서 프로그램을 계산 그래프로 표현한 성능 예측 데이터셋인 TpuGraphs를 소개합니다. 데이터셋의 각 그래프는 머신 러닝 작업의 주요 계산, 예를 들어 학습 에포크나 추론 단계를 나타냅니다. 각 데이터 샘플은 계산 그래프, 컴파일 설정, 그리고 해당 설정으로 컴파일된 그래프의 실행 시간을 포함하고 있습니다. 데이터셋의 그래프들은 오픈 소스 머신 러닝 프로그램들로부터 수집되었으며, ResNet, EfficientNet, Mask R-CNN, Transformer와 같은 인기 있는 모델 아키텍처를 특징으로 합니다. TpuGraphs는 가장 큰 그래프 속성 예측 데이터셋(비교 가능한 그래프 크기)보다 25배 더 많은 그래프를 제공하며, 기존 머신 러닝 프로그램용 성능 예측 데이터셋에 비해 평균적으로 770배 더 큰 그래프를 포함하고 있습니다. 이와 같은 대형 그래프에서 이루어지는 그래프 수준의 예측 작업은 학습의 확장성, 훈련 효율성, 모델 품질 등 다양한 새로운 과제들을 제시합니다.