[arXiv’ 24] Galaxy: A Resource-Efficient Collaborative Edge AI System for In-situ Transformer Inference 논문 리뷰

< 문서를 읽기 위한 배경 지식>
(1) Transformer 아키텍쳐에서 진행되는 연산 이해
(2) Collective Communication 에 대한 이해

원본 문서 링크: https://arxiv.org/abs/2405.17245

Galaxy: A Resource-Efficient Collaborative Edge AI System for In-situ Transformer Inference

 

Transformer 아키텍처와 Edge AI의 통합: 구조적 이해부터 협력적 추론까지

현재의 Language 관련 응용을 위한 딥러닝 모델은 Transformer 가 자리하고 있다는 사실은 틀림없다. 본 섹션에서는 Transformer 아키텍쳐에 대한 구조를 이해하고 해당 구조를 가지는 딥러닝 모델 추론을 어떻게 수행해왔는지에 대한 기존 패러다임을을 이해해본다.

전형적인 Transformer 아키텍쳐 구조

일반적으로 Transformer 는 Encoder 와 Decoder 로 나눌 수 있는데 BERT 와 GPT-2 처럼 Encoder 와 Decoder 둘 중 하나를 사용하는 모델만을 집중해서 Transformer 아키텍쳐를 그려보면 아래 그림과 같다.

Transformer 의 전형적인 Layer 구성

메인 구성 요소로 Multi-head Attention (MHA) 라는 블록이 있고, Multilayer Perceptron (MLP) 라는 블록이 있는 것을 볼 수 있다. 이러한 두 파란 블록 간에는 element-wise 연산인 Dropout, Residual Addition, 그리고 Layer Normalization 이 존재하는 것을 확인할 수 있고 이러한 요소가 한데 묶여서 Transformer Layer 하나의 블록을 만들게 된다. 조금 더 구체적으로, MHA 블록 안에서는 Head 마다의 Query (Q), Key (K), Value (V) 행렬을 만들어내는 Linear 연산이 있고, 만들어진 Q, K, V 를 가지고 Self-attention 로직을 수행하는 Multi Self-Attention 연산이 있고, 최종 출력물을 만들기 위해 다시 Linear 연산을 거치게 된다. MLP 블록 안에는 두 개의 Linear 연산과 Gelu 활성화 함수가 있다.

Edge AI 와 On-device AI 의 시작

대부분의 Transformer 기반의 응용 프로그램들은 클라우드 서비스에 의존한다고 보아도 과언이 아니다. 그 이유는 다음과 같다.

1. 높은 계산 자원 요구: Transformer 모델은 수억에서 수십억 개의 파라미터를 가질 수 있어 막대한 계산 자원을 필요로 한다. 이러한 큰 계산 자체 때문에 사용자는 자연스레 Cloud Service Provider 에서 제공하는 고성능 자원을 빌려 사용하는 방식에 의존하고 있다.
2. 메모리 요구 사항: 대규모 Transformer 모델일수록 높은 메모리 용량을 필요로 하기 때문이다.
3. 사전 학습된 모델 및 API 제공: 많은 Cloud Service Provider 는 사전 학습된 Transformer 기반 모델을 API 형태로 사용할 수 있게끔 제공하기 때문에 사용자가 더욱 쉽게 고성능 모델을 활용할 수 있다는 장점도 존재하기 때문이다.

그러나, 이러한 방식은 추론 프로세스 자체가 Wide-area Network (WAN) 을 타기 때문에 네트워크 지연으로 인한 Quality-of-Service (QoS) 저하도 우려가 되면서 수많은 유저들의 추론 요청을 동시 다발적으로 처리하는 것이 데이터센터에서 조차 성능 병목이 유발되며, 유저 개인 정보를 포함한 사적이고 민감한 정보에 대한 보안 우려 또한 있다.

