# torch_geometric

**Repository Path**: onescience-ai/torch_geometric

## Basic Information

- **Project Name**: torch_geometric
- **Description**: 本项目为基于flagos统一中间层实现的torch_scatter库
- **Primary Language**: Unknown
- **License**: MIT
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2026-10-12
- **Last Updated**: 2026-11-12

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# PyTorch Geometric 2.9.0 — CUDA / FlagOS 兼容版

这是基于 PyTorch Geometric（PyG）2.9.0 的设备兼容版本，保留 `torch_geometric` 的公开导入名和原有 CPU/CUDA 用法，并适配 Torch-FL 提供的 `flagos` 设备。

本版为常用图算子提供多后端 Triton 加速，并默认启用 C++ 主机辅助。新增 GPU 数值内核均由 Triton 实现；C++ 仅负责内存分配、元数据检查、内核启动和 Autograd 连接，不执行张量数值计算。

文档参考：[PyG Documentation](https://pytorch-geometric.readthedocs.io)；

## 库介绍

PyG 是基于 PyTorch 的图神经网络库，用于图、异构图、点云等结构化数据的表示学习。

- **数据表示**：`Data`、`HeteroData` 和 `Batch` 表示单图、异构图及多个图的批次。
- **图神经网络**：提供 GCN、GraphSAGE、GAT、GIN 等层与模型，可通过 `MessagePassing` 定义自己的消息传递层。
- **数据加载与采样**：提供 `DataLoader`、`NeighborLoader`、GraphSAINT 等入口，服务于小图批处理和大图采样训练。
- **数据集与变换**：提供常用图数据集及特征处理、图结构变换工具。
- **图操作与分析**：提供 scatter、稀疏计算、池化、图解释和性能分析工具。

以上是 PyG 的功能范围，并不表示每个 API 都已在 FlagOS 上验证；可选扩展、编译和多卡功能的边界见后文。

## 环境要求

- Python >= 3.10，与本仓库的包元数据保持一致。
- 已安装适合目标硬件的 PyTorch；PyG 不负责安装或切换 PyTorch 后端。
- Hygon 主机扩展编译需要 Linux、支持 C++17 的编译器、Python 开发头文件，以及当前 PyTorch 的 C++ 头文件和库；可通过 `CXX` 指定编译器。
- 原生 CUDA：使用匹配的 PyTorch CUDA 环境和驱动，不需要 `import torch_fl`。
- FlagOS：预先配置 Torch-FL / FlagGems / FlagTree 环境，在程序中先 `import torch_fl`，再导入 `torch` 和 PyG。

## 安装

### 原生 CPU / CUDA 环境

源码安装和构建 wheel 时，通过 `FLAGOS_BACKEND` 选择后端，可选 `hygon`、`nvidia`、`mthreads`，默认 `hygon`：

```bash
# Hygon/DCU（默认，安装时编译 C++ 主机辅助扩展）
python -m pip install . --no-deps --no-build-isolation
FLAGOS_BACKEND=hygon python -m pip install . --no-deps --no-build-isolation

# NVIDIA/GPU
FLAGOS_BACKEND=nvidia python -m pip install . --no-deps --no-build-isolation

# MThreads/MUSA
FLAGOS_BACKEND=mthreads python -m pip install . --no-deps --no-build-isolation
```

后端在安装/构建时写入包配置，运行时修改变量不会切换已安装的后端；切换需重新安装或使用对应 wheel。只有 Hygon 构建 C++ 主机辅助，GPU 数值计算均使用已有 Triton 实现。不同后端的 editable 安装请使用独立源码目录，避免共用后端配置。

先配置好 PyTorch，再安装本仓库。Hygon 源码安装默认编译并安装 C++ 主机辅助扩展；编译失败则安装报错，运行时仅加载已编译的扩展。编译主机扩展不需要 `nvcc`/`hipcc`，GPU 运行仍需匹配的驱动与后端运行时。

```bash
cd /path_to_dir/torch_geometic
python -m pip install . --no-deps --no-build-isolation
```

### Torch-FL / FlagOS 环境

先激活 Torch-FL 环境，再执行安装。C++ 扩展使用当前环境的 PyTorch 头文件和 ABI 构建。

```bash
cd /path_to_dir/torch_geometic
FLAGOS_BACKEND=hygon python -m pip install . --no-deps --no-build-isolation
```

开发时可使用 editable 安装：

```bash
FLAGOS_BACKEND=hygon python -m pip install -e . --no-deps --no-build-isolation
```

`--no-build-isolation` 使用当前环境的 PyTorch 编译；`--no-deps` 保留现有依赖，请先按 `pyproject.toml` 准备运行依赖。安装后无需设置启用开关；GPU 数值内核仍由 Triton 按输入配置 JIT 编译。

以下开关默认启用，运行时应在启动 Python 前设置；取消变量或设为 `1` 恢复：

| 环境变量 | 设为 `0` 的作用 |
| --- | --- |
| `PYG_GRAPH_HOST` | 安装/构建时跳过 C++ 扩展编译；运行时禁用扩展，使用已有 Python/Triton 回退。跳过编译后需重新安装才能启用扩展。 |
| `PYG_GRAPH_NATIVE_LAUNCH` | 仅关闭直接 HIP 启动优化，保留 C++ 辅助，使用 Triton 原有启动器；不影响安装编译。 |
| `PYG_GRAPH_NATIVE_AUTOGRAD` | 仅改用 Python Autograd 连接，GPU 梯度仍由 Triton 计算；不影响安装编译。 |

检查版本和导入位置，确认没有误用旧安装：

```bash
python -c "import torch_geometric; from torch_geometric._ops._build_config import BACKEND; print(torch_geometric.__version__, BACKEND); print(torch_geometric.__file__)"
```

### 构建 wheel 包

构建 wheel 时使用相同的后端变量；以下以 Hygon 为例，编译后的 C++ 主机扩展会打入 wheel：

```bash
cd /path_to_dir/torch_geometic
FLAGOS_BACKEND=hygon python -m pip wheel . --no-deps --no-build-isolation --wheel-dir dist
```

wheel 文件名和包版本包含 `+triton.<backend>` 标记，例如 Hygon 为 `2.9.0+triton.hygon`，其他后端为 `+triton.nvidia`、`+triton.mthreads`。`torch_geometric.__version__` 与安装元数据使用相同版本。

生成的 wheel 位于 `dist/`。选择与目标 Python、PyTorch ABI、操作系统和架构匹配的包安装；Hygon wheel 已包含 C++ 编译产物，安装时无需再次编译。

```bash
python -m pip install dist/torch_geometric-2.9.0+triton.hygon-*.whl
```

Torch-FL 与 DAS 环境的 PyTorch ABI 可能不同，应分别在对应环境构建 wheel。GPU 算子由 Triton 生成，不将 CUDA/HIP C++ 数值内核打入扩展。

### 可选依赖

基础的 `Data`、`GCNConv` 和下方示例不要求额外安装 `pyg-lib`、`torch-scatter` 或 `torch-sparse`。按实际使用的功能再安装相关依赖：

| 依赖 | 主要用途 |
| --- | --- |
| `pyg-lib` | 图采样及部分图算子 |
| `torch-scatter` | 扩展 scatter / segment 归约 |
| `torch-sparse` | `SparseTensor` 及相关稀疏操作 |
| GraphGym / Lightning 相关依赖 | 实验配置及训练框架集成 |
| Transformers 等依赖 | LLM / RAG 相关功能 |


可安装本仓库定义的依赖集合，例如：

```bash
python -m pip install '.[test]' --no-build-isolation
python -m pip install '.[graphgym]' --no-build-isolation
```

这不会自动完成 Torch-FL 或所有可选图扩展的配置。

## 基本用法：GCN 前向与训练一步

`x` 的形状为 `[节点数, 特征数]`；`edge_index` 是形状为 `[2, 边数]` 的整型张量，第一行是源节点、第二行是目标节点。数据、索引和模型应放在同一设备。

以下两个示例都使用小型人工图，不下载数据集，可分别保存为 Python 文件直接运行。它们演示前向、损失、反向和参数更新，不代表模型已经完成训练。

### 原生 CUDA

```python
import torch
import torch.nn.functional as F
from torch_geometric.data import Data
from torch_geometric.nn import GCNConv

device = torch.device('cuda:0')
data = Data(
    x=torch.randn(4, 3),
    edge_index=torch.tensor([[0, 1, 1, 2, 2, 3], [1, 0, 2, 1, 3, 2]],
                            dtype=torch.long),
    y=torch.tensor([0, 1, 0, 1], dtype=torch.long),
).to(device)
model = GCNConv(3, 2).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

model.train()
optimizer.zero_grad()
out = model(data.x, data.edge_index)
loss = F.cross_entropy(out, data.y)
loss.backward()
optimizer.step()
print('输出形状:', tuple(out.shape))  # (4, 2)
print('设备:', out.device)           # cuda:0
print('loss:', loss.item())
```

原生 CPU 使用时，将设备改为 `torch.device('cpu')` 即可。

### Torch-FL / FlagOS

Python 程序中显式加载 Torch-FL，所有加速设备使用 `flagos:0`：

```python
import torch_fl
import torch
import torch.nn.functional as F
from torch_geometric.data import Data
from torch_geometric.nn import GCNConv

device = torch.device('flagos:0')
data = Data(
    x=torch.randn(4, 3),
    edge_index=torch.tensor([[0, 1, 1, 2, 2, 3], [1, 0, 2, 1, 3, 2]],
                            dtype=torch.long),
    y=torch.tensor([0, 1, 0, 1], dtype=torch.long),
).to(device)
model = GCNConv(3, 2).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

model.train()
optimizer.zero_grad()
out = model(data.x, data.edge_index)
loss = F.cross_entropy(out, data.y)
loss.backward()
optimizer.step()
print('输出形状:', tuple(out.shape))  # (4, 2)
print('设备:', out.device)           # flagos:0
print('loss:', loss.item())
```

两种用法的 PyG API 相同，主要区别是 Torch-FL 的显式导入和目标设备名称。不要把 FlagOS 张量与 CUDA 模型或 CPU 标签混在同一次加速计算中。

## 多图批处理与池化

`DataLoader` 将多张图合并为一个 `Batch`，其中 `batch` 标记每个节点属于哪张图；`global_mean_pool` 将节点特征聚合为图特征。下面在 CPU 构造图，在主进程将整个批次转到目标设备。

### 原生 CUDA

```python
import torch
from torch_geometric.data import Data
from torch_geometric.loader import DataLoader
from torch_geometric.nn import global_mean_pool

graphs = [
    Data(x=torch.ones(2, 3),
         edge_index=torch.tensor([[0, 1], [1, 0]], dtype=torch.long)),
    Data(x=torch.full((2, 3), 2.0),
         edge_index=torch.tensor([[0, 1], [1, 0]], dtype=torch.long)),
]
loader = DataLoader(graphs, batch_size=2, num_workers=0)
batch = next(iter(loader)).to('cuda:0')
out = global_mean_pool(batch.x, batch.batch)
print(out.cpu())  # tensor([[1., 1., 1.], [2., 2., 2.]])
print(out.device)
```

### Torch-FL / FlagOS

```python
import torch_fl
import torch
from torch_geometric.data import Data
from torch_geometric.loader import DataLoader
from torch_geometric.nn import global_mean_pool

graphs = [
    Data(x=torch.ones(2, 3),
         edge_index=torch.tensor([[0, 1], [1, 0]], dtype=torch.long)),
    Data(x=torch.full((2, 3), 2.0),
         edge_index=torch.tensor([[0, 1], [1, 0]], dtype=torch.long)),
]
loader = DataLoader(graphs, batch_size=2, num_workers=0)
batch = next(iter(loader)).to('flagos:0')
out = global_mean_pool(batch.x, batch.batch)
print(out.cpu())  # tensor([[1., 1., 1.], [2., 2., 2.]])
print(out.device)
```

数据集的读取与预处理通常留在 CPU，再将训练批次迁移到目标设备。更完整的模型、异构图和采样训练示例见 [examples](examples/README.md)；已有示例中的硬编码 CUDA 设备或专用依赖仍需按实际环境检查，不能认为所有示例都已自动适配。


## 正确性测试

在本仓库目录执行；测试依赖可通过 `python -m pip install '.[test]' --no-build-isolation` 安装，部分测试另有可选依赖。

原生 CPU / CUDA 环境：

```bash
python -m pytest test -q
```

Torch-FL 环境 ：

```bash
python -m pytest test/test_flagos_compat.py --torch-fl -q -o addopts=
python -m pytest test/test_flagos_share_memory.py --torch-fl -q -o addopts=
python -m pytest test --torch-fl -q
```

`--torch-fl` 是本仓库测试入口提供的参数，用于提前加载 Torch-FL 并验证 CPU / FlagOS 路径。默认测试配置排除 `dataset` 和 `rag` 标记；上述全目录命令不是无条件执行所有联网数据集和 RAG 测试，也不是全量已通过的声明。

需要覆盖这些额外分组时，可在准备好依赖、网络和数据后执行：

```bash
python -m pytest test --torch-fl -q -m 'dataset or rag'
```

## 性能测试与更多资料

本仓库保留 [benchmark](benchmark/README.md) 和 [examples](examples/README.md)。运行前按对应说明准备数据、依赖及设备参数；其中 CUDA 专用入口不应直接当作 FlagOS benchmark。

性能比较应使用相同输入、设备和路由，区分首次编译与稳定运行，并计入 CPU 回退的数据传输开销。本适配版本没有宣称 FlagOS 与原生 CUDA 性能等价。

- [算子实现与后端支持范围](torch_geometric/_ops/README.md)
- [测试入口说明](test/README_TORCH_FL.md)
- [完整文档源码](docs/README.md)
- [上游贡献指南](.github/CONTRIBUTING.md)

## 许可证与引用

本仓库保留 PyG 上游的 [MIT 许可证](LICENSE)。研究中使用 PyG 时，请引用上游工作；

```bibtex
@inproceedings{Fey/Lenssen/2019,
  title={Fast Graph Representation Learning with {PyTorch Geometric}},
  author={Fey, Matthias and Lenssen, Jan E.},
  booktitle={ICLR Workshop on Representation Learning on Graphs and Manifolds},
  year={2019},
}
```