[논문 리뷰] Boosting Distributed Full-graph GNN Training with Asynchronous One-bit Communication

Meng Zhang, Qinghao Hu외 3명|arXiv (Cornell University)|2023. 03. 02.

Advanced Graph Neural Networks인용 수 4

한 줄 요약

이 논문은 전체 그래프 GNN 학습에서 GPU 간 통신 오버헤드를 크게 줄이기 위해 비동기 1비트 통신을 사용하는 분산 GNN 학습 프레임워크인 Sylvie를 제안한다. 저비트 양자화 및 오차 보정을 통한 복원을 적용함으로써 Sylvie는 기준 방법 대비 최대 28.1배의 속도 향상을 이룩하면서도 모델 정확도를 유지한다.

ABSTRACT

Training Graph Neural Networks (GNNs) on large graphs is challenging due to the conflict between the high memory demand and limited GPU memory. Recently, distributed full-graph GNN training has been widely adopted to tackle this problem. However, the substantial inter-GPU communication overhead can cause severe throughput degradation. Existing communication compression techniques mainly focus on traditional DNN training, whose bottleneck lies in synchronizing gradients and parameters. We find they do not work well in distributed GNN training as the barrier is the layer-wise communication of features during the forward pass & feature gradients during the backward pass. To this end, we propose an efficient distributed GNN training framework Sylvie, which employs one-bit quantization technique in GNNs and further pipelines the curtailed communication with computation to enormously shrink the overhead while maintaining the model quality. In detail, Sylvie provides a lightweight Low-bit Module to quantize the sent data and dequantize the received data back to full precision values in each layer. Additionally, we propose a Bounded Staleness Adaptor to control the introduced staleness to achieve further performance enhancement. We conduct theoretical convergence analysis and extensive experiments on various models & datasets to demonstrate Sylvie can considerably boost the training throughput by up to 28.1x.

연구 동기 및 목표

• 분산 전체 그래프 GNN 학습에서 높은 GPU 간 통신 오버헤드 문제를 해결하여 확장성과 처리량을 향상시키기 위해.
• 기존 압축 기법의 한계를 극복하기 위해, 이는 DNN에 최적화되어 있으나 GNN의 계층별 특성 및 기울기 전송에 효과적이지 않기 때문이다.
• 저비트 양자화 및 비동기 파이프라인 실행을 통해 통신 비용을 극적으로 줄이면서도 높은 모델 정확도를 유지하기 위해.
• 수렴성이나 성능을 희생시키지 않고 다양한 GNN 모델과 데이터셋에 일반화 가능한 시스템을 설계하기 위해.

제안 방법

• 전송을 위해 특성과 기울기를 1비트로 양자화하고, 오차 보정을 통한 복원을 사용해 수신자 측에서 전체 정밀도로 복원하는 저비트 모듈을 도입한다.
• 계산과 통신을 겹치기 위해 비동기 파이프라인(Sylvie-A)을 구현하여 지연을 숨기고 자원 활용도를 향상시킨다.
• 비동기 학습에서 오래된 특성과 기울기의 영향을 제어하기 위해, 수렴 안정성을 향상시키는 경계가 설정된 능동적 조절기(Bounded Staleness Adaptor)를 제안한다.
• GPU 간 그래프를 분할하고 통신이 중요한 노드를 식별하여 중복 데이터 전송을 최소화하는 그래프 엔진을 설계한다.
• 양자화된 데이터를 사용해 순전파 및 역전파를 수행하는 트레이너 모듈을 적용하여, 양자화에도 불구하고 모델 일관성을 유지한다.
• 전체 정밀도 계산을 유지하기 위해 양자화를 통신되는 텐서에만 선택적으로 적용한다.

실험 결과

연구 질문

• RQ1 1비트 양자화가 모델 정확도를 떨어뜨리지 않으면서도 분산 GNN 학습에서 통신 오버헤드를 효과적으로 줄일 수 있는가?
• RQ2 저비트 통신과 함께 비동기 파이프라인 실행이 GNN의 학습 처리량과 수렴에 어떤 영향을 미치는가?
• RQ3 특성과 기울기의 오래됨(staleness)이 모델 수렴에 어느 정도 영향을 미치며, 이를 제한함으로써 성능 유지가 가능한가?
• RQ4 제안된 프레임워크가 다양한 GNN 아키텍처와 실제 그래프 데이터셋에 일반화 가능한가?
• RQ5 저비트 GNN 학습에서 통신 감소, 계산 오버헤드, 메모리 점유율 간의 상호 상충 관계는 어떠한가?

주요 결과

• Reddit 데이터셋에서 Sylvie는 FP32 기준 대비 통신 오버헤드를 89.8% 감소시키고, 에포크당 학습 시간을 84.2% 단축시켰다.
• 1비트 양자화를 적용한 Sylvie는 그림 1에 나타낸 바와 같이 이중 서버 환경에서 최대 28.1배의 학습 처리량 향상을 달성했다.
• 모든 임bedding과 기울기를 1비트로 양자화하면 심각한 정확도 저하가 발생한다(예: Ogbn-products에서 70.6% 대비 Sylvie-S의 78.86%), 이는 선택적 양자화가 필수적임을 입증한다.
• 이중 서버 환경에서 저비트 모듈은 총 학습 시간의 5.9%에 불과하여 계산 오버헤드가 극히 낮음을 시사한다.
• 경계가 설정된 능동적 조절기는 비동기 업데이트가 성능을 떨어뜨릴 수 있는 상황에서 수렴성을 향상시킨다.
• Sylvie는 GraphSAGE, GCN, GAT 등 여러 모델과 Reddit, Yelp, Ogbn-products 등의 데이터셋에서 전체 정밀도 학습 대비 최소한의 정확도 손실로 높은 정확도를 유지한다.