🤖 machine learning

FORGE: Fused On-Register Gradient Elimination for Memory-Efficient LLM Training

FORGE는 그래디언트를 레지스터 내에서 완전히 처리함으로써 이를 실체화된 그래디언트 없이 제거하기 위해 옵티마이저 단계를 역전파 과정에 융합함으로써, 근본적인 업데이트 로직을 변경하지 않고도 옵티마이저 메모리 사용량을 절반으로 줄이고 학습 속도를 높이며 전체 정밀도 충실도를 보존하는 메모리 효율적인 LLM 학습 방법입니다.

원저자: Dikshant Kukreja, Kritarth Prasad, Avinash Anand, Zhengkui Wang, Erik Cambria, Timothy Liu, Aik Beng Ng, Simon See, Bapi Chatterjee

게시일 2026-06-23

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Dikshant Kukreja, Kritarth Prasad, Avinash Anand, Zhengkui Wang, Erik Cambria, Timothy Liu, Aik Beng Ng, Simon See, Bapi Chatterjee

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

큰 문제: "어질러진 작업실"

당신이 컴퓨터에서 거대한 AI 모델(예: 로봇의 뇌)을 훈련시키고 있다고 상상해 보세요. 이 로봇을 가르치기 위해 컴퓨터는 엄청난 양의 수학 계산을 수행해야 합니다.

표준적인 방식(이를 "역방향 미분(Reverse-mode differentiation)"이라고 부릅니다)에서 컴퓨터는 엄격한 2단계 과정을 따릅니다:

실수 계산하기: 데이터를 살펴보고 로봇이 어디서 틀렸는지 파악한 뒤, 로봇의 뇌 각 부분에 대한 "수정 노트(그레이디언트/gradient)"를 거대한 목록으로 작성합니다. 이 노트들을 컴퓨터의 메모리라는 거대한 화이트보드에 적어둡니다.
수정 사항 적용하기: 화이트보드가 완전히 채워진 후에야 컴퓨터는 지우개와 마커를 들고 그 노트들을 바탕으로 로봇의 뇌를 실제로 업데이트합니다.

병목 현상: 문제는 컴퓨터가 로봇의 뇌와 다음 단계를 위한 노트를 붙잡고 있는 동안, 그 거대한 화이트보드의 노트들을 메모리에 계속 보관하고 있어야 한다는 점입니다. 이 "화이트보드"는 너무 많은 공간을 차지해서, 훈련이 시작되기도 전에 컴퓨터의 메모리를 가득 채워버리곤 합니다. 이는 마치 차고에서 자동차를 수리하려고 하는데, 자동차의 설계도, 도구, 그리고 수리 매뉴얼 전체를 바닥에 한꺼번에 펼쳐놓아야 하는 것과 같습니다. 만약 차고가 너무 작다면, 자동차를 넣을 공간조차 없을 것입니다.

해결책: FORGE ("즉시 수정" 방식)

이 논문의 저자들은 FORGE(Fused On-Register Gradient Elimination)를 통해 이렇게 말합니다. "왜 굳이 화이트보드에 노트를 적어두나요?"

그들은 거대한 노트 목록이 학습 과정에 실제로 필요한 것이 아니라, 단지 우리가 수학을 처리하는 방식 때문에 발생하는 부산물일 뿐이라고 주장합니다.

FORGE의 작동 원리:
노트를 적어두었다가 나중에 다시 읽는 대신, FORGE는 컴퓨터의 가장 빠르고 작은 저장 공간(이를 "레지스터(registers)"라고 부릅니다)에서 수학 계산과 수정을 동시에 수행합니다.

비유: 복잡한 요리를 만드는 숙련된 요리사를 상상해 보세요.
- 기존 방식: 요리사가 양파를 다지고, 메모장에 "양파 다짐"이라고 적은 뒤, 양파를 그릇에 담아둡니다. 그다음 재료로 넘어갑니다. 모든 재료를 다 다지고 목록을 다 작성하고 나서야, 요리사는 메모장을 읽고 재료들을 모두 섞습니다. 이때 메모장은 조리대 공간을 차지합니다.
- FORGE 방식: 요리사가 양파를 다지자마자 즉시 냄비에 던져 넣습니다. 그다음 당근을 다지자마자 바로 던져 넣습니다. 요리사는 아무것도 적지 않습니다. 메모장이 필요 없습니다. 조리대는 깔끔하게 유지됩니다.

이것이 왜 중요한가?

논문은 이 "노트 없는" 접근 방식이 세 가지 주요 이점을 제공한다고 주장합니다.