이러한 맥락에서 스마트폰이 보급되기 시작한 2000 년대 중반쯤부터 스마트폰을 포함한 유저 근방에 존재하는 컴퓨팅 기기들을 활용하여 연산을 수행하자는 엣지 컴퓨팅 (Edge Computing) 접근법이 대두되었고, 현재에 와서는 AI 추론을 Edge 환경에서 수행한다는 의미를 담아 온디바이스 AI (On-device AI) 나 Edge AI 등의 표현으로 많은 산업과 학계가 주목하고 있다. 이러한 새로운 패러다임은 낮은 지연 시간과 데이터 프라이버시 보호 등의 이점을 제공하여 기존의 중앙 집중식 클라우드 컴퓨팅의 한계를 극복하려는 목적이 현재는 메인이다.

Edge 디바이스를 활용한 AI 추론의 한계

그런데 사실 지금 설명된 Transformer 모델을 연산하기에는 너무 많은 계산 자원을 필요하고, 이를 Edge 디바이스에서 동작시키기에는 Edge 디바이스 자체가 자원이 한정적이다. (그래서 사실 Edge 환경은 Resource-constrained 또는 Resource-hungry 한 환경이라고 표현한다.)

Motivation 실험에 사용된 모델들

 

본 연구 논문의 저자들은 Edge 디바이스의 제한된 성능 환경이 온디바이스 Transformer 추론에 얼만큼의 영향을 미치는지에 대한 간단한 실험을 진행했다. Off-the-shelf Edge 디바이스 시제품 (E.g., Jetson Nano) 과 NVIDIA GPU 플랫폼 (E.g., A100) 을 두고서 오른쪽 Table 에 나와있는 Model 들을 가지고 Sequence Length 가 30 일 때의 모델 추론 지연 시간 (Latency) 와 메모리 Footprint 를 측정했다. 측정된 결과는 아래 표와 같이, A100 과 Nano 에서의 지연 시간 성능 차이가 121 배 까지도 관찰이 되었으며 또한, FP16 기반 GPT2-L 모델 동작 시에는 1.6GB 의 메모리 Footprint 를 가지게 되어 메모리 용량이 1.5GB 밖에 되지 않는 Nano 에서 Out-of-Memory (OOM) 문제가 발생하였다.

여러 모델들에 대해 Nano 와 A100 간의 성능과 메모리 사용량 측면에서의 비교 평가

여러 개의 Edge 디바이스들을 활용한 협력적 추론 접근법

Edge 디바이스 자체가 연산을 동작시키기에는 자원이 너무 부족하니 그래서 나온 대안이 여러 개의 근접한 Edge 디바이스를 함께 묶어 쓰자는 패러다임이 발현되기 시작했다. 그러면 딥러닝 연산 워크로드가 분할되고 각 디바이스가 병렬 처리한다면 빠르게 추론이 진행될 것으로 기대한 것이다. 아래 그림과 같이, 여러 컴퓨팅 기기들 (E.g., 태블릿, 컴퓨터, 스피커 등등) 이 분산되어 있고 이러한 분산된 기기들을 함께 병렬 활용하게 된다면 Transformer 모델 가속화가 가능할 것이라는 패러다임이다.

본 논문에서 제안된 Galaxy 라는 프레임워크를 통해 스마트 홈에서의 AI 목소리 비서 역할을 수행하는 응용 시나리오 시각화

본 논문에서도 이러한 패러다임에 탑승한 연구를 전개하고 있다. 그래서 기존에는 어떻게 협력적으로 여러 컴퓨팅 기기를 사용해서 Transformer 기반 딥러닝 모델 추론 가속화를 이룰 수 있었는지에 대한 선행 기술을 검토했고 그 정리는 다음과 같다. 각 병렬화에 대한 그림 형상화는 아래와 같다.

Transformer 딥러닝 동작을 가속화 하기 위한 여러 병렬화 기법들에 대한 그림 형상화

1. Data & Pipeline Parallelism: 우선 Data Parallelism 은 입력 데이터 샘플들 자체를 나눠서 각 디바이스가 완전히 독립적으로 다른 입력 샘플들에 대해서 병렬적으로 수행하는것을 의미한다. Pipeline Parallelism 은 각 디바이스들이 딥러닝 모델 연산의 특정 단계들 (E.g., A single layer or multiple layers) 을 담당하도록 맵핑하여 파이프라이닝 방식으로 동작하는 기법이다.
2. Model Parallelism: 딥러닝 모델의 한 Layer 안에서 수행시킬 워크로드를 나눠 처리하는 동작을 의미한다. Model Parallelism 은 더욱 분류되는데, 딥러닝 모델의 Weight 들을 각 디바이스별로 나눠서 갖도록 하는 Tensor Parallelism (TP) 와 초기 입력 자체를 Sequence 차원으로 분할하여 각 디바이스가 병렬 동작 시키는 Sequence Parallelism (SP) 로 분류된다.

