新型材料纳米成像 (XANI) 的加速 X 射线分析

大规模 X 射线自由电子激光（XFEL）能够跟踪新型系统中的结构和电子动力学，包括聚变材料、半导体、电池和催化。它产生超短 X 射线脉冲，可以记录原子和电子的运动。这些仪器能够探测由缺陷和其他影响引起的材料结构中最微小的变化。这些明亮 X 射线脉冲的高重复频率可达到每秒 100 万次拍摄，并配备 3500 万像素相机。

采集到的多维数据集包含关于电子和原子最快微观运动的丰富物理信息，可帮助识别材料中的缺陷。按照传统方法，处理和分析这些数据集以提取物理信息通常需要超过九个月的计算时间。

XFEL 研究设施包括瑞士的 SwissFEL、日本的 Spring-8 Angstrom Compact free-electron Laser（SACLA）、SLAC 的 Linac Coherent Light Source（LCLS-II）、德国的 European XFEL，以及韩国的 Pohang Accelerator Laboratory（PAL）。

本文重点介绍 Accelerated X-ray Analysis for Nanoscale Imaging（XANI）工作流程的新技术突破。NVIDIA 团队展示了其在量子材料表征中的应用，用于从超快飞秒激光泵浦/硬 X 射线探测实验中重建声子色散。

具体而言，该团队加速了 XANI 工作流，并将处理和分析 42 太字节（TB）数据的计算时间从九个月缩短到在 32 个 NVIDIA GB200 Grace Blackwell Superchips 上不到四小时，同时保持所获取数据的相同精度。XANI 项目已被从量子物理到材料化学等不同领域的社区采用，展示了 CUDA Python 和分布式计算加速科学发现的能力。

单节点 Python 在百亿亿次级科学计算中面临哪些挑战？

大规模 XFEL 设施可以以高达兆赫（MHz）的速率运行，并生成数百 TB 到 PB 级的数据。如此海量的数据必须实时处理和分析，以指导科学实验并加速发现。

传统的受 CPU 限制的流水线需要大量手动参数调优和子采样，在实验期间通常只处理数据集的 10%。对于量子材料中新相的高分辨率成像，非线性拟合和 3D 重建的计算成本此前使分析只能放到实验后阶段进行。一次实验可能需要九个月的计算时间。

XANI 如何加速数值计算和 I/O 性能？

基于最初向量化的 NumPy 和 SciPy，NVIDIA 团队在 GB200 Grace Blackwell Superchip 的单个 GPU 上将 XANI 工作流加速了 43 倍，在 64 个 GPU 上加速了 1,000 倍。因此，处理和分析 42 TB 数据的计算时间缩短到不到四小时，同时保持所采集数据的相同精度。

为实现这一改进，开发了新的 cuPyNumeric 库，包括 LMFIT 和多线程 Hierarchical Data Format 5 (HDF5)。这些库进一步提升了 GPU 在数值计算中的利用率，并通过 GPUDirect Storage (GDS) 和多线程 HDF5 将 I/O 吞吐量加速了 165 倍。

XANI 架构有哪些优势？

XANI 使用 cuPyNumeric，促进从 CPU 编排的工作流迁移到以 GPU 为中心的分布式模型。通过最大限度缩短高分辨率 X 射线材料表征的求解时间，这种方法可实现实时反馈和自动化实验引导。图 2 展示了 XANI 工作流编排。新的基于非线性最小二乘最小化和曲线拟合的实现，直接最小化具有物理动机的阻尼余弦模型的最小二乘残差，从而获得更低的拟合残差和逐像素频率细化。

面向复杂多维数据的高吞吐量 I/O 与 GPUDirect Storage

XANI 应用以 HDF5 格式存储和访问数据。cuPyNumeric 支持将 HDF5 数据集加载到内存中，并将数据分布到在多块 GPU 上运行的并行任务中。通过在最新 GPU 和高性能 Lustre 存储系统上进行大量实验，执行了三项关键优化以实现峰值 I/O 性能：GDS、多线程 HDF5 和数据布局（详情如下）。

实现了最高 165 倍的加速；具体而言，在一个节点和两块 GB200 Grace Blackwell Superchip 上达到 76 GB/秒，在 16 个节点和 32 块 GB200 Grace Blackwell Superchip 上达到 700 GB/秒。

例如：

from legate.io.hdf5 import from_file

data = from_file("/path/to/file.h5", "dataset_name")

