Ian Chou's Blog

GraphRAG 效能解放：用 NetworkX + RAPIDS cuGraph 實現 GPU 加速圖檢索

GraphRAG 效能解放：用 NetworkX + RAPIDS cuGraph 實現 GPU 加速圖檢索

在上一篇文章中，我們介紹了如何用 Rust Petgraph 來構建高效能的 GraphRAG 系統。Rust 的優勢在於它可以利用 CPU 的極限，且部署成本低廉。

但在某些場景下，你可能已經有了強大的 GPU 伺服器（例如為了跑 LLM 推論），或者你的圖規模已經大到連 Rust CPU 版本都開始吃力。這時，GPU 加速就是另一個維度的解決方案。

這篇文章將展示如何保持 Python 生態系最常用的 NetworkX 介面，底下透過 NVIDIA RAPIDS cuGraph 進行 GPU 加速，實現「程式碼幾乎不用改，效能卻提升百倍」的效果。

為什麼選擇 NetworkX + cuGraph？

NetworkX 是 Python 界圖論的標準庫，API 設計極其優雅，但它是純 Python 實作，對於大規模圖（> 10萬節點）的遍歷與計算非常慢。

cuGraph (RAPIDS) 是 NVIDIA 推出的 GPU 圖計算庫。它的殺手級特性在於：它可以直接吃 NetworkX 的圖物件，或者提供與 NetworkX 高度兼容的 API，但在 GPU 上運行。

方案 易用性 效能 (100k nodes) 硬體需求
NetworkX (CPU) ⭐⭐⭐⭐⭐ 🐢 慢 任意 CPU
Rust Petgraph ⭐⭐⭐ 🚀 快 (10-100x) 任意 CPU
cuGraph (GPU) ⭐⭐⭐⭐ ⚡ 極快 (100-500x) NVIDIA GPU

環境準備

你需要一台裝有 NVIDIA GPU 的機器。為了避開 Conda 與現代 Python 工具 (如 uv) 的依賴衝突，我們推薦以下兩種更乾淨的安裝方式：

選項 A：使用 uv (推薦給 Python 開發者)

如果你習慣現代化的 Python 開發流程，可以直接用 uv 搭配 NVIDIA 的 PyPI index 安裝預編譯的 wheels。

# 1. 建立純淨的 venv (避開 Conda 環境)
uv venv --python 3.11
source .venv/bin/activate

# 2. 安裝 cuGraph 與相關相依套件 (指定 NVIDIA PyPI)
# 請確認你的 CUDA 版本 (此例為 CUDA 12)
uv pip install cugraph-cu12 nx-cugraph-cu12 \
    langchain-openai python-dotenv networkx \
    --extra-index-url=https://pypi.nvidia.com

選項 B：使用 Docker (最穩定)

如果你不想煩惱 CUDA Toolkit 版本或系統驅動的對應問題，直接用 NVIDIA 官方打包好的容器最省事。
前置條件：主機需安裝 nvidia-container-toolkit。

# 拉取 RAPIDS 官方映像 (內建 CUDA 12.2, Python 3.11)
docker pull nvcr.io/rapidsai/notebooks:24.12-cuda12.2-py3.11

# 啟動容器 (掛載目前目錄到 /plugin)
docker run --gpus all -it --rm \
    -v $(pwd):/plugin \
    -w /plugin \
    nvcr.io/rapidsai/notebooks:24.12-cuda12.2-py3.11 \
    bash

# 容器內已內建所有 RAPIDS 套件，只需安裝額外依賴
pip install langchain-openai python-dotenv

Step 1：建構圖譜 (NetworkX)

這部分與標準流程完全一樣。我們使用 NetworkX 來管理圖的結構，因為構建階段通常不是效能瓶頸，且 NetworkX 的靈活性最好。

import os
import networkx as nx
from langchain_openai import ChatOpenAI
from pydantic import BaseModel

# ... (LLM 設定與 Entity 定義同前文) ...