두 병렬화의 가장 중요한 차이점은 분할 대상이 다르다는 것이다. TP 의 경우엔 Layer 내부 연산 단위를 분할 (결국 특정 연산의 Weight 를 분할하는것이 Layer 내부 연산 단위를 분할하는 것.) 하는 반면, SP 는 특정 연산에 들어오는 입력 자체를 분할하는 방식이다. TP 의 경우엔 입력 자체를 나누지 않고 Layer 내부 연산을 분할하기 때문에 Layer 의 Weight 가 클 때, 적용하는 것이 좋기에 높은 메모리 효율성을 갖는다. 반면, SP 는 입력 자체를 분할하기 때문에 데이터 크기가 큰 상황에서 적용하는 것이 좋다. TP 의 경우엔 Megatron-LM 와 DeepSpeed 가 예제이며, SP 의 경우엔 gPipe 와 PipeDream 이 있으니 살펴보는 것이 이해에 도움이 될 것이다.

본 연구 역시 Model Parallelism 의 일환이며, 여기선 TP 와 SP 를 Transformer 아키텍쳐에 적절하게 혼용해서 사용하여 장점만 뽑아낸다면 각 디바이스간의 통신 오버헤드를 최소화하면서 빠른 병렬 처리가 가능해질 것이란 접근법으로 아이디어를 풀려고 하고 있다. (그러나 같은 해에 나온 Hepti 연구의 Weight-Stationary 기법에 비해서 제안된 프레임워크가 성능이 낮을 것으로 보여진다.)

 

저자들이 제안하는 Galaxy 프레임워크

제안된 Galaxy 프레임워크의 시스템 오버뷰

Galaxy 전체 워크플로우 설명

위의 그림은 제안된 Galaxy 프레임워크의 전반적인 동작 흐름을 나타낸 시스템 오버뷰이다. 여러 단계로 나뉘어져 있는데 구체적으로 Preprocessing Phase, Parallelism Planning Phase, Execution Phase 로 구분되어 있다.

1. Preprocessing Phase: 오프라인으로 Galaxy 에서 각 이기종 (Heterogeneous) 디바이스의 성능을 프로파일링하는 단계이다. 그림을 보게 되면, 여러 디바이스들에게 다양한 모델 들을 입력으로 주었을 때 각각의 Commputation Latency, Communication Latency, Memory Budget 등의 평가 메트릭이 어떻게 나왔는 지 정보를 수집한다. (프로파일링을 어떻게 하는 지에 대해 구체적인 설명은 없다.)
2. Parallelism Planning Phase: 각 디바이스마다 수행할 워크로드가 정해져서 각각 병렬처리를 수행하게 되는데 Galaxy 에서는 Hybrid Model Parallelism (HMP) 를 주장하고 있다. 그리고 이 단계에서 오프라인 프로파일링 단계에서 수집된 정보를 바탕으로 각 디바이스가 얼만큼의 워크로드를 정해서 수행하면 최적일지를 고려해서 워크로드 크기가 결정된다.
3. Execution Phase: 실제로 병렬 Transformer 추론 연산을 수행하는 단계이다. 여러 디바이스를 사용해서 병렬처리를 하게 된다면 각 디바이스끼리 불가피하게 데이터를 교환해야 다음 연산을 진행하게 되는 경우가 있는데 이러한 데이터 교환을 위한 통신 오버헤드를 End-to-end 성능 관점에서 낮추기 위해, Communication-level Optimization 을 추가로 제공한다는 점에서 다른 병렬 처리 연구와 차이가 있다.

Galaxy 에서 제시한 Hybrid Model Parallelism (HMP)

