PyTorch 프로그램을 한층 빠르게 만들어주는 소스 대 소스 컴파일러, 
Lightning Thunder를 소개합니다. PyTorch 모델을 더욱 빠르게 만들기 위해 다양한 하드웨어 실행자(nvFuser, torch.compile, cuDNN, TransformerEngine FP8 등)를 결합해 사용합니다. 단일 가속기는 물론 멀티-GPU 설정에서도 작동하여 사용하기 쉽고, 이해하기 쉬우며, 확장 가능한 것을 목표로 합니다.
주의: Lightning Thunder는 현재 알파버전입니다.
소개
Thunder는 PyTorch에 대한 소스 대 소스 컴파일러로, 최적화의 복합 효과와 최고의 실행자 사용을 통해 표준 PyTorch 즉시 코드보다 상당한 속도 향상을 달성할 수 있습니다. 예를 들어, Llama 2 7B의 사전 훈련 처리량에서 H100을 사용하여 기존의 즉시 코드에 비해 40%의 속도 향상을 이루었습니다. Thunder는 DDP 및 FSDP(ZeRO2 및 ZeRO3)와 같은 분산 전략도 지원합니다.
주요 기능
Thunder는 Python 호출 가능 객체 또는 PyTorch 모듈을 최적화된 프로그램으로 변환할 수 있으며, 이는 다음을 수행합니다:
- 순전파 및 역전파 계산
- 연산을 효율적인 융합 영역으로 결합
- 최적화된 커널로 계산 전달
- 기계 간에 계산을 최적으로 분배
Thunder는 다양한 최적화 실행자를 통해 디바이스 코드를 선택하거나 생성하는 것이 가능하게 해주며, 이는 Thunder를 PyTorch와의 완벽한 호환성과 함께 사용할 수 있게 합니다.
사용 방법
Thunder를 사용하기 위해서는 먼저 nvFuser를 설치하고 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)
이 예제에서 jfoo 컴파일된 함수는 PyTorch 텐서를 받아들이고 반환하므로, Thunder로 컴파일된 모듈과 함수를 더 큰 PyTorch 프로그램의 일부로 사용할 수 있습니다.
비교
Thunder는 기존 PyTorch 즉시 코드에 비해 상당한 성능 향상을 제공합니다. 특히, 다양한 하드웨어 실행자를 결합하여 사용함으로써 단일 가속기 및 멀티-GPU 환경 모두에서 뛰어난 성능을 발휘합니다.
더 읽어보기
Lightning Thunder GitHub 저장소
Excutors
Thunder는 가속기용 코드를 직접 생성하지 않습니다. 사용자 프로그램을 수집하고 변환하여 다음과 같은 빠른 실행기를 사용하여 디바이스 코드를 최적으로 선택하거나 생성할 수 있도록 합니다:
- nvFuser: nvFuser
- torch.compile: torch.compile
- cuDNN: cuDNN
- TransformerEngine: TransformerEngine
이 글은 GPT 모델로 정리한 글을 바탕으로 한 것으로, 원문의 내용 또는 의도와 다르게 정리된 내용이 있을 수 있으니 원문도 함께 참고해주세요! 읽으시면서 어색하거나 잘못된 내용을 발견하시면 덧글로 알려주시기를 부탁드립니다. 
파이토치 한국 사용자 모임
이 정리한 이 글이 유용하셨나요? 회원으로 가입하시면 주요 글들을 이메일
로 보내드립니다! (기본은 Weekly지만 Daily로 변경도 가능합니다.)
아래
쪽에 좋아요
를 눌러주시면 힘이 됩니다~ 