图 4 将启用 GDS 的集群上的优化结果与具有相似存储带宽的非 GDS 集群的初始基线进行了比较。对于两种集群配置，读取吞吐量在 16 个节点上实现扩展：一种是优化后的加速计算配置（两个 GB200 Grace Blackwell Superchips，186 GB HBM3e；四个 200 Gb/s 存储 NIC；四个 400 Gb/s 计算 NIC；两个 480 GB RAM，启用 GDS），另一种是基线 NVIDIA H100 配置（八个 NVIDIA H100 80 GB GPU，2 TB RAM，两个 400 Gb/s 存储 NIC，未启用 GDS）。

正文：GDS

GDS 提供了一种新的存储技术，使数据能够绕过主机 CPU 和内存直接读入 GPU。HDF5 通过 vfd-gds 插件支持 GDS，该插件使用 GDS cuFile 库将数据直接读取到 GPU 内存中，绕过主机。这消除了通过主机内存暂存数据并执行单独设备拷贝所带来的开销。与通过 CPU 主机的传统 POSIX 读取相比，GDS 始终表现出更高的吞吐量。

在现代集群上充分利用 GDS 需要调优 cuFile 配置参数，以实现更高的读取并行度。默认的 cuFile 设置通常较为保守，在面向高性能 Lustre 文件系统时会留下大量带宽未被利用。

多线程 HDF5

迄今为止，HDF5 库仍然是单线程的。当应用程序从多个线程发起并行读取调用时，该库会在内部将它们串行化。具体而言：

• cuFile 会将一次 HDF5 读取拆分为多个子请求。然而，同一时间只有一个 HDF5 读取会被激活，这严重限制了有效的 I/O 并发性。
• 由此产生的序列化读取流通常不足以使高性能 Lustre 文件系统达到饱和，而这类文件系统能够持续提供数百 GB/s 的吞吐量。

为了解决这一瓶颈，我们为 HDF5 开发了多线程支持，并将其与 cuPyNumeric 集成。这项工作目前已在一个公开分支中提供，并正在积极准备合并到 HDF5 主分支中。

数据布局

数据应以与其磁盘上的 HDF5 布局一致的方式读取。如果没有这种对齐，对于具有许多维度的数据集，读取会变得小而非顺序，从而严重降低吞吐量。

cuPyNumeric 得到改进，可沿变化最慢的维度对数据进行分块，从而实现高效的多维读取，并兼容标准 HDF5 数据布局。这包括连续布局、分块布局和虚拟数据集（VDS）。我们选择使用包含多个源文件的 VDS 数据布局，以促进更高的读取并行度。同时，以不妨碍优化读取的方式在磁盘中保持数据布局也很重要。

使用 cuPyNumeric 进行分布式计算

cuPyNumeric 通过在集群的聚合内存中划分数组，充当 XANI 的分布式引擎。除了提供 NumPy 和 SciPy API 外，它还作为一个用于大规模分发基于 NumPy 的应用程序的库。

这种设计使高级流水线代码保持简单——一次函数调用即可替代数百次手动任务提交——同时让 cuPyNumeric 能够全面了解数据依赖关系，以便在 GPU 之间进行自动调度和负载均衡。

分布式分区

通过将导入从 numpy 更新为 cupynumeric，运行时会主动构建任务依赖图，以跨节点并行化操作。该库将 NumPy 调用转换为任务，这些任务以异步方式处理数据移动和内核执行。计算工作流将在以下章节中详细说明，包括隐式并行、数据分区和任务重叠（图 4）。

# Transitioning to distributed memory architecture
import cupynumeric as np
from legate.io.hdf5 import from_file

# Array is logically contiguous but physically partitioned across the cluster
raw_frames = num.array(from_file("lcls_detector_data.h5", 'r')) 

# Reductions are performed via distributed tree-reduction algorithms
# This executes across multiple GPUs/nodes without manual MPI code
intensity_map = np.sum(raw_frames, axis=0)

显式并行：用于专用内核的 tiled_task

在 XANI 流水线中，信号检测需要对每个探测器像素拟合阻尼振荡——这是一个非线性最小二乘问题。我们用批量 GPU 实现（JAXfit）替换了基于 SciPy 的顺序求解器，该实现完全在 GPU 上运行 Levenberg-Marquardt（LM）算法，在一次批量求解中同时拟合一个分块中的所有像素。

