🤖 machine learning

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

LongStraw는 자동 미분(autograd) 없는 프롬프트 공유 평가와 순차적 응답 리플레이를 통해 메모리 사용량을 최적화함으로써 고정된 GPU 예산 내에서 백만 토큰 규모의 강화 학습 사후 학습을 가능하게 하는 아키텍처 인식 실행 스택으로, Qwen3.6-27B 및 GLM-5.2와 같은 모델에서 200만 토큰 이상의 컨텍스트 길이에 대한 실행 능력을 성공적으로 입증하였다.

원저자: Changhai Zhou, Kieran Liu, Yuhua Zhou, Qian Qiao, Jun Gao, Harry Zhang, Irvine Lu, Nolan Ho, Lucian Li, Andrew Lei, Cleon Cheng, Steven Chiang, Yihang Zeng, Di Zhang, Rio Yang, Kaijie Chen, Andrew Che

게시일 2026-07-17

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Changhai Zhou, Kieran Liu, Yuhua Zhou, Qian Qiao, Jun Gao, Harry Zhang, Irvine Lu, Nolan Ho, Lucian Li, Andrew Lei, Cleon Cheng, Steven Chiang, Yihang Zeng, Di Zhang, Rio Yang, Kaijie Chen, Andrew Chen, Pony Ma, Weizhong Zhang, Cheng Jin

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 초지능 로봇에게 거대하고 다단계적인 미스터리를 해결하는 법을 가르치려 한다고 상상해 보세요. 이를 위해 로봇은 해결책을 추측하기 시작하기도 전에 도서관 한 권 분량의 단서(즉, "컨텍스트")를 읽어야 합니다. 인공지능의 세계에서 이것을 "롱 컨텍스트 트레이닝(long-context training)"이라고 부릅니다. 보통 우리가 이 로봇들을 훈련시킬 때는, 엄청나게 긴 단서가 담긴 책을 로봇에게 먹여준 다음 짧은 답변을 쓰라고 요구합니다. 문제는 로봇의 뇌(컴퓨터 메모리)가 그 책의 세부 사항들로 너무 가득 차버려서, 답변을 완성하기도 전에 공간이 부족해진다는 점입니다. 이는 마치 도서관 전체를 손에 들고 미스터리를 풀려고 하는 것과 같습니다. 결국, 당신의 팔은 움직일 힘조차 남지 않게 될 것입니다.

이를 해결하기 위해 과학자들은 모든 정보를 담아둘 수 있도록 점점 더 큰 컴퓨터 칩(GPU)의 도서관을 구축해 왔습니다. 하지만 만약 당신에게 슈퍼컴퓨터가 없다면 어떨까요? 만약 당신이 8개나 32개처럼 정해진 적은 수의 칩만을 가지고 있다면 어떻게 될까요? 이것이 바로 "LongStraw"라는 논문이 해결하고자 하는 퍼즐입니다. 이 논문은 다음과 같이 묻습니다. 우리는 더 많은 칩을 구매하지 않고도, 고정된 적은 수의 컴퓨터 칩만을 사용하여 로봇이 200만 단어짜리 이야기를 읽고 해결책을 쓰도록 가르칠 수 있을까요? 그 답은 "네, 우리는 칩을 더 크게 만들었습니다"가 아니라, "네, 우리는 책을 잡는 방식을 바꿨습니다"입니다.

연구진은 LongStraw라고 불리는 영리한 시스템을 구축했습니다. 이것을 매우 정리 정돈을 잘하는 사서라고 생각해보세요. 로봇이 답변을 쓰는 동안 200만 단어짜리 이야기를 책상 위에 통째로 펼쳐 놓는 대신(그렇게 하면 책상이 무너질 것입니다), 사서는 이야기를 한 번 쭉 읽고 나서 이야기의 핵심적인 "분위기"나 "상태"만을 기억한 뒤 무거운 책을 치워버립니다. 이제 로봇은 그 작은 "분위기" 메모만을 보면서 짧은 답변을 쓰기만 하면 됩니다. 만약 로봇이 특정 세부 사항을 확인해야 한다면, 사서는 재빨리 그 페이지 하나만을 꺼내서 로봇에게 보여준 뒤 즉시 다시 집어넣습니다. 이렇게 하면 이야기가 아무리 방대하더라도 책상이 어지러워지지 않습니다.

