HyperAI초신경

PyTorch 2.0 실습: HuggingFace와 TIMM 모델 속도 향상!

2년 전
정보
Jiaxin Sun
特色图像

PyTorch 2.0은 간단한 torch.compile() 줄을 통해 모델 학습 속도를 30%-200%까지 높일 수 있습니다. 이 튜토리얼에서는 실제로 이러한 속도 향상을 재현하는 방법을 보여드립니다.

토치.컴파일()  할 수 있다다양한 컴파일러 백엔드를 쉽게 시도해 보세요.이렇게 하면 PyTorch 코드의 실행 속도가 빨라집니다. 그것은 토치.jit.스크립트()  소스 코드를 수정하지 않고도 nn.Module에서 직접 실행할 수 있는 대체 프로그램입니다.

이전 글에서는 torch.compile이 임의의 PyTorch 코드, 제어 흐름, 변형을 지원하고 어느 정도 동적 모양을 지원한다는 것을 소개했습니다.

163개의 오픈 소스 모델을 테스트한 결과, torch.compile()을 사용하면 30%~200%의 속도 향상이 가능하다는 것을 발견했습니다.

opt_module = torch.compile(module)

테스트 결과는 다음에 자세히 나와 있습니다.

https://github.com/pytorch/torchdynamo/issues/681

이 튜토리얼에서는 사용 방법을 보여줍니다. 토치.컴파일()  모델 학습 속도를 높입니다.

요구 사항 및 설정

GPU의 경우(최신 GPU일수록 성능이 더욱 크게 향상됨):

pip3 install numpy --pre torch[dynamo] --force-reinstall --extra-index-url https://download.pytorch.org/whl/nightly/cu117

CPU의 경우:

pip3 install --pre torch --extra-index-url https://download.pytorch.org/whl/nightly/cpu

선택 사항: 설치를 확인하세요

git clone https://github.com/pytorch/pytorch
cd tools/dynamo
python verify_dynamo.py

선택 사항: Docker 설치

모든 필수 종속성은 PyTorch Nightly Binaries 파일에 제공되며, 다음을 통해 다운로드할 수 있습니다.

docker pull ghcr.io/pytorch/pytorch-nightly

임시 실험의 경우,컨테이너가 모든 GPU에 액세스할 수 있는지 확인하세요.

docker run --gpus all -it ghcr.io/pytorch/pytorch-nightly:latest /bin/bash

시작

  간단한 예

먼저 간단한 예를 살펴보겠습니다. 최신 GPU일수록 속도 향상이 더 두드러진다는 점에 주목하세요.

import torch
def fn(x, y):
a = torch.sin(x).cuda()
b = torch.sin(y).cuda()
return a + b
new_fn = torch.compile(fn, backend="inductor")
input_tensor = torch.randn(10000).to(device="cuda:0")
a = new_fn()

이 예제는 실제로 속도를 증가시키지는 않지만 시작하는 데 사용할 수 있습니다.

이 예에서,토치.cos()  그리고 토치.sin()  는 점별 연산의 예이며, 벡터의 요소별로 연산을 수행합니다. 더 유명한 지점별 op는 다음과 같습니다. 토치.relu().

열망 모드에서의 지점별 작업은 각 연산자가 메모리에서 텐서를 읽고, 몇 가지 변경을 한 다음, 이 변경 사항을 다시 써야 하기 때문에 최적이 아닙니다.

PyTorch 2.0의 가장 중요한 최적화 중 하나는 퓨전입니다.

따라서 이 경우 2번의 읽기와 2번의 쓰기를 1번의 읽기와 1번의 쓰기로 바꿀 수 있습니다. 이는 병목 현상이 컴퓨팅(GPU가 부동 소수점 연산을 얼마나 빨리 수행할 수 있는지)이 아니라 메모리 대역폭(GPU로 데이터를 얼마나 빨리 보낼 수 있는지)인 최신 GPU에 중요합니다.

PyTorch 2.0의 두 번째 중요한 최적화는 CUDA 그래프입니다.

CUDA 그래프는 Python 프로그램에서 개별 커널을 실행하는 오버헤드를 제거하는 데 도움이 됩니다.

torch.compile()은 다양한 백엔드를 지원하는 데, 그 중 가장 주목할 만한 것은 Triton 커널을 생성할 수 있는 Inductor입니다.

https://github.com/openai/triton

