Meta 내에서 사용되는 FX 도구에 대한 전체적인 살펴보기: 그래프 변환을 통한 PyTorch 모델 최적화

PyTorch의 그래프 모드는 더 성능이 좋습니다. 이 글에서는 PyTorch 프로그램의 그래프를 캡처하고 최적화할 수 있는 강력한 도구인 Torch.FX를 소개합니다.

1. 서론

PyTorch는 두 가지 실행 모드, 즉 즉시 실행 모드와 그래프 모드를 지원합니다.

즉시 모드에서는 모델의 연산자가 읽히는 즉시 실행됩니다. 머신 러닝 실무자에게 사용하기 쉽고 사용자 친화적이므로 기본 실행 모드로 설정됩니다.

그래프 모드에서는 연산자를 먼저 그래프로 합성한 다음 전체적으로 컴파일하여 실행합니다. 성능이 더 뛰어나 실제 생산에 널리 사용됩니다.

특히 그래프 모드는 연산자 융합을 지원합니다. 두 연산자를 병합하면 메모리 읽기와 커널 실행의 총 오버헤드를 줄이거나 지역화할 수 있습니다.

퓨전은 수평적일 수 있습니다:단일 연산(예: BatchNorm)을 여러 피연산자에 적용하여 이를 단일 배열로 병합합니다.

퓨전은 수직적일 수도 있습니다.첫 번째 커널의 출력을 필요로 하는 다른 커널과 커널을 결합합니다(예: ReLU에 이은 합성곱).

Torch.FX(FX로 약칭)는 PyTorch 패키지의 일부로 그래프 모드 실행을 지원하는 공개적으로 사용 가능한 툴킷입니다. 다음과 같은 작업이 가능합니다.

1. PyTorch 프로그램에서 그래프를 가져옵니다.

2. 개발자가 얻은 그래프에 변환을 작성할 수 있도록 허용

Meta는 이전에 FX를 사용하여 프로덕션 모델의 학습 처리량을 최적화했습니다. 이 글에서는 그래프 변환을 사용하여 PyTorch에 배포된 모델의 성능을 최적화하는 방법을 보여주기 위해 Meta에서 개발한 FX 기반 최적화를 소개합니다.

2. 배경

임베딩 테이블은 추천 시스템에서 널리 사용됩니다.이 섹션에서는 FX와 임베딩 테이블에 대한 배경 지식을 소개합니다.

2.1. FX

그림 1은 FX를 사용하여 PyTorch 프로그램을 변환하는 방법을 보여주는 간단한 예입니다.이는 세 단계로 구성되어 있습니다.

• 프로그램에서 그래프를 가져옵니다
• 그래프를 수정합니다(이 경우 RELU 대신 GELU를 사용합니다)
• 수정된 그래프에서 새로운 프로그램을 생성합니다.

FX API는 PyTorch 프로그램 그래프를 검사하고 변환하기 위한 다양한 기능을 제공합니다.

2.2. 테이블 임베딩

추천 시스템에서는희소 기능(예: 사용자 ID, 스토리 ID)은 임베딩 테이블로 표현됩니다.

임베딩 테이블 E는 HxD 행렬이고, 여기서 H는 해시 크기이고 D는 임베딩 벡터 차원입니다. E의 각 행은 부동 소수점 숫자의 벡터입니다.

피처 해싱의 기능은 희소 피처를 [S1, S2, …, Sk]와 같은 E의 인덱스 목록에 매핑하는 것입니다. 여기서 0 ≤ Si입니다.

GPU를 최대한 활용하려면 일반적으로 희소한 기능을 일괄 처리합니다.배치의 각 엔터티에는 고유한 인덱스 목록이 있습니다. 배치에 B개의 엔터티가 있는 경우, 간단히 B개의 인덱스 목록을 갖는 표현으로 이해할 수 있습니다.

좀 더 엄격한 표현은 B 인덱스 목록을 하나의 인덱스 목록으로 병합하고 인덱스 길이 목록(배치의 각 엔터티에 대해 하나씩)을 추가하는 것입니다.

예를 들어, 배치에 3개의 엔터티가 포함된 경우 해당 인덱스 목록은 다음과 같습니다.

