indiscriminately clearing checkpoints with `SingleMemoryStorageSchedule` corrupts the adjoints

[sghelichkhani]

Second part of an issue we have encountered g-adopt/g-adopt#237, consider the case of this reproducer. When using no scheduler by:

python minimal_clearing_cache.py  none

we get:

using scheduler: NoneType
        J1: 12.520709541433629
        dJdm 1: 3205.301642607008

But when using a scheduler by:

python minimal_clearing_cache.py memory

we get:

using scheduler: SingleMemoryStorageSchedule
        J1: 12.520709541433629
        dJdm 1: 73765.11424588156

Cause

SingleMemoryStorageSchedule explicitly deletes every checkpoint that did not appear in the previous step’s adjoint_dependencies:

pyadjoint/pyadjoint/checkpointing.py

Lines 362 to 364 in a4f940a

	if isinstance(self._schedule, SingleMemoryStorageSchedule):
	if step > 1 and var not in self.tape.timesteps[step - 1].adjoint_dependencies:
	var._checkpoint = None

If a variable is part of a long-range dependency (e.g. is used every third step), it can be missing from the immediately-preceding dependency set even though it is still required later. The checkpoint is therefore cleared and the reverse pass reconstructs an incorrect value, giving a wrong gradient.

Considering the comment saying “Handle the case for SingleMemoryStorageSchedule”, I would have thought there is a reason for this, but no explanation is given.

Ideas (!?)

• Remove the special-case clearing and rely on the revised dependency machinery to decide what can be freed, or

• Only clear if the variable is guaranteed not to re-appear later (e.g. by inspecting timesteps[step:].adjoint_dependencies) ?

[dham]

I'm less clear what the issue here is. Long-range dependencies are already explicitly catered for when adjoint dependencies are recorded. See:

pyadjoint/pyadjoint/tape.py

Line 307 in a4f940a

for step in self.timesteps[last_used + 1:]:

.

If there is an adjoint dependency missing then the thing to look for is why it's not being recorded there.

[angus-g]

I wonder if it's something about the eager addition of the block's output to the tape? When the first BlockVariable inside the skipped block (i.e. the ProjectBlock in the case of the reproducer above) is added as an output, it gets added as an adjoint dependency via BlockVariable.will_add_as_output():

pyadjoint/pyadjoint/block_variable.py

Lines 73 to 83 in a4f940a

	def will_add_as_output(self):
	tape = get_working_tape()
	self.creation_timestep = tape.latest_timestep
	self.last_use = self.creation_timestep
	overwrite = self.output._ad_will_add_as_output()
	overwrite = bool(overwrite)
	if not overwrite:
	self._checkpoint = None
	if tape._eagerly_checkpoint_outputs:
	self.save_output()
	tape.add_to_adjoint_dependencies(self, self.last_use - 1)

From what I can see, self.last_use can only ever be self.creation_timestep, so it will only be an adjoint dependency of the current step (kind of obviously...).

When we're on step 2 and we hit this loop:

pyadjoint/pyadjoint/checkpointing.py

Lines 354 to 365 in a4f940a

	for var in current_step.checkpointable_state - to_keep.union(self._global_deps):
	# Handle the case where step is 0
	if step == 0 and var not in current_step._checkpoint:
	# Ensure initialisation state is kept.
	self._keep_init_state_in_work = True
	break
	# Handle the case for SingleMemoryStorageSchedule
	if isinstance(self._schedule, SingleMemoryStorageSchedule):
	if step > 1 and var not in self.tape.timesteps[step - 1].adjoint_dependencies:
	var._checkpoint = None
	continue

The block variable associated with r (from step 0) is in the checkpointable state (because it's needed to project onto u). But because it was added in step 0, it's not in the adjoint dependencies of step 1, so its checkpoint gets cleared.

It kind of seems like the logic for managing the adjoint dependencies is correct, since it doesn't make sense to push r as an adjoint dependency past its creation. It just seems like the checkpoint clearing is too eager, and it isn't documented why it needs to apply specifically to the SingleMemoryStorageSchedule...

[jrmaddison]

pyadjoint/pyadjoint/tape.py

Lines 274 to 275 in 8a9db39

	for step in self.timesteps[last_used + 1:]:
	step.checkpointable_state.add(block_var)

only loops over steps which exist on the tape, but those are added dynamically, I think here when checkpointing is enabled

pyadjoint/pyadjoint/checkpointing.py

Line 223 in 8a9db39

self.tape.get_blocks().append_step()

Even if the steps did all exist, looping over them would be quite inefficient, and I'm not sure it would have the intended behaviour in any case.

By my reading, with SingleMemoryStorageSchedule the code is attempting to free storage for variables which are not needed by any later steps. Maybe that's possible using block variable version information: free if the variable is not an adjoint dependency and if a later block variable for the same backend variable already exists on the tape? If that's not possible the behaviour of SingleMemoryStorageSchedule would need to be the same as 'no schedule'.

[sghelichkhani]

Yeah, I’ve been scratching my head trying to figure out a sensible way to re-add or revise dependencies after the tape has been generated, but I’m honestly not seeing how that could work in practice — and this is definitely drifting into stuff that I don't know much about.

But it does seem like we’re all circling around the same core issue? this particular clearing logic in SingleMemoryStorageSchedule feels like the main culprit. Unless I’m missing something, that behaviour might just be too aggressive?