🤖 machine learning

rl-triton: High-Performance Triton GPU Kernels for Reinforcement Learning Credit Assignment

이 논문은 Triton으로 구현된 통합 결합 연산 스캔(unified associative scan) 프레임워크를 활용하여 GPU 상에서 7가지의 서로 다른 강화 학습 신용 할당 알고리즘을 가속화함으로써, 메모리 오버헤드를 줄이고 $O(\log T)$ 병렬 계산을 가능하게 하여 벡터화된 베이스라인 대비 1.6~5.70배의 속도 향상을 달성하는 오픈 소스 라이브러리인 rl-triton을 소개한다.

원저자: Lars Simon Zehnder

게시일 2026-08-20

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Lars Simon Zehnder

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

인공지능의 세계에는 컴퓨터에게 어떻게 하면 좋은 결정을 내릴 수 있는지 가르치기 위한 끊임없는 투쟁이 존재합니다. 로봇이 걷는 법을 배우거나 프로그램이 게임을 배우는 상황을 상상해 보십시오. 성능을 개선하기 위해, 시스템은 어떤 구체적인 행동이 성공으로 이어졌고 어떤 행동이 실패로 이어졌는지 파악해야 합니다. 이 과정을 '신용 할당(credit assignment)'이라고 부릅니다. 이는 일련의 사건들을 되돌아보며 "이 단계는 좋았다" 또는 "저 단계는 나빴다"라고 결정하여, 시스템이 미래의 행동을 조정할 수 있도록 하는 작업입니다. 로봇은 세상을 탐험하거나 다음 행동을 결정하기 위해 복잡한 계산을 수행하는 데 대부분의 시간을 보낼 수도 있지만, 실수를 통해 배워야 하는 순간에는 특정한 유형의 수학을 수행해야 합니다. 이 수학은 긴 단계의 목록을 살펴보고 그들 사이의 점들을 연결하는 작업을 포함하며, 여기서 한 단계의 가치는 그 다음에 오는 단계에 의존합니다. 오랫동안 GPU라고 불리는 강력한 컴퓨터 칩에서 이 수학을 수행하는 것은 느렸는데, 하드웨어는 한 번에 여러 페이지를 읽을 수 있는 능력이 있음에도 불구하고 컴퓨터가 마치 책의 페이지를 한 장씩 넘기듯 목록의 단계를 하나씩 순차적으로 처리해야 했기 때문입니다.

Lars Simon Zehnder라는 연구자는 이 병목 현상을 해결하는 rl-triton이라는 새로운 도구를 개발했습니다. 이 도구는 강화 학습에서의 신용 할당 작업을 위해 특별히 설계된 매우 효율적인 컴퓨터 명령어들의 모음입니다. 컴퓨터가 단계의 목록을 느린 순차적 사슬 방식으로 처리하도록 강요하는 대신, 이 새로운 방법은 수천 개의 단계를 동시에 계산할 수 있도록 작업 방식을 재구성합니다. 핵심 아이디어는 전체 사건의 시퀀스를 하나의 통합된 수학적 구조로 취급하여, 이를 분해하고 병렬로 해결할 수 있게 만드는 것입니다. 이렇게 함으로써 컴퓨터는 수천 개의 서로 다른 시나리오가 동시에 발생하는 상황에서도 이전보다 훨씬 짧은 시간 안에 계산을 마칠 수 있습니다.

연구진은 이 새로운 접근 방식을 현재 분야에서 사용되는 표준 방식들과 비교 테스트했습니다. 그들은 수천 개의 환경이 동시에 시뮬레이션되는 가장 일반적이고 까다로운 시나리오에서 이 새로운 도구가 현저히 빠르다는 것을 발견했습니다. 어떤 경우에는 이전의 최고 방법보다 거의 6배 더 빠르게 작업을 완료했습니다. 이러한 속도 향상은 데이터가 컴퓨터 메모리를 통과하는 방식의 영리한 변화에서 비롯되었습니다. 기존 방식에서는 컴퓨터가 시퀀스의 매 단계마다 메인 메모리 뱅크에서 데이터를 계속 가져와야 했기에 교통 체증이 발생했습니다. 새로운 방법은 데이터를 계산 엔진 가까이에 유지하여, 컴퓨터가 이러한 지속적인 멈춤 없이 전체 시퀀스를 처리할 수 있도록 합니다. 이는 시스템이 수백 개의 단계와 함께 수천 개의 시뮬레이션을 병렬로 실행할 수 있는 현대적인 AI 훈련에서 특히 중요합니다.

