가중치 폴딩, CUDA 스트림, 그리고 모델을 거꾸로 말하게 만든 버그 — 필립 마크라둘리(Filip Makraduli)

AAI Engineer
컴퓨터/소프트웨어AI/미래기술

스크립트

00:00:00.
00:00:13안녕하세요 여러분, 와주셔서 감사하며
00:00:18이제 발표를 시작하겠습니다.
00:00:22이번 발표는 제가 작성한 논문에 관한 내용인데,
00:00:27매우 간단한 주제입니다.
00:00:33제시하는 바는 아주 명확합니다.
00:00:36기본적으로 단 두 줄의
00:00:39대수 연산으로 트랜스포머의
00:00:42RMS Norm 레이어 비용을 줄이고
00:00:45속도를 높이며, 트랜스포머
00:00:48아키텍처 내 레이어 성능을
00:00:51개선하는 방식입니다. 예전에
00:00:54Layer Norm이 표준으로 쓰이다가
00:00:57RMS Norm으로 대체된 것과 비슷하게,
00:01:00이 연구도 그런 사고방식을 따릅니다.
00:01:04그리고 감사하게도 오픈소스 진영의
00:01:09몇몇 분들을 만날 기회가 생겼고,
00:01:11이 아이디어를 처음 제안한
00:01:14닐스 그라프(Nils Graf)와 함께
00:01:17이 논문을 공동 집필하게 되었습니다.
00:01:21작업은 거기서부터 시작되었죠.
00:01:23이 논문은 arXiv에 공개되어 있습니다.
00:01:26확인해 보시고, 읽고 테스트해 보세요.
00:01:29리포지토리도 준비되어 있습니다.
00:01:31이 개념, 즉 핵심 아이디어와
00:01:36사고방식은 FlashAttention에 비유해
00:01:38설명하는 것이 가장 쉽습니다.
00:01:40FlashAttention이 행렬 곱셈이
00:01:45있을 때까지 기다렸다가
00:01:48메모리 간의 통신을 제한함으로써
00:01:52전체 프로세스를 가속하는 것과
00:01:54비슷한 접근법이죠.
00:01:55이 역시 동일한 맥락의 아이디어입니다.
00:01:58이를 통해 특정 개선을 이루어 내어
00:02:01RMS Norm 연산 속도를 훨씬 높이고
00:02:08결과적으로 전체 트랜스포머 성능을 향상시킵니다.
00:02:15여기서 한 가지 질문이 생기죠. “왜 하필 RMS Norm인가?”
00:02:19그 레이어는 수학적 연산을 거의 하지 않는데 말이죠.
00:02:24맞는 말씀입니다.
00:02:26실제 연산량이 차지하는 비율만 따져보면
00:02:29아주 적은 편입니다.
00:02:31하지만 실제 소요 시간(Wall time)을 보면
00:02:34꽤 큰 비중을 차지합니다.
00:02:36예를 들어, 단 한 번의 디코드 단계,
00:02:40즉 추론이 실행되는 시점에
00:02:43RMS Norm이 무려 33번이나 호출될 수 있습니다.
00:02:47물론 모델 구조 등에 따라 달라지지만요.
00:02:50논문에는 이에 사용된 구체적인 모델과
00:02:52테스트 방식이 기재되어 있습니다.
00:02:54핵심 질문은 이 과정을 어떻게 개선할 것인가,
00:02:59그리고 행렬 곱셈을 위한 이 가중치 대기 시간을
00:03:03어떻게 회피할 수 있는가입니다.
00:03:06이 프로세스가 느린 이유는
00:03:09GPU가 수학 연산 자체를 못하거나 느려서가 아니라,
00:03:15실제 연산 외의 다른 모든 작업에 취약하기 때문입니다.
00:03:20즉, 작업(실제 연산)을 개시하는 과정 말입니다.
00:03:24일부 실험에서처럼 프로세스를
00:03:3033번이나 새로 시작하는 것은
00:03:32많은 시간을 소모하게 됩니다.
00:03:35예를 들어, 각 정규화(Normalization) 과정을
00:03:40행렬 곱셈에 결합(Fusion)하면 이를 방지하는 데 도움이 됩니다.
00:03:44또한 Weight Folding 기법을 적용하면
00:03:46메모리 간 데이터 이동을 줄일 수 있는데,
00:03:50이 데이터 이동 역시 GPU에 병목이 되는 과정입니다.
00:03:54대기 시간 문제도 마찬가지입니다.
00:03:56예를 들어 RMS Norm 레이어에서 수행되는
00:04:00나눗셈 연산을 뒤로 미루는 것(Deferring)도
00:04:02이러한 대기 단계를 회피하는 한 방법입니다.
00:04:06요컨대 이 논문이 다루는 핵심은
00:04:09RMS Norm 계산 방식에
00:04:13몇 가지 대수학적 트릭을 적용하여
00:04:16이 세 가지 측면을 모두 개선하는 것입니다.
00:04:19그게 전부입니다.
00:04:21수학적으로 보면 트릭은 다음과 같습니다.
00:04:25주로 처음 두 가지 명제에 관한 내용인데요.
00:04:29첫 번째는 Weightless Normalization입니다.
00:04:31여기서 확인하실 수 있습니다.
00:04:33그리고 Deferred Normalization입니다.
00:04:35이것이 두 번째 명제입니다.
00:04:37최신 아키텍처에서는
00:04:39RMS가 두 번 등장하는 경우가 있습니다.
00:04:44예를 들어 Gemma 4 모델에서 이런 현상이 나타나죠.
00:04:48따라서 사전 정규화(Pre-normalization)를 상쇄하는 방법도 효과가 있습니다.
00:04:52이 모든 내용은 논문에서 대수적으로 증명되어 있습니다.
00:04:58첫 번째 명제는 바로 이것인데,
00:05:00이득(Gain)과 가중치 폴딩이
00:05:04하나의 행렬로 흡수되는 방식입니다.
00:05:06여기 별표가 붙은 W 행렬로 표시되어 있습니다.
00:05:09그리고 이건 오프라인으로 미리 계산됩니다.
00:05:12FlashAttention에서 매번
00:05:14메모리 통신이 발생하지 않도록
00:05:16일부 연산을 미리 사전에
00:05:18처리해 두는 것과 비슷하죠.
00:05:20이것이 적용된 단계 중 하나인
00:05:24Weight Folding 기법입니다.
00:05:26다른 단계는 행렬 곱셈의 스칼라 나눗셈을
00:05:32지연시켜 병렬로 처리할 수 있게 만드는 것입니다.
00:05:35일반적인 경우라면 먼저 계산을 수행한 뒤
00:05:38대기했다가 다시 계산해야 합니다.
00:05:41하지만 이 방식은 이를 분리하여
00:05:44병렬 처리할 수 있도록 만드는 것이 핵심입니다.
00:05:48그리고 일종의 변형 형태인 세 번째 명제는,
00:05:51두 개가 존재할 경우,
00:05:56스케일 불변성(Scale Invariance) 덕분에
00:05:58하나를 생략해도 정상 작동한다는 것입니다.
00:06:01이는 이러한 아키텍처와 구현을
00:06:04갖춘 최신 모델들에 적용 가능합니다.
00:06:10이를 실제 구현하기 위해서는,
00:06:13특히 두 번째 명제의 경우 더욱 그렇습니다.
00:06:16첫 번째 명제 같은 경우는 간단합니다.
00:06:20'transformer tricks'라는 리포지토리가 있는데,
00:06:22어떤 모델에든 적용하면 바로 잘 동작합니다.
00:06:25하지만 두 번째 명제는 커널 작업이 필요합니다.
00:06:29따라서 구현이 그리 간단하지 않습니다.
00:06:31그래서 제가 직접
00:06:35구현을 진행하다가 이런 실험 결과를 얻은 적이 있습니다.
00:06:42전반적으로는 괜찮아 보입니다.
00:06:46프롬프트가 “트랜스포머 아키텍처는
00:06:48NLP에 혁명을 일으켰다, 왜냐하면~” 이고
00:06:51그 뒤에 예상되는 출력 내용이 이어집니다.
00:06:54하지만 실제로 출력된 결과에서는
00:06:56반복 현상과 1단계 지연(Lag)이 관찰되었습니다.
00:06:59보시다시피 “because”라는 단어가 다시 나타나죠.
00:07:01GPU 스트림 쪽에 문제가 발생하고 있었고
00:07:07원인이 무엇인지 파악하려고 애썼습니다.
00:07:101단계 지연과 함께 일종의 이전 시점의
00:07:15출력값들이 뒤늦게 나오는 증상이었죠.
00:07:18이 모든 것을 디버깅하는 과정에서
00:07:21이러한 구조를 구축하는 절차에 대해 깨달았습니다.
00:07:26명제 2에서 설명한 것처럼 두 연산을 지연/미루는 경우,
00:07:32CUDA에서는 두 가지 처리가 가능합니다.
00:07:35행렬 곱셈의 일부를 담당하는 텐서 코어(Tensor Cores)를 쓸 수도 있고,
00:07:39요소별 연산, 리덕션, 제곱근 등을 실행하는
00:07:44CUDA 코어를 사용할 수도 있습니다.
00:07:47아이디어는 이 연산들을 병렬로 처리하여
00:07:52논문에서 설명한 이득을 얻고 실제 개념을 검증하는 것이었습니다.
00:07:58원래는 이런 모습이어야 했습니다.
00:08:00순차적으로 작업을 처리하게 되면
00:08:03벡터 유닛이 RMS 및 스케일링을 연산할 때 발생하는 유휴 대기 시간이 있고,
00:08:10그다음에 행렬 곱셈이 이어집니다.
00:08:12아이디어는 논문에 나온 기법인 Flash Norm을 사용해서
00:08:16이 두 작업을 병렬로 처리하는 것이었습니다.
00:08:19행렬 유닛은 행렬 곱셈을 연산하고, 벡터 유닛은 RMS를 연산하는 거죠.
00:08:23그런 방식으로 시간을 절약할 수 있습니다.
00:08:26하지만 이건 파이썬 수준에서만 처리할 수는 없고,
00:08:28더 하위 수준으로 내려가야 합니다.
00:08:30그래서 이런 식으로 CUDA 코드를 작성해 구현했습니다.
00:08:35당시에는 전반적으로 괜찮아 보였습니다.
00:08:42하지만 제가 약간의 실수를 했다는 걸 깨달았습니다.
00:08:47마지막에 두 스트림을 합쳐야(join) 하는 부분에서 암시적(implicit) 처리가 되었다는 것이었죠.
00:08:57테스트해 봤을 때는 단위 테스트도 통과했고,
00:09:01퍼플렉시티 검사 등에서도 비슷한 수준의 텍스트가 생성되어 성능이 유사하게 보였습니다.
00:09:08하지만 기나긴 생성 과정 속에서 이 문제를 포착할 수 있었습니다.
00:09:11처음엔 이게 무슨 원인인지 전혀 알지 못했죠.
00:09:14원인은 바로 암시적 결합을 할 때 한쪽 스트림의 작업이 아직 끝나지 않았던 것이었습니다.
00:09:24결국 아직 끝나지 않은 행렬 곱셈 결과에서 이전 데이터를 읽어오는 레이스 조건이 발생한 것이죠.
00:09:31이 문제를 해결하기 위해 결합(join) 과정을 명시적으로 처리해야 한다는 아이디어를 얻었습니다.
00:09:40작업 중 하나가 끝날 때까지 대기하도록 해서, 결합 시 과거 데이터가 읽히지 않음을 확실히 보장하는 것이었죠.
00:09:48CUDA 스트림을 탐구하면서 얻은 깨달음이 바로 그것이었습니다.
00:09:54기존에는 이렇게 처리가 되어 있었습니다.
00:09:57결합이 암시적으로 이루어졌던 거죠.
00:10:00그래서 후속 스케일링 연산이 이전 버퍼 값을 읽어오게 되었습니다.
00:10:06해결책은 우선 행렬 곱셈의 종료 시점을 명시적으로 표시하고,
00:10:14RMS의 종료 시점도 표시하는 것입니다.
00:10:17그리고 후속 스케일링 작업이 첫 번째 스트림을 기다리게 하고,
00:10:21이어서 두 번째 스트림을 기다리도록 만든 것이죠.
00:10:24그 덕분에 버그가 수정되었고, 논문의 방식대로 모델이 거꾸로 엉뚱한 말을 하지 않고 똑바로 말을 출력할 수 있게 되었습니다.
00:10:33여기까지가 흥미로운 학술적 관점의 이야기였습니다.
00:10:38하지만 직접 다양하게 시도해 보고 싶기도 했죠.
00:10:40이를 배포하고 테스트하여 실제 프로덕션 환경에서 어떻게 잘 작동시킬지 확인하는 것 말이죠.
00:10:47논문을 읽어보시면 제반 테스트 결과를 확인하실 수 있습니다.
00:10:50대부분 Llama 모델 중심으로 진행되었지만, 다른 구조에서도 동일하게 잘 작동합니다.
00:10:56이 논문과 관련해서는, 앞서 설명해 드린 명제 1의 가중치 폴딩 같은 경우 저장소에 있는 코드만으로 바로 적용할 수 있습니다.
00:11:10'flashify'라는 명령 하나면 자동으로 처리되죠.
00:11:13하지만 제가 언급했던 두 번째 부분의 경우, 적용하려면 조금 커널 작업이 필요합니다.
00:11:19앞선 예시에서 설명해 드린 바와 같이 말이죠.
00:11:22이것은 Llama 모델 기반의 몇 가지 결과물이며, 세부적인 항목들도 함께 확인해 보실 수 있습니다.
00:11:29지연 정규화(Deferred Normalization)만 수행할 때와 완전 융합 커널(Fused Kernel)을 적용할 때 각각 어떤 변화가 일어나는지 말이죠.
00:11:36모든 명제들을 검증하기 위해 이렇게 하위 단에서 수많은 실험이 이뤄졌습니다.
00:11:41그리고 다양하게 정밀함과 세부 조건을 달리하여 산출해 낸 결과들이 이것입니다.
00:11:48단순히 가중치 폴딩만 적용한 간단한 방식조차도 분명한 성능 향상을 보여줍니다.
00:11:54또한 모델 개발에 상시 사용하는 일상적인 도구들과도 완벽히 호환됩니다.
00:11:59바퀴를 새로 재발명하거나 처음부터 모든 걸 다시 구축할 필요가 없다는 것이죠.
00:12:05일종의 새로운 체크포인트 개념이기 때문에 torch compile과도 잘 맞물려 작동합니다.
00:12:13Flash Attention 역시 다른 레이어에서 이와 비슷한 기법을 활용합니다.
00:12:16양자화된 모델과도 문제없이 함께 동작하죠.
00:12:18따라서 이 기법을 도입해 혁신적인 정규화 레이어가 적용된 멋진 모델을 구축하는 것도 얼마든지 가능합니다.
00:12:27이를 구동하기 위한 세부 정보와 코드는 'transformer tricks' 저장소에서 확인하실 수 있습니다.
00:12:34설명해 드린 여러 대수적 기법들과 함께 언급했던 논문 내용이 함께 담겨 있습니다.
00:12:41또한 깃허브가 아니라, 제가 몇몇 모델에 이 작업을 적용해 둔 허깅페이스 모델 저장소도 마련되어 있습니다.
00:12:50해당 모델의 허깅페이스 링크를 통해 직접 테스트해보실 수 있습니다.
00:12:54그리고 이 허깅페이스 모델들을 이용해 실제 프로덕션 환경에 배포하는 것도 가능합니다.
00:13:00이 작업을 구상할 때, 연구 단계가 마무리되고 허깅페이스 모델 링크가 준비된 시점에서
00:13:11Superlink의 추론 엔진이 어떤 허깅페이스 모델이든 배포하기에 정말 훌륭한 방법이라는 걸 알게 되었습니다.
00:13:19해커톤에서도 참가자들이 직접 파인튜닝한 커스텀 허깅페이스 모델이나 체크포인트를 가져와서 활용하곤 했습니다.
00:13:27모델을 개선하고 본인만의 연구 아이디어를 검증하기 위해 이러한 대수적 트릭을 적용해 보려는 경우에도,
00:13:37모델 배포에 필요한 번거로운 글루 코드(glue code) 걱정 없이 클러스터 상에 배포된 버전으로 직접 테스트해 볼 수 있습니다.
00:13:48정말 유용하죠. 핵심은 전체 클러스터와 모델 추론 환경이 오픈소스로 제공된다는 점입니다.
00:13:58덕분에 커널 조작이나 Flash Norm 같은 신규 연구 아이디어들을 자유롭게 실험해 볼 수 있습니다.
00:14:12추론 환경 제어 권한이 없는 임대형 엔드포인트에서는 이런 시도를 하기가 훨씬 어렵기 때문이죠.
00:14:18따라서 이러한 실험을 허용하는 이식성과 유연성을 갖추면서도,
00:14:23대규모 환경에서도 안정적으로 검증할 수 있을 만큼 상용화 준비가 된 도구가 필요합니다.
00:14:27예를 들어 Sy를 사용하면 이를 다른 모델들과 조합할 수도 있습니다.
00:14:33왼쪽 상단에 보이듯, 원한다면 Flashify 처리된 모델과 다른 모델들을 결합해 에이전트(agentic) 작업을 수행할 수 있죠.
00:14:43그렇게 엔드투엔드의 더 큰 유스케이스를 구현하는 것입니다.
00:14:46Sy가 작동하는 방식은 프로덕션 클러스터가 모델 배포를 돕는 구조입니다.
00:14:52자세한 사항은 Sy 저장소를 참고하셔도 좋습니다.
00:14:57또한 더 스마트한 큐잉 메커니즘이 있어, 특히 소형 모델로 작업할 때 큰 도움이 됩니다.
00:15:03Flash Norm 작업 당시에 저 역시 비교적 작은 Llama 모델 및 허깅페이스의 소형 에이전트들을 활용했기 때문입니다.
00:15:11동일한 GPU 상에서 소형 모델들을 배포할 수 있으면 불필요한 GPU 비용 지출을 줄일 수 있을 뿐만 아니라,
00:15:24특히 소형 모델 간의 전환을 유연하게 처리할 수 있어 매우 유용했습니다.
00:15:30클러스터는 물론 API를 통해 모델 설정을 자유롭게 제어할 수 있다는 점도,
00:15:35오픈소스 연구 과정에서 전담 인프라 엔지니어의 도움 없이도 처리할 수 있어 상당히 편리합니다.
00:15:40이 점 또한 대단히 매력적이죠.
00:15:43클라우드 환경에 대한 완전한 소유권을 가질 수 있어 오픈 웨이트, 오픈 모델, 오픈소스를 추구할 때 매우 유용합니다.
00:15:50또한 Sy에는 제가 언급한 모델 외에도 다양한 모델 카탈로그가 구비되어 있으니
00:15:57한번 둘러보시기 바랍니다.
00:15:59유사한 시스템을 구축 중이시라면 리랭킹 임베딩 모델 등도 확인해보실 수 있습니다.
00:16:03이것으로 트랜스포머를 개선하는 기법에 대한 논문을 공동 집필했던 저의 연구 여정 이야기를 마치고자 합니다.
00:16:15연구 결과물을 상용 배포 환경에 적용하고 테스트하며, 오픈소스 모델들을 자유롭게 다룰 수 있는 계기가 되었습니다.
00:16:24질문이 있거나 기여하고 싶으신 부분이 있다면 언제든 링크드인으로 편하게 연락해 주세요.
00:16:30말씀드린 내용 중 상당수는 vLLM이나 허깅페이스 등에 제출한 Pull Request 형태로도 존재합니다.
00:16:36여기저기서 찾아보실 수 있을 겁니다.
00:16:38논문도 꼭 한번 확인해 보세요.
00:16:40화면에 보이는 arXiv 링크를 참고하시면 됩니다.
00:16:43Sy 저장소 및 제 링크드인 주소도 함께 확인하실 수 있습니다.
00:16:47참석해 주셔서 정말 감사합니다.
00:16:58질문이 있으신 분은 찾아와 주세요.
00:16:59저희는 계속 여기에
00:17:00가까이 머물 예정입니다.
00:17:13감사합니다.

