D-CHAG：解决多通道模型训练挑战的新方法

美国能源部橡树岭国家实验室的科学家们提出了一种创新方法——分布式跨通道分层聚合（D-CHAG），该方法通过分布式处理和分层策略，解决了极大规模模型在多通道数据集上的运行难题。

基础科学模型在推动科学发现与创新方面潜力巨大，这主要得益于其能够聚合多样化来源的图像数据，并利用 Transformer 架构学习时空相关性。然而，图像的 token 化与聚合过程计算成本高昂，现有分布式方法如张量并行（TP）、序列并行（SP）或数据并行（DP）尚未充分应对这一挑战。

在此背景下，橡树岭国家实验室的研究团队引入了 D-CHAG 方法，该方案对 token 化进行分布式处理，并通过分层策略优化通道聚合，从而支持极大规模模型在多通道数据集上的高效运行。在高光谱成像与天气预测任务中，D-CHAG 与张量并行和模型分片结合，在 Frontier 超级计算机上实现了最多 75% 的内存占用降低，并在 1,024 块 AMD GPU 上实现了超过 2 倍的持续吞吐量提升。

相关研究成果以「Distributed Cross-Channel Hierarchical Aggregation for Foundation Models」为题，已在 SC25 上发表。

研究亮点：

* D-CHAG 解决了多通道基础模型训练中的内存瓶颈和计算效率问题

* 与仅使用 TP 相比，D-CHAG 可实现最高 70% 的内存占用降低，助力大规模模型训练

* D-CHAG 在天气预测与高光谱植物图像掩码预测任务中展现了卓越性能

论文地址：https://dl.acm.org/doi/10.1145/3712285.3759870

两类典型多通道数据集的应用验证

本研究采用两类典型多通道数据集验证了 D-CHAG 方法的有效性：植物高光谱图像（Hyperspectral Images）和气象 ERA5 数据集。

植物高光谱图像数据由 Oak Ridge National Laboratory（ORNL）高级植物表型实验室（APPL）收集，包含 494 张杨树高光谱图像，每张图像包含 500 个光谱通道，覆盖波长从 400nm 到 900nm。该数据集主要用于生物质研究，是植物表型分析和生物能源研究的重要资源。

此外，研究团队在气象预测实验中使用了 ERA5 高分辨率再分析数据集，选取了 5 个大气层变量和 3 个地表层变量，生成 80 个输入通道。这些数据被重网格化为 5.625°，用于模型训练。

D-CHAG：结合分布式 Token 化与层级聚合

D-CHAG 方法融合了分布式 token 化与层级跨通道聚合两种独立方法。具体而言：

分布式 token 化方法

在前向传播过程中，每个 TP rank 仅对输入通道的子集进行 token 化。在通道聚合前执行 AllGather 操作以实现跨通道注意力。该方法理论上能降低每块 GPU 的 token 化计算开销。

层级跨通道聚合

通过减少每层的通道数量，层级跨通道聚合能显著降低内存占用。然而，增加层数会导致模型规模增大和内存使用增加。D-CHAG 将这两种方法结合，通过分布式方式优化：

D-CHAG 方法示意图

D-CHAG 方法结合了分布式 token 化和层级通道聚合的优势，同时缓解其不足。通过分布式方式将层级通道聚合分布到各个 TP rank 上，减少 AllGather 通信开销；同时增加模型深度保留每层处理较少通道的优势。

成果：D-CHAG 支持高通道数数据集上的更大模型训练

在构建 D-CHAG 后，研究人员评估了其在高光谱成像与天气预测任务上的表现：

模型性能分析

图中展示了不同部分通道聚合模块配置下的性能表现。结果显示，对于 512 通道数据，使用单层 cross-attention 层的性能略低于基线，但对 1024 通道数据可提升约 60%。随着层次结构加深，性能提升更为显著。

植物高光谱图像的自监督掩码预测

结果显示，单 GPU 实现与 D-CHAG 方法在训练损失表现上高度一致。

天气预测

在天气预测实验中，D-CHAG 方法在训练损失及三个测试变量的 RMSE 方面表现优异。

随模型规模扩展的性能

结果显示，对于不同规模的模型，D-CHAG 方法均能实现显著的性能提升。

...