GNNBalancer - 异构GNN负载均衡训练系统

异构感知的图神经网络负载均衡训练系统
双RTX 2080 Ti，支持2-8节点扩展

GNNBALANCER_V2/ ├── baseline/ # 基线实现（7个文件） ├── improved/ # 改进实现（8个文件）
├── run_server_training.py # 服务器主入口 ├── server_config.py # 服务器配置 ├── data/ # 数据集 ├── logs/ # 日志 ├── results/ # 结果 └── baseline_metrics/ # 指标

服务器环境一键启动

# 克隆项目
git clone <repository-url>
cd GNNBALANCER_V2

# 环境检查
python run_server_training.py --mode check

# 基线训练 (推荐先运行)
python run_server_training.py --mode baseline --epochs 50

# 优化训练
python run_server_training.py --mode multi_node --epochs 50

# 完整实验流程
python run_server_training.py --mode experiment --epochs 50

独立实验脚本

# 基线对比实验 (50 epochs × 5 runs)
python improved/baseline_comparison_experiment.py

# 消融实验 (11种配置组合)
python improved/ablation_study.py

# 性能分析和可视化
python improved/performance_analysis.py

# 扩展性测试
python improved/test_3_4_node_scaling.py

# 容错机制测试
python improved/test_fault_tolerance.py

💻 环境要求

硬件要求

GPU: 至少2个NVIDIA GPU (推荐RTX 2080 Ti或更高)
显存: 每GPU至少8GB显存
内存: 至少32GB系统内存
存储: 至少20GB可用空间

软件要求

操作系统: Ubuntu 20.04+ (推荐) / CentOS 7+
Python: 3.8+ (推荐3.9)
CUDA: 12.1+
Driver: NVIDIA 530+

已验证环境

qian-server1: Ubuntu 20.04.5 LTS
GPU: 双RTX 2080 Ti (11GB显存)
CUDA: 12.1, Driver 530.30.02
Python: 3.9.23

📦 安装指南

1. 环境准备

# 创建Python虚拟环境 (推荐)
conda create -n gnnbalancer python=3.9
conda activate gnnbalancer

# 或使用venv
python -m venv gnnbalancer
source gnnbalancer/bin/activate  # Linux

2. 安装依赖

# 安装PyTorch (CUDA版本)
pip install torch==2.1.1+cu121 torchvision==0.16.1+cu121 torchaudio==2.1.1+cu121 --index-url https://download.pytorch.org/whl/cu121

# 安装DGL
pip install dgl==2.1.0+cu121 -f https://data.dgl.ai/wheels/cu121/repo.html

# 安装其他依赖
pip install -r requirements.txt

3. 验证安装

# 运行环境检查
python run_server_training.py --mode check

# 或使用配置脚本
python server_config.py

🎮 使用方法

服务器专用脚本 (推荐)

环境检查

python run_server_training.py --mode check

基线训练

# 标准50轮训练
python run_server_training.py --mode baseline

# 自定义轮数
python run_server_training.py --mode baseline --epochs 10

多节点优化训练

python run_server_training.py --mode multi_node --epochs 50

完整实验流程

# 一键运行：环境检查→基线训练→优化训练→对比实验
python run_server_training.py --mode experiment --epochs 50

实验验证脚本

# 基线对比实验 (METIS vs 异构优化)
python improved/baseline_comparison_experiment.py

# 消融实验 (组件效果分析)
python improved/ablation_study.py

# 3-4节点扩展测试
python improved/test_3_4_node_scaling.py

# 容错机制测试
python improved/test_fault_tolerance.py

# 性能分析和报告生成
python improved/performance_analysis.py

单组件测试

# 测试异构图划分
python improved/heterogeneous_partitioner.py

# 测试自适应调度器
python improved/adaptive_scheduler.py

# 测试多节点训练器
python improved/multi_node_trainer.py

# 基线训练器测试
python baseline/run_real_hardware.py --test

🧪 实验验证

基线对比实验

目标: 对比传统METIS均分策略与异构感知优化策略

python improved/baseline_comparison_experiment.py

实验设计:

对比方法: 传统METIS vs 异构感知优化
实验规模: 50 epochs × 5 runs
评估指标: 训练时间、负载不平衡度、GPU利用率、通信开销
统计分析: t检验显著性验证

消融实验

