2초 이내의 강화 학습 사후 학습 가중치 전송

당사는 최근 Kimi-K2(1T 파라미터)를 대상으로 256개의 학습용 GPU(BF16)에서 128개의 추론용 GPU(FP8)로 가중치를 전송하여 1.3초 만에 머신 간 파라미터 업데이트를 완료했습니다.

비동기식 강화 학습 미세 조정에서는 학습과 추론이 서로 다른 GPU에서 실행됩니다. 각 학습 단계가 끝나면 새로운 가중치를 추론 노드로 푸시해야 합니다. 기존의 많은 프레임워크는 조 단위 파라미터 모델을 처리하는 데 수 초에서 수 분까지 소요됩니다.

RDMA 지점 간 통신(Point-to-Point communication)을 활용함으로써 추론 엔진을 변경하지 않고도 가중치 전송을 매우 빠르게 수행할 수 있으며, 코드 작성 및 유지보수도 더 수월해졌습니다.

당사의 솔루션은 소스 노드가 대상 노드의 GPU 메모리에 직접 데이터를 기록하는 단방향 기본 요소인 RDMA WRITE를 기반으로 구축되었습니다.

python

def rdma_write(src_ptr, dst_ptr, size, src_mr, dst_mr):
    # Write from local [src_ptr, src_ptr+size) to remote [dst_ptr, dst_ptr+size).
    # src_mr and dst_mr contains Memory Region metadata of both sides.
    ...

대상 측은 전송 사실조차 인지하지 못합니다. 이를 통해 학습 노드가 주도하여 추론 노드의 제어 로직 없이도 저지연, 고처리량의 제로 복사(Zero-copy) 전송을 수행할 수 있습니다.

메타데이터 수집 – 컨트롤러가 모든 학습 및 추론 GPU에서 파라미터 메타데이터를 수집합니다.

일정 계산 – 컨트롤러가 어떤 학습 GPU가 어떤 파라미터를 어떤 추론 GPU로, 어떤 순서로 전송할지 매핑하는 정적 가중치 전송 일정을 계산합니다.

일정 배포 – 컨트롤러가 모든 학습 GPU에 일정을 전송합니다.

실행 – 각 학습 단계 후 컨트롤러가 학습 GPU에 전송 시작 신호를 보냅니다.

워크플로우가 정의되면 핵심 과제는 조 단위 파라미터 규모에서 가중치 전송을 얼마나 효율적으로 실행하느냐입니다. 여기서는 실행 경로의 세부 정보를 설명합니다.

학습 중인 파라미터는 FSDP 배치에 따라 분산됩니다. full_tensor()를 사용하면 DeviceMesh 내의 모든 GPU가 전체 파라미터를 재구성할 수 있으므로, 모든 GPU가 가중치 전송의 소스 역할을 할 수 있습니다.

서로 분리된 여러 DeviceMesh가 메시 그룹(mesh group)을 형성합니다. 같은 그룹 내의 DeviceMesh는 분리되어 있으므로 전송 시 서로 간섭하지 않으며 완전히 병렬로 실행될 수 있습니다. 메시 그룹 간에는 순서를 강제하기 위해 전역 배리어를 삽입합니다.

가중치 전송 파이프라인은 호스트-디바이스 복사, 텐서 조립, 융합 및 양자화, 그리고 MeshGroups 전반의 RDMA 전송을 겹쳐서 처리합니다.

당사는 각 파라미터 텐서의 전송을 하나의 작업으로 취급합니다. 가중치 전송 프로세스는 여러 유형의 하드웨어 리소스를 활용하므로, 가중치 전송 작업을 시간상으로 중첩되는 여러 파이프라인 단계로 나눕니다.

호스트-장치 간 memcpy — FSDP가 가중치를 CPU로 오프로드하는 경우

파라미터 준비 — full_tensor()로 전체 가중치를 재구성하고, 프로젝션 퓨전을 적용하며, 필요 시 양자화를 수행합니다.

RDMA 전송 — 원격 추론 GPU 메모리에 제로 복사 방식으로 기록

전역 배리어 — 모든 full_tensor() 호출이 완료된 후, 이더넷을 통해 GLOO를 사용하여 메시 그룹 간 동기화를 수행합니다.

구현 시 각 파이프라인 단계에 대해 작업 FIFO 큐를 유지합니다. 큐의 맨 앞에 있는 작업이 단계를 완료하면 다음 단계 큐의 맨 뒤로 이동합니다.

full_tensor() 및 기타 GPU 작업은 추가적인 GPU 메모리 사용을 유발합니다. 메모리 부족 오류를 방지하기 위해, 현재 진행 중인 작업이 차지하는 임시 GPU 메모리가 설정된 워터마크보다 적을 때만 작업을 시작합니다.

