lightning thunder
Initial release
Сделайте модели PyTorch молниеносными.
Lightning.ai • Производительность • Начало работы • Установка • Примеры • Внутри Thunder • Принимайте участие! • Документация
Thunder делает модели PyTorch Lightning быстрыми.
Thunder — это компилятор исходного кода для PyTorch. Это ускоряет работу программ PyTorch за счет комбинирования и одновременного использования различных аппаратных исполнителей (например, nvFuser, torch.compile, cuDNN и TransformerEngine FP8).
Он поддерживает конфигурации как с одним, так и с несколькими графическими процессорами. Thunder стремится быть удобным, понятным и расширяемым.
Примечание
Lightning Thunder находится в стадии альфа-версии. Не стесняйтесь принимать участие, но ожидайте некоторых препятствий на этом пути.
Thunder может добиться значительного ускорения по сравнению со стандартным некомпилируемым кодом PyTorch («PyTorch нетерпеливым») за счет комплексного эффекта оптимизации и использования лучших в своем классе исполнителей. На рисунке ниже показана пропускная способность предварительного обучения для Llama 2 7B, реализованная в LitGPT.
Как показано на графике выше, Thunder достигает 40%-ного ускорения производительности обучения по сравнению с энергичным кодом на H100, используя комбинацию исполнителей, включая nvFuser, torch.compile, cuDNN и TransformerEngine FP8.
Thunder также поддерживает распределенные стратегии, такие как DDP и FSDP, для обучения моделей на нескольких графических процессорах. На следующем графике показана нормализованная пропускная способность, измеренная для Llama 2 7B без смешанной точности FP8; поддержка FSDP продолжается.
Самый простой способ начать работу с Thunder, не требующий дополнительных установок или настроек, — это использовать нашу студию учебных пособий от Zero to Thunder.
Thunder находится в стадии альфа-версии, и последние разработки происходят в main ветке. Вы можете установить последнюю версию Thunder из main ветки следующим образом:
pip install git+https://github.com/Lightning-AI/lightning-thunder.git@mainДля достижения максимальной производительности вы можете установить Thunder со следующими дополнительными зависимостями:
# install nvFuser built for the matching stable PyTorch
pip install --pre nvfuser-cu121-torch25 # install cudnn
pip install nvidia-cudnn-frontendЕсли вы хотите поработать с Thunder и внести свой вклад в его развитие, мы рекомендуем клонировать репозиторий Thunder и установить его в редактируемом режиме pip:
git clone https://github.com/Lightning-AI/lightning-thunder.git
cd lightning-thunder
pip install -e .После клонирования репозитория Lightning-Thunder и установки его в виде редактируемого пакета, как описано выше, вы можете настроить свою среду для разработки Thunder, установив требования для разработки:
pip install -r requirements/devel.txtТеперь вы запускаете тесты:
pytest thunder/testsThunder очень тщательно протестирован, так что ожидайте, что это займет некоторое время.
Ниже приведен простой пример того, как Thunder позволяет компилировать и запускать код PyTorch:
import torch
import thunder
def foo ( a , b ):
return a + b
jfoo = thunder . jit ( foo )
a = torch . full (( 2 , 2 ), 1 )
b = torch . full (( 2 , 2 ), 3 )
result = jfoo ( a , b )
print ( result )
# prints
# tensor(
# [[4, 4]
# [4, 4]]) Скомпилированная функция jfoo принимает и возвращает тензоры PyTorch, как и исходная функция, поэтому модули и функции, скомпилированные Thunder, можно использовать как часть более крупных программ PyTorch.
Thunder находится на ранней стадии разработки и пока не должен использоваться в производственных целях.
Тем не менее, он уже может обеспечить выдающуюся производительность для предварительной подготовки и точной настройки LLM, поддерживаемых LitGPT, таких как Mistral, Llama 2, Gemma, Falcon и других.
Ознакомьтесь с интеграцией LitGPT, чтобы узнать о совместной работе LitGPT и Thunder.
Учитывая вызываемый модуль Python или модуль PyTorch, Thunder может создать оптимизированную программу, которая:
Для этого в комплект поставки Thunder входят:
grad , fusion, распределенные (например, ddp , fsdp ), функциональные (например, vmap , vjp , jvp )Thunder полностью написан на Python. Даже его трассировка представляется как валидный Python на всех этапах трансформации. Это обеспечивает беспрецедентный уровень самоанализа и расширяемости.
Thunder не генерирует код для ускорителей, таких как графические процессоры, напрямую. Он получает и преобразует пользовательские программы, чтобы можно было оптимально выбирать или генерировать код устройства с помощью быстрых исполнителей, таких как:
Модули и функции, скомпилированные с помощью Thunder, полностью совместимы с ванильным PyTorch и поддерживают автоградацию PyTorch. Кроме того, Thunder работает вместе с torch.compile, чтобы использовать его самые современные оптимизации.
Доступна онлайн-документация. Для локальной сборки документации вы можете использовать
make docs и укажите в браузере сгенерированные документы по адресу docs/build/index.html .
Мы ценим ваши отзывы и вклад. Если у вас есть пожелания по функциям, вопросы или вы хотите добавить код или файлы конфигурации, не стесняйтесь использовать систему отслеживания проблем GitHub.
Мы приветствуем всех индивидуальных участников, независимо от их уровня опыта и оборудования. Ваш вклад ценен, и мы рады видеть, чего вы можете достичь в этой среде сотрудничества и поддержки.
Lightning Thunder выпускается под лицензией Apache 2.0. Подробности смотрите в файле ЛИЦЕНЗИИ.
