Distributed Cross-Channel Hierarchical Aggregation for Foundation Models

TL;DR

This paper introduces D-CHAG, a scalable method for efficient distributed aggregation in vision-based scientific models, significantly reducing memory and increasing throughput on large GPU clusters.

Contribution

The paper presents D-CHAG, a novel distributed hierarchical aggregation technique compatible with various transformer architectures, improving efficiency for multi-channel image datasets.

Findings

Achieved up to 75% memory reduction.

More than doubled throughput on 1024 GPUs.

Effective for hyperspectral and weather data.

Abstract

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.