Periodic Asynchrony: An On-Policy Approach for Accelerating LLM Reinforcement Learning

이 논문은 추론과 학습을 분리하여 비동기 파이프라인을 구축하고, 오프-폴리시 편향 없이 온-폴리시 정확도를 유지하며 NPU 환경에서 기존 RL 프레임워크 대비 3~5 배의 학습 처리량을 달성하는 주기적 비동기 방식을 제안합니다.

핵심

이 논문은 거대한 인공지능 모델 (LLM) 을 더 똑똑하게 만드는 과정인 '강화 학습'을 훨씬 더 빠르고 효율적으로 만드는 새로운 방법을 소개합니다.

핵심 아이디어를 한 마디로 요약하면: **"학습 (Training) 과 추론 (Inference) 이 서로 기다리는 시간을 없애고, 마치 공장과 창고가 따로 움직이면서도 완벽하게 같은 결과를 내는 시스템을 만든 것"**입니다.

이 복잡한 기술을 일상적인 비유로 쉽게 설명해 드릴게요.

---

1. 기존 방식의 문제점: "혼자서 일하는 비효율적인 요리사"

기존의 AI 학습 방식은 **한 명의 요리사 (GPU)**가 모든 일을 혼자 하는 것과 같습니다.

1. 재료 준비 (추론): 요리사가 손으로 재료를 다듬고 요리합니다. (AI 가 문제를 풀고 답을 생성)
2. 맛보기 (평가): 그 요리가 맛있는지 심사위원이 평가합니다.
3. 레시피 수정 (학습): 평가 결과를 보고 요리사가 레시피를 고칩니다.
4. 다음 요리: 고친 레시피로 다시 재료를 준비합니다.

문제점: 요리사가 재료를 다듬는 동안은 레시피를 고칠 수 없고, 레시피를 고치는 동안은 재료를 다듬을 수 없습니다. 즉, 한 작업이 끝날 때까지 다른 작업은 완전히 멈춰서 기다려야 합니다. (이를 '동기식'이라고 합니다.)

2. 이 논문이 제안한 해결책: "주기적 비동기 (Periodic Asynchrony)"

이 논문은 "요리사 (학습)"와 "조수 (추론)"를 분리해서 동시에 일하게 하되, 실수하지 않는 방법을 제안합니다.

🏭 비유: 공장 (학습) 과 창고 (추론) 의 협력

이 시스템은 두 개의 팀으로 나뉩니다.

• 조수 팀 (추론): 끊임없이 재료를 다듬고 요리를 만들어 창고에 쌓아둡니다.
• 요리사 팀 (학습): 창고에서 완성된 요리를 하나씩 꺼내서 맛을 보고 레시피를 고칩니다.

기존의 비동기 방식과의 차이점 (핵심!):
기존의 다른 비동기 방식들은 "조수가 만든 요리를 요리사가 바로바로 고쳐서 다음 요리에 쓰면 되잖아?"라고 생각하다가, 과거의 레시피로 만든 요리를 가지고 현재 레시피를 고치는 실수를 저지르기도 합니다. (이걸 '오프-폴리시'라고 하며, AI 가 엉뚱한 방향으로 학습할 위험이 있습니다.)

하지만 이 논문은 **"완벽한 타이밍"**을 잡습니다.

1. 조수 팀이 한 번에 한 번의 레시피로 만든 요리들만 모아서 창고에 쌓습니다.
2. 요리사 팀은 그 요리들을 하나씩 꺼내서 레시피를 고칩니다.
3. 모든 요리가 다 고쳐진 후에야 비로소 새로운 레시피를 정하고, 그 레시피를 조수 팀에게 알려줍니다.

이렇게 하면 조수 팀이 요리하는 동안 요리사 팀은 레시피를 고칠 수 있어 (기다리는 시간 0), 속도가 3~5 배 빨라지지만, 결과는 완벽하게 같은 레시피로 만든 것과 같습니다.

3. 추가적인 꿀팁: "공통된 주문서 (Shared-Prompt Attention)"

AI 가 문제를 풀 때, **문제 (프롬프트)**는 같지만 **답 (응답)**은 여러 개입니다.

• 기존 방식: 문제 10 개를 풀 때, 문제 텍스트를 10 번이나 다시 읽어서 계산합니다. (비효율적)
• 이 논문의 방식: 문제 텍스트는 한 번만 읽고, 그 결과를 10 개의 다른 답을 계산할 때 공유합니다.

이는 마치 10 명이 같은 영화를 볼 때, 스크린을 10 개 켤 필요 없이 한 개의 스크린을 공유하는 것과 같습니다. 특히 문제가 길고 답이 짧을 때 효과가 어마어마하게 큽니다.

