PyTorch 팀은 원래 구현보다 8배 빠르게 "모든 항목 분할" 모델을 다시 구현했습니다.

연초부터 지금까지 제너레이티브 AI는 빠르게 발전해왔습니다. 그러나 우리는 특히 PyTorch를 사용할 때 생성 AI의 훈련, 추론 등의 속도를 어떻게 높일 것인가라는 어려운 문제에 직면할 때가 많습니다.

이 기사에서는 PyTorch 팀의 연구원들이 솔루션을 제공합니다. 이 기사에서는 순수 네이티브 PyTorch를 사용하여 생성 AI 모델을 가속화하는 방법에 중점을 두고 있으며, 새로운 PyTorch 기능과 이를 결합하는 방법에 대한 실제 사례도 소개합니다.

결과는 어땠나요? PyTorch 팀은 Meta의 "모든 것을 분할"(SAM) 모델을 다시 작성하여 정확성을 잃지 않고 원래 구현보다 8배 빠른 코드를 생성했으며 모두 기본 PyTorch를 사용하여 최적화되었다고 말했습니다.

블로그 주소: https://pytorch.org/blog/accelerating-generative-ai/

이 글을 읽고 나면 다음과 같은 이해를 얻게 될 것입니다:

• Torch 컴파일. : PyTorch 모델 컴파일러인 PyTorch 2.0에는 한 줄의 코드로 기존 모델을 가속화할 수 있는 torch.compile()이라는 새로운 기능이 추가되었습니다.
• GPU 양자화: 계산 정확도를 줄여 모델을 가속화합니다. (Scaled Dot Product Attention): 메모리 효율적인 주의 구현
• 반구조적(2:4) 희소성: GPU에 최적화된 희소 메모리 형식
• Nested Tensor: 중첩된 Tensor 팩 {텐서, 마스크 } 서로 다른 크기의 이미지와 같이 균일하지 않은 크기의 데이터를 단일 텐서로 일괄 처리합니다.
• Triton 사용자 정의 작업: Triton Python DSL을 사용하여 GPU 작업을 작성하고 사용자 정의 연산자를 통해 PyTorch의 다양한 구성 요소에 쉽게 통합합니다. 등록.

PyTorch 기본 기능은 처리량을 높이고 메모리 오버헤드를 줄입니다.

본 연구에 대한 자세한 내용은 Meta에서 제안한 SAM을 참조하세요. 자세한 기사는 "CV는 더 이상 존재하지 않습니까? 메타가 "모든 것을 분할" AI 모델 출시, CV가 GPT-3 순간을 가져올 수도 있습니다"

다음으로 성능을 포함한 SAM의 최적화 프로세스를 소개하겠습니다. 분석, 병목 현상 식별 및 이러한 새로운 기능을 PyTorch에 통합하여 SAM이 직면한 문제를 해결하는 방법. 또한 torch.compile, SDPA, Triton 커널, Nested Tensor 및 반구조적 희소성(반구조적 희소성)을 포함하여 PyTorch의 몇 가지 새로운 기능도 소개합니다.

콘텐츠는 계층별로 심층적으로 설명됩니다. 이 글의 마지막 부분에서는 빠른 버전의 SAM을 소개하겠습니다. 관심 있는 독자는 GitHub에서 다운로드할 수 있습니다. 또한 이러한 데이터는 Perfetto UI를 사용하여 시각화되어 PyTorch

의 다양한 기능의 응용 가치를 보여줍니다. 이 프로젝트는 GitHub 주소: https://github.com/pytorch-labs/segment-anything-fast에서 확인할 수 있습니다.

의 소스 코드는 분할된 모든 모델 SAM

을 다시 작성합니다. 연구에 따르면 이 기사에서 사용된 SAM 기준 데이터 유형은 float32 dtype이고 배치 크기는 1이며 PyTorch Profiler를 사용하여 코어 추적 결과는 다음과 같습니다.

이 기사에서는 SAM에 최적화할 수 있는 두 곳이 있음을 발견했습니다.

