lightning thunder

PyTorch 모델을 빛처럼 빠르게 만듭니다.

---

Lightning.ai • 성능 • 시작하기 • 설치 • 예제 • Inside Thunder • 참여해 보세요! • 문서

⚡ 라이트닝 썬더에 오신 것을 환영합니다.

Thunder는 PyTorch 모델을 빛처럼 빠르게 만듭니다.

Thunder는 PyTorch용 소스-소스 컴파일러입니다. 다양한 하드웨어 실행기(예: nvFuser, torch.compile, cuDNN 및 TransformerEngine FP8)를 한 번에 결합하고 사용하여 PyTorch 프로그램을 더 빠르게 만듭니다.

단일 및 다중 GPU 구성을 모두 지원합니다. Thunder는 사용 가능하고 이해 가능하며 확장 가능하다는 것을 목표로 합니다.

메모

라이트닝 썬더는 알파 버전입니다. 자유롭게 참여하세요. 하지만 그 과정에서 몇 가지 어려움을 겪을 수 있습니다.

단일 GPU 성능

Thunder는 최적화와 동급 최고의 실행기 사용의 복합 효과를 통해 컴파일되지 않은 표준 PyTorch 코드("PyTorch 열성")에 비해 상당한 속도 향상을 달성할 수 있습니다. 아래 그림은 LitGPT에서 구현된 Llama 2 7B의 사전 훈련 처리량을 보여줍니다.

위 플롯에서 볼 수 있듯이 Thunder는 nvFuser, torch.compile, cuDNN 및 TransformerEngine FP8을 포함한 실행기 조합을 사용하여 H100의 Eager 코드에 비해 훈련 처리량에서 40% 속도 향상을 달성합니다.

다중 GPU 성능

Thunder는 또한 여러 GPU에서 모델을 훈련하기 위해 DDP 및 FSDP와 같은 분산 전략을 지원합니다. 다음 플롯은 FP8 혼합 정밀도 없이 Llama 2 7B에 대해 측정된 정규화된 처리량을 표시합니다. FSDP 지원이 진행 중입니다.

시작하기

추가 설치나 설정이 필요 없이 Thunder를 시작하는 가장 쉬운 방법은 Zero to Thunder Tutorial Studio를 사용하는 것입니다.

썬더 설치

Thunder는 알파 버전이며 main 브랜치에서 최신 개발이 진행되고 있습니다. 다음과 같이 main 브랜치에서 최신 버전의 Thunder를 설치할 수 있습니다:

pip install git+https://github.com/Lightning-AI/lightning-thunder.git@main

최상의 성능을 얻으려면 다음과 같은 추가 종속성을 사용하여 Thunder를 설치할 수 있습니다.

• 다음과 같이 PyTorch 2.5.1용으로 빌드된 사전 릴리스 nvFuser를 설치합니다.

 # install nvFuser built for the matching stable PyTorch
pip install --pre nvfuser-cu121-torch25

• 다음과 같이 cudnn을 설치합니다:

 # 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/tests

Thunder는 매우 철저한 테스트를 거쳤으므로 시간이 좀 걸릴 것으로 예상됩니다.

안녕하세요 세계

다음은 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는 초기 단계에 있으므로 아직 프로덕션 실행에 사용해서는 안 됩니다.

그러나 이미 Mistral, Llama 2, Gemma, Falcon 등과 같이 LitGPT에서 지원하는 LLM 사전 훈련 및 미세 조정에 탁월한 성능을 제공할 수 있습니다.

LitGPT와 Thunder를 함께 실행하는 방법을 알아보려면 LitGPT 통합을 확인하세요.

Inside Thunder: 핵심 기능 간략히 살펴보기

Python 호출 가능 또는 PyTorch 모듈이 제공되면 Thunder는 다음과 같은 최적화된 프로그램을 생성할 수 있습니다.

• 정방향 및 역방향 패스를 계산합니다.
• 운영을 효율적인 융합 영역으로 통합합니다.
• 최적화된 커널에 계산을 전달합니다.
• 컴퓨터 전체에 계산을 최적으로 분산합니다.

이를 위해 Thunder는 다음을 제공합니다.

• PyTorch 및 사용자 지정 작업을 대상으로 하는 Python 프로그램을 획득하기 위한 JIT
• 축소된 작업 집합의 추적으로 작업을 표현하는 다중 레벨 중간 표현(IR)
• grad , fusions, 분산( ddp , fsdp 등), 기능적( vmap , vjp , jvp )과 같은 계산 그래프 추적에 대한 확장 가능한 변환 세트
• 확장 가능한 실행자 컬렉션에 작업을 디스패치하는 방법

Thunder는 전적으로 Python으로 작성되었습니다. 그 추적조차도 변환의 모든 단계에서 유효한 Python으로 표시됩니다. 이를 통해 전례 없는 수준의 내성과 확장성이 가능해졌습니다.

Thunder는 GPU와 같은 가속기에 대한 코드를 직접 생성하지 않습니다. 다음과 같은 빠른 실행기를 사용하여 장치 코드를 최적으로 선택하거나 생성할 수 있도록 사용자 프로그램을 획득하고 변환합니다.

• 토치.컴파일
• nvFuser
• cuDNN
• 꼭대기
• 변압기 엔진
• PyTorch 열망
• PyCUDA, Numba, CuPy를 통한 맞춤형 CUDA 커널
• OpenAI Triton으로 작성된 맞춤형 커널

Thunder로 컴파일된 모듈과 기능은 바닐라 PyTorch와 완벽하게 상호 운용되며 PyTorch의 autograd를 지원합니다. 또한 Thunder는 torch.compile과 함께 작동하여 최첨단 최적화를 활용합니다.

선적 서류 비치

온라인 문서를 사용할 수 있습니다. 로컬에서 문서를 작성하려면 다음을 사용할 수 있습니다.

make docs

브라우저에서 docs/build/index.html 에 생성된 문서를 가리키도록 하세요.

참여해보세요!

여러분의 피드백과 기여에 감사드립니다. 기능 요청이나 질문이 있거나 코드 또는 구성 파일을 제공하고 싶다면 주저하지 말고 GitHub 문제 추적기를 사용하세요.

우리는 경험이나 하드웨어 수준에 관계없이 모든 개별 기여자를 환영합니다. 귀하의 기여는 가치가 있으며, 우리는 귀하가 이 협력적이고 지원적인 환경에서 무엇을 성취할 수 있는지 기대하고 있습니다.

특허

Lightning Thunder는 Apache 2.0 라이센스에 따라 출시됩니다. 자세한 내용은 LICENSE 파일을 참조하세요.