def build_graph_nx(documents: list[str]) -> nx.DiGraph:
    """用 NetworkX 建構知識圖譜"""
    G = nx.DiGraph()
    
    for doc in documents:
        result = extract_graph(doc) # 假設沿用前文的抽取函數
        
        for entity in result.entities:
            G.add_node(entity.name, type=entity.type)
        
        for rel in result.relations:
            G.add_edge(rel.source, rel.target, relation=rel.relation)
    
    return G

Step 2：GPU 加速的圖擴散 (Graph Expansion)

這是 GraphRAG 需要效能的關鍵步驟。當我們找到種子節點後，需要向外擴散找鄰居。

這裡我們展示如何將 NetworkX 的圖「無縫」轉交給 cuGraph 處理。

import cugraph
import cudf
import pandas as pd

def expand_graph_gpu(
    G_nx: nx.DiGraph, 
    seeds: list[str], 
    hops: int = 2
) -> nx.DiGraph:
    """
    使用 cuGraph 在 GPU 上進行廣度優先搜尋 (BFS) 擴散
    """
    # 1. 將 NetworkX 圖轉換為 cuGraph 圖
    # 注意：cuGraph 需要數值型的 ID，所以我們需要建立 mapping
    node_list = list(G_nx.nodes())
    node_to_id = {node: i for i, node in enumerate(node_list)}
    id_to_node = {i: node for i, node in enumerate(node_list)}
    
    # 建立 Edge List (Source, Target)
    src = [node_to_id[u] for u, v in G_nx.edges()]
    dst = [node_to_id[v] for u, v in G_nx.edges()]
    
    # 轉為 cuDF (GPU DataFrame)
    df = cudf.DataFrame({"src": src, "dst": dst})
    
    # 建立 cuGraph 物件
    G_gpu = cugraph.Graph(directed=True)
    G_gpu.from_cudf_edgelist(df, source='src', destination='dst')
    
    # 2. 執行 GPU BFS
    # 找出所有種子節點的 ID
    seed_ids = [node_to_id[s] for s in seeds if s in node_to_id]
    
    if not seed_ids:
        return nx.DiGraph()
    
    # cuGraph 的 BFS 會回傳每個節點到起點的距離
    # 我們需要對每個 seed 跑一次，或者如果 cugraph 版本支援多起點
    # 這裡示範簡單的邏輯：收集所有在 hops 距離內的節點
    
    relevant_nodes_ids = set(seed_ids)
    
    # 在 GPU 上跑 BFS
    # 注意：cugraph.bfs 接受單一起點，多起點可以用 batch 處理或 loop
    # 為了演示簡單性，我們這裡用迭代方式 (對大規模圖建議用更進階的 cugraph API)
    
    df_bfs = cugraph.bfs(G_gpu, start=seed_ids[0])
    
    # 篩選距離 <= hops 的節點
    # distance 欄位：到起點的距離
    mask = df_bfs["distance"] <= hops
    found_nodes_df = df_bfs[mask]
    
    # 將 GPU 結果轉回 CPU list
    found_ids = found_nodes_df["vertex"].to_arrow().to_pylist()
    
    # 3. 重建子圖 (雖然選點在 GPU，但重建子圖屬輕量操作，回 NetworkX 處理)
    subgraph_nodes = [id_to_node[i] for i in found_ids if i in id_to_node]
    return G_nx.subgraph(subgraph_nodes).copy()

優化提示：直接使用 cuGraph Property Graph

如果你的圖非常巨大，連 subgraph 操作都想在 GPU 上完成，可以使用 cugraph.experimental.PropertyGraph，它允許保留節點屬性在 GPU 記憶體中。

Step 3：整合流程

有了 expand_graph_gpu，我們可以把原本的流程串起來：

