From 38f68b4155e506f11b08b2a56b7d5ad03c3af8b0 Mon Sep 17 00:00:00 2001 From: Linlin Cui Date: Wed, 10 Dec 2025 16:03:21 +0000 Subject: [PATCH 1/4] add time logs; remove Bfloat16cast --- oper/run_graphcast.py | 39 ++++++++++++++++++++++++++++++++++----- 1 file changed, 34 insertions(+), 5 deletions(-) diff --git a/oper/run_graphcast.py b/oper/run_graphcast.py index 59fe158..a611b37 100644 --- a/oper/run_graphcast.py +++ b/oper/run_graphcast.py @@ -9,6 +9,7 @@ ''' import os import argparse +from time import time from datetime import timedelta import dataclasses import functools @@ -161,11 +162,17 @@ def construct_wrapped_graphcast(model_config, task_config): # Modify inputs/outputs to `graphcast.GraphCast` to handle conversion to # from/to float32 to/from BFloat16. - predictor = casting.Bfloat16Cast(predictor) - - # Modify inputs/outputs to `casting.Bfloat16Cast` so the casting to/from - # BFloat16 happens after applying normalization to the inputs/targets. - predictor = normalization.InputsAndResiduals(predictor, diffs_stddev_by_level=self.diffs_stddev_by_level, mean_by_level=self.mean_by_level, stddev_by_level=self.stddev_by_level,) + # NOTE: Do not use float16 for prediction. Reduced precision greatly increases run-to-run variance + # even without perturbed initial conditions + # predictor = casting.Bfloat16Cast(predictor) + + # Applying normalization to the inputs/targets. + predictor = normalization.InputsAndResiduals( + predictor, + diffs_stddev_by_level=self.diffs_stddev_by_level, + mean_by_level=self.mean_by_level, + stddev_by_level=self.stddev_by_level, + ) # Wraps everything so the one-step model can produce trajectories. predictor = autoregressive.Predictor(predictor, gradient_checkpointing=True,) @@ -176,8 +183,11 @@ def run_forward(model_config, task_config, inputs, targets_template, forcings,): predictor = construct_wrapped_graphcast(model_config, task_config) return predictor(inputs, targets_template=targets_template, forcings=forcings,) + t0 = time() jax.jit(self._with_configs(run_forward.init)) self.model = self._drop_state(self._with_params(jax.jit(self._with_configs(run_forward.apply)))) + elapsed_time = time() - t0 + print(f"Elapsed time for compiling the model: {elapsed_time} seconds") def get_predictions(self): @@ -280,11 +290,30 @@ def upload_to_s3(self, keep_data): args = parser.parse_args() runner = GraphCastModel(args.weights, args.input, args.case_name, args.config, args.output, int(args.pressure), int(args.length)) + t0 = time() runner.load_pretrained_model() + elapsed_time = time() - t0 + print(f"Elapsed time for loading model: {elapsed_time} seconds") + + t0 = time() runner.load_gdas_data() + elapsed_time = time() - t0 + print(f"Elapsed time for loading input data: {elapsed_time} seconds") + + t0 = time() runner.extract_inputs_targets_forcings() + elapsed_time = time() - t0 + print(f"Elapsed time for extracting inputs, targets, and forcings: {elapsed_time} seconds") + + t0 = time() runner.load_normalization_stats() + elapsed_time = time() - t0 + print(f"Elapsed time for loading normalization stats: {elapsed_time} seconds") + + t0 = time() runner.get_predictions() + elapsed_time = time() - t0 + print(f"Elapsed time for running the model: {elapsed_time} seconds") upload_data = args.upload.lower() == "yes" keep_data = args.keep.lower() == "yes" From ef6cac03313fd768133d1d09a98ae143c875c500 Mon Sep 17 00:00:00 2001 From: Linlin Cui Date: Wed, 10 Dec 2025 18:02:53 +0000 Subject: [PATCH 2/4] revert to bfloat16 --- oper/run_graphcast.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/oper/run_graphcast.py b/oper/run_graphcast.py index a611b37..e4e9f8c 100644 --- a/oper/run_graphcast.py +++ b/oper/run_graphcast.py @@ -162,9 +162,7 @@ def construct_wrapped_graphcast(model_config, task_config): # Modify inputs/outputs to `graphcast.GraphCast` to handle conversion to # from/to float32 to/from BFloat16. - # NOTE: Do not use float16 for prediction. Reduced precision greatly increases run-to-run variance - # even without perturbed initial conditions - # predictor = casting.Bfloat16Cast(predictor) + predictor = casting.Bfloat16Cast(predictor) # Applying normalization to the inputs/targets. predictor = normalization.InputsAndResiduals( From a130db3689a23c532522e657c130106bcf9616c6 Mon Sep 17 00:00:00 2001 From: Linlin Cui Date: Wed, 10 Dec 2025 20:32:21 +0000 Subject: [PATCH 3/4] replace print with absl.logging --- oper/run_graphcast.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/oper/run_graphcast.py b/oper/run_graphcast.py index e4e9f8c..7e10269 100644 --- a/oper/run_graphcast.py +++ b/oper/run_graphcast.py @@ -21,6 +21,7 @@ import boto3 import pandas as pd import pickle +from absl import logging from graphcast import autoregressive from graphcast import casting @@ -185,7 +186,7 @@ def run_forward(model_config, task_config, inputs, targets_template, forcings,): jax.jit(self._with_configs(run_forward.init)) self.model = self._drop_state(self._with_params(jax.jit(self._with_configs(run_forward.apply)))) elapsed_time = time() - t0 - print(f"Elapsed time for compiling the model: {elapsed_time} seconds") + logging.info(f"Elapsed time for compiling the model: {elapsed_time} seconds") def get_predictions(self): @@ -291,27 +292,27 @@ def upload_to_s3(self, keep_data): t0 = time() runner.load_pretrained_model() elapsed_time = time() - t0 - print(f"Elapsed time for loading model: {elapsed_time} seconds") + logging.info(f"Elapsed time for loading model: {elapsed_time} seconds") t0 = time() runner.load_gdas_data() elapsed_time = time() - t0 - print(f"Elapsed time for loading input data: {elapsed_time} seconds") + logging.info(f"Elapsed time for loading input data: {elapsed_time} seconds") t0 = time() runner.extract_inputs_targets_forcings() elapsed_time = time() - t0 - print(f"Elapsed time for extracting inputs, targets, and forcings: {elapsed_time} seconds") + logging.info(f"Elapsed time for extracting inputs, targets, and forcings: {elapsed_time} seconds") t0 = time() runner.load_normalization_stats() elapsed_time = time() - t0 - print(f"Elapsed time for loading normalization stats: {elapsed_time} seconds") + logging.info(f"Elapsed time for loading normalization stats: {elapsed_time} seconds") t0 = time() runner.get_predictions() elapsed_time = time() - t0 - print(f"Elapsed time for running the model: {elapsed_time} seconds") + logging.info(f"Elapsed time for running the model: {elapsed_time} seconds") upload_data = args.upload.lower() == "yes" keep_data = args.keep.lower() == "yes" From aa7642c680e78ad508ba1af4af8ef394836bf8ee Mon Sep 17 00:00:00 2001 From: Russell Manser Date: Wed, 10 Dec 2025 21:18:23 +0000 Subject: [PATCH 4/4] Revert to print statements, use unbuffered output --- oper/run_graphcast.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/oper/run_graphcast.py b/oper/run_graphcast.py index 7e10269..1ae3b03 100644 --- a/oper/run_graphcast.py +++ b/oper/run_graphcast.py @@ -1,4 +1,4 @@ -#!/usr/bin/env python3 +#!/usr/bin/env -S python3 -u ''' Description: Script to call the graphcast model using gdas products @@ -21,7 +21,6 @@ import boto3 import pandas as pd import pickle -from absl import logging from graphcast import autoregressive from graphcast import casting @@ -186,7 +185,7 @@ def run_forward(model_config, task_config, inputs, targets_template, forcings,): jax.jit(self._with_configs(run_forward.init)) self.model = self._drop_state(self._with_params(jax.jit(self._with_configs(run_forward.apply)))) elapsed_time = time() - t0 - logging.info(f"Elapsed time for compiling the model: {elapsed_time} seconds") + print(f"Elapsed time for compiling the model: {elapsed_time} seconds") def get_predictions(self): @@ -292,27 +291,27 @@ def upload_to_s3(self, keep_data): t0 = time() runner.load_pretrained_model() elapsed_time = time() - t0 - logging.info(f"Elapsed time for loading model: {elapsed_time} seconds") + print(f"Elapsed time for loading model: {elapsed_time} seconds") t0 = time() runner.load_gdas_data() elapsed_time = time() - t0 - logging.info(f"Elapsed time for loading input data: {elapsed_time} seconds") + print(f"Elapsed time for loading input data: {elapsed_time} seconds") t0 = time() runner.extract_inputs_targets_forcings() elapsed_time = time() - t0 - logging.info(f"Elapsed time for extracting inputs, targets, and forcings: {elapsed_time} seconds") + print(f"Elapsed time for extracting inputs, targets, and forcings: {elapsed_time} seconds") t0 = time() runner.load_normalization_stats() elapsed_time = time() - t0 - logging.info(f"Elapsed time for loading normalization stats: {elapsed_time} seconds") + print(f"Elapsed time for loading normalization stats: {elapsed_time} seconds") t0 = time() runner.get_predictions() elapsed_time = time() - t0 - logging.info(f"Elapsed time for running the model: {elapsed_time} seconds") + print(f"Elapsed time for running the model: {elapsed_time} seconds") upload_data = args.upload.lower() == "yes" keep_data = args.keep.lower() == "yes"