Slowdown in GPU runs with 0.9.0 (Oceananigans 0.104.2) vs 0.8.7 (Oceananigans 0.100.7)

[GodotMisogi]

I'm noticing a slowdown when running the one-degree global ocean sea-ice example code (with appropriate function name updates) on ClimaOcean v0.9.0 with Oceananigans v0.104.2, compared to v0.8.7 with v0.100.7, respectively.

This test was run on TACC Vista using Julia 1.12.4 with both configurations on CUDA. The log files are attached, which indicate that the older version ran 2-3x faster. Is it because a patch to WindowedTimeAverage was pushed that works concurrently with checkpointing?

older_versions.txt
newer_versions.txt

[GodotMisogi]

Probably has to do with fixes related to WindowedTimeAverage and checkpointing updates. My (wild) guess is that the previous version was running faster due to not performing the time-averaging as expected; need to diagnose further.

[simone-silvestri]

From version 0.8.7 to version 0.9.0 we switched the ocean model to use a Runge-Kutta timestepper instead of the default (quasi) Adams Bashforth second order. You can see it in the output of your model

[ Info: We've built an ocean simulation with model at resolution (360 x 180 x 50):
ocean.model = HydrostaticFreeSurfaceModel{CUDAGPU, ImmersedBoundaryGrid}(time = 0 seconds, iteration = 0)
├── grid: 360×180×50 ImmersedBoundaryGrid{Float64, Periodic, RightConnected, Bounded} on CUDAGPU with 5×5×4 halo
├── timestepper: SplitRungeKuttaTimeStepper # <- Here
├── tracers: (T, S, e)
├─ . . .

old version

[ Info: We've built an ocean simulation with model at resolution (360 x 180 x 50):
ocean.model = HydrostaticFreeSurfaceModel{GPU, ImmersedBoundaryGrid}(time = 0 seconds, iteration = 0)
├── grid: 360×180×50 ImmersedBoundaryGrid{Float64, Periodic, RightConnected, Bounded} on CUDAGPU with 5×5×4 halo
├── timestepper: QuasiAdamsBashforth2TimeStepper # <- Here
├── tracers: (T, S, e)
├─ . . .

RK is a bit more accurate, much more stable, but it costs roughly a factor 3 more than AB2 because it is intended to run at a higher CFL. Probably for very coarse simulations as the 1 degree which is not really CFL limited AB2 is good enough.
If you want to go back to using the AB2 you can pass the timestepper = :QuasiAdamsBashfort2 keyword argument to ocean_simulation

[GodotMisogi]

Right, I should have seen that. Thank you!

[GodotMisogi]

Still dealing with slowdowns when changing the ocean timestepper to AB2. Upon some digging, to confirm, the sea-ice model now uses Split RK3 in v0.9.0 instead of forward Euler as in v0.8.7? I am unable to find an option to modify the timestepper in the sea-ice simulation construction in ClimaOcean.jl; should this argument be added as a PR?

https://github.com/CliMA/ClimaOcean.jl/blob/c70bba88976af798930a03bef12912786ce47f53/src/SeaIces/sea_ice_simulation.jl#L12C1-L25C81

...
function sea_ice_simulation(grid, ocean=nothing;
                            Δt = 5minutes,
                            ice_salinity = 4, # psu
                            advection = nothing, # for the moment
                            tracers = (),
                            ice_heat_capacity = 2100, # J kg⁻¹ K⁻¹
                            ice_consolidation_thickness = 0.05, # m
                            ice_density = 900, # kg m⁻³
                            dynamics = sea_ice_dynamics(grid, ocean),
                            # ADD? timestepper = :SplitRungeKutta3,
                            bottom_heat_boundary_condition = nothing,
                            top_heat_boundary_condition = nothing,
                            phase_transitions = PhaseTransitions(; ice_heat_capacity, ice_density),
                            conductivity = 2, # kg m s⁻³ K⁻¹
                            internal_heat_flux = ConductiveFlux(; conductivity))
...

[GodotMisogi]

Just implemented the optional argument in my fork and ran it with the AB2 timestepper for the ocean and Euler for sea-ice, which I believe matches my older version's config (although the log doesn't state the timestepping scheme for sea-ice, I checked the code, which indicates it runs forward Euler). Still running at 2-3x time compared to v0.8.7 over the 1-day intervals.

new_version_old_steppers.txt

Here's the constructor for reference:

function global_ocean_simulation(arch, Nx, Ny, Nz, depth, scalar, k_skew, k_symmetric, date, stop_time, dt, bathymetry_modifications=Dict())
    # %% Construct grid
    z = ExponentialDiscretization(Nz, -depth, 0; scale=scalar * depth)
    z = Oceananigans.Grids.MutableVerticalDiscretization(z)
    underlying_grid = TripolarGrid(arch; size=(Nx, Ny, Nz), halo=(5, 5, 4), z)

    # %% Bathymetry
    # NOTE: This function is implemented elsewhere for a future PR, but doesn't play a role in this timestepper issue.
    # Can be replaced with regrid_bathymetry.
    bottom_height = process_bathymetry(underlying_grid;
        essential_straits=bathymetry_modifications,
        minimum_depth=10,
        interpolation_passes=10,
        major_basins=2) 

    grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bottom_height);
        active_cells_map=true)

    # %% Closures
    eddy_closure = Oceananigans.TurbulenceClosures.IsopycnalSkewSymmetricDiffusivity(κ_skew=k_skew,
        κ_symmetric=k_symmetric)
    vertical_mixing = ClimaOcean.Oceans.default_ocean_closure()

    # %% Ocean model
    substeps = 70
    free_surface = SplitExplicitFreeSurface(grid; substeps)
    momentum_advection = WENOVectorInvariant(order=5)
    tracer_advection = WENO(order=5)

    @time ocean = ocean_simulation(grid; momentum_advection, tracer_advection, free_surface,
        closure=(
            eddy_closure,
            vertical_mixing,
            # horizontal_viscosity,
        ),
        timestepper=:QuasiAdamsBashforth2, # NOTE: Instead of SplitRungeKutta3
    )

    @info "We've built an ocean simulation with model at resolution ($Nx x $Ny x $Nz):"
    @show ocean.model

    # %% Sea ice model
    @time sea_ice = sea_ice_simulation(grid, ocean; 
        advection=tracer_advection,
        timestepper=:ForwardEuler, # NOTE: Needs fork at https://github.com/GodotMisogi/ClimaOcean.jl/tree/sea-ice-timestep-fix
    )

    @info "We've built a sea-ice simulation with model at resolution ($Nx x $Ny x $Nz):"
    @show sea_ice.model

    # %% Initial condition (data from ECCO)
    date = DateTime(date)
    dataset = ECCO4Monthly()
    ecco_temperature = Metadatum(:temperature; date, dataset)
    ecco_salinity = Metadatum(:salinity; date, dataset)
    ecco_sea_ice_thickness = Metadatum(:sea_ice_thickness; date, dataset)
    ecco_sea_ice_concentration = Metadatum(:sea_ice_concentration; date, dataset)

    # %% Set initial conditions
    set!(ocean.model, T=ecco_temperature, S=ecco_salinity)
    set!(sea_ice.model, h=ecco_sea_ice_thickness, ℵ=ecco_sea_ice_concentration)

    # %% Atmospheric forcing
    @time radiation = Radiation(arch)
    @time atmosphere = JRA55PrescribedAtmosphere(arch;
        # dataset=MultiYearJRA55(),
        # start_date=DateTime(date),
        # end_date=DateTime(date)+Day(Second(stop_time)),
        backend=JRA55NetCDFBackend(80),
        include_rivers_and_icebergs=false,
    )


    # %% Coupled model
    @time coupled_model = OceanSeaIceModel(ocean, sea_ice; atmosphere, radiation)
    @time sim = Simulation(coupled_model; Δt=dt, stop_time=stop_time)

    return sim
