内存占用最高降低75%，美国能源部科学家提出跨通道分层聚合方法D-CHAG，实现极大规模模型多通道数据集运行

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

基于视觉的科学基础模型在推动科学发现与创新方面具有巨大潜力，主要源于其能够聚合多样化来源的图像数据（例如不同的物理观测场景），并利用 Transformer 架构学习时空相关性。然而，图像的 token 化与聚合过程计算开销巨大，而现有的分布式方法如张量并行（TP）、序列并行（SP）或数据并行（DP），尚未充分解决这一挑战。

在此背景下，来自美国能源部橡树岭国家实验室的研究人员提出了一种面向基础模型的分布式跨通道分层聚合方法（Distributed Cross-Channel Hierarchical Aggregation, 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 张杨树（Poplar）高光谱图像，每张图像包含 500 个光谱通道，覆盖波长从 400nm 到 900nm。

此数据集主要用于生物质研究，是植物表型分析和生物能源研究的重要资源。这些图像用于掩码自监督训练，即将图像切片作为 token 进行 mask，模型的任务是预测缺失的内容，从而学习图像的潜在数据分布。值得注意的是，该数据集未使用任何预训练权重，完全基于自监督学习进行训练，这也凸显了 D-CHAG 在高通道自监督任务中的适用性。

此外，在气象预测实验中，研究团队使用了 ERA5 高分辨率再分析数据集。研究选择了 5 个大气层变量（位势高度、温度、风速 u 分量、风速 v 分量、比湿度）和 3 个地表层变量（2 米温度、10 米 u 分量风速、10 米 v 分量风速），覆盖超过 10 个压力层，总共生成 80 个输入通道。为了适配模型训练，原始分辨率为 0.25° 的数据（770 × 1440）被重网格化为 5.625°（32 × 64），采用 xESMF 工具包 和双线性插值算法完成。

模型任务是进行未来时间步的气象变量预测，例如 500 hPa 位势高度（Z500）、850 hPa 温度（T850）、10 米 u 分量风速（U10），从而验证 D-CHAG 方法在时间序列预测任务上的性能。

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

简单而言，D-CHAG 方法来自两种独立方法的融合，分别是：

分布式 token 化方法

在前向传播过程中，每个 TP rank 仅对输入通道的子集进行 token 化。在进行通道聚合步骤之前，需要执行一次 AllGather 操作，以便在所有通道之间实现跨通道注意力（cross-attention）。理论上，该方法能够降低每块 GPU 的 token 化计算开销。

层级跨通道聚合

这种方法的主要优势在于每个跨通道注意力层的内存占用减少，因为每层处理的通道数量更少。然而，增加层数会导致整体模型规模增大、内存使用增加。对于通道数量庞大的数据集而言，这种权衡更为有利，因为标准跨通道注意力的二次内存开销更高。

这两种方法虽然各有优势，但也存在一些不足，比如分布式 token 化方法在 TP rank 之间存在较高的通信开销，并未解决通道维度大内存占用的问题；而层级跨通道聚合方法会增加每块 GPU 上的模型参数数量。D-CHAG 方法通过分布式方式将两种方法结合起来，整体架构如下图所示：

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

具体而言，每个 TP rank 对总通道子集中的二维图像进行 token 化。由于每块 GPU 仅持有全部通道的一部分，在这些通道上本地执行通道聚合——该模块称为部分通道聚合模块（partial-channel aggregation module）。在每个 TP rank 内完成通道聚合后，收集输出并使用跨通道注意力进行最终聚合。前向传播过程中仅需执行一次 AllGather 操作；在反向传播时，只收集每块 GPU 的相关梯度，从而避免额外通信。

D-CHAG 方法能够充分利用分布式 token 化和层级通道聚合的优势，同时缓解它们的不足。通过将层级通道聚合分布到 TP rank 上，研究人员将 AllGather 通信减少为每个 TP rank 仅需处理单个通道，在反向传播过程中无需任何通信。此外，通过增加模型深度保留了每层聚合处理通道数量减少的优势，同时通过部分通道聚合模块将额外模型参数分布到各 TP rank 上。

研究对比了两种实现策略：

* D-CHAG-L（Linear Layer）：层级聚合模块使用线性层，内存占用低，适合通道数较多的情况。

* D-CHAG-C（Cross-Attention Layer）：使用交叉注意力层，计算成本较高，但在超大模型或极高通道数时性能提升显著。

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

在构建 D-CHAG 后，研究人员对模型性能进行了验证，然后进一步评估了其在高光谱成像与天气预测任务上的表现：

模型性能分析

下图展示了 D-CHAG 在不同部分通道聚合模块配置下的性能表现：

图中展示了针对 1.7B 参数模型，在不同部分通道聚合模块配置下，每块 GPU 相对于仅使用 TP 基线的性能提升

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

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

结果显示：

对于 512 通道数据，使用单层 cross-attention 层的性能略低于基线，但对 1024 通道数据可提升约 60%。

随着层次结构加深，即便是 512 通道数据，也能获得明显性能提升，而 1024 通道数据的性能保持相对稳定。

使用 linear 层时，即使层次结构较浅，也能在 512 和 1024 通道图像上获得性能提升。实际上，最佳性能出现在 D-CHAG-L-Tree0，即仅包含一层通道聚合层。增加聚合层会增加模型参数，引入额外内存开销。虽然对于 512 通道情况，增加层数似乎有益，但对于两种通道规模，仅使用一层 linear 层的性能优于更深的配置。