이 논문은 이 아이디어를 Qwen과 GLM이라는 두 가지 유형의 로봇 뇌에 테스트했습니다. 연구진은 이 로봇들이 (프롬프트/이야기가) 2,097,152 토큰(약 200만 단어 또는 글자)인 내용을 읽고 응답을 작성하도록 했습니다. 이 과정은 엄격한 예산 하에 진행되었습니다: Qwen 로봇에는 8개의 H20 GPU를, GLM 로봇에는 32개의 H20 GPU를 사용했습니다.

결과는 다음과 같습니다:

"읽기" 부분은 성공적이었습니다: 연구진은 컴퓨터 메모리가 충돌하지 않고도 200만 토큰의 전체 이야기를 성공적으로 처리해 냈습니다. 정보를 언제 붙잡고 있고 언제 놓아줄지를 영리하게 조절함으로써, 고정된 수의 칩에 이 거대한 작업을 담아낼 수 있음을 증명했습니다.
공간은 절약하지만, 시간은 더 걸립니다: 무거운 책을 치워두었다가 필요할 때만 페이지를 꺼내는 방식을 통해 엄청난 양의 메모리를 아꼈습니다. 예를 들어, 로봇의 "추측" 횟수(그룹 크기)를 2에서 8로 늘렸을 때, 메모리 사용량은 거의 변하지 않았지만(약 0.21 GB 증가), 페이지를 한꺼번에 꺼내는 대신 하나씩 꺼내야 했기 때문에 작업을 마치는 데 시간이 더 오래 걸렸습니다.
이것은 "개념 증명"이지, 완벽한 결말은 아닙니다: 이 논문은 자신들이 하지 못한 부분에 대해서도 매우 솔직합니다. 시스템이 시뮬레이션을 성공적으로 실행했고 로봇이 긴 이야기를 읽고 답변을 쓸 수 있다는 것을 보여주긴 했지만, 로봇이 자신의 뇌를 업데이트하는 "완벽한 방식"을 완전히 배웠다는 것까지 입증하지는 못했습니다. 로봇의 서로 다른 "뇌 부위"(그래디언트)를 연결하는 수학적 계산 중 일부가, 완벽한 훈련 세션에서 요구되는 방식만큼 모든 칩 사이에서 완전히 동기화되지 않았기 때문입니다. 이는 마치 자동차 엔진이 아주 적은 양의 연료로도 작동할 수 있다는 것은 보여주었지만, 아직 경주 트랙 전체를 멈춤 없이 달릴 수 있다는 것까지는 증명하지 못한 것과 같습니다.

연구진은 이 결과를 "실행 영수증(execution receipt)"이라고 부릅니다. 이는 "네, 우리는 특정 하드웨어에서 이 거대한 작업을 충돌 없이 성공적으로 실행했습니다"라고 말해주는 영수증입니다. 그들은 레이스에서 승리했거나 가장 빠른 차를 만들었다고 주장하는 것이 아닙니다. 단지 그 차가 특정 양의 연료를 가지고 특정 트랙을 달릴 수 있다는 것을 증명했을 뿐입니다.

요약하자면, LongStraw는 AI를 거대한 이야기에 대해 훈련시키기 위해 반드시 슈퍼컴퓨터가 필요한 것은 아니라는 점을 보여줍니다. 대신, 시간을 조금 더 투자하더라도 훨씬 더 많은 공간을 절약할 수 있는 더 똑똑한 메모리 관리 방식이 필요합니다. 이는 더 적은 예산을 가진 작은 팀이나 연구자들이 거대한 AI 모델을 실험할 수 있는 길을 열어주며, 충분히 영리하다면 고정된 작은 도구 세트로도 큰 일을 해낼 수 있음을 증명합니다.

기술 요약: LongStraw: 고정된 GPU 예산 하에서의 2M 토큰 이상 장기 문맥 RL

1. 문제 정의

추론 시스템이 지원하는 문맥 길이(백만 토큰 수준에 근접)와 강화 학습(RL) 사후 학습(post-training)에서 활용되는 문맥 길이(종종 256K 토큰 이하로 제한됨) 사이에 상당한 격차가 발생했습니다. 이 격차는 에이전트가 긴 시퀀스에 걸쳐 관찰, 도구 출력, 이전 결정을 축적하는 AI 에이전트에게 특히 치명적입니다.

고정된 GPU 예산 하에서 장기 문맥 RL(특히 Group Relative Policy Optimization, GRPO)을 백만 토큰 문맥으로 확장할 때의 핵심 과제는 메모리입니다. 추론에서는 상태 캐싱 후 포워드 그래프를 폐기할 수 있는 것과 달리, 학습은 여러 응답을 점수화하고 이를 통해 역전파(backpropagation)를 수행해야 합니다. 기존 방식들은 프롬프트 전체와 모든 응답 그래프를 동시에 메모리에 유지해야 하므로, 이는 이차적인 어텐션 연산과 막대한 활성화 저장 공간(activation storage)을 초래합니다. FlashAttention, LoRA, 또는 Ring Attention과 같은 기존 기술들은 특정 병목 현상을 해결하지만, 더 큰 클러스터로 확장하지 않고도 고정된 GPU GRPO 실행이 2M 이상의 토큰 문맥을 수용할 수 있게 하지는 못합니다.

