vLLM RL-Kernel v0.1.0: ROCm에서 비트 단위 훈련-롤아웃 일관성
vLLM와 RL-Kernel은 AMD Instinct MI300X GPU에서 훈련 및 롤아웃 로그확률 간 비트 단위 수치 일관성을 보장하는 시스템을 구현했습니다. 이 통합은 훈련과 추론을 처리하는 서로 다른 실행 엔진이 존재할 때 일반적으로 발생하는 수치 정책 이동을 제거하여, 강화학습(RL) 후처리 단계에서 중요도 비율과 KL 진단이 동일한 부동소수점 계산에 기반을 두도록 보장합니다.
훈련-롤아웃 수치 불일치 문제
강화학습(RL) 워크플로우에서 훈련과 롤아웃은 일반적으로 서로 다른 실행 엔진과 연산자를 사용합니다. 동일한 모델 가중치와 입력을 사용하더라도, 커널 구현 방식, 병렬화 전략, 축소 순서의 차이로 인해 로그확률이 다를 수 있습니다.
이 불일치는 중요도 비율, 롤아웃-훈련 KL 진단, 클리핑 메커니즘으로 이어지는 수치 정책 이동을 유발합니다. 이전의 해결 노력은 주로 NVIDIA 플랫폼에 집중되었지만, 이번 작업은 이러한 보장을 ROCm 생태계로 확장합니다.
기술 아키텍처: vime와 RL-Kernel
일관성을 달성하기 위해서는 데이터 흐름의 시점과 커널의 실제 수치 실행을 모두 일치시켜야 합니다. 작업은 두 구성 요소로 나뉩니다:
- vime: 훈련 타임라인을 일치시킵니다. 어떤 토큰 배치가 어느 단계에 속하는지 관리하고, 특정 토큰을 생성한 가중치 버전을 추적하며, 새로운 가중치를 vLLM에 동기화합니다.
- RL-Kernel: 수치 실행을 일치시킵니다. 동일한 값이 계산에 참여하고, 동일한 방식으로 분할, 병합, 반올림되도록 보장합니다.
수치 실행 계약
비트 단위 일관성을 달성하기 위해 RL-Kernel은 전방 경로 전반에 걸쳐 엄격한 수치 실행 계약을 적용합니다. 이 계약은 여섯 가지 핵심 관측 가능한 동작을 정의합니다:
- 입력 의존성: 출력이 의존하는 특정 요소와 축소 도메인
- 분할: 축소 도메인을 부분 결과로 나누는 방식
- 병합 순서: 부분 결과를 결합하는 순서
- 정밀도: 입력, 누산기, 중간 상태, 출력의 정밀도
- 반올림: 반올림 또는 다운캐스팅이 발생하는 정확한 시점
- 기초 연산: 사용되는 특정 수치 기초 연산 (예: exp, log, rsqrt, SiLU, FMA)
ROCm에서의 구현
AMD Instinct MI300X에 이러한 보장을 적용하기 위해, 여러 내장 ROCm 실행 경로를 수정하여 훈련 측의 동작과 일치시켜야 했습니다. 주요 기술적 수정 사항은 다음과 같습니다:
- GEMM 커널: 결정론적 MFMA 기반 GEMM 커널 구현
- 어텐션: AITER/CK 어텐션의 실행 스케줄링 수정 및 페이지 기반 KV 레이아웃 상태 문제 해결
- 통신: HIP IPC 통신 구현 및 분산 집합적 랭크 축소 순서 수정
- 수학 및 융합: 컴파일러 융합과 특정 수학 함수 구현으로 인한 마지막 비트 차이 해결
- 로그확률: 어휘 축소 구현 및 크로스-TP 어휘 병합 문제 수정
RL-Kernel의 엄격 경로에서의 경계 수정
| 계산 경계 | 엄격 경로 수정 |
|---|---|
| RMSNorm | 축소 도메인, 엡실론, 잔여 추가, 출력 경계 |
| 어텐션 | 위치, 마스크, 논리적 페이지 기반 KV, 분할 정책, LSE 정밀도, 최종 캐스트 |
| GEMM 및 SwiGLU | K차원 축소, 누산 정밀도, 에필로지, 활성화, 재료화 경계 |
| 선형 로그확률 | 실제 어휘 범위, 대상 소유권, 로컬 축소, 크로스-랭크 LSE 병합 |
| 분산 집합적 연산 | 페이로드 소유권, dtype, 고정된 랭크 축소 순서 |
AMD Instinct MI300X에서의 검증 결과
단일 노드에 8개의 AMD Instinct MI300X GPU를 사용하여 Qwen3-8B 모델(BF16)을 기반으로 엔드투엔드 실험을 수행했습니다. 훈련에는 Megatron(TP4/CP2/PP1)을, 롤아웃에는 두 개의 vLLM 엔진(TP4 각각)을 사용했습니다.
주요 발견
- 완전한 일치: 200번 이상의 연속된 롤아웃 및 훈련 단계에서
mismatch_count와max_abs_diff는 정확히 0을 유지했습니다. - 내장 vime와의 비교: 내장 vime는 각 단계에서 일치하지 않는 결과를 보였지만, vime + RL-Kernel 구성은 전체 트래잭터리에서 비트 단위 일치를 유지했습니다.
- 안정성: 200단계 실행 동안 평균 절대 훈련-롤아웃 로그확률 차이는 항상 0을 유지했습니다.
성과 요약
ROCm에서 vime와 RL-Kernel의 통합은 RL 후처리에 다음과 같은 핵심 기능을 제공합니다:
- 비트 단위 일관성: ROCm 하드웨어에서 훈련 및 롤아웃 로그확률의 정확한 일치
- 검증 가능한 재현성: 안정적인 일관성 보장을 통해 RL 훈련 트래잭터리의 검증을 용이하게 함
- 실패 위치 식별: 연산자 아브레이션을 통해 수치 분리가 시작되는 정확한 시스템 경계를 식별할 수 있음
- 실행 증거: 런타임에 사용되는 커널, HIP Graph 실행, 페이지 기반 KV 경로에 대한 포괄적인 검증