🌟 LayerSkip: метод ускорения инференса в LLM.

LayerSkip - это способ ускорить LLM, уменьшая количество вычислений, необходимых для обработки каждого слова (токена) текста.

Основная идея LayerSkip - научить модель "угадывать" результат раньше, не проходя все слои. Для этого во время обучения модели некоторые слои "исключаются" (layer dropout) случайным образом. Помимо исключения слоев, LayerSkip использует специальную функцию потерь, которая помогает модели "понимать" результаты вычислений на более ранних слоях.

В отличие от других методов, LayerSkip использует одну и ту же LM head для всех слоев модели. Это упрощает обучение и уменьшает потребление памяти при инференсе.

Во время инференса LayerSkip обрабатывает текст только первыми "E" слоями, а затем сразу переходит к LM head, чтобы получить результат. Это называется "ранний выход" (early exit).

Чтобы повысить точность при раннем выходе, LayerSkip использует метод "самоспекулятивного декодирования". Модель сначала "угадывает" несколько следующих токенов, используя ранний выход. Затем она проверяет эти токены, используя оставшиеся слои, и исправляет ошибки.

LayerSkip был протестирован на различных наборах данных: Llama, CodeLlama и TOPv2. Результаты показали, что LayerSkip может ускорить работу LLM до 2 раз без значительного снижения точности.

Чтобы попробовать LayerSkip локально, разработчики предлагают использовать любую из 6 предобученных моделей:

🟢Llama2 - 7B и 13B;
🟢Codellama-7B или 34В;
🟢Llama3-8B:
🟢Llama3.2-1B.

⚠️ Для локального запуска будет нужен Huggingface API KEY.

▶️Локальный запуск:

# Clone repo
git clone git@github.com:facebookresearch/LayerSkip.git
cd LayerSkip

# Create env
conda create --name layer_skip python=3.10
conda activate layer_skip

# Install requirements
$ pip install -r requirements.txt

#Inference with self speculative 
$ torchrun generate.py --model facebook/layerskip-llama2-7B \
    --sample True \
    --max_steps 512 \
    --generation_strategy self_speculative \
    --exit_layer 8 \
    --num_speculations 6

▶️Ключи запуска:

--model: имя модели на HuggingFace;
--sample: включение/выключение семплирования (по умолчанию: True);
--max_steps: максимальное количество генерируемых токенов;
--generation_strategy: стратегия генерации (по умолчанию: 'greedy', для LayerSkip: 'self_speculative');
--exit_layer: номер слоя для раннего выхода;
--num_speculations: количество спекулятивных токенов;


🟡Коллекция моделей на HF
🟡Arxiv
🖥GitHub


@ai_machinelearning_big_data

#AI #ML #LLM #MetaAI #LayerSkip