4. 실제 성과: "속도만 3~5 배 빨라진 마법"

이 방법을 실제 AI 칩 (NPU) 에서 테스트한 결과:

• 속도: 기존 방식보다 3 배에서 5 배 더 빠르게 학습이 완료되었습니다.
• 정확도: 속도가 빨라졌다고 해서 AI 가 멍청해지거나 엉뚱한 답을 내놓지는 않았습니다. 기존 방식과 완전히 똑같은 성능을 냈습니다.
• 확장성: 컴퓨터를 더 많이 붙여도 속도가 거의 선형적으로 빨라집니다.

📝 한 줄 요약

"AI 가 학습할 때, '생각하는 시간'과 '계산하는 시간'이 겹쳐서 기다리는 시간을 없애고, 똑똑한 조수 시스템을 도입하여 속도는 5 배로, 정확도는 그대로 유지하게 만든 혁신적인 방법입니다."

이 기술은 앞으로 거대 AI 모델을 더 저렴하고 빠르게 발전시키는 데 큰 역할을 할 것으로 기대됩니다.

기술 요약

1. 문제 정의 (Problem)

대규모 언어 모델 (LLM) 의 강화 학습 (RL) 후학습 (Post-training) 에서 GRPO와 같은 알고리즘의 등장으로 RL 에 대한 관심이 높아졌으나, 여전히 훈련 효율성이 주요 병목 현상으로 남아있습니다.

• 동기식 실행의 한계: 기존 주요 RL 프레임워크는 추론 (Inference) 과 훈련 (Training) 을 동일한 장치에서 동기적으로 실행합니다. 이로 인해 추론이 완료될 때까지 훈련이 대기하게 되어, 두 작업이 동시에 수행되지 못하며 자원 활용도가 낮아집니다.
• 계산 오버헤드: RL 훈련은 정책 모델, 이전 정책 모델 (Old-policy), 참조 모델 (Reference) 등 여러 모델의 순전파 (Forward pass) 를 필요로 하며, 특히 긴 문맥 (Long-context) 에서 많은 수의 사고 과정 (CoT) 추론을 생성해야 하므로 메모리 사용량과 계산 비용이 급증합니다.
• 비동기식 접근의 단점: 기존 비동기식 방법들은 추론과 훈련을 분리하여 처리 속도를 높였으나, 이는 **Off-policy 편향 (Off-policy bias)**을 유발하거나 알고리즘을 수정해야 하는 문제가 있었습니다. 이는 GRPO 와 같은 엄격한 On-policy 알고리즘의 수렴성을 보장하기 어렵게 만듭니다.

2. 방법론 (Methodology)

이 논문은 추론과 훈련을 분리하되, On-policy 정확성을 해치지 않는 주기적 비동기 (Periodic Asynchrony) 프레임워크를 제안합니다.

가. 주기적 비동기 프레임워크 (Periodic Asynchronous Framework)

• 프로듀서 - 컨슈머 파이프라인: 데이터 로더와 훈련기 사이에 '임시 데이터 생성기 (Temporary Data Generator)'를 도입하여 비동기 파이프라인을 구성합니다.
  • 프로듀서 (배경 스레드): 데이터 로더에서 배치의 프롬프트를 가져와 추론 워커 (Inference Workers) 로 분배합니다.
  • 컨슈머 (메인 프로세스): 추론 및 보상 평가가 완료된 샘플을 큐에서 가져와 훈련 엔진에 투입합니다.
• 주기성 (Periodicity): 전체 배치의 샘플이 모두 소비될 때까지 가중치 업데이트를 수행하지 않습니다. 즉, **배치 단위 (Batch-level)**로 동기화되므로, 각 배치 내의 모든 샘플은 동일한 최신 정책 ($\pi_{\theta_t}$) 으로 생성됩니다. 이는 On-policy 조건을 엄격히 유지합니다.

나. 통합 3-모델 아키텍처 (Unified Tri-Model Architecture)

• GRPO 훈련 시 필요한 정책 (Policy), 이전 정책 (Old Policy), 참조 (Reference) 세 가지 모델의 로짓 (Logits) 계산을 효율화합니다.
• 세 모델이 동일한 병렬 배치 (Parallel Layout) 를 공유하도록 설계하여, 단일 마이크로 스텝에서 세 모델의 로짓을 동시에 계산합니다.
• 가중치 동기화는 배치 처리 완료 후 한 번에 이루어지며, 이는 메모리 효율성과 일관성을 보장합니다.

