이전 포스트에서는 KV Cache에 대해 더 깊게 알아보는 내용을 다뤘었습니다. 주로 Hugging Face Transformers의 구현체와 vLLM 프로젝트의 구현체를 탐구해 보았죠. 더 자세한 내용은 아래 링크에서 확인하실 수 있습니다:
LLM 인퍼런스 훑어보기 (3) - KV Cache (deep dive)
이전 포스트에서는 트랜스포머의 구조와 어텐션 메커니즘이 무엇인지 다루었다. 그리고 LLM 인퍼런스에서 KV Cache가 무엇이며, 왜 도입되어 사용되는지에 대해서 다뤘다. 지난 포스트는 아래 링
dytis.tistory.com
오늘은 LLM 인퍼런스 가속화 방법 중 하나로 알려진 "커널 퓨전 (kernel fusion)"에 대해 알아보려고 합니다.
Kernel fusion
GPU는 주로 다양한 수학적 연산, 예를 들면 MatMul, Softmax, BatchNorm 등을 수행합니다. 이러한 연산들을 하나의 커널로 통합(fusion)함으로써 메모리 액세스 시간과 커널 실행 시간을 상각(amortize)하여 최적화할 수 있습니다.
GPU에서 연산이 실행되면, 필요한 데이터를 메모리로 불러와 연산을 수행한 후, 결과를 다시 메모리에 저장합니다. kernel fusion은 기본적으로 연속된 독립적인 계산 작업들을 단일 하드웨어 작업으로 통합합니다. 결과적으로 여러 독립적인 계산을 하나로 통합하여 메모리 이동을 최소화할 수 있습니다.
이 기법은 학습과 평가 단계 모두에 적용 가능하므로, 딥러닝 모델 사용 시 유용한 최적화 전략 중 하나로 간주됩니다.
아래 그림은 FlashAttention에서 제공하는 fused Kernel의 성능 향상을 시각적으로 보여줍니다. 이 kernel은 MatMul, Mask, SoftMax, Dropout, 그리고 MatMul 연산을 single kernel로 fusion하여, 더 뛰어난 계산 성능을 눈에 띄게 나타냅니다.
Kernel fusion 방법
kernel fusion을 하는 방법으로는 직접 CUDA C++ kernel 코드를 작성하는 방법과 PyTorch JIT, TensorRT 등과 같은 툴킷을 사용하여 연산을 자동으로 인식하여 kernel fusion하는 방법이 있습니다.
PyTorch JIT
PyTorch JIT는 kernel을 자동으로 융합할 수 있습니다. 그러나 아직 컴파일러에 구현되지 않은 추가적인 fusion 가능성이 있을 수 있고, 모든 장치 유형이 동일하게 지원되는 것은 아닙니다.
가장 간단한 경우에는 torch.jit.script 데코레이터를 함수 정의에 적용하여 융합을 활성화할 수 있습니다. 예를 들면 다음과 같습니다.
@torch.jit.script
def fused_gelu(x):
return x * 0.5 * (1.0 + torch.erf(x / 1.41421))
이러한 기능은 TorchScript를 통해 제공됩니다. TorchScript는 PyTorch 코드를 직렬화하고 최적화 가능한 모델로 만드는 방법을 제공합니다. 모든 TorchScript 프로그램은 Python 프로세스에서 저장하고, Python 종속성 없이 다른 프로세스에서 로드할 수 있습니다.
PyTorch JIT 및 TorchScript에 대해 더 자세히 알아보시려면 아래 링크를 참조하시면 좋을 것 같습니다.
Performance Tuning Guide — PyTorch Tutorials 2.1.0+cu121 documentation
Note Click here to download the full example code Performance Tuning Guide Author: Szymon Migacz Performance Tuning Guide is a set of optimizations and best practices which can accelerate training and inference of deep learning models in PyTorch. Presented
pytorch.org
TorchScript — PyTorch 2.1 documentation
TorchScript TorchScript is a way to create serializable and optimizable models from PyTorch code. Any TorchScript program can be saved from a Python process and loaded in a process where there is no Python dependency. We provide tools to incrementally tran
pytorch.org
Introduction to TorchScript — PyTorch Tutorials 2.1.0+cu121 documentation
Note Click here to download the full example code Introduction to TorchScript Authors: James Reed (jamesreed@fb.com), Michael Suo (suo@fb.com), rev2 This tutorial is an introduction to TorchScript, an intermediate representation of a PyTorch model (subclas
pytorch.org
TensorRT
TensorRT는 네트워크 계산 그래프를 구문 분석하고 그래프 최적화의 기회를 찾습니다. 이러한 그래프 최적화는 그래프의 기본 계산을 변경하지 않고, 대신 작업을 훨씬 빠르고 효율적으로 수행하기 위해 그래프를 재구성합니다. 그래프 재구성 단계에서, kernel fusion을 사용하여 연산 성능도 최적화합니다.
위의 그림은 GoogLeNet 아키텍처의 Inception 모듈을 TensorRT에서 최적화한 결과를 시각적으로 보여줍니다. 왼쪽 네트워크에서는 오른쪽과 같이 다양한 크기의 컨볼루션, 바이어스, 그리고 ReLU 레이어를 CBR이라고 불리는 단일 커널로 효과적으로 결합할 수 있습니다.
TensorRT kernel fusion에 대해서 더 자세히 알아보시려면 아래 링크를 참조하시면 좋을 것 같습니다.
TensorRT 3: Faster TensorFlow Inference and Volta Support | NVIDIA Technical Blog
NVIDIA TensorRT™ is a high-performance deep learning inference optimizer and runtime that delivers low latency, high-throughput inference for deep learning applications. NVIDIA released TensorRT last…
developer.nvidia.com
NVIDIA TensorRT – Inference 최적화 및 가속화를 위한 NVIDIA의 Toolkit
TensorRT는 학습된 Deep Learning 모델을 최적화하여 NVIDIA GPU 상에서의 Inference 속도를 수배 ~ 수십배 까지 향상시켜 Deep Learning 서비스 TCO (Total Cost of Ownership) 를 개선하는데 도움을 줄 수 있는 모델 최
developer.nvidia.com
Kernel Fusion의 예
kernel fusion은 딥러닝 모델 사용 시 가장 유용한 최적화 전략 중 하나로 간주되며, 이러한 이유로 다양한 딥러닝 프레임워크에서 이에 대한 연구와 구현을 진행하고 있습니다.
Deepspeed Inference (deep fusion)
DeepSpeed Inference는 여러 연산자를 단일 커널로 융합하여 커널 호출 수와 커널 전체의 기본 메모리 액세스 대기 시간을 줄일 수 있습니다. kernel fusion은 PyTorch JIT 컴파일러, Tensorflow XLA 등에서 사용되는 일반적인 기술이지만 DeepSpeed의 deep fusion은 다릅니다. 요소별 연산을 주로 융합하는 기존 fusion 기술과 달리 DeepSpeed의 deep fusion은 element-wise operation, matmul, transposition, reduction을 모두 단일 커널로 융합하여 커널 호출 횟수와 메인 메모리 액세스를 크게 줄여줍니다.
FlashAttention
FlashAttention은 GPU 고대역폭 메모리(HBM)와 GPU 온칩 SRAM 사이의 메모리 읽기/쓰기 횟수를 줄이기 위해 Tiling을 사용하는 "IO-aware" 알고리즘입니다. Tiling을 통해, 한번의 HBM load 이후에 MatMul, Mask, Softmax, Dropout, Matmul이 융합된 kernel을 실행함으로써 한번에 HBM에 저장할 수 있습니다. 특히, pytorch의 torch.nn.functional.scaled_dot_product_attention (SDPA) 함수는 FlashAttention을 기반으로 구현되었습니다. FlashAttention은 이후에 별도의 포스트로 다루어 볼 예정입니다.
이번 포스트에서는 커널 퓨전의 의미, 구현 방법, 그리고 활용 사례에 대해 알아보았습니다. 커널 퓨전은 딥러닝 프로젝트에서 널리 활용되는 가장 유용한 최적화 방법 중 하나이므로, 이 내용을 숙지하면 LLM 인퍼런스(추론) 시에 큰 도움이 될 것입니다.
다음 포스트에서는 다양한 배치 방법에 대해 다뤄보겠습니다.
'AI > 모델 인퍼런스' 카테고리의 다른 글
| LLM 인퍼런스 훑어보기 (6) - quantization (0) | 2023.11.17 |
|---|---|
| LLM 인퍼런스 훑어보기 (5) - continuous batching (0) | 2023.11.10 |
| LLM 인퍼런스 훑어보기 (3) - KV Cache (deep dive) (0) | 2023.09.23 |
| LLM 인퍼런스 훑어보기 (2) - KV Cache (1) | 2023.09.19 |
| LLM 인퍼런스 훑어보기 (1) - LLM을 이용한 문장 생성 (0) | 2023.09.14 |