DOCSEARCHbyAlgolia
Submit Docs

google-deepmind/weathernext

repository·main·Indexed 27 days ago

https://github.com/google-deepmind/weathernext

Implementation and example code for Google DeepMind's GraphCast and GenCast weather forecasting models. The repository provides pretrained weights, data utilities, and notebooks for medium-range weather forecasting simulations. It includes detailed guides for running models on Google Cloud TPU and GPU, memory requirements for different model resolutions (0.25deg and 1deg), and instructions for configuring GPU inference using SparseTransformerConfig.

Tokens
6.8K
Snippets
24
Records
29
Agent score
92%

What's inside weathernext

  1. Transfer predictions from TPU to Cloud Storage

    main

    After running inference, save your predictions to the TPU filesystem and then upload them to your Cloud Storage bucket.

    1. Save to TPU: In the notebook, run:
      predictions.to_zarr("predictions.zarr")
    2. Upload to Bucket: In the command line, use gcloud storage cp to move the directory to your bucket.
    gcloud storage cp --recursive predictions.zarr gs://<bucket_name>/
    gcloud storage cp --recursive predictions.zarr gs://<bucket_name>/

  2. Configure GenCast for GPU inference

    main

    Because splash attention is currently not supported for GPU in JAX, you must run inference using a model with a specific SparseTransformerConfig. To use a pre-trained checkpoint, you must load the model and manually override the attention_type and mask_type settings.

    Required Configuration:

    • attention_type: "triblockdiag_mha"
    • mask_type: "full"
    with ... as f:
    ckpt = checkpoint.load(f, gencast.CheckPoint)
    ...
    denoiser_architecture_config = ckpt.denoiser_architecture_config
    denoiser_architecture_config.sparse_transformer_config.attention_type = "triblockdiag_mha"
    denoiser_architecture_config.sparse_transformer_config.mask_type = "full"

  3. Provision a Cloud VM TPU via gcloud CLI

    main

    You can provision TPU resources using the gcloud CLI. Use the following patterns depending on whether you are requesting a Spot TPU or using the queuing system.

    Requesting a Spot TPU

    To request a Spot device (which offers significant discounts but can be preempted), append the --spot flag. If enabling queuing, use queued-resources create.

    Example (Queued Spot TPU):

    gcloud compute tpus queued-resources create node-1 --node-id=node-1 --zone=us-south1-a --accelerator-type=v5litepod-4 --runtime-version=v2-tpuv5-litepod --spot

    Configuration Summary for GenCast

    ModelRecommended AcceleratorZoneSoftware Version
    GenCast 0p25deg (Oper)v5p-Nus-east5-a or europe-west4-bv2-alpha-tpuv5
    GenCast 1deg (Mini)v5litepod-Nus-south1-a or asia-southeast1-bv2-tpuv5-litepod

  4. Run GraphCast demo in Colab

    main

    To start with GraphCast, use the graphcast_demo.ipynb notebook in Colaboratory. This provides an example of:

    • Loading data
    • Generating random weights or loading a pre-trained snapshot
    • Generating predictions
    • Computing loss and gradients
    https://colab.research.google.com/github/deepmind/graphcast/blob/master/graphcast_demo.ipynb

  5. Set up a Cloud Storage bucket for GenCast predictions

    main

    To store GenCast predictions, create a bucket and grant the Cloud TPU service account permission to write to it.

    1. Create Bucket: Use a location corresponding to your TPU zone (e.g., us-south1-a maps to US-SOUTH-1).
    2. Identify TPU Service Account: Create the service identity and note the returned service account name (format: service-<project_number>@cloud-tpu.iam.gserviceaccount.com).
    3. Grant Permissions: Assign the roles/storage.objectCreator role to the service account for the bucket.
    # Create bucket
gcloud storage buckets create gs://<bucket_name> --location <location>

# Create TPU service account
gcloud beta services identity create --service tpu.googleapis.com --project <project>

