LoGRA: 저랭크 기울기로 대형 언어모델 강화학습의 메모리 줄이기

배경과 문제

LLM을 강화학습으로 학습시키면 모델이 스스로 생성한 답변에 대한 피드백을 통해 추론 능력이 크게 향상된다. 하지만 이 과정은 메모리 요구량이 매우 크다. 모델 가중치 자체보다 기울기와 옵티마이저 상태(optimizer state)가 몇 배 더 많은 메모리를 차지하기 때문이다. 예를 들어 7B 파라미터 Qwen 모델의 BF16 가중치는 약 14GB인데, Adam 옵티마이저의 두 FP32 모멘트 버퍼만 56GB를 추가로 사용한다. 여기에 기울기와 활성화(activation)까지 더하면, 모델을 실행할 수 있는 GPU로도 RL 학습은 불가능한 경우가 생긴다.

이 문제는 단순히 큰 GPU를 쓰면 해결되는 것이 아니다. RL 학습은 응답 생성, 보상 평가, 정책 최적화를 번갈아 수행해야 하므로 각 단계의 메모리와 계산을 동시에 관리해야 한다. 기존 시스템 최적화 기법인 FSDP 샤딩, 활성화 체크포인팅, CPU 오프로딩은 메모리를 줄여주지만 통신 오버헤드나 속도 저하가 있을 수 있다. LoRA 같은 파라미터 효율적 방법은 학습할 파라미터 수를 줄이지만, 원래 가중치 행렬을 고정하므로 전체 모델을 업데이트하지 못한다. 따라서 선택한 원래 가중치 행렬을 직접 갱신하면서 메모리를 줄이는 새로운 접근이 필요하다.

기존 접근의 한계

기존 메모리 절약 기법은 크게 시스템 수준과 알고리즘 수준으로 나뉜다. 시스템 수준에서는 FSDP가 학습 상태를 여러 GPU에 분산시키고, 활성화 체크포인팅은 중간 활성화를 다시 계산해 저장을 줄이며, CPU 오프로딩은 일부 상태를 호스트 메모리로 옮긴다. 하지만 이런 방법들은 통신 비용이 늘거나 재계산으로 인해 속도가 느려질 수 있다.

알고리즘 수준에서는 LoRA가 사전 학습된 가중치를 고정하고 저랭크 어댑터만 학습하여 기울기와 옵티마이저 상태를 줄인다. LoRA와 LoGRA의 차이는 업데이트를 매개하는 방식이며, 이 논문의 비교가 LoRA의 표현력 열등을 입증하는 것은 아니다. GaLore나 Flora 같은 기울기 압축 방법은 원래 가중치를 학습 가능하게 유지하면서 기울기와 옵티마이저 상태를 압축하지만, 주로 사전 학습이나 지도 학습에 초점을 맞추었고 RL의 생성-학습 루프에서 정책 동기화까지 압축을 확장하지는 않았다. LoGRA는 이 틈을 메우기 위해 RL 후학습에 특화된 저랭크 기울기 스케치를 제안한다.

연구방법을 단계별로

LoGRA의 핵심 아이디어는 전체 기울기 행렬 G(크기 d×k)를 저장하는 대신, 랜덤 투영 행렬 A(크기 r×k, r은 d와 k보다 훨씬 작음)를 이용해 스케치 S = G A^T (크기 d×r)만 저장하는 것이다. 이렇게 하면 저장할 값이 d×k에서 d×r로 줄어든다. 예를 들어 d=k=4096, r=64일 때 FP32 기울기 버퍼가 64MiB에서 1MiB로 감소한다. 투영 행렬 A는 학습되지 않으며 한 업데이트의 마이크로배치 누적 동안 고정되지만, 다음 업데이트에서는 새로 뽑을 수 있다. 기본적으로 각 원소가 ±1/√r인 랜덤 부호 행렬이다.

실제 구현에서는 전체 기울기를 먼저 계산한 뒤 압축하지 않는다. 대신 역전파 과정에서 각 마이크로배치의 기여분을 바로 스케치 S에 누적한다. 구체적으로, 가중치 기울기 계산에 사용되는 입력을 먼저 투영한 후 역전파 미분값과 곱하는 방식으로, 전체 기울기 버퍼를 만들지 않는다. 이렇게 얻은 스케치 S에 옵티마이저를 적용해 조정된 스케치 U를 만들고, 최종 업데이트는 W ← W - αη U A로 수행한다. 여기서 η는 기본 학습률, α는 predicted-KL 제어기가 정하는 추가 배율이다.

