Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Pretraining Billion-scale Geospatial Foundational Models on Frontier

Aristeidis Tsaris, Philipe Ambrozio Dias, Abhishek Potnis, Junqi Yin, Feiyi Wang, Dalton Lunga

TL;DR

The paper tackles the challenge of scaling geospatial foundation models by pretraining billion-parameter Vision Transformer models on publicly available remote sensing data using the Frontier HPC system. It provides a practical, PyTorch FSDP-based workflow for distributed training, analyzes bottlenecks, and evaluates a range of ViT architectures from base to multi-billion parameter sizes. Through MAE-based self-supervised pretraining on MillionAID and linear probing on UCM, AID, NWPU-RESISC45, and MillionAID downstream splits, the authors demonstrate up to a 30% gain in top1 accuracy over smaller baselines, establishing that geospatial FMs at scale can yield meaningful gains. The work offers concrete guidance on system configurations, sharding strategies, and throughput optimizations, highlighting the compute-communication trade-offs and setting a foundation for broader downstream tasks in remote sensing and geospatial analysis.

Abstract

As AI workloads increase in scope, generalization capability becomes challenging for small task-specific models and their demand for large amounts of labeled training samples increases. On the contrary, Foundation Models (FMs) are trained with internet-scale unlabeled data via self-supervised learning and have been shown to adapt to various tasks with minimal fine-tuning. Although large FMs have demonstrated significant impact in natural language processing and computer vision, efforts toward FMs for geospatial applications have been restricted to smaller size models, as pretraining larger models requires very large computing resources equipped with state-of-the-art hardware accelerators. Current satellite constellations collect 100+TBs of data a day, resulting in images that are billions of pixels and multimodal in nature. Such geospatial data poses unique challenges opening up new opportunities to develop FMs. We investigate billion scale FMs and HPC training profiles for geospatial applications by pretraining on publicly available data. We studied from end-to-end the performance and impact in the solution by scaling the model size. Our larger 3B parameter size model achieves up to 30% improvement in top1 scene classification accuracy when comparing a 100M parameter model. Moreover, we detail performance experiments on the Frontier supercomputer, America's first exascale system, where we study different model and data parallel approaches using PyTorch's Fully Sharded Data Parallel library. Specifically, we study variants of the Vision Transformer architecture (ViT), conducting performance analysis for ViT models with size up to 15B parameters. By discussing throughput and performance bottlenecks under different parallelism configurations, we offer insights on how to leverage such leadership-class HPC resources when developing large models for geospatial imagery applications.

Pretraining Billion-scale Geospatial Foundational Models on Frontier

TL;DR

The paper tackles the challenge of scaling geospatial foundation models by pretraining billion-parameter Vision Transformer models on publicly available remote sensing data using the Frontier HPC system. It provides a practical, PyTorch FSDP-based workflow for distributed training, analyzes bottlenecks, and evaluates a range of ViT architectures from base to multi-billion parameter sizes. Through MAE-based self-supervised pretraining on MillionAID and linear probing on UCM, AID, NWPU-RESISC45, and MillionAID downstream splits, the authors demonstrate up to a 30% gain in top1 accuracy over smaller baselines, establishing that geospatial FMs at scale can yield meaningful gains. The work offers concrete guidance on system configurations, sharding strategies, and throughput optimizations, highlighting the compute-communication trade-offs and setting a foundation for broader downstream tasks in remote sensing and geospatial analysis.

Abstract

