[Paper] 超越单GPU：将 PDLP 扩展到分布式多GPU系统

Source: arXiv - 2601.07628v1

概述

本文介绍了一种分布式内存实现的原始‑对偶混合梯度（PDHG）算法，该实现可以在多块 GPU 上运行，突破了现有单 GPU 线性规划求解器的内存和速度限制。通过将约束矩阵划分到二维 GPU 网格上并使用高性能 NCCL 通信，作者展示了此前因规模过大而无法在单个加速器上求解的大规模线性规划，能够以完整的双精度精度并获得显著的加速效果。

关键贡献

• 多GPU PDHG 框架：将已验证的 PDHG 求解器（cuPDLPx）从单 GPU 扩展到任意数量的 GPU，使用 LP 约束矩阵的二维网格划分。
• 负载均衡的数据分配：引入 非零感知 划分方案以及 块级随机洗牌 策略，使每个 GPU 都保持忙碌，避免出现拖慢的情况。
• 通信高效的设计：利用 NVIDIA 的 NCCL 库和融合的 CUDA 内核，以最小开销同步原始‑对偶更新。
• 完整的 FP64 数值保真度：即使在分布式执行下也保证双精度结果，这是许多工业优化流水线的要求。
• 广泛的实证验证：在来自 MIPLIB、Hans 实例和真实数据集的基准上展示出强大的可扩展性，相比最佳单 GPU 基准实现数倍加速。

方法论

1. Problem formulation – The LP is expressed in standard form, and the PDHG algorithm iteratively updates primal (variables) and dual (constraints) vectors using gradient steps and proximal operators.
2. 2‑D grid partitioning – The large sparse constraint matrix (A) is sliced both row‑wise and column‑wise, assigning each sub‑matrix to a distinct GPU in a logical grid. This spreads both memory usage and compute work.
3. Nonzero‑aware distribution – Instead of naïvely splitting rows/columns evenly, the algorithm first counts non‑zero entries per block and redistributes blocks so that each GPU receives roughly the same number of non‑zeros, which correlates with actual compute load.
4. Block‑wise random shuffling – Within each GPU, blocks are processed in a random order each iteration, reducing correlation‑induced load imbalance and improving convergence stability.
5. Communication layer – After each PDHG step, GPUs need to exchange boundary primal/dual values. This is done with NCCL’s all‑reduce and point‑to‑point primitives, wrapped in fused kernels that combine multiple small messages into a single GPU launch, cutting latency.
6. Implementation – Built on top of the existing cuPDLPx codebase, the authors added the distributed data structures, communication routines, and kernel fusions while keeping the core algorithm unchanged.

结果与发现

• Memory breakthrough：在 32 GB GPU 上出现内存不足错误的问题，在每块 GPU ≤ 12 GB 的 4‑GPU 集群上能够轻松运行。
• Communication overhead：即使在 16 GPUs 时，NCCL 通信也仅占总运行时间的 < 12 %，验证了融合 kernel 和 2‑D 划分的有效性。
• Numerical accuracy：所有实验均保持 FP64 解的质量，残差在相对容差 1e‑12 内与单 GPU 求解器匹配。

Practical Implications

• 可扩展的优化服务: 云服务提供商现在可以将 LP 求解作为多 GPU 服务提供，而无需担心内存上限，从而实现对大型物流、金融或能源规划模型的按需求解。
• 与现有流水线的集成: 由于实现基于 cuPDLPx 并使用标准的 CUDA/NCCL API，开发者可以将其直接嵌入当前的 GPU 加速数据科学栈（如 RAPIDS、PyTorch），几乎无需修改代码。
• 成本效益的硬件利用: 与其采购单个大型 GPU（可能稀缺且昂贵），组织可以通过聚合多块普通 GPU 获得相当的性能，从而提升硬件投资回报率。
• 其他求解器的基础: 2‑D 划分和负载均衡的思路可以迁移到其他一阶方法（如 ADMM、随机梯度），这些方法同样受到内存瓶颈的限制。

限制与未来工作

• 稀疏矩阵结构依赖：当前的划分器假设稀疏模式相对均匀；高度不规则的矩阵仍可能导致负载不平衡。
• 网络拓扑敏感性：实验在高速 NVLink 互连的集群上进行；在基于以太网的多节点设置上，由于更高的延迟，性能可能下降。
• 算法扩展：作者计划探索自适应步长方案和混合精度变体，以在保持精度的同时进一步缩短运行时间。
• 更广泛的基准套件：未来工作包括在机器学习中新兴的线性规划基准（例如大规模支持向量机）上进行测试，并与更高级的建模语言集成。

作者

• Hongpei Li
• Yicheng Huang
• Huikang Liu
• Dongdong Ge
• Yinyu Ye

论文信息

• arXiv ID: 2601.07628v1
• 分类: math.OC, cs.DC
• 出版日期: 2026年1月12日
• PDF: 下载 PDF