SwapAdvisor: Pushing Deep Learning Beyond the GPU Memory Limit via Smart Swapping

 

SwapAdvisor: Pushing Deep Learning Beyond the GPU Memory Limit via Smart Swapping

ACM Reference Format: Chien-Chin Huang, Gu Jin, and Jinyang Li. 2020. SwapAdvisor: Pushing Deep Learning Beyond the GPU Memory Limit via Smart Swapping. In Proceedings of the Twenty-Fifth International Conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS ’20), March 16–20, 2020, Lausanne, Switzerland. ACM, New York, NY, USA, 15 pages. https://doi.org/10.1145/ 3373376.3378530 

 

Abstract

 더 깊고 넓은 신경망이 더 나은 정확도를 달성할 수 있다고 알려져 있습니다. 그러나 제한된 GPU 메모리로 인해 모델 크기를 늘리는 추세를 지속하기는 어렵습니다. 한 가지 유망한 솔루션은 GPU와 CPU 메모리 간의 스와핑을 지원하는 것입니다. 그러나 기존의 스와핑 작업은 특정 모델만 다루며 만족스러운 성능을 내지 못하고 있습니다.

 딥러닝 연산은 일반적으로 스와핑을 개선하기 위해 분석할 수 있는 데이터 흐름 그래프로 표현됩니다. 주어진 데이터 플로우 그래프를 기반으로 연산자 스케줄링, 메모리 할당 및 스왑 결정의 3차원에 따라 공동 최적화를 수행하는 SwapAdvisor를 제안합니다. SwapAdvisor는 맞춤 설계된 유전자(genetic) 알고리즘을 사용하여 방대한 검색 공간을 탐색합니다. 다양한 대형 모델을 사용한 평가에 따르면 SwapAdvisor는 GPU 메모리 제한의 최대 12배까지 모델을 훈련하는 동시에 infinite GPU 메모리로 가상 기준선 처리량의 53-99%를 달성할 수 있습니다.

 

1 Introduction

 딥 러닝 커뮤니티는 지난 몇 년 동안 더 복잡한 작업에서 더 높은 정확도를 달성하기 위해 더 큰 DNN(심층 신경망) 모델을 사용했습니다[16, 47, 55, 58, 60]. 경험적 증거에 따르면 80년대 이후 최첨단 신경망의 매개변수 수는 하드웨어 개선과 대규모 데이터 세트의 가용성으로 인해 약 2.4년마다 두 배로 증가했습니다[10]. 그러나 오늘날 탐색할 수 있는 DNN 모델의 크기는 제한된 GPU 메모리(예: NVIDIA의 V100 GPU의 경우 16GB)로 인해 제한됩니다.

 기존 작업에서는 정밀도가 낮은 부동 소수점[12, 24]을 사용하거나 양자화 및 희소화를 통해 모델 매개변수를 압축하여 메모리 소비를 줄이려고 했습니다[3, 9, 13–15, 20]. 그러나 이러한 기술은 모델 정확도에 영향을 줄 수 있으며 과도한 하이퍼 매개변수 조정이 필요합니다. 일부 다른 솔루션은 중간 데이터를 버리고 필요할 때 나중에 다시 계산합니다[2, 11, 29]. 그러나 모델 매개변수를 쉽게 다시 계산할 수 없기 때문에 대형 모델을 지원할 수 없습니다.

 

 정확도에 영향을 주지 않고 GPU 메모리 제한을 해결하는 유망한 접근 방식은 DNN 계산 중에 GPU와 CPU 메모리 간에 텐서 데이터를 교환하는 것입니다[26, 30, 40, 49]. 몇 가지 기술 동향으로 인해 스와핑이 매력적입니다: 1) CPU 메모리가 GPU 메모리보다 훨씬 크고 저렴하며, 2) 최신 GPU 하드웨어가 컴퓨팅과 통신과 효과적으로 겹칠 수 있음, 3) GPU와 CPU 간의 통신 대역폭이 현재 충분히 양호하고 PCIe 5.0의 도래[37]와 NVLink의 광범위한 적용으로[26, 36] 향후 크게 성장할 것입니다.

 DNN 계산을 위한 스와핑은 DNN 계산 구조가 일반적으로 데이터 흐름 그래프의 형태로 실행 전에 알려져 있다는 점에서 기존의 스와핑(CPU 메모리와 디스크 간)과 다릅니다. 이러한 지식은 계산과 통신을 최대한 겹치게 하여 스와핑 성능을 최적화할 수 있는 엄청난 기회를 제공합니다. 안타깝게도 기존 작업은 이 정보(예: TensorFlow의 스왑 확장[57])를 활용하지 않거나 수동 휴리스틱[26, 30, 49]에 기반한 기초적인 방식으로만 사용합니다. 예를 들어, TFLMS[26] 및 vDNN[40]은 그래프의 토폴로지 정렬 순서에 따라 활성화 텐서만 교환합니다. SuperNeurons[49]는 컨볼루션 연산을 위한 데이터만 교환합니다. 결과적으로 이러한 작업은 제한된 유형의 DNN만 지원할 뿐만 아니라 스와핑의 전체 성능 잠재력을 달성하지도 못합니다.

 

 본 논문에서는 제한된 GPU 메모리로 다양한 종류의 대형 모델 훈련 및 추론을 지원할 수 있는 일반적인 스와핑 시스템인 SwapAdvisor를 제안합니다. 주어진 DNN 계산에 대해 SwapAdvisor는 계산 및 통신 중첩을 최대화하기 위해 실행 전에 정확히 무엇을 언제 스왑할지 계획합니다.

 데이터 플로우 그래프만으로는 이러한 정밀한 플래닝에 충분하지 않으며, 이는 오퍼레이터가 실행하도록 스케줄하는 방법과 메모리 할당이 수행되는 방법에 따라 달라집니다. 더 중요한 것은 메모리 할당 및 오퍼레이터 스케줄링도 달성 가능한 최상의 스와핑 성능에 결정적인 영향을 미친다는 것입니다. SwapAdvisor는 맞춤형 유전자 알고리즘을 사용하여 모든 메모리 할당 및 운영자 일정의 공간을 검색하여 최종 스와핑 계획이 운영자 일정, 메모리 할당 및 교체에 대한 공동 최적화의 결과를 나타내도록 합니다.

 

 우리의 평가에 따르면 SwapAdvisor는 GPU 메모리 제한의 최대 12배까지 매우 큰 모델의 교육을 지원할 수 있습니다. SwapAdvisor가 달성한 처리량은 GPU 메모리가 무한대일 때 가상 기준선의 53%-99%입니다. SwapAdvisor는 모델 추론에도 사용할 수 있습니다. 추론은 훈련보다 메모리 사용량이 적습니다. 그러나 비용을 절약하기 위해 여러 모델이 단일 GPU를 사용하도록 할 수 있습니다. 이 설정에서 SwapAdvisor를 사용하여 모델 간에 전체 메모리를 공유하는 시간과 반대로 각 모델이 메모리의 일부만 사용하도록 제한할 수 있습니다. 우리의 실험에 따르면 SwapAdvisor는 다른 시간 공유 접근 방식에 비해 최대 4배 더 낮은 지연 시간을 달성합니다.

 우리가 아는 한, SwapAdvisor는 다양한 대형 DNN 모델을 지원하는 최초의 일반 스와핑 시스템입니다. 유망하지만 SwapAdvisor에는 몇 가지 제한 사항이 있습니다(Sec 8). 특히, 제어 흐름 프리미티브가 없는 정적 데이터 흐름 그래프가 필요하며 단일 GPU만 교체할 계획입니다. 이러한 제한을 제거하려면 추가 연구가 필요합니다.

 