D-CHAG-C-Tree0 在两块 GPU 时对性能略有负面影响，但扩展至 8 块 GPU 时可获得 60% 提升。

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

下图比较了基线方法与 D-CHAG 方法在高光谱植物图像掩码自编码器应用中的训练损失，结果显示：在训练过程中，单 GPU 实现与 D-CHAG 方法（在两块 GPU 上运行）的训练损失表现高度一致。

基线方法与 D-CHAG 方法在高光谱植物图像掩码自编码器应用中的训练损失

橡树岭国家实验室分子与细胞成像组的高级研究员拉里·约克表示，D-CHAG 可以帮助植物科学家快速完成诸如直接从图像中测量植物光合作用活性等任务，从而取代费时费力的手动测量。

天气预测

研究人员在 ERA5 数据集上进行 30 天气象预测实验，下图比较了基线方法与 D-CHAG 方法在天气预测应用中的训练损失及三个测试变量的 RMSE：

基线方法与 D-CHAG 方法在天气预测应用中的训练损失及三个测试变量的 RMSE

下表则展示了模型在 7、14 和 30 天预测任务上的最终对比，包括 RMSE、MSE 以及 Pearson 相关系数（即 wACC）

D-CHAG 方法相较于单 GPU 训练在 7、14 和 30 天预测任务中的 MSE、RMSE 及 wACC 的百分比变化（% Δ）

结合图和表总体来看，训练损失与基线模型高度一致，各项指标的偏差极小。

随模型规模扩展的性能

下图显示了 3 种模型规模在需要使用 TP 的通道配置下，D-CHAG 方法相较于仅使用 TP 的性能提升：

D-CHAG 方法结合 TP 的情况下，相较于仅使用 TP 时，7B、15B 和 26B 参数模型每个 GPU 的性能提升情况

结果显示，对于 7B 参数模型，使用部分通道聚合模块中的线性层（linear layers）可获得 30% 至 70% 的性能提升，而使用交叉注意力层（cross-attention layers）可获得 10% 至 60% 的提升；对于 15B 参数模型，性能提升超过 20% 至 50%；而 26B 参数模型的性能提升在 10% 至 30% 之间。

此外，在固定模型规模下，随着通道数增加，性能提升更明显，这是因为在给定架构下，增加通道数不会增加 transformer block 的计算量，但会增加 tokenization 和 channel-aggregation 模块的工作量。

另一方面，仅使用 TP 无法训练 26B 参数、256 通道图像，但使用 D-CHAG 方法时，可以训练 26B 参数、512 通道的模型，仅使用不到 80% 的可用内存——这表明该方法能够支持高通道数数据集上更大模型的训练。

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

过去十年，计算机视觉模型主要围绕「单任务优化」展开——分类、检测、分割、重建各自独立发展。然而，随着 Transformer 架构在自然语言领域催生出 GPT、BERT 等基础模型（Foundation Models），视觉领域也正在经历类似的范式转移：从任务特化模型走向通用视觉基础模型。在这一趋势下，Vision Transformer（ViT）被视为视觉基础模型的关键技术基石。

Vision Transformer（ViT）首次将 Transformer 架构完整引入计算机视觉任务，其核心思想是：将图像视为一系列 patch token 序列，用自注意力机制替代卷积神经网络的局部感受野建模。具体而言，ViT 将输入图像划分为固定大小的 patch，并将每个 patch 映射为 embedding token，然后通过 Transformer Encoder 建模 patch 之间的全局关系。

与传统 CNN 相比，ViT 对科学数据尤其具有优势：适合高维多通道数据（如遥感、医学影像、光谱数据），可处理非欧几里得空间结构（如气候格点、物理场），适用于跨通道建模（不同物理变量之间的耦合关系），这也正是 D-CHAG 论文所关注的核心问题。

除了上文研究中提及的场景，ViT 正在更多场景发挥核心价值。2025 年 3 月，北京大学国际医院皮肤科主任医师韩钢文携其团队开发出一种名为 AcneDGNet 的深度学习算法，这是一种融合视觉 Transformer 与卷积神经网络，能获取更高效的分层特征表，让分级更精准。经前瞻性评估表明，AcneDGNet 的深度学习算法不仅比初级皮肤科医生更准确，而且与高级皮肤科医生的准确性相当，能够在不同的医疗保健场景中同时准确地完成痤疮病变检测并判断严重程度，有效帮助皮肤科医生和患者在在线问诊和线下就医场景中诊断和管理痤疮。

论文标题：Evaluation of an acne lesion detection and severity grading model for Chinese population in online and offline healthcare scenarios

论文地址：https://www.nature.com/articles/s41598-024-84670-z

从产业视角看，Vision Transformer 标志着视觉 AI 从感知模型走向通用视觉基础模型的关键拐点。其统一的 Transformer 架构为跨模态融合、规模化扩展与系统级优化提供了通用底座，使视觉模型成为 AI for Science 的核心基础设施。未来，围绕 ViT 的并行化、内存优化与多通道建模能力，将成为决定视觉基础模型产业落地速度与规模的关键竞争点。

参考文献：

1.https://phys.org/news/2026-01-empowering-ai-foundation.html

2.https://dl.acm.org/doi/10.1145/3712285.3759870

3.https://mp.weixin.qq.com/s/JvKQPbBQFhofqlVX4jLgSA

本文来自微信公众号“HyperAI超神经”，作者：梅菜，36氪经授权发布。