Taking Reinforcement Learning Cross Datacenter — Nan Jiang, Modal
한 줄 핵심
강하게 결합된 RL 훈련 클러스터에서 롤아웃만 분리해 전역으로 흩어진 GPU를 쓰고, 전체 체크포인트 대신 거의 변하지 않는 서빙 가중치의 sparse diff만 전송하면 크로스 데이터센터 RL이 현실적인 아키텍처가 된다.
시간별 요약
-
0:01
이 영상은 post-training RL의 병목이 알고리즘이 아니라 GPU가 어디에 있고 지금 바로 쓸 수 있느냐는 물리적 제약이라고 본다. 기본적인 trainer-rollout 구조는 같은 클러스터와 RDMA에 묶여 있어 확장성이 떨어진다.
-
2:01
훈련은 all-reduce와 model parallel 통신이 필요한 하나의 강결합 작업이지만, rollout은 trajectories와 metadata를 주고받는 serving job들의 집합이다. 따라서 backpropagation은 클러스터 안에 남기고, 이동 가능한 단위는 rollout serving island으로 봐야 한다.
-
5:01
핵심은 전체 checkpoint를 매번 보내는 대신, 서빙에 보이는 weight 변화만 작은 patch로 전송하는 것이다. BF16 같은 유한 정밀도에서는 optimizer update가 rounding boundary를 넘지 못하면 서빙 값이 바뀌지 않기 때문에, 실제로는 대부분의 step에서 변화가 매우 sparse해진다.
-
10:02
실험과 관측은 이 가설을 뒷받침한다. dense gradient와 dense master update가 있어도 rollout이 보는 weight는 약 99%가 step마다 bit-identical로 유지될 수 있고, patch는 additive delta가 아니라 bitwise하게 정확히 재구성되는 lossless diff여야 한다.
-
15:01
실행 방식은 trainer가 immutable rollout version을 게시하고, 외부 rollout engine이 각자 최신 버전을 pull해 로컬에서 materialize하는 구조다. 이렇게 하면 version-aware proxy가 뒤처진 엔진을 따라잡게 하거나, 불가능하면 not ready를 돌려주면서도 전역의 유휴 GPU를 elastic rollout fleet으로 활용할 수 있다.
-
17:01
Stitch는 이 아이디어를 구현한 예시로, trainer와 engine, transport에 비교적 독립적인 비동기-first 구조를 취한다. 전역 분산 GPU를 rollout에 직접 쓰게 만들며, full checkpoint 전송 대비 수백 GB를 수백 MB 수준의 delta로 줄여 cross-datacenter RL을 현실화한다.
실무 인사이트
- RL 시스템의 병목은 모델 품질보다도 동기화 단위가 너무 크다는 데 있으므로, 아키텍처를 설계할 때 '무엇이 실제로 이동해야 하는가'를 먼저 분리해야 한다.
- 서빙 정밀도와 optimizer step 크기를 함께 보면, 많은 업데이트는 실제 서빙 값에 반영되지 않으므로 체크포인트 전체가 아니라 versioned diff 전송으로 비용을 크게 줄일 수 있다.
- 전역 GPU 활용을 노릴 때는 trainer와 rollout을 같은 클러스터에 묶기보다, version-aware protocol과 immutable artifact를 중심으로 비동기 운영하는 편이 확장성에 유리하다.