2.1 Best-practice principles

To achieve this, we propose the following set of best-practice requirements and features for a hybrid implementation:

1. It can be readily adapted to replace any portion of the GCM, focusing on (but not limited to) physics parameterizations.

2. It offers a concise and scalable design pattern, combining ease of use and run-time performance.

3. It offers a rapid “plug-and-play” replacement of previous physics surrogates with new ones.

   • This should take the form of a plugin that is attached to a GCM, which, once installed, allows ML and AI models to be loaded into the GCM without requiring to recompile the GCM.

   • This requirement does not include any changes that need to be done to the GCM, namely the need to pipe inputs and outputs from the surrogate model back to the GCM's workflow.

4. It allows multiple ML and AI models to coexist in the same run or as an overlay on top of TPs – for example, to replace two different TPs or for data blending, online learning, or to use side-by-side with TPs.

5. It allows ML and AI parameterizations to coexist in the same source code branch and execution flow of a GCM alongside TP approaches.

6. It uses existing parallelization frameworks and infrastructure (i.e., Message Passing Interface (MPI) over CPUs) but without limiting the ability of ML and AI approaches to use GPUs on demand inside a GCM or during the learning stage.

7. It allows ML and AI approaches to be activated under specific conditions, for example, in a given time period and region.

8. It offers a workflow and supporting tools to boost the learning process of ML and AI models. Since most GCMs are written in Fortran, offer Fortran in addition to C/C$++$ interface implementation.

9. It allows the use of scripted languages during the learning process while avoiding the need to implement an ML and AI model in Fortran.

These requirements and features are achieved here via two deliverables:

1. The TorchClim plugin. This is a library (shared object) that is largely agnostic to a specific GCM. It exposes an interface by which one can access surrogate models from a GCM.

2. A reference implementation. This is a case study that demonstrates the above best-practice principles for incorporating the plugin into a GCM.

TorchClim bundles these into a repository with a GPL v3.0 license (Fuchs et al., 2023a). Our approach relies on PyTorch (Paszke et al., 2019), a deep-learning neural network framework that offers both Python and C/C$++$ interfaces.

2.2 Architecture overview

The workflow of a GCM generally cycles through repeated time steps, where the output from one step feeds into the next. Prognostic variables (e.g., temperature, winds, humidity, and usually one or more cloud quantities) are integrated into the dynamical core and carried from one time step to the next, while diagnostic ones (e.g., radiation fluxes) are recalculated from scratch at each time step inside various parameterization modules (Fig. 1a). A parameterization module can “see” prognostic variables from the dynamical core and, in some cases, diagnostic variables from other parameterizations that have already executed.

The recommended mode of use of TorchClim is one where the surrogate model is exposed to the GCM as a parameterization module that is indistinguishable from others in the GCM. Like any other parameterization, a surrogate model that is implemented in this fashion can use prognostic, diagnostic variables and spatiotemporal dimensions that the GCM provides at the point of insertion of the new parameterization. The role of the parameterization module that wraps the surrogate model boils down to piping inputs and outputs from the GCM into the surrogate model.

In order to call a surrogate model, the new parameterization relies on calls to the TorchClim plugin (hereafter “the plugin”). This plugin is a lean implementation with the main purpose of hiding the details of the underlying ML and AI library implementation from the GCM. The plugin handles initialization and configuration of LibTorch, loads and calls ML and AI models, and, if required, keeps state and alignment between MPI ranks, GPUs, and stateful ML and AI models (see Sect. 5.2).

The first release of the plugin is focused more on functionality than on the computational performance of a hybrid model. It is shipped with a Fortran layer that packs variables according to the input and output specifications of the underlying ML and AI model that is used in the reference implementation (Sect. 3). Users can choose to change this code to match their objective, which will require recompiling the plugin, or implement this layer in their GCM and call the C/C$++$ interface directly from the GCM (calling “model_predict_c” under torch-wrapper/src/interface/torch-wrap-cdef.f90). An example of how to export a PyTorch Model into a TorchScript format that can be uploaded by the TorchScript plugin is also provided (examples/torchscript_example.py in the TorchClim repository). This implementation is relatively lean, so much of the computational complexity that the hybrid model incurs depends on the complexity of the ML and AI surrogate. One advantage of implementing the packing layer in the TorchClim plugin rather than the GCM is that it allows the user to test the calls to the surrogate in a test application without running the GCM. That is, the plugin can be loaded as a separate executable, calling an ML and AI model using a set of validation input test profiles that the user supplies. For example, the reference implementation offers a standalone test application that feeds prescribed atmospheric profiles into ML and AI models via a test executable of this plugin.

