Optimal Kernel Orchestration for Tensor Programs with Korch

논문 : https://arxiv.org/abs/2406.09465

아래글은 Scholar GPT와 대화를 나눈 내용입니다.

Q : 논문의 핵심 내용과 강점, 알고리즘 설명, 그리고 한계점

Answer

업로드된 논문은 DNN(딥 뉴럴 네트워크)의 성능을 개선하기 위해 커널 오케스트레이션을 최적화하는 텐서 프로그램 최적화 도구인 Korch에 대해 다루고 있습니다. 아래는 논문의 강점, 독창성, 핵심 알고리즘 예시와 한계점에 대한 설명입니다.

강점 및 독창성

  1. 운영자 분해와 세밀한 최적화: Korch는 “운영자 분해(operator fission)”라는 새로운 접근 방식을 통해 높은 수준의 텐서 연산자를 더 작고 관리하기 쉬운 기본 단위(프리미티브)로 나눕니다. 이를 통해 기존의 단순 합치기(fusion) 기법보다 다양한 최적화를 가능하게 하여, 다양한 하드웨어에 걸쳐 성능을 미세하게 조정할 수 있습니다.

  2. 이진 선형 프로그래밍(BLP)을 활용한 최적화: Korch는 제한된 이진 선형 프로그래밍(BLP) 모델을 사용하여 최적의 오케스트레이션 전략을 찾아냅니다. 기존의 규칙 기반 프레임워크가 탐욕적(fusion) 규칙에 제한되는 반면, Korch는 체계적인 방법으로 더 넓은 최적화 공간을 탐색할 수 있습니다.

  3. 다양한 플랫폼에서의 성능 향상: Korch는 V100 및 A100 GPU에서 TVM이나 TensorRT와 같은 프레임워크와 비교했을 때 다양한 DNN 작업에 대해 성능을 크게 향상시켜, 여러 하드웨어 구성에서 적응성을 보여줍니다.

핵심 알고리즘 예제와 함께 설명

Korch의 커널 오케스트레이션 최적화 과정은 다음과 같이 설명할 수 있습니다.

  1. 입력 표현: 입력 DNN은 DAG(유향 비순환 그래프)로 표현되며, 각 노드는 텐서 연산(예: 행렬 곱셈)을, 각 간선은 데이터 종속성을 나타냅니다.

  2. 운영자 분해: “소프트맥스” 연산을 예로 들면, 이 연산은 여러 단계(지수화, 합산, 나눗셈)가 필요한데, Korch는 이를 개별 프리미티브로 분해합니다.
    • 지수화: 각 요소에 대해 ( e^{x_i} )를 계산합니다.
    • 합산: 벡터의 지수화 결과를 합산합니다.
    • 브로드캐스트 나눗셈: 모든 요소를 합산값으로 나누어 확률을 계산합니다.
  3. 커널 식별: Korch는 깊이 우선 탐색(DFS)을 사용하여 이러한 프리미티브를 실행할 수 있는 모든 커널 후보를 찾고, 각 커널의 성능을 측정하여 최적의 조합을 파악합니다.

  4. BLP 최적화: Korch는 BLP 솔버를 통해 프리미티브와 제약 조건을 고려하여 전체 실행 시간을 최소화하는 최적의 실행 전략을 도출합니다.

  5. 실행 파일 생성: Korch는 최적화된 커널 오케스트레이션을 통해 실행 가능한 GPU 파일을 생성하며, 최적의 전략을 적용하여 성능을 극대화합니다.

예제

소프트맥스 연산 후 행렬 곱셈이 포함된 DNN 블록을 고려할 때, Korch는 소프트맥스를 세 가지 프리미티브(지수화, 합산, 나눗셈)로 나누고, 이를 주변 층의 연산과 결합하여 실행 성능을 높입니다. 이처럼 세밀한 오케스트레이션은 커널 실행 오버헤드를 줄이고 메모리 사용을 최적화할 수 있습니다.

한계점

  1. 하드웨어별 최적화 라이브러리 의존성: Korch는 다양한 하드웨어에 적응할 수 있지만, TVM이나 cuDNN과 같은 라이브러리에 의존하여 커널 성능을 측정하기 때문에 특정 작업에서 최적 성능을 내지 못할 수 있습니다.

  2. 튜닝 시간과 계산 오버헤드: BLP 최적화 및 커널 프로파일링 과정은 대형 DNN 모델에 대해 많은 시간이 소요될 수 있습니다. 특히 복잡한 네트워크의 경우, 튜닝 시간이 몇 시간에 달해 신속한 모델 배포에 어려움이 있을 수 있습니다.

  3. 커널 실행 전략의 제한적 유연성: Korch는 아직 CUDA 멀티스트리밍과 같은 고급 커널 병렬화 최적화를 지원하지 않아, 병렬 실행 효율성을 더욱 높이는 데는 한계가 있습니다.

Korch는 DNN의 GPU 실행을 최적화하기 위해 연산을 세분화하고 고급 오케스트레이션 기법을 사용하여 성능을 크게 향상시킬 수 있지만, 일부 복잡하거나 하드웨어에 특정된 상황에서는 한계가 있을 수 있습니다【5†source】.