• 엔티티 1: 인덱스 = [10, 20]
• 엔티티 2: 인덱스 = [5, 9, 77, 81]
• 엔티티 3: 인덱스 = [15, 20, 45]

그러면 전체 배치 크기의 인덱스와 길이는 다음과 같습니다.

• 지수 = [10, 20, 5, 9, 77, 81, 15, 20, 45]
• 길이 = [2, 4, 3]

전체 배치에 대한 임베딩 테이블 쿼리의 출력은 BxD 행렬입니다.

3. 3 FX 변환

PyTorch는 임베딩 테이블에 대한 액세스 속도를 높이기 위해 세 가지 FX 변환을 업데이트했으며, 이 섹션에서는 이에 대해 하나씩 소개합니다.

아래는 여러 개의 작은 입력 텐서를 하나의 큰 텐서로 결합하는 변환에 대한 3.1입니다. 3.2 여러 개의 병렬 컴퓨팅 체인을 하나의 컴퓨팅 체인으로 융합하는 것에 대한 변환에 관하여; 3.3 중복되는 커뮤니케이션과 컴퓨팅의 변형에 관한 내용입니다.

3.1 입력 희소 기능 결합

배치의 각 입력 희소 특징은 인덱스 목록과 B 길이 목록의 두 목록으로 표현될 수 있습니다. 여기서 B는 배치 크기를 나타냅니다.

PyTorch에서는 두 리스트 모두 텐서로 존재할 수 있습니다.PyTorch 모델이 GPU에서 실행되는 경우, 임베딩 테이블은 일반적으로 GPU 메모리에 저장됩니다(이 메모리는 GPU에 더 가깝고 CPU 메모리보다 읽기 및 쓰기 대역폭이 더 높습니다).

입력 희소 특징을 사용해야 하는 경우 두 텐서를 먼저 CPU에서 GPU로 복사해야 합니다. 그러나 호스트에서 장치로 메모리를 복사하려면 커널을 실행해야 하므로 실제 데이터 전송보다 시간이 더 많이 걸립니다.

모델이 많은 입력 희소 특성을 사용하는 경우 이러한 복사는 성능 병목 현상이 될 수 있습니다(예: 1000개의 입력 희소 특성에는 호스트에서 장치로 2000개의 텐서를 복사해야 함).

호스트-장치 memcpy의 수를 줄이기 위한 한 가지 최적화 방법은 장치로 전송하기 전에 여러 개의 입력 스파스 기능을 결합하는 것입니다.

예를 들어, 다음 세 가지 입력 기능이 주어졌습니다.

• Feature_A: 인덱스 = [106, 211, 7], 길이 = [2, 1]
• Feature_B: 인덱스 = [52, 498, 616, 870, 1013], 길이 = [3, 2]
• Feature_C: 인덱스 = [2011, 19, 351, 790], 길이 = [1, 3]

결합된 형태는 다음과 같습니다.

특징_A_B_C: 인덱스 = [106, 211, 7, 52, 498, 616, 870, 1013, 2011, 19, 351, 790], 길이 = [2, 1, 3, 2, 1, 3]

따라서 호스트에서 장치로 3×2=6개의 텐서를 복사하는 대신, 2개의 텐서만 복사하면 됩니다.

그림 3(b)는 두 가지 구성 요소로 구성된 이 최적화의 구현을 보여줍니다.

• CPU 측:입력 파이프라인은 모든 희소 피처의 인덱스를 하나의 텐서로 결합하고, 모든 길이를 다른 텐서로 결합하도록 수정되었습니다. 그런 다음 이 두 텐서는 GPU에 복사됩니다.
• GPU 측:FX를 사용하여 모델 그래프에 Permute_and_Split 연산자를 삽입하여 병합된 텐서에서 개별 피처 인덱스와 길이 텐서를 복구하고 이를 다운스트림의 해당 노드로 전송합니다.

3.2 임베딩 테이블 접근부터 시작하는 계산 체인의 수평 융합

실제 운영 모델에서는 GPU당 10개의 임베딩 테이블을 갖는 것이 일반적입니다. 성능상의 이유로,이러한 테이블에 대한 쿼리는 함께 그룹화되어 출력이 하나의 큰 텐서로 연결됩니다.(그림 4(a)의 빨간색 부분 참조).