Figure 1An overview of the implementation of TorchClim into a GCM. (a) An idealized GCM, where a hybrid parameterization module calls TorchClim. To the GCM, it appears like any other parameterization module. (b) An elaboration of panel (a) for the reference implementation into CESM CAM, where the hybrid model replaces convection, radiation, and cloud parameterizations (adapted from Wang et al., 2022).

Download

2.3 Parallelization considerations

One of the main dilemmas in introducing an ML and AI surrogate into the GCM involves parallelization. We note that most GCMs rely on CPUs to parallelize their execution and scale via MPI to multiple CPUs and compute nodes, while ML and AI frameworks tend to scale using GPUs. This creates a duplication of parallelization frameworks in the GCM, leading to the possibility that the CPUs will be idle while the GPUs access the surrogate model, and vice versa. Furthermore, there is an overhead in copying state from the CPU to the GPU, which could reduce the efficiency of this approach. It is possible that a better use of resources would be to use more CPUs to parallelize the GCM via MPI and keep the surrogate model on the same CPU core that is bound with the MPI rank. This is especially relevant to implementations that access the GPU serially or use a global MPI gather to push data to the GPU. At the least, users should be advised to benchmark their hybrid GCM against several architectures before running expensive workloads to ensure the best use of resources.

TorchClim offers the user the flexibility to choose between a pure CPU configuration (on the same CPU that the MPI rank resides) and one where the surrogate model is pushed to a GPU. Note that GPU functionality will be tested in the next release of TorchClim following the addition of vectorization of the MPI tile and configuration of multiple surrogate models (see Sect. 5). To achieve this flexibility, an instance of the TorchClim plugin is loaded at run-time to the thread-local storage (TLS) inside an MPI rank. This means that TorchClim's recommended best practice naturally follows the spatial discretization used to parallelize GCMs via CPU or MPI. This is illustrated in Fig. 2 for CAM, where the CPU or MPI parallelization discretizes the spatial domain into tiles and allocates each tile to a CPU. The instance of TorchClim lives inside the MPI rank associated with this CPU and can offload the requests to TorchClim to a range of parallelization options as supported by the underlying ML and AI framework.

Figure 2An architecture diagram depicting the relationship between a GCM's physics geographical tile (chunk), representing a set of adjacent geographical grid locations, and a LibTorch instance. In a GCM with MPI support, both coexist within the MPI rank. LibTorch ML and AI models can subsequently execute on a local GPU or a thread pool.

Download

3.1 Alternate physics workflows

We recognize the potential of ML and AI to serve a range of modes of use. For example, users may wish to run an ML and AI surrogate only at a certain time step or within a particular geographical region. It is also desirable to be able to choose among multiple physics parameterizations that point to different surrogate model combinations. Alternatively, it may be desired to run multiple models side-by-side or as a replacement for different physics parameterizations. This functionality could also be used to diagnose instabilities in the surrogate model or do online corrections (e.g., Rasp, 2020), for example, by comparing the outputs from the surrogate model with the TP. The reference implementation supports this need by placing a mode selector before the call to the physics parameterization (Fig. 3). An advantage of this approach is that it offers users a way to extend CAM with new capabilities without removing existing ones. The reference implementation still offers the original parameterization of CAM alongside an ML and AI mode that runs combined moist and radiative parameterizations as an ML and AI surrogate and a dual mode that runs both of these side by side. Our implementation allows the user to achieve that out-of-the-box (without directly changing CESM CAM code base).

Figure 3A sketch of the mode selector, which allows the reference implementation to add new models of CAM physics alongside the original parameterization. The mode selector is placed before the entry point to CAM's physics parameterization, allowing the selector to choose the workflow for a given spatial and temporal location and use case. The selector facilitates fast turnover to extend CAM with additional physics workflows without compromising other workflows.

Download

3.2 Assisting the learning phase