논문은 이 방식이 동일한 근본적인 수학적 패턴을 공유하는 7가지 서로 다른 학습 알고리즘에 어떻게 적용되는지 상세히 설명합니다. 이 새로운 도구는 이 모든 것을 하나의 통합된 프레임워크로 처리합니다. 또한 에피소드가 갑작스럽게 종료되거나 시뮬레이션이 중단되는 경우와 같은 실제 데이터의 복잡한 현실을 세심하게 다룹니다. 연구진은 이 방법이 경계 조건을 올바르게 처리하여, 학습 신호가 적절한 곳에서 멈추고 한 시나리오에서 다른 시나리오로 실수로 흘러 들어가지 않음을 증명했습니다. 그들은 새로운 도구를 느리고 구식인 방식 및 표준 프로그래밍 도구를 사용하는 더 현대적이고 최적화된 버전 모두와 비교함으로써 결과를 검증했습니다. 새로운 도구는 두 방식 모두에서 일관되게 우수한 성능을 보였으며, 이는 속도 향상이 단순히 더 나은 코딩 기술의 결과가 아니라 실제적인 것임을 보여주었습니다.

가장 흥러운 발견 중 하나는 문제의 크기에 따라 속도 이점이 어떻게 변하는가 하는 점입니다. 단계의 시퀀스가 짧을 때는 새로운 도구가 여전히 더 빠르지만 그 차이는 작습니다. 그러나 시퀀스가 길어질수록 그 이점은 커집니다. 이는 기존 방식은 목록이 길어질수록 메모리 호출 과정을 훨씬 더 많이 반복해야 하는 반면, 새로운 방식은 훨씬 더 효율적으로 확장되기 때문입니다. 연구진은 이것이 AI 에이전트의 전체 훈련 과정에 어떤 영향을 미치는지도 살펴보았습니다. 그들은 신용 할당 단계 자체가 훨씬 빨라졌음에도 불구하고 전체 훈련 속도의 향상은 때때로 완만하다는 것을 발견했습니다. 이는 신용 할당이 전체 훈련 파이프라인의 일부일 뿐이며, 만약 나머지 과정이 느리다면 단 하나의 부분만 빠르게 만든다고 해서 전체가 극적으로 빨라지지는 않기 때문입니다. 하지만 신용 할당 단계가 전체 시간에서 차지하는 비중이 큰 특정 설정에서는 전체 훈련 속도가 눈에 띄게 개선되었습니다.

이 연구는 몇 가지 한계점도 강조합니다. 매우 긴 시퀀스의 경우, Retrace라고 불리는 특정 유형의 알고리즘은 컴퓨터 칩이 특정 유형의 빠른 저장 공간을 다 써버리는 하드웨어 제약에 부딪혀 속도가 저하됩니다. 연구진은 이 문제를 식별하고 이것이 설계상의 알려진 트레이드오프임을 언급했습니다. 또한 현재의 도구가 표준 데이터 형식에 가장 잘 작동하며, 일부 특화된 변형에는 추가적인 개발이 필요할 수 있다고 언급했습니다. 이러한 제한에도 불구하고, 이 논문은 AI 훈련의 고질적인 문제에 대한 명확하고 실용적인 해결책을 제시합니다. 순차적인 단계별 계산을 병렬적인 동시 계산으로 전환함으로써, 연구진은 강화 학습을 훨씬 더 효율적으로 만드는 것이 가능하다는 것을 보여주었습니다. 이러한 효율성은 AI 시스템이 점점 더 커지고 복잡해짐에 따라, 더 짧은 시간 안에 방대한 양의 데이터로부터 학습해야 하는 상황에서 매우 중요합니다. 이 도구는 이제 다른 사람들이 사용할 수 있도록 공개되어, 학습의 근본적인 방식을 바꾸지 않고도 지능형 시스템의 훈련을 가속화할 수 있는 방법을 제공합니다.

기술 요약: rl-triton

문제 정의

심층 강화 학습(Deep Reinforcement Learning, RL) 훈련 파이프라인은 매 업데이트 단계마다 어드밴티지(advantages), 가치 타겟(value targets), 리턴 추정치(return estimates)와 같은 크레딧 할당(credit assignment) 양을 계산해야 합니다. 이러한 계산은 $[num\_envs, seq\_len]$ 형태의 텐서 상에서 동작하며, 일반적으로 1차 선형 재귀(first-order linear recurrence, 예: GAE, V-Trace, Retrace)를 포함합니다.

