EasyLM

EasyLM

EasyLM — это универсальное решение для предварительного обучения, точной настройки, оценки и обслуживания LLM в JAX/Flax. Большие языковые модели (LLM) стали проще. EasyLM может масштабировать обучение LLM на сотни ускорителей TPU/GPU, используя функциональность JAX pjit.

Этот репозиторий, основанный на преобразователях и наборах данных Hugginface, предоставляет удобную в использовании и легко настраиваемую базу кода для обучения больших языковых моделей без сложностей, присущих многим другим платформам.

EasyLM построен на JAX/Flax. Используя утилиту pjit JAX, EasyLM может обучать большие модели, которые не умещаются на одном ускорителе, распределяя веса моделей и данные обучения по нескольким ускорителям. В настоящее время EasyLM поддерживает обучение нескольким хостам TPU/GPU на одном хосте, а также обучение на нескольких хостах на модулях Google Cloud TPU.

На данный момент поддерживаются следующие модели:

• ЛАМА
• ЛЛАМА 2
• ЛЛАМА 3

Дискорд-сервер

У нас есть неофициальное сообщество Discord (не связанное с Google) для обсуждения вопросов обучения LLM использованию JAX. Перейдите по этой ссылке, чтобы присоединиться к серверу Discord. У нас есть выделенные каналы для нескольких фреймворков LLM на базе JAX, включая EasyLM, JaxSeq, Alpa и Levanter.

Модели, обученные с помощью EasyLM

OpenLLaMA

OpenLLaMA — это наше лицензированное воспроизведение LLaMA, которое можно использовать в коммерческих целях. Посетите главную страницу проекта здесь. OpenLLaMA может служить заменой весов LLaMA в EasyLM. Более подробную информацию можно найти в документации LLaMA.

Коала

Koala — наш новый чат-бот, доработанный на базе LLaMA. Если вас интересует наш чат-бот Koala, вы можете ознакомиться с публикацией в блоге и документацией по его локальному запуску.

Установка

Метод установки различается на хостах GPU и хостах Cloud TPU. Первый шаг — загрузить с GitHub.

git clone https://github.com/young-geng/EasyLM.git
cd EasyLM
export PYTHONPATH= " ${PWD} : $PYTHONPATH " 

Установка на хост графического процессора

Среду графического процессора можно установить через Anaconda.

conda env create -f scripts/gpu_environment.yml
conda activate EasyLM

Установка на облачный хост TPU

Виртуальная машина хоста TPU поставляется с предустановленными Python и PIP. Просто запустите следующий сценарий, чтобы настроить хост TPU.

./scripts/tpu_vm_setup.sh

Документация

Документацию EasyLM можно найти в каталоге docs.

Ссылка

Если вы нашли EasyLM полезным в ваших исследованиях или приложениях, пожалуйста, цитируйте, используя следующий BibTeX:

 @software{geng2023easylm,
  author = {Geng, Xinyang},
  title = {EasyLM: A Simple And Scalable Training Framework for Large Language Models},
  month = March,
  year = 2023,
  url = {https://github.com/young-geng/EasyLM}
}

Кредиты

• Реализация LLaMA взята из JAX_llama.
• Реализация JAX/Flax GPT-J и RoBERTa взята из преобразователей.
• Большинство утилит JAX взяты из mlxu.
• Кодовая база во многом вдохновлена ​​JAXSeq.