2 Background

 DNN 훈련 및 추론은 일반적으로 PCIe 및 NVLink와 같은 고성능 버스를 통해 호스트 CPU에 연결된 GPU에서 수행됩니다. 예를들어 NVIDIA V100의 16GB GPU는 대역폭은 높지만 용량이 제한된 다른 메모리 기술을 사용합니다. 대조적으로 CPU에는 수백 기가바이트의 메모리가 장착되는 것이 일반적입니다. 따라서 GPU 메모리 제약 조건에서 불가능했을 훈련 및 추론을 지원하기 위해 GPU와 CPU 메모리 간에 데이터를 교환하는 것이 매력적입니다.

 최신 DNN은 일반적으로 정교한 비선형 토폴로지에서 함께 구성되는 최대 수백 개의 레이어로 구성되도록 진화했습니다. TensorFlow/MXNet과 같은 프로그래밍 프레임워크는 DNN 계산을 텐서 연산자의 데이터 흐름 그래프로 표현합니다. DNN의 메모리 소비는 3가지 범주로 나뉩니다:

1. 모델 매개변수(Model parameters). DNN 훈련에서 매개변수는 반복(iteration)이 끝날 때 업데이트되고 다음 반복에서 사용됩니다. 매개변수 텐서는 DNN 모델의 "depth"(the number of layers)와 "width"(the size of a layer)에 비례합니다. 대형 모델의 경우 이것이 메모리 사용을 지배합니다.
2. 중간 결과(Intermediate results). 여기에는 활성화(activation), 기울기(gradient) 및 오류 텐서(error tensors)가 포함되며, 그 중 후자 2개는 학습(training)에만 존재하지만 추론(inference)에는 존재하지 않습니다.
3. 스크래치 공간(Scratch space). 특정 연산자 구현(예: 컨볼루션)에는 최대 1GB의 스크래치 공간이 필요합니다. 스크래치 공간은 전체 메모리 사용량의 작은 부분입니다.

 

 기존 작업은 다양한 범주의 메모리 사용 패턴을 기반으로 수동 휴리스틱을 사용합니다. 예를 들어, 이전 작업은 매개변수1을 교환하지 않고 CPU에 대한 액티베이션만 교환합니다[26, 40]. 매개변수 스와핑이 없으면 이전 작업에서 매개변수가 GPU 메모리에 맞지 않는 DNN을 지원할 수 없습니다. 게다가, 수동 휴리스틱을 기반으로 하는 설계는 최신 DNN 데이터 흐름 그래프가 사람이 분석하기에는 너무 복잡하기 때문에 성능 향상 기회를 놓치고 있습니다.

 이 논문에서 우리는 메모리 압박에서 모든 텐서를 스왑 인/아웃할 수 있는 일반적인 스와핑 메커니즘을 제안합니다. 더 중요한 것은 수동 휴리스틱에서 벗어나 임의로 복잡한 데이터 흐름 그래프가 주어지면 최상의 스와핑 계획을 위해 자동으로 최적화하는 것입니다. 우리는 단일 GPU에 대한 스와핑 논의에 중점을 두지만 데이터 병렬 처리를 사용하여 다른 GPU에서 모델을 복제하는 다중 GPU 교육 설정에서 설계를 사용할 수 있습니다.

 

3 Challenges and Our Approach

좋은 스와핑 계획은 최대한 통신과 계산을 오버랩해야 합니다. 중복 가능성은 (일시적으로) 사용하지 않는 텐서를 교체하여 연산자 실행에 필요한 메모리 부족 텐서 이전에 메모리 부족 텐서를 교체할 수 있는 공간을 만들 때 발생합니다. 데이터 흐름 그래프의 도움으로 무엇을 언제 교환할지 신중하게 계획하여 이러한 중복을 최대화하는 것을 목표로 합니다.

