lightning thunder
Initial release
Torne os modelos PyTorch extremamente rápidos.
Lightning.ai • Desempenho • Primeiros passos • Instalação • Exemplos • Por dentro do Thunder • Participe! • Documentação
Thunder torna os modelos PyTorch extremamente rápidos.
Thunder é um compilador fonte-a-fonte para PyTorch. Ele torna os programas PyTorch mais rápidos combinando e usando diferentes executores de hardware ao mesmo tempo (por exemplo, nvFuser, torch.compile, cuDNN e TransformerEngine FP8).
Ele suporta configurações de GPU única e múltipla. O Thunder pretende ser utilizável, compreensível e extensível.
Observação
Lightning Thunder está em alfa. Sinta-se à vontade para se envolver, mas espere alguns obstáculos ao longo do caminho.
O Thunder pode alcançar acelerações significativas em relação ao código PyTorch não compilado padrão ("PyTorch ansioso"), por meio dos efeitos compostos de otimizações e do uso dos melhores executores da categoria. A figura abaixo mostra o rendimento de pré-treinamento para Llama 2 7B conforme implementado em LitGPT.
Conforme mostrado no gráfico acima, o Thunder atinge uma aceleração de 40% no rendimento do treinamento em comparação com o código ansioso no H100 usando uma combinação de executores, incluindo nvFuser, torch.compile, cuDNN e TransformerEngine FP8.
Thunder também suporta estratégias distribuídas como DDP e FSDP para modelos de treinamento em múltiplas GPUs. O gráfico a seguir exibe o rendimento normalizado medido para Llama 2 7B sem precisão mista do FP8; o apoio ao FSDP está em curso.
A maneira mais fácil de começar a usar o Thunder, sem precisar de instalações ou configurações extras, é usar nosso Zero to Thunder Tutorial Studio.
Thunder está em alfa e o desenvolvimento mais recente está acontecendo no branch main . Você pode instalar a versão mais recente do Thunder a partir do branch main da seguinte maneira:
pip install git+https://github.com/Lightning-AI/lightning-thunder.git@mainPara obter o melhor desempenho, você pode instalar o Thunder com as seguintes dependências adicionais:
# install nvFuser built for the matching stable PyTorch
pip install --pre nvfuser-cu121-torch25 # install cudnn
pip install nvidia-cudnn-frontendSe você estiver interessado em mexer e contribuir com o Thunder, recomendamos clonar o repositório Thunder e instalá-lo no modo editável do pip:
git clone https://github.com/Lightning-AI/lightning-thunder.git
cd lightning-thunder
pip install -e .Após clonar o repositório relâmpago-thunder e instalá-lo como um pacote editável conforme explicado acima, você pode configurar seu ambiente para desenvolver o Thunder instalando os requisitos de desenvolvimento:
pip install -r requirements/devel.txtAgora você executa testes:
pytest thunder/testsO Thunder foi testado exaustivamente, então espere que demore um pouco.
Abaixo está um exemplo simples de como o Thunder permite compilar e executar código 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]]) A função compilada jfoo pega e retorna tensores PyTorch, assim como a função original, então módulos e funções compilados pelo Thunder podem ser usados como parte de programas PyTorch maiores.
O Thunder está em seus estágios iniciais e ainda não deve ser usado para produção.
No entanto, ele já pode oferecer excelente desempenho para pré-treinamento e ajuste fino de LLMs suportados por LitGPT, como Mistral, Llama 2, Gemma, Falcon e outros.
Confira a integração do LitGPT para saber como executar o LitGPT e o Thunder juntos.
Dado um módulo Python chamável ou PyTorch, o Thunder pode gerar um programa otimizado que:
Para fazer isso, o Thunder vem com:
grad , fusões, distribuídas (como ddp , fsdp ), funcionais (como vmap , vjp , jvp )Thunder é escrito inteiramente em Python. Até mesmo seu rastreamento é representado como Python válido em todos os estágios de transformação. Isso permite níveis sem precedentes de introspecção e extensibilidade.
O Thunder não gera código para aceleradores, como GPUs, diretamente. Ele adquire e transforma programas de usuário para que seja possível selecionar ou gerar código de dispositivo de maneira otimizada usando executores rápidos como:
Módulos e funções compilados com Thunder interoperam totalmente com o vanilla PyTorch e suportam a autogradação do PyTorch. Além disso, o Thunder trabalha junto com torch.compile para aproveitar suas otimizações de última geração.
A documentação on-line está disponível. Para construir documentação localmente você pode usar
make docs e aponte seu navegador para os documentos gerados em docs/build/index.html .
Agradecemos seus comentários e contribuições. Se você tiver solicitações de recursos, perguntas ou quiser contribuir com código ou arquivos de configuração, não hesite em usar o rastreador de problemas do GitHub.
Damos as boas-vindas a todos os colaboradores individuais, independentemente do seu nível de experiência ou hardware. Suas contribuições são valiosas e estamos entusiasmados em ver o que você pode realizar neste ambiente colaborativo e de apoio.
Lightning Thunder é lançado sob a licença Apache 2.0. Consulte o arquivo LICENSE para obter detalhes.