Transformer 아키텍쳐를 여러 개의 디바이스가 병렬로 수행한다고 했을 때 Galaxy 프레임워크의 Hybrid Model Parallelism (HMP) 가 어떻게 진행시키는 지를 논문에서는 예제와 함께 설명하고 있다. 특별히, 두 개의 디바이스가 존재할 때를 예를 들어 아래 그림과 같이 병렬 처리를 진행한다고 설명한다.

Galaxy 에서 제안한다고 하는 Hybrid Model Parallelism (HMP) 예제이며 구체적으로 TP 와 SP 가 혼재되어 있는 상태이다.

위 HMP 그림에서 초록색 블록이 Tensor Parallelism (TP) 방법론을 적용한 것이고, 파란색 블록이 Sequence Parallelsim (SP) 방법론을 적용한 그림이다.

• 첫 번째 초록색 블록 (Multi-Head Attention, MHA): Self-Attention 부분에서 애초에 Head 단위로 독립적인 연산할 수 있으니, Head 단위로 데이터가 각 디바이스에 나뉘어지도록 하기 위해, $W_0^V, W_0^K, W_0^Q$이 세 개의 Weight 행렬의 Column 차원 (Column-partitioned)을 분할한다. 그렇게 한다면, 각 디바이스가 일정 Head 만큼에 대해서 동기화를 위한 데이터 교환이 일어나지 않으니 독립적으로 연산이 가능해진다. (아래 그림은 Hepti 의 병렬화 그림인데, Parallel Multi-head Self-Attention 부분에서 A, B, C 행렬을 Column 차원으로 분할한 그림이고 Head 단위로 나누게 되면 각각 다른 행렬로 나뉘게 되는 모습을 볼 수 있다.)

• 두 번째 초록색 블록 (MLP): MLP 부분을 보면, 두 개의 GeMM 연산을 진행하는 것을 볼 수 있다. 가장 처음 GeMM 연산에서는 마찬가지로 $W_0^D, W_1^D$을 Column 차원 분할을 한다. 그러고 난 후, 두 번 째 GeMM 연산에서는 Column 단위로 잘려진 데이터가 입력으로 오게 되고 다음 $W_0^D, W_1^D$와 연산을 진행해야 하니, 이 Weight 행렬은 Row 차원에 대해서 분할 (Row-partitioned) 하여 연산을 수행한다. 이렇게 하면 첫 번째 GeMM 의 출력물을 두 번째 GeMM 의 입력으로 특별한 데이터 교환 없이 곧 바로 진행할 수 있다는 장점이 있다.
• MHA 블록과 MLP 블록 사이의 파란색 부분: Galaxy 그림에서 가운데 파란색 부분을 보면, Dropout, Residual Addtition, Layer Normalization 연산이 있다. 이들은 간단한 Element-wise 연산이다. 저자들은 이 부분에 대해서는 입력 데이터를 Row 방향으로 분할하는 Sequence Parallelism 을 적용했다.

각 블록에 대해서는 병렬화를 어떻게 진행하는지 알아보았고 이제 각 블록의 병렬 처리를 위해 블 록간의 데이터 교환이 언제 필요하고 어떻게 일어나는지 알아보자.

우선 Galaxy 에서는 SP 와 TP 를 사용한 블록 이후에 항상 데이터 동기화를 위한 통신을 수반한다.

• TP 가 적용된 초록색 블록이 끝난 직후: ReduceScatter 연산을 진행한다. TP 가 끝나고 각 디바이스가 점유한 데이터들을 모아주어야 하고 이들을 SP 를 위해서 다시 Row-partitioned 되도록 분할해서 나뉘어져야 하는 두 단계가 필요하다.
• SP 블록이 끝난 이후: AllGather 연산을 진행한다. SP 가 끝나고 다시 TP 를 수행하기 위해서 각 디바이스의 데이터를 모아서 각 디바이스가 갖도록 하기 위해 AllGather 연산을 수행한다.

Galaxy 에서 제시한 워크로드 파티셔닝 방법론

