Nereus가 LLM 강화학습 후학습의 병렬화 계획을 실행 중에 바꾼다
Nereus: Adaptive Parallelism for LLM Post-Training
무엇인가
LLM 강화학습(RL) 후학습은 액터, 크리틱, 리워드, 레퍼런스 모델을 생성·추론·학습 세 단계에 걸쳐 GPU 클러스터 위에서 함께 돌리는 분산 작업이다. 문제는 실행 중에 조건이 계속 변한다는 점이다. 논문은 이를 드리프트(drift)라고 부르며 세 가지 원인을 든다. 첫째 자원 공급이다. AWS 트레이스 분석에서 16-GPU 작업이 10시간 동안 10회 이상 가용성이 바뀌었고, Llama-3.1-8B의 최적 TP/PP 배치는 32 GPU와 256 GPU에서 서로 다르다. 둘째 작업 부하다. 액터가 학습하면서 생성 시퀀스 길이가 1,000 스텝 안에 500토큰에서 8,000토큰으로 늘어나 KV 캐시와 활성값 메모리를 압박한다. 16 GPU에서 (4,2,2) 배치는 2K 토큰에서 가장 빠르지만 4K에서 OOM이 나고, (4,4,1)은 4K에서 살아남는 대신 2K에서 통신 오버헤드가 붙는다. 셋째 실효 하드웨어 효율이다. 네트워크 혼잡과 열 스로틀링 때문에 같은 배치의 스텝 지연이 달라지며, 오프라인 예측기는 8 GPU에서 실행 불가능한 계획을 고르고 64~256 GPU에서 최적 실측 대비 최대 1.56배 느린 계획을 고르기도 한다. 기존 RL 프레임워크는 시작 시점의 할당과 병렬화를 고정하고, 탄력적 학습 시스템은 모델 하나만 다루며, 체크포인트 기반 재샤딩은 작업을 재시작한다. DynaRL은 고정 풀 안에서만 단계별 자원을 재할당한다.
어떻게 동작하나
Nereus의 첫 번째 축은 비용 인식 적응 정책이다. 모니터가 시퀀스 길이, 가용 GPU 풀, 피크 메모리, 달성된 연산·통신 효율을 읽고 두 종류의 트리거를 건다. GPU 풀 변화나 메모리 위반 예측은 즉시 재계획을 부르고, 드리프트 신호가 마지막 재계획 시점의 기준값과 상대 임계값 δ 이상 차이 나면 재계획을 건다. 실험에서 스텝당 시퀀스 길이에 δ_len=30%를 썼다. 재계획 단계에서는 모델-스테이지마다 TP/PP/DP와 GPU 집합 후보를 두고, 메모리 용량과 단계별 GPU 예산 제약 아래 예측 스텝 지연을 최소화하는 전역 계획을 고른다. 비용 모델은 NanoFlow와 Alpa의 계획 기법을 다중 모델로 확장한 것으로, 연산 시간은 피크 연산량에서, 집합 통신 시간은 경로 대역폭에서 추정하고, GEMM·어텐션·all-reduce 같은 연산 클래스별로 실측 대비 효율을 온라인 학습한다. 실행 단계에서는 실행 가능성과 수익성을 따로 본다. 수익성 조건은 ΔL>0 이면서 전환 비용 대비 회수 기간 L_tran/ΔL이 γH_x 이하일 때다. γ는 기본 0.5이고, H_x는 해당 신호가 다음 임계값을 넘을 때까지의 예상 스텝 수다. 즉 전환 비용을 절약분으로 갚을 수 있다고 판단될 때만 계획을 바꾼다.
무엇과 다른가
두 번째 축은 의존성에 맞춘 상태 추상화인 EMU(Elastic Model Unit)다. EMU는 모델-스테이지 복제본 하나를 나타내는 상태 단위로, 내부에 TP/PP 레이아웃과 파라미터 샤드, 옵티마이저 샤드를 담고 DP는 복제본 개수로 표현한다. DP 샤딩은 단위 안에 넣지 않으므로(ZeRO-0) DP 확장은 단순히 단위를 늘리거나 줄이는 일이 된다. 계획 차이는 Split, Merge, Extend, Destroy 네 가지 원시 연산으로 컴파일된다. 세 번째 축은 전환 오케스트레이션이다. 전환 엔진은 전역 전환 DAG를 만들어 EMU 변환과 GPU 전송 순서를 정하고, 일시적으로 GPU가 겹치는 경우 자원 의존성 간선을 추가한다. GPU 여유가 없으면 한 모델-스테이지가 자원을 놓아야 다른 쪽이 확장할 수 있으므로, 자원을 해제하는 연산을 우선순위로 올려 교착을 피한다. 구현은 Python, C++, CUDA/HIP 43,000줄이며 vLLM, DeepSpeed, Megatron-LM과 통합되고 NCCL/RCCL 집합 통신을 쓴다.
어떻게 쓰나
계획 품질과 전환 비용은 별도로 측정됐다. Cluster #2에서 8~256 GPU, 시퀀스 길이 1K/2K/4K 조합 18개 설정의 유효 계획 480개를 전부 실행해 실측 최적과 비교했는데, 61.1%에서 정확히 최적을 맞추고 모든 설정에서 5% 이내에 들었다(중앙값 격차 0%, p95 4.1%). 비용 모델은 별도 실행에서 얻은 커널·집합 통신 타이밍으로만 보정한 72개 계획에 대해 평균 절대 백분율 오차 4.96%, 최대 오차 14.61%, R²=0.9993을 기록했다. 재계획 결정 오버헤드는 32 GPU에서 0.17ms, 1,024 GPU에서 338ms로, SCIP 기반 솔버의 3.4ms~511초와 대비된다. 재계획 메모리는 100MB 미만이다. 전환 DAG 계획은 1,024 GPU에서 1.22ms로 SCIP의 120.10초보다 98,443배 빠르면서 스케줄 완료 시간 격차는 6.9% 이내였다. Extend 비용은 4→8 GPU에서 32→64 GPU까지 2.45~8.48초로, Oobleck보다 3.8~9.9배, Gemini보다 6.7~10.8배, Tenplex보다 9.3~16.2배, UCP보다 115.7~284.1배 빠르다. 70B 모델 64 GPU에서 재샤딩 비용은 MCP 대비 99.1%, Tenplex 대비 96.7% 줄었다. 8B에서는 8~64 GPU 구간에서 15.7~16.1초로 안정적인 반면 Tenplex는 62.7~65.0초, MCP는 329.9~486.5초가 걸렸다. GPU가 겹치는 전환 조정 실험에서 Nereus는 100회 시행 전부 성공했고, DynaRL식 컴포넌트별 마이그레이션은 34~62%만 성공했다.
전제와 한계
종단 간 성능은 8B PPO 기준으로 OpenRLHF 대비 2.14~7.27배(중앙값 3.99배), Verl 대비 1.10~1.47배(중앙값 1.21배)다. 32 GPU에서 1,024 GPU로 늘릴 때 스텝 지연이 15배 줄어 OpenRLHF의 10배를 앞서고, 1,024 GPU에서 OpenRLHF 대비 3.20배다. 실데이터로 만든 트레이스에서 온라인 TP/PP 적응은 고정 TP/PP 레이아웃 대비 평균 스텝 지연을 27.7% 줄였고(861.9초 대 1,191.6초), DynaRL식 수용보다 8.3%, 3스텝마다 무조건 전환하는 정책보다 2.5% 낮았다. 홀드아웃 트레이스 3개에서는 평균 858.7초/스텝으로 3-step adapt의 873.8초, DynaRL식의 928.3초보다 좋았다. 1,024 GPU까지 확장하는 1,000 스텝 실행에서 전환 6회는 전체 62,353초 중 49.5초, 0.079%만 썼다. ReMax는 최대 13.3%, GRPO는 최대 15.8% 지연을 줄였고, 완전 비동기 RL에서는 Laminar보다 최대 37.3% 낮은 스텝 지연을 보였다. 학습 동작도 유지됐다. 50 스텝 동안 보상과 PPO KL 추정치가 Verl과 비슷하게 따라갔고(스텝 50에서 보상 0.9199 대 0.9160, 평균 KL 4.9e-5 대 7.4e-5), 같은 50 스텝을 3,141초에 끝내 Verl의 3,646초보다 13.9% 빨랐다.
실무 관점에서 이 논문이 주는 것은 "언제 계획을 바꿀 것인가"에 대한 판단 규칙이다. 전환 자체는 GPU 상주 상태를 재사용하므로 체크포인트 재시작보다 훨씬 싸지만, 공짜는 아니다. Nereus는 절약분 ΔL과 전환 비용 L_tran의 비율을 회수 기간으로 환산해 γH_x와 비교하고, γ=1로 두면 전환이 14회로 늘어나 오히려 지연이 939.5초로 나빠진다는 것을 보여준다. 즉 무조건 적응하는 정책은 손해다. 또한 적응은 동기 실행에서는 RL 스텝 완료 시점, 비동기 실행에서는 가중치 동기화 시점이라는 안전 경계에서만 일어나며, 진행 중인 롤아웃은 자기 정책 버전을 유지한다. 자체 RL 파이프라인에 이런 적응 계층을 얹으려면 프레임워크 계측으로 커널·집합 통신 시간을 노출하고, 안전 경계를 정의하고, 상태 단위를 모델-스테이지 복제본 수준으로 잡을 수 있는지부터 확인해야 한다.
저자들이 밝힌 전제와 한계도 분명하다. EMU는 DP 샤딩을 단위 내부에 두지 않는 ZeRO-0 가정 위에 서 있고, 적응은 안전 경계에서만 가능하다. 메모리·GPU 풀 검사를 통과하는 후보가 하나도 없으면 더 작은 복제본 배치와 재계산을 고려하고, 그래도 실패하면 체크포인트/재시작으로 물러난다. GPU 상실 시에는 죽은 랭크를 포함한 통신자를 중단하고 살아남은 단위에서 재계획하지만, 마지막 사본까지 잃은 상태는 별도 결함 허용 시스템이 복구해야 한다. 확장 방향으로는 컨텍스트/전문가 병렬화, 외부 도구 서비스 오토스케일링, 다른 RL 프레임워크 지원을 남겨 두었다.