大规模 X 射线自由电子激光 (XFEL) 能够追踪新型系统中的结构和电子动力学，包括聚变材料、半导体、电池和催化剂。它产生的超短 X 射线脉冲可以记录原子和电子的运动。这些仪器可以检测由缺陷和其他影响引起的材料结构的微小变化。使用 3500 万像素的摄像头，这些明亮的 X 射线爆发的高重复率可达到每秒 100 万次。
所获取的多维数据集包含有关电子和原子最快速微观运动的丰富物理信息，有助于识别材料中的缺陷。处理和分析这些数据集以提取物理数据通常需要 9 个月以上的计算时间。
XFEL 研究机构包括瑞士的 SwissFEL、日本的 Spring-8 皮米级紧凑型自由电子激光装置（SACLA）、美国 SLAC 的直线加速器相干光源（LCLS-II）、德国的欧洲 XFEL，以及韩国浦项加速器实验室（PAL）。
本文重点介绍了加速 X 射线分析纳米级成像 (XANI) 工作流程的新技术突破。NVIDIA 团队演示了量子材料的表征，以重建超快飞秒激光泵/ 硬 X 射线探针实验中的声子色散。
具体而言，该团队加速了 XANI 工作流，在 32 个 NVIDIA GB200 Grace Blackwell 超级芯片上，将处理和分析 42TB 数据的计算时间从 9 个月缩短至不到 4 小时，同时保持了原始数据的同等精度。XANI 项目已被量子物理和材料化学等多个领域广泛采用，充分展现了 CUDA Python 与分布式计算在推动科学发现方面的强大能力。
单节点 Python 为百亿亿次级 (Exascale) 科学带来了哪些挑战？
大规模 XFEL 设施可以以高达兆赫 (MHz) 的速率运行，并生成数百 TB 到 PB 的数据。必须实时处理和分析如此庞大的数据量，才能引导科学实验并加速发现。
传统的受 CPU 限制的工作流需要大量的手动参数调整和二次采样，在实验期间通常只处理 10% 的数据集。对于量子材料新相位的高分辨率成像，非线性拟合和 3D 重建的计算成本过去将分析降级到后实验阶段。单个实验可能需要九个月的计算时间。
XANI 如何加速数值计算和 I/ O 性能？
与最初的向量化 NumPy 和 SciPy 相比，NVIDIA 团队在 GB200 Grace Blackwell 超级芯片上将 XANI 工作流的速度提高了 43 倍，在 64 个 GPU 上提高了 1000 倍。因此，处理和分析 42 TB 数据的计算时间缩短到不到 4 小时，同时保留了与所采集数据相同的精度。
为实现这一改进，我们开发了新的 cuPyNumeric 库，集成了 LMFIT 和多线程分层数据格式 5（HDF5）。借助 GPUDirect Storage（GDS） 和多线程 HDF5，这些库可进一步提升数值计算中的 GPU 利用率，并将 I/O 吞吐量加速达 165 倍。
XANI 架构有哪些优势？
使用 cuPyNumeric，XANI 有助于从 CPU 编排的工作流迁移到以 GPU 为中心的分布式模型。这种方法可更大限度地缩短高分辨率 X 射线材料表征的解决方案时间，从而实现实时反馈和自动实验转向。图 2 显示了 XANI 工作流编排。新的非线性最小二乘最小化和基于曲线拟合的实施可直接最小化物理动机阻尼余弦模型的最小二乘残差，从而生成更低的拟合残差和每像素频率优化。
借助 GPUDirect Storage 实现复杂多维数据的高吞吐量 I/ O
XANI 应用程序以 HDF5 格式存储和访问数据。cuPyNumeric 支持将 HDF5 数据集加载到内存中，从而将数据分发到多个 GPU 中运行的并行任务中。通过在最新的 GPU 和高性能 Lustre 存储系统上进行广泛的实验，执行了三项关键优化以实现峰值 I/ O 性能：GDS、多线程 HDF5 和数据布局 (详情如下) 。
速度提升高达 165 倍；具体而言，单节点和两块 GB200 Grace Blackwell 超级芯片可加速 76 GB/s，16 节点和 32 GB200 Grace Blackwell 超级芯片可加速 700 GB/s。
例如：
from legate.io.hdf5 import from_filedata = from_file("/path/to/file.h5", "dataset_name") |
图 4 将启用 GDS 的集群的优化结果与具有相似存储带宽的非 GDS 集群的初始基准进行比较。读取吞吐量可跨 16 个节点扩展，适用于两种集群配置：优化的加速计算设置 (两个 GB200 Grace Blackwell 超级芯片，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/ 秒的速度。
为解决这一瓶颈，我们为 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 architectureimport cupynumeric as npfrom legate.io.hdf5 import from_file# Array is logically contiguous but physically partitioned across the clusterraw_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 codeintensity_map = np.sum(raw_frames, axis=0) |
显式并行：专用内核的 tiled_task
在 XANI 工作流中，信号检测需要为每个检测器像素拟合阻尼振荡，这是一个非线性最小二乘问题。我们将基于 SciPy 的顺序求解器替换为批量 GPU 实现 (JAXfit) ，该实现完全在 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对应的分片中。在np.sum规约操作（涉及跨GPU数据传输）完成后，最终结果被广播回所有节点。需要注意的是，整个过程中，原始分片数据始终保留在各自GPU上，并可由对应GPU直接访问。
用于材料数据处理和分析的 XANI 工作流性能和结果
将分布式 Python 高性能计算堆栈迁移至 NVIDIA GB200 NVL72 集群，显著提升了 X 射线科学领域的数据处理规模，使大规模数据集可用于材料表征模型的 AI 训练。关键技术影响包括计算速度提升、系统可扩展性增强，以及整个优化过程中的 I/O 性能优化。
计算加速
借助 32 个 NVIDIA GB200 Grace Blackwell 超级芯片，处理时间缩短了 1000 倍（从六个月缩短至不到 4 小时）。新的批量 GPU 实现方案达到的 GPU 利用率比以往任何结果高出 3 倍，详情请参阅使用加速计算在大型研究设施中进行实时科学实验。
基于 LM 的新实现可直接更大限度地减少物理动机阻尼余弦模型的最小二乘残差，从而产生更低的拟合残差和每像素频率优化。工作流将检测器划分为 760 个固定大小的图块，在 GPU 上独立处理。每个图块的计算 (拟合 16 像素追踪) 太小，无法满足现代 GPU 的需求，因此执行时间主要由每个图块的固定用度 (任务调度、数据移动、设置) 决定。除了约 32 个超级芯片之外，每个 GPU 的图块数量都会下降到低于此开销摊销点的水平，而额外的 GPU 也无法带来任何好处。
系统扩展
在使用专用多节点/ 多 GPU 配置的 NVIDIA GB200 和 NVIDIA H100 架构中，可观察到良好的扩展性。
I/ O 吞吐量优化
GDS 支持的 I/ O，具有数据布局感知优化和多流读取，旨在使现代 GPU 和高性能存储饱和。
了解详情
XANI 项目表明，Python 主导的社区 (从量子物理学到材料化学) 可以成功采用 CUDA Python 和分布式计算。通过专注于批处理、显式 Python 任务和 GDS 支持的 I/ O，科学工作流可以进行扩展，以满足新一代设施的 PB 级需求，而无需科学家成为低级 C++/ MPI 专家。
查看以下资源以开始使用，并了解更多信息：
- 针对您的工作流程测试新开发的多线程 HDF5 库 (测试版)
- 观看 GTC 会议：加速 HPC + AI 工作流，实现 Vera C. Rubin 天文台和无 X 射线自由电子激光的实时调控
- 联系我们 如果您有兴趣协作开展 XANI 项目，以加速用于实验指导和科学发现的多维数据处理和分析
致谢
感谢 NVIDIA 贡献者 Malte Foester 和 Quincey Koziol。