要将其与 cuPyNumeric 集成，请使用 tiled_task API，该 API 会自动将大型探测器阵列划分到可用的 GPU 上，并为每个空间分块调度一个 GPU 任务。每个任务都会以原生 CuPy 数组的形式接收其对应的探测器数据切片——无需显式数据移动或手动分发/汇总。时间延迟轴会广播到所有分块，而空间维度（高度乘以宽度）会被划分：

@tiled_task(partitioned_dims=("i", "j"), primary_array="onn")
def roi_task_cp(
    # Broadcast: All tiles see the full time axis
    delay: Input(cupy.ndarray, partitioning=None),
    
    # Partitioned: Each tile gets its unique 16×16 slice
    onn: Input(
        cupy.ndarray, 
        partitioning=(None, "i", "j"), 
        tile_dims=(None, 16, 16)
    ),
    
    # Results are written back to the sharded output array per tile
    amp_out: Output(
        cupy.ndarray, 
        partitioning=("i", "j", None), 
        tile_dims=(16, 16, None)
    ),
    ...
):
    # Inside the task, arrays are already local CuPy instances.
    # We can run the batched LM solver directly on the GPU.
    fits, _, freq_lis, amp_lis, ... = jaxfit_linearpred(
        onn, delay, amp_threshold, ...
    )

# A single call triggers the cuPyNumeric runtime to schedule 
# one GPU task per 16×16 tile across the entire cluster.
roi_task_cp(
    onn=detector_on, 
    delay=delay_axis, 
    amp_out=amp_all, 
    tile_shape=(16, 16)
)

图 5 展示了基于隐式数据划分和通信的 cuPyNumeric 分布式执行流程。数据从 NVMe/Lustre 存储异步加载到特定 GPU 的分片中。在涉及跨 GPU 数据传输的 np.sum 归约之后，最终结果会广播回所有节点。请注意，在整个过程中，原始分片数据会保留在每个 GPU 上并可被访问。

用于材料数据处理与分析的 XANI 工作流性能和结果

在 NVIDIA GB200 NVL72 集群上向分布式 Python HPC 栈的过渡，已经改变了 X 射线科学的数据计算规模，使得为材料表征模型的 AI 训练准备海量数据集成为可能。主要技术影响包括计算加速、系统扩展以及 I/O 吞吐优化。

计算加速

使用 32 个 NVIDIA GB200 Grace Blackwell Superchips，处理时间缩短了 1,000 倍（从六个月缩短到不到四小时）。新的批处理 GPU 实现实现了比以往任何结果高约 3 倍的 GPU 利用率，详见 Using Accelerated Computing to Live-Steer Scientific Experiments at Massive Research Facilities。

新的基于 LM 的实现直接最小化一个基于物理动机的阻尼余弦模型的最小二乘残差，从而获得更低的拟合残差和逐像素频率细化。该流水线将探测器划分为 760 个固定大小的图块，并在各个 GPU 上独立处理。每个图块的计算——拟合 16 条像素轨迹——规模太小，无法让现代 GPU 达到饱和，因此执行时间主要由每个图块的固定开销（任务分派、数据移动、设置）主导。超过约 32 个 superchips 后，每个 GPU 分配到的图块数会降至无法摊销这些开销的程度，增加更多 GPU 也不会带来收益。

系统扩展

在 NVIDIA GB200 和 NVIDIA H100 架构上，使用专用的多节点/多 GPU 配置观察到了良好的扩展性。

I/O 吞吐量优化

基于 GDS 的 I/O，结合数据布局感知优化和多流读取，旨在充分利用现代 GPU 和高性能存储的能力。

了解更多

XANI 项目表明，以 Python 为主导的社区——从量子物理到材料化学——能够成功采用 CUDA Python 和分布式计算。通过聚焦于批处理、显式 Python 任务以及 GDS 支持的 I/O，科学流水线可以扩展以满足下一代设施在 PB 级规模上的需求，而不需要科学家成为更底层的 C++/MPI 专家。

查看以下资源以开始使用并了解更多信息：

• 测试新开发的多线程 HDF5 库（测试版）用于你的工作流
• 观看 GTC 会议环节“Accelerated HPC+AI Workflow Enables Live-Steering of Vera C. Rubin Observatory and X-ray Free Electron Laser”
• 如果你有兴趣合作参与 XANI 项目，以加速用于实验引导和科学发现的多维数据处理与分析，请联系我们

致谢

感谢 NVIDIA 贡献者 Malte Foester 和 Quincey Koziol。