RL 가중치 리핏 가속화 위한 Reshard Collective 공개

새로운 리핏 통신 방식 도입

엔비디아는 강화 학습(RL) 후처리 훈련 과정에서 모델 가중치(weight)를 훈련 GPU에서 추론 GPU로 전송하는 '리핏(refit)' 작업의 효율성을 높이기 위해 새로운 통신 방식인 Reshard Collective를 공개했다. 이 새로운 방식은 기존의 비효율적인 통신 오버헤드를 줄여 대규모 모델의 리핏 성능을 크게 향상하는 것을 목표로 한다. NCCL M2N Reshard는 Reshard Collective의 오픈소스 참조 구현으로, NVIDIA/nccl-extensions GitHub 저장소를 통해 제공된다.

기존 리핏 구현은 훈련 및 추론 레이아웃의 차이를 처리하기 위해 '수집 및 브로드캐스트(gather-and-broadcast)' 방식을 주로 사용했다. 이 방식은 모든 모델 파라미터에 대해 훈련 GPU가 로컬 샤드를 모두 수집하여 완전한 파라미터를 재구성한 다음, 이 완전한 파라미터 텐서를 추론 GPU에 브로드캐스트하고 각 추론 GPU는 필요한 샤드만 유지하는 구조이다. 그러나 이 과정에서 추론 GPU가 필요한 샤드 외에 전체 파라미터를 수신하게 되어 수집 및 브로드캐스트 단계 모두에서 중복 통신이 발생한다.

대규모 모델의 리핏 과제

리핏 과정은 두 가지 주요 특성으로 인해 어려움을 겪는다. 첫째, 추론과 훈련은 서로 다른 병렬화 매핑(데이터, 컨텍스트, 텐서, 전문가, 파이프라인 병렬화)을 사용하므로 가중치가 훈련 및 추론 인스턴스에 다르게 분산된다. 둘째, 리핏은 추론 GPU의 모든 모델 가중치를 업데이트해야 하므로 통신 페이로드(payload)가 모델 전체 가중치 크기에 달한다. 최신 대규모 모델의 경우 이 페이로드가 이미 매우 크며, 앞으로도 계속 증가할 것으로 예상된다.

예를 들어, Kimi K3 모델은 2.8조 개의 파라미터를, Qwen 3.8 2.4T A95B 모델은 2.4조 개의 파라미터를 포함한다. BF16 가중치를 기준으로 파라미터당 2바이트를 가정하면, 이들 모델의 가중치 크기는 각각 약 5.6TB와 4.8TB에 달한다. 기존의 수집 및 브로드캐스트 방식은 올-개더(all-gather) 및 브로드캐스트 단계를 모두 고려할 때 모델 가중치 크기의 약 두 배에 해당하는 논리적 페이로드를 발생시킨다. 이는 Kimi K3의 경우 약 11.2TB, Qwen 3.8 2.4T A95B의 경우 약 9.6TB에 이르는 통신량을 의미하며, 이러한 통신은 매 훈련 단계마다 수행되어야 한다. 따라서 다양한 가중치 레이아웃을 처리하면서 실제 페이로드 크기를 줄이는 정교한 리핏 접근 방식이 필요하다.

DTensor 기반 Reshard Collective 동작

Reshard Collective는 리핏을 DTensor-to-DTensor 전송 문제로 재구성하여 이러한 한계를 극복한다. PyTorch 분산 텐서(DTensor) 추상화는 GPU에 걸쳐 샤딩되거나 복제된 텐서를 표현하도록 설계되었으며, Reshard Collective는 이 DTensor를 활용하여 훈련 레이아웃에서 추론 레이아웃으로 텐서를 재샤딩하는 방식으로 리핏을 정의한다. 이 추상화는 리핏 문제를 단순화하고 통신 로직을 라이브러리에 캡슐화하여 애플리케이션과 독립적으로 알고리즘을 개선할 수 있도록 한다.

