IsoExec: SkyRL에서 트레이너-인퍼런스 불일치 제거하기

vLLM는 강화학습(RL) 워크로드에서 학습 및 인퍼런스 엔진 간 수치적 불일치를 제거하기 위해 설계된 IsoExec을 도입했습니다. 이는 엄격한 실행 계약을 적용하고 병렬성에 영향을 받지 않는 커널을 활용함으로써, 롤아웃 엔진과 트레이너가 비트 단위 일관성으로 동일한 정책을 평가하도록 보장합니다. 이로 인해 평균 엔드투엔드 롤아웃 대 훈련 로그확률 차이가 $10^{-7}$ 이하로 감소했으며, 성능 오버헤드는 약 25% 수준입니다.

문제: 트레이너-인퍼런스 불일치

온폴리시 RL에서는 동일한 정책을 두 번 실행해야 합니다. 하나는 토큰 샘플링을 위한 롤아웃 엔진에서, 다른 하나는 로그확률 재계산을 위한 트레이너에서 수행됩니다. 실제로는 두 단계가 종종 다른 엔진(예: vLLM로 롤아웃, Megatron으로 훈련)을 사용하며, 이는 서로 다른 커널, 배치 형태, 병렬화 레이아웃을 사용하게 됩니다.

부동소수점 연산은 결합법칙이 성립하지 않기 때문에, 이러한 체계적인 차이로 인해 축약 순서가 달라져 모델 파라미터가 동일하더라도 토큰 확률이 달라집니다. 이 불일치는 REINFORCE 및 GRPO와 같은 RL 알고리즘을 불안정하게 만들고 보상 붕괴를 유발할 수 있습니다. 예를 들어, Fireworks에서 보고한 GLM-5.2 실행 결과는 트레이너-인퍼런스 KL이 약 0.013로, 클리핑을 통해 45%의 토큰이 제거되었으며, 약 20단계에서 보상 붕괴가 발생했습니다.

IsoExec 통합 실행 계약

IsoExec은 실행 계약을 사용하여 이러한 불일치를 해결합니다. 이 계약은 두 런타임이 따라야 할 모든 비트 관련 실행 선택을 명시합니다. 이 프레임워크 독립적인 계약은 엔진 간 반올림 민감한 세부 사항이 동일하게 유지되도록 보장합니다.

계약 구조 및 강제화

• 케이스와 영역: 계약은 전방 연산자를 "영역"(한 커널로 구현된 산술 범위)으로 나누고, 각 영역이 어떻게 처리될지를 지정하는 "케이스"(예: engine_prefill 대 trainer_fwd)를 정의합니다.
• 구성: 각 영역-케이스 쌍에 대해 계약은 특정 구현을 선택하고, 누적 데이터형, 축약-분해 파라미터(예: split-K 분할 수)와 같은 상수로 고정합니다.
• 청구사항: 런타임에서 강제되는 조건으로, 예를 들어 특정 병렬 크기에서 축약 트리가 비트 단위 불변임이 입증된다는 토폴로지 청구사항이 포함됩니다.
• 정체성: SHA-256 해시(semantic, numerical_policy, deployment)는 트레이너와 롤아웃 엔진이 동일한 검증된 수치 정책을 실행하고 있는지 확인하는 데 사용됩니다.

각 런타임에 대한 계약 어댑터는 이러한 사양을 프레임워크의 확장 지점에 연결하고 런타임을 모니터링하여 준수 여부를 확인합니다.

통합 모델 및 병렬성에 영향을 받지 않는 커널

IsoExec는 텐서, 전문가, 시퀀스 병렬화를 포함한 다양한 분산 전략에서도 비트 단위 일관성을 유지하는 통합 모델 정의를 구현합니다.

병렬성 불변성 달성

다른 GPU 레이아웃 간 수치 일관성을 유지하기 위해 IsoExec는 고정된 이진 축약 트리 방식을 사용합니다:

• 텐서 병렬화(TP): pik 구현을 사용하여 K 차원을 연속적인 리프로 나눕니다. 각 리프는 FP32 누적을 사용하는 결정론적 Tensor Core MMA를 사용하며, 계약은 랭크-리프 매핑과 산술 스케줄을 고정합니다.
• 전문가 병렬화(EP): 전문가 출력은 랭크 순서가 아닌 고정된 라우팅 순서로 결합됩니다.
• 시퀀스 병렬화(SP): IsoExec는 SP가 활성화되었는지 여부와 관계없이 동일한 축약 트리를 재사용하며, 각 랭크는 자신의 출력 슬라이스를 유지하므로 트레이너 로짓이 동일하게 유지됩니다.

청크 단위 병렬 순환(CPR) 게이트드 델타넷

Gated DeltaNet(GDN)과 같은 선형 어텐션 아키텍처는 일반적으로 훈련/프리필에 청크 단위 병렬 형태를 사용하고 디코딩에 순환 형태를 사용하여 수치적 불일치를 유발합니다. 이 문제를 해결하기 위해 이전에 순환 형태를 전부 사용하려는 시도가 있었지만, 일부 워크로드에서 최대 5배의 성능 저하가 발생했습니다.

IsoExec는 청크 단위 병렬 순환(CPR) 을 도입하여, 순환을 주 함수로 유지하면서도 청크별로 병렬로 평가합니다. 디코딩 시, $C$개의 토큰마다(여기서 $C$는 청크 크기) 숨겨진 상태를 재동기화합니다. 이는 비트 정확성을 달성하면서도 높은 처리량을 유지합니다:

단계	네이티브 믹스드	청크 단위 전부	순환 전부	CPR
비트 정확성	아니요	예	예	예
트레이너 전방+역전파	5.177 ms	5.177 ms (1.00x)	22.863 ms (4.42x)	7.386 ms (1.43x)
롤아웃 프리필	0.844 ms	0.844 ms (1.00x)	3.639 ms (4.31x)	1.412 ms (1.67x)
롤아웃 디코딩	0.0612 ms	2.2374 ms (36.6x)	0.0612 ms (1.00x)	0.0846 ms (1.38x)

실험 결과

IsoExec는 동기 RL을 사용하여 8xH100 노드 하나에서 Qwen3.5-35B-A3B 모델을 DAPO-Math-17k 데이터셋으로 훈련하여 테스트되었습니다.

수치 정확도

50단계 동안 평균 업데이트 전 롤아웃 대 훈련 절대 로그확률 차이가 네이티브 SkyRL의 $10^{-3}$에서 IsoExec의 $10^{-7}$ 이하로 감소했습니다. 평균 단계별 최대 차이는 5.073에서 $10^{-4}$로 감소했습니다.

성능 오버헤드

불일치를 제거함으로써 네이티브 SkyRL 스택보다 성능 저하가 발생했습니다:

메트릭	네이티브	IsoExec	오버헤드
생성	591.3 s	776.6 s	31.3%
정책 훈련	498.6 s	591.3 s	18.6%
전체 RL 단계	1224.6 s	1534.0 s	25.3%

수치 일관성은 달성되었지만, 단 50단계의 짧은 테스트 창에서 의미 있는 보상 향상은 관측되지 않았습니다.