데이터 흐름 그래프 구조만을 기반으로 휴리스틱한 좋은 스와핑 계획을 찾으려는 이전 작업 [26, 40, 49]. 그러나 이것으로 충분하지 않습니다. 특히 스왑 계획에 영향을 미치는 두 가지 중요한 요소가 있다고 주장합니다:

• 메모리 할당(Memory allocation). DNN 계산은 몇 KB에서 수백 MB에 이르는 광범위한 텐서 크기를 사용합니다. 속도를 개선하고 내부 파편화를 줄이기 위해 MXNet과 같은 프레임워크는 다양한 크기 클래스의 고정 크기 텐서 개체를 미리 할당하는 메모리 풀을 사용합니다. 결과적으로 스와핑은 GPU 메모리가 가득 찼을 때뿐만 아니라 특정 크기 클래스에 여유 개체가 없을 때 발생합니다. 따라서 할당을 위해 메모리 풀을 구성하는 방법은 스와핑 성능에 중대한 영향을 미칠 수 있습니다.
• 오퍼레이터 스케줄링(Operator scheduling). 최신 DNN은 레이어가 더 이상 체인을 형성하지 않지만 분기(branch), 조인(join) 및 풀린(unrolled) 루프를 포함하므로 복잡한 데이터 흐름 그래프를 갖는 경향이 있습니다. 결과적으로 오퍼레이터를 실행하기 위한 다양한 잠재적 스케줄들이 있습니다. 실행 순서는 메모리 사용량에 큰 영향을 미치므로 스와핑 성능에 영향을 줄 수 있습니다.

 

 Example. 메모리 할당과 스케줄링이 스와핑에 어떤 영향을 미치는지 보여주기 위해 예제를 사용합니다. 이 예는 그림 1(a)에 나와 있는 것처럼 토이 신경망의 데이터 흐름 그래프의 일부를 기반으로 합니다. 단순화를 위해 데이터 흐름 그래프는 순방향 전파만 표시하고 역방향 부분은 생략합니다. 이 분기 구조는 현대 CNN에서 일반적입니다[47, 60].

 그림 1(a)에서 파란색 둥근 사각형은 연산자를 나타내고 작은 노란색 사각형은 텐서를 나타냅니다. 텐서는 Ax(활성화 텐서) 또는 Wx(매개변수 텐서)로 레이블이 지정됩니다. 2MB인 A2와 A5를 제외하고 모든 텐서가 1MB라고 가정합니다(A2,A5는 두 경로를 결합하는 데 사용됨). 따라서 메모리 소비는 12MB입니다. GPU의 메모리 용량이 10MB이고 오퍼레이터를 실행하거나 GPU와 CPU 간에 1MB 데이터를 전송하는 데 1단위시간이 걸린다고 가정합니다.

 매개변수 텐서는 처음에 CPU 메모리에 있으며 사용하기 전에 GPU 메모리로 스왑해야 합니다. 우리는 매개변수 텐서를 CPU 메모리에 복사하지 않고 교체할 수 있습니다. 이는 순방향 전달 동안 변경되지 않기 때문입니다. 반면, 활성화 텐서는 (오퍼레이터가 생성하기 때문에) 스왑할 필요가 없지만 역방향 패스에서 필요하기 때문에 스왑 아웃 시 CPU 메모리에 복사해야 합니다.

 

 그림 1(a)에 대해 메모리를 할당하고 실행을 예약하는 방법에는 여러 가지가 있습니다. 두 가지 예제 스케줄을 보여줍니다. left-first는 왼쪽 분기의 연산자를 먼저 실행하고 right-first는 오른쪽 분기를 먼저 실행합니다. 우리는 2개의 메모리 할당 예시를 보여줍니다. 대단위(coarse-grained)는 각각 2MB의 메모리 개체 5개를 할당하고, 세분화(fine-grained)는 각각 1MB의 메모리 개체 8개와 2MB의 개체 1개를 할당합니다. 함께 4가지 일정/할당 조합이 있으며 그림 1(b)~(e)에서 각 조합에서 최상의 스와핑 계획을 보여줍니다. GPU-CPU 통신은 GPU 실행과 이중(duplex) 및 동시(concurrent)이므로 각 테이블의 상위 3개 행은 각각 GPU 계산, 스왑 인(CPU에서 GPU로) 및 스왑 아웃(GPU에서 CPU로)에 대한 작업 타임라인을 제공합니다. 마지막 테이블 행은 현재 GPU 메모리에 상주하는 텐서 객체를 보여줍니다. 

 

 메모리 할당이 스왑 계획에 영향을 미치는 이유를 알아보기 위해 그림 1(c)와 (d)를 대조해 보겠습니다. (c)와 (d)는 모두 동일한 오른쪽 우선 스케줄링을 갖는다. 그러나 (c)의 총 실행 시간은 (d)보다 1 단위 시간이 더 깁니다. 특히, 그림 1(c)에서 GPU는 시간 슬롯 t5에서 유휴 상태에 있고 연산자 Conv1은 매개변수W1이 교체될 때까지 대기합니다. 2MB의 메모리 개체 5개가 t4에서 가득 찼기 때문에 W1을 스왑 인 할 수 없습니다. 앞의 5개 GPU 개체 중 어떤것도 스왑 아웃 할 수 없습니다. A3, W4 및 A4는 현재 실행 중인 연산자 Conv4에 필요한 입력/출력 텐서이고 A0은 다음 연산자 Conv1의 입력으로 필요합니다. A2를 교체하는 중이지만 크기가 더 커서 통신에 2 단위 시간이 걸립니다. 그림 1의 (d)는 8개의 1MB 개체와 1개의 2MB 개체가 있는 세분화된 메모리 풀을 사용합니다. 결과적으로 메모리 풀에 공간이 아직 남아 있기 때문에 연산자 Conv1이 필요로 하는 W1을 t4에서 1 단위 시간 먼저 교환할 수 있습니다. 

 

 스케줄링이 스왑 계획에 영향을 미치는 이유를 알아보기 위해 그림 1(d)와 (e)를 대조합니다. (d)와 (e)는 모두 동일한 세분화된 메모리 풀을 사용합니다. 그러나 그림 1(e)는 왼쪽 우선 일정으로 인해 (d)보다 한 단위 시간이 더 걸립니다. 그림 1(e)에서 operator Concat이 2MB 출력 A5를 위한 공간을 만들기 위해 2MB 텐서 A2가 스와핑을 완료할 때까지 기다리기 때문에 GPU는 시간 슬롯 t6 동안 유휴 상태입니다. A2는 시간 t4에서 실행되는 연산자 Conv3의 입력이므로 더 빨리 A2를 교체할 수 없습니다. 대조적으로, 그림 1(d)의 rightfirst 일정은 t3에서 더 일찍 Conv3를 실행할 수 있으므로 A2가 더 일찍 교체될 수 있습니다.

 

 Our approach. 메모리 할당과 오퍼레이터 스케줄링은 스와핑 성능에 결정적인 영향을 미치므로 주어진 데이터 흐름 그래프와 해당 메모리 할당 방식 및 오퍼레이터 스케줄(Sec 4)을 추정하여 스와핑 계획(즉, 어떤 텐서가 언제 스왑 인/아웃할지)을 도출합니다. 특히, 스왑 계획은 미래에 가장 오랫동안 필요하지 않은 텐서를 교체하고 이전에 교체된 텐서를 가능한 한 빨리 프리페치하여 계산 및 통신 중첩을 최적화합니다.

 

 가능한 메모리 할당 공간과 오퍼레이터 스케줄을 검색하여 최상의 스와핑 성능을 가진 조합을 찾습니다. 검색을 제한하고 안내하기 위해 수동 휴리스틱을 사용하는 대신, 메모리 할당과 운영자 스케줄링의 적절한 조합을 검색하기 위해 Genetic Algorithm(GA)[5, 8, 18]을 채택합니다. GA는 NP-하드 조합 문제[28, 39] 및 병렬 시스템의 스케줄링[43]에 사용되었습니다. 우리는 GA가 빠르기 때문에 다른 검색 휴리스틱(예: 시뮬레이션된 어닐링) 중에서 GA를 선택했습니다. GA는 멀티코어 CPU에서 병렬화되고 효율적으로 계산될 수 있습니다.

 

 방대한 검색 공간을 효과적으로 탐색하려면 메모리 할당/스케줄링의 모든 조합에서 스왑 계획의 전체 성능(즉, 종단 간 실행 시간)을 빠르게 평가할 수 있어야 합니다. 실제 프레임워크에서 실제 실행을 수행하기에는 너무 느립니다. 따라서 데이터 흐름 엔진 시뮬레이터에서 스왑 계획을 실행하여 성능을 추정합니다. 시뮬레이터는 주어진 스케줄링, 메모리 할당 및 스왑 계획에서 데이터 흐름 그래프의 실행 시간을 추정할 수 있도록 GPU-CPU 통신 대역폭뿐만 아니라 각 연산자에 대해 측정된 계산 시간을 사용합니다. CPU 코어에서 시뮬레이터의 실행 시간은 실제 실행보다 몇 배나 빠르므로 모델 검색 시간을 1시간 미만으로 줄입니다. 시뮬레이터를 통해 SwapAdvisor의 GA는 종단간 실행 시간을 직접 최적화할 수 있습니다.

 