So far we assumed that an ML surrogate is available, without addressing where the data that would be used to train it would come from. A good starting point for training a surrogate model is simply to use the TP of the host GCM itself. In the reference implementation, this is the combined effect of the moist-physics and radiation parameterizations of CAM. This approach has the benefit that it allows us to benchmark a surrogate model both computationally and scientifically against an ideal synthetic dataset. It can also serve as a starting point for further training from other datasets that have missing or insufficient data, as is frequently the case with climate data. The user could mitigate learning biases in this using the TP as a starting point in different ways, for example, using an ensemble of TPs, potentially from different GCMs, or increasing the learning rate when training with subsequent data sources. To this end, the reference implementation of TorchClim is shipped with the export_state Fortran module. The user is provided with two subroutines: the first is inserted in the TP before the location where the surrogate model is to be called, while the second is inserted after that point. These produce a dataset of inputs and labels required for supervised learning. For example, our reference implementation replaces the parameterizations of moist physics and radiation, so these subroutines are placed in the TP before and after these parameterizations (Fig. 4). These add additional history variables to CAM's output, recording state and accumulated tendencies before and after the desired section of TPs.

This functionality could also be used to study instabilities during online runs of the GCM with the surrogate model. In our reference implementation, this can be done by extracting samples before and after the call to the surrogate model and tracking the locations where the surrogate model diverges from the TP offline.

Figure 4The process of producing inputs and labels for supervised learning from the TP parameterization. Here we exemplify it using the TP of moist and radiative parameterization in CAM. The export_state module allows the user to take snapshots of CAM's physics state and tendencies before and tendencies after these TPs. These are saved to CAM's history files and subsequently used to train a baseline DNN (e.g., for our reference implementation).

Download

4.1 DNN surrogate model description

The DNN model for our case study was trained offline via PyTorch's Python interface. The input and output variables that the DNN used are discussed in Sect. 3 and shown in Table 1. These are composed of scalars such as radiative fluxes and vertical model level profiles (26 levels in this version of CAM).

The DNN is composed of seven fully connected layers. Each layer (except the last) is followed by a batch normalization, dropout, and a linear rectifier. We use the zero-mean or 1 standard deviation normalization per feature before the first layer, i.e., separately normalizing each model level (z) globally for each 3-D input variable, such as water vapor, and each 2-D variable (Table 1). Correspondingly, we use the 1 standard deviation re-normalization after the last layer for re-scaling the output during inference.

Equation (1) defines the loss function used to be minimized. The mean squared error (MSE) on the normalized output variables is the initial loss function to be minimized (L_MSE). We also introduce additional terms in the loss function to address the biases described below. We add L2 regularization with a weight of one to all parameters except the biases and batch norm (L_L2). Several constraints on range, equality, and conservation to prevent unphysical predictions are encoded as either parameterization of the outputs in the DNN model or as additional regularization terms in the optimization objective. These are defined for the three constraint groups discussed in Sect. 4.2 and Table 3. The loss trade-off coefficients α_(⋅) are chosen experimentally.

Optimization is done using the Adam method (Kingma and Ba, 2014) with a learning rate of $\mathrm{1}\times {\mathrm{10}}^{-\mathrm{3}}$. We train for 100 epochs, with a large batch size of 24 × 96 × 144 atmospheric columns that are randomly selected in time and space from the original 10-year dataset. We determined 100 epochs to be a good compromise between model accuracy and training time on a separate validation set. Training beyond 100 epochs resulted only in marginal gains. Finally, we use linear warm-up for the first 10 % of epochs and cosine cool-down learning rate scheduler to zero.

During training, the validation data are used to monitor optimization progress and perform early stopping if necessary.

In practice, we tested dozens of versions of the DNN with various bug fixes and tweaks. This was made easy by the TorchClim interface which enables newly trained DNNs to be dropped in with no recompile of CAM required.

4.2 Constraining the target solution

Past studies and our own efforts have found that to obtain good performance, it is necessary to apply physical constraints to ML surrogates to prohibit nonphysical predictions (Beucler et al., 2021; Brenowitz et al., 2020; Mooers et al., 2021; Karniadakis et al., 2021).

In training the model for our reference implementation, we found a set of rules that helped constrain the target solution.

Here, we propose a classification of such constraints into three types, according to the nature of the applied constraint (Table 3).

Table 3Classification of constraints of the solutions' space.

Download Print Version | Download XLSX