핵심 요약

단 두 줄의 대수적 트릭과 Weightless Normalization 및 Deferred Normalization 기법을 적용하면 RMS Norm 연산 비용을 절감하고 트랜스포머 아키텍처의 성능을 향상시킬 수 있다.

하이라이트

  • 단 두 줄의 대수 연산으로 트랜스포머의 RMS Norm 레이어 비용을 줄이고 속도를 높일 수 있다.

  • 추론 시 단 한 번의 디코드 단계에서 RMS Norm이 최대 33번까지 호출된다.

  • 가중치 폴딩(Weight Folding) 기법을 적용하면 GPU 병목 현상의 주요 원인인 메모리 간 데이터 이동을 줄일 수 있다.

  • CUDA 스트림에서 암시적 결합(implicit join)을 사용할 경우 아직 끝나지 않은 행렬 곱셈 결과에서 데이터를 읽어오는 레이스 조건이 발생한다.

  • 작성된 코드는 'transformer tricks' 저장소에서 'flashify' 명령어로 바로 적용할 수 있다.

타임라인

RMS Norm 최적화를 위한 대수적 접근법

  • 트랜스포머 아키텍처 내 RMS Norm 레이어의 비용을 줄이고 속도를 높이는 두 줄의 대수 연산 방식이 제안되었다.
  • FlashAttention과 유사하게 메모리 간 통신을 제한하여 전체 프로세스를 가속하는 원리를 따른다.
  • 해당 연구 결과는 arXiv에 공개된 논문과 오픈소스 리포지토리를 통해 확인할 수 있다.