1. 엄청난 공간 절약 (메모리)
컴퓨터가 거대한 노트 목록을 메인 메모리에 기록하지 않기 때문에, 엄청난 양의 공간을 확보할 수 있습니다.

결과: 표준 칩(H200)에서 80억 개의 파라미터를 가진 모델을 훈련할 때 메모리를 53% 적게 사용할 수 있었습니다.
비유: 복도에 여분의 의자를 쌓아둘 필요가 없어져서, 스튜디오 아파트 안에 10인용 식탁을 놓을 수 있게 된 것과 같습니다.

2. 실제로 더 빠름
컴퓨터가 노트를 적고, 기다리고, 다시 읽는 과정을 멈출 필요가 없기 때문에 전체 과정이 빨라집니다.

결과: 훈련 단계가 약 1.5배 더 빠르게 실행됩니다.
비유: 배달원이 물건을 하나 내려놓고 다시 트럭으로 달려가 다음 물건을 가져오는 과정(기존 방식)과, 물건을 집는 동시에 컨베이어 벨트를 통해 트럭에 바로 실어 보내는 방식(FORGE)의 차이와 같습니다.

3. 정확도를 잃지 않음
보통 공간을 아끼기 위해 단계를 건너뛰면 정밀도가 떨어지기 마련입니다(로봇이 약간 덜 똑똑하게 학습됨). 하지만 저자들은 데이터를 버리기 전, 컴퓨터의 최고 품질인 "레지스터"(풀 프리시전/full precision)에서 수학 계산을 수행하기 때문에, 이 학습 방식이 기존의 느린 방식과 수학적으로 동일함을 증명했습니다.

결과: 로봇은 똑같이 잘 학습하면서도 훨씬 더 효율적으로 학습합니다.

"타일(Tile)" 기법

어떻게 혼란 없이 이 일을 해낼까요? 그들은 거대한 수학 문제를 "타일(tiles)"(예: 128x128 크기의 사각형)이라는 작고 관리 가능한 덩어리로 나눕니다.

타일 하나를 처리합니다: 에러를 계산하고, 뇌를 수정하고, 에러를 즉시 버립니다.
그다음 타일로 넘어갑니다.
이렇게 타일 단위로 작업하기 때문에, 컴퓨터는 전체 에러 목록을 한꺼번에 들고 있을 필요가 없습니다.

이것이 가능하게 하는 것들

논문은 FORGE를 통해 다음과 같은 일이 가능함을 보여줍니다:

동일한 컴퓨터에서 더 큰 모델을 훈련할 수 있습니다.
메모리 부족 문제 없이 더 큰 "배치(batch)"(로봇에게 한 번에 더 많은 예시를 가르치는 것)를 사용할 수 있습니다.
특정 테스트에서, 기존 방식으로는 8개의 GPU가 필요했을 작업을 4개의 GPU만으로 훈련할 수 있었습니다.

요약

FORGE는 컴퓨터가 임시 "수정 노트"로 메모리를 어지럽히는 것을 방지하는 새로운 AI 훈련 방식입니다. 오류를 계산하고 칩의 가장 빠른 부분에서 즉시 모델을 수정함으로써, 학습의 지능을 떨어뜨리지 않으면서도 메모리 사용량을 절반으로 줄이고 훈련 속도를 높입니다. 이는 복잡하고 느린 작업실을 능률적이고 고속인 조립 라인으로 바꾸는 것과 같습니다.

기술 요약: FORGE (Fused On-Register Gradient Elimination)

1. 문제점: 그래디언트 플로어 (The Gradient Floor)

현재의 메모리 효율적인 대규모 언어 모델(LLM) 학습 방법들은 주로 GPU의 네 가지 메모리 풀 중 세 가지(압축을 통한 옵티마이저 상태, 재계산을 통한 활성화 함수, 모델 가중치)를 대상으로 합니다. 그러나 마지막 네 번째 풀인 **실체화된 가중치 그래디언트(materialized weight gradient)**는 건드리지 않은 채로 남겨둡니다.

표준 역전파 미분(reverse-mode differentiation)에서는, 백워드 패스(backward pass)가 모든 레이어에 대한 전체 그래디언트 텐서( $\nabla W$ )를 계산하여 고대역폭 메모리(HBM)에 기록한 다음, 옵티마이저가 가중치를 업데이트하기 위해 이를 다시 읽어 들입니다. 이 2단계 스케줄은 백워드 패스와 옵티마이저 단계 사이의 경계에서 모든 레이어의 그래디언트가 메모리에 공존해야 하는 "그래디언트 플로어" 현상을 만들어냅니다. 8B 파라미터 모델의 경우, 이 일시적인 그래디언트 풀만으로도 약 16 GiB를 소비하며, 이는 종종 특정 학습 설정이 단일 GPU에 적합할지 여부를 결정짓는 요소가 됩니다. 기존 방식들(예: LOMO, FlashOptim)은 전체 파라미터나 레이어 단위로 백워드와 옵티마이저 단계를 융합하지만, 이는 파라미터에 대한 전체 그래디언트 텐서가 업데이트가 일어나기 전에 여전히 메모리에 실체화되어 있음을 의미합니다.

