SSD 가중치층으로 35B MoE를 라우팅 예측으로 24GB 데스크톱에서 서빙한다

무엇인가

이 논문은 소비자용 하드웨어에서 대형 MoE를 서빙할 때의 병목이 계산이 아니라 가중치 메모리라는 문제를 다룬다. 35B급 모델은 4비트로 줄여도 19.5GB이고, MoE의 희소성은 토큰당 계산량만 줄인다(35B MoE가 토큰당 약 3B 파라미터를 활성화). 40층 256-expert 구성에서 expert 가중치는 층당 453MB, expert 하나당 1.77MB로 총 18GB(16.9GiB)이며, 이는 19.5GB 체크포인트의 4분의 3에 해당한다. 24GB 데스크톱에서 OS와 KV 캐시를 빼면 남는 공간이 없다. 그렇다고 SSD로 단순 오프로딩하면 빨라지지 않는다. 층 N+1의 expert 선택이 층 N의 출력에 의존하는데, 그 출력은 층 N+1의 가중치를 읽어야 할 시점에 아직 존재하지 않기 때문에 디스크 읽기를 계산 뒤에 숨길 수 없다. 저자들은 이를 메모리 월의 나머지 절반이라고 부른다.

어떻게 동작하나

Edge0는 expert 가중치를 SSD에 두고 mmap으로 필요한 바이트 범위만 온디맨드로 읽는 스트리밍 추론 엔진이다. 피크 메모리는 파라미터 수가 아니라 활성 집합에 비례한다. 실행 경로는 네 가지다. exact는 중복 제거 온디맨드 번들 빌드·스택·양자화 gather로 모든 경로의 정확성 기준이고, staged는 고정 슬롯 더블 버퍼링으로 디코드의 주력이며, hot은 층별 LRU 상주 hot expert 집합(PowerInfer, Mixtral-offloading, MoE-Infinity와 같은 hot/cold 분할), whole-layer는 프리필용으로 층의 모든 expert를 한 번에 읽는다. 모든 경로는 down(silu(gate(x))·up(x))를 같은 양자화 gather 커널로 계산하고, 테스트 스위트가 각 경로를 역양자화 참조와 원소 단위로 비교한다(상대 L2 1% 미만, 측정 잔차 약 0.24%는 커널 자체의 bf16 내부 정밀도). 이 계약 덕분에 층과 단계별로 경로를 자유롭게 바꿔도 출력이 변하지 않는다.

무엇과 다른가

핵심은 prerouter다. 층 N이 소유한 작은 head가 토큰 t에서 층 N의 post-attention norm 출력을 입력받아 층 N+1의 토큰 t+1 라우팅을 예측하고, 층 N+1은 한 토큰 전에 만들어진 그 예측을 소비한다. 예측은 라우팅 자체로 소비되므로 스테이징된 expert 집합과 실제 라우팅 집합이 구성상 동일해지고, 아무것도 버려지지 않는다. Pre-gated MoE처럼 같은 토큰 안에서 다음 블록의 expert를 고르는 방식은 층마다 동기화와 head 평가가 필요해 GPU 파이프라인을 비우고(저자 측정으로 스텝당 30~100ms), 측정한 모든 same-token 변형이 단순 LRU 베이스라인보다 느렸다. head는 fc1 → erf-gelu → fc2에 선형 잔차 경로를 더한 형태로, 잔차는 다음 층 라우터 가중치로 warm-start되고 기본값은 0이라 훈련이 다음 라우터를 이 은닉 상태에 직접 적용하는 것에서 시작한다. 입력은 은닉 상태에 두 개의 top-k one-hot(이 층이 이 토큰에서 실제로 라우팅한 expert, 그리고 직전 토큰에서 라우팅한 expert)을 이어 붙인 것으로, 35B 티어는 2048 + 2×256 = 2560 차원, head 은닉 폭 512다. 35B 티어는 head 33개(소유 층 6~38), 8B 티어는 16개(7~22)를 가지며, 35B는 40층 중 32층, 8B는 24층 중 16층이 자기 gate 대신 예측으로 스트리밍한다. 8B는 K=8로 층당 8개 슬롯, 35B는 K=4로 점진적 sticky-slot 스택을 쓰고, 프리필은 두 티어 모두 whole-layer 경로를 쓴다.