end

sim = global_ocean_simulation(...)
run!(sim, 10minutes, 10days) # Warmup
run!(sim, 20minutes, 600days) # Full run

[simone-silvestri]

This is not expected. I ll try to profile in the coming days.

[glwagner]

Probably the biggest change was to time-stepping the hydrostatic model? Are there other changes that should be on our minds?

[simone-silvestri]

I think that is the largest change indeed, but it should have not changed performance, I mean if the performance changed it must be a bug.

[glwagner]

There was also the change to WENO advection, but it would be surprising is that produced a major performance decrease or change at all.

[simone-silvestri]

yeah it should not impact too much since it is only switching an UpwindBiased with a Centered scheme.

[GodotMisogi]

I forgot to mention that I was running writers as well (excluded that from the code). I can confirm that running the above code snippets yields no noticeable differences in performance. If I include the file-writing callback, the performance differs across the two versions. So, it is (likely) related to WindowedTimeAverage. Logs attached; writer snippet below.

new_version_no_writers.txt
old_version_with_writers.txt
old_version_no_writers.txt

function prepare_writers!(sim, date, stop_time, dirpath)
    ocean = sim.model.ocean
    sea_ice = sim.model.sea_ice

    # %% Output quantities
    ocean_outputs = merge(ocean.model.tracers, ocean.model.velocities)

    sea_ice_outputs = merge((
            h=sea_ice.model.ice_thickness,
            ℵ=sea_ice.model.ice_concentration,
            T=sea_ice.model.ice_thermodynamics.top_surface_temperature
        ),
        sea_ice.model.velocities)

    free_surface_outputs = (; η=ocean.model.free_surface.η,)

    # %% Transient fields
    ocean.output_writers[:surface] = JLD2Writer(ocean.model, ocean_outputs;
        schedule=TimeInterval(1days),
        overwrite_existing=true,
        filename=joinpath(dirpath, "fields/ocean"),
    )

    ocean.output_writers[:free_surface] = JLD2Writer(ocean.model, free_surface_outputs,
        schedule=TimeInterval(1days),
        with_halos=false,
        overwrite_existing=true,
        filename=joinpath(dirpath, "fields/free_surface"),
    )

    sea_ice.output_writers[:surface] = JLD2Writer(sea_ice.model, sea_ice_outputs;
        schedule=TimeInterval(1days),
        overwrite_existing=true,
        filename=joinpath(dirpath, "fields/sea_ice"),
    )

    # %% Averages
    schedule = AveragedTimeInterval(30days; window=30days)
    # dates = DateTime(date):Month(1):DateTime(date)+Day(stop_time)
    # secs = datetime2unix.(dates)
    # secs .= secs .- secs[1]  # Make times start at zero
    # schedule = AveragedSpecifiedTimes(secs; window=31days)
    ocean.output_writers[:avg_ocean] = JLD2Writer(ocean.model, ocean_outputs,
        filename=joinpath(dirpath, "fields/monthly_average_ocean"),
        schedule=schedule,
        overwrite_existing=true,
    )

    ocean.output_writers[:avg_free_surface] = JLD2Writer(ocean.model, free_surface_outputs,
        filename=joinpath(dirpath, "fields/monthly_average_free_surface"),
        with_halos=false,
        schedule=schedule,
        overwrite_existing=true,
    )

    sea_ice.output_writers[:avg_ice] = JLD2Writer(sea_ice.model, sea_ice_outputs,
        filename=joinpath(dirpath, "fields/monthly_average_sea_ice"),
        schedule=schedule,
        overwrite_existing=true,
    )

    # %% Checkpoints
    # dates = DateTime(date):Month(1):DateTime(date)+Day(Second(stop_time))
    ocean.output_writers[:checkpointer] = Checkpointer(ocean.model,
        schedule=TimeInterval(1days), 
        # schedule=TimeInterval(182.5days),
        dir=joinpath(dirpath, "checkpoint/"),
        # prefix="",
    )
    
    return nothing
end

sim = global_ocean_simulation(...)
prepare_writers!(sim, ...) # Writer callback
run!(sim, 10minutes, 10days) # Warmup
run!(sim, 20minutes, 600days) # Full run

[glwagner]

Inspecting your code, I see two time averages: one for ocean outputs, and one for sea ice outputs. Do both of them affect performance, or only one?