4 SwapAdvisor Design

 Overview. 그림 2는 기존 DNN 프레임워크(우리 구현에서는 MXNet)와 통합된 SwapAdvisor의 아키텍처를 보여줍니다. 데이터 흐름 그래프가 주어지면 SwapAdvisor는 그래프를 기반으로 합법적인 일정과 메모리 할당을 초기 값으로 선택하고 이를 스왑 플래너에 전달하여 어떤 텐서를 언제 스왑 인/아웃할지 결정합니다. 스왑 플래너의 결과는 추가 스왑 인 및 스왑 아웃 연산자와 추가 제어 흐름 에지를 포함하는 증가된 데이터 흐름 그래프입니다. 최종 실행 순서가 주어진 일정과 플래너의 스왑 타이밍을 준수하도록 하기 위해 추가 에지가 있습니다.

 최적화를 위해 증가된 그래프가 SwapAdvisor의 데이터 흐름 시뮬레이터로 전달되어 전체 실행 시간을 추정합니다. SwapAdvisor의 GA 기반 검색은 많은 메모리 할당/스케줄 조합의 성능을 측정하고 스왑 플래너에 대한 새로운 할당/스케줄 후보를 제안합니다. 스왑 계획이 충분히 최적화되면 실제 실행을 위해 최종 증강 데이터 플로우 그래프가 프레임워크에 제공됩니다.

 이 섹션에서는 스왑 플래너의 입력(Sec 4.1)을 설명하고 플래너가 특정 일정과 메모리 할당(Sec 4.2)에서 성능을 최대화하는 방법을 설명합니다. 이후 섹션(Sec 5)에서는 GA 기반 최적화에 대해 설명합니다.

 

4.1 Operator schedule and memory allocation

