RL 지연시간을 줄이는 PyTorch Custom Ops의 역할

강화학습(RL) 학습 루프의 병목은 모델 연산에만 있지 않습니다. 환경 실행, 샘플 전송, 정책 업데이트 사이의 동기화가 GPU를 기다리게 만들 수 있습니다. 이때 필요한 질문은 “커널을 얼마나 빠르게 만들 것인가”를 넘어섭니다. 더 중요한 것은 “Python 밖에서 만든 비동기 흐름을 PyTorch 학습 루프와 어떻게 안전하게 연결할 것인가”입니다.

Custom Op은 빠른 코드가 아니라 PyTorch가 이해하는 코드다

PyTorch Custom Op은 C/CUDA나 외부 라이브러리로 만든 연산을 PyTorch 연산처럼 등록하는 방법입니다.¹ 네이티브 코드에 텐서 포인터를 넘겨 직접 호출할 수도 있지만, 그렇게 하면 torch.compile, autograd, vmap 같은 PyTorch 하위 시스템과 자연스럽게 맞물리지 않을 수 있습니다.

핵심은 “등록”입니다. 커스텀 커널을 학습 루프에 넣으려면 PyTorch에 두 가지를 알려야 합니다.

• 이 연산을 어떻게 실행하는가

• 이 연산을 어떻게 미분하는가

즉 Custom Op은 단순한 성능 최적화 기법이 아닙니다. 최적화한 네이티브 코드를 PyTorch의 실행, 컴파일, 미분 체계 안으로 들여오는 인터페이스입니다.²

RL에서 성능을 좌우하는 것은 겹치기다

RL에서는 rollout과 training이 번갈아 실행되기 쉽습니다. 이 구조에서는 한쪽이 작업하는 동안 다른 쪽이 쉬게 되고, 전체 처리량은 동기화 지점에 묶입니다.

VERL의 fully async PPO 문서는 Trainer와 Rollouter를 분리해 샘플 생성과 학습을 겹치는 방식을 설명합니다. 이 접근은 Qwen2.5-7B, 128 GPU 학습에서 2.35~2.67배 성능 향상을 보고했습니다.³

같은 문제는 LLM 추론에서도 나타납니다. CPU가 배치를 준비하고 GPU가 계산하는 과정이 번갈아 실행되면 GPU 유휴 시간이 생깁니다. Hugging Face의 분석은 한 실험에서 전체 시간의 24%가 GPU가 CPU를 기다리는 시간이었다고 설명합니다.⁴

RL의 환경 처리와 샘플 전송도 비슷한 병목을 만납니다. 따라서 C 기반 비동기 전송이나 환경 실행 코드를 PyTorch Custom Op으로 감싸는 전략은 의미가 있습니다. Python 루프에서 발생하는 동기화 비용을 줄이고, 데이터 이동과 GPU 연산을 더 잘 겹치게 만들 수 있기 때문입니다.

torch.compile과 autograd는 별도로 설계해야 한다

Custom Op을 학습 루프에 넣을 때 가장 흔한 함정은 “네이티브 구현만 있으면 충분하다”고 생각하는 것입니다. 실제로는 torch.compile과 autograd가 그 연산을 이해할 수 있어야 합니다.

torch.compile 아래에서 비합성 커스텀 커널은 fake 구현이 필요합니다.² fake 구현은 실제 값을 계산하지 않고 shape, dtype, device 같은 메타정보를 알려주는 구현입니다. 컴파일러가 그래프를 분석하려면 실제 실행 없이도 텐서의 형태와 속성을 추론할 수 있어야 하기 때문입니다.

학습에 사용한다면 backward 구현도 필요합니다. 비합성 커널은 autograd가 미분 규칙을 자동으로 알 수 없으므로, PyTorch에 역전파 방법을 명시해야 합니다.²

반대로 연산의 의미를 Aten 연산 조합으로 표현할 수 있다면 CompositeImplicitAutograd로 등록할 수 있습니다. 다만 이 방식도 만능은 아닙니다. 디바이스별 C++ 구현을 추가했는데 autograd 커널이 등록되지 않으면 해당 디바이스의 backward가 실패할 수 있고, "autograd kernel was not registered" 경고가 발생할 수 있습니다.²

적용 기준: 병목이 어디에 있는지 먼저 확인하라

RL 지연시간을 줄이기 위해 모든 코드를 Custom Op으로 옮길 필요는 없습니다. 먼저 병목이 어디에서 생기는지 확인해야 합니다.

Custom Op을 검토할 만한 경우는 다음과 같습니다.

• Python 루프의 환경 실행 또는 전송 로직이 자주 동기화를 만든다

• CPU 작업과 GPU 연산이 충분히 겹치지 않는다

• 네이티브 비동기 처리로 옮겼을 때 처리량 개선 여지가 크다

• 해당 연산을 torch.compile 또는 autograd가 포함된 학습 루프 안에서 계속 사용해야 한다

이 조건에 해당한다면 Custom Op은 단순한 “빠른 커널”이 아니라, 비동기 실행을 PyTorch 생태계 안에 통합하는 실용적인 선택지가 됩니다.

결론은 명확합니다. RL 지연시간을 줄이려면 Python 밖으로 코드를 빼는 것만으로는 부족합니다. PyTorch가 이해할 수 있는 실행 경로, fake 구현, backward 경로까지 함께 설계해야 합니다. 병목이 환경 실행·샘플 전송·동기화에 있다면, 네이티브 비동기 처리와 Custom Op 등록을 함께 검토할 가치가 있습니다.

¹PyTorch Custom Operators -- PyTorch Tutorials

²Custom Ops Under torch.compile: autograd.Function vs torch.library.custom_op

³Recipe: Fully Async Policy Trainer

⁴Unlocking asynchronicity in continuous batching