Sigrid Jin (Jin Hyung Park) | ML Backend Engineer
개론
최근 대규모 Transformer 모델인 GPT-3나 Codex가 널리 활용되면서, 한 번의 훈련이나 추론에 기하급수적으로 늘어난 연산을 처리하기 위해서는 GPU 자원을 최대한 효율적으로 활용해야 합니다. 이때 하드웨어 아키텍처적 세부 사항을 간과하면 막대한 비용과 전력이 낭비될 수 있습니다.
이번에 리뷰할 Nvidia의 The Case for Co-Designing Model Architectures with Hardware 논문의 핵심 주장은 모델 파라미터 수가 동일하더라도, 하이퍼파라미터를 GPU가 원하는 차원으로 맞추면 훈련과 추론 속도를 크게 높일 수 있다는 점입니다. 특히 Transformer 구조는 대부분의 연산이 GEMM(General Matrix Multiplication)에 의존하므로, GPU 내부의 병렬 구조와 텐서 코어가 선호하는 타일 크기에 잘 맞춰 모델을 설계해야 합니다.
GPU는 여러 SM(Streaming Multiprocessor)에서 수천 개의 스레드를 동시에 실행하며, 행렬 곱을 일정 타일 크기로 쪼개 처리합니다. 이때 차원이 텐서 코어가 선호하는 배수에 정렬되지 않으면 padding으로 인한 타일 단위 낭비와 wave quantization 문제 때문에 성능이 떨어집니다. wave quantization은 스레드 블록이 SM 수의 배수가 아닐 때 남은 소수의 블록이 새 웨이브를 시작해야 하는 상황을 의미하는데, 이로 인해 전체 GPU가 빈 블록을 처리하는 데도 똑같은 시간을 소모하게 됩니다. 타일 크기가 행렬 차원을 정확히 나누지 않는다면 tile quantization 문제가 발생해, 마지막 타일이 부분적으로만 유효 연산을 수행하더라도 전체 타일 크기만큼 자원이 소모됩니다. 결국 행렬 차원과 스레드 블록 수가 SM 개수의 배수가 되도록 잘 정렬하면 이러한 비효율을 크게 줄일 수 있습니다.
Transformer에서 핵심인 Multi Head Attention 연산은 특히 h(hidden dimension)과 a(head 수)의 비율, 즉 h/a가 텐서 코어가 선호하는 배수(예를 들어 FP16 기준 64 배수)가 되도록 설계하면 성능이 현저히 좋아집니다. 예컨대 GPT-3 2.7B 모델에서 hidden size를 2560, head 수를 32로 두면 h/a가 80으로 64 배수에 맞지 않아 비효율이 생깁니다.
하지만 hidden dimension을 더 키우거나 head 수를 다른 값으로 조정해 h/a가 64 배수에 가까워지면 GPU 내부 병렬성을 보다 온전히 활용할 수 있습니다. 이 밖에도 vocab size 역시 64 배수에 맞춰 패딩하면 로짓 레이어의 (v, h) 형태 행렬 곱에서 텐서 코어 활용도가 높아집니다. batch size 또한 메모리가 허용하는 한도를 고려해 최대한 크게 잡으면 memory-bound를 넘어서 compute-bound 상태로 전환되어 처리량을 끌어올릴 수 있습니다.
Tile Quantization vs. Wave Quantization
Tile Quantization은 GPU가 행렬 곱을 처리할 때 발생하는 현상입니다. GPU는 행렬을 일정 크기의 타일로 나누어 각 스레드 블록이 처리하는데, 행렬 차원이 타일 크기로 나누어떨어지지 않으면 마지막 타일이 일부만 사용되어도 스레드 블록은 전체 크기로 동작하게 됩니다. 이는 타일 단위와 실제 행렬 크기가 맞지 않아 생기는 비효율을 의미합니다.
Wave Quantization은 SM(Streaming Multiprocessor) 단위의 스레드 블록 배정과 관련된 현상입니다. 전체 스레드 블록 수가 SM 개수의 배수가 아닐 경우, 첫 번째 웨이브에서 최대한 많은 블록을 동시에 처리하고 남은 블록은 두 번째 웨이브로 처리하게 됩니다. 이때 두 번째 웨이브도 거의 같은 지연시간을 소모하므로, 남은 작업량이 적더라도 전체 처리 시간이 길어지는 비효율이 발생합니다.
발생 메커니즘을 보면, Tile Quantization의 경우 GPU 커널은 보통 특정 타일 크기로 행렬을 나누어 병렬화하는데, 행렬의 크기가 타일 크기의 배수가 아니면 마지막 타일에 빈 칸이 포함되어도 스레드 블록은 전부 계산해야 합니다. 이 부분 타일도 독립된 스레드 블록으로 할당되어, 실제 유효 연산이 없는 공간에도 같은 비용이 듭니다.
Wave Quantization의 경우, GPU에는 여러 SM이 있어 한 번에 SM 수만큼 블록을 동시에 처리할 수 있습니다. 스레드 블록 수가 SM 개수의 배수가 아니면, 첫 웨이브에서 대부분을 처리하고 남은 소수의 블록이 두 번째 웨이브로 넘어가게 됩니다. 두 번째 웨이브도 동일한 지연시간을 요구하므로 효율이 급격히 떨어질 수 있습니다.
Tile Quantization은 일부 타일이 부분적으로만 유효 연산을 수행해도 GPU가 풀 타일로 인식해 계산하며, 이러한 불일치가 클수록 스레드 블록 단위의 낭비가 누적되어 처리량이 감소합니다. Wave Quantization은 한 웨이브 안에서는 타일이 동시에 실행되지만, 남은 블록이 소수여도 새로운 웨이브를 시작해야 하며, 이때 SM 자원의 재할당과 스케줄링이 필요해 전체 GPU가 거의 빈 블록만 처리하면서도 비슷한 시간을 소모하게 됩니다.
예를 들어, 타일 크기가 128×256일 때 행렬 크기가 1024×1024면 딱 맞아떨어져 빈 타일이 없지만, 1025×1024라면 마지막 행열에 빈 공간이 생겨 한 타일이 거의 무의미한 연산만 하게 됩니다. 웨이브의 경우, A100 GPU에 108개의 SM이 있다면, 스레드 블록 총 수가 108의 배수일 때는 한 번의 웨이브로 끝나지만, 109개라면 첫 웨이브에 108개, 두 번째 웨이브에 1개만 남아 비효율이 발생합니다.
이러한 문제들을 완화하기 위해서는 모델이나 행렬 차원을 128 배수 등으로 맞추거나 부분 타일 전용 처리 커널을 사용할 수 있으며, 행렬 크기나 타일 크기를 조정하거나 동적 타일 크기를 사용하여 wave quantization 영향을 줄일 수 있습니다.
실험
실험에서는 AWS p4d 노드(A100 GPU), ORNL Summit(V100 GPU), SDSC Expanse(V100 GPU) 등 다양한 HPC 시스템에서 PyTorch를 이용해 GEMM 성능을 측정했습니다. 행렬 크기가 작을 때는 메모리로부터 데이터를 로드하고 store하는 비용이 상대적으로 커져 memory-bound 상태가 되지만, 행렬 크기가 충분히 커지면 GPU의 실질적인 연산 능력에 의해 처리량이 좌우됩니다.
따라서 (행렬 크기 × 배치 크기)가 충분히 커질수록 A100은 V100 대비 2배 이상의 처리량을 낼 수 있습니다. 더 나아가 특정 차원이 64 배수를 만족하지 않을 때는 throughput이 파동(peaks and valleys)처럼 변동해, wave quantization과 tile quantization 영향이 확인되었습니다. 결국 행렬 크기와 batch size, h와 a를 조정해 텐서 코어의 128바이트 정렬을 만족하면 높은 처리 효율을 달성할 수 있습니다.
이러한 원리는 Transformer 전체 레이어에도 그대로 적용됩니다. 트랜스포머 레이어는 Q, K, V 변환부터 Attention Score 계산, Attention over Value, MLP 확장(4배) 등 여러 행렬 곱 연산을 순차적으로 수행하므로, 각 부분에서 발생하는 타일, wave quantization 문제와 텐서 코어 정렬 상태가 누적되어 최종 처리량을 좌우합니다. 논문에서는 예시로, GPT-NeoX나 Pythia, OPT 등 다양한 Transformer 아키텍처를 GPU 위에서 벤치마크한 결과를 제시했습니다.
Pythia-1B 모델이 Pythia-410M 모델보다 실제 추론 시간이 짧은 이유가 “더 큰 hidden dimension과 적절한 head 수”로 인해 행렬 곱 효율이 좋아졌기 때문이라는 점도 주목할 만합니다. 다시 말해서 FlashAttention 같은 최신 고성능 커널을 사용할 때도 유효합니다. FlashAttention v2는 어텐션 연산을 여러 커널로 나누지 않고 한 커널로 통합해 중간 데이터를 재활용하며, h/a가 64 배수가 아니더라도 내부 최적화로 비교적 빠른 속도를 보여줍니다. 그럼에도 하드웨어가 선호하는 배수 조건을 만족할 때 가장 안정적인 속도를 얻을 수 있다는 사실은 여전히 유효하다는 거죠.
MLP 블록에서 4배 확장이 아니라 8/3처럼 다른 배수를 쓰는 SwiGLU 기법도 자주 거론되는데, 이 경우 텐서 코어 정렬 조건을 예전과 똑같이 맞추기가 까다로워집니다. 예를 들어 LLaMA-2-7B 모델은 2.6875라는 배율을 썼고, 70B 모델은 3.5를 쓰면서 파라미터를 조정했습니다. 이렇게 배수를 미세하게 조정해도 정확도나 모델 크기를 만족하려면 하드웨어 정렬(64 배수)을 고려해야 성능 저하를 줄일 수 있습니다. 또, 데이터센터 노드 구성이 8-GPU가 아니라 6-GPU로 설계된 Summit 같은 슈퍼컴퓨터를 사용할 때는 텐서 병렬 방식이 변형되어, 기존 8-GPU 병렬 설계를 그대로 적용하기 어려울 수도 있습니다. 이처럼 하드웨어 환경에 따라 효율적인 설정이 달라질 수 있기 때문에, 모델 개발자는 훈련과 추론에서 모두 최적의 하이퍼파라미터를 찾는 과정을 면밀히 거쳐야 합니다.
논문에서는 Transformer를 구성하는 GEMM 연산들을 처음부터 끝까지 나열해, GPU 내부에서 어떻게 행렬을 나누고 배치하며 스레드를 할당해 연산을 수행하는지를 상세히 설명하였습니다. 또한 h/a가 2의 거듭제곱 배수를 얼마나 포함하는지, vocab 차원이 64 배수인지, batch 크기가 충분히 크므로 compute-bound 상태에 도달하는지 등을 종합적으로 살피면 최대 20% 이상의 훈련 시간 단축이 가능하다는 실험 결과를 제시합니다. 실제로, GEMM 단위 벤치마크와 Transformer 전 레이어 벤치마크 결과가 거의 일치한다는 점도 흥미롭습니다. BERT나 GPT 시리즈를 비롯해 Transformer 기반 모델이 하드웨어 벤치마크의 표준으로 자리 잡아가는 배경이 바로 여기에 있습니다.
추가로, PyTorch에서 3차원 텐서를 2차원처럼 처리할 때 처리량 차이가 없었다는 사실도 공개했습니다. 예컨대 (2048, 4, n) 형태든 (4, 2048, n) 형태든 성능에 큰 차이가 없었기 때문에, Transformer 내부에서 배치 차원을 어떻게 배치할지에 대해서는 그다지 민감하지 않다는 결론도 얻을 수 있습니다. 그보다는 wave quantization과 tile quantization 등으로부터 자유로워지도록 하이퍼파라미터를 GPU의 타일 크기와 SM 개수의 배수에 가깝게 조정하는 편이 훨씬 중요합니다. 마지막 Output 레이어(로짓 레이어) 역시 (b⋅s, v) × (v, h)의 행렬 곱으로 표현되므로, vocab size v가 64 배수인지 아닌지에 따라 처리량이 크게 달라집니다. 따라서 실제 vocab이 64 배수가 아니면 패딩을 통한 정렬을 권장합니다.
결론
결국 하드웨어와 모델 아키텍처를 별개로 다루는 기존 관행을 지적하고, Transformer 모델을 GPU 친화적으로 설계한다는 교훈을 줍니다. 트랜스포머가 사실상 여러 GEMM 연산으로 구성되어 있으니, 행렬 곱 벤치마크만 잘 살펴봐도 GPU의 실효 성능을 손쉽게 파악할 수 있고, 각 GPU(V100, A100, H100 등) 세대별 텐서 코어 특성을 고려해 하이퍼파라미터를 조정하면 작은 수치 변화만으로도 처리 속도가 크게 개선된다는 점을 실험으로 보여줍니다. 나아가 FlashAttention, parallel layers, rotary나 ALiBi 같은 위치 임베딩, SwiGLU 등 최근 자주 도입되는 최적화 기법도 근본적으로는 GPU 내부 구조와 어떻게 맞물리는지 면밀히 살펴야 한다고 역설합니다.
앞으로 하드웨어 성능이 더 강력해지고 Transformer 모델이 더욱 커지면, 이런 GPU 특징을 고려한 모델의 디자인이 필수적인 연구 방향이 될 것입니다. 하드웨어가 새로운 메모리 계층이나 더 효율적인 타일링 방식을 도입할 때, 모델 역시 이에 최적화해 설계되어야 대규모 모델을 빠르고 저렴하게 훈련·추론할 수 있기 때문입니다.