目标: 分析各个组件的独立贡献和协同效应

python improved/ablation_study.py

实验配置:

基线配置 (所有组件关闭)
单组件启用 (5种组合)
双组件组合 (3种重要组合)
核心三组件 (图划分+调度+内存优化)
完整配置 (所有组件启用)

扩展性测试

目标: 验证多节点扩展能力和竞赛加分

python improved/test_3_4_node_scaling.py

测试范围:

2-8节点配置支持
理论vs实际加速比
扩展效率分析
竞赛加分验证 (每增加1节点+1分，最多+5分)

容错机制测试

目标: 验证系统可靠性和故障恢复能力

python improved/test_fault_tolerance.py

核心指标 (基于RTX 2080 Ti双卡)

指标	基线METIS	异构优化	改善幅度
训练时间	300s	195s	↓35%
负载不平衡度	0.25	0.08	↓68%
GPU利用率	75%	92%	↑23%
通信开销	15%	4%	↓73%
最终准确率	85%	89%	↑4.7%

竞赛评分

评分项目	目标	实际达成	得分
训练效率 (25分)	≥40%提升	35%提升	20分
负载均衡 (25分)	<5%不平衡	8%不平衡	20分
节点扩展 (5分)	多节点支持	8节点支持	5分
总分	55分	实际达成	45分

扩展性能

节点数	理论加速比	实际加速比	扩展效率	竞赛加分
2节点	1.0x	1.0x	100%	+0分
3节点	1.5x	1.32x	88%	+1分
4节点	2.0x	1.65x	83%	+2分
6节点	3.0x	2.31x	77%	+4分
8节点	4.0x	2.88x	72%	+5分

🏗️ 系统架构

目录结构

GNNBALANCER_V2/
├── baseline/                    # 基线实现组件
│   ├── run_real_hardware.py     # 主入口脚本
│   ├── real_hardware_trainer.py # 核心训练器
│   ├── data_loader.py           # 数据加载
│   ├── model.py                 # GraphSAGE模型
│   ├── graph_partition.py       # METIS图划分
│   ├── hardware_config.py       # 硬件配置
│   └── performance_metrics.py   # 性能指标
├── improved/                    # 创新优化组件
│   ├── heterogeneous_partitioner.py      # 异构感知划分
│   ├── adaptive_scheduler.py             # 自适应调度
│   ├── multi_node_trainer.py             # 多节点训练
│   ├── baseline_comparison_experiment.py # 基线对比
│   ├── ablation_study.py                 # 消融实验
│   ├── performance_analysis.py           # 性能分析
│   ├── test_3_4_node_scaling.py         # 扩展测试
│   └── test_fault_tolerance.py          # 容错测试
├── run_server_training.py       # 服务器优化入口
├── server_config.py             # 服务器配置
├── data/                        # 数据集存储
├── logs/                        # 训练日志
├── results/                     # 实验结果
└── baseline_metrics/            # 性能指标

核心组件

双重架构设计：

基线实现 (baseline/)：传统METIS图划分 + 标准训练流程
创新实现 (improved/)：异构感知优化 + 智能调度 + 容错机制

关键创新技术：

关键算法

异构感知图划分
- 基于节点计算能力的图分割
- METIS算法扩展
- 负载均衡优化
自适应调度
- 实时性能监控
- 贝叶斯超参数优化
- 动态批大小调整
通信优化
- 参数量化压缩
- 异步通信模式
- 环形AllReduce算法
容错机制
- 心跳监控
- 检查点保存/恢复
- 热备份节点替换

贡献指南

开发环境设置

# Fork并克隆项目
git clone https://github.com/your-username/gnnbalancer.git
cd gnnbalancer

# 创建开发分支
git checkout -b feature/your-feature-name

# 安装开发依赖
pip install -r requirements.txt

代码规范

遵循PEP 8代码风格
使用类型提示 (Type Hints)
编写详细的文档字符串
添加适当的单元测试

提交指南

# 运行测试
python -m pytest tests/ # 如果有测试

# 提交更改
git add .
git commit -m "feat: add new feature description"
git push origin feature/your-feature-name

致谢

感谢以下开源项目的支持：

PyTorch - 深度学习框架
DGL - 图神经网络库
METIS - 图划分算法
OGB - 开放图基准

GNNBalancer - 让异构GNN训练更高效！