LLM을 빠르게 만들 때 가장 먼저 떠올리는 건 거대한 행렬 곱이야. 그런데 그 행렬 곱 사이사이에 끼어 있는 정규화가 발목을 잡을 때가 있어. 계산 자체가 무거워서가 아니라, 데이터를 GPU 메모리에서 읽고 다시 쓰느라 기다리기 때문이지.
PyTorch 팀은 이 기다림을 줄이는 방법을 공개했어. 핵심은 정규화 연산을 따로 끝내지 않고, 앞뒤의 행렬 곱(GEMM) 안에 섞어 넣는 거야. 정규화를 공짜로 만든다는 뜻이 아니라, 이미 돌아가는 계산 뒤에 최대한 숨긴다는 뜻이야.1
왜 작은 연산이 시간을 먹을까
LayerNorm과 RMSNorm은 각 벡터의 크기나 분포를 일정하게 맞추는 연산이야. 모델 학습을 안정시키고, LLM의 거의 모든 블록에 반복해서 들어가지. 하지만 이 연산은 GPU의 Tensor Core를 바쁘게 쓰기보다 메모리에서 값을 옮기는 시간이 큰 편이야.
PyTorch 글은 이런 정규화가 메타의 추천 모델 학습 지연의 약 20%를 차지했고, 일반적인 LLM에서도 약 10% 수준일 수 있다고 설명해. 거대한 행렬 곱 바로 앞이나 뒤에서 계속 멈춰 서는 셈이야.1
문제는 둘의 일하는 방식이 다르다는 데 있어. 행렬 곱은 큰 행렬을 작은 타일로 쪼개 병렬 처리한다. 반면 정규화는 한 행 전체를 봐야 평균이나 크기를 계산할 수 있어. 행렬 곱의 타일 하나만으로는 그 행 전체를 알 수 없으니, 결과를 메모리에 한번 저장하고 정규화가 다시 읽는 경로가 생긴다.
먼저 곱하고, 나중에 나눈다
RMSNorm에는 예외적으로 쓸 수 있는 요령이 있어. 행 전체에 같은 값으로 곱하는 부분은 행렬 곱의 앞에서 하든 뒤에서 하든 결과가 같아.
RMSNorm(A) @ B를 곧바로 계산하는 대신 A @ B를 먼저 계산하고, 행마다 필요한 보정값을 마지막에 곱할 수 있다는 뜻이야. 이 방식에서 커널은 행렬 곱을 하는 동안 제곱합도 함께 모으고, 끝에서 한 번만 보정한다. PyTorch는 이를 Lazy Pre-Norm이라고 불렀어.1
다만 아무 정규화에나 통하는 마술은 아니야. 열마다 다른 값을 곱하는 affine 항이 있으면 이 교환이 깨지고, 평균을 빼는 LayerNorm에도 그대로 적용할 수 없어. 역전파에서는 중간 결과를 다시 구성해야 해 구현도 복잡해진다.1
여러 블록이 한 행을 나눠 본다
행렬이 커지면 한 블록이 행 전체를 품을 수 없다. PyTorch의 다른 방법은 여러 GPU 실행 블록(CTA)이 같은 행의 다른 조각을 맡고, 작은 요약값만 서로 주고받게 하는 거야. 이렇게 하면 각 블록이 큰 결과를 전역 메모리에 내보냈다가 다시 읽지 않고도 정규화를 끝낼 수 있다.
발표에 따르면 이 방식은 특정 정규화 커널 지연의 최대 90%를 다른 계산과 겹칠 수 있었고, attention 주변의 여러 정규화를 합친 실험에서는 커널 속도가 최대 35% 빨라졌어.1 숫자는 인상적이지만, bfloat16·NVIDIA B200·750W 전력 제한의 메타 데이터센터 환경에서 측정했고 일부 실험은 affine 항을 끈 조건이야. 모든 모델과 GPU에서 같은 폭으로 빨라진다는 뜻은 아니다.
다음에 볼 것
이 기법이 의미 있으려면 정규화가 실제 병목이어야 해. 행렬 곱이 전체 시간을 거의 다 차지하는 큰 형상에서는 숨길 정규화 시간이 작아져 이득도 줄어든다. PyTorch도 차원이 커질수록 정규화 비중이 5% 아래로 내려간다고 적었어.1
그래서 다음 질문은 단순해. 이 커널을 실제 LLM 추론과 학습 경로에 넣었을 때, 모델 구조·형상·정밀도별로 얼마나 재현되는가. 공개된 구현과 다른 하드웨어의 벤치마크가 쌓이면, 이 아이디어가 특수한 커널 기교인지 LLM 기본 경로의 변화인지 더 분명해질 거야.
댓글