Distributed Cross-Channel Hierarchical Aggregation for Foundation Models

📄 arXiv: 2506.21411v1 📥 PDF

作者: Aristeidis Tsaris, Isaac Lyngaas, John Lagregren, Mohamed Wahib, Larry York, Prasanna Balaprakash, Dan Lu, Feiyi Wang, Xiao Wang

分类: cs.LG

发布日期: 2025-06-26

💡 一句话要点

提出D-CHAG以解决图像通道聚合计算效率低下问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 图像聚合 跨通道聚合 计算效率 高光谱成像 天气预报 视觉变换器 模型并行 深度学习

📋 核心要点

  1. 现有的分布式方法在处理图像聚合时计算效率低，难以满足大规模数据集的需求。
  2. 提出的D-CHAG方法通过跨通道层次聚合，优化了图像数据的处理效率，兼容多种变换器架构。
  3. 在高光谱成像和天气预报任务中，D-CHAG实现了75%的内存减少和超过两倍的吞吐量提升。

📝 摘要（中文）

基于视觉的科学基础模型在推动科学发现和创新方面具有重要潜力。这种潜力源于其能够聚合来自不同来源的图像，并利用变换器架构学习时空相关性。然而，图像的标记和聚合计算密集，现有的分布式方法未能充分解决这一挑战。本文提出了分布式跨通道层次聚合（D-CHAG）方法，旨在处理具有大量通道的图像数据集。该方法兼容任何模型并行策略和视觉变换器架构，显著提高计算效率。我们在高光谱成像和天气预报任务上评估了D-CHAG，结果显示在Frontier超级计算机上，结合张量并行和模型分片，内存使用减少了75%，持续吞吐量提高了两倍以上。

🔬 方法详解

问题定义：本文旨在解决现有分布式方法在图像通道聚合时计算效率低下的问题，尤其是在处理大规模多通道数据集时的挑战。

核心思路：D-CHAG方法通过跨通道层次聚合技术，优化了数据处理流程，能够有效减少计算资源的消耗，并提高数据处理的速度。

技术框架：D-CHAG的整体架构包括数据预处理、层次聚合模块和输出生成模块。数据预处理负责将输入图像进行标准化，层次聚合模块则通过并行处理不同通道的数据，最后输出生成模块将聚合结果整合为最终的模型输入。

关键创新：D-CHAG的主要创新在于其跨通道聚合策略，能够在保持模型性能的同时显著降低内存使用和计算时间，这与传统方法的串行处理方式形成鲜明对比。

关键设计：在设计中，D-CHAG采用了张量并行和模型分片技术，优化了内存管理，并通过调整网络结构和损失函数，确保了模型的高效训练和推理。具体参数设置和网络结构细节在实验部分进行了详细描述。

📊 实验亮点

在Frontier超级计算机上，D-CHAG方法结合张量并行和模型分片技术，实现了高达75%的内存使用减少，并且在处理能力上提升了超过两倍，展现出显著的性能优势。

🎯 应用场景

该研究的潜在应用领域包括高光谱成像、气象预测、遥感技术等，能够为科学研究提供更高效的数据处理工具。未来，D-CHAG方法有望在更广泛的视觉任务中得到应用，推动相关领域的技术进步与创新。

📄 摘要（原文）

Vision-based scientific foundation models hold significant promise for advancing scientific discovery and innovation. This potential stems from their ability to aggregate images from diverse sources such as varying physical groundings or data acquisition systems and to learn spatio-temporal correlations using transformer architectures. However, tokenizing and aggregating images can be compute-intensive, a challenge not fully addressed by current distributed methods. In this work, we introduce the Distributed Cross-Channel Hierarchical Aggregation (D-CHAG) approach designed for datasets with a large number of channels across image modalities. Our method is compatible with any model-parallel strategy and any type of vision transformer architecture, significantly improving computational efficiency. We evaluated D-CHAG on hyperspectral imaging and weather forecasting tasks. When integrated with tensor parallelism and model sharding, our approach achieved up to a 75% reduction in memory usage and more than doubled sustained throughput on up to 1,024 AMD GPUs on the Frontier Supercomputer.