LongStraw - Long-Context RL Beyond 2M Tokens under a Fixed GPU Budget

C. Zhou, K. Liu, Y. Zhou, Q. Qiao, et al., "LongStraw: Long-Context RL Beyond 2M Tokens under a Fixed GPU Budget," arXiv:2607.14952, 2026.

저자

논문을 낸 조직은 Andrew Chen이 창업한 Macaron AI의 연구 부문 MindLab입니다. Andrew Chen은 MIT 출신으로, 개인화 AI 에이전트 서비스를 만들다가 "실제 경험에서 학습하는 AI"를 연구 목표로 2025년 말 Mind Lab을 출범시켰습니다. 최고연구책임자는 마샤오텅(马骁腾, Xiaoteng Ma, 공개명 Pony Ma)으로, 청화대 자동화과에서 강화학습을 전공한 뒤 합류했습니다.

1저자 Changhai Zhou는 Fudan University와 MindLab을 함께 적을 두고 있습니다. LongStraw는 Macaron의 대규모 RL 훈련 인프라 MinT 위에서 동작하는데, MinT가 "모델을 계속 띄운 채로 LoRA 어댑터만 교체하며 정책을 관리"하는 시스템이라면, LongStraw는 그 훈련 트랜잭션 안에서 2M 토큰 이상의 긴 컨텍스트를 실제로 처리하는 실행 스택을 맡습니다.

배경

추론 시스템은 이미 100만 토큰짜리 컨텍스트를 처리합니다. 그런데 RL 사후 훈련(post-training)은 보통 256K 이하에 머물러 있습니다. 왜 간극이 이렇게 벌어졌을까요?

추론과 훈련의 메모리 모델이 근본적으로 다릅니다. 추론 서버는 프롬프트를 prefill한 뒤 KV 캐시만 남기고 forward 그래프를 버립니다. 훈련에서는 같은 프롬프트에 조건화된 여러 응답을 채점하고 역전파까지 해야 합니다.

GRPO(Group Relative Policy Optimization)의 손실 함수를 그대로 보면 이 문제가 드러납니다.

\[\mathcal{L}_{policy} = -\frac{1}{G}\sum_{i=1}^{G}\frac{1}{R_i}\sum_{t=1}^{R_i}\min\!\left(\rho_{i,t}A_i,\; \text{clip}(\rho_{i,t},\, 1-\epsilon,\, 1+\epsilon)A_i\right)\]

\(\rho_{i,t}(\theta) = \exp(\log\pi_\theta(y_{i,t}|x_{1:P}, y_{i,<t}) - \log\pi_{old}(y_{i,t}|x_{1:P}, y_{i,<t}))\)

\(G\)개의 응답 각각이 \(P\)개의 프롬프트 토큰에 조건화됩니다. 기존 방식은 \(G\)개의 응답 그래프 전부를 프롬프트 그래프와 함께 메모리에 유지했습니다. 여기에 FlashAttention과 QLoRA, 그래디언트 체크포인팅을 다 동원해도, 프롬프트 그래프와 응답 그래프 \(G\)개가 동시에 살아 있는 구조 자체가 GPU 메모리를 포화시킵니다.

2M 토큰 수준에서는 이 문제가 단순한 최적화 대상이 아닙니다. 스케일아웃 방식(GPU 수를 늘리는 Ring Attention, DeepSpeed-Ulysses 등)은 이미 존재합니다. LongStraw가 묻는 것은 다른 질문입니다. GPU 수를 고정하면 GRPO 실행이 어디까지 갈 수 있는가.

어떻게 만들었나

LongStraw의 핵심 원칙은 하나입니다. "이후 응답 토큰이 필요로 하는 텐서만 프롬프트 경계를 넘겨라."

이를 구현하면 하나의 훈련 트랜잭션이 네 단계로 나뉩니다.

1단계 (프롬프트 캡처): 프롬프트 \(x_{1:P}\)를 autograd 없이(no-grad) 실행합니다. 각 레이어마다 모델 고유의 조건부 상태만 저장하고, 트랜지언트 히든, FFN 중간값, attention scratch, MoE 라우팅 버퍼는 즉시 해제합니다.

2단계 (사전 채점): 캡처한 프롬프트 상태를 읽기 전용으로 고정합니다. 응답마다 old-policy와 reference-policy의 로그 확률을 미리 계산해 둡니다. 파라미터는 이 단계 내내 변하지 않습니다.

3단계 (응답 재생): 각 응답을 순서대로 하나씩 꺼내, autograd 활성화 상태에서 짧은 응답 그래프만 재구성합니다. 역전파 후 그 그래프는 즉시 해제하고 그래디언트만 누적합니다.

4단계 (옵티마이저 호출): \(G\)번 역전파가 끝나면 한 번의 옵티마이저 호출로 파라미터를 업데이트합니다. 이 시점 이후 캡처된 프롬프트 상태는 오래된(stale) 것이 됩니다.

이 설계가 바꾸는 것은 라이브 autograd 그래프의 지배적 규모입니다.

\[M_{live} \approx M_{fixed} + M_{prompt}(P) + M_{grad} + \max_i M_{branch}(R_i) + M_{score}\!\left(\textstyle\sum_i R_i\right)\]

\(M_{branch}\)가 모든 응답을 동시에 잡는 대신 가장 긴 응답 하나에만 묶이는 구조입니다. 직렬 재생이 시간은 늘리지만, 프롬프트 그래프와 응답 그래프 \(G\)개가 동시에 메모리를 점유하는 상황을 피합니다.

두 모델, 두 구현

LongStraw는 두 가지 아키텍처에 대해 각각 구현됩니다. 차이는 "프롬프트 경계에서 무엇을 유지하는가"에서 나옵니다.

