[论文解读] DistStat.jl: Towards Unified Programming for High-Performance Statistical Computing Environments in Julia

Seyoon Ko, Hua Zhou+2 more|arXiv (Cornell University)|Oct 30, 2020

Gene expression and cancer classification参考文献 42被引用 4

一句话总结

DistStat.jl 是一个 Julia 软件包，通过分布式数组抽象，在 CPU 集群和多 GPU 环境中统一实现高性能统计计算，支持统计算法的透明、可扩展执行。它加速了大规模基因组分析——在 AWS 上，使用 400,000 名 UK Biobank 受试者对包含 500,000 个变量的 Cox 模型进行拟合，耗时不足 50 分钟，相比等效的 Python 实现，在 CPU 上提速 20%–30%，在 GPU 上性能相当，即使使用未优化的内核。

ABSTRACT

The demand for high-performance computing (HPC) is ever-increasing for everyday statistical computing purposes. The downside is that we need to write specialized code for each HPC environment. CPU-level parallelization needs to be explicitly coded for effective use of multiple nodes in cluster supercomputing environments. Acceleration via graphics processing units (GPUs) requires to write kernel code. The Julia software package DistStat.jl implements a data structure for distributed arrays that work on both multi-node CPU clusters and multi-GPU environments transparently. This package paves a way to developing high-performance statistical software in various HPC environments simultaneously. As a demonstration of the transparency and scalability of the package, we provide applications to large-scale nonnegative matrix factorization, multidimensional scaling, and $\ell_1$-regularized Cox proportional hazards model on an 8-GPU workstation and a 720-CPU-core virtual cluster in Amazon Web Services (AWS) cloud. As a case in point, we analyze the on-set of type-2 diabetes from the UK Biobank with 400,000 subjects and 500,000 single nucleotide polymorphisms using the $\ell_1$-regularized Cox proportional hazards model. Fitting a half-million-variate regression model took less than 50 minutes on AWS.

研究动机与目标

• 为解决在统计计算中，针对多 CPU 集群和多 GPU 系统等不同高性能计算（HPC）环境分别编写低级别代码的挑战。
• 通过单一高层接口，统一异构 HPC 环境（CPU 集群、多 GPU 工作站和基于云的虚拟集群）的编程范式。
• 在多种 HPC 平台中，实现广泛使用的统计优化算法（如非负矩阵分解、多维缩放和 ℓ₁-正则化 Cox 回归）的高效、可扩展执行。
• 证明 Julia 的多分派和原生并行机制在大规模数据统计工作负载中，可超越现有基于 Python 的 HPC 工具。
• 支持分析从 TB 到 PB 级别的数据集（如全基因组关联研究），使用统一的高性能代码。

提出的方法

• 该软件包实现了一种分布式数组数据结构，其底层基于数组类型（如 CPU 数组、CuArrays）进行抽象，能自动适配目标 HPC 环境。
• 利用 Julia 的多分派和内置向量化能力，为跨分布式数据的元素级操作和线性代数运算提供高层、表达性强的语法。
• 该软件包使用 MPI 作为底层通信层，支持在任何 MPI 支持的环境中运行，包括基于云的虚拟集群和本地集群。
• 通过透明的分派机制，同时支持 CPU 和 GPU 执行路径：用户通过配置切换后端，无需修改代码。
• 在 GPU 加速方面，集成 CUDA.jl，并使用 CUDA 内核加速计算密集型操作（如 ℓ₁-正则化回归中的矩阵-向量乘积）。
• 包含针对操作如 P_{(n+1)}δ 的优化、内存高效的内核，CPU 上使用 AVX 指令集，GPU 上使用 CUDA 内核。

实验结果

研究问题

• RQ1 单一高层编程接口能否统一异构 HPC 环境（包括多 CPU 集群和多 GPU 系统）中的统计计算？
• RQ2 基于 Julia 的分布式数组库在统计优化任务中的性能，与现有基于 Python 的 HPC 工具相比如何？
• RQ3 统一代码库在多大程度上可扩展至 TB 级基因组数据集（如 400,000 名受试者和 500,000 个 SNP），使用多变量生存模型？
• RQ4 该抽象能否降低为统计算法分别维护 GPU 和 CPU 代码库的开发负担？
• RQ5 当结合 MPI 和 GPU 加速时，高级语言 Julia 在大规模统计推断任务中的性能上限是什么？

主要发现

• 在包含 400,000 × 500,000 个基因组数据的 500,000 个 SNP 数据集上，ℓ₁-正则化 Cox 比例风险模型在 20 个节点的 AWS 虚拟集群（720 个 CPU 核心）上，于 50 分钟内成功拟合完成。
• 相同分析在 CPU 上相比基于 Python 的 dist_stat 实现提速 20%–30%，归因于更优的线程控制和 Julia 更灵活的线程模型。
• 在 8 块 GPU 的工作站上，尽管使用未优化的 GPU 内核，DistStat.jl 仍实现了与 dist_stat 相当的性能，展现出显著的 GPU 加速潜力。
• 该软件包成功实现了迄今最大规模的全基因组数据多变量生存模型分析，受试者数量是先前工作的两倍。
• 分布式数组抽象实现了在 CPU 和 GPU 后端之间透明切换，仅需极少配置更改，证明了在统计计算中统一 HPC 编程的可行性。
• 使用 AVX 和 CUDA 实现的 P_{(n+1)}δ 内存高效内核，显著提升了性能，尤其在 CPU 和 GPU 优化的线性代数操作中表现突出。