첫 번째는 aten::index에 대한 긴 호출입니다. 텐서 인덱스 작업(예: []) 생성된 기본 호출로 인해 발생합니다. 그러나 GPU가 aten::index에 소비하는 실제 시간은 상대적으로 낮습니다. 그 이유는 두 코어를 시작하는 과정에서 aten::index가 둘 사이의 cudaStreamSynchronize를 차단하기 때문입니다. 즉, CPU는 두 번째 코어가 시작될 때까지 GPU가 처리를 완료할 때까지 기다립니다. 따라서 본 논문에서는 SAM을 최적화하기 위해서는 유휴 시간을 유발하는 GPU 동기화 차단을 제거하기 위해 노력해야 한다고 믿습니다.

두 번째 문제는 SAM이 트랜스포머 모델에서 흔히 볼 수 있는 행렬 곱셈(그림의 진한 녹색 부분)에 GPU 시간을 많이 소비한다는 것입니다. 행렬 곱셈에서 SAM 모델의 GPU 시간을 줄일 수 있다면 SAM의 속도를 크게 향상시킬 수 있습니다.

다음으로 SAM의 처리량(img/s)과 메모리 오버헤드(GiB)를 비교하겠습니다. 기준. 그런 다음 최적화 프로세스가 있습니다

다시 작성해야 하는 문장은 다음과 같습니다. Bfloat16 반 정밀도(GPU 동기화 및 일괄 처리 포함)

위 문제를 해결하려면 즉, 행렬 곱셈에 필요한 시간을 줄입니다. 이 기사에서는 bfloat16으로 전환합니다. bfloat16은 일반적으로 사용되는 반정밀도 유형으로, 각 매개변수 및 활성화의 정밀도를 줄여 컴퓨팅 시간과 메모리를 많이 절약할 수 있습니다. 게다가 이 기사에서 발견한 GPU 동기화 제거를 최적화할 수 있는 곳은 두 군데가 있습니다

구체적으로 위의 그림을 보면 이해가 더 쉽다는 연구 결과가 나왔습니다. SAM의 이미지 인코더에는 좌표 스케일러 역할을 하는 변수 q_coords와 k_coords가 있으며, 이 변수들은 CPU에 할당되어 처리됩니다. 그러나 이러한 변수를 사용하여 rel_pos_resize에서 인덱싱하면 인덱싱 작업이 자동으로 이러한 변수를 GPU로 이동하므로 GPU 동기화 문제가 발생합니다. 이 문제를 해결하기 위해 연구에서는 위와 같이 torch.where 함수를 사용하여 다시 작성하면 이 부분을 해결할 수 있다고 지적했습니다. 특히 소규모 배치(여기서는 1)의 경우 개별 커널 호출 간의 간격이 발생합니다. 이 현상에 대한 더 깊은 이해를 얻기 위해 배치 크기 8로 SAM 추론의 성능 분석을 시작했습니다. 요소별 커널 및 소프트맥스 연산에 소요됩니다.

이제 행렬 곱셈의 상대적 오버헤드가 훨씬 작다는 것을 알 수 있습니다.

GPU 동기화와 bfloat16 최적화를 결합하여 SAM 성능이 3배 향상되었습니다.

Torch.compile (+그래프 나누기 및 CUDA 그래프)

SAM을 연구하면서 많은 작은 작업을 발견했습니다. 연구원들은 이러한 작업을 통합하기 위해 컴파일러를 사용하는 것이 매우 유익하다고 생각하므로 PyTorch는 torch.compile에 다음과 같은 최적화를 수행했습니다.

nn.LayerNorm 또는 nn.GELU와 같은 작업 시퀀스를 단일 GPU 커널로 융합합니다. 행렬 곱셈 커널 바로 뒤에 연산을 융합하여 GPU 커널 호출 수를 줄입니다.

이러한 최적화를 통해 연구에서는 GPU 전역 메모리 왕복 횟수를 줄여 추론 속도를 높였습니다. 이제 SAM의 이미지 인코더에서 torch.compile을 사용해 볼 수 있습니다. 성능을 최대화하기 위해 이 문서에서는 몇 가지 고급 컴파일 기술을 사용합니다.

코어 추적

결과에 따르면 torch.compile은 매우 잘 수행됩니다

softmax가 시간 , 그리고 각 GEMM 변형. 다음 측정값은 배치 크기 8 이상에 대한 것입니다.

SDPA: scaled_dot_product_attention

