GPU memory usage not consistent with estimated usage

[liuchihl]

Hi,

I've conducted a few 3D experiments to examine the relationship between actual GPU memory usage and estimated GPU memory usage. The estimated GPU memory usage is calculated using the formula: Nx x Ny x Nz x 8 x N_state, where Nx, Ny, Nz are the grid dimensions, the factor of 8 represents 64-bit floating-point numbers, and N_state is the number of state variables. Since buoyancy is the only tracer, the state variables are u, v, w, b, p, du/dt, dv/dt, dw/dt, db/dt.

The figure below illustrates that the GPU usage can be scaled linearly with my estimation, which is great. However, the actual GPU memory usage is approximately 3.7 times larger than my initial estimate (y-axis). This suggests that there may be additional state variables that were not accounted for in my estimation. I'm interested in identifying any overlooked variables. Alternatively, if my estimation is correct, I would like to understand the reasons for this discrepancy. Any thoughts would be appreciated!

More information that might help:

1. running longer does not seem to change the GPU memory usage, indicating there is no memory leak
2. The above experiments include 12 variables that are written in netcdf output writer. But writing 4 variables (another experiment) as output is only 14% less than writing 12 variables. Therefore, I think the number of variables does not explain why there is a discrepancy between estimated GPU memory usage and actual GPU usage.
3. I am running on Oceananigans@v0.90.11 with Julia 1.10.2.

Chih-Lun

[glwagner]

Each prognostic variable has the cost of 3 fields --- 1 for the state, and 2 for the tendencies (it looks like you're only allowing for one tendency term per field). You may also want to account for halo regions (so each field has size Nx + 2Hx, etc). For the NonhydrostaticModel there are two pressure fields (one for hydrostatic pressure, one for nonhydrostatic pressure anomaly). The closures sometimes allocate fields. Also, the nonhydrostatic pressure solver allocates at least one additional complex-valued field for scratch memory (I think there can occasionally be a second, depending on the topology and also whether its on GPU vs CPU --- take a look at the FFTBasedPoissonSolver).

Let's see how far off the estimate is after accounting for all that...

[francispoulin]

All very interesting to see and I just have two thoughts.

1. What if you try saving just one field? I know that's not as exicting but if you get there's a descrptency when you only save the buoyancy, then maybe that will shed some light on what happens with other fields.
2. Would you consider doing this with the (experimental) shallow water model? It does not involve any pressure solvers so that might be easier to sort out.

[liuchihl]

Thanks for your thoughts.

What if you try saving just one field? I know that's not as exicting but if you get there's a descrptency when you only save the buoyancy, then maybe that will shed some light on what happens with other fields.

I agree, that's a good idea. I can maybe even run without saving any fields, excluding any memory usage for output writer (though small), the discrepancy would be smaller.

Would you consider doing this with the (experimental) shallow water model? It does not involve any pressure solvers so that might be easier to sort out.

That's a good point, it would be a great experiment to do.

[glwagner]

I agree, that's a good idea. I can maybe even run without saving any fields, excluding any memory usage for output writer (though small), the discrepancy would be smaller.

This is also all an approximation because there are many 1D arrays that are allocated, eg for the stretched grid, the tridiagonal solver, etc.

[glwagner]

You might just remove the grid stretching to see how much usage changes. If you do that, you'll use a simpler pressure solver and the grid metrics will be constant rather than arrays. And the grid nodes will be ranges rather than arrays.

[liuchihl]

Thanks for the suggestions, I will try that out!