원문 캡처 · modal.com
원문 캡처 · modal.com

대규모 강화학습(RL)에서 샘플러와 학습기는 보통 같은 가속기와 정밀도를 사용한다. 이번 실험은 서로 다른 가속기·정밀도·커널을 쓰는 샘플러를 하나의 BF16 학습기에 연결해도, 오프폴리시 보정으로 학습을 안정화할 수 있는지 살폈다. 혼합 샘플러는 동종 샘플러보다 샘플링 비용을 업데이트당 대부분 약 15% 줄였고, SWE-bench Pro V2 평가 성능은 동종 샘플러와 비슷했다.

실험에서는 Qwen3.6-35B-A3B를 최대 256K 토큰의 에이전트 코딩 작업으로 학습했다. 학습 데이터는 Xiaomi가 공개한 MiMo-V2.6 저장소 작업 2,700개에서 참조 수정사항이 새는 기록 등을 제거해 구성했다. 각 에피소드에서 모델은 여러 차례 대화하며 이슈를 수정하고, 제출한 패치가 새 샌드박스에서 테스트를 통과하면 보상 1, 그렇지 않으면 0을 받았다. 학습과 겹치지 않는 SWE-bench Pro V2와 하위 집합 HARD-51로 평가했으며, 모든 체크포인트는 BF16으로 평가했다.

혼합 구성은 B200·B300의 NVFP4, H100·H200의 FP8 및 BF16, A100·RTX PRO 6000의 BF16 샘플러로 이뤄졌다. 각 샘플러는 자체 커널과 FP8 KV 캐시를 사용했다. 비교 대상인 동종 구성은 BF16 가중치와 KV 캐시를 쓰는 B200만 사용했다. 두 구성 모두 Megatron-LM 학습기를 B200에서 BF16으로 실행했고, 학습은 비동기 방식으로 진행했다. 샘플러의 지연 정도에는 제한을 두지 않았다.

저자들은 GRPO에 세 가지 보정을 적용했다. 마스크 중요도 샘플링(MIS)은 각 토큰의 중요도 비율을 정해진 범위 안에서만 반영해 극단적인 가중치가 업데이트를 지배하지 못하게 한다. 점수 중심화(SC)는 샘플링 정책에서 생기는 기울기 편향을 제거한다. top-p 샘플링과 마스크 재생은 샘플러와 학습기가 같은 후보 토큰 집합을 사용하도록 맞춘다.

보정 없이 학습하면 혼합 구성의 GRPO는 시작 직후 붕괴했고, 동종 구성도 더 오래 버틴 뒤 결국 붕괴했다. 세 방법을 각각 단독으로 적용해도 붕괴를 막지 못했다. SC와 MIS를 함께 적용한 경우에는 top-p 사용 여부와 관계없이 학습이 안정적으로 끝났다. SC+MIS+top-p 구성에서 HARD-51 pass@4는 약 100~140단계 동안 올라 65~75%에서 안정됐고, 전체 평가 집합에서는 89%에서 약 94%로 상승했다.

안정적인 혼합 구성의 pass@4는 동종 구성과 학습 내내 비슷했다. 샘플러와 학습기 사이의 KL 발산은 평균적으로 동종 구성보다 4.7배 컸지만, 학습 보상과 기울기 크기는 비슷한 흐름을 보였다. 저자들은 불안정성의 핵심이 불일치의 크기 자체보다 그 불일치가 계속 커지는지에 있다고 해석한다. 실제로 붕괴한 실험에서는 KL 발산이 3배 이상 커졌고, 학습 보상이 떨어지며 에피소드 다수가 형식 오류를 반복해 끝났다.

다만 보정은 토큰 단위로 작동하므로 문맥과 어드밴티지는 여전히 샘플러가 만든 것을 사용한다. 이 실험은 시퀀스 단위 보정이나 지연 제한을 적용하지 않았으며, 평균 지연은 약 4단계였다. top-p와 마스크 재생은 학습 대상도 바꾼다. 업데이트는 top-p 후보 집합으로 제한된 정책을 최적화하고, 집합 밖 토큰은 해당 위치에서 기울기를 받지 않는다. 평가 때는 전체 어휘에서 샘플링했으며, top-p를 쓴 구성의 성능은 학습에서도 전체 어휘를 쓴 SC+MIS와 비슷했다.

이번 결과는 하나의 모델과 실험 구성에서 얻은 것이다. 저자들은 보정이 다른 모델과 가속기에도 적용될 것으로 예상하지만, 이를 일반적으로 입증한 결과로 제시하지는 않는다. 이 접근은 샘플러와 학습기의 장비·정밀도 선택을 분리해 각각 조정할 수 있게 하며, 가속기와 숫자 형식이 다양해지는 상황에서 샘플링 자원을 학습기와 별도로 확장할 수 있는 가능성을 보여준다.

참고 자료: Qwen3.6-35B-A3B, MiMo-V2.6 학습 데이터, SWE-bench Pro V2, HARD-51, Miles, SGLang, Megatron-LM, GRPO 논문, MIS 논문, SC 논문, top-p 마스크 재생 논문.