MoE 전문가 그래디언트를 역전파 중 즉시 갱신해 학습 메모리를 줄인다

무엇인가

MoE 모델은 토큰마다 일부 전문가만 활성화해 연산량을 늘리지 않고 파라미터를 키우지만, 메모리 비용은 그대로 치른다. 학습 중에는 활성화되지 않은 전문가를 포함한 모든 가중치의 그래디언트와 옵티마이저 상태를 유지해야 하기 때문이다. 논문이 든 예로 OLMoE-1B-7B는 6.92B 파라미터 중 6.44B가 전문가이고, 이들의 bf16 그래디언트만 12.9GB다. 표준 학습 루프는 역전파가 전부 끝난 뒤에야 옵티마이저를 호출하므로, 마지막 MoE 레이어의 그래디언트는 첫 레이어까지 역전파가 도달할 때까지 GPU에 남아 있는다. AdamW 같은 원소별 옵티마이저에는 이 대기가 필요 없다.

어떻게 동작하나

Sonar-MoE의 방법은 PyTorch의 register_post_accumulate_grad_hook을 각 MoE 레이어의 배치 전문가 텐서 두 개, 즉 gate/up 투영을 담은 W_up(E×2n×d)과 W_down(E×d×n)에 거는 것이다. 텐서의 그래디언트가 완성되는 순간 훅이 fused AdamW 업데이트를 적용하고 그래디언트를 None으로 해제해 메모리를 반납한다. 어텐션·임베딩·라우터 등 나머지 파라미터는 기존처럼 스텝 끝에서 글로벌 노름 클리핑과 함께 업데이트한다. 역전파가 MoE 레이어를 한 번에 하나씩 지나가므로 남아 있는 전문가 그래디언트는 한 레이어 분량으로 제한되고, 이는 전체 전문가 그래디언트의 약 1/L이다. OLMoE 기준 12.9GB에서 0.8GB로 16배 줄어든다. 훅이 역전파 중에 학습률을 읽기 때문에, 스텝 t의 학습률은 역전파 전에 설정해야 한다는 조건이 붙는다.

무엇과 다른가

설계에서 중요한 선택은 훅을 어디에 거는가다. 전문가별로 훅을 걸면 레이어당 전문가 수 E와 MoE 레이어 수 L에 대해 스텝당 2EL번 옵티마이저가 호출되지만, 배치 텐서에 걸면 2L번이면 된다. OLMoE(E=64, L=16)에서는 2,048번 대신 32번이다. AdamW는 원소별 연산이라 배치 텐서 한 번 업데이트와 전문가별 E번 업데이트의 결과는 같고, 달라지는 것은 고정 비용이 붙는 호출 횟수뿐이다. 실제로 전문가별 훅은 OLMoE에서 스텝당 23.6ms(+20%), Granite(E=40, 32층)에서 61.1ms(+72%)를 추가한다. 다른 제약은 두 가지다. 전문가 그래디언트를 스텝 끝까지 들고 있지 않으므로 전체 파라미터에 대한 글로벌 노름 클리핑은 불가능하고, 그래디언트 누적(A>1)에서는 마이크로배치마다 훅이 발동해 2차 모멘트가 평균 그래디언트의 제곱이 아닌 개별 그래디언트 제곱으로 A번 갱신되는, 즉 수식은 같아도 다른 옵티마이저가 된다. 그래서 주 실험은 모두 A=1이다.

어떻게 쓰나

실험은 OLMoE-1B-7B-0924(16개 MoE 레이어, 64 전문가, top-8)와 granite-3.1-3b-a800m-base(32개 MoE 레이어, 40 전문가)를 Alpaca 8,000개 예제로 2,000스텝 파인튜닝해 수행했다. 배치 크기 2, 최대 256 토큰, 그래디언트 누적 없음, 학습률 2e-5에서 100스텝 워밍업 후 2e-6까지 코사인 감쇠, AdamW β=(0.9, 0.95), bf16에 fp32 마스터 가중치 없음, 시드 42 고정, 별도 500개 홀드아웃 예제 사용. 전문가 커널 백엔드는 SonicMoE와 ScatterMoE 두 가지를 썼고, 측정은 96GB RTX PRO 6000 Blackwell 2장에서 이뤄졌다.

전제와 한계

