lightning thunder
Initial release
Haga modelos de PyTorch a la velocidad del rayo.
Lightning.ai • Rendimiento • Comenzar • Instalar • Ejemplos • Dentro de Thunder • ¡Participe! • Documentación
Thunder hace que los modelos PyTorch sean increíblemente rápidos.
Thunder es un compilador de fuente a fuente para PyTorch. Hace que los programas PyTorch sean más rápidos al combinar y usar diferentes ejecutores de hardware a la vez (por ejemplo, nvFuser, torch.compile, cuDNN y TransformerEngine FP8).
Admite configuraciones de GPU única y múltiple. Thunder pretende ser utilizable, comprensible y extensible.
Nota
Lightning Thunder está en alfa. Siéntete libre de involucrarte, pero espera algunos obstáculos en el camino.
Thunder puede lograr aceleraciones significativas sobre el código PyTorch estándar no compilado ("PyTorch ansioso"), mediante los efectos compuestos de las optimizaciones y el uso de los mejores ejecutores de su clase. La siguiente figura muestra el rendimiento del preentrenamiento para Llama 2 7B implementado en LitGPT.
Como se muestra en el gráfico anterior, Thunder logra una aceleración del 40 % en el rendimiento del entrenamiento en comparación con el código ansioso en H100 utilizando una combinación de ejecutores que incluyen nvFuser, torch.compile, cuDNN y TransformerEngine FP8.
Thunder también admite estrategias distribuidas como DDP y FSDP para entrenar modelos en múltiples GPU. El siguiente gráfico muestra el rendimiento normalizado medido para Llama 2 7B sin precisión mixta del FP8; el soporte para FSDP está en progreso.
La forma más sencilla de empezar a utilizar Thunder, que no requiere instalaciones ni configuraciones adicionales, es utilizando nuestro Tutorial Studio de Zero to Thunder.
Thunder está en alfa y el último desarrollo se está produciendo en la rama main . Puede instalar la última versión de Thunder desde la rama main de la siguiente manera:
pip install git+https://github.com/Lightning-AI/lightning-thunder.git@mainPara lograr el mejor rendimiento, puedes instalar Thunder con las siguientes dependencias adicionales:
# install nvFuser built for the matching stable PyTorch
pip install --pre nvfuser-cu121-torch25 # install cudnn
pip install nvidia-cudnn-frontendSi está interesado en jugar y contribuir con Thunder, le recomendamos clonar el repositorio de Thunder e instalarlo en el modo editable de pip:
git clone https://github.com/Lightning-AI/lightning-thunder.git
cd lightning-thunder
pip install -e .Después de clonar el repositorio de Lightning-Thunder e instalarlo como un paquete editable como se explicó anteriormente, puede configurar su entorno para desarrollar Thunder instalando los requisitos de desarrollo:
pip install -r requirements/devel.txtAhora ejecutas pruebas:
pytest thunder/testsThunder se prueba minuciosamente, así que espere que esto lleve un tiempo.
A continuación se muestra un ejemplo simple de cómo Thunder le permite compilar y ejecutar 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]]) La función compilada jfoo toma y devuelve tensores de PyTorch, al igual que la función original, por lo que los módulos y funciones compilados por Thunder se pueden usar como parte de programas PyTorch más grandes.
Thunder se encuentra en sus primeras etapas y aún no debería usarse para tiradas de producción.
Sin embargo, ya puede ofrecer un rendimiento sobresaliente para la capacitación previa y el ajuste de LLM respaldados por LitGPT, como Mistral, Llama 2, Gemma, Falcon y otros.
Consulte la integración de LitGPT para aprender cómo ejecutar LitGPT y Thunder juntos.
Dado un módulo Python invocable o PyTorch, Thunder puede generar un programa optimizado que:
Para ello, Thunder se entrega con:
grad , fusiones, distribuidas (como ddp , fsdp ), funcionales (como vmap , vjp , jvp ).Thunder está escrito completamente en Python. Incluso su rastro se representa como Python válido en todas las etapas de transformación. Esto permite niveles sin precedentes de introspección y extensibilidad.
Thunder no genera código para aceleradores, como GPU, directamente. Adquiere y transforma programas de usuario para que sea posible seleccionar o generar código de dispositivo de manera óptima utilizando ejecutores rápidos como:
Los módulos y funciones compilados con Thunder interoperan completamente con PyTorch básico y son compatibles con la autograduación de PyTorch. Además, Thunder trabaja junto con torch.compile para aprovechar sus optimizaciones de última generación.
La documentación en línea está disponible. Para crear documentación localmente puedes usar
make docs y apunte su navegador a los documentos generados en docs/build/index.html .
Agradecemos sus comentarios y contribuciones. Si tiene solicitudes de funciones, preguntas o desea contribuir con código o archivos de configuración, no dude en utilizar el rastreador de problemas de GitHub.
Damos la bienvenida a todos los contribuyentes individuales, independientemente de su nivel de experiencia o hardware. Sus contribuciones son valiosas y estamos entusiasmados de ver lo que puede lograr en este entorno de colaboración y apoyo.
Lightning Thunder se lanza bajo la licencia Apache 2.0. Consulte el archivo de LICENCIA para obtener más detalles.
