Taso : graph optimization 방법론

TASO

Define 1. Equivalent Transformation

Define 2. Data Layout

• Data layout은 multi-dimensional Array형태를 저장할 수 있는 방법을 의미한다.

• Tensor 가 가질 수 있는 Data layout은 크게 2가지가 있다.
1. Row-Major
2. Column-Major

• Data layout은 하드웨어와 operator의 성능에 따라 달라진다.
• V100에서의 예시를 살펴보자.
• 4X4 matrix operation을 제공하는 tesla V100같은 경우, 4X4로 tensor를 자르는 것이 효율적이다.

이전 Equivalent transformation

Problem 1 : Manual Definition은 scalability가 낮다.

1. Maintainability
• 많이 복잡한 알고리즘 혹은 operator가 생기면 새로운 가이드를 생성해야해서 유지보수가 크다.

2. Correctness
• 인간이 만든 Rule 이어서 오류가 발생할 수 밖에 없다.

Problem 2 : 현존하는 프레임워크에서 optimization 방법은 최적의 optimization opportunity를 잃을 수 있다.

• 제공된 operator들을 사용하여 substitution 후보군을 생성

• operator specification 리스트를 사용하여 생성된 substitution graph verification 진행 ( Theorem prover Z3)

• cost-based backtracking algorithm을 사용하여 optimized graph를 찾는다.

1. tensorflow의 manual optimization rules가 53,000줄 구현된 것을 TASO에서는 1,400 줄로 구현된 것만 갖고 진행

2. 이전 방법론들보다 인력이 덜 쓰이면서 2.8X 성능 향상

접근 방법

1. Substitution 생성
• Taso에서 Graph substitution Generator는 Computation Graph의 모든 가능한 substitution candidate를 생성한다.
• 이때 제공된 DNN operator들을 가지고 Fixed size(최대 4개)까지 생성한다.
• 이후 random input tensor를 가지고 같이 fingerPrint(Output)에 기반하여 substitution graph를 hash table에 저장한다.

2. Substitution 검증 with theorem prover
• 1번 단계에서 생성된 substitution graph들을 검증하는 과정이다.
• 이때 위 자료처럼 operator specifications들을 활용한다.
• operator들의 property 중 mathematical computation 등을 참고한다.

3. Joint optimization
• Graph substitution과 data layout transformation을 같이 묶어서 진행한다.
• Metaflow에서 제안된 cost-based backtracking search algorithm을 사용한다.
• 하드웨어마다 달라지는 data layout cost도 고려한다.

2. Graph substitution Generator

• primitive operator들이 제공되면 potential한 substitution들을 자동으로 만든다.
• 이때 특정 size만큼의 substitution들을 만드는데 여기서는 보통 4개의 operator들을 묶음으로 만든다.

• 가장 적합한 substitution을 만드는데, naive한 방법은 두 개씩 pair로 enumerate하면서 비교하는 방법이다.
• 하지만 이렇게 되면 quadratic한 time consumption이 발생하기에 더 효과적인 방법으로 아래와 같이 진행한다.
💡
superoptimization[7]의 아이디어와 substitution graph마다fingerprint라는 값을 계산한다.
→ 이후 같은 fingerprint 값을 갖는 graph들끼리 비교하여 비교 시간을 줄인다.
→ TASO에서 구현한대로 같은 fingerprint를 갖는 graph들을 서로 equivalent한 graph라는 전제가 깔려있다.

2-1. Graph Substitution Definition

• Graph Substitution 은 3가지의 구성요소로 나뉜다.
1. Source Graph
• 원본 computation graph의 subgraph
2. Target Graph
• Source graph와 Equivalent한 새로운 subgraph
3. Mapping Relation
• source graph와 target graph의 input, output tensor 정보

1. Tensor Shape



• 위 figure처럼 A,B,C에 대해 어떤 tensor shape인지들을 제한하지 않는다.

2. Configuration Parameter
• i.e. Convolution operator를 사용할 때 strides, padding, activation 등이 Configuration parameter로 사용된다.

2-1-2. Concatenation and split operators

• 위 figure처럼 같은 input을 공유하는 경우에 concatenation 과 split operator들이 사용된다.

1. split operator
• 하나의 tensor를 2개의 disjoint한 tensor로 분리한다.
• 이때 parameter로 제공된 dimension을 기준으로 분리한다.
• 하지만 split operator를 언제 사용해야 하는지 결정해야 하는데 이는 input tensor, parameter들로 알 수 없다.
💡
그런데 실험을 통해 split operator는 이전에 concatenation operator가 불러진 곳에 사용된다는 것을 알게 되었다.→ 가장 최근의 concatenation operator를 ‘undo’하기 위해서 split operator를 사용하기 때문이다.
→ 이 점을 착안하여 concatenation history를 split tree 를 사용

2. split tree