데이터 흐름 그래프 외에도 스왑 플래너는 오퍼레이터 스케줄 및 메모리 할당을 입력으로 사용합니다. 

 

 Operator schedule. 비순환 데이터 흐름 그래프 G가 주어지면 오퍼레이터 스케줄은 G에 있는 노드의 토폴로지 정렬 순서입니다. 단일 GPU를 사용할 때 프레임워크는 스케줄에 따라 GPU에 연산자를 보내 GPU를 계속 사용하도록 할 수 있습니다. 실제로 MXNet과 같은 프레임워크는 일반적으로 연산자를 예약하기 위해 토폴로지 정렬을 수행합니다.

 NVIDIA의 최신 GPU는 여러 "스트림"을 지원합니다. SwapAdvisor는 3개의 스트림을 사용합니다. 하나는 GPU 실행을 수행하기 위한 것, 하나는 텐서를 CPU로 교체하기 위한 것, 다른 하나는 CPU에서 텐서를 교체하기 위한 것입니다. GPU-CPU 통신은 이중적이므로 이러한 방식으로 사용할 경우 세 가지 스트림이 모두 동시에 진행될 수 있습니다. 반대로 계산에 여러 스트림을 사용하는 경우 병렬 실행을 위한 GPU 계산 리소스가 충분하지 않으면 해당 스트림을 동시에 실행할 수 없습니다. 우리는 테스트한 모든 DNN 모델에 대해 계산을 위해 둘 이상의 스트림을 사용할 때 성능상의 이점을 관찰하지 못했습니다. 이 관찰은 다른 사람들도 공유합니다[40].

 

Memory allocation. 메모리 풀을 구성하고 주어진 데이터 플로우 그래프에 대한 메모리 할당을 지정해야 합니다. 메모리 풀은 다양한 크기 클래스로 구성되며 각 클래스에는 고정 크기 텐서 개체의 특정 수가 할당됩니다. 각 연산자에 필요한 모든 입력/출력 텐서의 크기를 포함하는 데이터 플로우 그래프 G가 주어지면 메모리 할당 체계는 다음 두 가지를 지정하여 정의할 수 있습니다. 1) G의 각 텐서 크기에서 메모리 풀을 지원하는 일부 크기 클래스로의 매핑. 2) 지원되는 크기 클래스 세트와 각 클래스에 할당된 텐서 객체의 수. 예를 들어, 그림 1(b)(c)의 대략적인 할당 방식은 5개의 개체가 있는 하나의 크기 클래스(2MB)만 가지며 각 1MB 또는 2MB 텐서를 2MB 크기 클래스에 매핑합니다. 그림 1(d)(e)의 세분화된 구성표에는 각각 8개 및 1개의 개체가 있는 두 가지 크기 클래스(1MB 및 2MB)가 있으며 각 1MB 텐서는 1MB 크기 클래스에 매핑하고 각 2MB 텐서는 2MB 크기 클래스에 매핑합니다.

 

4.2 Swap planning

 스왑 플래너에는 데이터 플로우 그래프와 유효한 오퍼레이터 스케줄 및 메모리 할당 체계가 제공됩니다. 그 역할은 주어진 스케줄/할당 조합에서 최고의 성능을 가진 스왑 계획을 찾는 것입니다. 특히 스왑 플래너는 1) 메모리 부족 상태에서 스왑 아웃할 메모리 상주 텐서, 2) 스왑 인 또는 스왑 아웃을 수행할 때를 결정합니다.

 

 Which tensors to swap out? 높은 수준에서 스왑 플래너는 Belady의 전략을 사용하여 스왑을 위해 미래에 가장 오랫동안 필요하지 않을 텐서를 선택합니다. 플래너에게 스케줄이 주어졌기 때문에 미래를 내다보는 것이 가능합니다. Belady의 전략은 캐시 교체에 최적이며 플래너가 다음에 사용하기 전에 텐서를 다시 교체할 수 있는 충분한 시간을 주기 때문에 우리의 맥락에서도 잘 작동합니다. 구체적으로, 플래너는 스케줄의 순서에 따라 각 연산자를 스캔하고 일련의 연산자를 실행한 결과 메모리에 상주하게 되는 입력/출력 텐서 개체 집합을 추적합니다. 크기가 s인 텐서를 추가할 때 메모리 부족이 발생하면(즉, 크기 클래스 s에 여유 객체가 없음) 플래너는 스왑할 s와 동일한 크기 클래스에서 텐서를 선택합니다. 후보가 여러 개인 경우 계획자는 가장 먼 미래에 사용할 후보를 선택합니다.

 

 우리 환경에서 Belady의 전략을 사용할 때 주의할 점이 있습니다. 연산자 op_i가 마지막에 사용하는 텐서 T_i가 현재 연산자 op_j에 대한 입력 텐서인 텐서 T_j를 위한 공간을 만들기 위해 선택되었다고 가정합니다. 따라서 Ti를 교체할 수 있는 가장 빠른 시간은 op_i가 종료될 때입니다. 오퍼레이터 op_i와 op_j가 일정 시간에 너무 가깝다면 오퍼레이터 op_j가 필요로 하기 전에 T_j에서 스왑할 시간이 거의 없습니다. T_i가 교체될 때까지 메모리 공간을 사용할 수 없기 때문입니다. 대안으로 스왑할 후보 텐서를 선택할 때 플래너는 최소 임계값 전에 가장 최근에 사용된 텐서를 선택합니다.

 

 DNN 트레이닝은 반복적이지만 스왑 플래너에는 단일 반복에 대해서만 데이터 흐름 그래프가 제공됩니다. 각 반복이 끝날 때 매개변수 텐서를 제외한 모든 텐서는 폐기될 수 있습니다. 그러나 많은 반복에 걸쳐 동일한 스왑 계획을 사용할 수 있도록 하려면 반복이 끝날 때 GPU 메모리의 매개변수 텐서 세트가 처음과 동일한지 확인해야 합니다. 이를 달성하기 위해 이중 통과를 수행합니다. 즉, 스케줄을 스캔하여 두 번 스와핑을 계획합니다. 첫 번째 단계에서는 GPU 메모리에 매개변수 텐서가 없다고 가정하고 처음 사용하기 전에 스와핑해야 합니다. 첫 번째 패스가 끝나면 매개변수 텐서의 하위 집합이 메모리에 상주하게 되며 이를 초기 상주 매개변수라고 합니다. 그런 다음 초기 상주 매개변수가 스케줄 시작 시 메모리에 있다고 가정하여 두 번째 단계를 수행합니다. 두 번째 패스에서는 첫 번째 패스에서 발생하지 않은 추가 메모리 부족이 있는 경우 초기 레지던트 집합에서 매개변수 텐서를 제거하여 부족을 해결합니다. 최종 스왑 계획의 초기 GPU 상주 매개변수에는 두 번째 단계에서 제거된 매개변수가 포함되지 않습니다.

 