2. 방법론: FORGE

FORGE(Fused On-Register Gradient Elimination)는 옵티마이저 단계를 타일 단위(구체적으로 $128 \times 128$ 타일)로 가중치 그래디언트 계산 과정에 직접 융합함으로써 실체화된 그래디언트 텐서를 완전히 제거합니다.

핵심 메커니즘

FORGE는 전체 그래디언트 행렬을 계산하여 저장하는 대신, 가중치 행렬을 한 번에 하나의 타일씩 처리합니다:

타일 단위 그래디언트 누적: 특정 가중치 타일에 대해, 상위 그래디언트( $\Delta Y$ )와 입력 활성화 값( $X$ )의 부분 곱(partial products)을 FP32 GPU 레지스터에 직접 누적합니다.
인-레지스터 최적화 (In-Register Optimization): 타일의 그래디언트가 레지스터 내에 완전히 누적되면, 옵티마이저 업데이트(예: AdamW)가 해당 타일에 대해 레지스터 내부에서 즉시 적용됩니다.
즉시 폐기: 그래디언트 타일은 생성되는 즉시 소비됩니다. 이는 텐서 형태로 HBM에 절대 기록되지 않습니다.
쓰기 작업 (Write-Back): 업데이트된 가중치 타일과 (선택적으로 양자화된) 옵티마이저 상태 타일(모멘트)만이 글로벌 메모리에 기록됩니다.

정밀도 및 충실도

정확성 (Exactness): 모든 원소별 분리 가능한 옵티마이저(weight의 업데이트가 해당 weight의 좌표에만 의존하는 경우, 예: AdamW, SGD, Lion)에 대해, 타일 단위 업데이트는 수학적으로 표준 전체 그래디언트 단계와 동일합니다. 유일한 차이점은 토큰에 따른 부동 소수점 합산 순서이며, 이는 모든 GEMM 구현에 내재된 특성일 뿐 옵티마이저 로직의 근사화가 아닙니다.
저정밀도 충실도 (Low-Precision Fidelity): BF16 또는 INT8 학습 환경에서도 FORGE는 두 가지 방식으로 완전 정밀도 충실도를 유지합니다:
- 그래디언트가 최적화 전 BF16으로 절단되지 않고, FP32 레지스터 내에서 계산되고 소비됩니다.
- BF16의 "데드 존(dead zone)"(작은 업데이트가 0으로 반올림되는 현상)을 방지하기 위해, FORGE는 가중치 쓰기 작업 시 **확률적 반올림(stochastic rounding)**을 사용하여 sub-ULP 업데이트의 편향 없는 누적을 보장합니다.
- INT8 옵티마이저 상태의 경우, FORGE는 양자화 오차를 최소화하기 위해 더 작은 블록 크기(표준 8비트 옵티마이저의 2048 대비 64)를 사용합니다.

분산 학습 호환성

FORGE는 텐서 병렬성(TP) 및 시퀀스 병렬성(SP) 하에서 증명 가능한 정확성을 가집니다. 이는 각 샤드의 로컬 그래디언트가 백워드 시점에 완결되기 때문입니다. 그러나 현재 형태로는 데이터 병렬성(DP) 또는 컨텍스트 병렬성과 호환되지 않습니다. 이러한 방식들은 업데이트 적용 전 부분 그래디언트의 all-reduce를 필요로 하며, 이는 그래디언트 텐서의 실체화를 강제하기 때문입니다.

3. 주요 기여

