PyTorch torch.compile CUDA 스트림 지원: 비동기 실행과 활성값 오프로딩으로 메모리 줄이기

핵심 메시지

• 스트림은 GPU 위의 서로 다른 실행 큐다. 여러 커널을 동시에 돌려 연산과 통신, 연산과 메모리 전송을 겹칠 수 있게 해주며 스트림 사이의 순서 맞추기는 이벤트로 한다.
• 지금까지 torch.compile은 스트림 컨텍스트를 만나면 그래프를 끊어버렸다. 우회하려면 스트림 조작을 커스텀 연산으로 감싸야 했고, 그만큼 코드가 복잡해졌다.
• 이번 작업은 컴파일 경로 세 곳을 모두 손봤다. Dynamo는 심볼릭 스택으로 현재 스트림을 추적하고, AOT Autograd는 역전파 노드에 같은 스트림을 물려주며 동기화를 끼워 넣고 재정렬을 막으며, Inductor는 서로 다른 스트림에 배정된 커널의 융합을 금지한다.
• 활용처로는 통신과 연산을 겹치는 마이크로배치 오버랩, 그리고 순전파 활성값을 CPU로 비동기 전송해 두었다가 역전파 직전에 미리 가져오는 활성값 오프로딩이 제시됐다.
• 실측에서는 모델이 커질수록 런타임 손해가 줄고 메모리 절감폭이 커졌다. 작은 모델은 절감이 10% 남짓인데 런타임 손해가 9% 정도라 이득이 적다. 이 기능은 2.12 릴리스에 포함될 예정이다.

쉽게 이해하기

메타의 소프트웨어 엔지니어 마이클 라조스가 torch.compile에 CUDA 스트림 지원을 넣은 작업을 발표했다. 스트림은 GPU 위의 별도 실행 큐로 이해하면 된다. 여러 스트림에 커널을 나눠 올리면 동시에 실행되므로 연산과 통신을 겹치거나 데이터 전송 시간을 감출 수 있다. 동시에 도는 만큼 순서를 맞출 장치가 필요한데 그 역할을 이벤트가 한다. 한 스트림의 진행 상황을 이벤트로 기록해 두면 다른 스트림이 그 이벤트를 기다렸다가 이어서 실행된다. 덜 알려진 쓰임새로는 같은 호스트에 붙은 여러 GPU 사이의 동기화도 가능하다. 문제는 지금까지 torch.compile이 스트림 컨텍스트를 만나면 그래프를 끊어버렸다는 점이다. 스트림을 쓰는 모델일수록 손해가 컸고, 스트림 조작을 커스텀 연산으로 감싸는 우회법은 있었지만 코드가 복잡해지고 사용 경험이 나빠졌다. 이번 작업의 목표는 eager 모드에서 쓰던 익숙한 스트림 API를 그대로 컴파일 안에서도 쓸 수 있게 만드는 것이다.

구현은 컴파일 경로 세 단계에 모두 걸쳐 있다. 그래프를 뜨는 Dynamo에서는 사용자 코드의 스트림 컨텍스트 매니저 중첩을 심볼릭 스택으로 따라가며 각 노드에 어느 스트림에서 돌지를 표시한다. 스트림과 이벤트 객체를 그래프 인자로 넘기는 일은 쉽지 않았다. AOT Autograd가 기울기를 계산할 수 있는 텐서만 인자로 받기 때문인데, 대대적인 개편 대신 전역 객체 테이블을 두고 그래프 안에서는 인덱스로 조회하는 방식을 택했다. 대신 이 조회를 이벤트 기록·대기 같은 익숙한 연산 뒤로 숨겨 그래프 가독성을 지켰다.

역전파를 만드는 AOT Autograd 단계에서는 순전파에서 사용자가 정한 병렬화 방식을 그대로 물려주는 것을 원칙으로 삼았다. eager 모드와 최대한 똑같이 동작하게 하려는 것이다. 순전파 노드의 스트림 배정을 대응하는 역전파 노드로 전파한 뒤, 어떤 커널의 입력이 다른 스트림에 있으면 그 자리에 동기화를 끼워 넣는다. 순전파에 대응 노드가 없는 기울기 누적 노드는 첫 사용자의 스트림을 따라간다. 또 컴파일러가 최적화를 위해 노드 순서를 바꾸다가 이벤트 경계를 넘어서면 경쟁 상태가 생기므로, 경계를 가로지르는 텐서를 가짜 의존성으로 명시해 재정렬을 막았다.

