OneScience-Group/PrithviWxC

Prithvi WxC

Model Introduction

Prithvi WxC (Weather and Climate) was proposed by NASA-IMPACT, IBM, and other teams. It is a weather and climate foundation model based on a Vision Transformer (alternating local/global attention with Hiera and MaxViT), supporting forecasting (6-hour-step rollout) and climate simulation (internal error growth).

Paper:Prithvi WxC: Foundation Model for Weather and Climate

https://arxiv.org/abs/2409.13598

Model Description

Prithvi WxC is a deterministic global weather foundation model: it takes atmospheric states at two consecutive 6-hour time steps, optionally with static fields, and outputs the target state. Longer lead times are obtained through autoregressive rollout.This repository is organized from the official NASA-IMPACT/Prithvi-WxC implementation and integrated with the OneScience data loading and training workflow.

Use Cases

Scenario	Description
Global weather and climate foundation model research	Train or fine-tune a Vision Transformer forecasting model on ERA5 data.
Long-horizon autoregressive rollout	Generate medium- to long-range forecasts autoregressively at 6-hour intervals.
Local quick validation	Use synthetic data to check data loading, training, inference, and result scripts.
ModelScope/OneCode execution	Download the model package, install dependencies, and run the scripts directly.
Multi-card training	Launch multi-process training with torchrun.

Usage

1. OneCode Usage

Use the OneCode online environment for intelligent one-click AI4S programming:

Try intelligent one-click AI4S programming

2. Manual Installation and Usage

Hardware Requirements

• GPU or DCU is recommended.
• CPU can be used for imports and small-configuration connectivity validation, but full training and inference are slower.
• DCU users must install DTK beforehand. DTK 25.04.2 or later, or the OneScience-recommended version matching the current cluster, is recommended.
• The paper-level configuration (embed_dim=2560, 25 encoder blocks, 5 decoder blocks, and approximately 2.3 billion parameters) requires substantial GPU memory.

Download the Model Package

hf download OneScience-Group/PrithviWxC --local-dir ./PrithviWxC
cd PrithviWxC

Install the Runtime Environment

DCU Environment

# Activate DTK and CONDA first
conda create -n onescience311 python=3.11 -y
conda activate onescience311
# uv installation is supported
pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai

GPU Environment

# Activate CONDA first
conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
conda activate onescience311
# uv installation is supported
pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai

Training Data

The OneScience community provides ERA5 data for training (the current repository contains complete data slices subject to data-file size limits). Download it with the command below and confirm that the data path in conf/config.yaml is correct:

hf download --repo-type dataset OneScience-Group/ERA5 --local-dir ./data

For a quick workflow validation, run the synthetic data script first:

python scripts/fake_data.py

Note: scripts/fake_data.py generates the [T, C, H, W] HDF5 data required by the two input time steps and generates data/static/static.npy (currently [4, 32, 64]) for training and inference.

Training

Single card：

python scripts/train.py

Multiple cards：

torchrun --nproc_per_node=8 --nnodes=1 --rdzv_id=1000 --rdzv_backend=c10d --max_restarts=0 --master_addr="localhost" --master_port=29500 scripts/train.py

Training outputs:

data/checkpoints/model_bak.pth
data/checkpoints/trloss.npy
data/checkpoints/valoss.npy

Training Weights

The weight/ folder is reserved for model weights. Official weights with approximately 2.3 billion parameters (such as PrithviWxC_160_13b_2t_0p5d_v1.pt) are published on Hugging Face, but their structure differs from this repository's small configuration. Align the channel count and grid size before loading; weights are not provided by default, and users may train the model using the paper configuration.

Inference

Inference reads data/checkpoints/model_bak.pth：

python scripts/inference.py

Prediction results are written to:

result/output/

Evaluation and Visualization

python scripts/result.py

Outputs include：

• result/rmse.npy
• result/acc.npy
• result/loss.png
• Forecast comparison plots for the specified date and variables

Official Source and Reproduction Notes

• The model implementation comes from the official NASA-IMPACT/Prithvi-WxC (MIT License). The official implementation is embedded unchanged in model/prithvi_wxc_official.py; model/prithvi_wxc.py is only a YAML-driven thin wrapper (with identity normalization parameters for small-configuration connectivity validation).
• Commit fetched for the current case directory: 79dabfcd17abe77e2d5c696707c0164a04f2ec01 (2026-02-05).
• conf/config.yaml uses a small configuration (embed_dim=32, n_blocks_encoder=1, n_blocks_decoder=1) for connectivity validation by default; paper-level reproduction requires a 0.5°×0.625° grid, 160 channels, embed_dim=2560, and 13+12 encoder blocks/3+2 decoder blocks as specified in the paper.
• The following details are not disclosed in the paper and are assumptions in this reproduction:data normalization statistics (identity normalization is currently used; real statistics will be supplied with the data), masked-training details and pretraining schedule, and some hyperparameters (such as the relative positional encoding implementation).

Official OneScience Information

Platform	OneScience Main Repository	Skills Repository
Gitee	https://gitee.com/onescience-ai/onescience	https://gitee.com/onescience-ai/oneskills
GitHub	https://github.com/onescience-ai/OneScience	https://github.com/onescience-ai/oneskills

Citation and License

• This repository is an independent organization and adaptation of Prithvi WxC. The model source is based on the official NASA-IMPACT/Prithvi-WxC implementation by Schmude et al. (2024) and follows the MIT License.
• Please cite:Schmude et al. Prithvi WxC: Foundation Model for Weather and Climate. arXiv:2409.13598, 2024.