추천 모델의 학습 단계를 GPU와 TPU로 옮기는 RecIR

추천 모델의 순전파·역전파·옵티마이저 갱신을 한 그래프로 묶어 GPU와 TPU로 옮기는 초기 컴파일러의 구조와 지원 범위를 소개합니다.

쉽게 말하면비유하자면 설계도 한 장으로 두 공장에서 물건을 만들게 해 주는 변환기다. 추천 모델의 학습 한 걸음을 통째로 도면 한 장에 옮긴 뒤, GPU와 TPU라는 서로 다른 계산 장치에 맞는 작업 지시서로 각각 바꾼다. 아직 시험판이고 DLRMv3라는 모델의 학습만 끝까지 지원한다.

왜 봐야 하나

DLRMv3 계열 추천 모델을 GPU와 TPU 양쪽에서 학습시켜 보려는 팀에게 참고가 되는 초기 도구다. 모델 코드는 그대로 두고 실행 방식만 바꾸는 구조를 볼 수 있다. 다만 초기 시험판(알파)이고 처리 속도 수치는 공개되지 않았다. 성능은 직접 확인해야 한다.

이 글은 2026년 10월 29일 올라온 RecIR의 새 한국어 해설이다. RecIR은 Meta의 추천 모델 학습용 컴파일러로, 한 번 작성한 PyTorch 학습 단계를 그래프로 묶어 GPU나 TPU용으로 변환한다. 첫 지원 경로는 DLRMv3 훈련이며, 알파 단계로 추론은 실험적이다. 공식 문서는 실행 경로와 수치 대조 도구를 제공하지만 GPU·TPU 처리량 비교는 싣지 않았다. 저장소의 최초 공개일은 확인되지 않았으며, 이 해설의 게시일을 새 논문의 발표일로 볼 수는 없다.

알아낸 것

• 일반 PyTorch 모델의 학습 단계 전체를 FX 그래프로 캡처하고 실행 전략을 모델 코드 밖에서 적용합니다.
• 첫 지원 경로는 DLRMv3 훈련이며 GPU는 Inductor·TorchRec, TPU는 torchax·JAX/XLA를 사용합니다.
• CPU 설치 점검과 수치 대조 도구를 제공하지만 GPU·TPU 처리량 수치는 공개 문서에 없습니다.

한계와 주의할 점

• RecIR은 아직 초기(알파) 단계다. 처음부터 끝까지 지원하는 것은 DLRMv3 모델의 학습 하나뿐이고, 학습된 모델로 결과를 내는 추론은 실험 단계여서 실제 서비스용으로 보장되지 않는다.
• GPU와 TPU에서 얼마나 빠른지 보여 주는 처리량 수치는 공개되지 않았다. 글에 나오는 18.3ms는 다른 프로젝트인 TorchTPU의 결과이며 RecIR의 성능이 아니다.
• CPU만으로 하는 점검은 설치와 학습 흐름이 도는지만 확인한다. 실제 TPU에서의 동작과 속도는 따로 확인해야 하고, 한 번에 처리하는 데이터 묶음 크기나 장치 수가 바뀌면 다시 변환(컴파일)해야 한다.

용어 풀이

make_fx

PyTorch 함수를 추적해 FX 그래프를 만드는 도구다. RecIR은 순전파, 역전파, 옵티마이저 갱신까지 포함한 전체 학습 단계를 make_fx로 캡처한다.

HybridParallel

RecIR의 병렬 모드 중 하나로, 밀집 매개변수에는 데이터 병렬을 쓰고 큰 임베딩 테이블은 EmbeddingShardingPlan을 따라 나눈다. GPU DLRMv3 예제가 이 모드를 사용한다.

GSPMD

XLA의 자동 분할 방식으로, 텐서 샤딩을 파티션 명세로 표현하고 표준 밀집 연산의 집합 통신을 XLA가 자동으로 삽입한다. RecIR의 TPU 경로가 이 모드를 사용한다.

Pallas 링 시프트

TPU에서 임베딩 테이블을 행 단위로 샤딩하고, 전역 테이블을 만들지 않고 ID와 결과를 칩 사이에서 돌려 보내는 커널이다. RecIR의 TPU 기본 임베딩 백엔드가 이 방식을 사용한다.

토큰 예산 배치

가변 길이 입력을 고정 크기 블록으로 맞추는 전략. 남는 자리는 학습 가중치가 없는 더미 행으로 채우며, 예산이 실제 이력보다 작으면 긴 이력의 오래된 이벤트를 잘라낸다.

수치 동등성

즉시 실행 PyTorch와 컴파일 결과를 같은 초기 상태와 결정적 배치에서 비교해 손실, 출력, 그래디언트, 매개변수, 옵티마이저 상태가 허용 오차 내에 있는지 확인하는 검증 방식이다.

더 깊이 읽기 읽는 데 약 8분

10월 30일(금) 정리 · AI가 쓰고 원문과 대조