Advancing Human Security with AI: A Collaborative Effort in Training Large Foundation Models

Dalton Lunga

Group Lead, GeoAI

AI for Science - NASA SMD Workshop 2024

Philipe Dias, Abhishek Potnis, Jacob Arndt, Jordan Bowman, Lexie Yang, Ben Swan, Aristeidis Tsaris, Dan Lu, Feiyi Wang, Prasanna Balaprakash

ORNL is managed by UT-Battelle LLC for the US Department of Energy

1

1

Resources needed to train LLM-like models

• GPT like models grows linearly with #Layers and quadratically with #Hidden dimensions
• Total parameters is roughly 12*L*d²
• Training one parameter needs 20 Bytes of memory and 120 FLOPs
• Need a minimum of 8, 55, and 313 MI250X GCDs to fit models of size 22B, 175B, and 1T

​

​

​

One layer has 11d² parameters

2

2

Optimizing distributed training strategies

• Ported SOTA LLM training frameworks and libraries to Frontier (Megatron-DeepSpeed, FlashAttention 1&2)
• GPT model can be trained with a combination of Tensor (TP), Pipeline (PP), and Data parallelism (DP)
• TP slices model horizontally, PP vertically, and DP replicates
• TP should limit within node, PP should use large #micro-batches

​

​

​

​

​

​

​

3

3

Training LLMs with 1 trillion parameters

• Automatically identify the best distribution strategies to achieve high GPU throughput and scaling efficiency
• Achieved 100% weak scaling efficiency and 87.05% strong scaling efficiency at 3072 GPUs.
• Trained the models for 10-100 iterations instead of training till completion.

​

https://arxiv.org/abs/2312.12705

4

4

Frontier first indicates path�for energy efficient scaling

• First time in open science setting and on non-NVIDIA hardware
• So far, large runs on only NVIDIA hardware
• No heroic effort in software engineering
• Largely dependent on existing software like Megatron-DeepSpeed
• Train massive LLMs with good efficiency using an entirely NVIDIA-free supercomputer
• Don't need an army of researchers at a national lab to make productive use of AMD GPUs for training large LLMs

Business Sensitive

5

5

Unique aspects of Earth observation image modalities

6

6

Quetzal Foundation models�Multimodal GeoAI models: toward rapid characterization of infrastructure and environmental change. Modalities: Optical, Synthetic Aperture Radar, Terrain, Text, GIS Vector data�

3-Billion parameter, 2PB High res. Imagery

Pic: Sun, X. et al. “RingMo: A remote sensing foundation model with masked image modeling”.

IEEE TGRS (2022)

Pic: Wang, Yi, et al. "DeCUR: decoupling common & unique representations for multimodal self-supervision." arXiv preprint arXiv:2309.05300 (2023).

300 million parameter, 2TB Optical med res. + SAR imagery

Flood Mapping

Land Use/Land Cover Segmentation

Glacier Mapping

7

7

Quetzal-HR – Open data

• MAE pretraining with 1M samples (MillionAID)
• ViT configurations up to 3B parameters
• Frontier, Pytorch DDP, 2048 global batch size, 100k iterations
• Semantic Segmentation (fine-tuning)
• more complex decoder (3B requires sharding)
• 64 nodes (512 GPUs), BS=1 for ViT-1B model
• limited gains with limited data

Image classification (linear probing)

Image segmentation (fine-tuning)

Tsaris, A.; Dias, P.; Potnis, A.; Yin, J.; Wang, F.; Lunga, D. “Pretraining Billion-scale Geospatial Foundational Models on Frontier” To be published at IEEE International Workshop on Parallel and Distributed Scientific and Engineering Computing (PDSEC 2024)

8

8

Quetzal-HR: buildings segmentation

9

9

Model scaling

• Multiples works taking place in Remote Sensing
• Contrastive learning, Masked Autoencoders
• But restricted to small scale (model sizes)
• Mostly conducted by academia