When to swap in and out? 이전에 논의한 선택 전략은 스케줄에 따라 각 연산자를 실행하기 위해 필요한 경우 스왑 아웃 및 스왑 인할 텐서 쌍을 결정했습니다. 연산과 통신 중첩을 극대화하기 위해 스케줄에서 해당 연산자의 실행을 차단하지 않도록 가능한 한 빨리 한 쌍의 스왑 아웃 및 스왑 인을 완료하려고 합니다. 그러나 스왑 인/아웃 시간이 안전한지 확인해야 합니다.

 

플래너가 3-node 데이터 플로우 그래프의 예(그림 3(a))와 스케줄 op1, op2, op3(그림 3(b))을 사용하여 스왑 타이밍을 제어하는 ​​방법을 보여줍니다. 단순화를 위해 예제의 모든 텐서는 크기가 1단위이고 총 GPU 메모리 크기는 4단위라고 가정합니다. op1을 실행하려면 매개변수 텐서 W1을 스왑해야 하므로 플래너는 스왑 인 전용 GPU 스트림에서 실행될 W1의 스와핑을 위한 새 데이터 플로우 노드를 추가합니다. 마찬가지로 W2용 스왑 인 노드가 추가됩니다. op1과 op2의 입력/출력 텐서를 보유하기에 충분한 메모리가 있습니다. 그러나 op3를 실행하려면 W3에서 스왑하고 A3를 위한 공간이 필요합니다. 플래너는 W1을 선택하여 W3을 위한 공간을 만들고(W1 → W3이라고 함), W2를 선택하여 A3을 위한 공간을 만듭니다(W2 → A3이라고 함). 먼저 W1 → W3의 경우를 살펴보겠습니다. 플래너는 두 개의 데이터 흐름 노드 W1(swap-out) 및 W3(swap-in)을 추가합니다. W3(스왑인)에서 op3으로의 제어 흐름 에지가 추가되어 W3이 GPU 메모리에 있어야만 연산자 실행이 시작됩니다. W1(swap-out)에서 W3(swap-in)으로 에지가 추가되어 해당 swap-out이 완료되어 메모리를 사용할 수 있게 된 경우에만 swap-in이 시작됩니다. 또한 op1이 사용을 마칠 때까지 메모리에서 W1을 제거할 수 없으므로 op1에서 W1(swap-out)까지의 에지가 포함됩니다. W2 →A3의 경우는 오퍼레이터에 의해 생성되기 때문에 플래너가 A3에 대한 스왑인 노드를 추가할 필요가 없다는 점을 제외하고 유사합니다. 결과로 생성된 증강 데이터 흐름 그래프는 실행을 위해 프레임워크의 데이터 흐름 엔진으로 전달될 수 있습니다.

 

5 Optimization via Genetic Algorithm

5.1 Algorithm overview

 유전 알고리즘(GA)은 교차, 돌연변이 및 선택과 같은 자연에서 영감을 받은 메커니즘을 통해 개인의 전체 인구를 진화시키고 개선하는 것을 목표로 합니다[5, 8, 18]. SwapAdvisor에서 개인의 염색체는 오퍼레이터 스케줄과 메모리 할당의 두 부분으로 구성됩니다. 1세대 개인은 무작위로 생성되고 인구의 크기는 하이퍼 매개변수 N_p에 의해 결정됩니다. 

 새로운 세대의 개체를 생성하기 위해 현재 세대의 염색체에 교차 및 돌연변이를 수행합니다. 교차는 한 쌍의 부모 염색체를 취하고 부모의 특징을 결합하여 새로운 개체를 생성하여 자녀(children)가 부모로부터 "좋은(good)" 특성(확률적으로)을 상속받을 수 있도록 합니다. SwapAdvisor에서 각 크로스오버는 2개의 새 스케줄과 2개의 메모리 할당을 생성하므로 4개의 자식이 생성됩니다. 그런 다음 GA가 로컬 미니멈(local minimum)을 탈출하고 조기 수렴(premature convergence)을 피하기 위해 필수적인 자식(children)에 대한 돌연변이를 수행합니다[5, 8, 18]. 결과적으로 변형된 자식은 스왑 플래너에 제공되어 스와핑 노드가 있는 증강 데이터 플로우 그래프를 생성합니다. 우리는 맞춤형 데이터 플로우 시뮬레이터를 사용하여 증강 그래프를 실행하고 개인의 품질을 측정하는 데 사용되는 실행 시간을 얻습니다. 마지막으로 GA는 현재 개체수 중 N_p 개체를 선택하여 다음 세대까지 생존합니다. 

 

