Run earth2studio fcn3 on Nvidia LS40s with 46Gb memory, but out of memory

Physics Simulation PhysicsNeMo-ML Model Framework Technical Support (PhysicsNeMo Only)
1

My code is as follows:

import os

import gc

import torch

import shutil

# 1. Set environment flag BEFORE any torch calls

os.environ[“PYTORCH_CUDA_ALLOC_CONF”] = “expandable_segments:True”

# 2. Force clear everything existing

gc.collect()

torch.cuda.empty_cache()

# 3. Handle the IO/Directory (avoiding the previous RuntimeError)

d_out = “outputs/fcn3_forecast.zarr”

if os.path.exists(d_out):

shutil.rmtree(d_out)

from earth2studio.models.px import FCN3

from earth2studio.data import GFS

from earth2studio.io import ZarrBackend

import earth2studio.run as run

package = FCN3.load_default_package()

#model = FCN3.load_model(package).to(device=“cuda”, dtype=torch.bfloat16)

#model = FCN3.load_model(package).to(device=“cuda”, dtype=torch.float32)

model = FCN3.load_model(package).to(device=“cuda”)

data = GFS()

io = ZarrBackend(d_out)

# 5. Run with no_grad to prevent memory doubling

with torch.no_grad():

\# Final check: clear cache one last time before execution

torch.cuda.empty_cache()

\# Run the forecast

\# Ensure nsteps isn't excessively high for the first test (e.g., 20)

io = run.deterministic(\["2024-01-01T00:00"\], 2, model, data, io)

But I get error:

torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 20.37 GiB. GPU 0 has a total capacity of 44.40 GiB of which 13.28 GiB is free. Including non-PyTorch memory, this process has 31.12 GiB memory in use. Of the allocated memory 30.54 GiB is allocated by PyTorch, and 66.87 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.

I tried to clean the cache, and data type like float16, float32, and bfloat16, all failed. The above error occurs when the default is data type is used. Plz help me to figure out how to resolve the issue. THX.