이 커널은 Python으로 작성되었습니다.하지만 대부분의 손으로 작성된 CUDA 커널보다는 낫습니다.위의 예제가 trig.py라고 가정하면 실제로 다음을 실행하여 Triton 커널을 생성하는 코드를 검사할 수 있습니다.

TORCHINDUCTOR_TRACE=1 python trig.py
@pointwise(size_hints=[16384], filename=__file__, meta={'signature': {0: '*fp32', 1: '*fp32', 2: 'i32'}, 'device': 0, 'constants': {}, 'configs': [instance_descriptor(divisible_by_16=(0, 1, 2), equal_to_1=())]})
@triton.jit
def kernel(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr):
xnumel = 10000
xoffset = tl.program_id(0) * XBLOCK
xindex = xoffset + tl.reshape(tl.arange(0, XBLOCK), [XBLOCK])
xmask = xindex < xnumel
x0 = xindex
tmp0 = tl.load(in_ptr0 + (x0), xmask)
tmp1 = tl.sin(tmp0)
tmp2 = tl.sin(tmp1)
tl.store(out_ptr0 + (x0 + tl.zeros([XBLOCK], tl.int32)), tmp2, xmask)

위의 코드에서 다음을 확인할 수 있습니다.   융합은 두 가지 때문에 발생했습니다.   이러한 작업은 Triton 커널에서 수행되고, 임시 변수는 레지스터에 저장되므로 액세스 속도가 매우 빠릅니다.

실제 모델 예제

PyTorch Hub의 resnet50을 예로 들어 보겠습니다.

import torch
model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True)
opt_model = torch.compile(model, backend="inductor")
model(torch.randn(1,3,64,64))

실제 운영에서는 모델이 컴파일되고 있기 때문에 첫 번째 실행이 매우 느리다는 것을 알 수 있습니다. 이후 작업 속도가 더 빨라질 것입니다.따라서 벤치마킹을 시작하기 전에 모델을 워밍업하는 것이 일반적입니다.

보시다시피, 여기서는 컴파일러 이름을 나타내기 위해 "inductor"를 사용했지만 사용 가능한 백엔드는 이것뿐이 아닙니다. REPL에서 실행할 수 있습니다 토치._다이나모.리스트_백엔드()  사용 가능한 백엔드의 전체 목록을 보려면 클릭하세요.

또한 시도할 수도 있습니다 aot_cudagraphs  또는 엔비퓨저  .

Hugging Face 모델 예시

PyTorch 커뮤니티는 종종 변환기 또는 TIMM의 사전 학습된 모델을 사용합니다.

https://github.com/huggingface/transformers
https://github.com/rwightman/pytorch-image-models

PyTorch 2.0의 디자인 목표 중 하나는 실제로 실행되는 대부분의 모델에서 모든 컴파일 스택을 바로 사용할 수 있어야 한다는 것입니다.

여기서는 HuggingFace 허브에서 사전 학습된 모델을 직접 다운로드하여 최적화합니다.

import torch
from transformers import BertTokenizer, BertModel
# Copy pasted from here https://huggingface.co/bert-base-uncased
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained("bert-base-uncased").to(device="cuda:0")
model = torch.compile(model) # This is the only line of code that we changed
text = "Replace me by any text you'd like."
encoded_input = tokenizer(text, return_tensors='pt').to(device="cuda:0")
output = model(**encoded_input)

모델에서 제거하면 to(장치="cuda:0")  그리고 인코딩된 입력 PyTorch 2.0은 CPU에서 실행되도록 최적화된 C++ 커널을 생성합니다.

위의 삼각 함수 예제보다 훨씬 복잡한 BERT의 Triton이나 C++ 커널을 확인해 보세요. 하지만 PyTorch를 알고 있다면 건너뛸 수 있습니다.

더 나은 결과를 얻으려면 동일한 코드를 다음 코드에 사용할 수 있습니다.

* https://github.com/huggingface/accelerate

* 동대문디자인플라자

다시 TIMM 예를 들어보세요.

import timm
import torch
model = timm.create_model('resnext101_32x8d', pretrained=True, num_classes=2)
opt_model = torch.compile(model, backend="inductor")
opt_model(torch.randn(64,3,7,7))

PyTorch의 목표는 더 많은 모델에 적응하고 대부분의 오픈 소스 모델의 작업 속도를 높일 수 있는 컴파일러를 만드는 것입니다.지금 HuggingFace 허브를 방문하세요,PyTorch 2.0으로 TIMM 모델의 속도를 높여보세요!

https://huggingface.co/timm