Selection methodology. 생존할 개체을 선택하는 방법은 GA에서 중요합니다. 생존할 최고의 개인만 선택하면 개체군이 다양성을 잃고 조기에 수렴할 수 있습니다.

 SwapAdvisor의 선택은 개인의 품질을 고려하여 생존 확률을 결정합니다. 개인의 실행 시간이 t라고 가정하고 정규화된 실행 시간을 t_norm = (T_Best - t)/T_Best로 정의합니다. 여기서 T_Best는 지금까지 본 모든 개인 중 최고의 시간입니다. 개인의 생존 확률은 softmax 함수에 의해 결정됩니다.

방정식에서 S는 선택 전의 모집단 크기(보통 N_p보다 큼)입니다. 우리의 실험은 유명한 토너먼트 선택에 비해 더 안정적인 결과에 도달한다는 것을 보여주기 때문에 softmax 기반 선택을 사용합니다[5, 8, 18].

5.2 Creating new schedules

Encoding. 스케줄은 데이터 플로우 그래프(G)의 토폴로지 순서이므로 스케줄을 목록 SCH로 인코딩하는 것은 당연합니다. 여기서 SCH의 각 요소는 G의 노드 ID입니다. 

 

Crossover. [43]에서 아이디어를 빌려서 두 개의 자식 스케줄을 만듭니다. 두 개의 상위 스케줄인 SCH1과 SCH2를 교차함으로써. 그림 4와 같은 예를 통해 설명합니다. 먼저 교차 지점 CR이 무작위로 선택됩니다. 예에서 CR = 3입니다. 자식 SCH_C1을 생성하기 위해 크로스오버는 SCH2의 슬라이스(SCH2[1 . . . .CR] = [2, 3, 4])를 SCH_C1의 첫 번째 부분으로 사용합니다. SCH_C1에 없는 노드는 SCH1의 순서에 따라 채워집니다. 예에서 노드 1과 5는 SCH_C1 = [2, 3, 4]에 없으므로 SCH1의 나머지 슬롯을 순서대로 채웁니다. SCH_C2는 동일한 접근 방식을 통해 생성될 수 있지만 그림 4의 하단에 표시된 것처럼 SCH1과 SCH2의 부분이 다릅니다. 알고리즘은 SCH_C1과 SCH_C2가 G의 토폴로지 순서임을 보장합니다[43]. 

Mutation. 스케줄을 변경하는 간단한 방법은 결과가 토폴로지 순서로 유지되는한 목록에서 노드의 위치를 무작위로 변경하는 것입니다[43]. 그러나 우리는 하나의 돌연변이에서 여러 노드를 변경하면 GA가 더 잘 작동한다는 것을 경험적으로 관찰했습니다(예: RNN의 경우 2배 이상의 성능 향상).

 

 SwapAdvisor의 돌연변이 알고리즘은 데이터 흐름 스케줄러를 모방합니다. 실행할 준비가 된 모든 노드를 포함하는 준비 세트를 유지합니다(모든 선행 노드가 실행됨). 돌연변이 알고리즘의 핵심 기능은 다음 두 가지 조건을 기반으로 준비된 집합에서 노드를 선택하는 것입니다. 첫째, 확률 P로 돌연변이가 발생합니다. 이러한 경우 알고리즘은 준비된 집합에서 노드를 무작위로 선택합니다. 그렇지 않으면 알고리즘은 원래 스케줄에서 가장 먼저 실행되는 준비 집합에서 노드를 선택합니다(변경되지 않음). 선택된 노드는 "실행"된 것으로 간주됩니다. 모든 노드가 "실행"되면 알고리즘이 종료됩니다. 돌연변이 알고리즘은 대부분 입력 일정을 따르지만 일부 노드는 다르게 스케줄되는 새 스케줄을 생성합니다. 

 

5.3 Creating new memory allocation

Encoding. 메모리 할당은 각 크기를 크기 클래스에 매핑하는 방법과 각 크기 클래스에 할당할 개체 수를 제어합니다. 해시 맵(hash map)을 사용하여 텐서 크기를 크기 클래스에 매핑하는 것은 당연하지만 그렇게 하면 서로 다른 텐서 크기 간의 상대적 크기 정보가 손실되어 효율적인 크로스오버를 수행하기가 더 어려워집니다. 텐서 크기-클래스 매핑을 나타내기 위해 두 개의 목록 CLS 및 CNT를 사용합니다.

TS를 데이터 흐름 그래프에서 관찰된 고유한 '텐서 크기의 정렬된 목록'이라고 합시다. CLS는 TS와 길이가 같은 목록이고, CLS(CLS[i])의 i번째 항목은 크기가 TS[i]인 텐서의 크기 클래스를 나타내는 양의 정수입니다. 따라서 크기 클래스의 수는 Max(CLS)입니다. CNT는 각 size-class에 할당된 텐서 객체의 수를 나타내는 목록입니다. 따라서 CNT의 길이는 Max(CLS)입니다. 예를 들어, 그림 1의 데이터 흐름 그래프에는 텐서 크기가 두 개뿐이므로 TS = [1MB, 2MB]입니다. 5개의 2MB 개체가 있는 대략적인 할당은 CLS = [0, 0]에 해당하며, 이는 1MB 및 2MB 크기가 모두 ID가 1이고 개체 크기가 2MB인 동일한 크기 클래스에 매핑됨을 나타냅니다. CNT = [5]는 id = 0...Max(CLS)에서 각 크기 클래스에 할당된 개체 수를 포함합니다.

 