• tensor의 차원마다 생성하는 metadata
• Figure3을 통해서 row dimension들에 대한 split tree를 생성하는 예시를 볼 수 있다.
• source graph와 다르게 target graph에서는 concat과 split operator가 존재하는 것을 볼 수 있다.
• 이때 split tree(gray box)를 살펴보면 concat이 row 차원에서 일어난 것을 확인할 수 있고, 이 정보를 활용해서 split은 matmul 연산 이후에 추가해준다.
• 이를 통해서 equivalent한 target graph가 생성되는 것을 볼 수 있다.
💡
split tree를 사용하여 split point를 하기 위한 추가 parameter들이 요구되지 않는다.

2-2 Generation Algorithm

Step 1. Violent enumeration for potential graph with fingerprint

1. inner hash의 경우 graph에 대한 output 결과들을 hash 형태로 만든 것이다.
• 이때 size, shape, tensor 정보를 사용해서 hashing을 진행한다.

2. outer hash의 경우, 어떤 순서로 output들이 나열되어도 같은 fingerprint라고 인식할 수 있게 하기 위해 사용된다.
• 이때 outer hash는 symmetric hash function으로 사용되어 unordered set of hash values에게 적용된다.

step 2. Testing graph with identical fingerprint

현재까지 제안한 알고리즘은 generic하지만, 논문에서 이 섹션에서 2가지 operator는 special 케이스라고 소개했다.

1. Relu
• Often Returns 0.
→ 때문에 불필요한 graph substitution이 후보군으로 생성된다.
• 해결 방법으로 아래 operator로 교체

2. Enlarge
• Zero padding: 다른 conv 연산끼리 fusing하는데 사용
• 0이 너무 많아서 불필요한 substitution candidate들 생성
• 이 문제를 해결하기 위해 input tensor를 Enlarge operator에 넣는 경우만 고려하고, 다른 operator의 Output이 input으로 들어오는 Enlarge의 경우는 모두 없앴다.

3. Graph substitution Verifier

First-order logic으로 표현한다? 아래의 예시를 살펴보자.

First-order logic이란?

• 위 예제 문장을 명제함수로 만든다면 : P(x) = x> 3

• 변수 x에 따라 참, 거짓을 판별할 수 없기 때문에 명제가 아니다.

• 변수에 특정 값을 대입하면 참, 거짓을 판정할 수 있다.

• 명제논리와 다르게 위와 같은 명제함수가 포함된 논리를 1차 논리, predicate logic이라고 부름

1. 변수에 특정 값을 넣는다.
• 만약 x=10이면 P(x)는 True일 것이고, x=1이면 False일 것이다.

2. quantification : 범위를 제한한다.
• 만약 x≥4 라면 p(x)는 항상 True다.

quantifier란?

• 한정자라고도 불리며, 명제함수에서의 변수의 범위를 제한하는 연산자다.
1. 전체한정자 ( Universal Quantifier)
2. 존재한정자 ( Existential Quantifier)
• 마지막으로 수학 정리의 기술 예시를 살펴보자.

Z3 prover란?

• Z3은 microsoft에서 개발한 수치 해석 도구로 방정식과 같은 수학적 요소들을 해결하는데 쓰인다.

from z3 import *
>>> x = Int('x')
>>> y = Int('y')
>>> solve(x>2, y<10, x+2*y==7)
[y = 0, x = 7]

• Z3 prover는 SMT solver의 한 종류라고 보면 된다.

그럼 SMT Solver는 뭘까?

더 자세한 예시는 없을까?

# 변수 선언
x = Int('x')
y = Int('y')

# SMT 솔버 초기화
solver = Solver()

# 조건 추가
solver.add(x + y > 5, x < 3)

# 충족 가능성 확인
is_satisfiable = solver.check()

# 결과 출력
if is_satisfiable == sat:
    model = solver.model()
    print("문제에 대한 해:", model)
else:
    print("해결 불가능한 문제")

Key point는

1. Solver 모델 초기화

2. Solver 모델에 조건을 추가

3. solver가 조건을 check한다.

4. sat object가 나오면 문제에 대한 해가 존재한다.

Operator Properties 개발

• 본래 operator property들은 생성된 graph substitution들이 잘 생성되었는지를 검증하기 위해 user단에서 제공해야 한다.

• 개발 과정에서 검증이 안되고 생성이 되는 경우들이 발생하여 validation step 을 추가했다.

4. Pruning Redundant Substitution

1. Input tensor naming
• TASO는 substitution 그래프 중 다른 input name이 붙여졌지만 사실은 같은 input tensor 값을 넣는 경우, naming을 통일한다.
• 가령 왼쪽에 제안된 tensor C가 사실은 A랑 같은 tensor라면 오른쪽처럼 A로 통일된다.

1. Common Subgraph
• TASO에서는 source graph와 target graph가 같은 subgraph가 존재하면 pruning을 진행한다.
1. case 1



