D-CHAG：破解大规模模型训练瓶颈的创新方法

源自美国能源部橡树岭国家实验室的科学家们创新性地提出了分布式跨通道分层聚合方法（D-CHAG），该方法对 token 化过程进行分布式处理，并采用分层策略进行通道聚合，使得极大规模模型能够在多通道数据集上高效运行。

基于视觉的科学基础模型在推动科学发现与创新方面潜力巨大，这主要得益于其能够聚合多样化来源的图像数据（例如不同的物理观测场景），并利用 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

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

其中，用于自监督掩码预测的植物高光谱图像数据由 Oak Ridge National Laboratory（ORNL）高级植物表型实验室（APPL）收集。数据集包含 494 张杨树（Poplar）高光谱图像，每张图像包含 500 个光谱通道，覆盖波长从 400nm 到 900nm。这些图像用于掩码自监督训练，模型任务是预测缺失内容，从而学习图像的潜在数据分布。

此外，研究团队在气象预测实验中使用了 ERA5 高分辨率再分析数据集。研究选择了 5 个大气层变量和 3 个地表层变量，生成 80 个输入通道。为了适配模型训练，原始数据被重网格化并采用 xESMF 工具包和双线性插值算法完成。

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

D-CHAG 方法结合了两种独立的方法：分布式 token 化方法和层级跨通道聚合。

分布式 token 化方法

在前向传播过程中，每个 TP rank 仅对输入通道的子集进行 token 化。理论上，该方法能够降低每块 GPU 的 token 化计算开销。

层级跨通道聚合

该方法的主要优势在于减少每层的内存占用，因为每层处理的通道数量更少。然而，增加层数会导致整体模型规模增大、内存使用增加。D-CHAG 通过分布式方式将这两种方法结合起来。

D-CHAG 方法在基础架构上的示意图

具体而言，每个 TP rank 对总通道子集中的二维图像进行 token 化。在前向传播过程中仅需执行一次 AllGather 操作；在反向传播时，只收集每块 GPU 的相关梯度，从而避免额外通信。

成果：D-CHAG 支持更大规模模型的训练

在构建 D-CHAG 后，研究人员对模型性能进行了验证：

模型性能分析

* Tree0 表示部分聚合模块中仅有一层聚合，Tree2 表示两层；

* 后缀 -C 和 -L 表示所用层的类型：-C 中所有层为 cross-attention，-L 中所有层为 linear

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

天气预测

随模型规模扩展的性能

ViT：视觉 AI 从感知模型走向通用视觉基础模型

过去十年，计算机视觉模型主要围绕「单任务优化」展开。然而，随着 Transformer 架构在自然语言领域催生出 GPT、BERT 等基础模型（Foundation Models），视觉领域也正在经历类似的范式转移。Vision Transformer（ViT）被视为视觉基础模型的关键技术基石。