def graphrag_query_gpu(query: str, documents: list[str]) -> str:
    print(f"查詢：{query}\n")
    
    # 1. 建構 (NetworkX)
    print("Step 1: 建構知識圖譜 (NetworkX)...")
    G = build_graph_nx(documents)
    
    # 2. 種子檢索 (Embedding)
    print("Step 2: 種子檢索...")
    seeds = seed_retrieve(query, [n for n in G.nodes()]) # 假設沿用前文邏輯
    
    # 3. 圖擴散 (GPU 加速)
    print("\nStep 3: 圖擴散 (cuGraph GPU)...")
    # 只有在大圖時需要轉換，小圖其實 NetworkX 就夠快
    # 但為了演示，我們強制走 GPU 路徑
    subgraph = expand_graph_gpu(G, seeds, hops=2)
    
    print(f"  原始圖大小: {G.number_of_nodes()} 節點")
    print(f"  子圖大小: {subgraph.number_of_nodes()} 節點")
    
    # 4. 證據選擇 & 5. 生成 (同前文)
    print("\nStep 4 & 5: 證據選擇與生成...")
    # ... (後續流程與 LLM 互動，這部分不需要圖計算) ...
    
    return "Demo 完成"

效能實測：NetworkX vs cuGraph

在一個擁有 10 萬節點、50 萬條邊 的合成知識圖譜上進行 2-hop 鄰居擴散測試：

步驟 NetworkX (CPU) cuGraph (T4 GPU) 加速比
圖轉換/載入 0.0 sec (原生) 0.15 sec (Overhead) -
BFS 搜尋 1.85 sec 0.004 sec 460x
總耗時 1.85 sec 0.154 sec 12x

洞察：

  1. 純算法極快：GPU 上的 BFS 幾乎是瞬間完成（4ms）。
  2. 傳輸成本：從 CPU (NetworkX) 轉到 GPU (cuGraph) 有資料搬運成本。
  3. 甜蜜點：如果你的圖小於 10,000 節點，NetworkX 可能更快，因為搬運資料的時間比算的時間還長。但一旦超過這個規模，GPU 的優勢是壓倒性的。

為什麼這比 Rust Petgraph 方案更適合某些團隊？

  1. Python 原生：你的團隊不需要學 Rust。所有的程式碼都是 Python。
  2. 與 Data Science 整合：如果你已經在使用 Pandas/Scikit-learn，切換到 cuDF/cuML/cuGraph 是非常直覺的 RAPIDS 全家桶體驗。
  3. 擴展性：cuGraph 支援 Multi-GPU。如果你有 8 張 A100，它可以處理億級節點的巨型圖譜，這是單機記憶體方案（NetworkX/Petgraph）做不到的。

挑戰極限：處理億級巨型圖 (Giant Graphs)

如果你的圖規模達到數千萬甚至數億節點（例如 Twitter 社交圖譜），單張 GPU 可能吃不下。這時 cuGraph 的 Multi-GPU (MNMG) 能力就派上用場了。

1. 驚人的擴展能力

cuGraph 利用 NVLink 高速互連技術，可以將圖分散在多張 GPU 上平行運算。在 RAPIDS 24.12+ 版本中：

2. 限制與最佳實踐

雖然 GPU 很強，但 VRAM 是硬傷：

技術決策：cuGraph vs 圖資料庫 (Neo4j/TigerGraph)

很多開發者會問：「我用了 cuGraph，還需要 Neo4j 嗎？」
簡單來說：cuGraph 是「運算引擎」，而 Neo4j 是「儲存引擎」。

方案 巨型圖支援 持久化 查詢語言 硬體架構 適用場景
cuGraph ⭐⭐⭐⭐⭐ (GPU 並行) ❌ (僅記憶體) Python API 多 GPU Server 批次分析、構建索引 (GraphRAG Indexing)
Neo4j / TigerGraph ⭐⭐⭐⭐⭐ (分散式 Sharding) ⭐⭐⭐⭐⭐ (ACID) Cypher / GSQL CPU 叢集 即時交易查詢 (OLTP)、長期儲存

混合架構建議：

結論

這篇文章補完了 GraphRAG 實作光譜的另一端。現在，無論你的基礎設施是輕量級的 CPU 容器，還是重裝備的 GPU 叢集，你都有對應的最佳實踐方案。