# Grant writing permissions
gcloud storage buckets add-iam-policy-binding gs://<bucket_name> --member=serviceAccount:<service_account> --role=roles/storage.objectCreator
    gcloud storage buckets create gs://<bucket_name> --location <location>

  6. Access pretrained models and data on Google Cloud

    main

    Pretrained model weights, normalization statistics, and example input data for both GraphCast and GenCast are hosted on a Google Cloud Bucket.

    To perform full model training, you must download the ERA5 dataset from ECMWF. For operational fine-tuning, you can access Weatherbench2's HRES 0th frame data.

    https://console.cloud.google.com/storage/browser/dm_graphcast

  7. Run GenCast Mini demo in Colab

    main

    The fastest way to start with GenCast is to use the GenCast 1p0deg Mini <2019 model, which is designed for low-cost demonstrations and can run in a free Colab notebook.

    Open the gencast_mini_demo.ipynb notebook to learn how to:

    • Load data
    • Generate random weights or load a GenCast 1p0deg Mini <2019 snapshot
    • Generate predictions
    • Compute loss and gradients
    https://colab.research.google.com/github/deepmind/graphcast/blob/master/gencast_mini_demo.ipynb

  8. Prepare an active Cloud VM TPU

    main

    To prepare a Cloud VM TPU for GenCast inference, follow these steps:

    1. SSH and Port Forward: Connect to your TPU VM and forward port 8081 to your local device to allow Colab to connect to the remote Jupyter server.
      gcloud compute tpus tpu-vm ssh --zone <zone> <name> --project <project> -- -L 8081:localhost:8081
    2. Transfer Required Files: Since mounting Google Drive or Cloud Buckets is not supported in local runtimes, you must manually copy datasets and parameters to the VM using gcloud storage cp.
    3. Install Dependencies: Install JAX with TPU support and Jupyter to enable the runtime and server.

    Example for a 30-step rollout of the 1-degree model:

    # Transfer files
gcloud storage cp gs://dm_graphcast/gencast/dataset/source-era5_date-2019-03-29_res-1.0_levels-13_steps-30.nc .
gcloud storage cp "gs://dm_graphcast/gencast/params/GenCast 1p0deg <2019.npz" .
gcloud storage cp --recursive gs://dm_graphcast/gencast/stats/ .

# Install software
pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
pip install jupyter
    gcloud compute tpus tpu-vm ssh --zone <zone> <name> --project <project> -- -L 8081:localhost:8081

  9. Connect Colab notebook to a Cloud VM TPU

    main

    To run gencast_demo_cloud_vm.ipynb using the remote TPU compute:

    1. Start Jupyter Server: On the TPU VM (while connected via SSH), launch the notebook server allowing Colab origins and binding to port 8081.
      python3 -m notebook --NotebookApp.allow_origin='https://colab.research.google.com' --port=8081 --NotebookApp.port_retries=0 --no-browser
    2. Get the URL: Copy the http://localhost:8081/... URL generated by the command (including the token).
    3. Connect in Colab: In the gencast_demo_cloud_vm.ipynb notebook, use the 'Connect to a local runtime' option and paste the URL.
    python3 -m notebook --NotebookApp.allow_origin='https://colab.research.google.com' --port=8081 --NotebookApp.port_retries=0 --no-browser

  10. Run GenCast on Google Cloud TPU or GPU

    main

    For running larger GenCast models (0.25deg resolution), follow these instructions:

    1. TPU VM: Use cloud_vm_setup.md for detailed instructions on launching a Google Cloud TPU VM. You can then run models via gencast_demo_cloud_vm.ipynb in Colaboratory.
    2. GPU: To run GenCast on a GPU, you must use a different attention implementation as described in the cloud_vm_setup.md guide.
    https://github.com/google-deepmind/graphcast/blob/main/docs/cloud_vm_setup.md

  11. Construct a wrapped GenCast predictor

    main

    A full GenCast predictor pipeline involves wrapping the base gencast.GenCast model with normalization and NaN cleaning layers.

    1. Initialize gencast.GenCast with the loaded configs.
    2. Wrap with normalization.InputsAndResiduals using provided statistics (diffs_stddev_by_level, mean_by_level, stddev_by_level).
    3. Wrap with nan_cleaning.NaNCleaner to handle missing values (e.g., for sea_surface_temperature) using min_by_level.
    def construct_wrapped_gencast():
  """Constructs and wraps the GenCast Predictor."""
  predictor = gencast.GenCast(
      sampler_config=sampler_config,
      task_config=task_config,
      denoiser_architecture_config=denoiser_architecture_config,
      noise_config=noise_config,
      noise_encoder_config=noise_encoder_config,
  )

predictor = normalization.InputsAndResiduals(
      predictor,
      diffs_stddev_by_level=diffs_stddev_by_level,
      mean_by_level=mean_by_level,
      stddev_by_level=stddev_by_level,
  )

predictor = nan_cleaning.NaNCleaner(
      predictor=predictor,
      reintroduce_nans=True,
      fill_value=min_by_level,
      var_to_clean='sea_surface_temperature',
  )

return predictor