어떻게 쓰나

서빙되는 모델은 (int4 베이스) + (라우팅 대체) + (LoRA)이고, LoRA는 앞의 두 가지로 잃은 품질을 되돌리기 위해 존재한다. 중요한 설계 결정은 LoRA를 병합하지 않는다는 것이다. y = W_4bit(x) + (α/r)BAx를 병렬 델타로 계산해 4비트 베이스 바이트를 건드리지 않는다. 병합 후 재양자화는 구현 문제가 아니라 산술 문제로 실패한다. LoRA 델타의 RMS가 10^-3 수준으로 4비트 그룹 스텝보다 작아서, 재양자화가 가중치 수준 델타를 대부분 지운다(어텐션 투영에서 효과의 34%, dense 투영에서 2%만 생존, 로짓 수준에서는 18% 생존). 병합하지 않은 경로는 어댑터 42MB를 쓰고 디코드 시간은 측정 가능할 만큼 늘지 않는다. 훈련은 3단계다. 1단계는 head만 훈련해 다음 층의 실제 라우터를 모방하게 하고, 2단계는 어텐션·linear-attention·공유 expert 투영에 LoRA를 붙이고(라우팅되는 expert에는 붙이지 않는다. 스트리밍되어 교체 가능해야 하므로) prerouter 라우팅이 켜진 student 경로로 약 200만 행의 교사 생성 텍스트에 SFT를 한다. 저자들에 따르면 이 순서는 타협 불가였다. head 먼저, SFT 다음, on-policy 증류 마지막이며, 증류만 한 체크포인트는 student 라우팅에서 반복적이고 붕괴된 텍스트를 낸다. 3단계는 student가 생성한 텍스트를 fp16 원본이 교사로 채점하고 teacher의 top-k 토큰에 reverse KL(top-k 밖 질량은 tail 항으로 처리)을 적용하며, SFT 코퍼스의 10분의 1인 약 20만 행에서 수렴한다.

전제와 한계

품질 평가는 OpenCompass로 계산 서버에서 Edge0(int4 + 어댑터 + prerouter 라우팅)와 원본 fp16 베이스를 동일 설정으로 비교했고, 벤치마크별 평균 격차는 35B 티어 3.9포인트, 8B 티어 2.8포인트로 두 서빙 티어를 fp16 베이스와 품질이 맞먹는 것으로 취급한다. 속도와 메모리는 단일 장치에서 측정했다. Mac mini M4 Pro 24GB에서 35B 티어는 피크 활성 메모리 2.9GiB로 20.4tok/s를 내는데, 같은 기기에서 19.5GB 가중치를 전부 상주시킨 바닐라 mlx-lm 서버는 18.2GiB를 점유하며 3.9tok/s에 그친다. Edge0는 상주 서버의 6~7분의 1 메모리로 5배 빠르다. 프리필은 반대 양상으로, 35B 티어가 3.1k 토큰 프롬프트에서 콜드 113 / 웜 140tok/s, 8B 티어가 500 / 1102tok/s다. 즉 제약된 단계는 프리필이 아니라 디코드다.

