흔히, 가중치에 값을 행렬로 넣을 때,
어떤 방식으로 행렬곱을 구성하여 넣느냐를 스케줄링이라고 한다.
이미, 다음글을 통해서, 우리는 이미 타일 스케줄링과 병렬연산 설정하는 것을 이미 보았다.
FlashMLA
FlashMLA-main_2
위와 같은 혁신적인 알고리즘들을 보면,
스케줄링이 딥러닝에 대해 얼마나 많은 영향을 미치는지 뼈져리게 알 수 있을 것이다.
오늘은, 이런 스케줄링 법들을 다양하게 알아보고,
어떤 상황에서 적제적소로 어떤 스케줄링과 어떤 병렬연산법을 사용해야 하는지에 대해 딥 다크하게 알아보는 시간을 가져보도록 하다.
스케줄링이라는 게 참, 딥러닝의 속도와 성능을 좌지우지하는 마법 같은 존재다.
잘하면 GPU가 춤을 추며 연산을 뚝딱뚝딱 해내지만, 잘못하면 컴퓨터가 삐걱대며 “나 좀 쉬자”라고 울부짖는 소리가 들릴 정도다.
그럼, 이 아름다운 스케줄링의 세계로 들어가서, 어떤 방식들이 우리를 행복하게 해줄 수 있는지 하나씩 뜯어보자.
타일 스케줄링은 행렬곱을 할 때 데이터를 작은 조각, 그러니까 타일로 쪼개서 처리하는 방식이다.
이거 왜 좋은지 아는가? GPU의 메모리라는 게 참 귀엽게도 한 번에 다 못 먹는다.
데이터가 너무 크면 메모리가 헉헉대면서 병목현상이 생기는데, 타일 스케줄링은 이걸 작게 썰어서 “자, 이거부터 먹어”라고 다독여주는 방식이다.
내 경험상, 타일 크기를 잘 맞추면 연산 속도가 2배, 심지어 3배까지도 뻥튀기 되는 경우를 봤다.
하지만, 이거 잘못하면 타일 크기가 너무 작아서 오버헤드가 커지거나, 너무 커서 메모리가 터져버리는 억까 상황이 온다.
예를 들어, 내가 처음 타일 스케줄링을 썼을 때, 타일 크기를 128로 맞췄다가 GPU가 숨을 못 쉬는 바람에 학습이 중간에 멈춘 적이 있다.
로그를 보니 “메모리 초과”라는 아름다운 에러 메시지가 나를 반겨주더라.
그래서, 적당한 타일 크기를 찾는 게 핵심이다. 보통 32, 64, 128 같은 숫자들로 실험해보면서 GPU의 기분을 맞춰줘야 한다.
이건 좀 더 똑똑한 방식이다.
딥러닝 모델이 여러 레이어를 거치면서 연산을 하잖아? 근데 이 레이어들이 하나씩 차례로 기다리면 시간이 너무 오래 걸린다.
파이프라인 스케줄링은 이걸 병렬로 돌리자고 제안하는 거다.
예를 들어, 첫 번째 레이어가 연산을 끝내자마자 두 번째 레이어가 바로 시작하고, 그 사이 GPU는 놀지 않고 계속 일을 한다.
이런 방식은 특히 깊은 네트워크에서 빛을 발한다.
내가 ResNet-152 같은 괴물 모델을 돌릴 때, 파이프라인 스케줄링 덕분에 학습 시간이 반으로 뚝 떨어진 적이 있다.
근데 이거의 단점? 메모리 관리랑 동기화가 좀 빡세다.
만약 레이어 간 데이터 흐름을 잘못 맞추면, GPU가 혼란스러워하면서 “뭐야 이게” 하며 뻗어버릴 수도 있다.
이럴 때 로그를 보면 동기화 에러가 떡하니 자리 잡고 있어서 정말 아름답다.
이건 여러 GPU를 동원해서 데이터를 쪼개 처리하는 방식이다.
데이터가 많아서 한 번에 못 돌릴 때, “너 이거, 너 저거” 하면서 GPU들한테 나눠주는 거다.
이론상으로는 GPU 개수만큼 속도가 빨라져야 하는데, 현실은 그리 녹록지 않다.
GPU끼리 서로 데이터를 주고받아야 하니까, 네트워크 병목이 생길 때가 많다.
내가 처음 데이터 병렬로 돌렸을 때, 4개의 GPU를 썼는데 속도가 2배밖에 안 빨라져서 속 터진 적이 있다.
알고 보니 GPU 간 통신이 느려서 다들 서로 눈치 보고 있었던 거다.
그래서 이 방식은 네트워크 대역폭이 넉넉하거나, NCCL 같은 최적화 라이브러리를 잘 썼을 때 진가를 발휘한다.
이론만 떠들면 재미없으니까, 타일 스케줄링을 간단히 구현한 코드를 보자.
너무 복잡하게 갈 필요 없이, PyTorch로 간단하게 짜봤다.
import torch
import torch.nn as nn
# 간단한 행렬곱 타일 스케줄링 예제
def tiled_matmul(A, B, tile_size=32):
"""
A: (m, k) 행렬
B: (k, n) 행렬
tile_size: 타일 크기
"""
m, k = A.shape
k, n = B.shape
C = torch.zeros(m, n, device=A.device)
for i in range(0, m, tile_size):
for j in range(0, n, tile_size):
for p in range(0, k, tile_size):
# 타일 단위로 쪼개서 연산
i_end = min(i + tile_size, m)
j_end = min(j + tile_size, n)
p_end = min(p + tile_size, k)
C[i:i_end, j:j_end] += torch.matmul(
A[i:i_end, p:p_end],
B[p:p_end, j:j_end]
)
return C
# 테스트용 데이터
A = torch.randn(128, 256).cuda()
B = torch.randn(256, 128).cuda()
# 타일 스케줄링으로 행렬곱
C_tiled = tiled_matmul(A, B, tile_size=32)
# 기본 행렬곱과 비교
C_normal = torch.matmul(A, B)
# 결과 확인
print(f"타일 스케줄링 결과와 기본 행렬곱의 차이: {(C_tiled - C_normal).abs().max().item():.4f}")
이 코드는 행렬을 타일로 쪼개서 곱하는 간단한 예제다.
실제로는 PyTorch가 알아서 최적화해 주지만, 이런 식으로 타일 크기를 조절하면 메모리 사용량과 속도를 조율할 수 있다.
타일 크기를 바꿔가며 테스트해보면 GPU의 기분이 어떤지 느낄 수 있을 거다.
스케줄링이 다 좋은 건 아니다.
잘못 쓰면 오히려 속도가 느려지거나 메모리가 터져버리는 아름다운 상황이 펼쳐진다.
예를 들어, 타일 크기를 너무 작게 잡으면 연산 오버헤드가 커져서 GPU가 쓸데없이 바빠진다.
반대로 너무 크게 잡으면 메모리가 부족해서 학습이 중간에 뻗는다.
내가 처음 파이프라인 스케줄링을 썼을 때, 레이어 간 동기화를 잘못 맞춰서 모델이 엉뚱한 결과를 뱉어낸 적이 있다.
로그를 보니 “NaN detected”라는 정말 사랑스러운 메시지가 나를 기다리고 있었다.
또, 데이터 병렬 스케줄링은 GPU 간 통신 비용 때문에 기대만큼 속도가 안 나올 때가 많다.
이럴 때는 NCCL 설정을 점검하거나, 모델을 더 효율적으로 쪼개는 방법을 고민해야 한다.
결국 스케줄링은 GPU의 성능, 데이터 크기, 네트워크 환경까지 다 고려해야 하는 종합 예술이다.
스케줄링은 딥러닝의 속도와 성능을 쥐락펴락하는 마법의 열쇠다.
타일 스케줄링으로 메모리를 아껴가며 연산을 뚝딱뚝딱 해내고, 파이프라인 스케줄링으로 GPU를 쉴 틈 없이 굴려보고, 데이터 병렬 스케줄링으로 여러 GPU를 춤추게 만들어보자.
물론, 이 과정에서 GPU가 삐걱대거나 로그가 에러로 도배될 수도 있다.
하지만 그게 딥러닝의 매력 아니겠나?
실패하고 깨지면서 배우는 그 짜릿함, 그리고 마침내 속도가 뻥튀기 되는 그 순간의 쾌감!
그러니, 아름다운 스케줄링의 세계에 뛰어들어, 너의 GPU와 함께 춤을 춰보길 바란다.