위 예시와 같이 src와 tgt graph 모두 같은 subgraph가 존재하는 것을 확인했다.이 subgraph를 새로운 input tensor로 표현해서 substitution graph를 간소화시킬 수 있다.
2. case 2



위 예시와 같은 경우 모든 output tensor들이 겹치는 것을 확인할 수 있다. 이 subgraph를 새로운 output tensor로 표현해서 substitution graph를 간소화시킬 수 있다.

5. Joint optimizer

• Joint optimizer는 data layout과 graph optimization을 동시에 진행

• 해당 자료를 통해 조금 더 살펴보자.

• Cost-based Backtracking search는 Joint optimizer에서 data layout의 cost와 optimizing substitution을 동시에 고려하는데 사용된다.

• 알고리즘에 대해서 설명해보겠다.

• 결과적으로 생성된 search space를 바탕으로 optimized graph를 찾는다고 생각

6. Implementation

7. Evaluation

1. Taso가 과연 자동적으로 graph를 optimize하고 graph substitution을 검증하는데 너무 느리진 않는가?

2. Taso는 Real-world DNN architecure와 미래의 operator들 또한 성능을 개선시킬 수 있는가?

3. Taso에서 제안된 Joint optimization은 각각 순차적으로 진행하는 것보다 더 효과적인가?

1. End-to-End Evaluation
우선 모든 DNN graph들을 가지고 실험했을 때, 모두 10분 내의 시간 안에 End-to-End generation이 완료되었다고 한다.
TensorFlow, TensorFlow XLA, TensorRT, TVM, MetaFlow, 그리고 TASO를 V100 GPU에서 비교했다. TensorFlow, TensorFlow XLA, TensorRT 는 cuDNN, cuBLAS 기반으로 최적화된 라이브러리를 사용하기에 TVM의 경우 2000 trial 을 통해서 각각의 DNN operator들의 최적의 configuration을 구한다.
End-To-End Evaluation 결과, Taso는 Resnet-50의 모델에 대한 최적화가 성능이 제일 높았다.TVM과 Tensorrt 성능과 비교했을 때 각각 1.1X에서 1.8X , 1.3X에서 2.8X 까지의 성능 향상이 나타났다.

2. Substitution Case Study
해당 study는 NaseNet-A와 ResNext-50을 다룬다.
• NasNet-A의 경우 unconventional한 graph 구조를 갖고 있어서 사람이 직접 optimization을 적용하는데 어려움을 느낀다. Taso에서는 이 아키텍처를 최적화하는데 있어 어떻게 적용되는지 한 번 확인해보자.
1. 2개의 average pooling과 element wise addition을 depth-wise convolution으로 transformation



• 위 그림을 이해하기 위해서 아래의 수식을 알아볼 필요가 있다.



• 상단에 있는 수식은 average pooling 수식이고 하단에 있는 수식은 depth wise convolution이다.
• 하단 수식에 있는 w(c, k_x, k_y)는 상수이기 때문에 figure b에서 1단계의 변환은 수식처럼 변환이 용이하다.
2. Depth wise convolution과 convolution으로 변환



• substitution을 통해 operator overhead를 줄이는 효과를 보았다고 한다.
• ResNext-50의 경우 large convolution를 사용하는 것보다 multiple branch를 사용하여 small convolution으로 분리해서 figure a처럼 사용한다고 한다. 하지만 이 경우, kernel launch overhead가 있다는 것이 단점이다. figure b에 소개된 것처럼 cuDNN library 중 tensorrt에서 지원하는 Grouped convolution이 존재하지만, 모든 convolution의 상태를 저장하는 cache가 크다는 단점으로 조금의 성능 저하가 있다. 하지만 이와 다르게 taso에서 자동으로 생성해주는 figure c의 graph 경우 figure b와 비교해서 2.8x의 성능을 보였다고 한다.

3. Joint optimization of Graph substitution, Data layout
• joint optimization의 성능을 확인하고자 3개의 baseline을 bert model에 사용하여 그 결과를 paper에서 공유하고 있다.
1. graph substitution만 진행
2. data layout optimization만 진행
3. 순차적으로 2개를 진행
• 아래 결과를 통해 동시에 optimization을 진행하는 joint optimization의 성능이 제일 좋은 것을 확인할 수 있다.

Limitations and Future work.

참고자료

1. sysml 블로그
SOSP '19：TASO - 计算图在特定体系结构/框架上，最优的等价变换
文章标题：TASO: optimizing deep learning computation with automatic generation of graph substitutions 文章出处：SOSP '19/Stanford University 文章地址：https://cs.stanford.edu/~matei/papers/2019/…

https://zhuanlan.zhihu.com/p/94613279

2. Taso graph optimization is for inference efficiency.
optimize pytorch computation graph for training · Issue #32 · jiazhihao/TASO
hi, thanks for the great work. But does it support optimize computation graph of pytorch for faster training? If supports, is there any benchmark?

https://github.com/jiazhihao/TASO/issues/32