Q : 튜닝 시간과 계산 오버헤드에 대해서 예시를 추가해서 자세하게 설명해 줄 수 있겠니? 그리고 커널 식별 단계에서 커널의 성능은 어떻게 측정하는지도 자세하게 설명해줘

A:

Korch의 튜닝 시간과 계산 오버헤드, 그리고 커널 식별 단계에서의 성능 측정 방식에 대해 더 구체적으로 설명하겠습니다.

튜닝 시간과 계산 오버헤드

Korch의 튜닝 단계에서는 최적의 커널 오케스트레이션을 찾기 위해 각 커널의 실행 시간을 측정하고, 다양한 후보 커널 조합을 탐색하는 작업이 포함됩니다. 이 과정에서 많은 시간이 소요될 수 있는데, 특히 복잡한 DNN 모델일수록 더 많은 시간이 필요합니다.

예시:

예를 들어, YOLO나 EfficientViT와 같은 모델을 튜닝할 때 Korch는 수백 개에서 수천 개의 프리미티브 노드를 생성하고, 이로부터 여러 후보 커널을 만듭니다. 각 커널의 성능을 측정하기 위해서는 실제로 커널을 실행해 보거나 특정 하드웨어에서의 실행 시간을 예측해야 합니다.

  • EfficientViT 모델을 튜닝한다고 가정할 때, 이 모델에는 약 380개의 프리미티브 노드가 있으며 약 2174개의 후보 커널이 생성될 수 있습니다. 각 커널은 TVM의 메타스케줄러를 통해 성능을 프로파일링하며, 이 과정에서 메모리 집약적 커널의 경우 약 2분이 소요될 수 있습니다.
  • 만약 2174개의 후보 커널이 있고 그중 절반이 메모리 집약적이라면, 전체 프로파일링에는 약 18시간 이상이 걸릴 수 있습니다.

이러한 튜닝 과정은 한 번 완료되면 재사용할 수 있으나, 초기 튜닝 시간은 매우 큰 오버헤드로 작용할 수 있습니다. 특히, 다양한 입력 크기나 배치 크기를 지원해야 하는 경우, 각 구성마다 별도의 튜닝이 필요하므로 실험이 더 오래 걸릴 수 있습니다.

커널 식별 단계에서의 성능 측정 방식

Korch의 커널 식별 단계에서는 특정 프리미티브 조합을 실행할 커널 후보를 찾고, 각 커널의 성능을 측정하여 최적의 오케스트레이션 전략을 선택합니다. 이를 위해 Korch는 각 후보 커널의 실행 시간을 평가하는데, 이 과정은 다음과 같이 이루어집니다.

  1. 프리미티브 그래프 탐색 및 후보 생성: Korch는 깊이 우선 탐색(DFS)을 사용해 각 프리미티브의 조합으로 구성된 여러 커널 후보를 식별합니다. 각 커널 후보는 특정 프리미티브 세트를 하나의 커널로 묶은 형태로 나타나며, 이 중 어떤 조합이 성능이 좋은지 알기 위해 성능을 측정하게 됩니다.

  2. 커널 프로파일링: 커널 성능을 측정하기 위해 Korch는 메모리 집약적 커널과 계산 집약적 커널을 구분합니다.
    • 메모리 집약적 커널: TVM의 메타스케줄러(TVM MetaScheduler)를 사용하여 성능을 측정합니다. 이 메타스케줄러는 여러 파라미터를 조정해 최적의 성능을 찾기 위해 하이퍼파라미터 튜닝을 수행합니다.
    • 계산 집약적 커널: 이러한 커널은 주로 복잡한 행렬 연산을 포함하며, Korch는 cuDNN이나 cuBLAS와 같은 고도로 최적화된 벤더 라이브러리를 활용하여 성능을 측정합니다. 이러한 라이브러리는 특정 하드웨어(GPU)에서 최적화된 성능을 제공하기 때문에, 자동 튜닝 없이도 고성능을 보장할 수 있습니다.
  3. 실제 실행 시간 측정: 프로파일링 과정에서 각 후보 커널을 GPU에서 직접 실행하여 실행 시간을 기록합니다. 이 실행 시간은 커널 오케스트레이션 최적화를 위한 BLP 문제의 입력으로 사용되며, 가장 빠른 커널 조합을 찾기 위해 활용됩니다.

  4. 실행 불가능한 커널 필터링: 특정 프리미티브 조합이 불가능하거나 실행 시간이 매우 긴 커널은 프로파일링 단계에서 필터링됩니다. 예를 들어, 두 개 이상의 계산 집약적 프리미티브가 하나의 커널로 실행되기 어려운 경우, 이러한 조합은 평가하지 않도록 합니다.

이러한 성능 측정 과정을 통해 Korch는 각 커널의 최적 실행 시간을 확보하고, 이를 기반으로 최종적으로 가장 효율적인 오케스트레이션 전략을 선택합니다.