Failed to allocate device memory in Neural GCM simulation

[pha96]

I am using an NVIDIA GPU to run the neural GCM climate simulation. However, I always get stuck at the horizontal interpolation due to memory issues. The program keeps failing to allocate memory, although the input variable is only 1 GB. Has anyone encountered this problem? Thank you for your help.
nvidia-smi

+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 610.57.04              KMD Version: 610.57.04     CUDA UMD Version: 13.3     |
+-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA RTX 4000 Ada Gene...    On  |   00000000:8A:00.0 Off |                  Off |
| 30%   29C    P8             12W /  130W |       2MiB /  20475MiB |      0%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+

+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI              PID   Type   Process name                        GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
|  No running processes found                                                             |
+-----------------------------------------------------------------------------------------+

It seems that my code can access the gpu, the cuda-toolkit and driver work properly.
nvidia-smi pmon -i 0

# gpu         pid   type     sm    mem    enc    dec    jpg    ofa    command
# Idx           #    C/G      %      %      %      %      %      %    name
    0          -     -      -      -      -      -      -      -    -
    0     101726     C      0      0      -      -      -      -    python
    0     101726     C     62     12      -      -      -      -    python
    0          -     -      -      -      -      -      -      -    -
    0          -     -      -      -      -      -      -      -    -

I also tried to specify some environment variables, but none of this ever worked.

import os
#os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
#os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = ".75"
#os.environ["XLA_PYTHON_CLIENT_ALLOCATOR"] = "platform"

Code:

import gcsfs
import jax
import numpy as np
import pickle
import xarray

from dinosaur import horizontal_interpolation
from dinosaur import spherical_harmonic
from dinosaur import xarray_utils
import neuralgcm
import matplotlib.pyplot as plt
import time

start_time = time.perf_counter()
# print("done import")

model_name = 'v1/deterministic_2_8_deg.pkl'  #@param ['v1/deterministic_0_7_deg.pkl', 'v1/deterministic_1_4_deg.pkl', 'v1/deterministic_2_8_deg.pkl', 'v1/stochastic_1_4_deg.pkl', 'v1_precip/stochastic_precip_2_8_deg.pkl', 'v1_precip/stochastic_evap_2_8_deg.pkl'] {type: "string"}

gcs = gcsfs.GCSFileSystem(token='anon')
with gcs.open(f'gs://neuralgcm/models/{model_name}', 'rb') as f:
  ckpt = pickle.load(f)

model = neuralgcm.PressureLevelModel.from_checkpoint(ckpt)
#print(model)

era5_path = 'gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3'
full_era5 = xarray.open_zarr(
    era5_path, chunks=None, storage_options=dict(token='anon')
)

demo_start_time = '2000-04-01'
demo_end_time = '2000-04-01'
data_inner_steps = 24  # process every 24th hour

sliced_era5 = (
    full_era5
    [model.input_variables + model.forcing_variables]
    .pipe(
        xarray_utils.selective_temporal_shift,
        variables=model.forcing_variables,
        time_shift='24 hours',
    )
    .sel(time=slice(demo_start_time, demo_end_time, data_inner_steps))
    .compute()
)
print(sliced_era5)

era5_grid = spherical_harmonic.Grid(
    latitude_nodes=full_era5.sizes['latitude'],
    longitude_nodes=full_era5.sizes['longitude'],
    latitude_spacing=xarray_utils.infer_latitude_spacing(full_era5.latitude),
    longitude_offset=xarray_utils.infer_longitude_offset(full_era5.longitude),
)
regridder = horizontal_interpolation.ConservativeRegridder(
    era5_grid, model.data_coords.horizontal, skipna=True
)
eval_era5 = xarray_utils.regrid(sliced_era5, regridder)
eval_era5 = xarray_utils.fill_nan_with_nearest(eval_era5)
print("done regrid")

inner_steps = 24  # save model outputs once every 24 hours
outer_steps = 5 * 24 // inner_steps  
timedelta = np.timedelta64(1, 'h') * inner_steps
times = (np.arange(outer_steps) * inner_steps)  # time axis in hours

print('initialize model state')
inputs = model.inputs_from_xarray(eval_era5.isel(time=0))
input_forcings = model.forcings_from_xarray(eval_era5.isel(time=0))
initial_state = model.encode(inputs, input_forcings)

# use persistence for forcing variables (SST and sea ice cover)
all_forcings = model.forcings_from_xarray(eval_era5.head(time=1))

print('make forecast')
final_state, predictions = model.unroll(
    initial_state,
    all_forcings,
    steps=outer_steps,
    timedelta=timedelta,
    start_with_input=True,
)
predictions_ds = model.data_to_xarray(predictions, times=times)
print("done prediction")