The variable x will usually be outputs of the surrogate but could also be inputs. Type I constraints stem from the fact that physical variables may be bounded in some way either by definition, a static threshold, or by another variable. For example, precipitation rate and shortwave radiation flux cannot be negative. Likewise, no component of shortwave radiation can exceed the incoming solar flux (given that plane-parallel radiation is assumed in CAM). The constraint, in this case, limits these variables to the space of solutions $\mathit{\{}x\mathrm{|}x\in R,x\mathit{>=}\mathrm{0}\mathit{\}}$ and $\mathit{\{}x\mathrm{|}x\in R,x\mathit{>=}\mathrm{SOLIN}\mathit{\}}$, where SOLIN is solar insulation variable in CAM. We have found that ML approaches struggle to consistently obey this type of constraint to the required accuracy, especially since even small violations are unphysical and can cause problems elsewhere in the model; hence, we enforce the positivity constraint by applying the rectifier function to the output variables max(x,0). This is a type of inductive biasing (Karniadakis et al., 2021). Type II or “redundancy” constraints formalize a relation among variables, which generally means that the target solution exists on a manifold inside the unconstrained set of target solutions, i.e., there is physical redundancy in the output variables. For example, the sum of direct and diffuse shortwave radiation at the surface must be equal to the total radiation at the surface. Type II constraints, therefore, take the form $x_j={\sum }_{i=\mathrm{1}}^n{x}_i$, where x_j is a surrogate output that consists of n components x_i also represented by other surrogate outputs. To improve overall model learning, we incorporate these physically dictated constraints by converting them to a residual loss term in the objective $\left|\right|x_j-{\sum }_{i=\mathrm{1}}^n{x}_i|{|}_{\mathrm{2}}$, such that when the constraint is satisfied the residual will be zero. Note that such constraints could also be nonlinear. Type III or “soft” constraints penalize solutions that disobey a desired physical property expressed via functions of the output variables. For example, we expect that the sum of all tendencies of water vapor in a column will balance those of condensed water plus precipitation, but this relationship may not be exact due to small terms in the conservation equation (such as storage of condensed water) that are not known or available. Type III constraints are examples of learning biasing, following Karniadakis et al. (2021). As listed in Table 3, we implement a soft constraint that penalizes increasing moistening tendencies when relative humidity is above a specified threshold with gradient regularization term $\left|\mathbf{1}\right[h\left(x_i\right)>\mathit{\beta }\left]\frac{\partial x_j(x_i,\mathrm{\dots })}{\partial x_i}\right|$. This term will penalize positive slopes, i.e., gradients of the output variable x_j with respect to input variable x_i but only when h(x_i)>β. In this example, x_j is the moistening tendency, x_i is the input humidity, and h() is the relative humidity function.

Naively, these constraints could just be applied at run-time in the GCM's integration. For example, one might truncate a precipitation variable to be non-negative at run-time. Adding these constraints to the learning stage, however, can assist in learning the joint solution rather than correcting a single variable. For example, if the learning process had predicted negative precipitation, this may imply non-physical predictions of other variables that are empirically related to precipitation even if no other constraints are violated. To ameliorate this issue, we bring Type I–III constraints into the learning stage, presenting them as additional terms in the loss function. Type I constraints are also applied at run-time, since learning does not completely eliminate violations, and even small Type I violations (for example negative solar fluxes) can crash the land model.

4.3 The skill of the hybrid CAM ML model

We evaluate the hybrid CAM ML model by comparing it with standard CAM, whose physics the surrogate emulates. We do not perform an extensive skill evaluation, since our main objective is to test the plugin rather than the skill of the NN physics emulator, which can easily be changed. Here we present results from a hybrid CAM ML run that starts from the same starting point that was used for extract the training data. That is the CESM Atmospheric Model Intercomparison Project (AMIP) scenario for CAM4 with default configuration parameters. The hybrid CAM ML model was able to run for 6 months before exhibiting numerical instabilities. These instabilities appear to be brought on by stratospheric drift which we have not yet attempted to correct. Here we present results for the fifth month into the run, before the instabilities became apparent. During this month, the hybrid model produces zonal mean moisture and temperature tendencies that closely resemble those of the original model (Fig. 5). The Inter-Tropical Convergence Zone (ITCZ) is slightly too narrow, and there is a double-ITCZ bias in the hybrid model and too much heating near the tropopause, but the differences overall appear modest. The hybrid model also shows similar temporal variations in tropical variability and waves when comparing them to the original CAM model (Fig. 6). This is a stronger test of the surrogate model, since the character of tropical waves is sensitive to the behavior of moist physics (Majda and Khouider, 2002; Kelly et al., 2017). Note that the initial behavior of the hybrid model matches that obtained with CAM parameterizations, diverging later, as expected, but retaining a similar character. Despite spatiotemporal pattern similarities, the hybrid model exhibits larger relative humidity extremes (Fig. 6). It is not clear why these extremes are permitted by the surrogate model, although earlier versions of the model that lacked the relative humidity (RH) sensitivity regularization (Sect. 4.2) showed a more severe manifestation of the problem, suggesting that even with regularization, the surrogate model is insufficiently quick in the condensation of water above saturation compared to the CAM physics. This is not too surprising, as the parameterized physics is extremely nonlinear, and we have not explicitly coded the Clausius–Clapeyron relation into the DNN, so it must learn the point at which condensation (rapidly) begins. Thus, we expect that better use of physics-informed architectures or other improvements would lead to further improvement in emulation skill.