과거 Layer Norm에서 RMS Norm으로 대체된 것처럼, 본 연구는 트랜스포머 내 정규화 레이어의 성능을 개선하기 위한 간단하지만 명확한 대수적 방식을 다룬다. FlashAttention이 행렬 곱셈 전 메모리 통신을 제한해 가속하는 것과 동일한 맥락에서 RMS Norm의 연산 속도를 높인다.

RMS Norm의 병목 현상과 세 가지 대수적 트릭

  • RMS Norm 자체의 수학적 연산량은 적지만, 추론 시 한 번의 디코드 단계에서 최대 33번 호출되어 실제 소요 시간에서 큰 비중을 차지한다.
  • GPU가 느린 이유는 연산 능력 부족이 아니라 잦은 작업 개시와 메모리 간 데이터 이동 때문이다.
  • Weightless Normalization, Deferred Normalization, 그리고 사전 정규화를 상쇄하는 기법이 핵심 대수적 트릭으로 활용된다.

실제 연산량이 적음에도 불구하고 빈번한 호출 횟수와 GPU 작업 개시 지연으로 인해 전체 성능에 영향을 미친다. 이를 해결하기 위해 정규화 과정을 행렬 곱셈에 결합하고, 나눗셈 연산을 뒤로 미루는(deferring) 등의 대수적 트릭을 적용하여 메모리 이동과 대기 시간을 줄인다.