An incomplete summary of FMs developed for EO

Dehghani, M., et al. "Scaling vision transformers to 22 billion parameters." ICML 2023.

Zhai, X., et al. "Scaling vision transformers." IEEE/CVF CVPR 2022.

10

10

Data scaling

• ORBITaL-Net (ORNL BFE) vs Ukraine only:
• larger volume, but worse results 🡪 diversification issues
• ORBITaL-Net (ORNL BFE) [1]
• North America, South America, Africa, Asia
• variety of viewing angles, vernacular architecture styles, LU/LC contexts, and atmospheric conditions

Xie, Z. et al. “On data scaling in masked image modeling”. IEEE/CVF CVPR 2023.

Larger models require more data to avoid MIM overfitting

Ineffective to just “dump” a bunch of data

11

11

Data biases

Schmitt, Michael, et al. "There are no data like more data: Datasets for deep learning in earth observation." IEEE Geoscience and Remote Sensing Magazine (2023).

12

12

Dataset needs for pretraining and benchmarking

Schmitt, Michael, et al. "There are no data like more data: Datasets for deep learning in earth observation." IEEE Geoscience and Remote Sensing Magazine (2023).

13

13

Energy-efficient AI foundation model for better climate and weather solution

• ClimaX: an AI foundation model for climate and weather simulations.
• It pretrains the deep vision transformer model on the large set of CMIP6 model simulations using an unsupervised objective and fine-tune it for various weather and climate tasks with a limited labelled data.
• ClimaX has the following strength in comparison to the single-task ML model:
• Integrate big data and knowledge which improves model performance and generalizability
• Use for a wide range of modeling tasks which saves energy and cost compared to the single-task ML
• Its performance scales well with data size, data resolution and model capacity

14

14

ClimaX forecasts weather 72 hours ahead

Surface air temperature

Wind speed

15

15

Superior scaling efficiency of ClimaX on Frontier supercomputer

• It achieves 92% scaling efficiency for synthetic case and 79% scaling efficiency for real case.

Scaling ClimaX model with 100M parameters up to 32 Frontier nodes

Scaling ClimaX model with 1B parameters up to 32 Frontier nodes

• It achieves 87% scaling efficiency;
• Compared to 100M-parameter model, the efficiency of the 1B model drops 70%, mainly due to the increased computing and communication;

• We are using data parallelism, model parallelism, and sequence parallelism to improve accuracy and efficiency for larger data size, finer data resolution and more complex model architecture.

16

16

A new grand challenge: AI security�

17

17

18

18

AI security research – understanding risks and threats�A new field of research at the intersection of AI and Cyber Security Research

AI SECURITY RESEARCH

​

Risks

Threats

Control

Robustness and Reliability

ARTIFICIAL

INTELLIGENCE

RESEARCH

​

Optimization

​

Testing and Validation

​

Scaling

​

Explainability

CYBER

SECURITY

RESEARCH

​

Vulnerability Research

​

Reverse Engineering

​

Cyber Physical Systems

19

19

Understanding performance and risk-benefit trade-offs �With the increasing economic and societal impact of GeoAI Foundation models, it is critical that we measure and understand their performance, threats, and risk-benefit trade-offs

Downstream:�TASK-SPECIFIC AND CONTEXTUAL

Multimodal Geospatial Foundation models: �BENCHMARKING

Downstream:�HUMAN ALIGNMENT

• Unification between task and domain specific benchmarks
• Example: Hugging Face LLM Leaderboard 
• AI2 Reasoning Challenge (25-shot)
• HellaSwag (10-shot)
• MMLU (5-shot)
• TruthfulQA (0-shot)

• Geospatial AI assistants for map products: harmless, helpful to analysts?

​

• Establish standardized transparency practices
• Evaluate with societal considerations in mind: fairness, bias, …
• Recommend better base models for different situations

​

Pic: FLASK: Fine-grained Language Model Evaluation based on

Alignment Skill Sets, arXiv, 2024

20

20