Figure 5Zonal mean specific humidity and temperature tendencies in the fifth month after the start of the simulation. The original parameterized CAM (a, c) and the hybrid CAM ML, where CAM is coupled with the surrogate model using TorchClim (b, d), are shown.

Download

Figure 6Predicted values vs. longitude and time during the fifth month after the start of the simulation at hybrid model level 233 at the Equator (approximately 233 hPa). The original parameterized CAM (a–c) and the hybrid CAM ML, where CAM is coupled with the surrogate model using TorchClim (d–f), are shown. (a–f) Relative humidity, specific humidity tendencies, and temperature tendencies. The y axis denotes the model physics time step (which is 30 min).

Download

Further insights into the behavior of the hybrid model are gained using the third mode of operation of TorchClim, where the ML model is called alongside the original parameterization of CAM. Here we call the ML model and output its response without assimilating it into the CAM workflow. Figures 7–8 compare specific humidity and temperature tendencies of the original CAM parameterization to the ones from the ML model at various vertical levels between 10° N and 10° S from the Equator. A perfect ML model would follow the 1:1 line. Both specific humidity and temperature tendencies exhibit a line of best fit that is smaller than one. The ML model generally predicts smaller positive specific humidity tendencies (Fig. 7). Interestingly, for both quantities, the lowest skill is found in the mid-troposphere between hybrid model levels 510 and 696. At these levels, the ML model learns the spurious features of specific humidity tendencies that do not exist in the original model. These discrepancies appear at vertical levels where shallow convection is active and are in line with Mooers et al. (2021), who found that their deterministic DNN had less skill in regions with fast stochastic convective activity. However, this issue could also be attributed to the fact that cloud liquid and ice are not treated by the current surrogate DNN model.

Although these biases could be discovered offline during training, we found that the ability of TorchClim to run the original and hybrid CAM models side by side enabled us to quickly diagnose errors in hybrid simulations, particularly where there may not be analogs in the training dataset.

Figure 7Specific humidity tendencies (kg kg⁻¹ d⁻¹ × 10) of the original cam parameterization (x axis) and the hybrid CAM ML parameterization (y axis) for various vertical levels on day 10 of the run. Calls to the ML model were done from within the integration loop of the original CAM parameterizations. Black lines denote the origins and the 1:1 line. The hybrid model level is denoted in the title, while the slope (sl) and r² are printed inside each panel.

Download

Figure 8As in Fig. 7 but for temperature tendencies (K d⁻¹).

Download

4.4 Computational performance

The evaluation of computing speed is undertaken on the normal queue of the Gadi infrastructure in the National Computational Infrastructure (NCI), supported by the Australian Government. The normal queue nodes are based on 2 × 24 core Intel Xeon Platinum 8274 (Cascade Lake) 3.2 GHz CPUs, with 192 GB RAM, two CPU sockets, each with two NUMA (non-uniform memory access) nodes, and no hyper-threading. Noticeably, MPI ranks are bound to the core, with each MPI rank hosting an instance of the DNN model.

For the sake of comparison, both original and hybrid CAM configurations were run for a year with minimal monthly mean output variables so that a relatively large fraction of the overall CPU would be spent on computation. Each run used three nodes on the normal queue (144 CPU cores) and 64 GB RAM. This configuration matched the number of longitudes of our AMIP scenario spatial discretization. The hybrid and standard configurations of the GCM required similar resources, unlike previous implementations we know of, where the DNN implementation required significantly more compute resources than the original model. The results of this experiment are detailed in Table 4. The hybrid model added 20 % to the wall-time of CAM, reducing its performance by 8 modeled years per wall-time day. The total CESM speed degraded by 10 %, which amounted to 4 modeled years per wall-time day. The addition of ML and the AI model added only about 5 % (2 GB) to the overall memory requirements.

