Training a convection emulator inside a differentiable GCM

Context in one page

ResCu is a neural-network replacement for the convection and cloud parameterisation of a climate model. It was trained offline on a super-parameterised model (SPCAM), where every grid column contains a small cloud-resolving model. The net learns to map the column state the host provides to the tendencies the embedded cloud model produced. It runs inside CAM5 (the hybrid model is called HyCAM) and now inside a JAX port of a spectral GCM (jcm, built on the dinosaur dynamical core), where the whole model, dynamics and physics, is differentiable.

The offline-trained net is good at short lead but has a persistent climatological error in the JAX host: precipitation forms a double band around the equator, with too much rain south of it and too little in the northern intertropical convergence zone. A ladder of nudging experiments showed that the host winds, not the net's column response, build that band: with the winds relaxed to reanalysis the bias mostly disappears. Offline training cannot see this feedback. Online training, where the loss is evaluated after the net and the host have interacted for some time, can in principle reach it.

Why the emulator is a one-dimensional CNN

The question from the computing side: of all the architectures, why a convolution along the vertical?

In a global model each atmospheric column is treated independently by the physics. The input to a convection scheme is a set of vertical profiles (temperature, humidity, the large-scale forcing on each) plus a few surface scalars, and the output is a set of vertical profiles of tendencies. So the natural data object is a multi-channel one-dimensional signal along the level axis: 30 levels, 32 input channels, 4 output channels. That is exactly what a 1-D CNN consumes.

The reasons, in order of weight

• Locality in the vertical. Convective transport is a vertical flux divergence, cloud formation depends on the layers just above and below, and stability is a local vertical gradient. A kernel of width 3 expresses those operators directly; a stack of blocks extends the receptive field to the full column, which deep convection needs.
• Weight sharing across levels. A plume or an inversion behaves similarly whether it sits at level 12 or level 20. Sharing the kernel over levels gives approximate translation equivariance in height and keeps the net small: a million parameters, versus a dense network over the flattened 960-element column that would spend most of its parameters on level-to-level coupling it must learn from scratch.
• Residual blocks. The net predicts tendencies, small corrections to a state; pre-activation residual blocks with identity shortcuts train stably at depth and start close to the identity, which matches the physics of a small tendency.
• Memory instead of recurrence. A cloud-resolving model has state that persists between host time steps. ResCu approximates that with a two-step history of inputs and its own previous outputs fed back as channels, which is cheaper and simpler to deploy than a recurrent cell and proved sufficient offline.
• Deployability. Column independence means the net batches over all columns of the globe in one call, in Fortran (through FTorch in CAM5) or in JAX. Nothing in the architecture couples neighbouring columns, so the host's parallel decomposition is untouched.

What we know it does not do

• It is deterministic. The cloud-resolving truth is stochastic; the net predicts a conditional mean. Cloud condensate in particular is under-dispersed, which is one reason the cloud net exists as a separate, smaller stage. A diffusion-model variant was explored for this reason.
• It has no horizontal context. Organised convection (the MJO, squall lines) is emergent from the host coupling, not represented inside the net.
• A vertical convolution-plus-attention variant was built and coupled. It did not remove the coupled biases either, which is part of why the emphasis moved from architecture to training procedure.

Online training setup

Each iteration has two steps. Generate: run the model with the current weights from many start dates and save the state after a spin-up. Optimise: from each saved state, roll the model forward inside the gradient for a window, compare with observations, and back-propagate through dynamics and physics into the net.

What is differentiated and what is not

The forward model is complete: dynamics, the net, the rectifiers, vertical diffusion, surface fluxes and radiation all run every step. In the backward pass the radiation scheme is detached: its heating is treated as a fixed forcing along the trajectory rather than differentiated with respect to the state. The reason is practical. The exact radiation backward needs about 300 GB of temporaries unless chunked, and the chunked version produces NaN on its own. Operational 4D-Var systems make the same kind of simplification in their adjoint physics.

Cost figures

Pass 1 results

Three epochs over the 200 episodes: 603 gradient evaluations, 355 finite updates, 248 lost to NaN. Because each epoch visits the same 200 windows, the epoch means are directly comparable.

The precipitation term fell 16 percent, improved in every one of the 200 episodes, and was still falling linearly at the end: the fit is far from saturated. The wind term, four times larger in the loss, did not move at all. Within 6 hours the winds are fixed by the initial condition, so that term cannot be learned by a convection net, yet it dominates the gradient.

Held-out free rollouts with the final weights

Problems, ranked

1. Half the gradient steps come back NaN

The forward pass is always finite. On a NaN step every leaf of the super net's gradient is NaN. What we established today:

