Issue with DiskCheckpointing via SingleDiskStorageSchedule

[drhodrid]

Following discussion on the G-ADOPT github page, I am porting this issue across to pyadjoint, on the advice of @Ig-dolci. For background, G-ADOPT discussion here -- g-adopt/g-adopt#190.

I have tested a simpler reproducer of our typical mantle convection adjoint cases in 2-D, where we invert for an unknown initial condition. Good news and bad news:

1. If I run a standard Taylor test for our typical style of case with all default checkpointing options (i.e. in memory -- nothing written to disk), I get convergence rates of 2.0 as expected. Good!
2. If I take this case and add SingleMemoryStorageSchedule, convergence remains at 2.0, but there is a substantial speedup. Again good!

So the SingleMemoryStorageSchedule seems to behave as expected. The issues arise with disk checkpointing:

3. If I run a standard Taylor test, now with only "enable_disk_checkpointing" triggered, convergence remains at 2.0.
4. If I modify this case to also use SingleDiskStorageSchedule as is advised, my convergence rate drops. In the reproducer I have, it drops to ~1.9, but I have seen some other cases go below 1.3.

This issue with SingleDiskCheckpointSchedule appears in both serial and parallel runs, but is pretty difficult to reproduce via a minimal reproducer. I have tried all day with no luck. With that in mind I share my setup here:

Serial -- Disk Checkpointing Enabled with SingleDiskCheckpointingSchedule: https://www.dropbox.com/scl/fi/6gmurfrlkj2hairiscijf/Taylor_Serial_SingleDiskCheckpointingSchedule.tar.gz?rlkey=lcs4slwgc12fshbw5s24xsch9

Serial -- Disk Checkpointing Enabled with no Checkpoint Schedule specified:
https://www.dropbox.com/scl/fi/z7e1d1xym6kkuzi8iwayx/Taylor_Serial_EnableDiskCheckpointing.tar.gz?rlkey=myd05wssurg808uurfrozp5pd

Serial -- SingleMemoryCheckpointingSchedule: https://www.dropbox.com/scl/fi/e7puz36orfbwn2oq1p6mj/Taylor_Serial_SingleMemoryCheckpointingSchedule.tar.gz?rlkey=gfby5hum7qmyewkhih1fv5up2

Note convergence rates are printed at the end of the output.dat files in each folder. Aside from the changes to memory/checkpointing, cases are otherwise identical.

To run these case it's as simple as: python inverse.py &> output.dat -- they pick up from a checkpoint, which is also included in each folder.

For me cases take roughly an hour to run. I know this isn't ideal -- but attempts to reproduce in a simpler setup have failed (which I find difficult to understand!).

[drhodrid]

I managed to make a somewhat quicker reproducer. The issue is clearly evident within 30 timesteps, so the 50 timesteps specified in my previous examples was overkill. Everything for the 30 timestep case is here (directory names should be self-explanatory):

https://www.dropbox.com/scl/fi/62g67dpuo0jr190cr7f18/30_Steps.tar.gz?rlkey=bz9am1kzegpr2qktnkxgmeml6

I get convergence of 2.000 with SingleMemoryStorageSchedule and 1.95-1.97 with SingleDiskStorageSchedule. The situation deteriorates further as simulation length (i.e. # timesteps) increases.

[dham]

OK so now you know there is a bug, Taylor test is actually overkill. The different checkpointing schemes should produce identical answers, so you can just compare them to each other. I would start by computing a single gradient and checking that it is different when using disk vs memory. Assuming that is the case, I would spike the adjoint computation to print out the norm of the adjoint variables as each block's adjoint executes, and look for where they start to differ. That way you can home in on the actual point at which something changes.

[Ig-dolci]

I have tracked this issue. Single disk checkpointing should work like single memory checkpointing, the latter is working correctly. I will verify where the single disk checkpointing deviates from the single memory schedule, leading to this issue.

[Ig-dolci]

First, keep in mind the SingleDiskStorageSchedule schedule stores adjoint dependencies on disk.

The issue comes from how tape.timesteps[step].adjoint_dependencies is populated.

To exemplify, consider a forward equation F(u^{n+1}, u^n, c) = 0
where:

• u^{n+1} and u^n are the solution states at timesteps n+1 and n, respectively.
• c is the control.

The adjoint-based derivative depends on the computation of ∂ F \ ∂ c*. Suppose,

∂ F \ ∂ c*= G(u^n, c)

Then, u^n should also be included in tape.timesteps[step].adjoint_dependencies.

Root Cause

Currently, tape.timesteps[step].adjoint_dependencies is populated with Block.outputs, but u^n is not an output of a solver block. It is a Block dependence, which means u^n was not included tape.timesteps[step].adjoint_dependencies. This creates an issue when using SingleDiskStorageSchedule:

• When checkpointing, tape.timesteps[step].checkpointable_state (or forward restart data) are cleaned, and only the block outputs (tape.timesteps[step].adjoint_dependencies) are stored on disk.
• When computing ∂ F \ ∂ c*, the code attempts u^n.block_variable.checkpoint, which is None. That leads to the usage of u^n.block_variable.output data. As a consequence, the Taylor test fails because the gradient is wrong.

A Quick Fix (Discarded)

The simplest fix is to modify SingleDiskStorageSchedule to enable the storage of all timesteps, as suggested in the checkpoint_schedules PR.

Rigtly Populating adjoint_dependencies

We can think about building correctly tape.timesteps[step].adjoint_dependencies during the taping process. The challenge is that, during annotation, the control variable is not yet known. The relevant nodes for adjoint computations are only determined when computing the gradient.

One possible approach is to introduce a Block method that identifies the relevant nodes when taping if the control is provided. This would help construct tape.timesteps[step].adjoint_dependencies correctly.

[jrmaddison]

Follow up to comment firedrakeproject/checkpoint_schedules#73 (comment) on firedrakeproject/checkpoint_schedules#73

The issue (by my reading of the above) is that pyadjoint is being asked to store adjoint dependencies, but instead stores block outputs. Agreed it's difficult to identify what the adjoint dependencies are because you can't apply an activity analysis until you know the controls -- but you can defensively store all potential dependencies of the adjoint, which e.g. could just be the union of all block dependencies. With checkpoint_schedules+pyadjoint that looks like what firedrakeproject/checkpoint_schedules#73 would do, but I think it would fix the bug for the wrong reason, by modifying the schedule to instruct pyadjoint to store more data.

[Ig-dolci]

I see. Yes, Your suggestion looks like the right thing to do.