핵심 과제는 이러한 재귀 구조에 내재된 **순차적 의존성(sequential dependency)**입니다. 표준 구현은 시퀀스를 계산하기 위해 $T$ 번의 직렬 단계를 요구하며, 이는 수천 개의 코어를 가진 GPU에서도 시간축에 따른 병렬화를 방-해하는 엄격한 시간적 체인을 형성합니다. torch.compile은 파이썬 인터프리터 오버헤드를 제거할 수는 있지만, 알고리즘 자체가 부과하는 데이터 의존성을 깨뜨릴 수는 없으며, 이로 인해 계산 깊이가 $O(T)$ 로 남게 됩니다. 또한, 문제를 병렬 결합 스캔(parallel associative scan)으로 재구성하는 기존의 벡터화된 구현들은 높은 메모리 트래픽 문제를 겪습니다. 즉, $O(\log T)$ 리덕션의 각 배가(doubling) 단계마다 중간 결과물을 고대역폭 메모리(HBM)에 실체화(materialize)함으로써 상당한 지연 시간과 대역폭 비용을 발생시킵니다.

방법론

본 논문은 Triton 언어를 사용하여 RL 크레딧 할당을 위한 고성능 GPU 커널을 구현한 오픈 소스 라이브러리인 rl-triton을 제시합니다. 본 방법론은 세 가지 핵심 축을 기반으로 합니다.

1. 통합 대수 프레임워크 (Unified Algebraic Framework)

저자는 일곱 가지의 서로 다른 RL 추정 알고리즘—GAE, V-Trace, Retrace( $\lambda$ ), TD( $\lambda$ ) 리턴, 할인된 리턴(discounted returns), 엘리기빌리티 트레이스(eligibility traces), 에피소드 접두사 합(episodic prefix sums)—이 모두 단일한 1차 선형 재귀의 인스턴스임을 입증합니다:
$A_t = \alpha_t + \beta_t \cdot A_{t+1}$
(또는 엘리기빌리티 트레이스의 경우 그 역방향 시간 미러 형태).
이 알고리즘들은 특정 결합 연산자 $\oplus$ 를 사용하는 튜플 $(\alpha_t, \beta_t)$ 에 대한 **병렬 결합 스캔(parallel associative scans)**으로 재구성됩니다:
$(\alpha_B, \beta_B) \oplus (\alpha_A, \beta_A) = (\alpha_B + \beta_B \alpha_A, \beta_A \beta_B)$
이 변환을 통해 재귀를 $O(T)$ 의 순차적 단계가 아닌 $O(\log T)$ 병렬 단계로 해결할 수 있습니다.

2. 융합된 Triton 커널 (Fused Triton Kernels)

중간 계수 텐서( $\alpha$ 및 $\beta$ )를 스캔하기 전에 HBM에 먼저 실체화하는 표준 PyTorch 구현과 달리, rl-triton은 **융합된 커널(fused kernels)**을 사용합니다.

온칩 구축 (On-Chip Construction): 커널은 로우 롤아웃 입력(보상, 가치, 종료 플래그 등)을 로드하고, 레지스터 상에서 직접 재귀 계수( $\alpha_t, \beta_t$ )를 구축합니다.
온칩 스캔 (On-Chip Scan): 결합 스캔은 tl.associative_scan을 통해 레지스터와 공유 메모리를 사용하여 온칩에서 완전히 수행됩니다. 중간 스캔 상태는 최종 출력이 나올 때까지 HBM에 기록되지 않습니다.
메모리 효율성: 이 접근 방식은 HBM 왕복(round-trips)을 최소화합니다. 길이 $T$ 인 시퀀스의 경우, 융합된 커널은 단 한 번의 HBM 읽기와 쓰기를 수행하지만, 실체화된 스캔은 $O(\log T)$ 번의 HBM 왕복을 필요로 합니다.

3. 명시적 경계 의미론 (Explicit Boundary Semantics)

본 라이브러리는 표준 구현에서 흔히 오류를 유발하는 복잡한 에피소드 경계 조건을 명시적으로 처리합니다.

종료(Termination) vs 절단(Truncation): 자연스러운 에피소드 종료(다음 가치가 0이 되는 경우)와 절단(공급된 부트스트랩 값이 사용되어야 하는 경우)을 구분합니다.
롤아웃 윈도우(Rollout Windows): 롤아웃 버퍼의 끝( $t=T-1$ )에서의 경계를 정확하게 관리하여, 특정 알고리즘에 따라 부트스트랩 값이 로컬 업데이트( $\alpha_{T-1}$ ) 또는 스캔 캐리( $A_T$ )에 중복 계산되지 않고 적절히 포함되도록 보장합니다.

