FLUX Fast Software-based Communication Overlap On GPUs Through Kernel Fusion

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

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

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

A :

물론입니다. 논문의 주요 내용과 강점, 알고리즘 과정, 한계점에 대해 한국어로 설명드리겠습니다.

강점 및 독창성

  1. 향상된 통신 중첩: Flux는 GPU에서 종속 계산과 통신 작업을 매우 세밀한 단위로 분해하고 결합하여 최대 96%까지 통신 지연을 숨깁니다. 이 방식은 대규모 언어 모델의 분산 학습 및 추론에서 효율성을 크게 높여줍니다.
  2. GPU에 최적화된 성능: 기존 통신 중첩 기법은 GPU에서 성능이 저하되는 경우가 많지만, Flux는 GPU의 아키텍처를 최대한 활용하도록 설계되었습니다. 이를 통해 커널 분해로 인해 발생하는 성능 저하를 최소화하고, 기존 기법 대비 최대 1.24배 학습 속도와 1.66배 추론 속도를 구현했습니다.
  3. 모듈화 및 유연성: Flux는 NVIDIA의 CUTLASS 라이브러리를 사용하여 구현되었으며 다양한 GPU 아키텍처와 인터커넥트에 적용할 수 있어 하드웨어 구성에 맞게 조정 가능한 장점이 있습니다.

핵심 알고리즘과 예시를 통한 과정 설명

Flux의 핵심은 텐서 병렬화에서 일반 행렬 곱셈(GEMM) 연산을 ReduceScatterAllGather와 같은 통신 작업과 결합하는 것입니다. 주요 프로세스는 다음과 같습니다.

  1. ReduceScatter 중첩 (알고리즘 1):
    • 입력: 행렬 (A), (B), 출력 포인터 (C)들.
    • 과정:
      • 스레드 블록 ID와 랭크에 따라 타일 좌표를 계산합니다.
      • GEMM 연산(예: (C = A \times B))을 단계별로 수행합니다.
      • 에필로그(마지막 단계)에서 결과를 출력하거나 로컬에서 Reduction 작업을 수행합니다.
    • 출력: 결합된 커널이 각 GPU에서 병렬로 데이터를 처리하고 통신하여 효율성을 높입니다.
  2. 입력 예시:
    • 예를 들어, ( A )와 ( B ) 행렬이 있는 ( 2 \times 2 ) 크기의 행렬 곱셈 설정에서 4개의 GPU가 있다고 가정합니다.
    • 각 GPU는 자신의 데이터에 대해 GEMM 연산을 수행하며, 에필로그 단계에서 ReduceScatter를 사용해 부분 결과를 교환합니다.
    • 이렇게 통신을 중첩하여 GPU들이 계산 작업으로 바쁘게 유지되면서 유휴 시간을 줄입니다.

한계점

  1. GPU 종속적 기능 의존: Flux는 CUTLASS와 CUDA 같은 GPU 특정 기능에 크게 의존하기 때문에 다양한 하드웨어 플랫폼에서의 적용이 제한적입니다. 또한 GPU 모델과 인터커넥트 종류에 따라 성능이 달라집니다.
  2. 작은 행렬에서의 성능 제한: Flux는 큰 행렬에서 성능이 뛰어나지만, 작은 행렬(예: 배치 크기 < 1024)에서는 세밀한 타일링과 GPU 실행 제약으로 인해 성능 향상이 크지 않습니다.
  3. 높은 튜닝 복잡도: Flux는 하드웨어 구성을 최적화하기 위해 자동 튜닝 메커니즘을 사용해야 하며, 이는 설정 시간을 증가시키고 특히 새로운 GPU나 혼합 GPU 클러스터에서 추가적인 복잡성을 초래할 수 있습니다.

Flux는 GPU 최적화 텐서 병렬화에서 상당한 개선을 제공하지만, 더 다양한 하드웨어와 문제 크기에 적용할 수 있도록 추가 개선이 필요할 수 있습니다【5†source】.

Q : 이 논문에서 제시하는 독창성은 어떤거라고 생각해? 내생각에는 NCCL에서 제공하는 communitve 연산과 compute kernel을 특정상황에서 묶은거라고 생각이 드는데 어떻게 생각해?

A :

맞습니다. Flux의 독창성은 말씀하신 것처럼 NCCL의 커뮤니케이션 연산과 컴퓨팅 커널을 특정 상황에서 결합한 점에 있습니다. 특히 Flux는 기존 방법들이 데이터 전송과 계산을 개별 커널로 나눠서 순차적으로 실행하는 방식에서 벗어나, 세밀하게 분해된 타일 단위로 통신과 계산을 결합하여 하나의 큰 커널로 처리하는 방식을 제안했습니다. 이로 인해, 커뮤니케이션을 계산과 동시에 수행하여 GPU의 유휴 시간을 줄이고 전체적인 계산 효율성을 높일 수 있었습니다.

Flux의 접근은 단순히 통신과 계산을 함께 실행하는 것을 넘어서, 세부적으로 최적화된 타일링 기법을 적용하고, GPU 아키텍처에 맞게 조정함으로써 GPU의 SMs(Streaming Multiprocessors)를 더 잘 활용합니다. 이렇게 커널 수준에서 세밀하게 병합하는 방식은 NVIDIA의 NCCL에서 제공하는 기본 통신 연산을 GPU의 계산 특성에 맞춰 더 효율적으로 활용한 것입니다. 따라서 Flux의 독창성은 통신과 계산의 미세 분해 및 결합을 통해 GPU에서 최적의 성능을 이끌어내는 방식에 있다고 할 수 있습니다.