Qwen3.6-27B (밀집 하이브리드, 8 H20): 64개 레이어 중 48개는 재귀 GDN(Gated DeltaNet), 16개는 풀-어텐션 레이어입니다. GDN 레이어의 경계 상태는 프롬프트 길이에 무관하게 고정 크기이고, 풀-어텐션 레이어의 KV 페이지만 \(O(P/8)\)로 선형 증가합니다. LongStraw는 이 KV 페이지를 물리적으로 제대로 소유된 컴팩트 페이지로 복사하여 GPU에 유지합니다. 응답 재생 시 CP8에 걸친 글로벌 LSE/출력 merge로 완전한 풀-어텐션 연산을 재구성합니다.

GLM-5.2 (MoE, 32 H20): 78개 레이어, 256개 라우팅 전문가(top-8 + 1 shared), MLA/DSA 어텐션 구조입니다. MoE 구조에서는 단일 BF16 히든 버퍼 하나([524,288, 6,144] 형태)가 \(6\) GiB를 차지합니다. 기존 풀-시퀀스 방식에서는 여기에 DSA scratch, expert 출력, permutation, LoRA 중간값이 겹쳐 메모리가 포화됩니다. LongStraw는 MLA latent 페이지와 인덱서-키 페이지(21개 레이어분)를 CPU에 두고, 응답 재생 시 레이어 하나씩 GPU로 스테이징한 뒤 전체 레이어 체크포인팅으로 역전파합니다.

결과

G를 2에서 8로 늘리면 메모리 증가는 +0.208 GB(+0.213%)에 불과하지만 시간은 +1,586초 늘어납니다. 이 증가분은 serial 재생이 응답마다 평균 264초씩 추가하기 때문입니다.

반면 프롬프트 캡처를 그룹 전체가 공유하면 응답당 평균 벽시계 시간이 크게 줄어듭니다.

\(G=8\)이 \(G=2\)보다 응답 하나를 처리하는 평균 비용이 67.4% 낮습니다. 4,653초짜리 프롬프트 캡처를 8개 응답이 나눠 부담하기 때문입니다.

4.25M 컨텍스트(= \(4.25 \times 2^{20} = 4{,}456{,}448\) 위치)는 8개 H20에서 prefix-frozen 모드로 64개 멤버 재생, 8번 옵티마이저 스텝을 완료했습니다. peak 메모리는 83.9 GB로 2M 경우(97.7 GB)보다 낮습니다. 응답 길이가 짧아 응답 그래프 자체가 작아지기 때문입니다.

GLM의 2M 실행은 캡처 구간에서 랭크별 peak가 112.571-145.148 GB로 최대 32.577 GB 편차가 납니다. 논문은 이를 부하 불균형 진단 지표로 해석하며, 2M 이상 GLM probe는 아직 시도하지 않았다고 명시합니다.

회고

이 논문이 드문 이유는 제한을 아주 명시적으로 적어놓기 때문입니다. 저자들은 네 단계 증거 사다리를 정의하고, 현재 구현이 어느 단계까지 도달했는지 항목별로 구분합니다.

완료된 것: 실행 용량(Execution capacity). 요청된 채점, 역전파, 집합 통신, 옵티마이저 이벤트가 모든 랭크에서 유한한 값으로 완료됨.

Qwen만 부분 완료: 응답 연산자 정확도(Response-operator fidelity). 글로벌 CP8 merge로 풀-어텐션 응답 forward를 재구성합니다. 단, BF16 분자 리덕션이므로 FP32 완전 동등성이 아닌 "파티션-정확 forward 시맨틱스"입니다.

둘 다 미완성: 분산 업데이트 일관성. Qwen은 K/V 어댑터 그래디언트를 CP 랭크 간에 동기화하지 않아 8개 AdamW 인스턴스가 제각각 스텝을 밟습니다. GLM은 Megatron의 finalize_model_grads를 우회해 CP-복제 어댑터 그래디언트가 리덕션 없이 옵티마이저에 들어갑니다.

둘 다 미완성: 풀 그래디언트 동등성. 프롬프트 상태를 분리(detach)한 시점에서, 정확한 역전파 식

\[\nabla_\theta \ell(\theta, z_P(\theta)) = \frac{\partial \ell}{\partial \theta}\bigg|_{z_P} + \frac{\partial \ell}{\partial z_P}\frac{\partial z_P}{\partial \theta}\]

의 두 번째 항은 포기합니다. 이것이 "정확한 어텐션 연산"과 "풀-시퀀스 그래디언트 동등성"이 다른 주장인 이유입니다.

반복 훈련 루프나 실제 태스크 평가는 없습니다. 옵티마이저를 한 번 호출한 뒤 프롬프트 상태는 오래된 것이 되고, 다음 루프를 위해 반드시 재캡처해야 합니다. 논문 스스로 이것을 "가속기 제약 시스템 실행 가능성 외피(accelerator-bounded systems feasibility envelope)"라고 부릅니다. 길이 기록이나 학습 성능 주장은 하지 않습니다.

정리

• 핵심 설계: 프롬프트를 no-grad로 한 번 캡처하고, 각 응답을 순서대로 재생한다. 라이브 autograd 규모가 \(P+R\)에서 \(R\)로 줄어든다.
• 측정 결과: Qwen3.6-27B가 8개 H20에서 2.1M, 4.25M 컨텍스트 GRPO 실행 완료. GLM-5.2는 32개 H20에서 2.1M 프롬프트 78레이어 역전파 완료.
• 현 상태: 실행 용량은 확인됐으나 분산 그래디언트 조합은 미완성. 반복 학습 루프와 태스크 평가는 후속 과제.