Hybrid Flow

HybridFlow: RLHF를 위한 하이브리드 실행 프레임워크

RLHF는 하나의 모델만 학습하는 과정이 아니다. actor가 응답을 만들고, critic과 reference policy, reward model이 그 결과를 계산한 뒤 actor와 critic이 다시 갱신된다. PPO를 기준으로 해도 생성, 준비, 학습 단계가 반복되고, Safe-RLHF나 ReMax처럼 알고리즘이 달라지면 필요한 모델과 데이터 흐름도 함께 달라진다.

이런 환경에서는 각 모델의 분산 계산 성능만 높여서는 부족하다. 모델 사이 데이터 이동, 실행 순서, GPU 배치까지 함께 다뤄야 한다. HybridFlow는 RLHF를 여러 분산 프로그램이 얽힌 데이터플로로 보고, 이 흐름을 더 유연하게 다루는 방법을 제안한다.

기존 방식에서 생기는 문제

기존 RLHF 시스템은 대체로 multi-controller 방식에 기대고 있다. 각 GPU worker가 자기 제어 흐름을 갖고 모델 계산과 데이터 전달을 직접 맞춘다. 분산 학습 엔진과 서빙 엔진에는 잘 맞지만, RLHF 전체 흐름을 바꾸기에는 손이 많이 간다.

actor에서 만든 응답을 critic과 reward model로 넘기는 과정, 학습 단계에 맞춰 실행 순서를 맞추는 과정, 모델마다 다른 병렬화 전략을 두는 과정이 같은 코드 안에 섞이기 쉽다. 알고리즘 하나를 바꾸면 관련 모델 코드까지 함께 손봐야 하는 이유다.

 

 

비교 그림에는 기존 multi-controller 구현과 HybridFlow가 나란히 놓여 있다. 기존 방식은 계산 코드와 데이터 전달 코드가 같은 프로그램 안에 얽혀 있고, HybridFlow는 모델 내부 계산과 모델 사이 조정을 다른 계층으로 나눠 둔다.

HybridFlow가 나누는 계층

모델 내부 계산과 모델 사이 흐름을 따로 다루는 구조다. 모델 내부 계산은 각 모델이 맡는다. 학습, 추론, 생성은 기존 분산 엔진을 그대로 활용한다. 모델 사이 흐름은 상위 컨트롤러가 관리한다. 사용자는 actor의 생성, critic의 value 계산, reward model의 점수 계산처럼 상위 단계만 적으면 되고, 데이터 재배치와 전송은 transfer protocol이 맡는다.

이 구조는 ResourcePool, ParallelWorker, Auto Mapping, 3D-HybridEngine으로 이어진다. 자원 풀은 실제 GPU를 추상화하고, ParallelWorker는 분산 모델 실행을 감싼다. Auto Mapping은 어떤 모델을 같은 장치에 둘지, 각 모델에 몇 개의 GPU를 줄지 고른다. 3D-HybridEngine은 actor의 학습 단계와 생성 단계 사이 전환 비용을 줄인다.

 

 

구성도에는 입력된 데이터플로 그래프와 모델 설정이 자원 풀, 자동 배치, 학습·생성 엔진으로 이어지는 흐름이 정리돼 있다. RLHF 전체를 하나의 데이터플로로 묶고, 필요한 실행 계층을 그 위에 얹는 방식이다.

3D-HybridEngine이 다루는 전환 비용

RLHF에서 actor는 생성과 학습을 모두 맡는다. 생성은 메모리 사용량의 영향을 크게 받고, 학습은 계산량의 영향을 크게 받는다. 같은 병렬화 전략을 두 단계에 그대로 쓰면 한쪽 효율이 내려가기 쉽다. 학습에서는 큰 model parallel이 유리하고, 생성에서는 더 작은 model parallel과 더 큰 data parallel 조합이 처리량을 끌어올릴 수 있다.

문제는 전환 비용이다. 기존 방식은 학습용과 생성용 actor를 따로 두거나, 단계가 바뀔 때 전체 가중치를 다시 모은 뒤 나눠야 한다. 70B actor에서는 반복마다 140GB의 가중치를 옮겨야 할 수 있어 전환 시간이 크게 늘어난다.

 

 

그림의 위쪽은 학습에 쓰던 가중치를 넓게 모은 뒤 생성용으로 다시 나누는 흐름이다. 아래쪽은 micro DP group 안에서 필요한 조각만 모아 전환 범위를 줄인 흐름이다. HybridFlow는 이 방식을 써서 메모리 중복을 없애고 통신량을 낮춘다.

배치도 고정하지 않는다

RLHF에서는 모든 모델을 같은 GPU에 모아 두는 편이 나은 경우가 있고, actor와 critic을 나눠 병렬로 돌리는 편이 나은 경우도 있다. 작은 클러스터에서는 모델을 함께 두는 쪽이 GPU를 쉬지 않게 쓰기 쉽다. GPU 수가 커지면 단계별 병렬 실행이 더 유리해진다. HybridFlow는 모델 크기와 workload를 같이 보고 이 배치를 자동으로 고른다.

 

 

배치 비교 그래프에서는 GPU 수가 늘수록 유리한 배치가 달라진다. 13B actor·reference와 70B critic·reward를 둔 실험에서는 64 GPU까지 colocation이 평균 44.8% 높은 처리량을 냈고, 96 GPU에서는 split이 더 높았다. 128 GPU에서는 actor·reference·reward를 64 GPU에 두고 critic을 나머지 64 GPU에 두는 조합이 가장 빨랐다.

실험에서 나온 처리량 변화

평가는 PPO, ReMax, Safe-RLHF를 대상으로 7B부터 70B까지의 모델과 8개부터 128개까지의 GPU 환경에서 진행됐다. PPO 기준으로는 DeepSpeed-Chat 대비 평균 3.67배, OpenRLHF 대비 평균 3.25배, NeMo-Aligner 대비 평균 12.52배 높은 처리량이 나왔다. 최대 차이는 20.57배였다.

속도 차이는 한 지점에서만 나오지 않는다. 모델별로 다른 병렬화 전략을 쓰고, actor 전환 비용을 줄이고, 클러스터 크기에 맞게 배치를 바꾼 결과가 겹친 값이다. 8 GPU에서도 최소 2.09배 앞섰고, 7B 모델을 128 GPU에 올린 경우에도 PPO에서 OpenRLHF보다 1.68배 높았다.

 

 

PPO 처리량 그래프에서는 GPU 수가 늘어날수록 기준 시스템도 함께 빨라지지만, HybridFlow의 증가 폭이 더 크게 유지된다. 70B 구간에서 평균 속도 차이가 크게 벌어진 이유도 학습과 생성 사이 전환 비용이 크게 줄었기 때문이다.

HybridFlow는 RLHF를 개별 모델 학습 코드의 묶음으로 다루지 않는다. 여러 모델의 분산 계산, 모델 사이 데이터 이동, 학습과 생성 전환, GPU 배치를 하나의 데이터플로 문제로 묶는다. 알고리즘이 바뀌어도 흐름을 다시 적기 쉽고, 분산 실행 효율도 함께 끌어올리는 방향이다.