핵심 요약
LinkedIn AI 팀은 GPT-OSS 모델을 에이전틱 애플리케이션의 백본으로 활용하기 위해 강화학습(RL) 학습 과정에서의 불안정성 문제를 해결했다. 초기 학습에서 발생한 그래디언트 폭주와 보상 정체 현상의 근본 원인이 추론 엔진과 학습 스택 간의 어텐션 싱크(Attention Sink) 지원 불일치임을 밝혀냈다. 이를 해결하기 위해 FlashAttention v3에 어텐션 싱크 역전파(Backward Pass)를 직접 구현하고, MoE 아키텍처의 비결정론적 라우팅 문제를 보완했다. 결과적으로 GSM8K, ReTool 등 복잡한 에이전트 작업에서 안정적인 학습과 빠른 수렴 성능을 확보했다.
배경
강화학습(PPO) 메커니즘, MoE(Mixture of Experts) 아키텍처, FlashAttention 커널 이해, FSDP(Fully Sharded Data Parallelism)
대상 독자
대규모 언어 모델(LLM) 강화학습 및 에이전트 시스템을 구축하는 머신러닝 엔지니어
의미 / 영향
이 연구는 오픈소스 MoE 모델을 실제 에이전트 환경에서 학습시킬 때 발생하는 실무적 난관과 해결책을 제시하여 고성능 에이전틱 모델 개발의 기술적 장벽을 낮추는 데 기여한다. 특히 어텐션 싱크의 학습 지원은 향후 다양한 모델의 안정성 향상에 표준적인 기법이 될 가능성이 높다.
섹션별 상세
에이전틱 RL은 단일 응답이 아닌 환경과의 상호작용을 통한 다단계 의사결정 과정을 최적화하며, GPT-OSS 모델의 적합성을 검증하기 위해 verl 프레임워크를 사용했다.
MoE 아키텍처의 특성상 동일 입력에 대해서도 라우팅 결과가 미세하게 달라져 PPO 학습 시 로그 확률 불일치가 발생하며, 이를 해결하기 위해 온폴리시(On-policy) 상황에서 중요도 샘플링 비율을 1로 강제 고정하는 패치를 적용했다.
학습 불안정성의 핵심 원인은 추론 시 사용되는 어텐션 싱크(Attention Sink)가 학습용 FlashAttention v2에서 지원되지 않는 학습-추론 불일치였으며, 이를 위해 FlashAttention v3 기반의 역전파 커널을 직접 구현했다.
어텐션 싱크는 소프트맥스 계산 시 학습 가능한 스칼라 파라미터를 가상 토큰처럼 활용하여 어텐션 가중치를 안정화하며, 이를 통해 GSM8K와 같은 수학 추론 및 도구 사용 작업에서 수렴 속도를 획기적으로 개선했다.
FSDP 환경에서 MoE 전문가(Expert)를 병렬로 구동할 때 발생하는 메모리 폭주 문제를 해결하기 위해, 전문가를 순차적으로 처리하는 메모리 효율적 실행 경로로 Hugging Face 코드를 수정했다.
긴 컨텍스트를 유지해야 하는 에이전트 특성을 고려하여 어텐션 싱크와 호환되는 시퀀스 병렬 처리(Sequence Parallelism)를 도입해 GPU당 활성화 메모리 점유율을 낮추고 확장성을 확보했다.
</> 코드 예제 포함
실무 Takeaway
- 추론 엔진(vLLM/SGLang)과 학습 프레임워크 간의 커널 구현 차이인 어텐션 싱크 유무가 RL 학습 붕괴의 치명적 원인이 될 수 있다.
- MoE 모델의 RL 학습 시 비결정론적 라우팅으로 인한 PPO 클리핑 오류를 방지하기 위해 로그 확률 계산 로직의 정밀한 조정이 필요하다.
- 대규모 MoE 모델 학습 시 메모리 효율을 위해 전문가 행렬 연산을 순차 처리 방식으로 변경하고 시퀀스 병렬 처리를 적극 활용해야 한다.
AI 분석 전체 내용 보기
AI 요약 · 북마크 · 개인 피드 설정 — 무료