다음으로 이 기사에서는 주의 메커니즘에 초점을 맞춰 SDPA(scaled_dot_product_attention)에 대한 실험을 수행했습니다. 일반적으로 기본 주의 메커니즘은 시간과 메모리의 시퀀스 길이에 따라 2차적으로 확장됩니다. PyTorch의 SDPA 작업은 Flash Attention, FlashAttentionV2 및 xFormer의 메모리 효율적인 주의 원칙을 기반으로 구축되어 GPU 주의 속도를 크게 높일 수 있습니다. torch.compile과 결합된 이 작업을 통해 MultiheadAttention의 변형에서 공통 패턴을 표현하고 융합할 수 있습니다. 약간의 변경 후에 모델은 이제 scaled_dot_product_attention을 사용할 수 있습니다.

Core Tracking

이제 GPU에서 많은 계산 시간을 차지하는 메모리 효율적인 Attention 커널을 볼 수 있습니다.

PyTorch의 기본 사용 scaled_dot _product_attention, 예 배치 크기가 크게 늘어납니다. 아래 그래프는 배치 크기가 32 이상인 경우의 변경 사항을 보여줍니다.

다음으로 연구에서는 Triton, NestedTensor, 배치 Predict_torch, int8 양자화, 반구조적(2:4) 희소성 및 기타 연산에 대한 실험을 수행했습니다.

예를 들어 이 기사에서는 Triton 커널에서는 배치 크기 32의 측정이 관찰되었습니다.

Nested Tensor 기술을 사용하여 배치 크기를 32 이상으로 조정

양자화 추가 후 배치 크기가 32 이상으로 변경된 측정 결과.

글의 끝은 반구조화된 희소성입니다. 이 연구는 행렬 곱셈이 여전히 직면해야 할 병목 현상임을 보여줍니다. 해결책은 희소화를 사용하여 행렬 곱셈을 근사화하는 것입니다. 희소 행렬(즉, 값을 0으로 설정)을 사용하면 가중치와 활성화 텐서를 저장하는 데 더 적은 비트를 사용할 수 있습니다. 텐서의 어떤 가중치를 0으로 설정하는 프로세스를 가지치기(pruning)라고 합니다. 더 작은 가중치를 잘라내면 정확도가 크게 떨어지지 않고 모델 크기가 잠재적으로 줄어들 수 있습니다.

완전히 구조화되지 않은 것부터 고도로 구조화된 것까지 가지치기 방법에는 여러 가지가 있습니다. 구조화되지 않은 가지치기는 이론적으로 정확도에 최소한의 영향을 미치지만, 희박한 경우 GPU는 대규모의 조밀한 행렬 곱셈을 수행할 때 매우 효율적임에도 불구하고 상당한 성능 저하를 경험할 수 있습니다. 최근 PyTorch에서 지원하는 가지치기 방법 중 하나는 균형을 찾는 것을 목표로 하는 반구조적(또는 2:4) 희소성입니다. 이 희소 저장 방법은 밀도가 높은 텐서 출력을 생성하는 동시에 원래 텐서를 50% 줄입니다. 설명은 아래 그림을 참조하세요

이 희소 저장 형식과 관련 빠른 커널을 사용하기 위해 다음으로 해야 할 일은 가중치를 잘라내는 것입니다. 이 기사에서는 2:4의 희소성에서 가지치기를 위해 가장 작은 두 개의 가중치를 선택합니다. 기본 PyTorch("스트라이드") 레이아웃에서 이 새로운 반구조적 스파스 레이아웃으로 가중치를 변경하는 것은 쉽습니다. apply_sparse(모델)를 구현하려면 32줄의 Python 코드만 필요합니다.

희소성이 2:4인 경우 vit_b 및 배치 크기가 32인 SAM 최고 성능을 관찰합니다.

마지막으로 , 이 기사의 요약은 다음과 같습니다. 이 기사는 지금까지 PyTorch에서 Segment Anything을 구현하는 가장 빠른 방법을 소개합니다. 이 기사는 공식적으로 출시된 일련의 새로운 기능을 통해 원래 SAM을 순수 PyTorch로 다시 작성하며, 손실이 없습니다. 정확성

관심 있는 독자는 원본 블로그에서 자세한 내용을 확인할 수 있습니다

위 내용은 PyTorch 팀은 원래 구현보다 8배 빠르게 "모든 항목 분할" 모델을 다시 구현했습니다.의 상세 내용입니다. 자세한 내용은 PHP 중국어 웹사이트의 기타 관련 기사를 참조하세요!