위 설명처럼, TP 와 SP 가 끝난 직후 데이터 교환을 위한 각 디바이스 간의 통신이 필요하다. 그런데 데이터 교환에 관여하는 디바이스가 여러 개가 있을 때, 데이터 교환 딜레이는 가장 늦게까지 데이터 교환을 수행하는 디바이스의 딜레이에 좌우된다 (bound by the completion time of the slowest device).

따라서, 각 디바이스의 데이터 교환 딜레이가 최대한 유사해지도록 데이터 파티셔닝을 하는 것이 중요하다. 그러나, 모든 디바이스마다 성능이 제각각인 (Heterogeneous devices) 점이 상당히 어려운 점이다. 또한 메모리 마저도 디바이스 마다 다르니 특정 디바이스가 성능이 좋다고 무조건 상당한 데이터를 분배하는 것이 메모리 오버플로우 (이로 인해 Program OOM 이슈 발생) 를 만들 수 있다. 그래서 각 디바이스 마다의 가용 메모리 공간 제약도 고려해야 최적의 파티셔닝을 결정할 수 있다.

저자들은 각 Transformer 의 블록들에 대해서 각 디바이스에 할당되는 워크로드를 다음과 같이 모델링했다.

• MHA 블록: Head 방향으로 잘려지는 나뉘어질 때의 워크로드
• MLP 블록: Row 방향으로 잘려지는 모델의 Weight 행렬에 대한 워크로드
• MHA 와 MLP 블록 사이 (Connective) 블록: Sequence 방향으로 입력 데이터가 잘려서 나뉘어질 때의 워크로드

그리고 각 블록에서의 각 디바이스 (D 개의 디바이스) 가 나뉘어지는 파티셔닝과 실행 딜레이를 아래와 같이 표기했다.

• MHA 블록에 대한 디바이스 워크로드 파티셔닝 $A = \{{a_0, a_1, ..., a_{D-1}}\}$
• MLP 블록에 대한 디바이스 워크로드 파티셔닝 $B = \{{b_0, b_1, ..., b_{D-1}}\}$
• MHA 와 MLP 블록 사이의 영역에 대한 디바이스 워크로드 파티셔닝 $C = \{{c_0, c_1, ..., c_{D-1}}\}$
• MHA 블록에 대한 딜레이 $L(MHA, A) = \max_{d \in \{ 0, 1, ..., D-1}\}Latency(MHA, a_d)$
• MLP 블록에 대한 딜레이 $L(MLP, B) = \max_{d \in \{ 0, 1, ..., D-1}\}Latency(MLP, b_d)$
• Connective 블록에 대한 딜레이 $L(CON, C) = \max_{d \in \{ 0, 1, ..., D-1}\}Latency(CON, c_d)$

이렇게 모델링 된 실행 딜레이를 통해 Galaxy 프레임워크의 총 병렬 연산 딜레이를 최적으로 줄이기 위해서 아래의 optimization objective 식을 풀어서 각 디바이스가 나뉘어져야 할 워크로드 파티셔닝 값들을 찾는다. (위 식에서 구체적인 Memory Footprint 조건식과 구체적인 알고리즘은 논문을 참고.)

$$ L(MHA, B) = \min_{\{A, B, C\}}(L(MHA, A) + L(MLP, B) + L(CON, C)) \\ s.t. \space Memory \space foorprint \space of \space each \space device < Memory \space Budget $$

위와 같은 식을 통해 Galaxy 프레임워크는 각 디바이스가 각 Transformer 의 블록들을 처리하는 시간과 처리하는데 필요한 가용 메모리를 고려해서 모든 디바이스가 유사한 시점에 연산이 종료되도록 하는 파티셔닝을 지원한다.

Galaxy 에서 제시한 Edge 디바이스 간의 데이터 통신 오버헤드 최소화 기법

Galaxy 에서는 Edge 환경을 가정하고 있기 때문에 각 디바이스가 사용하는 네트워크 대역폭의 제한도 신경쓰고 있다. 따라서, 각 디바이스가 네트워크를 사용해서 데이터 교환하는 딜레이에 대해서도 최적화를 한 단계 진행했다. 논문에서는 타일(Tile) 기반 통신 기법이라고 명명하고 있다.