2. 방법론: LongStraw 실행 스택

LongStraw는 추가적인 재현(replay) 시간을 대가로 GPU 메모리 사용량을 줄이기 위해, 긴 프롬프트와 짧은 응답 계산을 분리하도록 설계된 아키텍처 인식형 실행 스택입니다. 핵심 설계 원칙은 나중의 응답 토큰이 의존하는 경우에만 프롬프트 경계를 가로질러 텐서를 유지하는 것입니다.

2.1 실행 스케줄

시스템은 그룹 업데이트를 위해 네 가지 단계로 작동합니다:

프롬프트 캡처 (No-Grad): 공유 프롬프트( $P$ 토큰)는 자동 미분(automatic differentiation)을 비활성화한 상태에서 한 번만 평가됩니다. 시스템은 미래의 토큰에 필요한 아키텍처별 조건부 상태(예: 순환 상태, KV 페이지, 잠재 표현)만을 저장하고, 모든 일시적 활성화(FFN 중간값, 어텐션 스크래치, MoE 라우팅 버퍼)를 해제합니다.
Pre-Step Scoring: 고정된 프롬프트 상태를 사용하여 모든 그룹 멤버에 대한 구 정책(old-policy) 및 참조 정책(reference-policy) 로그 확률을 평가합니다.
직렬 응답 재현 (Serial Response Replay): $G$ 개의 응답 그래프를 동시에 유지하는 대신, 시스템은 autograd 하에서 한 번에 하나의 응답 브랜치를 재현합니다. 시스템은 읽기 전용으로 캡처된 프롬프트 상태를 재사용하여, 손실(loss)을 계산하고, 역전파를 수행한 후 즉시 응답 그래프를 해제합니다.
그래디언트 누적 및 최적화: 모든 $G$ 멤버로부터 로컬 그래디언트를 누적한 다음, 워커당 단 한 번의 옵티마이저 호출을 수행합니다.

2.2 아키텍처별 구현

LongStraw는 이 전략을 두 가지 서로 다른 모델 제품군에 적응시킵니다:

Qwen3.6-27B (하이브리드 순환/풀-어텐션):
- 상태 유지: 압축된 Gated DeltaNet (GDN) 순환 상태와 16개의 풀-어텐션 레이어를 위한 컨텍스트 병렬(CP) 샤딩된 Key/Value (KV) 페이지를 GPU에 유지합니다.
- 메커니즘: NF4 QLoRA를 사용합니다. 48개의 GDN 레이어는 고정 크기 상태를 제공하며, 16개의 풀-어텐션 레이어는 물리적 페이지 압축을 사용합니다. 시스템은 포워드 패스를 위해 글로벌 LSE/출력 병합을 수행하지만, K/V 어댑터 그래디언트 동기화가 누락되었음을 인지합니다.
- 하드웨어: 8개의 H20 GPU (CP8)에서 실행되었습니다.
GLM-5.2 (압축 어텐션/MoE):
- 상태 유지: 긴 프롬프트 상태(MLA 잠재 페이지 및 DSA 인덱서-키 페이지)를 CPU 메모리로 이동합니다. 응답 재현 중에 한 번에 하나의 디코더 레이어만 GPU에 스테이징됩니다.
- 메커니즘: 32개 랭크(CP32) 및 전문가 병렬성(EP32)에 걸쳐 "Zigzag" 컨텍스트 페이지 레이아웃을 사용합니다. 응답 그래프를 제한하기 위해 전체 레이어 활성화 체크포인팅을 채택합니다. 시스템은 레이어별로 IndexShare 스케줄을 재구성합니다.
- 하드웨어: 32개의 H20 GPU (TP1/CP32/EP32)에서 실행되었습니다.

3. 주요 기여