주요 기여

통합 프레임워크: 7가지 표준 RL 크레딧 할당 알고리즘이 단일 아핀 스캔(affine scan) 연산자의 인스턴스라는 증명.
융합 구현: 계수를 온칩에서 구축하고 중간 HBM 실체화 없이 스캔을 실행하는 알고리즘별 Triton 커널.
강건한 경계 처리: 종료된 단계, 절단된 단계, 롤아웃 윈도우 경계에 대한 검증된 명시적 로직.
성능 벤치마크: 순차 루프 및 벡터화된 torch.compile 베이스라인 대비 속도 향상을 보여주는 종합적인 평가.
오픈 소스 공개: 정답성 테스트, 성능 회귀 테스트 및 문서화를 포함한 Python 패키지.

결과

라이브러리는 다양한 환경 수와 시퀀스 길이에 대해 NVIDIA H100 (80GB HBM3) 및 RTX 2000 Ada GPU에서 평가되었습니다.

표준 루프 대비 속도 향상: 대규모 병렬 시뮬레이션 환경(예: 4,096개 환경, 128 스텝)에서 rl-triton은 컴파일되지 않은 순차 루프 대비 16–317배의 속도 향상을 달면합니다. 이는 주로 의존성 깊이를 $O(T)$ 에서 $O(\log T)$ 로 줄였기 때문입니다.
벡터화된 베이스라인 대비 속도 향상: 중간 결과물을 실체화하는 벡터화된 torch.compile 베이스라인과 비교했을 때, rl-triton은 1.6–5.70배의 전체 호출(full-call) 속도 향상을 달성했습니다.
- 속도 향상은 단계당 계산 복잡도가 낮은 알고리즘(예: Discounted Returns, $\lambda$ -returns)과 시퀀스 길이가 긴 경우에서 가장 두드러졌는데, 이는 베이스라인의 반복적인 HBM 왕복이 병목 현상이 되기 때문입니다.
- Retrace( $\lambda$ )의 경우, 단계당 계산량이 더 많고(행동 확률 로드) 매우 긴 시퀀스 길이( $>2048$ )에서 레지스터 압박 문제가 발생하여 속도 향상이 1.6–1.9배로 낮았습니다.
엔드 투 엔드(End-to-End) 영향: 합성 PPO 업데이트에서, GAE가 전체 업데이트 시간의 10~~15%를 차지할 때 커널 속도 향상은 1.11–1.6배의 엔드 투 엔드 훈련 속도 향상으로 이어졌습니다. GAE의 비중이 작을 때(~~2.2%)는 엔드 투 엔드 이득이 미미했습니다(~1.02배). 이는 암달의 법칙(Amdahl's Law)과 일치합니다.
PufferLib와의 비교: rl-triton은 테스트된 40개 형태 중 31개에서, 특히 시퀀스 길이가 길어질수록 대규모 병렬 시뮬레이션 환경(짧은 롤아웃, 높은 환경 수)에서 PufferLib의 수동 작성 CUDA 커널보다 우수한 성능을 보였습니다.

의의 및 주장

본 논문은 rl-triton이 RL의 핵심 경로인 크레딧 할당을 위한 통합된 고성능 솔루션을 제공한다고 주장합니다. 결합 스캔 프레임워크와 Triton의 융합 능력을 활용함으로써, 표준 구현 및 일반적인 벡터화 PyTorch 방식이 겪는 메모리 대역폭 병목 현상을 제거합니다.

저자는 이 라이브러리가 현대의 분산 훈련 및 RLHF 파이프라인에서 흔히 볼 수 있는 대규모 병렬 시뮬레이션 환경(수천 개의 환경, 짧은 롤아웃)에 특히 유의미하다고 강조합니다. 이 환경에서는 HBM 트래픽 감소가 실질적인 실제 시간(wall-clock time) 절감으로 이어집니다. 비록 이 라이브러리가 전체 훈련 파이프라인을 가속화한다고 주장하는 것은 아니지만(네트워크 순전파/역전파가 지배적인 경우가 많으므로), 매 업데이트 단계의 엄격한 의존성인 크레딧 할당 구성 요소에 대해 상당한 최적화를 제공합니다. 또한, 이 연구는 융합된 스캔이 긴 컨텍스트의 LLM 사후 학습(post-training) 및 스파이킹 신경망(spiking neural networks)과 같이 분절된 아핀 재귀를 포함하는 다른 도메인에도 적용될 수 있음을 강조합니다.