Dimuon: LoRA 에 맞춰 다시 짠 Muon

옵티마이저 상태: 이동 평균 하나를 16비트로

LoRA 로 학습할 때 GPU 메모리에는 얼린 모델, 어댑터(두 인자), 어댑터의 그래디언트, 그리고 옵티마이저가 스텝 사이에 들고 다니는 상태가 올라간다. AdamW 는 파라미터 하나마다 그래디언트의 이동 평균 m 과 그래디언트 제곱의 이동 평균 v 를 32비트로 들고 다녀 8바이트를 쓴다. 32비트로 둔 어댑터 가중치가 4바이트이니, 옵티마이저 상태가 학습하는 숫자 자체보다 두 배 크다. 얼린 모델은 여러 학습이 나눠 쓸 수 있지만, 옵티마이저 상태는 학습마다 따로 있어야 한다. 이 상태를 줄일 수는 없을까?

역사: 2차 모멘트를 행과 열로 줄인 Adafactor

옵티마이저 상태를 줄이려는 시도는 Dimuon 이 처음이 아니다. 노엄 셰이저(Noam Shazeer)와 미첼 스턴(Mitchell Stern)은 2018년 ICML 에 「Adafactor: Adaptive Learning Rates with Sublinear Memory Cost」를 냈다. 이들은 Adam 의 v 를 칸마다 들고 다니는 대신, 행렬의 행마다 합과 열마다 합만 들고 다니며 v 를 그 둘의 곱으로 다시 만들었다. n×m 칸이 n + m 개로 준다. 그 대가로 칸마다의 정보를 잃는데, 그래도 번역 모델 학습에서 Adam 과 비슷하게 학습된다고 보고했다.

AdamW 1.61 GB, Muon 0.81 GB, Dimuon 0.40 GB

재현 결과의 Dimuon 은 이동 평균을 모멘텀 M 하나만 들고 다니고, 그것도 저장할 때 16비트 실수(bfloat16)로 한 번 반올림해 둔다. 계산은 32비트로 하고 이번 스텝에는 반올림 전 값을 쓰며, 다음 스텝을 위해 넣어 두는 사본만 줄인다. 파라미터 하나에 2바이트다. 원래 Muon 은 모멘텀 하나를 32비트로 들고 다녀 4바이트였다.

독립된 재현자 한 사람(catid)이 2026년 10월 말에 이 차이를 실제 크기에서 재었다. 파라미터 320억 개 규모의 언어 모델 Qwen2.5-32B 를 얼려 두고, 64층의 피드포워드 사영 셋(gate, up, down)에만 랭크 32 LoRA 를 달았다. 모델 설정의 폭(은닉 5120, 피드포워드 27648)으로 세면 LoRA 파라미터는 층마다 3,145,728개, 모두 201,326,592개다(sources/dimuon/s7_memory.py).

이 셈은 catid 가 GPU 에서 잰 옵티마이저 상태(1.61, 0.81, 0.40 GB)와 소수 둘째 자리까지 같다. Dimuon 의 상태는 AdamW 의 4분의 1이고, 반은 v 를 버린 데서, 나머지 반은 16비트로 저장한 데서 온다.

v 가 하던 일은 누가 하나

Adam 의 v 는 칸마다 그래디언트의 크기를 재어 큰 칸은 덜, 작은 칸은 더 움직이게 했다. v 를 버리면 이 일을 다른 무엇이 해야 한다. Muon 에서는 직교화가 특이 방향마다 크기를 1 근처로 맞추는 것으로 대신했다. Dimuon 은 여기에 행·열 조건화를 더하는데, 이 조건화는 따로 쌓아 두는 통계가 아니라 이미 저장해 둔 지난 모멘텀에서 스텝마다 새로 계산한다. Adafactor 가 v 를 줄여 행·열 합으로 들고 다녔다면, Dimuon 은 그 행·열 합마저 들고 다니지 않는다. 스텝 크기도 노름을 학습률과 랭크로 못 박아 정하므로, 크기를 재려고 들고 다닐 통계가 없다.

ML에서: 한 얼린 모델 위의 여러 학습