다. 공유 프롬프트 어텐션 (Shared-Prompt Attention)

• GRPO 는 동일한 프롬프트에서 여러 응답을 생성하므로, 프롬프트 부분의 계산 중복을 제거합니다.
• 메커니즘:
  • 여러 응답을 하나의 시퀀스로 연결하고, 프롬프트 토큰은 공유합니다.
  • 커스텀 어텐션 마스크: 각 응답 토큰이 프롬프트와 자신의 이전 응답 토큰에만 어텐션 하도록 제한하여, 다른 응답 간의 정보 유출을 방지합니다.
  • 손실 계산: 프롬프트 토큰은 제외하고 응답 토큰에 대해서만 손실을 계산합니다.
• 효과: 긴 프롬프트 환경에서 어텐션 계산 복잡도를 $O(K)$배 감소시킵니다.

3. 주요 기여 (Key Contributions)

1. 이론적 동등성 증명: 제안한 주기적 비동기 방식이 동기식 훈련과 알고리즘적으로 완전히 동등함을 수학적으로 증명했습니다.
   • Proposition 1: 모든 롤아웃 샘플이 동일한 정책에서 생성됨 (On-policy 조건 유지).
   • Proposition 2: 샘플 소비 순서가 그라디언트 업데이트에 영향을 주지 않음 (그라디언트 교환 법칙).
   • Theorem 1: 비동기식 업데이트가 동기식 업데이트와 동일하므로, On-policy RL 알고리즘의 수정 없이도 정확성을 보장합니다.
2. 시스템 최적화: 3D 병렬 아키텍처 기반의 통합 3-모델 설계와 공유 프롬프트 어텐션 메커니즘을 통해 연산 중복을 제거하고 하드웨어 활용도를 극대화했습니다.
3. 알고리즘 중립성: 특정 RL 알고리즘에 의존하지 않으며, GRPO 를 포함한 모든 On-policy 알고리즘에 적용 가능합니다.

4. 실험 결과 (Results)

NPU (Ascend-910B) 플랫폼에서 수행된 실험 결과는 다음과 같습니다.

• 처리량 (Throughput) 향상:
  • 기존 주요 RL 프레임워크 (MindSpeed-RL, VERL 등) 대비 3 배에서 5 배의 엔드 - 투 - 엔드 훈련 처리량 (Tokens Per Second Per Device) 향상을 달성했습니다.
  • 동기식 베이스라인 대비 약 1.9 배 ~ 2 배의 속도 향상 (이론적 상한선에 근접) 을 확인했습니다.
  • 공유 프롬프트 어텐션 (SPA) 적용 시, 훈련 토큰 수 감소와 패딩 오버헤드 제거로 인해 8 배에 달하는 처리량 향상을 보였습니다.
• 정확도 (Accuracy):
  • AIME24, GSM8K 등 수학 추론 태스크에서 동기식 훈련과 동등한 정확도를 유지했습니다.
  • 훈련 단계별 보상 점수 (Reward Score) 곡선이 동기식 및 비동기식 방법 간 거의 완전히 겹쳐, 이론적 동등성이 실증적으로 입증되었습니다.
• 확장성 (Scalability):
  • 장치 수 (16, 32, 64 NPU) 가 증가함에 따라 **거의 선형 (Near-linear)**으로 확장되는 것을 확인했습니다. 추론과 훈련 인스턴스를 독립적으로 스케일링할 수 있어 병목 현상을 효과적으로 해결했습니다.

5. 의의 및 결론 (Significance)

이 연구는 On-policy 강화 학습의 효율성 병목을 해결하는 획기적인 접근법을 제시합니다.

• 이론적 안전성: 기존 비동기 방식이 가진 'Off-policy 편향'이나 '알고리즘 수정'의 리스크 없이, 엄격한 On-policy 조건을 유지하면서 비동기 가속을 가능하게 했습니다.
• 실용적 가치: 추론과 훈련의 분리를 통해 하드웨어 자원의 활용도를 극대화하여, 대규모 LLM 의 RL 후학습 비용을 획기적으로 절감할 수 있습니다.
• 범용성: 제안된 프레임워크는 GRPO 뿐만 아니라 다른 On-policy 알고리즘에도 적용 가능하며, 커뮤니티가 이론적 보장을 유지한 채 비동기 가속의 이점을 누릴 수 있는 기반을 마련했습니다.

결론적으로, 이 논문은 **주기적 비동기성 (Periodic Asynchrony)**을 통해 추론과 훈련의 대기 시간을 최소화하면서도 알고리즘적 정확성을 완벽하게 보존하는 새로운 RL 훈련 패러다임을 정립했습니다.