이번 논문은 ViT를 가속화하는 기법을 소개하는 논문이다. 우선 ViT는 Vision Transformer이다. NLP에서 Self Attention을 이용한 Transformer가 비약적인 성능 향상을 가져왔고, 그 구조를 Vision 영역에 맞게 고친 것이 ViT이다. ViT에 대해서 자세하게 설명해둔 글이 있는데 조만간 업로드하려고 한다.
트랜스포머 구조의 가장 큰 문제점은 Self Attention 연산 속도가 매우 느리다는 것이다. 논문에서는 ViT의 트랜스포머 구조는 사실 NLP의 트랜스포머보다 가속화하기에 유리하다고 설명하고 있다. 그 이유는 ① ViT는 고정된 사이즈의 입력 토큰이라는 것과 ② ViT 어탠션의 90% 이상은 sparse한 패턴이기 때문이다. 이제 어떻게 ViT의 Self Attention을 가속화 했는지 자세하게 알아보자.
1. divide-and-conquer 알고리즘
ViT 어탠션의 90% 이상은 sparse한 패턴이다. 어탠션 연산은 matrix multiplication을 통해 계산되는데 이렇게 연산한 효과가 그렇게 크지 않는 부분, 즉 어탠션 연산의 효과가 미비한 부분의 weight은 0으로 치환해서 matrix 자체를 sparse하게 바꿀 수 있다. 성능을 크게 회손하지 않는 범위 내에서 matrix 자체를 sparse하게 바꿔도 무방하다는 것이다.
그러면 성능을 크게 회손하지 않는다는 것을 어떤 기준으로 판단할까? 논문에는 아래와 같이 알고리즘이 기술돼 있다.
2번 줄에 보면 어떤 임계치 값을 정해두고 그 임계치를 넘지 않을 때까지 계속해서 인덱스를 기록해두고 그 인덱스에 해당하는 부분만을 사용하도록 하게 한다.
이렇게 어탠션의 weight 값을 임계치 기반으로 sparse하게 바꾼 후에 reordering을 한다. sparse한 부분과 dense한 부분으로 나누는 것이다. 나눈 후에 어탠션 값을 plotting 하면 아래와 같은 패턴으로 보여진다. ViT의 입력 토큰이 고정적이기 때문에 NLP의 트랜스포머에 비하여 가속화하기에 유리하다고 말했는데, 그 이유를 여기에서 조금 설명할 수 있다. sparse한 부분과 dense한 부분의 비율이 어느 정도 고정적으로 정해질 수 있기 때문에 가속기를 하드웨어적으로 구현할 때 PE(Processing Elements)의 워크로드 사이즈를 어느 정도 정할 수 있다고 한다.
dense한 부분의 연산은 일반적인 행렬연산(GEMM)이므로 이는 병렬로 처리하면 빠르게 할 수 있다. sparse한 부분의 연산은 sparse matrix multiplication을 수행하면 된다. 이는 sparse한 인덱스를 따로 테이블로 가지고 있어서 그 부분에 대해서만 연산을 수행하면 되므로 훨씬 적은 양의 곱셈 연산으로 원하는 결과를 얻을 수 있다.
2. 하드웨어 가속기
소프트웨어만을 다루는 사람으로서 이해하기가 조금 힘든 영역이었다. 그래도 조금 어설프게나마 이해한 부분을 정리해보자. divide and conquer를 통해서 구분한 sparse/dense한 영역을 계산하기 위해서 따로 하드웨어 모듈을 만들었다고 한다. 아래의 그림에서 sparser engine과 denser engine이 바로 그 부분이다.
위 그림에서 PE Array라고 하는 부분을 잠깐 설명해보자. neural network accelerator에서 핵심 연산은 곱셈과 덧셈(MAC)이다. MAC을 레지스터, 메모리와 함께 어떻게 잘 패킹하는지가 neural network accelerator를 디자인하는데 중요한 부분이다. MAC, 레지스터, 메모리를 함께 패킹해둔 것이 하나의 PE라고 하고 이것을 여러 개 연결해둔 것이 PE Array이다.
다시 VitCoD 논문으로 돌아와보자. 위 그림을 보면 denser engine 부분과 sparser engine의 버퍼 구조가 조금 다르다. denser engine은 딱히 특별한 것 없이 일반적인 행렬 곱셈(GEMM)을 수행하면 되고, sparser engine은 sparse matrix multiplication을 수행하기 위해 IdxBuf 등이 추가적으로 필요하다. sparse matrix multiplication은 행렬 자체가 sparse하기 때문에 non-zero 부분만 연산을 하면 되고 그 non-zero 인덱스를 가지고 있기 위한 버퍼가 IdxBuf이다.
실제 칩의 그림을 보면 위와 같다고 한다.
3. Auto Encoder/Decoder 모듈
어탠션 값을 sparse/dense한 부분으로 나누고 이를 위한 각각의 하드웨어 가속기를 만들었으니 이제 빨라졌겠네. 라고 생각하면 오산이다. 위 그림에서 sparse/denser engine 전까지의 연산은 Q, K를 만들때까지이다. sparser/denser engine은 Q, K 값이 나오기 전까지는 사용되지 않는다는 것이다. 그러면 Q, K를 만들고 그 값을 sparser/denser engine으로 값을 이동시켜야 하는데, 그 과정에서 발생하는 overhead가 엄청나다. 이 overhead를 줄이기 위해서 Auto encoder/decoder 모듈을 사용한다. Q, K 값을 auto-encoder로 압축해서 작게 만든 후 sparser/denser engine으로 이동시키고, 이동시킨 후에 decoder를 통해서 다시 Q', K'를 만들어서 그 값으로 이 후 연산을 진행하게 된다. Auto Encoder/Decoder 모듈을 통해서 데이터 이동 시간을 줄이므로 인해서 sparser/denser engine의 사용율이 높아지게 된다.
4. 결과
ViTCoD를 적용한 결과 전체적인 연산의 속도는 아래와 같이 개선됐다.
speed-up
CPU
Edge-GPU
GPU Platform
Attention Only
x 235
x 160
x 86.0
Overall
x 33
x 5
NaN
또한 AE 모듈을 삽입한 부분을 보면 AE를 삽입했음에도 불구하고 정확도 측면에서는 거의 차이가 나지 않는다는 것을 알 수 있다.
AE를 삽입함으로서 Data Movements 비용도 절반으로 줄어들었다. 아래의 그림을 보면 Sanger에 비해서 ViTCoD(1)은 Data Movements 비용이 상당히 높지만, ViTCoD(2)는 Data Movements 비용이 훨씬 낮아진 것을 확인할 수 있다. 또한 Computation 비용도 ViTCoD가 훨씬 낮아졌다.
이번에 ViTCoD 논문을 리뷰하면서 SW를 하는 사람으로서 HW의 벽을 실감할 수 있었다. 2번 하드웨어 가속기 부분을 정말 이해하기 힘들더라. 조금씩 조금씩 하나 하나 이해하다보면 빨라지겠지..?
TVM은 딥러닝 컴파일러이다. 이번 글에서는 왜 TVM이 필요한지부터 TVM 내에서 일어나는 과정들에 대해서 소개하려고 한다.
딥러닝 모델을 학습하는 프레임워크는 너무 많고, 학습된 모델을 실행시킬 서버/디바이스도 너무 많다.
각각의 조합을 모두 만족하기 위해서 Tensorflow가 저 많은 서버/디바이스 별로 코드를 개발해야하고, Pytorch도 Pytorch 나름대로 코드를 개발하는 것은 비효율적이다.
그래서 딥러닝 컴파일러가 필요했고, 딥러닝 컴파일러 중에서 Apache에서 개발 중인 TVM이라는 컴파일러가 있다.
기본적으로 컴파일러란 사람이 이해할 수 있는 코드를 기계가 이해할 수 있는 언어(기계어)로 바꿔주는 것이다.
딥러닝 컴파일러가 중간에 들어가므로서 고려해야하는 가지 수가 크게 줄어들었다.
여기에서 TVM의 역할은 크게 두가지이다.
다양한 프레임워크로부터의 모델을 통합된 어떤 형태로 컴파일하기
통합된 모델을 타겟 하드웨어에 최적화된 형태로 실행시키기
TVM은 아래의 과정을 통해서 위 두 가지 역할을 수행한다.
Graph-Level로 최적화하는 것은 여러가지 딥러닝 프레임워크로부터 학습한 모델을 하나의 공통된 표현으로 바꾸기 위해서이다. 그 표현을 IR(Intermediate Representation)이라고 한다. 모델을 IR로 변형할 때 단순히 변형만 하지 않는다. 최적화를 한다.
| Optimizing Computational Graphs
1. Operator Fusion
Operator Fusion은 Layer Fusion과 같다고 생각하면 된다.
Layer Fusion을 수학적으로 이해하고 연산 결과가 100% 동일함을 예제 소스코드이다.
여기에서는 간단하게 Conv2D + BatchNorm을 단일 Conv2D 레이어로 Fusion하는 예제를 소개하려고 한다.
우선 이론적인 내용부터 알아보자. Conv2D + BatchNorm의 연산을 수식으로 전개하면 아래와 같다.
즉, Conv2D + BatchNorm 레이어는 단일 Conv2D 레이어로 변형 가능하다.
예제 코드를 봐보자. 아래의SampleModel레이어는 단일Conv2D로 변형 가능하다.
class Conv_Bn_Activation(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size, stride, activation, bn=True, bias=False):
super().__init__()
pad = (kernel_size - 1) // 2
self.conv = nn.ModuleList()
if bias:
self.conv.append(nn.Conv2d(in_channels, out_channels, kernel_size, stride, pad))
else:
self.conv.append(nn.Conv2d(in_channels, out_channels, kernel_size, stride, pad, bias=False))
if bn:
self.conv.append(nn.BatchNorm2d(out_channels))
# self.conv.append(MyBatchNorm(out_channels))
if activation == "mish":
self.conv.append(Mish())
elif activation == "relu":
self.conv.append(nn.ReLU(inplace=True))
elif activation == "leaky":
self.conv.append(nn.LeakyReLU(0.1, inplace=True))
elif activation == "linear":
pass
else:
print("activate error !!! {} {} {}".format(sys._getframe().f_code.co_filename,
sys._getframe().f_code.co_name, sys._getframe().f_lineno))
def forward(self, x):
for l in self.conv:
x = l(x)
return x
class SampleModel(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size, stride, activation, bn=True, bias=False):
super().__init__()
conv_1 = Conv_Bn_Activation(in_channels, 1024, kernel_size, stride, activation)
conv_2 = Conv_Bn_Activation(1024, 1024, kernel_size, stride, activation)
conv_3 = Conv_Bn_Activation(1024, 1024, kernel_size, stride, activation)
conv_4 = Conv_Bn_Activation(1024, 1024, kernel_size, stride, activation)
conv_5 = Conv_Bn_Activation(1024, out_channels, kernel_size, stride, activation)
self.convs = nn.ModuleList([conv_1, conv_2, conv_3, conv_4, conv_5])
def forward(self, x):
for conv in self.convs:
x = conv(x)
return x
위의 모델을 fuse하는 함수는 아래와 같다.
def fuse_layer_cbr(layer, inputs):
c_layer = layer[0]
b_layer = layer[1]
assert type(c_layer) == nn.Conv2d, 'first layer is not nn.Conv2d'
assert type(b_layer) == nn.BatchNorm2d, 'second layer is not nn.BatchNorm2d'
kernel_size = c_layer.kernel_size[0]
stride = c_layer.stride[0]
pad = (kernel_size - 1) // 2
conv = nn.Conv2d(
c_layer.in_channels,
c_layer.out_channels,
kernel_size,
stride,
pad,
bias=True)
with torch.no_grad():
for i in range(b_layer.num_features):
# update new bias and weight
precomputed = b_layer.weight[i] / (np.sqrt(b_layer.running_var[i] + b_layer.eps))
new_bias = b_layer.bias[i] + precomputed * (0 - b_layer.running_mean[i])
# update
conv.weight[i,:,:,:] = c_layer.weight[i,:,:,:] * precomputed
conv.bias[i] = new_bias
return conv
def fuse_model(model, inputs):
# print(f'original={model.conv.conv} length={len(model.conv.conv)}')
if [type(m) for m in model.conv.conv] == [nn.Conv2d, nn.BatchNorm2d, nn.ReLU]:
fused_conv = fuse_layer_cbr(model.conv.conv, inputs)
return fused_conv
SampleModel을 하나 만들어서 그 함수에 (1, 3, 112, 112) 사이즈의 입력을 넣었을 때와 SampleModel을 fuse한 모델의 속도 연산속도 차를 비교해보면 아래와 같다. (100번 수행함)
x = torch.randn(1, 3, 112, 112)
run x 100 times
>>> SampleModel : duration=4.4151 params=30707200
>>> FusedModel : duration=3.2266 params=30698496
이런 속도 차이가 나는 이유는 Conv2D 단일 레이어로 바꿈으로서 파라미터 수가 줄어들기 때문이 아니다. Conv2D와 BatchNorm을 실행시킬 때 GPU와 CPU의 메모리 복사 과정이 생략되기 때문이다.
2. Data Layout Transformation
모델의 연산을 그래프로 표현할 때 텐서들을 그래프에 여러가지 방법으로 표현할 수 있다.
예를 들어서 2D의 텐서가 있을 때 이 텐서의 데이터를 row-major로 표현할 지 column-major로 표현할지 정할 수 있다.
데이터의 레이아웃은 메모리에 접근하는 횟수를 결정할 수 있기 때문에 코드의 가속화에 큰 역할을 할 수 있다.
| Generating Tensor Operation
TVM의 핵심은 Pytorch/Tensorflow/ONNX 등등의 서로 다른 DL 프레임워크에서 학습된 모델을 Arm, Intel 등등의 서로 다른 하드웨어 아키텍쳐에 맞게 최적화 및 컴파일하는 것이다.
그렇게 하기 위해서 TVM에서는 모델 내에서 행해지는 연산들을 텐서로 만들어내는 과정을 거친다. 이때 메모리의 캐시를 조금 더 효율적으로 사용하기 위한 Cache-Friendly한 접근이 사용된다.
Cache-Friendly한 코드가 어떤 코드인지 간단하게 알아볼 수 있는 소스코드를 보자.
// cache-friendly approach
void matsum_v1(float** a, int n) {
float sum = 0;
for(int i = 0; i < n; i++)
for(int j = 0; j < n; j++)
sum += a[i][j];
cout << sum << endl;
}
// cache-unfriendly approach
void matsum_v2(float** a, int n) {
float sum = 0;
for(int i = 0; i < n; i++)
for(int j = 0; j < n; j++)
sum += a[j][i];
cout << sum << endl;
}
행렬의 모든 원소의 합을 구하는 함수이다. matsum_v1은 2차원에서 loop가 실행되기 때문에 2차원의 연속된 공간의 메모리가 캐시에 올라가게 된다. 따라서 캐시를 최대한 효율적으로 사용할 수 있다. 반면에 matsum_v2는 1차원의 공간을 loop가 iteration하고 있다. 1차원의 공간은 메모리에서 연속적으로 있지 않기 때문에 cache는 계속 miss하게 되어 결국 메모리로부터 계속 값을 읽어오므로서 속도가 느려진다. 위의 코드를 실행해보면 속도차이는 아래와 같다.
$ ./cache_friendly
3.60288e+16
matsum_v1: elapse time 201 ms
3.60288e+16
matsum_v2: elapse time 550 ms
matsum_v2가 두배 이상 느린 것을 확인할 수 있다. 캐시를 사용하는 방법만 바꿨을 뿐인데 속도차이는 어마어마하다.
Tensor Expression도 마찬가지이다. 같은 연산을 수행하는데 cache-friendly한 방식으로 코딩을 하면 속도 차이가 훨씬 빨라진다. 행렬 연산을 수행하는 tensor expression을 보면 그 차이를 명확하게 알 수 있다.
# run time: 1.832 sec
k = te.reduce_axis((0, K), "k")
A = te.placeholder((M, K), name="A")
B = te.placeholder((K, N), name="B")
C = te.compute((M, N), lambda x, y: te.sum(A[x, k] * B[k, y], axis=k), name="C")
evaluate_operation(s, [A, B, C], target=target, name="mmult", optimization="none", log=log)
# run time: 0.130 sec
xo, yo, xi, yi = s[C].tile(C.op.axis[0], C.op.axis[1], bn, bn)
(k,) = s[C].op.reduce_axis
ko, ki = s[C].split(k, factor=4)
s[C].reorder(xo, yo, ko, ki, xi, yi)
s[C].vectorize(yi)
evaluate_operation(s, [A, B, C], target=target, name="mmult", optimization="blocking & vectorization", log=log)
# run time: 0.030 sec
s = te.create_schedule(C.op)
xo, yo, xi, yi = s[C].tile(C.op.axis[0], C.op.axis[1], bn, bn)
(k,) = s[C].op.reduce_axis
ko, ki = s[C].split(k, factor=4)
s[C].reorder(xo, yo, ko, xi, ki, yi)
s[C].vectorize(yi)
evaluate_operation(s, [A, B, C], target=target, name="mmult", optimization="loop permutation", log=log)
똑같이 행렬 곱을 수행하는 Tensor Expression인데 수행하는 구현하는 방법에 따라서 1.83 → 0.13 → 0.03 초로 가속화시킬 수 있다.
Tensor Expression은 TVM에서 schedule primitive의 집합이다. 위 코드에서 tile, split, vectorize, reorder 등이 TVM의 schedule primitive이다. 하나의 모델을 schedule primitive로 표현할 때 여러가지 조합이 나올 수 있다. 이 중에서 TVM은 가장 효율적인 조합을 선택한다. 어떻게 할까?
| Automating Optimization
주어진 schedule primitive들을 이용해서 각각의 레이어에 대해서 어떻게 최적화할 수 있을까?
가령, Conv2D를 schedule primitive로 구현할 수 있는 방법은 여러가지이다. 그 중에서 가장 빠른 방법을 찾는 과정을 알아보자.
1. Schedule Space Specification
TVM은 어떤 high-level tensor expression이 주어졌을 때, 그것과 연산 결과는 같지만 다른 expression들을 만들어내는 generic master template이 구현돼 있다.
generic master template을 통해서 생성된 tensor expression들이 하나의 큰 search space가 되는데, 이 search space에서 optimizer는 최적의 tensor expression을 찾아낸다.
쉽게 말해서, TE-1부터 TE-4까지의 tensor expression을 search space에 넣어두면 optimizer가 가장 빠른 연산을 찾아내는 것이다.
# TE-1
C = tvm.compute((n,), lambda i: A[i] + B[i])
s = tvm.create_schedule(C.op)
----------------------------------------
for (int i = 0; i < n; ++i) {
C[i] = A[i] + B[i];
}
# TE-2
C = tvm.compute((n,), lambda i: A[i] + B[i])
s = tvm.create_schedule(C.op)
----------------------------------------
xo, xi = s[C].split(s[C].axis[0], factor=32)
for (int xo = 0; xo < ceil(n / 32); ++xo) {
for (int xi = 0; xi < 32; ++xi) {
int i = xo * 32 + xi;
if (i < n) {
C[i] = A[i] + B[i];
}
}
}
# TE-3
C = tvm.compute((n,), lambda i: A[i] + B[i])
s = tvm.create_schedule(C.op)
xo, xi = s[C].split(s[C].axis[0], factor=32)
s[C].recorder(xi, xo)
----------------------------------------
for (int xi = 0; xi < 32; ++xi) {
for (int xo = 0; xo < ceil(n / 32); ++xo) {
int i = xo * 32 + xi;
if (i < n) {
C[i] = A[i] + B[i];
}
}
}
# TE-4
C = tvm.compute((n,), lambda i: A[i] + B[i])
s = tvm.create_schedule(C.op)
xo, xi = s[C].split(s[C].axis[0], factor=32)
s[C].recorder(xi, xo)
s[C].bind(xo, tvm.thread_axis(“blockIdx.x”)
s[C].bind(xi, tvm.thread_axis(“threadIdx.x”)
----------------------------------------
int i = threadIdx.x * 32 + blockIdx.x;
if (i < n) {
C[i] = A[i] + B[i];
}
2. ML-Based Cost Model
머신러닝 모델을 이용해서 최적의 tensor expression을 찾아낸다. 더 정확히 말하면 머신러닝 모델을 이용해서 최적의 후보 몇 개를 선택한다.
search space에 있는 tensor expression은 몇 억개 정도로 많다. 그 모든 것을 타겟 하드웨어에서 실행해서 제일 빠른 expression을 찾으면 너무 느릴 수 밖에 없다.
따라서 각각의 expression을 입력으로 해서 타겟 하드웨어에서의 실행시간을 예측하는 머신러닝 모델을 통해서 최적의 후보 몇 개를 선택한다.
TVM의 ML-Based Cost 모델은 XGBoost를 이용해서 학습돼 있다.
3. Schedule Exploration
ML-Based Cost 모델을 통해서 몇 개의 후보 expression을 구한 후에 각 expression의 실행시간을 측정한다.
그런데 expression을 실행할 때 주어진 설정대로 그대로 실행시키지 않는다. 처음에는 설정을 임의로 지정해서 실행한다. 그러면서 조금씩 주어진 설정대로 변경해가면서 실제로 실행시간이 단축되는지 확인해보는 과정을 거친다.
이 과정에서 실행시간이 실제로 단축이 되면 pass이고 그렇지 못하면 reject된다.
4. Distributed Device Pool and RPC
TVM은 RPC 프로토콜을 지원한다. 로컬에서 모델을 컴파일 한 후 컴파일 한 모델을 리모트(타겟 하드웨어)에 업로드해서 실행시키고, 그 결과를 다시 로컬로 가져올 수 있다.
우선 RPC가 무엇인지 알아보자. RPC는 Remote Procedure Call의 약자이다. 즉 원격 프로시져 콜이다. 내 컴퓨터에 있는 어떤 프로그램을 원격에 있는 다른 컴퓨터에서 실행시키고 결과값을 받는 것이다.
# on local machine(alienware)
$ python rpc-example.py
# on remote target(RaspberryPi)
$ python -m tvm.exec.rpc_server --host 0.0.0.0 --port=9090
INFO:root:If you are running ROCM/Metal, fork will cause compiler internal error. Try to launch with arg ```--no-fork```
2022-12-06 18:00:36.393 INFO bind to 0.0.0.0:9090
2022-12-06 18:10:43.965 INFO connection from ('192.168.0.5', 53314)
2022-12-06 18:10:44.342 INFO load_module /tmp/tmp0e37s1ac/lib.tar
2022-12-06 18:10:44.440 INFO Finish serving ('192.168.0.5', 53314)
위의 예시는 노트북 컴퓨터를 로컬로 두고 라즈베리파이를 리모트 디바이스로 두고 서로를 RPC로 연결 시킨 후 노트북에서 모델을 컴파일 후 라즈베리에서 실행시킨 결과이다.
마지막으로 정리를 하고 마치려고 한다.
TVM은 딥러닝 컴파일러다. TVM은 우선 여러 딥러닝 프레임워크로부터 학습된 모델을 Intermediate Representation으로 변형을 한 후, 이 변형된 그래프 내에서 최적화를 진행한다. 최적화된 그래프는 다시 Tensor Expression으로 바뀌는데(schedule primitive의 집합) 변형 가능한 Tensor Expression들을 모두 하나의 search space에 모아두어서 머신러닝 모델로 제일 빠를 것 같은 Tensor Expression을 몇 개 찾고, 찾은 것들은 실제로 실행시켜보면서 최적의 속도를 보이는 Tensor Expression을 선택한다.
같은 결과를 내는 함수인데도 더 빠르게 실행시키기 위해서 코드 연산을 항상 Cache-Friendly하게 고치는 과정을 전체적인 TVM 컴파일 과정에서 보여주고 있다.
