Edge0, 학습형 사전 라우팅으로 SSD에서 35B MoE 실행

들어가며

대형 Mixture-of-Experts(MoE) 모델을 소비자용 하드웨어에서 실행할 때 병목은 연산 능력만이 아니다. MoE는 토큰마다 일부 전문가만 활성화하므로 계산량은 줄일 수 있지만, 모델에 포함된 전체 전문가 가중치를 저장해야 한다는 문제는 남는다. 자료가 제시한 35B급 모델은 4-bit 기준 약 19.5GB이며, 이는 많은 소비자 시스템에서 메모리 용량 자체가 제약이 된다는 뜻이다.

Edge0는 일부 전문가 가중치를 SSD에 보관하고, 계산이 진행되는 동안 다음 가중치를 미리 읽는 방식으로 이 문제에 접근한다. 제공된 설명에 따르면 단일 24GB 머신에서 35B급 MoE를 20 token/s로 서비스하고, 피크 활성 메모리를 3GiB 이하로 유지한다. 5개 공개 벤치마크에서는 FP16 교사 모델과 평균적으로 몇 포인트 이내의 차이를 보였다고 한다.

핵심 설계

• 오프로딩에는 순서 문제가 있다. 일반적인 실행에서는 N+1 레이어가 필요로 하는 전문가를 N 레이어의 출력이 나온 뒤에야 알 수 있다. 실제 라우팅을 기다린 후 SSD를 읽으면, 읽기 시간을 계산과 겹치기 어렵다.
• 사전 라우터가 다음 레이어를 예측한다. 각 레이어에 붙은 경량 예측 헤드가 한 토큰 앞서 다음 레이어의 전문가 선택을 추정한다. 엔진은 이 결과를 이용해 SSD에서 가중치를 더 일찍 스테이징한다.
• 예측 결과를 곧바로 라우팅으로 사용한다. 예측한 경로를 원래 라우터와 비교한 뒤 틀리면 수정하는 구조가 아니다. 예측된 전문가를 그대로 실행한다. 따라서 프로젝트가 말하는 “100% 정확도”는 예측 결과와 실행 경로가 내부적으로 일치한다는 의미이지, 원래 라우터의 선택과 100% 일치한다는 뜻은 아니다.
• 복구 LoRA를 적용한다. 4-bit 양자화와 라우팅 교체로 발생할 수 있는 품질 저하를 보완하기 위해, 실제 학생 경로에서 학습한 비병합 LoRA를 사용한다. 이 방식은 배포 환경의 양자화·사전 라우팅 동작에 맞춰 적응하는 것을 목표로 한다.

의미와 남은 질문

보고된 측정이 재현된다면 Edge0는 엣지 추론의 설계 방향을 넓힌다. 단순히 모델을 줄이는 대신 SSD 대역폭, 전문가 스케줄링, 라우팅, 양자화 후 보정을 하나의 시스템 문제로 다루기 때문이다. 메모리는 제한적이지만 빠른 NVMe SSD를 가진 장치라면 더 큰 모델을 실행할 가능성이 생긴다.

동시에 핵심적인 비용도 더 자세히 공개돼야 한다. 현재 자료에는 레이어별 예측 통계, 잘못된 예측의 장尾 분포, 잘못된 전문가를 읽었을 때의 비용, SSD와 시퀀스 길이에 따른 지연 분포가 없다. 예측을 직접 실행하는 구조에서는 원래 라우터로 자동 복귀해 오류를 제거할 수 없으므로, 처리량은 유지되더라도 출력 품질 저하로 나타날 수 있다. 평균 벤치마크 점수와 평균 토큰 속도만으로는 충분하지 않으며, 첫 토큰 지연, 토큰별 지연 변동, 긴 문맥에서의 안정성, SSD 대역폭 압박도 함께 봐야 한다.

따라서 Edge0의 의의는 SSD만으로 MoE 문제를 해결했다는 데 있지 않다. 학습형 라우팅 예측, 스트리밍 가중치 로딩, 양자화 복구, 저메모리 실행을 결합했다는 점이 핵심이다. 프레임워크와 체크포인트, 어댑터가 오픈소스로 제공된다고 소개된 만큼, 향후에는 코드 재현과 세부 시스템 측정이 주장의 신뢰도를 결정할 것이다.

출처: Hugging Face Daily Papers