Sequence Parallelism: Long Sequence Training fromSystem Perspective

AI 모델이 거대화되면서, GPU parallel 학습에 대한 연구가 많이 진행되고 있다.

그중에서, sequence parallel 방법에 대한 논문을 리뷰하고자 한다.

 

1. Background

Transformer 기반 모델을 사용해본 사람은 모두 알것이다.

긴 문장, 문서가 있으면, max_vocab을 정해서 짜른 후 input으로 사용한다.

왜 일부를 잘라서 사용할까? 당연하게도 메모리 문제이다.

data batch, layer, parameter를 쪼개서 multi-GPU로 올려 연산하고 있는 상황에서 ( data parallel, pipe parallel, tensor parallel)

Input(sequence)도 쪼개서 학습시킬수 없을까?

 

2. Method (Ring self-Attention , RSA)

논문에서 제안하는 sequence parallel (RSA)은 그림과 같은 N개의 토큰으로 이루어진 sequence 를 쪼개서 각 GPU에 넣는다.

sequence parallel 은 Transformer 기반의 모델에서만 사용할 수 있는 방법이다.

 

Transformer 모델에서 우리에게 필요한건? Q,K,V 벡터이다.

 

1. sequence를 n개의 덩어리(chunk)로 쪼개보자.

각 GPU에는 (Q^1..Q^n), (K^1..K^n), (V^1..V^n) 씩의 벡터들이 들어있다.

여기서 Q^n벡터는 n번째 GPU에 있는 Q벡터의 일부를 의미한다.

 

2. n번째 GPU의 attention score를 구해보자.

Attention^n = (Q^n, K )

Q^n 벡터에 대한 attention score는 모든 GPU에 있는 K벡터와 내적해야한다.

Q^n벡터에 대해 K^1..K^n벡터를 내적한 결과가 Ring communication을 통해 전달된다.

 

 

3. n번째 GPU 의 output을 구해보자.

Output^n = (Attention^n, V )

2에서 구한 Attention 벡터에 대해서 모든 V벡터와의 내적이 필요하다.

2와 마찬가지로 V1..Vn의 백터와 각 GPU에서 내적을하고, 그 결과가 Ring communication을 통해 전달되어 concat된다.