GPU 학습 성능 최적화: CPU-GPU 동기화 제거부터 커널 퓨전, 커스텀 CUDA 커널까지

핵심 메시지

• GPU는 연산 하나만 놓고 보면 CPU보다 1000배 느리지만, 병렬성 덕분에 2048×2048 행렬곱에서는 CPU의 28밀리초를 200마이크로초로 줄인다.
• 파이토치는 CPU가 커널을 큐에 쌓고 GPU가 비동기로 꺼내 쓰는 구조라, 프로파일러에서 GPU 실행 막대가 빈틈없이 채워져 있는지가 첫 번째 점검 지점이다.
• `.cpu()` 호출뿐 아니라 `if` 문 안에서 텐서를 불리언으로 바꾸는 순간에도 동기화가 걸려 파이프라인이 통째로 멈춘다.
• 성능 병목은 연산, 메모리, 커널 실행 오버헤드 세 가지로 나뉘며 각각 텐서 코어, 커널 퓨전, CUDA 그래프라는 처방이 다르다.
• torch.compile이 손대지 못하는 패턴에서는 직접 쓴 CUDA 커널이 같은 작업을 20밀리초에서 30마이크로초로 줄였다.

쉽게 이해하기

금융 트레이딩 회사 제인스트리트의 엔지니어 두 명이 사내 기술 강연에서 GPU로 신경망을 학습시킬 때 성능이 어디서 새는지를 처음부터 짚었다. 이들은 먼저 상식을 뒤집는 질문을 던진다. 부동소수점 연산 하나를 처리하는 데 CPU는 약 1나노초가 걸리지만 GPU는 약 1마이크로초, 즉 1000배가 더 걸린다는 것이다. 그럼에도 GPU를 쓰는 이유는 단 하나, 병렬성이다. 2048×2048 행렬곱처럼 출력 원소마다 독립적으로 계산할 수 있는 작업에서는 CPU가 28밀리초를 쓰는 동안 GPU가 200마이크로초 만에 끝낸다.

파이토치 코드에 `.to("cuda")` 몇 줄을 더하면 연산이 GPU로 넘어가지만, 그 뒤에서 실제로 벌어지는 일은 단순하지 않다. 커널 하나를 실행하려면 CPU가 실행 파라미터를 계산해 드라이버에 넘기고, 드라이버가 GPU에 명령을 전달한다. 다행히 이 과정은 비동기라서 CPU는 GPU의 완료를 기다리지 않고 다음 커널을 계속 큐에 쌓는다. 발표자들은 엔비디아 Nsight Systems 프로파일러를 띄워놓고, GPU 실행을 나타내는 막대가 끊김 없이 채워져 있어야 한다고 강조했다. 막대에 틈이 보이면 그만큼 비싼 GPU가 놀고 있다는 뜻이다.

이 파이프라인을 무너뜨리는 주범은 CPU와 GPU 사이의 동기화다. 강연에서는 학습 루프 안에 숨은 동기화 세 곳을 관객과 함께 찾아냈는데, 지표를 모으려고 `.cpu()`를 부르는 곳, 손실이 NaN인지 검사하며 텐서를 `if` 조건에 넣는 곳, 그리고 의외로 CPU에서 GPU로 데이터를 복사하는 곳이다. 두 번째가 특히 까다롭다. `.isNaN()` 자체가 아니라 그 결과를 불리언으로 강제 변환하는 `if` 문이 동기화를 일으키기 때문에 코드만 봐서는 알아채기 어렵다. 해결책은 각각 달라서, 지표는 GPU 커널로 계산해 루프 끝에서만 가져오고, NaN 검사는 비동기로 돌리며, CPU에서 GPU로의 복사는 `pin_memory`로 메모리를 고정한 뒤 `non_blocking=True`를 넘겨 GPU가 DMA로 직접 읽어가게 만든다.

실행 막대를 가득 채웠다고 GPU를 잘 쓰는 것은 아니다. 여기서부터는 하드웨어 구조를 알아야 한다. H100에는 132개의 SM(스트리밍 멀티프로세서)이 있고, SM끼리는 사실상 대화하지 못하므로 문제를 SM 단위 조각으로 쪼개는 것이 프로그래밍의 출발점이다. 메모리 계층의 속도 차이도 결정적이다. 글로벌 메모리가 초당 3테라바이트, L2 캐시가 7테라바이트, 공유 메모리가 26테라바이트로 벌어지기 때문에 가능한 한 위쪽 계층에 데이터를 붙잡아두는 것이 효율적인 커널의 핵심이 된다.

마지막으로 두 사람은 실제 사례를 보여줬다. 연구자가 모든 그래디언트의 평균을 기록하려고 수천 개 텐서의 합을 파이썬 for 루프로 돌렸더니 20밀리초가 걸렸다. 커널 하나를 큐에 넣는 데 CPU가 5마이크로초를 쓰는데 정작 GPU 실행은 1마이크로초도 걸리지 않아 GPU가 굶주린 것이다. torch.compile을 붙이자 오히려 조금 더 느려졌는데, 컴파일러가 for 루프를 펼쳐볼 뿐 병렬 커널로 재구성하지는 못하는 데다 트라이톤 템플릿이 엔비디아 손수 최적화 커널의 80~90% 성능에 그치기 때문이다. 결국 텐서 하나를 SM 하나에 배정하는 CUDA 커널을 직접 써서 30마이크로초, 약 1000배 빠른 결과를 얻었다.

주요 인사이트

• "파란 막대가 꽉 찼다"는 것은 GPU를 쉬지 않고 쓴다는 뜻일 뿐, 잘 쓴다는 뜻은 아니다. 두 지표를 분리해서 봐야 최적화 순서를 잘못 잡지 않는다.
• 동기화는 코드에 이름을 드러내지 않는다. 텐서를 조건문·출력·형변환에 넣는 모든 지점이 후보이며, 전체 그래프 컴파일 옵션으로 컴파일해 오류를 내보는 것이 숨은 동기화를 찾는 실용적인 방법이다.
• torch.compile은 만능이 아니다. 그래프 브레이크가 많으면 퓨전 기회가 사라지는데, 파이토치가 모르는 문법이나 커스텀 연산자, 숨은 CPU-GPU 동기화가 모두 그래프를 끊는다.
• 커널을 직접 쓸 때는 스레드가 메모리에서 서로 붙어 있는 값을 읽도록 배치해야 한다. 한 스레드에 연속 구간을 통째로 맡기는 직관적인 분할이 오히려 느리다.
• 발표자들은 최적화 순서를 순수 파이토치 → 트라이톤 → CUDA로 제시했다. 커스텀 커널은 초보자에게 일주일, 숙련자에게 반나절이 드는 비용이므로 마지막 수단이다.

자주 묻는 질문

관련 AI 소식