NCCL M2N Reshard는 M개의 소스 GPU에 분산된 텐서를 N개의 대상 GPU가 요구하는 레이아웃으로 전송한다. 이 구현은 먼저 소스 및 대상 샤드 간의 정확한 교차점을 도출한 다음, 이러한 전송을 물리적 GPU 네트워크 토폴로지에 매핑한다. 모든 소스 기여를 모든 대상 복제본에 독립적으로 보내는 대신, NCCL M2N은 복제본을 NVIDIA NVLink 도메인별로 그룹화한다. 각 도메인 내에서 들어오는 기여를 복제본 랭크에 최대한 균등하게 분배하여 각 수신 랭크가 네트워크를 통해 전체 데이터의 고유한 하위 집합을 수신하도록 한다. 이후 수신 랭크는 NVLink를 통해 해당 하위 집합을 교환하여 모든 복제본이 중복 네트워크 트래픽 없이 완전한 샤드를 조립할 수 있도록 한다. 네트워크 전송과 NVLink 복제는 서로의 지연 시간을 숨기기 위해 파이프라인 방식으로 처리된다.

NeMo RL 통합 및 성능 향상

Reshard Collective의 효과를 입증하기 위해 NVIDIA NeMo RL 프레임워크에 참조 구현이 제공된다. NeMo RL은 RL 후처리 훈련을 간소화하고 확장하도록 설계된 오픈소스 프레임워크이다. Reshard Collective의 명확한 추상화 덕분에 프레임워크 내 구현은 비교적 간단하다. 주어진 병렬화 구성에 대해 프레임워크는 훈련 및 추론 인스턴스에 대한 각 파라미터의 DTensor 설명을 제공하기만 하면 된다.

실제 구현에서는 트랜스포머 블록의 피드포워드 프로젝션 가중치가 리핏 페이로드의 대부분을 차지한다는 점에 착안하여 하이브리드 접근 방식을 사용한다. 최신 대규모 언어 모델(LLM)에서 이 가중치는 모델 파라미터의 약 95%를 차지한다. 따라서 지배적인 피드포워드 프로젝션 가중치는 Reshard 경로를 사용하고, 다른 작은 가중치(예: QKV 프로젝션 또는 맘바 가중치)는 광범위한 기능 범위를 위해 최적화된 기존의 수집 및 브로드캐스트 리핏 경로를 사용한다. 이 방식은 대부분의 페이로드에 대해 효율적인 Reshard 통신과 기존 방식과 동일한 기능 범위를 결합한다.

NeMo RL 구현을 사용한 벤치마크 결과에 따르면, Reshard 기반 리핏은 기존 수집 및 브로드캐스트 리핏 대비 2.7배에서 5.7배의 속도 향상을 달성했다. NVIDIA DGX H100 및 NVIDIA DGX GB200 시스템에서 Qwen 3 235B 모델의 리핏 지연 시간은 6.4초에서 1.9초로, DeepSeek v3 671B 모델은 14.9초에서 2.6초로 단축되었다. 특히 2.4조 개의 파라미터를 가진 Qwen3.8 BF16 모델의 리핏은 4.3초 만에 완료되어, Reshard Collective가 설계된 조 단위 파라미터 규모에서 효과적임을 입증했다.

적용 조건 및 향후 계획

Reshard Collective를 NeMo RL에서 활성화하려면 설정 파일에서 `policy.generation.colocated.enabled`를 `false`로, `refit_transport`를 `nccl_reshard`로 설정해야 한다. NeMo RL 리핏 가이드와 NCCL Reshard 설정 요구사항은 이 기능에 대해 지원되는 백엔드, 레이아웃 및 정밀도 모드를 설명한다.

GPU 멤버십이 변경될 수 있는 런타임 환경에 대응하기 위해, 프레임워크는 통신 그룹을 파괴하고 재구성하거나 NCCL shrink 및 grow API를 사용하여 기존 통신 그룹을 조정할 수 있다. 현재 NeMo RL은 복원력 지원을 위해 첫 번째 접근 방식을 사용하며, NCCL shrink 및 grow API를 사용한 복원력 지원은 향후 업데이트에서 계획하고 있다. 다양한 레이어 유형과 각 병렬화 구성에서 훈련 및 추론 엔진이 파라미터를 샤딩하는 방식의 다양성으로 인해 모든 엣지 케이스를 Reshard Collective를 통해 처리하는 것은 여전히 과제로 남아 있다.