Since our setup is similar to Wang et al. (2022), it is convenient to compare the performance of these approaches. Wang et al. (2022) used a similar 1.9 × 2.5° grid resolution and 30 instead of 26 vertical levels but used CAM5, which is about 4 times slower than CAM4. Noting uncertainties regarding the computational complexities of the DNNs of each implementation, their configuration was able to produce 10 modeled years per wall-time day, which is comparable to the performance of our reference implementation. However, their approach used 25 % more CPU cores, Intel's Math Kernel Library optimization, and additional dedicated hardware (GPU and storage). We also note that our own results were achieved before any optimization (such as vectorizing the calls per geographical tile, pre-allocating of memory buffers, and using GPUs). In developing future versions of TorchClim, we anticipate vectorizing the calls to the surrogate DNN model, pre-allocating and re-using data structures, and considering automatically offloading to a GPU when available. It is expected that, despite each MPI rank hosting a single thread, the performance improvements in vectorization and pre-allocation will be significant (even without GPUs or multi-threading). For example, assuming that vectorization will match the spatial tile (chunk in CAM's terminology), the reduction in heap memory allocations calls will be of the order of $O\left(\right|\mathrm{tile}\left|\right)$, where $\left|\mathrm{tile}\right|$ denotes the number of atmospheric columns of the tile (thus being more important for larger tiles in terms of the number of grid locations per tile).

Table 4Run-time performance for CESM total run-time and the atmospheric component (CAM) only.

Download Print Version | Download XLSX

4.5 How might TorchClim help with stability issues

Stability issues take a considerable portion of surrogate model development. Past efforts have found that stability can vary unpredictably among similarly constructed surrogate models, indicating a need to understand this for hybrid modeling to progress (Wang et al., 2022; Brenowitz et al., 2020). We list the ways in which the TorchClim approach could help.

1. Our use of data from a TP to train a surrogate model seems like an ideal test bed to study online stability issues that arise due to imperfect emulation.

2. The mode selector (Sect. 3.1) allows the TP and the surrogate model to be run side by side, which could help catch or correct issues online during the GCM run.

3. The data extraction tool (Sect. 3.2) also allows training samples to be extracted with both the surrogate and TP predictions side by side, providing the option to study instabilities offline. These could also be used to re-train the surrogate model to mitigate specific instabilities.

4. The TorchClim plugin can be used as a separate executable, allowing changes to the interface to be tested outside the GCM.

We note that the last two of these possibilities helped us ensure that the integration between the plugin and surrogate models matches to inference via Python.

5.1 ONNX support

The current version of TorchClim is able to import surrogate models that were saved through PyTorch's scripting infrastructure. A natural extension of that is to allow the use of the ONNX interface to imports. With the maturity of ONNX, this would improve performance while allowing us to use the ML and AI framework beyond PyTorch.

5.2 Stateful ML and AI models

A stateful AI and ML model preserves an internal state between one call to another, for example, in the case of agent-based algorithms such as reinforcement learning. In the case of GCMs, a specific instance of surrogate model, compute node, and GPU may be associated with an MPI rank. However, a standing issue with MPI applications is the need to optimize their performance on a given infrastructure. In the age of non-uniform memory architectures (NUMA), this tends to be achieved via binding (affinity) instructions that communicate with the underlying infrastructure the required binding of MPI ranks. For example, many applications require binding to cores, memory, or cache to achieve the best performance. The case of the stateful ML and AI model brings in an additional challenge. In this case, an instance of the ML and AI model holds an internal state that is bound to a given geographical location. This has two consequences.

1. A geographical location in the GCM needs to be mapped to an instance of the ML and AI model, which is mapped to a specific thread.

2. The instance of the ML and AI model (and state) is bound to a given GPU if the ML and AI model uses GPUs at the GCM run-time.

The consequence of this is the binding of a geographical location or tile to a resource (thread and GPU). An implementation of this feature requires TorchClim to retain a mapping between a GPU and a geographical location and for the GCM to bind the geographical location (tile) to a CPU. Since the latter requirement is GCM-specific, its implementation of these is out of scope in the current reference implementation.

5.3 Vectorization and pre-allocation

The current implementation is shipped without optimization. Despite that, our experiments with the reference implementation into CESM CAM exhibit good performance. Beyond that, the framework is desired to automate the vectorization of calls to underlying surrogates. In addition, the framework should pre-allocate all the buffers required to interact with LibTorch. These improvements are expected to have a significant effect on performance even without the use of GPUs.

5.4 Multiple surrogate models

The reference implementation does not support multiple surrogate models within the same run. However, the extension of the TorchClim framework to support multiple surrogate models is straightforward.