구체적으로 Galaxy 의 HMP 아키텍쳐에서 TP 블록의 시작이 AllGather 통신을 통해서 입력 데이터에 준비 되고, GeMM 연산을 수행 그리고 TP 블록의 끝이 GeMM 연산을 마지막으로 수행하고 ReduceScatter 통신을 통해서 다음 연산의 입력 데이터를 준비한다. 이러한 지점에서 저자들은 Computation 과 Communication 간의 Overlap 을 통한 Data 통신 지연을 Hiding 시키고자 했다. 구체적인 최적화 방법은 아래와 같다.

AllGather 통신과 GeMM 연산 간의 중첩

위 왼쪽 그림은 HMP 에서 MLP 블록의 GeMM 을 수행되는 모습을 잘라낸 이미지이다. 여기서 각 디바이스는 다른 디바이스들로부터 Row-partitioned 로 잘려진 입력 데이터를 모아주는 작업이 필요한데, 사실 각자 디바이스가 이미 점유하고 있는 데이터를 가지고도 먼저 GeMM 연산을 해서 전체 결과의 일부 (Partial) 를 미리 계산할 수 있다. 모든 디바이스가 가능하다. 이러한 계산 동안 인접한 (각자 디바이스를 인덱스 $i$ 와 전체 디바이스 개수를 $D$ 로 표현할 때, 인접한 디바이스는 $(i+1)\%D \space or \space (i-1)\%D$ 의 인덱스를 갖는 디바이스) 디바이스로부터 다음 계산에 필요한 일부 데이터를 전달받는다. 이러면 데이터 교환 딜레이가 실질적으로는 각 디바이스의 연산 과정 속에 숨어지는 (Hiding) 현상이 되어 사실 상 전체 지연 시간에 반영이 덜 되어 성능이 올라가는 최적화 전략이다.

GeMM 연산과 ReduceScatter 연산 간의 중첩

설명 자체가 위 AllGather 와 GeMM 연산 간의 중첩이랑 본질이 같다.

 

메인 성능 평가 실험

실험 환경 셋업

• 테스트를 위한 타겟 디바이스와 모델 리스트
• 비교군: Local Inference (Serialized Execution with one device), Megatron-LM (M-LM) 와 같은 TP 방법론이 적용된 Parallel Inference , SP 방법론이 적용된 Parallel Inference

모델 추론을 위한 평균 지연 시간 비교

• Implication 1: Galaxy 는 HMP 아키텍처와 타일 기반 통신 최적화 덕분에 기존 방법보다 우수한 성능을 보여주는데, M-LM 대비 최대 1.46배, SP 대비 최대 1.11 배 성능 향상 관측.
• Implication 2: 모델 크기가 커질수록 통신-계산 비율이 감소해 Galaxy 의 통신 최적화 효과가 줄어들지만, 참여 장치 수가 증가하면 비율이 높아져 최적화의 이점이 확대.
• Implication 3: SP는 동기화 통신이 적어 속도 향상이 제한적이지만, Sequence 차원으로분할하기 때문에 각 장치가 전체 모델 가중치를 보유해야 하고 이에따라 메모리 사용량이 높아져 OOM(메모리 부족) 문제가 발생함
• Implication 4: 다양한 네트워크 대역폭에서 Galaxy는 항상 우수한 성능을 보였으며, 1.04배에서 1.45배까지 추론 지연 시간 감소를 달성.

여러 이기종 환경에 따른 성능 변화

계산 능력과 메모리 용량이 다양한 여러 이기종 엣지 환경에서 Galaxy 와 기존 방법(M-LM, SP)을 비교한 실험에 대한 해석은 아래와 같다.

• Implication 1: Galaxy 는 항상 뛰어난 성능을 보여주며, 1.3배에서 2.5배의 추론 지연 시간 감소를 달성.
• Implication 2: Galaxy 는 Edge 디바이스의 메모리 제약을 반영한 작업 분배가 가능해 OOM 오류 방지할 수 있음.
• Implication 3: M-LM 과 SP 는 병렬 처리 설계 시 메모리 제약을 간과하여 Edge 디바이스의 비효율성이 생기고 OOM 오류 방지가 불가능함.