Edge0: 35-миллиардную MoE-модель научили работать с SSD на домашнем ПК

Edge0: 35-миллиардную MoE-модель научили работать с SSD на домашнем ПК

MoE-модели (mixture-of-experts) на потребительском железе упираются не в вычисления, а в память под веса: разреженность экономит вычисления на токен, но не сокращает объём весов, которые нужно держать под рукой. Модель класса 35 млрд параметров при 4-битном квантовании весит 19,5 ГБ, это уже не помещается в память обычной видеокарты. Простая выгрузка весов на SSD не спасает: чтобы понять, каких экспертов слоя N+1 нужно подгрузить, сначала нужен выход слоя N, поэтому чтение с диска не успевает начаться заранее и спрятаться за время вычислений.

Авторы предложили Edge0, потоковый движок инференса MoE с «предроутером»: отдельная небольшая голова на каждом слое предсказывает маршрутизацию следующего слоя на один токен вперёд, и это предсказание используется как сама маршрутизация, то есть набор заранее подгруженных с SSD экспертов всегда совпадает с реально выбранными, и ни один эксперт не подгружается впустую.

Чтобы компенсировать потерю качества от 4-битного квантования и от замены настоящей маршрутизации предсказанной, авторы обучили отдельный «восстанавливающий» LoRA-адаптер (recovery LoRA), он не сливается с основными весами и обучается на «пути ученика».

На одной машине с 24 ГБ памяти Edge0 обслуживает MoE-модель на 35 млрд параметров со скоростью 20 токенов в секунду, при этом пиковое потребление активной памяти, всего 3 ГиБ. По качеству результат в среднем отстаёт от исходной fp16-версии («учителя») лишь на несколько пунктов на пяти открытых бенчмарках. Тот же фреймворк работает и с моделью на 8 млрд параметров. Framework, чекпойнты и адаптеры выложены в открытый доступ.

Ключевые факты

  • MoE-модель на 35 млрд параметров весит 19,5 ГБ при 4-битном квантовании и не помещается в память обычной видеокарты, а простая выгрузка весов на SSD не работает из-за того, что маршрутизация следующего слоя зависит от выхода текущего.
  • Edge0 использует «предроутер», отдельную голову, которая предсказывает маршрутизацию следующего слоя на один токен вперёд, и заранее подгружает с SSD именно тех экспертов, что будут реально выбраны.
  • Восстанавливающий LoRA-адаптер (recovery LoRA), обучаемый отдельно и не сливаемый с основными весами, компенсирует потери качества от 4-битного квантования и от замены реальной маршрутизации предсказанной.
  • На машине с 24 ГБ памяти Edge0 обслуживает 35-миллиардную MoE-модель со скоростью 20 токенов/с при пиковом расходе всего 3 ГиБ активной памяти.
  • Качество отстаёт от fp16-версии модели лишь на несколько пунктов в среднем по пяти открытым бенчмаркам; framework, чекпойнты и адаптеры для 35B- и 8B-моделей выложены в открытый доступ.

Почему это важно

MoE-модели растут в размере, но их разреженность экономит только вычисления, а не память под веса, это и есть «вторая половина стены памяти», о которой говорит заголовок работы. Edge0 показывает, что для запуска больших MoE-моделей необязательна дорогая видеокарта с десятками гигабайт памяти: узкое место, не вычислительная мощность GPU, а то, что чтение весов с диска не успевает подстроиться под динамическую маршрутизацию экспертов. Предсказание маршрутизации на шаг вперёд снимает именно это ограничение.

Кому это важно

Разработчикам и энтузиастам, которые хотят запускать большие MoE-модели на собственном железе, на одной видеокарте с 24 ГБ памяти, а не арендовать облачные GPU; инженерам, которые занимаются инференсом на ограниченной памяти и оптимизацией edge-развёртывания моделей.

Как это применить

Framework, чекпойнты и LoRA-адаптеры выложены в открытый доступ, поэтому подход можно опробовать напрямую: развернуть 35-миллиардную MoE-модель на одной домашней видеокарте с 24 ГБ памяти. Для менее мощного железа тот же фреймворк поддерживает и модель меньшего размера, на 8 млрд параметров.

Можно ли доверять

Материал, препринт на HuggingFace Papers, метод проверен авторами на пяти открытых бенчмарках, и потеря качества относительно fp16-версии измерена и заявлена как «в пределах нескольких пунктов». При этом в тексте нет ни имён авторов, ни организации, стоящей за работой, ни названий конкретных бенчмарков, ни сравнения с другими системами выгрузки весов на SSD, оценка результатов целиком со стороны самих авторов, независимой проверки нет.

Риски и подводные камни

Формулировка «в пределах нескольких пунктов» расплывчата и не даёт точных цифр деградации качества. Сравнения с конкурирующими системами оффлоада на SSD или с другими базовыми показателями в тексте нет, поэтому судить о преимуществе Edge0 перед альтернативами по этому материалу нельзя. Не указаны ни стоимость, ни конкретная модель видеокарты, только объём памяти (24 ГБ).