RL-Align/RL-Kernel

해결하는 문제

RL-Kernel은 강화학습 (RL) 후학습 중에 발생하는 수치적 불일치 문제를 해결합니다. 표준 설정에서 롤아웃 엔진(샘플 생성)과 학습 엔진(모델 업데이트)은 종종 서로 다른 커널이나 계산 순서를 사용합니다. 이는 로그 확률의 불일치로 이어져 RL 학습의 안정성과 품질을 저하시킬 수 있습니다. RL-Kernel은 이 두 단계가 동일한 결과를 생성하도록 보장하여 전체 학습 과정에서 수치적 일관성을 유지합니다.

작동 원리

이 라이브러리는 실행 엔진(롤아웃을 위한 vLLM 및 학습을 위한 Megatron-LM 등)과 기본 하드웨어 사이의 미들웨어 계층으로 작동합니다. attention, dense FFN, 그리고 GRPO 및 PPO와 같은 특정 RL 목적 함수를 포함한 결정론적이고 고성능인 연산자 세트를 제공함으로써, 롤아웃 엔진과 학습 엔진 모두가 수학적으로 동일한 방식으로 계산을 수행하도록 강제합니다. 또한 지원되는 가속기에서 성능을 최적화하고 메모리 사용량을 줄이기 위한 하드웨어 전용 커널도 포함되어 있습니다.

대상 사용자

이 프로젝트는 LLM을 위한 대규모 RL 후학습을 수행하며, 고성능이고 재현 가능하며 하드웨어 간 호환성이 뛰어난 학습 파이프라인이 필요한 머신러닝 엔지니어 및 연구자를 위해 설계되었습니다.

주요 특징

• 결정론적 일관성: 롤아웃 엔진과 학습 엔진 간의 비트 단위로 정확한 로그 확률 일치를 보장합니다.
• 하드웨어 간 지원: NVIDIA (CUDA) 및 AMD (ROCm) GPU를 기본적으로 지원하며, Huawei Ascend 하드웨어에 대한 부분적인 지원을 제공합니다.
• 성능 최적화: 융합 연산(fused computation)과 커스텀 커널을 사용하여 롤아웃 시간과 동기화 오버헤드를 최소화합니다.
• 엔진 통합: vime와 통합되어 vLLM 및 Megatron-LM을 오케스트레이션하며, RL 워크플로우를 위한 통합된 경로를 제공합니다.