RL 후학습에 적용할 때는 어텐션과 MLP 투영 행렬에만 스케치를 적용하고 나머지 파라미터는 고정한다. 각 행렬마다 고유한 A, S, U를 가지며 α는 공유된다. 업데이트는 가중치에 직접 병합되며 별도의 어댑터 브랜치를 남기지 않는다. 매 스텝마다 A를 새로 샘플링하면 연속된 업데이트들이 서로 다른 부분공간을 사용하므로, 업데이트 합이 반드시 rank r 이하로 제한되지 않는다.

정책 동기화 측면에서도 이점이 있다. 롤아웃 엔진은 시드로부터 A를 재생성하고, 수신한 αηU와 재생성한 A의 곱을 가중치에 적용한다. 전송할 데이터는 d×r 값과 시드, 메타데이터뿐이므로 전체 d×k 업데이트를 보내는 것보다 업데이트 전송량이 줄어든다. 이것이 전체 RL 실행시간의 같은 비율 단축을 뜻하지는 않는다.

결과와 숫자의 의미

주요 실험은 Qwen2.5-Math-1.5B/7B와 Qwen3.8-27B 모델을 DAPO-Math-7.5K 데이터셋으로 RL 학습시키고 MATH-500 벤치마크로 평가했다. 1.5B 모델에서 LoGRA는 평균 학습 메모리를 9.18 GiB에서 7.18 GiB로 21.8% 줄였고, Pass@1 정확도는 63.77%에서 67.87%로 오히려 향상되었다. 7B 모델에서는 평균 메모리가 31.82 GiB에서 17.29 GiB로 45.7% 감소했고, Pass@1은 72.33%로 Dense의 72.48%와 거의 같았다. 메모리는 학습용 GPU 한 장당 수치이며, 평균은 전체 실행의 업데이트별 메모리를 평균한 값이다. 정확도는 세 시드의 결과이며 표의 최고값은 초기 200개 업데이트 구간에서 산출했다. 따라서 평균 메모리 절감률과 최종 체크포인트 정확도를 혼동하면 안 된다.

27B 모델에서는 Dense Adam이 첫 Adam 상태 할당에서 OOM으로 실패했지만, LoGRA는 평균 51.54 GiB, 피크 53.88 GiB로 학습을 완료하여 Pass@1 71.52%, Pass@4 81.87%를 달성했다. 27B에서는 Dense가 한 업데이트도 끝내지 못해 실제 절감률이나 처리량을 비교할 수 없었다. 두 방식이 모두 실행된 1.5B·7B에서는 처리량이 비슷했다. 1.5B에서 LoGRA 55.44 vs Dense 54.28 updates/h, 7B에서 25.73 vs 25.28 updates/h로, LoGRA의 주요 이점은 메모리 절약과 27B 학습 가능성에 있다.

장기 학습 안정성 실험에서는 Qwen3.8-27B를 Reasoning-Gym Hard 혼합 데이터셋으로 1,100스텝 이상 학습시켰다. 헬드아웃 매크로 검증 점수가 초기 39.69%에서 62.94%로 상승했고, 피크는 1,060스텝에서 65.52%였다. 평균 응답 길이는 8,312토큰에서 4,705토큰으로 줄었고, 16,384토큰 컨텍스트 제한으로 잘린 응답 비율은 34.72%에서 12.95%로 감소했다. 이 점수는 8개 작업의 헬드아웃 문제 193개에서 계산한 작업별 검증 점수의 매크로 평균이다. 결과는 하나의 학습 궤적에서 관찰된 것으로, 모든 시드나 RL 작업에서 장기 안정성이 보장된다는 의미는 아니다.

구체적인 예시

다음은 논문의 행렬 크기를 대입해 계산한 설명용 예시다. 저랭크 스케치의 동작을 이해하기 위해, 가중치 행렬 W가 4096×4096 크기라고 가정하자. 전체 기울기 G를 저장하려면 16,777,216개의 숫자가 필요하다. LoGRA는 rank r=64를 선택하고, 64×4096 크기의 랜덤 투영 행렬 A를 만든다. 스케치 S = G A^T는 4096×64 크기이므로 262,144개의 숫자만 저장하면 된다. 이는 원래의 1/64에 해당한다. 업데이트할 때는 S에 A를 곱해 4096×4096 기울기를 저랭크로 근사하며, 원래 기울기를 정확히 복원하는 것은 아니다. 또한 실제로 전체 임시 행렬을 저장할 필요 없이 두 인자로 표현한다. 이 근사 행렬의 각 열은 S의 64개 열의 선형 결합이므로 rank가 최대 64이다.

