파이토치 사전 컴파일과 SPMD 타입: 분산 학습 병목과 잘못된 그래디언트를 줄이는 새 기능

핵심 메시지

• torch.compile은 실행 시점에 컴파일하는 방식이라 분산 학습에서 모든 노드가 각자 같은 작업을 반복하고, 작업 재시작이 느려지는 원인이 된다.
• 랭크마다 입력 형상이 달라 한 노드만 재컴파일하면 나머지가 기다리다 통신 타임아웃이 나고, 컴파일러가 랭크별로 다른 선택을 하면 교착 상태로 이어질 수 있다.
• 사전 컴파일은 한 노드에서 만든 결과물을 모든 랭크에 배포해, 실행 중 컴파일을 없애고 모든 랭크가 같은 코드를 돌리도록 보장한다.
• 파이토치에서 컬렉티브 연산에는 미분이 정의돼 있지 않아, 텐서 병렬 구현에서 특정 연산을 빠뜨리면 오류 없이 조용히 잘못된 그래디언트를 얻는다.
• SPMD 타입은 각 랭크의 로컬 관점으로 짠 코드에 타입 주석을 달아, 빠뜨린 컬렉티브를 타입 검사기가 위치까지 짚어 알려주게 만든다.

쉽게 이해하기

파리에서 처음 열린 파이토치 유럽 콘퍼런스의 키노트는 프로젝트 현황 소개로 시작했다. 발표자가 인용한 2026년 3월 기준 수치에 따르면 코어 라이브러리에 한 해 동안 들어온 기여가 1만 5,000건이고, 상위 머신러닝 학회의 연구 프로젝트 열 중 아홉, 오픈소스 머신러닝 프로젝트의 90%가 파이토치를 쓴다. 분산 학습용 기능인 디텐서와 디바이스 메시를 채택한 오픈소스 프로젝트도 1,000곳을 넘었다.

키노트가 집중한 첫 주제는 사전 컴파일이다. 현재 torch.compile은 함수를 처음 실행할 때 입력을 보고 컴파일하는 방식인데, 분산 학습에서는 모든 노드가 같은 파이썬 프로그램을 돌리므로 결국 모두가 같은 컴파일을 반복한다. 캐시가 유효한지 확인하는 데도 추적 과정을 전부 거쳐야 해서 작업을 재시작할 때마다 시간이 든다.

더 성가신 건 정확성 쪽 위험이다. 추천 시스템처럼 랭크마다 입력 개수가 달라지는 경우, 운 나쁜 한 노드만 재컴파일에 들어가면 다른 노드들이 기다리다 통신 타임아웃이 난다. 그래프가 랭크마다 다르면 컴파일러가 서로 다른 선택을 하면서 통신 순서를 어긋나게 바꿔 교착에 빠질 수도 있다. 해법은 단순하다. 한 노드에서 미리 컴파일하고 그 결과물을 나눠 주면 실행 중 컴파일이 사라지고 모든 랭크가 같은 코드를 돌린다.

구현이 늦어진 이유도 설명됐다. 사전에 컴파일하려면 모델 전체를 그래프로 잡아야 하는데, 이해하지 못하는 부분을 건너뛰고 나머지만 컴파일하는 유연함이야말로 torch.compile을 쉽게 만들어 준 특성이었다. 또 한 번 컴파일한 결과를 어디서나 쓰려면 코드가 특정 랭크를 들여다보지 않아야 한다. 현재는 torchtitan에 전체 모델 캡처를 전제로 한 그래프 트레이너가 들어갔고, 라마 3에서 비트 단위로 같은 결과와 동일한 성능이 확인됐다. 발표자는 아직 매우 초기라고 못 박으며, 마지막 랭크의 데이터 개수가 다르거나 특정 랭크에서만 로그를 남기는 흔한 패턴도 지원해야 한다고 덧붙였다.

두 번째 주제인 SPMD 타입은 다른 각도의 문제를 겨냥한다. 파이토치는 어떤 프로그램이든 미분할 수 있다는 보장을 내세우지만 컬렉티브 연산만은 예외다. 메가트론식 텐서 병렬 구현에는 순전파에서는 아무것도 하지 않고 역전파에서만 올 리듀스를 수행하는 연산이 등장하는데, 이걸 빠뜨리면 오류 없이 틀린 그래디언트가 나온다. SPMD 타입은 텐서가 랭크들 사이에서 어떻게 흩어져 있는지를 타입으로 적어 두고, 필요한 컬렉티브가 빠졌을 때 타입 검사기가 잡아내게 하려는 시도다.

주요 인사이트

• 디텐서가 이미 있는데 왜 또 만드느냐는 질문에 발표자는 겨냥하는 문제가 다르다고 답한다. 디텐서는 병렬 처리를 잊고 전체 텐서를 다루는 전역 관점을 주지만, 일부 숙련 사용자는 랭크에서 실제로 무엇이 도는지를 직접 보며 코드를 짜고 싶어 한다는 것이다. 이거 연산의 오버헤드도 이유로 꼽혔다.
• 타입 구분에서 핵심은 replicate와 invariant의 차이다. 둘 다 순전파에서는 모든 랭크가 같은 값을 갖지만, replicate는 역전파에 아직 처리되지 않은 리덕션이 남아 있다는 뜻이다. DDP나 FSDP가 훅으로 뒤에서 처리해 주는 바로 그 올 리듀스가 여기에 해당한다.
• 시퀀스 병렬을 켜서 입력을 텐서 병렬 축으로 쪼개면 invariant와 varying이 섞이며 타입 오류가 난다. 오류 메시지가 '여기서 invariant를 replicate로 바꾸는 컬렉티브가 필요하다'고 제안해, 빠진 연산의 위치까지 알려주는 식이다.
• 이 타입 체계는 JAX의 샤딩 타입 설계를 거의 그대로 가져와 파이토치에 맞게 기본값만 조정한 것이라고 발표자가 직접 밝혔다. 타입 검사는 선택 사항이라 유닛 테스트에서만 돌려도 되고, 런타임 오버헤드는 없다.

자주 묻는 질문

관련 AI 소식