단일 기능 출력을 계산하려면Split 연산자를 사용하여 큰 텐서를 N개의 작은 텐서로 분할합니다.(여기서 N은 특징의 개수) 그런 다음 각 텐서에 원하는 계산을 적용합니다.

그림 4(a)에 표시된 것처럼 각 피처 출력 O에 적용된 계산은 Tanh(LayerNorm(O))입니다. 모든 계산 결과는 큰 텐서로 연결된 다음 하위 연산자(그림 4(a)의 Op1)로 전달됩니다.

여기서 주요 런타임 비용은 GPU 커널 실행 오버헤드입니다.예를 들어, 그림 4(a)의 GPU 커널 실행 횟수는 2*N+3입니다(그림의 각 타원은 GPU 커널을 나타냄). LayerNorm과 Tanh의 GPU에서의 실행 시간은 커널 실행 시간에 비해 매우 짧기 때문에 이는 성능에 영향을 미칩니다.

또한, Split 연산자는 임베딩 벡터 출력 텐서의 추가 복사본을 생성하여 추가 GPU 메모리를 소모할 수 있습니다.

FX를 사용하여 수평 융합이라는 최적화를 구현하면 GPU 커널 실행 횟수를 크게 줄일 수 있습니다.(이 예에서 최적화 후 GPU 커널이 실행되는 횟수는 5회입니다. 그림 4(b) 참조).

명시적인 Split 대신 Add_middle_dim 연산자를 사용하여 (B, NxD) 모양의 2D 임베딩 텐서를 (B, N, D) 모양의 3D 텐서로 재구성합니다. 다음으로 단일 LayerNorm이 마지막 차원에 적용됩니다. LayerNorm의 결과에 Tanh를 적용합니다. 마지막으로 Remove_middle_dim 연산자는 Tanh 결과를 2D 텐서로 복원하는 데 사용됩니다.

Add_middle_dim과 Remove_middle_dim은 텐서를 재형성하기만 하므로추가 복사본이 생성되지 않으므로 GPU 메모리 사용량도 줄일 수 있습니다.

3.3 계산과 통신의 중복

생산을 위한 추천 모델의 훈련은 일반적으로 분산 GPU 시스템에서 수행됩니다.각 GPU의 장치 메모리 용량은 모델의 모든 임베딩 테이블을 보관하기에 충분하지 않으므로 여러 GPU에 분산해야 합니다.

학습 단계에서 GPU는 다른 GPU의 임베딩 테이블에서 기능 값을 읽거나 써야 합니다.이를 전체 대 전체 의사소통이라고 하며, 이는 상당한 성능 저하로 이어질 수 있습니다.

FX를 통해 변환을 구현함으로써 계산을 전체 대 전체 통신과 겹치게 할 수 있습니다.그림 5(a)는 임베딩 벡터 테이블 액세스(EmbeddingAllToAll) 및 기타 연산자를 사용한 모델 그래프의 예를 보여줍니다. 그림 5(b)에서 볼 수 있듯이, 어떠한 최적화도 없이 하나의 GPU 스트림에서 순차적으로 실행됩니다.

FX를 사용하여 EmbeddingAllToAll을 EmbeddingAllToAll_Request와 EmbeddingAllToAll_Wait로 분할하고, 그 사이에 독립적인 연산자를 배치합니다.

3.4 요약

이러한 변환을 통해 어떤 모델이 이점을 얻을 수 있는지 알아내기 위해 개발자는 MAIProf가 메타 데이터 센터에서 실행되는 모델에 대해 수집한 성능 데이터를 분석했습니다.우리는 이러한 변환이 즉시 모드와 비교했을 때 일련의 생산 모델에서 2~3배의 속도 향상을 달성한다는 것을 보여줍니다.

4. 결론

성능 관점에서 볼 때 PyTorch의 그래프 모드는 프로덕션 환경에서 사용되는 즉시 모드보다 선호됩니다. FX는 PyTorch 프로그램 그래프를 캡처하고 최적화하는 강력한 도구입니다. 이 문서에서는 Meta 내에서 프로덕션 추천 모델을 최적화하기 위한 세 가지 FX 변환을 소개합니다.

마지막으로, 더 많은 PyTorch 개발자가 그래프 변환을 사용하여 모델의 성능을 개선할 수 있기를 바랍니다.

​—— 끝 ——