• The same episode with the same weights gives a finite gradient in a fresh process, and NaN on every retry within the process that first produced NaN.
• Two differently compiled executables of the same loss (a straight-line unrolled form and a nested-group form) fail on independent subsets of the same 25 episodes: 9 and 7 failures, 2 shared, where independence predicts 2.5.
• Ruled out: the episode, the weights, the time step, jax 0.11.2 versus 0.11.1, CUDA command buffers, XLA autotuning.

Our reading is a memory-state-dependent kernel in the compiled backward, most likely an uninitialised read, whose effect is deterministic within a process because the allocation sequence repeats. The mitigation now installed is a second executable used only when the first gives NaN, which should cut the loss from a third of the steps to under a tenth. It is a workaround, not an understanding.

2. Compile time and memory scale with the window

The 12-step backward compiles in 30 to 40 minutes per forcing year, and the compile grows with the number of unrolled steps: a 5-day window (240 steps) unrolled would compile for hours. A scan over checkpointed 12-step groups, which compiles one group and loops it, is written and being tested. That form was blamed for NaN early on; given problem 1, that blame is probably misattributed.

3. The loss is badly proportioned

The wind term is unlearnable at 6 hours and dominates the gradient. A precipitation-only pass (wind weight zero, global-mean penalty kept) is running now. Winds become a meaningful target only at 2-day or 5-day windows, where the net's heating has had time to spin the circulation up.

4. Short windows cannot reach the climatology

This is the central scientific problem and the reason for everything above. The double band is a host-circulation feedback with a time scale of days. Training must integrate over that time scale, which means longer windows and the cost that comes with them.

Plan and success criteria

The goal for the next weeks is a demonstration that online training produces a significant precipitation improvement in free runs of days, weeks and a month, convincing enough to justify a larger GPU allocation.

Criteria we propose to hold ourselves to

• Five-day rollouts: RMSE and tropical pattern correlation better than the original net at every lead day. Pass 1 meets this, by 3 to 9 percent.
• Thirty-day rollouts: the 30-day-mean pattern error against IMERG clearly down (pass 1: 2.39 to 2.30, not enough), the southern zonal peak at least halfway from 5.8 toward IMERG's 4.4 mm/day, the northern peak not weakened.
• No drift and a stable global mean over 30 days.

What more GPU would buy

• 80 GB cards: radiation live in the backward pass at manageable chunking, so the cloud net can be trained and the adjoint is closed.
• Several cards per job: one episode per card with gradient accumulation across cards, so a batch of episodes per update instead of one.
• Dedicated nodes: compiles of long windows once, without the 2-hour shared-queue slot as the unit of work.

Questions for discussion

Where a computing-science view would change what we do.

1. Localising the NaN. A gradient that is NaN under one XLA executable and finite under another, deterministic within a process, for the same inputs. What is the right tool: compute-sanitizer initcheck on the compiled kernels, XLA's HLO dumps with buffer assignment, deterministic-ops flags (which failed to compile here), or bisecting the backward by stopping gradients at intermediate points? Is the two-executable fallback an acceptable engineering answer for now?
2. Memory versus compile time for long windows. Straight-line unrolling with per-step checkpointing fits 12 steps in 29 GB; nesting groups extends it at the price of compile time. Is there a better remat policy or a scan structure that keeps the compile at one group while preserving the checkpoint structure? Is pipeline parallelism over the window across GPUs realistic in JAX?
3. Batching. One episode per update saturates a 40 GB card. Gradient accumulation across cards is the obvious route; is there anything smarter than data parallelism over episodes for a model where each sample is a global state?
4. Loss design. Square-root transform of daily precipitation, a global-mean penalty, a wind term that only becomes learnable at longer windows. Is a curriculum over window length (6 hours, then 2 days, then 5 days, warm-starting each) the right structure, and how would one detect that a longer window has started to matter?
5. Architecture. Given a differentiable host, is the 1-D CNN still the right emulator, or does online training favour something with horizontal context, a learned memory state, or stochastic outputs? What would it take to make the cloud net trainable without radiation gradients?
6. Precision. The dynamical core defaulted to bfloat16 matrix multiplies and had to be forced to float32; the finite-difference closure check is at the edge of float32 noise. Where would mixed precision be safe in this pipeline and where not?

Code: the trainer runs/train/train_episodes.py (episode loop, loss, two-executable fallback), the generator runs/train/generate_episodes.py, the evaluator runs/train/rollout_eval.py, all in the jcm_rescu project on Perlmutter. The design note with the full derivation is analysis/online_training_design_20260927.pdf.