논문Hugging Face
LongStraw: Long-Context RL Beyond 2M Tokens under a Fixed GPU Budget
논문 ID: 2607.14952166 추천
#LLM#RL#Long-Context#Efficiency#Agent#Evaluation#Inference
핵심 요약
고정된 GPU 자원 내에서 수백만 토큰 규모의 RL 포스트 트레이닝을 가능하게 하는 아키텍처 중심의 실행 스택 제안
상세 내용
추론 시스템의 컨텍스트 길이는 급격히 늘어나는 반면, RL 포스트 트레이닝은 여전히 상대적으로 짧은 길이에 머물러 있어 에이전트 학습에 제약이 발생하고 있습니다. 이를 해결하기 위해 LongStraw는 고정된 GPU 예산 내에서 백만 토큰 이상의 RL 학습을 지원하는 아키텍처 인지형 실행 스택을 제안합니다. 이 방식은 공유 프롬프트를 autograd 없이 평가하고, 이후 토큰에 필요한 모델 상태만 유지하며, 짧은 응답 브랜치를 순차적으로 재실행하여 활성 학습 그래프 크기를 줄입니다. 실험 결과, 8대의 H20 GPU 환경에서 2.1M 위치의 그룹 스코어링 및 응답 역전파를 성공적으로 수행했습니다. 또한 32대의 H20 GPU를 통해 2.1M 토큰 프롬프트에 대한 엔드투엔드 실행 경로를 검증했습니다.
주요 내용
- 고정된 GPU 예산 내에서 수백만 토큰 규모의 RL 학습을 가능케 하는 실행 스택 설계
- 공유 프롬프트의 autograd 제외 및 모델 상태 압축을 통한 메모리 효율화
- 응답 브랜치의 순차적 재실행(Replay)을 통한 활성 학습 그래프 크기 최적화
실무에서 볼 만한 점
에이전트와 같이 매우 긴 컨텍스트를 다루는 모델을 학습시킬 때, 하드웨어 증설 없이도 메모리 문제를 해결하며 초장거리 컨텍스트 학습을 시도할 수 있는 방법론을 제시합니다.
참고할 행동
- GRPO와 같은 RL 알고리즘에 LongStraw의 상태 유지/재실행 메커니즘 적용 실험
- 모델 아키텍처(Attention vs Recurrent)에 따른 메모리 절감 효과 비교
- 프롬프트 상태 분리(Detached state) 방식의 학습 수렴도 영향 분석