CUDA 스트림 병렬화와 레이스 조건 버그 해결

  • 두 번째 명제인 Deferred Normalization을 구현하는 과정에서 텐서 코어와 CUDA 코어를 병렬로 사용하는 커널 작업이 필요하다.
  • 암시적 결합 방식으로 인해 행렬 곱셈이 끝나지 않은 상태에서 데이터를 읽어오는 레이스 조건이 발생했다.
  • 행렬 곱셈과 RMS의 종료 시점을 명시적으로 표시하고 후속 작업이 각 스트림을 차례로 대기하도록 수정하여 버그를 해결했다.

파이썬 수준을 넘어 하위 수준의 CUDA 코드를 작성하는 과정에서 스트림 결합 방식의 문제로 인해 모델 출력이 반복되거나 지연되는 현상이 발생했다. 두 스트림의 종료 시점을 명시적으로 동기화하는 코드로 수정한 뒤 정상적인 텍스트 출력을 회복하였다.

실제 배포 환경 적용 및 오픈소스 생태계 활용

  • 가중치 폴딩은 'transformer tricks' 저장소의 'flashify' 명령어로 쉽게 적용할 수 있다.
  • torch compile, 양자화 모델 등 기존 도구들과 완벽히 호환된다.
  • 오픈소스 추론 엔진과 클러스터를 활용하면 커널 조작 및 Flash Norm 같은 신규 연구 아이디어를 실무 환경에서 자유롭게 실험할 수 있다.

Llama 모델 등을 중심으로 검증된 해당 기법들은 허깅페이스 모델 저장소와 연동하여 실제 프로덕션 환경에 배포할 수 있다. 인프라 제어 권한을 가진 오픈소스 클러스터 환경을 통해 개발자는 전담 엔지니어의 도움 없이도 고도화된 실험을 안정적으로 수행할 수 있다.

커뮤니티 글

아직 글이 없습니다. 이 영상에 대한 첫 번째 글을 작성해 보세요!

이 영상에 대해 글쓰기