싱킹 머신스 랩은 팅커를 소개하면서 LoRA 를 쓰는 이유를 「같은 계산 자원을 여러 학습이 나눠 쓰게 해 비용을 낮추려고」라고 적었다. 얼린 모델 한 벌 위에서 여러 사람의 어댑터가 함께 학습된다면, 학습마다 따로 드는 옵티마이저 상태가 한 장비에 몇 개의 학습을 올릴 수 있느냐를 정하는 데 들어간다. 장비 하나에서 옵티마이저 상태에 20 GB 를 쓸 수 있다고 하면, 위 크기의 어댑터를 AdamW 로는 12개, Dimuon 으로는 49개 올릴 수 있다. 회사가 이것을 Dimuon 을 만든 이유로 밝힌 적은 없으므로, 이 연결은 이 책의 짐작이다.

같은 catid 벤치마크는 품질과 속도도 쟀다. 검증 세트로 학습률을 고른 뒤 씨앗 셋으로 돌린 시험에서, 여섯 가지 설정 가운데 Dimuon 이 가장 낮은 시험 손실을 낸 것은 인자 A 를 순전파마다 정규화하는 변형(NoRA)의 두 칸이었고, 표준 LoRA 에서는 AdamW 나 Muon 이 조금 나았다. 재현자도 「Dimuon 이 품질에서 고르게 이긴다고는 볼 수 없다」고 적었고, 초당 처리한 토큰은 Muon 보다 많고 AdamW 보다 조금 적었다. 이 결과에서 분명하게 남는 차이는 손실이 아니라 옵티마이저 상태의 크기다.

문제 2. 한 장비에 어댑터 몇 개

Qwen2.5-32B 의 FFN 에 단 랭크 32 LoRA(파라미터 201,326,592개)를 여러 사람이 같은 얼린 모델 위에서 따로 학습시킨다. 장비 하나에서 옵티마이저 상태에 20 GB 를 쓸 수 있다. (가) AdamW, Muon, Dimuon 으로 각각 몇 개의 학습을 올릴 수 있는가? (나) Dimuon 이 모멘텀을 32비트로 저장하면? (다) 모멘텀을 bfloat16 으로 저장하면 무엇을 잃는가? 재현 결과처럼 「계산은 32비트, 이번 스텝은 반올림 전 값, 저장할 때만 반올림」으로 하면 무엇을 지키는가?

(가)는 나누기예요. 2억 개에 8바이트면 1.61 GB 라 20 GB 에 12개, Muon 은 0.81 GB 라 24개, Dimuon 은 0.40 GB 라 49개요.

(나)는 Muon 이랑 같아져서 24개야. 네 배 가운데 두 배는 v 를 버린 데서, 두 배는 16비트로 저장한 데서 오는 거네.

(다)로 가 볼까요. bfloat16 은 유효 숫자가 세 자리가 채 안 돼요. β = 0.9 인 모멘텀을 bfloat16 으로 들고 다니면 무엇이 걱정되죠?

음… 모멘텀에 새로 더해지는 건 0.1 × G 잖아요. 모멘텀이 크고 이번 그래디언트가 작으면 그 0.1 × G 가 반올림에 묻힐 수 있겠네요.

그래서 계산은 32비트로 하고 이번 스텝은 반올림 전 값을 쓰는구나. 그러면 이번 스텝의 방향은 정확하고, 반올림 오차는 다음 스텝부터 0.9 배씩 줄어들며 남고.

재현한 사람도 β = 0.99 로 열 스텝을 돌려 이 차이를 쟀어요. 저장을 32비트로 두면 서버와 점점 멀어졌고, 「쓴 뒤 bfloat16 으로 한 번 반올림」이 서버와 가장 잘 맞았대요. 반올림한 뒤에 쓰거나, 반올림 대신 잘라 버리는 방법은 덜 맞았고요.

옵티마이저 상태를 줄이는 방법이 식이 아니라 숫자를 저장하는 방식에도 들어 있는 거네요. 서버 쪽에서는 그 저장 방식까지 정해 둔 거고요.