당사의 시스템은 몇 가지 설계 선택을 통해 일반적인 오픈 소스 솔루션보다 실행 속도가 훨씬 빠르고 유지보수가 용이합니다.

Rank0 기반 가중치 전송은 추론 노드에 분배하기 전에 가중치를 수집합니다.

포인트 투 포인트 통신은 두 학습 노드의 가중치를 추론 노드의 네트워크 인터페이스로 전송합니다.

일반적인 패턴은 모든 파라미터를 rank-0 GPU를 통해 집중시키는 방식입니다. 학습 rank-0에서 수집하고, 추론 rank-0으로 전송한 다음, 다시 분산합니다. 이는 단일 GPU의 PCIe 대역폭과 NIC(예: 400Gbps ≈ 50GB/s)에 의해 제한되어 금방 병목 현상이 발생합니다.

반면, 당사의 지점 간 설정은 모든 학습 GPU가 모든 추론 GPU로 직접 전송할 수 있게 하여 단일 링크가 아닌 전체 네트워크 패브릭을 포화시킵니다.

RPC 기반 전송은 추론 측 버퍼와 워커 호출을 사용하는 반면, 단방향 전송은 추론 워커의 참여 없이 가중치를 직접 씁니다.

일부 시스템은 각 텐서에 대해 추론 엔진의 update_weight() 메서드를 호출하는 방식에 의존합니다. 이는 추론 코드에 대한 침투적 변경을 의미하며, RPC, 직렬화 및 제어 평면 조정으로 인한 오버헤드가 발생합니다.

RDMA WRITE 기본 요소를 사용하면 추가 복사 없이 추론 GPU 메모리에서 가중치를 조용히 업데이트합니다. 제어 평면 메시지나 CPU 제어 로직이 관여하지 않으며, 추론 엔진 수정도 필요하지 않습니다.

가중치 전송 프로세스는 네 가지 유형의 하드웨어 리소스를 활용할 수 있습니다. (1) 호스트-장치 데이터 이동 (2) 프로젝션 퓨전 및 양자화를 위한 GPU 연산 (3) 데이터 평면을 위한 RDMA 네트워크 (4) 제어 평면을 위한 이더넷.

당사의 설계는 가중치 전송 작업을 파이프라인 단계로 나누어 서로 다른 하드웨어 리소스 간의 중첩 실행을 원활하게 합니다.

일부 구현에서는 매 학습 단계마다 전송 일정을 재계산하여 메타데이터를 반복적으로 수집하고 명령을 배포합니다. 이는 불필요한 제어 평면 지연 시간을 추가합니다.

당사의 일정은 초기화 시 한 번만 계산됩니다. 각 학습 반복마다 컨트롤러가 "실행" 신호를 보내면 GPU가 미리 할당된 경로를 따르는 방식으로 계획을 재생합니다. 실행은 예측 가능하고 가볍습니다.

메타데이터 수집, 이름 매칭, 노드 내 수집, 프로젝션 퓨전, 양자화, 서브슬라이싱 통신 월드, 노드 간 네트워크 전송 등 전체 가중치 업데이트 프로세스를 하나의 거대한 함수에 얽히게 만들고 싶은 유혹이 있습니다. 이는 프로그래밍하기 어려울 뿐만 아니라 최적화하기는 더욱 어렵습니다.

엔지니어링 과정에서 당사는 이러한 단계를 개별 구성 요소로 분리했습니다. 각 구성 요소는 독립적으로 단위 테스트를 수행하고, 논리적 검증을 거치며, 최적화할 수 있습니다.

빠르고 신뢰할 수 있는 가중치 전송은 대규모 강화 학습 미세 조정을 위한 핵심 구성 요소입니다. RDMA WRITE 기본 요소, 정적 전송 일정, 파이프라인 실행을 결합하여 Kimi-K2에서 조 단위 파라미터 업데이트를 단 1.3초로 단축했습니다. 이 접근 방식은 논리적으로 이해하기 쉽고 유지보수가 간편하며 전통적인 설계의 병목 현상을 방지합니다.

BibTeX

@online{perplexity2025WeightTransferForRlPostTrainingInUnder2Seconds,
  author = {{Perplexity} and Lequn Chen and Erin McKnight and Vivian Wang},
  title = {{Weight transfer for RL post-training in under 2 seconds}},
  date = {2025-09-24},
  url = {https://www.perplexity.ai/hub/blog/weight-transfer-for-rl-post-training-in-under-2-seconds}
}

RDMA WRITE: 단방향 전송

워크플로우 개요

가중치 전송 실행

DeviceMesh 및 Mesh 그룹

작업 파이프라인

GPU 메모리 사용량 제어

빠르고 단순한 이유

지점 간 통신

단방향 데이터 전송

파이프라이닝

정적 일정

깔끔한 분리

결론

인용