for v in predictions_ds.data_vars:
    da = predictions_ds[v]
    da.to_netcdf(f"{v}.nc")
    print(f"Exported to {v}...")

Error:

E1016 16:21:17.024215  101726 cuda_executor.cc:1365] [0] Failed to allocate device memory of 14.68GiB (15766388736 bytes): RESOURCE_EXHAUSTED: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
=== Source Location Trace: ===
external/xla+/xla/stream_executor/cuda/cuda_status.cc:45
external/xla+/xla/stream_executor/cuda/cuda_device_allocator.cc:243
external/xla+/xla/stream_executor/cuda/cuda_device_allocator.cc:487

E1016 16:21:17.371251  101726 cuda_executor.cc:1365] [0] Failed to allocate device memory of 13.21GiB (14189749248 bytes): RESOURCE_EXHAUSTED: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
=== Source Location Trace: ===
external/xla+/xla/stream_executor/cuda/cuda_status.cc:45
external/xla+/xla/stream_executor/cuda/cuda_device_allocator.cc:243
external/xla+/xla/stream_executor/cuda/cuda_device_allocator.cc:487

E1016 16:21:17.675444  101726 cuda_executor.cc:1365] [0] Failed to allocate device memory of 11.89GiB (12770774016 bytes): RESOURCE_EXHAUSTED: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
=== Source Location Trace: ===
external/xla+/xla/stream_executor/cuda/cuda_status.cc:45
external/xla+/xla/stream_executor/cuda/cuda_device_allocator.cc:243
external/xla+/xla/stream_executor/cuda/cuda_device_allocator.cc:487

Traceback (most recent call last):
  File "/home/pdang/research/neural_gcm/test_run.py", line 63, in <module>
    eval_era5 = xarray_utils.regrid(sliced_era5, regridder)
  File "/home/pdang/.local/lib/python3.13/site-packages/dinosaur/xarray_utils.py", line 1233, in regrid_horizontal
    data = xarray.apply_ufunc(
        regridder,
    ...<4 lines>...
        vectorize=True,  # loop over time & level, for lower memory usage
    )
  File "/home/pdang/.conda/envs/clim/lib/python3.13/site-packages/xarray/computation/apply_ufunc.py", line 1279, in apply_ufunc
    return apply_dataset_vfunc(
        variables_vfunc,
    ...<7 lines>...
        on_missing_core_dim=on_missing_core_dim,
    )
  File "/home/pdang/.conda/envs/clim/lib/python3.13/site-packages/xarray/computation/apply_ufunc.py", line 527, in apply_dataset_vfunc
    result_vars = apply_dict_of_variables_vfunc(
        func,
    ...<4 lines>...
        on_missing_core_dim=on_missing_core_dim,
    )
  File "/home/pdang/.conda/envs/clim/lib/python3.13/site-packages/xarray/computation/apply_ufunc.py", line 452, in apply_dict_of_variables_vfunc
    result_vars[name] = func(*variable_args)
                        ~~~~^^^^^^^^^^^^^^^^
  File "/home/pdang/.conda/envs/clim/lib/python3.13/site-packages/xarray/computation/apply_ufunc.py", line 820, in apply_variable_ufunc
    result_data = func(*input_data)
  File "/home/pdang/.conda/envs/clim/lib/python3.13/site-packages/numpy/lib/_function_base_impl.py", line 2530, in __call__
    return self._call_as_normal(*args, **kwargs)
           ~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
  File "/home/pdang/.conda/envs/clim/lib/python3.13/site-packages/numpy/lib/_function_base_impl.py", line 2523, in _call_as_normal
    return self._vectorize_call(func=func, args=vargs)
           ~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pdang/.conda/envs/clim/lib/python3.13/site-packages/numpy/lib/_function_base_impl.py", line 2604, in _vectorize_call
    res = self._vectorize_call_with_signature(func, args)
  File "/home/pdang/.conda/envs/clim/lib/python3.13/site-packages/numpy/lib/_function_base_impl.py", line 2643, in _vectorize_call_with_signature
    results = func(*(arg[index] for arg in args))
  File "/home/pdang/.local/lib/python3.13/site-packages/dinosaur/horizontal_interpolation.py", line 285, in __call__
    not_nulls = jnp.logical_not(jnp.isnan(field))
                                ~~~~~~~~~^^^^^^^
MemoryError: std::bad_alloc
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

[Erfouni]

JAX preallocates about 75% of GPU memory by default, so it can fail even when the input is small. Try XLA_PYTHON_CLIENT_PREALLOCATE=false. The 1 GB input can also create much bigger intermediate arrays during interpolation, so the full error message will show the size it tried to allocate. If you post that line, we can check whether it's the input or an intermediate array.

[pha96]

It works in my case. Thank you very much for the quick response.

[Erfouni]

yw, plz mark it as an answer.