LongStraw: 고정 GPU 예산으로 백만 토큰급 RL 포스트트레이닝을 실행하는 아키텍처 인식 스택

대규모 에이전트·대화·문서 누적 관찰을 요구하는 응용에서 추론 컨텍스트 길이는 이미 백만 토큰에 근접하고 있으나 RL 포스트트레이닝 워크로드는 여전히 훨씬 짧은 컨텍스트를 전제로 설계되는 경우가 많아 실전 적용에 갭이 생겼다. LongStraw는 이 갭을 메모리 예산 내에서 메우기 위해 실행 스택 차원에서 프롬프트 평가와 역전파 그래프를 축소하는 전략을 도입해 고정 GPU 자원으로도 백만 토큰급 학습 단계를 가능하게 했다. 이로 인해 에이전트의 긴 궤적을 학습하는 환경에서 기존 방식으로는 불가능했던 규모의 포스트트레이닝 실행이 실험적으로 확인되었다.

딥러닝에서 Self-Attention 기반 모델은 시퀀스 길이에 따라 메모리와 연산이 크게 증가하는 문제가 있으며, RL 포스트트레이닝에서는 길이 일반화에 의존하는 대신 실제 긴 관찰 궤적을 직접 처리해야 하는 수요가 늘어났다. 이 논문은 긴 컨텍스트를 그대로 autograd로 유지하면 라이브 트레이닝 그래프가 메모리 한계를 초과한다는 점을 출발점으로 삼아, 공유된 프롬프트의 중간 계산을 autograd에서 분리하고 역전파 시 필요한 부분만 재연산하는 방식으로 메모리 부담을 줄이는 해결책을 제시했다. 이렇게 하면 전체 시퀀스의 모든 중간 상태를 메모리에 저장하지 않아도 되고 대신 재연산으로 그래디언트를 복원함으로써 고정된 GPU 예산으로 훨씬 긴 컨텍스트를 처리할 수 있다.

전체 접근은 아키텍처 인식 실행 스택을 설계해 처리 단위를 그룹으로 묶고 그룹 단위로 연산 경로를 분리하는 것이다. 구체적으로 공유된 프롬프트에 대한 전방 계산은 autograd에서 분리되어 저장되지 않으며, 모델별로 이후 토큰 계산에 필수적인 내부 상태만 보존한다. 응답은 짧은 브랜치 단위로 순차적으로 재생되어 각 브랜치에서 필요한 역전파를 수행하므로 라이브 그래프의 크기는 줄어들지만 재연산 시간이 증가하는 트레이드오프가 발생한다.

주요 실험에서 LongStraw는 8개의 H20 GPU 환경에서 Qwen 기반 파이프라인의 그룹화된 스코어링과 응답 역전파를 2.1M 포지션에서 성공적으로 완료했으며, 그룹 크기를 늘려도 피크 할당 메모리가 그룹 크기 증가에 따라 0.21GB만 추가로 증가하는 것을 관찰했다. 스트레스 테스트에서는 4.46M 포지션까지 실행 용량이 확장되었고, 32개의 H20 GPU에서는 GLM-5.2의 78개 레이어 전역에 걸쳐 2.1M 토큰의 end-to-end 실행 경로가 검증되었다. 논문은 이러한 실험 결과가 '실행 용량(execution capacity)'을 확립한 것이지 완전한 학습적 정합성(training correctness)을 보장한 것은 아니며, 캡처된 프롬프트 상태가 분리(detached)되어 있고 일부 분산된 전방 및 그래디언트 구성 경로가 미완성 상태라고 명시했다.

전체 아키텍처는 '아키텍처 인식 실행 스택' 개념을 중심으로 구성되며, 입력 시퀀스를 그룹으로 묶어 공유된 프롬프트 부분과 모델별로 필요한 내부 상태를 분리 관리하는 실행 파이프라인을 포함한다. 공유된 프롬프트의 전방 연산은 autograd에서 분리되어 저장되지 않으며, 역전파를 위해서는 각 응답 브랜치를 별도로 재연산해 필요한 그래디언트를 복원한다. 이 구조는 메모리 사용량을 감소시키는 대신 재연산(리플레이) 비용을 도입하는 명확한 트레이드오프를 수반한다.

논문은 제시된 실험이 '실행 용량'을 확립하는 데 초점을 맞추었으며 완전한 학습적 정합성을 보장하지 않는다고 명시했다. 그 이유는 캡처된 프롬프트 상태가 autograd에서 분리(detached)되어 있고 일부 분산된 전방 및 그래디언트 구성 경로가 아직 완성되지 않았기 때문이다. 따라서 LongStraw의 결과는 대규모 컨텍스트의 실행 가능성을 입증했지만 실제 학습 결과의 완전성 검증을 위해 추가적인 분산 경로 보강과 학습적 정합성 평가가 필요하다.