까다로운 경계 사례도 그대로 공개됐다. 컴파일러는 제자리 연산을 바깥 연산으로 바꾸고 변경 내용을 그래프 맨 뒤에서 복사해 반영하는데, 이 복사가 이벤트 기록 뒤로 밀리면 바깥에서 기다리던 쪽이 옛 값을 보게 된다. 이 경우를 조용히 넘기는 대신 정확히 찾아내 컴파일 에러를 내고 사용자가 코드를 고치도록 했다. 메모리 쪽에도 함정이 있다. CUDA 캐싱 할당자는 스트림마다 메모리 풀을 따로 두고 참조가 사라지면 곧바로 재사용하는데, 다른 스트림이 아직 그 텐서를 쓰는 중이라면 문제가 된다. 그래서 이 경우에는 사용자가 텐서를 지운 위치를 존중해 동기화 연산을 그래프에 끼워 넣는다. 코드를 생성하는 Inductor에서는 융합 규칙이 바뀐다. 평소라면 여러 연산을 한 커널로 합치지만, 사용자가 서로 다른 스트림에 올린 커널은 병렬로 돌리겠다는 의도이므로 합쳐서는 안 된다. 버퍼 재사용도 그 텐서가 만들어진 스트림 안에서만 일어나도록 제한했다. 활용 사례로는 분산 학습에서 흔한 all-reduce나 all-gather 같은 통신을 옆 스트림으로 빼 연산과 겹치는 방식과, 순전파 활성값을 CPU로 비동기 전송했다가 역전파 직전에 미리 되가져오는 활성값 오프로딩이 소개됐다.

트랜스포머 모델로 측정한 결과는 규모에 따라 갈렸다. 작은 모델에서는 피크 메모리 절감이 10% 정도인데 런타임 손해가 9% 가까이 나서 쓸 이유가 크지 않다. 모델이 커질수록 전송을 감춰 줄 연산량이 늘어 런타임 손해는 줄고, 활성값이 커진 만큼 메모리 절감폭은 커진다. 질의응답에서는 이 기능이 CUDA 전용이 아니라 디바이스에 종속되지 않는 스트림 위에 올라가 있어 AMD에서도 쓸 수 있다는 답이 나왔다. 다만 노드가 여러 대로 늘어나면 이벤트가 아니라 실제 통신을 써야 한다.

주요 인사이트

• 컴파일러가 사용자의 병렬화 의도를 '보존해야 할 의미'로 취급한 점이 흥미롭다. 서로 다른 스트림의 커널을 합치지 않는 것은 성능을 포기하는 선택처럼 보이지만, 사용자가 병렬로 돌리라고 지시한 것을 컴파일러가 되돌리지 않겠다는 원칙이다.
• 설계 기준이 반복해서 'eager 모드와 똑같이'로 돌아온다. 역전파의 스트림 배정도, 대응 노드가 없는 기울기 누적 노드의 처리도 모두 eager가 하던 방식을 따랐다.
• 안전하지 않은 경우를 조용히 넘기지 않고 컴파일 에러로 드러낸 선택이 눈에 띈다. 발표자는 이를 근본 해결이 아닌 회피라고 인정하면서도, AOT Autograd를 크게 뜯어고치는 대신 택한 절충임을 분명히 했다.
• 메모리 재사용 함정은 스트림을 쓰는 사람이 흔히 놓치는 부분이다. 캐싱 할당자가 스트림별 풀을 쓰기 때문에, 옆 스트림이 아직 쓰는 텐서를 지우면 그 메모리가 곧바로 다른 용도로 재사용될 수 있다.
• 활성값 오프로딩의 이득은 결국 '전송 시간을 가릴 만큼 연산이 충분한가'로 결정된다. 그래서 작은 모델에서는 손해가 이득을 거의 상쇄하고, 큰 모델에서만 순수한 메모리 절감에 가까워진다.

자주 묻는 질문

관련 AI 소식