多圖形處理器環境下訓練樹狀模型的高效能資料分群演算法

Abstract

圖神經網路是深度學習中處理結構化數據的重要工具，其中圖由節點與邊組成，用於表示實體及其關係。當圖神經網路呈樹狀結構時會面臨各種挑戰，例如不規則的結構與變化的深度。將其分配與並行執行於多張圖形處理器上極具挑戰性。此外，樹狀結構的相依性要求先處理父節點再處理子節點，這嚴重限制了執行的並行性。

本研究旨在提升樹狀得圖神經網路在多圖形處理器系統上的訓練速度。首先，我們引入一個成本模型來估算多圖形處理器訓練的執行時間。接著，我們證明在該成本模型下，尋找最佳的樹狀數據分配方式是一個NP完全問題。因此，我們提出一種啟發式方法來分配數據，以提高效率並維持訓練品質。此方法首先基於成本模型將數據分配至不同批次，然後在每個批次內將數據分配至多個圖形處理器。我們進一步證明，我們的分配演算法是一個四倍近似演算法，確保其在實務上具有效率。

我們實作了該演算法並進行實驗。結果顯示，我們的演算法在訓練時間上獲得顯著提升。當使用兩張圖形處理器時加速比達 1.86，四張圖形處理器達 3.43，而八張圖形處理器達 7.25。

Graph Neural Network (GNN) is an important tool in deep learning to handle structured data, where graphs with nodes and edges represent entities and their relationships. Various challenges arise when GNN is tree-shaped, with irregular connectivity patterns and varying depth. It is difficult to distribute and process the dynamic structure for parallel execution on multiple GPUs. In addition, tree data dependency demands the processing of parent nodes before their children, severely limiting execution parallelism.

This research aims to improve the training speed of tree-shaped GNN on multi-GPU systems. First, we introduce a cost model that estimates the running time of the training across multiple GPUs. Then, we demonstrate that finding an optimal way to distribute tree-structured data across GPUs is an NP-complete problem on this cost model. We then propose a practical heuristic method for distributing data that improves efficiency while maintaining training quality. The heuristic method first assigns data to batches based on our cost model and then assigns data in each batch to the devices. We also show that our device-assigning algorithm is a 4-approximation algorithm. That is, it guarantees that its cost is four times the optimal running time in each training batch, ensuring that it performs effectively in practice.

We implement the algorithm and conduct the experiments. The results show that our algorithm achieves a significant increase in training time. The speedup is up to 1.86 for two GPUs, 3.43 for four GPUs, and 7.25 for eight GPUs.