GPU 활용률 끌어올리는 스케줄링: 서울대 서지원 교수의 딥러닝 학습·LLM 추론 최적화 강연

핵심 메시지

• 모델을 학습시키는 동안 GPU의 병렬 연산 자원은 절반 가까이 놀고 있는 경우가 흔하며, 원인은 커널 실행 준비 비용과 연산 말미의 자원 공백, GPU 간 동기화 대기, 파이프라인 공백으로 나뉜다.
• 서로 달라 보이는 이 문제들은 하나의 스케줄링 최적화 문제로 정식화할 수 있고, 그 형태는 산업공학의 잡숍 스케줄링과 같아 대부분 NP-난해에 속한다.
• 역전파에서 가중치에 대한 그래디언트는 앞쪽 레이어 계산에 쓰이지 않는다. 이 사실에 착안해 계산 순서를 레이어 역순에서 풀어 준 것이 아웃오브오더 백프롭이다.
• 이 원리를 단일 GPU, 데이터 병렬, 파이프라인 병렬 학습에 각각 적용해 당시 최고 성능 시스템 대비 학습 처리량을 최대 두 배까지 끌어올렸다.
• LLM 추론에서는 토큰 하나를 만드는 지연시간에 제약이 주어졌을 때 그 제약을 지키면서 처리량을 최대화하는 문제로 바꿔 풀었고, 제어 변수의 단조성을 이용해 탐색 공간을 잘라냈다.

쉽게 이해하기

서울대학교 AI연구원 콜로퀴움에서 서지원 교수가 머신러닝 학습과 대형 언어 모델 추론의 시스템 최적화를 주제로 강연했다. 발표자는 서울대 첨단융합학부와 전기정보공학부에 소속돼 있으며, 연구실에서는 머신러닝 시스템과 빅데이터 시스템을 함께 다룬다. 이날 다룬 두 축은 학습 스케줄링과 추론 스케줄링으로, 각각 시스템 분야 학회에 발표된 연구다.

출발점은 단순한 관찰이다. GPU 한 장으로 학습하든 여러 장으로 데이터 병렬이나 파이프라인 병렬을 쓰든, GPU 안의 병렬 연산 장치가 절반 가까이 놀고 있다는 것이다. 단일 GPU에서는 커널을 실제로 계산하는 시간보다 CPU가 그 커널을 준비해 내려보내는 시간이 더 걸리는 구간이 생기는데, 텐서를 할당하고 반납하는 비용과 연산별 사전 준비, 프레임워크가 커널마다 수행하는 방대한 오류 검사가 겹치기 때문이다. 여기에 커널 끝자락에서 스트리밍 멀티프로세서 대부분이 노는 이른바 테일 이펙트가 더해진다.

여러 장을 쓸 때는 원인이 달라진다. 데이터 병렬에서는 각 GPU가 계산한 그래디언트를 평균 내는 동기화 시간이 계산 시간보다 길어져, 동기화가 끝날 때까지 GPU가 대기한다. 파이프라인 병렬에서는 앞 구간 계산이 끝나야 뒤 구간을 계산할 수 있는 의존 관계 때문에 버블이 생기고, 발표자는 이 구간에서 30~40% 가량의 시간이 낭비된다고 설명했다. 기존 대응책들은 상황별로 따로 나왔지만, 발표자는 이 모두가 하나의 최적화 문제로 묶인다고 봤다.

핵심 아이디어인 아웃오브오더 백프롭은 역전파 계산을 두 갈래로 쪼개는 데서 나온다. 각 레이어의 역전파는 입력 뉴런에 대한 그래디언트와 가중치에 대한 그래디언트를 계산하는데, 앞쪽 레이어로 이어지는 것은 앞의 것뿐이고 가중치 그래디언트는 그다음 반복 전까지만 끝내면 된다. 그래서 임계 경로에 있는 계산은 메인 스트림에서 순서대로 돌리고 가중치 그래디언트와 파라미터 갱신은 서브 스트림으로 빼, 비어 있는 연산 자원을 메운다. 파이프라인 병렬에서는 연속된 레이어를 한 GPU에 몰아 주는 관행을 깨고 레이어를 GPU에 번갈아 배치해 버블을 줄여 1.5~1.6배를 얻었고, 전체적으로는 당시 최고 성능 시스템 대비 최대 두 배의 학습 처리량을 보였다.

후반부는 LLM 추론이다. 추론은 입력 전체를 한 번에 처리하는 프리필과 토큰을 하나씩 만들어 되먹이는 디코딩으로 나뉘는데, 두 단계의 연산량 차이가 커서 파이프라인 버블이 생기고 먼저 끝난 요청이 빠지면서 배치가 작아져 GPU를 못 채운다. 발표자는 애플리케이션마다 토큰 하나를 내놓는 데 허용되는 지연시간이 다르다는 점에 착안해, 그 제약을 지키면서 처리량을 최대화하는 문제로 다시 세웠다. 레이어를 라운드로빈으로 나눠 배치 크기를 크게 유지하는 방식과, 프리필용 GPU와 디코딩용 GPU를 아예 분리해 버블을 없애는 방식을 함께 제시했다.

주요 인사이트

• 이 최적화 문제는 1990년대 전후 마이크로아키텍처와 컴파일러가 다뤘던 명령어 수준 병렬성 문제와 구조가 같다. 당시 축적된 기법을 변형해 딥러닝 스케줄링에 가져올 여지가 있다는 지적이 이 연구의 출발점이다.
• 제안 기법은 계산 의존 관계를 지키므로 학습의 의미를 바꾸지 않는다. 그래서 끝까지 학습시키지 않고 처리량만 측정해도 유효한 비교가 되는데, 역전파가 끝나기 전에 다음 반복을 시작하는 부류의 기법과는 성격이 다르다.
• 프리필과 디코딩을 서로 다른 GPU에 분리하는 설계는 발표 시점에는 새로운 제안이었지만, 이후 여러 기업과 대학이 유사한 논문을 내면서 표준에 가까워졌다.
• 다만 분리 방식은 디코더 전용 모델에서 모델을 두 벌 올려야 해 메모리 압박이 커진다. 실제로 340B급 큰 모델 실험에서는 메모리 부족 오류가 나 이 방식을 쓰지 못했다.
• 발표자는 학계의 자원 제약도 솔직히 밝혔다. 실험 상당 부분을 AWS 클라우드에서 돌렸고 비용이 1억원 가까이 들었으며, 고용량 GPU는 비용을 지불해도 곧바로 할당받지 못해 두 번째 논문 제출 때에야 쓸 수 있었다고 한다.

자주 묻는 질문

관련 AI 소식