[이화여대 강의] 2. Deep Learning Software Stack 1 (3)

이전 글 1 : https://mari970.tistory.com/83

이전 글 2 : https://mari970.tistory.com/84

[이화여대 강의] 2. Deep Learning Software Stack 1 (2)

 

CUDA

c 의 extension. low level 코드이다.

GPU 하드웨어 아키텍처와 깊게 연관되어 GPU 프로그래밍을 구현되어있다.

 

Building PyTorch Custom CUDA Kernel

cuda kernel 을 직접 사용자가 설계하여 Pytorch 에 연결할 수 있다.

 

a = torch.randn(3,5)
b = torch.randn(3,1)

b_sized_a = b.expand_as(a)
a += b_sized_a

 

torch 가 broadcasting sum 이 안된다고 가정해보면,

dimension 이 다른 행렬에 대해 복사해서 더할 수 있도록 처리해야 한다. (expand_as())

 

__global__ void broadcast_sum_kernel(float *a, float *b, int x, int y, int size)
{
	int i = (blockIdx.x + blockIdx.y * gridDim.x) * blockIdx.x + threadIdx.x;
    if (i >= size) return;
    int j = i % x; i = i/x;
    int k = i % y;
    a[IDX2D(j, k, y)] += b[k];
 }

 

그러면 위와 같이 custom cuda kernel 을 만들 수 있다. (참조 : src/mathutil_cuda_kernel.cu)

** gpu 는 여러 thread 가 병렬적으로 동작할 수 있도록 구현되어있다. 각 thread 마다 id 가 존재한다.

thread 마다 행렬의 row 를 쪼개서 각각 더할 수 있도록 구현된다.

 

위처럼 만든 함수를 커널을 호출하는 wrapper 함수를 짜주고, nvcc 를 통해(.cu 파일을 컴파일해서) build 한다.

 

→ 더 알아보고싶다면 토치에 custom cuda extension 사용하는 방법을 찾을 수 있다.

 

CUDA semantics

• tensorfloat-32 (TF32) : Ampere 디바이스(이상?) 에서 사용할 수 있는 precision
  • FP32 는 IEEE754 의 호환 가능한 precision 기준이다.
  • 위의 기준을 따르지 않고 nvidia 에서 개발한 precision 으로, exponent 와 mentissa bit 가 다르다.
  • 이를 사용할 때 머신러닝에서 accuracy 가 좀 더 좋다고 한다.
• Reduced Precision Reduction in FP16 GEMMs
• Reduced PRecision Reduction in BF16 GEMMs
• CUDA streams

 

TensorFloat 32 (TF32)

matmul 이나 conv 등의 연산을 더 빨리 할 수 있도록 디자인됨.

 

# TF32 로 matmul 을 할 수 있도록
torch.backends.cuda.matmul.allow_tf32 = True

# TF32 를 cuDNN 에서 사용할 수 있도록
torch.backends.cudnn.allow_tf32 = True

 

10240x10240 을 행렬곱을 할 때 7배 빠른것을 확인할 수 있다.

 

Reduced precision reduction in FP16 GEMMs

 

Reduced precision reduction 이란? : 매트릭스 행렬곱에서 A 의 행과 B 의 열 (여러 값) 을 하나로 줄이는 연산이기 때문에 reduction 연산이라고 한다.

 

즉 precision 을 reduce (낮춘) 행렬곱 연산 이라는 뜻이다.

 

A (mxk) 와 B (kxn) 행렬을 계산 때 : k dimension 이 커질수록 reduced precision 을 켯을 때 속도 개선이 빨라진다.

→ 우리가 matmul 할 때 기본적으로 켜져있고, full precision 을 하려면 아래처럼 False 할 수 있다.

 

torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False

 

2024-1 인공지능융합기반시스템개론 ❘ 이화여자대학교 심재형 교수님 ppt

 

CUDA streams

cuda stream 이란 GPU 에서 코드에 기술된 순서대로 동작하는 연산 sequence 를 의미한다.

 

cuda 를 사용하면 default stream 을 하나 사용하게 되고, 새로운 stream 이 생기면 이 디폴트 스트림과 병렬적으로 연산이 진행된다.

 

cuda = torch.device('cuda')

s = torch.cuda.Stream() # 새로운 stream 생성
A = torch.empty((100,100), device=cuda).normal_(0.0, 1.0)

with torch.cuda.stream(s):
	B = torch.sum(A)

 

위 코드에서는 새로운 s stream 을 생성하고 이 stream 에서 A 행렬을 더하는 동작이 일어나도록 한다.

하지만 위에서는 default stream 에서 A 행렬을 normalization 하고 있어서 다 끝나기 전에 새로운 stream 에서 작업하게되어 에러가 난다.

 

그래서 밑의 코드처럼

 

cuda = torch.device('cuda')

s = torch.cuda.Stream() # 새로운 stream 생성
A = torch.empty((100,100), device=cuda).normal_(0.0, 1.0)