결론적으로, Flux는 NCCL의 기본 커뮤니케이션 연산을 단순히 사용하는 것이 아니라, GPU 아키텍처와 통신 요구사항에 맞춰 통신과 계산의 세부적 통합과 동시 실행을 가능하게 하는 방식으로 독창성을 발휘하고 있습니다.

Q: 세밀하게 분해된 타일 단위라는게 어떤건지 궁금한데 어떠한 방식으로 compute kernel을 tiling하고있니? 그리고 GPU 아키텍처와 통신 요구사항에 맞춰 통신과 계산의 세부적 통합과 동시 실행을 가능하게 하는 방식에 대해서도 자세하게 설명해줘

A : Flux에서 말하는 세밀하게 분해된 타일 단위(tiled granularity)는 GEMM(General Matrix Multiply) 연산을 수행할 때 행렬을 타일이라는 작은 블록 단위로 나누고, 이 타일 단위에서 통신과 계산을 결합하여 커널을 병합하는 것을 의미합니다. Flux는 기존 방식보다 훨씬 더 세밀하게 타일을 분해함으로써 통신과 계산의 병렬화를 극대화했습니다.

1. Compute Kernel의 Tiling 방식

Flux는 GEMM 연산을 수행할 때 타일링 기법을 활용하여 연산을 최적화합니다. 구체적으로, Flux는 GPU의 스레드 블록과 워프 수준의 타일링을 적용하여 작은 블록 단위에서 통신과 계산을 결합합니다. 이 방식은 다음과 같은 단계로 이루어집니다:

  • 타일 분할: 행렬을 GPU 아키텍처에 맞는 작은 블록(타일)으로 나누어, 각 타일이 GPU의 스레드 블록에 할당되도록 합니다. 예를 들어, 4개의 GPU에서 병렬로 GEMM을 수행한다면, Flux는 각 GPU에서 타일을 담당하여 부분 행렬 곱을 계산할 수 있도록 합니다.

  • 타일 매핑: 타일의 좌표는 GPU의 스레드 블록 인덱스와 로컬 랭크(예: GPU ID)를 기반으로 결정됩니다. 이를 통해, GPU의 병렬 연산 자원을 최대한 활용하며 각 스레드 블록이 특정 타일을 처리하도록 합니다.

  • 타일 내 통신과 계산의 결합: 타일 단위에서 각 GPU는 필요한 데이터를 계산하고, 이 데이터는 ReduceScatterAllGather와 같은 통신 연산을 통해 타 GPU로 분배됩니다. 이때 타일 간 데이터 전송이 필요할 경우, 각 타일에서 직접 통신이 이루어지며, Flux는 이를 커널의 에필로그(후처리 단계)에 결합하여 통신이 계산과 겹쳐 수행되도록 합니다.

2. 통신과 계산의 세부적 통합 및 동시 실행

Flux의 진정한 독창성은 타일링된 통신과 계산의 동시 실행을 GPU 아키텍처와 통신 요구사항에 맞게 조정한 데 있습니다. 이를 가능하게 하는 요소들은 다음과 같습니다.

  • 타일 좌표 스위즐링(Tile Coordinate Swizzling): Flux는 메모리 접근 충돌을 줄이고, 통신과 계산의 병렬성을 최대로 유지하기 위해 타일 좌표를 스위즐링합니다. 이는 타일의 메모리 접근 순서를 바꾸어 메모리 컨텐션(memory contention)을 최소화하는 방법입니다. 예를 들어, GPU의 메모리 컨트롤러가 특정 시점에 동일한 메모리 위치에 접근하는 것을 피하기 위해 타일의 위치를 미리 조정하여 통신 병목을 줄입니다.

  • AllGather와 ReduceScatter의 융합: Flux는 AllGather와 ReduceScatter 연산을 각각 GEMM 커널의 프롤로그(전처리 단계)에필로그에 결합합니다. 예를 들어, AllGather는 GEMM 연산의 시작 부분에서 필요한 데이터를 가져오도록 설계되어 있으며, WaitSignal 함수를 통해 각 타일이 통신 신호를 기다리며 GEMM 연산을 차단합니다. 반면, ReduceScatter는 계산이 완료된 후 데이터가 분산되는 에필로그에 결합되어 통신과 계산이 서로 중첩되도록 합니다.

  • 타일 크기와 통신 순서 조정: Flux는 GPU 아키텍처에 따라 타일 크기를 조정하고 통신 순서를 최적화하여 각 타일의 통신이 네트워크 대역폭을 효율적으로 사용하도록 합니다. 예를 들어, PCIe와 NVLink 인터커넥트가 다른 클러스터 환경에서 각기 다른 최적의 타일 크기와 통신 순서를 사용하도록 설계되어 있습니다. 이를 통해 GPU의 성능을 최대로 끌어올리고 통신 지연을 최소화합니다.

이와 같은 방식으로 Flux는 GPU 아키텍처와 네트워크의 특성에 맞춘 세부 조정이 가능하여, 최대한의 통신 중첩과 높은 효율성을 확보하게 됩니다.