공유-프롬프트 실행 설계: 프롬프트 상태를 한 번 저장하고(detatched 상태로) 응답 작업을 직렬화하여, 라이브 훈련 그래프를 $P + \sum R_i$ 에서 $\max(R_i)$ 로 줄이는 새로운 스케줄을 제안합니다.
이중 아키텍처 구현: 이 설계를 두 가지 구조적으로 다른 모델(밀집 하이브리드 순환/어텐션 스택인 Qwen과 압축 어텐션 MoE 스택인 GLM)에 성공적으로 적용했습니다.
고정 예산 실행 영수증 (Execution Receipts): GPU 수를 늘리지 않고도 2M 이상의 토큰 문맥(Qwen의 경우 최대 4.25M)을 위해 GRPO 실행 경로를 완료할 수 있음을 보여주는 경험적 증거를 제시합니다. 또한 분산 그래디언트 구성의 한계점을 명확히 문서화했습니다.

4. 실험 결과

4.1 Qwen3.6-27B (8 H20s)

문맥 길이: 2,097,152-토큰 문맥(2,088,960 프롬프트 + 8,192 응답)을 성공적으로 실행했습니다.
메모리 효율성: 그룹 크기를 $G=2$ 에서 $G=8$ 로 늘렸을 때 피크 할당 메모리는 단 0.21 GB(0.213%)만 증가했습니다. 이는 메모리가 그룹 크기가 아닌 가장 큰 응답 브랜치에 따라 스케일링됨을 확인시켜 줍니다.
확장 능력: 동일한 8개의 H20에서 4.25M 포지션(4,448,256 프롬프트 + 8,192 응답) 스트레스 테스트를 통과했으며, 랭크당 약 83 GB에 도달했습니다.
한계점: 분산 업데이트가 완전히 일관되지 않습니다. K/V 어댑터 그래디언트는 랭크 간에 동기화되지 않으며, 프롬프트 상태 그래디언트는 분리(detached)되어 있습니다.

4.2 GLM-5.2 (32 H20s)

문맥 길이: 78개 레이어 전체에 대해 2,097,152-토큰 프롬프트에 대한 엔드 투 엔드 실행 경로를 검증했습니다.
메모리 전략: CPU 상주 프롬프트 상태(랭크당 약 5.8 GiB)를 통해 응답 재현 중의 78개 레이어 역전파 과정에서 GPU가 예산 내에 머물 수 있었습니다.
한계점: 수용된 영수증의 DSA(Dynamic Sparse Attention) 연산자는 각 CP 샤드에 로컬로 존재합니다(글로벌 2M 문맥이 아닌 65k 토큰 중 상위 2048개를 선택). 또한, 이 경로는 Megatron의 finalize_model_grads를 우회하여, 옵티마이저 단계 이전에 CP 복제된 non-expert 어댑터 그래디언트가 축소(reduce)되지 않은 채 남게 됩니다.

5. 의의 및 주장

본 논문은 LongStraw를 "세계 기록"으로서가 아니라, **고정된 예산 내에서의 실행 가능 범위(feasibility envelope)**로 포지셔닝합니다.

상태 수명 vs. 희소성 (State Lifetime vs. Sparsity): 저자들은 고정 예산 용량이 단순히 희소성(MoE)이나 메모리 효율적인 어텐션 커널에 의해서가 아니라, **텐서의 수명(lifetime)과 물리적 소유권(physical ownership)**에 의해 결정된다고 주장합니다. 프롬프트 중간값을 캡처 직후 즉시 소멸시킴으로써, 시스템은 전체 시퀀스 autograd의 메모리 폭발을 방지할 수 있습니다.
물리적 소유권: 본 연구는 논리적 샤딩만으로는 부족하며, 메모리를 효과적으로 해제하기 위해서는 물리적 메모리 할당이 샤딩 전략(예: 페이지를 적절한 크기의 텐서로 복사)과 일치해야 함을 강조합니다.
겸허한 주장: 저자들은 이것들이 단순한 **실행 영수증(execution receipts)**이지, 훈련의 정확성에 대한 증명이 아님을 명시합니다. 프로프트 상태는 분리되어 있으며( $\partial z_P / \partial \theta$ 항이 누락됨), 분산 그래디언트 구성은 두 모델 모두에서 불완전합니다. 이 결과들은 실행 경로가 하드웨어 제약 내에 들어온다는 것을 입증하여, 연구자들이 대규모 확장을 하지 않고도 장기 문맥 RL을 탐색할 수 있게 해주지만, 아직 분산 그래디언트 일치성이나 정책 개선 측면을 해결했다고 주장하는 것은 아닙니다.

요약하자면, LongStraw는 상태 수명과 물리적 소유권을 세심하게 관리함으로써, 상대적으로 작은 GPU 클러스터에서도 백만 토큰 단위의 RL 사후 학습이 가능하다는 것을 보여줍니다. 이는 병목 현상을 순수 메모리 용량에서 스케줄링 및 상태 관리로 전환합니다.