As AI workloads increase in scope, generalization capability becomes challenging for small task-specific models and their demand for large amounts of labeled training samples increases. On the contrary, Foundation Models (FMs) are trained with internet-scale unlabeled data via self-supervised learning and have been shown to adapt to various tasks with minimal fine-tuning. Although large FMs have demonstrated significant impact in natural language processing and computer vision, efforts toward FMs for geospatial applications have been restricted to smaller size models, as pretraining larger models requires very large computing resources equipped with state-of-the-art hardware accelerators. Current satellite constellations collect 100+TBs of data a day, resulting in images that are billions of pixels and multimodal in nature. Such geospatial data poses unique challenges opening up new opportunities to develop FMs. We investigate billion scale FMs and HPC training profiles for geospatial applications by pretraining on publicly available data. We studied from end-to-end the performance and impact in the solution by scaling the model size. Our larger 3B parameter size model achieves up to 30% improvement in top1 scene classification accuracy when comparing a 100M parameter model. Moreover, we detail performance experiments on the Frontier supercomputer, America's first exascale system, where we study different model and data parallel approaches using PyTorch's Fully Sharded Data Parallel library. Specifically, we study variants of the Vision Transformer architecture (ViT), conducting performance analysis for ViT models with size up to 15B parameters. By discussing throughput and performance bottlenecks under different parallelism configurations, we offer insights on how to leverage such leadership-class HPC resources when developing large models for geospatial imagery applications.
Paper Structure (17 sections, 6 figures, 3 tables)

This paper contains 17 sections, 6 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Background and Related Work
  3. Experimental Setup
  4. Model Architecture Variants
  5. Hardware and Software
  6. Sharding Billion Scale Models
  7. Performance Evaluation: Scaling Model Size
  8. MAE Workload Bottlenecks
  9. Communication Optimizations with FSDP
  10. Weak scaling: ViT-Base to ViT-3B
  11. Sharding Strategies for Larger ViT Models
  12. Performance Observations
  13. Downstream evaluation as model scales
  14. Datasets
  15. Pretraining
  16. ...and 2 more sections

Figures (6)

  • Figure 1: Weak scaling plot of the MAE 3B parameter model. The solid blue line shows the measured average image-per-second (ips) of the real application, while the dash blue line (ideal) shows the ips drawn from an ideal linear scaling scenario without any additional costs as we scale. The "syn" plot, are runs on synthetic data, represent the compute and communication performance, while the "syn no comm" is synthetic on runs as well but without any communication. The "IO" plot shows the IO performance of the application. All hyper-parameters are kept constants for the tests, with a local batch-size 32 with "NO_SHARD" FSDP strategy.
  • Figure 2: Average image-per-second (ips) for the ViT-5B model architecture as a function of three sharding strategies for various FSDP configurations. All hyper-parameters are kept constants for the tests, with a local batch-size 32 on 8 Frontier nodes.
  • Figure 3: The plots show the weak scaling for four model architectures: base (top left), huge (top center), 1B (top right) and 3B (bottom left). All four models can fit on a single GPU on Frontier. Also, the bottom center plot shows the memory usage of the FULL_SHARD mode, and the bottom right plot for the rest of the FSDP modes. The measured average image-per-second (ips) has been measured with a local batch-size of 32 for different FSDP sharding strategies as well as the distributed-data-parallel (DDP) strategy. The dash lines (ideal) show the ips drawn from an ideal linear scaling scenario without any additional costs as we scale. The memory usage per GPU for the DDP and HYBRID is constant, as we use more nodes, while the FULL sharding strategy is not.
  • Figure 4: The plots show the weak scaling of the 5B (top left), and 15B (top center) model architectures. The 5B does not fit on a single GPU on Frontier, while the 15B needs at least 4 GPUs to fit. Also, the top right and the bottom left plots show the memory usage of the HYBRID and FULL_SHARD FSDP modes respectively, for the two model architectures. The bottom left plot shows the GPU power, memory, and utilization trace for the different sharding strategies, on 32 node runs, for the 5B model using the rocm-smi utility. The average image-per-second (ips) is measured with a local batch-size of 32 for different FSDP sharding strategies. The dash lines (ideal) show the ips drawn from an ideal linear scaling scenario without any additional costs as we scale. The memory usage for HYBRID schemes is constant as we use more nodes, while the for FULL_SHARD and SHARD_GRAD it is not.
  • Figure 5: The MAE pretraining loss of the four ViT models, base, huge, 1B, and 3B, as a function of steps, for an approximate total of 100 epochs.
  • ...and 1 more figures