Weight Folding, CUDA Streams, and the Bug That Made My Model Speak Backwards — Filip Makraduli

한 줄 핵심

RMS norm을 몇 가지 대수적 변형으로 가볍게 만들고 CUDA 스트림을 병렬화해 모델 속도를 높이되, 커널 수준의 동기화 실수를 잡아야 실제로 정상 동작한다는 점을 보여준다.

시간별 요약

1. 0:00
   RMS norm을 위한 두 줄의 대수 RMS norm을 더 싸고 빠르게 만드는 간단한 대수적 변형이 핵심이며, 이는 Layer Norm이 RMS norm으로 대체된 흐름처럼 트랜스포머 내부의 표준을 다시 개선하려는 시도다. FlashAttention처럼 메모리 왕복과 통신을 줄이는 사고방식으로 접근한다.
2. 2:07
   연산이 거의 없어도 시간이 많이 드는 이유 RMS norm은 실제 수학 연산량은 적지만 디코드 단계에서 여러 번 실행되며, 시작 오버헤드와 메모리 이동 때문에 wall time이 크게 늘어난다. 따라서 정규화와 matmul을 더 촘촘하게 묶고, 실행 대기 시간을 줄이는 방향이 중요하다.
3. 3:44
   가중치 접기, 지연된 나눗셈, 제거된 사전 정규화 가중치 접기(weight folding), 나눗셈 지연(deferred division), 사전 정규화 제거(pre norm dropped)라는 세 가지 대수적 트릭으로 RMS norm 비용을 줄인다. 특히 일부 연산은 오프라인에서 미리 합치고, 일부는 병렬화하며, 최신 구조처럼 RMS가 두 번 등장하는 경우에는 하나를 생략할 수 있다.
4. 6:33
   과거에서 나온 출력 직접 구현한 실험에서 출력이 한 단계 늦게 나오고 같은 단어가 반복되는 이상 현상이 나타나면서 GPU 스트림 문제를 의심하게 된다. 이 현상은 단순한 품질 저하가 아니라 이전 단계의 값이 섞여 들어오는 동기화 오류의 징후다.
5. 7:25
   Tensor cores와 CUDA cores의 병렬 실행 Tensor cores는 matmul을, CUDA cores는 원소별 연산과 reductions, sqrt 같은 보조 계산을 맡겨 동시에 돌릴 수 있다. 순차 실행으로 생기는 유휴 시간을 줄이려면 Python이 아니라 CUDA 코드 수준에서 병렬 구조를 직접 짜야 한다.
6. 8:45
   암묵적 조인과 레이스 컨디션 두 스트림의 종료를 암묵적으로 처리한 탓에 한쪽 작업이 끝나기 전에 결과를 읽는 race condition이 발생했다. 단위 테스트와 짧은 생성에서는 드러나지 않았지만, 긴 생성에서 과거 값을 읽는 문제가 나타나면서 명시적 join이 필요함이 드러난다.
7. 9:58
   양쪽 스트림이 끝날 때까지 기다리게 만들기 matmul과 RMS 작업의 끝을 각각 표시한 뒤, post scale이 두 스트림 모두를 기다리도록 바꾸면 오래된 버퍼 값을 읽지 않게 된다. 이 수정으로 모델이 거꾸로 말하던 버그가 사라지고, 제안한 병렬화가 정상적으로 작동한다.
8. 10:42
   오픈 모델에서의 결과와 바로 쓸 수 있는 것 llama 계열을 중심으로 여러 실험을 했고, weight folding 같은 단순한 변형만으로도 개선이 보였다. 일부는 Transformer Tricks의 코드나 flashify로 바로 적용할 수 있고, torch compile이나 quantized models에도 맞물리며, Hugging Face 모델과 배포용 클러스터까지 연결할 수 있다.
9. 13:54
   임대 엔드포인트에서 커널 작업을 할 수 없는 이유 kernel manipulation이나 flash norm 같은 작업은 inference를 소유하지 못한 rented endpoint에서는 다루기 어렵다. 오픈 소스 클러스터와 모델 서빙을 함께 소유해야 작은 모델을 유연하게 바꾸고, 큐잉·API 제어·생산 배포까지 한 번에 실험할 수 있다.
10. 16:00
    논문, 저장소, 그리고 연락처 이 작업은 논문과 repo, Hugging Face 모델, 그리고 여러 PR로 이어진 연구·구현 여정의 결과물이다. 오픈 소스 모델 개선과 배포에 관심이 있으면 LinkedIn과 관련 저장소를 통해 추가 질문이나 기여를 이어갈 수 있다.

실무 인사이트

• 트랜스포머 최적화는 알고리즘 자체보다도 메모리 이동, 시작 오버헤드, 스트림 동기화 같은 시스템 비용을 줄이는 데서 체감 성능이 크게 난다.
• CUDA 수준의 병렬화는 단위 테스트를 통과해도 장기 실행에서 깨질 수 있으므로, 짧은 검증 외에 긴 생성과 race condition 점검이 필수다.
• 연구용 최적화라도 오픈 소스 클러스터와 배포 경로를 함께 갖추면, 실험 결과를 바로 제품 수준으로 끌어올리기 쉬워진다.