다음은 실제 실험 결과가 아니라 predicted-KL 공식에 숫자를 대입한 예시다. 현재 정책과 제안된 업데이트 후 정책 사이의 KL 발산을 추정한다. 만약 KL 예산 δ가 2×10^-4이고, 제안된 업데이트의 예측 KL q(D)가 8×10^-4라면, 비율 δ/q(D) = 0.25이므로 제곱근을 취해 상한 α_max가 허용하면 α = 0.5가 되고, 상한이 더 작으면 그 값까지만 적용한다. 즉, 업데이트 크기를 절반으로 줄여 예측 KL을 예산에 맞춘다. 반대로 q(D)가 δ/4라면 상한이 2 이상일 때에만 α = 2로 업데이트를 두 배로 키울 수 있다. 실제 배율은 α_max와 계산값 중 작은 값이다.

한계와 오해하기 쉬운 점

LoGRA의 메모리 절약은 기울기와 옵티마이저 상태에 국한된다. 모델 가중치, 캐시된 투영 행렬, 활성화, 기타 버퍼는 여전히 메모리를 차지한다. 따라서 전체 학습 메모리 감소율은 기울기 버퍼의 감소율보다 작다. 논문의 45.7% 절약은 평균 학습 메모리 기준이며, 피크 메모리 기준으로는 7B에서 31.85 GiB에서 22.26 GiB로 약 30% 감소에 그친다.

또한 LoGRA가 항상 Dense보다 정확도가 높은 것은 아니다. 7B 모델에서는 Pass@1이 Dense 72.48% 대비 LoGRA 72.33%로 약간 낮았다. LoRA와의 비교에서도 LoRA가 평균 Pass@1에서 70.45%로 LoGRA의 68.47%보다 높았고, 업데이트 시간도 LoRA가 더 빨랐다. 따라서 LoGRA는 메모리가 제한된 환경에서 유용한 대안이지, 모든 경우에 우월한 방법은 아니다.

predicted-KL 제어는 국소 근사에 기반하므로 실제 KL 발산과 다를 수 있다. 논문의 절제 실험에서 큰 상수 KL 예산(1.5×10^-3)은 성능 저하를 보였지만, KL 제어 없이도 학습이 안정적으로 유지된 경우가 있었다. 저자들은 두 시드 결과만으로는 KL 제어가 안정성에 필수적이라고 단정할 수 없다고 명시한다. 또한 27B 장기 학습 실험은 단일 시드로 진행되어 일반화에 한계가 있다.

실제로 어떻게 활용할까

LoGRA는 GPU 메모리가 부족하여 RL 후학습을 시도하지 못하는 연구자나 소규모 팀에게 특히 유용하다. 예를 들어 8개의 H100 80GB GPU가 있는 단일 노드에서 27B 모델을 Dense Adam으로 학습하면 OOM이 발생하지만, LoGRA를 사용하면 1,100스텝 이상 학습할 수 있다. 논문의 기본 자원 배분은 H100 8장 중 4장을 FSDP 학습에, 4장을 독립 응답 생성에 사용한다. 이는 이 설정의 기존 하드웨어로 더 큰 모델을 실험할 수 있게 해준다.

구현 측면에서 LoGRA는 코드가 Molt 라이브러리에 공개되어 있다. 적용 시에는 rank, 투영 행렬 분포, KL 예산 스케줄 등을 조정해야 한다. 논문의 절제 실험에 따르면 rank가 성능에 가장 큰 영향을 미치며, rank 256이 rank 4보다 Pass@1을 3.02~3.95%p 향상시켰다. 투영 행렬을 매 스텝 새로 샘플링하는 것과 고정하는 것 사이에는 큰 차이가 없었으므로, 단기 학습에서는 고정 투영도 충분할 수 있다. predicted-KL 예산은 보수적으로 작게 유지하는 것이 안정적 학습에 도움이 된다.