결과는 SonicMoE 기준 OLMoE 피크 메모리 55.6GB에서 43.6GB로 12.0GB 감소(21.6%), Granite는 5.8GB 감소(21.4%)다. ScatterMoE에서도 21.1%와 19.1%로 같은 경향이다. 스텝 시간은 글로벌 노름 클리핑을 쓰는 표준 학습보다 17~23% 짧다. 다만 전문가 그래디언트를 클리핑하지 않는 표준 학습과 비교하면 OLMoE에서는 115.9ms 대 116.1ms로 사실상 동일하고, Granite에서만 11.6% 빠르다. OLMoE의 이득은 전부 클리핑 커널을 건너뛰는 데서 나오며(클리핑에 스텝당 약 25ms, 26.9ms 대 1.8ms), Granite는 역전파가 커널 론치 오버헤드로 놀던 유휴 시간이 65ms에서 42ms로 줄어든 효과가 더해진다. 배치가 커지면 클리핑 비용이 스텝당 고정이므로 이득이 줄어 배치 2에서 18%, 배치 32에서 8%가 되고, 메모리 절감도 12.0GB(배치 2), 9.4GB(8), 5.2GB(16), 0GB(32)로 사라진다.

정확성 측면의 주장도 구체적이다. 그래디언트 누적이 없을 때 Sonar-MoE는 전문가 그래디언트를 클리핑하지 않는 표준 학습과 동일한 산술을 수행하며, 3스텝 학습 후 6.44B 전문가 가중치가 비트 단위로 일치했다. 나머지 477M 가중치 중 6개가 달랐는데, 같은 베이스라인을 두 번 돌려도 6개가 달라지므로 실행 간 잡음 수준이다. 홀드아웃 손실은 모든 베이스라인 대비 0.003 이내였다. 흥미로운 부수 결과로, 표준 학습처럼 전문가 그래디언트를 클리핑하면 3스텝 후 전문가 가중치의 47%가 달라지지만 2,000스텝 후 홀드아웃 손실 차이는 0.0014에 그친다. AdamW가 그래디언트 스케일 추정으로 갱신을 나누기 때문에 균일한 재스케일링이 상쇄되기 때문이다. OLMoE 체크포인트의 다운스트림 평가에서는 ARC-Challenge가 3~5점 올랐고 HellaSwag와 MMLU는 거의 변하지 않았으며, 두 학습 방식의 차이는 ARC에서 0.016으로 표준오차 한 개 수준이다.

실무적으로 이 방법이 유효한 조건은 논문이 명시한다. 옵티마이저가 원소별이고, 마이크로배치 누적이 없고, 전문가가 배치 텐서로 저장돼 있고, 전문가 그래디언트에 글로벌 노름 클리핑이 필요 없을 때다. 메모리 절감은 배치가 모델에 비해 작을수록 크고, 속도 이득은 베이스라인이 전문가 그래디언트를 클리핑하거나 역전파가 GPU를 놀릴 때만 나타난다. 8-bit AdamW와는 절감 영역이 달라 결합되며, 둘을 함께 쓰면 약 30GB로 30.5GB 상한 아래에서 돌아가지만 8-bit AdamW 단독으로는 OOM이 난다(다만 스텝당 186ms 대 141ms로 느려진다). 배치 2에서는 그래디언트 체크포인팅이 피크 메모리를 줄이지 못하고 시간만 10~16% 늘었다. 모든 파라미터에 역전파 중 업데이트를 적용하는 방식은 메모리와 속도가 같지만 조밀 파라미터조차 클리핑할 수 없다.

저자가 밝힌 한계는 분명하다. 글로벌 노름 클리핑과의 호환은 업데이트를 미루거나 그래디언트를 오프로드·재계산하거나 훅 안에서 텐서별 노름으로 클리핑하는 방식으로 복원할 계획이지만 아직 평가되지 않았다. 그래디언트 누적 A≥2에서는 스텝당 A번 업데이트하는 다른 옵티마이저가 되어 글로벌 클리핑 표준 학습보다 A=2에서 5%, A=4에서 23%, A=8에서 35% 느려진다(추가 업데이트마다 약 37ms). 모든 실험이 단일 GPU에서 이뤄졌고, 데이터 병렬에서는 그래디언트 all-reduce 이후에 업데이트해야 하므로 통신 훅과의 통합이 남아 있다. 모델은 최대 7B, 데이터는 Alpaca, 시드는 하나뿐이며, 배치가 큰 영역에서는 활성화가 피크를 결정해 절감이 사라진다.