잠재적인 CLS 목록의 수는 O(N^N )이고 여기서 N은 고유한 텐서 크기의 수입니다[52]. 이러한 거대한 검색 공간은 GA의 성능을 심각하게 손상시킬 수 있습니다. 우리는 CLS가 단조 증가 시퀀스여야 하고 각 CLS[i]가 CLS[i - 1]보다 크거나 같다는 제한을 부과하여 검색 공간을 제거합니다. 직관은 연속 크기가 동일하거나 인접한 크기 클래스에 매핑될 때 할당이 더 효율적이라는 것입니다. 이 제한은 검색 공간을 O(N^N )에서 O(2^N )으로 줄입니다.

Crossover. 먼저 그림 5에 표시된 예를 사용하여 두 CLS를 교차하는 방법을 설명합니다. 교차 이전에 두 개의 CLS(CLS1 및 CLS2)와 네 가지 다른 텐서 크기(1MB, 2MB, 3MB 및 4MB)가 있습니다. CLS1 = [1, 1, 1, 1]은 4개의 텐서 크기가 모두 동일한 크기 클래스인 4MB에 속한다는 것을 의미하고 CNT1 = [4]는 4개의 4MB 텐서 객체가 할당되었음을 나타냅니다. CLS2 = [1, 2, 2, 2]는 1MB와 4MB의 두 가지 크기 등급이 있음을 의미합니다. CNT2는 [8, 1]과 같이 8개의 1MB 텐서 객체와 1개의 4MB 텐서 객체가 있습니다.

 

교차점은 교차점 CR을 무작위로 선택하여 상위 목록 CLS1 및 CLS2를 분할합니다. 첫 번째 자식 크기 클래스 매핑 CLS_C1은 CLS2[1. . .CR] 및 CLS1[CR+1 . . . N]을 연결하여 만듭니다. 두 번째 자식 크기 클래스 매핑 CLS_C2는 CLS1[1. . .CR] 및 CLS2[CR+1 . . . N]을 연결하여 만듭니다. 그림 5(a)는 CLS에 대한 크로스오버를 보여줍니다. 그림에서 CR은 2이고 결과는 CLS_C1([1, 2, 1, 1])과 CLS_C2([1, 1, 2, 2])입니다. CLS_C1은 단조 증가하지 않습니다. 따라서 새로운 크기 클래스 매핑이 여전히 유효한지 확인하기 위해 수정합니다. 우리의 수정은 결과 시퀀스가 유효하도록 최소한의 양만큼 문제가 있는 CLS의 요소를 증가시킵니다. 그림 5(a)에서는 CLS_C1의 3번째와 4번째 요소를 1만큼 증가시켜 복구된 CLS_C1이 [1, 2, 2, 2]가 되도록 합니다. 

 

길이가 해당 CLS의 내용에 따라 달라지므로 동일한 크로스오버 체계를 CNT에 직접 사용할 수 없습니다. 결과적으로 우리는 CNT를 CLS(및 TS)와 동일한 길이를 갖는 확장된 CNT_EXT로 확장합니다. CNT는 각 크기 클래스에 할당된 텐서 개체 수를 캡처하는 반면 CNT_EXT는 각 텐서 크기에 사용할 수 있는 텐서 개체 수를 나타냅니다. 예를 들어 그림 5에서 CLS2는 [1, 2, 2, 2]로 1MB가 1MB 크기 등급에 속하고 2MB, 3MB, 4MB가 4MB 크기 등급에 속한다는 의미입니다. CNT2는 [8, 1]이고 CNT_EXT2는 [8, 1, 1, 1]로 볼 수 있습니다. CNT_EXT1은 [4, 4, 4, 4]입니다. 우리는 확장된 CNT를 교차하기 위해 동일한 기술을 적용합니다. 그림 5(b)는 child1과 child2에 대한 확장된 CNT 결과를 보여줍니다. 예를 들어 CNT_EXT_C1은 [8, 1, 4, 4]이고 마지막 세 요소는 CLS_C1에 따라 동일한 크기 클래스에 속합니다. 동일한 크기 클래스의 모든 요소를 평균화하여 해당 크기 클래스의 개수를 얻습니다. 따라서 결과 CNT_C1은 [8, 3]입니다. 

 

CLS와 유사하게 새 CNT를 수정해야 할 수도 있습니다. 예를 들어, CNT_C1은 메모리 사용량이 12MB 메모리 용량을 초과하는 20MB(8 * 1 + 3 * 4)이므로 유효하지 않습니다. 각 요소를 요소의 해당 크기 등급에 반비례하여 감소시켜 CNT를 수리합니다. 예를 들어 원본 CNT_C1은 [8, 3]이고 복구된 CNT_C1은 [4, 2]입니다. 첫 번째 크기 클래스는 1MB이고 두 번째 크기 클래스는 4MB입니다. 

 

 

Mutation. 스케줄링 돌연변이와 유사하게, 우리는 CLS(및 CNT)에서 둘 이상의 요소를 변경합니다. 요소는 P의 확률로 돌연변이됩니다. 

CLS에서 i번째 요소를 변경하려면 1만큼 늘리거나 1만큼 줄일 수 있습니다. CLS[i]가 CLS[i -1]과 같으면 CLS[i]를 1만큼 증가시킬 수 있습니다. CLS[i]가 CLS[i − 1] + 1과 같으면 CLS[i]를 1만큼 줄입니다. 단조 증가 기능을 유지하려면 i 번째 요소 이후의 모든 요소도 증가 또는 감소해야 합니다. 

CNT 의 요소를 변경하기 위해 원래 값을 평균으로 하는 가우스 분포를 사용합니다. 가우스 분포를 사용하면 변형된 값이 대부분의 경우 원래 값에 가깝지만 작은 기회로 큰 변동이 있을 수 있습니다. 돌연변이된 CNT 및 CLS는 메모리 제한을 초과할 수 있습니다. 크로스오버와 동일한 방법을 사용하여 수리합니다.