s.wait_stream(torch.cuda.default_stream(cuda)) # 추가해준 코드

with torch.cuda.stream(s):
	B = torch.sum(A)
    
A.record_stream(s)

 

디폴트 스트림이 끝날 때 까지 기다렸다가 다른 스트림을 만들어서 sum 한다.

: 이를 synchronization 이라고 한다.

 

마지막에 해주는 record_stream 은 아래 문제를 해결하기 위해 사용한다.

처음에 디폴트 스트림에서 A 를 가지고있었고

새로운 스트림에서 B 가 A 를 이용하기 때문에 A 가 디폴트 스트림에서 de-allocation 될 수 있다.

 

record_stream 은 tensor 가 stream 에서 사용되고 있음을 표시한다.

그래서 tensor 가 deallocate 되었을 때 stream queue 에 있는 모든 work 가 끝날 때 까지 tensor memory 가 다른 텐서가 재사용하지 않도록 해준다. 

 

---

NCCL

NVIDIA Collective Communications Library 라고 한다. 

 

이 라이브러리에는 GPU 간 communication preimitives (기본 연산) 들을 제공한다. : AllReduce, Broadcast, AllGather, ReduceScatter 등..

NIVIDA gpu 에 최적화 되어있고, torch 에 bulit-in 되어있어 사용자가 직접 쓸 일은 없다고 한다.

 

torch.distributed.broadcast 와 같은 함수로 직접 내가 다른 지피유로 전송을 해줄 수 있기도 하다.

 

---

OpenAI Triton

openai triton 은 전문가가 아니어도 GPU 코드를 비교적 쉽게 쓸 수 있는 파이썬 언어이다.

즉, cuda 프로그래밍에서 파이썬을 사용할 수 있도록 하는 라이브러리로, 이는 c 파일로 래핑하고 바인딩하고 컴파일링할 필요가 없다는 뜻이다.

GPU 에서 병렬동작을 할 수 있도록 할 수 있다.

 

예를 들면 아래와 같은 새로운 activation 함수를 구현했을 때

 

def Act_new(x):
	x = torch.add(x,1)
	x = torch.sqrt(x)
	return x

 

torch.add 함수만으로도 많은 cuda 코드를 실행해야하므로 속도가 많이 느려질 수 있다.

 

하지만 사용자는 Act_new 함수가 2개의 커널로 쪼개지는 것이 아니라 하나의 kernel 안에서 add 와 sqrt 연산을 다 처리함으로서 데이터 이동에 대한 퍼포먼스 오버헤드가 사라지도록 짤 수 있다.

 

CUDA GPU 프로그래밍의 단점

쿠다 병렬 프로그래밍에는 3가지 관점으로 아키텍처를 알아야 한다.

 

DRAM (지피유 (밖의) 메모리) 로부터 데이터를 전송받을 때 memory coalesce 를 한다.

memory coalesce 란 global memory 에 여러개 thread 가 동시에 연속된 메모리 공간에 접근하도록 하는 것이다.

이는 데이터를 더 빨리 가져올수 있다.

프로그래머는 자신이 사용하는 알고리즘의 메모리 access 방식을 알고 coalesce 할 수 있도록 해줘야 한다.

 

사용자가 직접 read write 가 가능한 GPU 안의 Shared Memory 사용도 알아야하고

 

cuda 코드가 GPU 구조에 따라 어떻게 분산처리될 지도 알아야 한다.

 

⇒ Triton 은 위의 모든 과정을 자동화해줄 수 있다.

 

Example

Triton 을 사용하는 예시를 확인해보자.

 

** 아래 코드그림은 import 가 필요하다.

import triton module as tl

2024-1 인공지능융합기반시스템개론 ❘ 이화여자대학교 심재형 교수님 ppt

 

위 코드는 Triton 을 이용하여 element wise sum 연산을 병렬화한다.

 

pid 는 각 병렬로 수행되는 연산의 id 를 받는다.

 

메모리로부터 x y 각각(4x1) 을 (x_ptr=x의 포인터) 포인터를 이용해  tl.load 함수를 통해 가져온다.

그리고 각 element 마다 pid 가 0 1 2 3 이라고하면 이를 offset 으로 사용한다.

 

tl.store() : 을 이용하여 output 을 포인터의 해당 offset에 저장한다.

 

이 연산이 pid 4개 가 병렬로 동작하고 있는 것이다.

 

(cuda 도 비슷하지만) triton 은 data indexing 이 대부분의 코드에 해당한다.

 

Use Case

FlashAttention: Fast and Memory-Efficient Exact Attention with IO Awareness

 

일반적으로 구현된 어텐션 연산은 메모리 hierarchy 가 고려되지 않는다.

 

속도가 빠른 SRAM 을 최대한 활용하는 방식이 Flash attention 이다.

 

OpenAI Triton 은 torch 에서는 맨 오른쪽 그림처럼 matmul-dropout-softmax-mask-matmul 이 따로따로 구현되어있었다면 이를 하나의 kernel 로 구현할 수 있다. = Kernel Fusion