NeMo-DCR, 초대형 에이전트 RL 가중치 동기화를 약 35배 가속

왜 가중치 동기화가 병목이 되는가

에이전트 강화학습에서는 정책을 업데이트하는 학습 클러스터와, 현재 정책으로 상호작용 데이터를 생성하는 rollout 클러스터를 분리할 수 있습니다. 이 구조는 자원 확장에 유리하지만, 학습이 끝날 때마다 새 정책을 rollout 측에 전달해야 한다는 과제를 만듭니다. 모델 규모가 1조 파라미터에 이르면 전체 체크포인트를 반복적으로 옮기는 비용이 다음 학습 주기를 지연시킬 수 있습니다.

NeMo-DCR(Delta-Compressed Refit)은 이 refit 과정을 변경분 전송으로 바꾸는 방법입니다. 논문이 제시한 관찰에 따르면 BF16 학습에서 한 스텝마다 저장된 값이 바뀌는 가중치 요소는 약 1%입니다. 그러나 단순히 변경된 값만 모으는 것으로는 충분하지 않습니다. 학습 클러스터와 rollout 클러스터의 샤딩 및 메모리 배치가 다를 수 있기 때문에, 최종 결과가 전체 체크포인트를 로드했을 때와 비트 단위로 같아야 합니다.

핵심 설계

• 정규 좌표로 투영: 고정된 아핀 매핑을 이용해 학습 샤드의 변경사항을 체크포인트 정규 좌표로 변환합니다. 약 96%의 변경은 직접 투영하고, 나머지는 잔여 변환으로 처리합니다.
• XOR 마스크와 덮어쓰기: 로딩 과정에서도 저장 비트 패턴을 보존할 수 있는 변경에는 압축 가능한 XOR 마스크를 적용합니다. 그 밖의 변경은 절대값 덮어쓰기로 전달해 산술 재구성에 따른 오차를 피합니다.
• 기존 로더를 활용한 배치: 모델별 배치 규칙을 새로 구현하거나 전체 텐서를 먼저 조립하지 않습니다. rollout 런타임의 네이티브 로더가 수행하는 메모리 복사 경로에 개입해 수신 영역에 업데이트를 직접 적용합니다.
• 실패 복구와 버전 관리: 일부 쓰기만 완료된 상태에서도 덮어쓰기를 통해 재시도할 수 있습니다. 공동 커밋은 설치된 정책과 다음 차분의 기준이 되는 baseline을 연결해 버전 불일치를 방지합니다.
• 통신 경로의 유연성: 객체 스토리지나 릴레이 트리를 사용할 수 있어 클러스터 간 collective에 의존하지 않습니다. 차분 생성, 전송, 적용을 겹쳐 실행할 수 있으며 비동기 학습에서는 rollout 요청 생성과 동시에 전송할 수도 있습니다.

결과와 의미

1조 파라미터 모델에서 요소 변경률을 3%로 설정한 스트레스 테스트에서 전체 체크포인트 전송은 87.5분이 걸렸지만, NeMo-DCR refit은 2.5분에 완료됐습니다. 이는 약 35배의 속도 향상입니다. 제공된 자료는 3%와 5% 변경률에서 30B부터 1T 규모 모델까지의 refit 실험도 언급합니다.

NeMo-DCR의 의미는 대역폭 절감에만 있지 않습니다. 희소 업데이트, 서로 다른 메모리 배치의 변환, 비트 정확성, 인플레이스 적용, 중단된 refit의 복구를 하나의 동기화 경로로 묶었습니다. 실제 시스템에서는 변경량이 적더라도 잘못된 위치에 기록되거나 반올림 차이로 비트가 달라지거나, 장애 뒤 불완전한 상태가 남으면 사용할 수 없습니다.

더 빠른 refit은 rollout 워커가 오래된 정책을 사용하는 시간을 줄이고 학습 파이프라인의 대기 시간을 낮출 수 있습니다. 다만 실제 효과는 변경률, 네트워크와 스토리지 처리량, 서빙 로더 구현에 따라 달라집니다. 차분 동기화를 안정적으로 운영하려면 기준 체크포인트와 정책 버전의 관계를 끝까지 일관되게 관리해야 합니다.

출처: Hugging Face Daily Papers