NVIDIA, RL 가중치 재조정을 최적화하기 위해 Reshard Collective 소개

NVIDIA는 강화 학습(RL) 워크플로우에서 가중치 재조정을 최적화하기 위해 설계된 새로운 통신 방법인 Reshard Collective를 발표했습니다. 가중치 재조정은 훈련 GPU에서 추론 GPU로 업데이트된 정책 모델 파라미터를 전송하는 과정으로, 이는 새로운 훈련 경로를 생성하는 데 필수적입니다. Reshard Collective는 특히 모델 크기가 계속 증가함에 따라 현재 방법의 비효율성을 해결하는 것을 목표로 하고 있습니다.

대규모 모델의 가중치 재조정에서의 과제

전통적인 가중치 재조정 방법은 훈련 GPU가 전체 모델 파라미터를 재구성한 다음 추론 GPU에 브로드캐스트하는 수집 및 브로드캐스트 접근 방식에 의존합니다. 이는 직관적이지만, 추론 GPU는 특정 샤드만 필요할 때에도 전체 텐서를 수신하기 때문에 불필요한 통신이 발생합니다. Kimi K3( 2.8 조 파라미터) 및 Qwen 3.8 2.4T A95B( 2.4 조 파라미터)와 같은 대규모 모델의 경우 이러한 비효율은 더욱 심화됩니다. BF16 가중치를 가정할 때, 이러한 모델은 각각 11.2 TB와 9.6 TB의 통신 페이로드를 필요로 합니다.

Reshard Collective: DTensor-to-DTensor 솔루션

Reshard Collective는 가중치 재조정을 DTensor에서 DTensor로의 전송 문제로 재정의하여, PyTorch의 분산 텐서 추상화를 활용합니다. 이 접근 방식은 훈련 GPU와 추론 GPU 간의 직접적인 샤드 간 전송을 가능하게 하여, 중간의 전체 텐서 표현의 필요성을 제거합니다. 겹치는 샤드 영역을 매핑하고 통신 경로를 최적화함으로써 Reshard Collective는 페이로드 크기를 최소화하고 GPU 간 데이터 분배의 균형을 맞춥니다.

성능 벤치마크

NVIDIA의 내부 벤치마크에 따르면 Reshard Collective는 리피트 대기 시간을 크게 줄입니다. 671억 파라미터 DeepSeek-V3 모델의 경우 DGX H100 시스템에서 대기 시간이 14.9초에서 2.6초로 줄어들고, DGX GB200 시스템에서는 6.1초에서 2.3초로 줄어들었습니다. NDR InfiniBand가 있는 256-GPU NVL72 클러스터에 대한 추가 테스트에서는 가중치 동기화 시간이 5.78초에서 2.77초로 줄어들어 총 RL 단계 시간이 12.7%만큼 단축되었습니다. 단일 피드포워드 전문가 혼합(layer) 전송은 평면 직접 전송과 비교하여 77.3 밀리초에서 9.8 밀리초로 개선되었습니다.

오픈 소스 구현

Reshard Collective는 NVIDIA NCCL 확장 리포지토리를 통해 오픈 소스 구현 형태로 제공됩니다. 이 패키지는 실험적 메쉬 간 구성 요소 및 CUDA 12 및 CUDA 13에 대한 종속성을 포함하며, NCCL 2.30.7는 사용이 검증되었습니다. NVIDIA의 문서에서는 과거에 가중치 전송 최적화를 탐색한 NeMo RL 및 vLLM과 같은 프레임워크에 통합할 가능성도 강조됩니다.

RL 워크플로에서 가중치 동기화의 커져가는 병목 문제를 해결함으로써 Reshard Collective는 대규모 AI 모델을 효율적으로 확장하는 데 중요한 발전을 의미합니다.