prerouter의 이득은 16GB MacBook M2에서 18.4GiB 체크포인트를 서빙하는, 가중치가 물리적으로 들어가지 않는 조건에서 A/B로 측정했다. 모든 구성이 SSD에서 expert를 폴트인하며, 엔진이 돌아가려면 회수 불가능한 MLX 할당이 들어가야 하는데 prerouter가 그 값을 올린다(K=2,4,8에서 온디맨드 1.72/1.73/1.79GiB 대비 2.33/2.60/3.22GiB). 대신 메인 스레드가 expert 로드에 막혀 있는 시간이 K=2에서 스텝당 154.9→46.5ms, K=4에서 244.0→101.9ms, K=8에서 575.0→211.6ms로 줄고, 디코드는 8.6/6.4/3.3tok/s 대 온디맨드 4.8/3.5/1.8tok/s로 +80%, +82%, +84%가 된다. 읽는 바이트 수는 거의 같다(K=8에서 125.1 대 126.9MiB, K=2에서 29.5 대 30.4MiB). K=4에서만 prerouter 쪽이 58.9 대 50.9MiB로 16% 더 읽는다. 즉 이득은 예측 품질이 아니라 임계 경로에서 걷어낸 콜드 로드 지연에서 나온다. 비용은 페이지 캐시다. K=8에서 prerouter 쪽은 회수 불가능한 MLX 메모리를 1.43GiB 더 붙잡고 페이지 캐시는 1.15GiB를 잃는다. 또 콜드 페이지가 더 적고 큰 로드로 도착한다(콜드 90%에서 로드당 1.40MiB, 콜드 20%에서 0.32MiB). 로드 비용은 1.17ms + 1.33ms × (콜드 비율)로 피팅되어 측정된 로드당 비용을 0.13ms 이내로 재현한다. 저자들은 아직 당기지 않은 두 레버로 번들과 스택 텐서 뷰의 중복 복사 제거(K=8에서 약 0.45GiB 회수)와 점진적 스태킹을 메인 스레드 밖으로 옮기기(K=8에서 스텝당 62.5ms 제거)를 꼽는다.

저자들이 밝힌 한계는 분명하다. Edge0는 한 번에 하나의 요청만 FIFO로 직렬 처리하며, 동시성은 서빙 계층의 몫이고 배칭은 expert 작업 집합을 바꾸므로 이 프로파일로는 모델링되지 않는다. 스트리밍 엔진의 디코드 비용은 저장장치가 아니라 CPU 쪽에 있다. 40층 포워드에서 그래프 빌딩에 스텝당 44ms가 바닥으로 깔리고 저장장치 최적화로는 줄지 않는다. prerouter 이득은 제거하는 대기 시간에 비례하므로 캐시가 뜨겁거나 저장장치가 빠르면 줄어들고, 재사용률에 묶인다. 인접 토큰은 층 expert 집합의 약 4분의 1만 일치해서 프리페치한 expert가 한 번 읽히고 두 번 값을 치르는 경우가 많다. 품질 손실은 긴 사고 연쇄에 집중되어 AIME에서 35B 티어 6.1포인트, 8B 티어 10.0포인트가 떨어진다(다른 8B 벤치마크는 6.7포인트 이내이고 MMLU-Pro는 student가 4.3포인트 앞선다). 구현은 MLX 백엔드 하나뿐이며 CUDA 슬롯은 코드가 아니라 아키텍처로 남겨 두었다.

실무 관점에서 이 논문이 주는 신호는 두 가지다. 첫째, 24GB급 소비자 기기에서 35B MoE를 돌리는 문제를 가중치를 어디에 둘 것인가라는 배치 결정으로 재정의하고, 그 대가로 라우팅 정확도를 훈련 시점에 미리 지불한다는 발상이다. 예측을 라우팅으로 그대로 소비하기 때문에 런타임 폴백 로드나 토큰 드롭 같은 커버리지-품질 트레이드오프가 구성상 사라진다. 둘째, 어댑터를 병합하지 않고 병렬 델타로 서빙하는 패턴이다. 4비트 베이스에서 LoRA를 병합·재양자화하면 효과의 대부분이 산술적으로 사라지므로, 어댑터를 별도 파일로 두고 베이스는 읽기 전용으로 유지하는 편이 낫다. 라우팅 폭 K를 바꾸는 것이 어댑터 파일 교체만으로 끝난다는 점도 운영 비용 면에서 실용적이다. 다만 이득이 저장장치 지연과 메모리 압박에 조건부라는 점, 즉 빠른 NVMe와 넉넉한 메모리에서는 이득이 줄어든다는 점은 도입 전에 확인해야 한다. 프레임워크, 체크포인트, 어댑터는 Apache-2.0으로 공개되어 있다.