구조적 제거, 압축이 아님: FORGE는 실체화된 그래디언트를 수학적 필연성이 아닌 학습 스케줄링의 산물로 식별합니다. 이는 $2P$ (파라미터 수)의 그래디언트 메모리 항을 제거하여, 그래디언트 풀을 모델 크기와 무관하게 $O(P)$ 에서 $O(T_R T_C)$ (타일 크기)로 줄입니다.
증명된 정확성: 논문은 원소별 옵티마이저에 대해 타일 융합 스케줄이 FP32 그래디언트에 대한 표준 AdamW와 동일한 업데이트를 생성함을 증명합니다. 이 정확성은 텐서 및 시퀀스 병렬성 하에서도 유지됩니다.
저정밀도에서의 완전 정밀도 충실도: 그래디언트를 최적화 전까지 FP32 레지스터에 유지함으로써, FORGE는 그래디언트를 BF16으로 절단할 때 발생하는 편향을 방지합니다. 확률적 반올림과 결합하여, BF16 또는 INT8 상태를 사용할 때도 FP32-master 수렴도에 부합하는 성능을 보여줍니다.
범용성: 이 방법은 선형 레이어(linear layers)에 대해 구조적 독립적(architecture-agnostic)이며, 8가지의 완전 융합된 원소별 옵티마이저(AdamW, NAdam, RAdam, Lion, RMSprop, AdaGrad, SGD, SGD-momentum)를 지원합니다.

4. 실험 결과

NVIDIA H200 및 B200 GPU에서 0.6B부터 14B 파라미터 규모의 모델(Llama-3.1-8B, Qwen3 제품군)을 대상으로 평가되었습니다:

메모리 감소: FORGE는 융합된 AdamW와 비교하여 피크 메모리를 38–53% 감소시킵니다. Llama-3.1-8B의 경우, 피크 메모리가 약 75 GiB에서 35.4 GiB로 감소합니다(53% 감소).
처리량 (Throughput): FORGE는 별도의 옵티마이저 커널 실행과 그래디언트의 HBM 왕복(round-trip)을 제거함으로써 스텝 시간에서 1.52배의 속도 향상(99.1 ms 대 150.9 ms)을 달야냅니다.
확장성:
- 단일 H200에서, FORGE는 표준 융합 AdamW 사용 시 OOM(Out-Of-Memory) 오류를 일으키는 Qwen3-14B를 63.2 GiB로 학습시킵니다.
- 4-GPU A100-40GB 설정에서, FORGE는 마이크로 배치 크기 4(24.6 GiB에 적합)를 가능하게 하는 반면, 표준 융합 AdamW는 마이크로 배치 1만 수용 가능합니다. 이를 통해 FORGE는 동일 하드웨어에서 표준 베이스라인보다 2.42배 높은 처리량을 달성할 수 있습니다.
수렴도: FORGE는 지속적 사전 학습(continued pretraining) 작업에서 융합된 AdamW의 수렴도를 0.001 nats(평균 편차) 이내로 추적하며, 이는 저정밀도 편차가 실행 간 노이즈 범위 내에 있음을 확인해 줍니다.
결합성: FORGE는 FP8 양자화 및 활성화 함수 체크포인팅과 깔끔하게 결합되며, 표준 옵티마이저가 흔히 겪는 역양자화 비용을 흡수합니다.

5. 의의 및 주장

본 논문은 FORGE가 점진적인 최적화가 아니라 **구조적 전환(structural shift)**이라고 주장합니다. 그래디언트 플로어를 제거함으로써, 옵티마이저 메모리 오버헤드를 더 큰 배치 사이즈, 더 긴 시퀀스, 또는 기존 하드웨어에서 더 큰 모델을 위한 가용 헤드룸으로 전환합니다.

능력의 경계: 저자들은 그래디언트 플로어 제거가 단순한 상수 배수 개선이 아니라 능력의 경계 문제라고 명시합니다. 예를 들어, 단일 H200에서 14B 모델을 위한 표준 학습 설정은 그래디언트 플로어 때문에 불가능하지만, FORGE는 이를 가능하게 만듭니다.
직교성 (Orthogonality): 이 방법은 상태 양자화나 활성화 함수 재계산과 같은 다른 메모리 절약 기술과 직교하며, 이들과 결합하여 메모리 풋프린트를 더욱 줄일 수 있습니다.
한계점: 이 방법은 현재 원소별로 분리 가능한 옵티마이저로 제한되며, 수정 없이는 데이터 병렬성을 지원하지 않습니다. 또한, 입력 그래디언트가 가중치 업데이트 전에 읽혀야 하므로, 단일 스텝 내에서 여러 그래디언트 소비자를 갖는 타이드 임베딩(tied embeddings)이나 MoE 전문가(experts)에는 적용할 수 없습니다.

요약하자면, FORGE는 "그래디언트 플로어"가 현재 학습 루프의 불필요한 산물임을 입증합니다. 타일 수준에서 백워드 패스에 옵티마이저를 융합함으로써, 광범위한 클래스의 옵티마이저에 대해 수학적 정확성을 유지하면서도 상당한 메모리 절감과 속도 향상을 달성합니다.