Make PipelineModule.set_checkpoint_interval actually change the interval

[vineethsaivs]

Problem

PipelineModule.set_checkpoint_interval() does not change the activation checkpoint interval.

def set_checkpoint_interval(self, interval):
    assert interval >= 0
    self.checkpoint_interval = interval

Every reader uses self.activation_checkpoint_interval: forward() branches on it and steps the layer loop by it (module.py L370-379), _precompute_checkpointable_values() keys its cache on it (L223-230), and PipelineEngine assigns it directly (pipe/engine.py L212). self.checkpoint_interval is read nowhere in the repo, so the call assigns a dead attribute and the schedule silently stays as it was.

Assigning the right attribute is not sufficient on its own, because the recompute path it feeds is also broken:

def _precompute_checkpointable_values(self):
    if self.activation_checkpoint_interval > 0 and self.is_checkpointable_results_interval != self.activation_checkpoint_interval:
        num_layers = len(self.forward_funcs)
        self.interval_was_zero = False
        for start_idx in range(0, num_layers, self.activation_checkpoint_interval):
            ...
            self.is_checkpointable_results.append(self._is_checkpointable(funcs))
        self.is_checkpointable_results_interval = self.activation_checkpoint_interval

The is_checkpointable_results_interval != activation_checkpoint_interval guard exists precisely so the values are recomputed when the interval changes, but the loop appends to self.is_checkpointable_results without clearing it, so results computed for the previous interval stay at the front of the list.

forward() then pairs the layer blocks with that list positionally:

for start_idx, is_checkpointable_result in \
    zip(range(0, num_layers, self.activation_checkpoint_interval), self.is_checkpointable_results):

so each block is checkpointed according to a decision made for a different partitioning of the layers.

Reproduction

8 layers, the first 4 without parameters and the last 4 with, so _is_checkpointable genuinely differs per block. Going from interval 4 to interval 1:

interval=4  results: [False, True]

set_checkpoint_interval(1) -> activation_checkpoint_interval = 4     # unchanged
  results: [False, True]                                            # never recomputed

# assigning activation_checkpoint_interval directly, the way PipelineEngine does:
  results: [False, True, False, False, False, False, True, True, True, True]   # 10 entries for 8 blocks
  expected:               [False, False, False, False, True, True, True, True]

forward blocks : [(0, False), (1, True), (2, False), (3, False), (4, False), (5, False), (6, True), (7, True)]
expected       : [(0, False), (1, False), (2, False), (3, False), (4, True), (5, True), (6, True), (7, True)]

Blocks 1, 4 and 5 are checkpointed against the wrong decision: block 1 holds no parameters and is checkpointed anyway, blocks 4 and 5 hold parameters and are not.

Fix

Clear the cached results before recomputing them, and have the setter assign activation_checkpoint_interval and rebuild the cache:

self.is_checkpointable_results = []

def set_checkpoint_interval(self, interval):
    assert interval >= 0
    self.activation_checkpoint_interval = interval
    self._precompute_checkpointable_values()

The setter has to do both. Assigning the interval alone would leave forward() zipping the new, longer block range against a list still sized for the old interval, and zip stops at the shorter one, so trailing layer blocks would be dropped from the forward pass entirely.

Nothing changes for the normal path: PipelineEngine assigns the interval once and calls _precompute_checkpointable_values() while the cache is still empty, so clearing an empty list is a no-op and the guard still skips the recompute when the interval is unchanged.

Testing

TestPipeModuleCheckpointInterval in tests/unit/pipe/test_pipe_module.py builds a PipelineModule at interval 4, calls set_checkpoint_interval(1), and asserts the interval is updated and the results match a module constructed at interval 1 directly. It fails on master on both assertions (the interval stays 4, and the results stay [False, True]) and passes with the fix. The mixed ReLU/Linear model is deliberate: with a uniformly parameterised model only the length of the list is wrong, and the misalignment would not show.

yapf and flake8 are clean on both changed files.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 6a0966a8dc

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Sync engine state when enabling checkpointing

When this setter is used to change a module from interval 0 to a positive interval while the PipelineEngine config still has the default interval 0, it now enables the reentrant checkpoint path but leaves PipelineEngine._config.pipeline['activation_checkpoint_interval'] unchanged. The load path in PipelineEngine._exec_load_micro_batch() only marks first-stage inputs as requires_grad when that config value is >0 (deepspeed/runtime/pipe/engine.py:904-916), so the default reentrant checkpoint function runs with detached inputs and the first checkpointed segment is not connected to autograd, dropping or failing gradients for that training run. Either keep the engine/config flag in sync or have the engine key this requires_grad decision off the module interval.

Useful? React with 👍 / 👎.

[tohtana]

Hi @vineethsaivs,
Thank you for submitting this PR! I think the comment from Codex bot is reasonable. Can you address it?

[vineethsaivs]

Thanks @tohtana. Agreed, the Codex comment is a real bug and I have addressed it in cb8dae6.

I checked it against the source before fixing, and it holds:

• _exec_load_micro_batch gates the requires_grad marking on self._config.pipeline['activation_checkpoint_interval'] > 0 (engine.py:904 and :914)
• but the pipeline config only ever seeds the module, once, in __init__ (engine.py:211-212)
• while PipelineModule.forward() branches on self.activation_checkpoint_interval, the module attribute (module.py:370)

So once the setter really changes the interval, the module starts checkpointing while the config still reads 0, and under reentrant checkpointing the first stage's inputs stay detached. The first checkpointed segment is then cut off from autograd.

Codex offered two options: keep the config in sync, or key the decision off the module. I went with the second, for two reasons.

The first is that the module is already the source of truth everywhere else. engine.py:221 decides whether to call _precompute_checkpointable_values() off self.module.activation_checkpoint_interval, and forward() branches on the same attribute. Syncing the config would add a third copy of the same state rather than remove one.

The second is that this hole is not actually new. PipelineModule(layers=..., activation_checkpoint_interval=N) under a pipeline config of 0 already checkpoints, via the module attribute, while the engine already declines to mark inputs. My setter made an existing inconsistency reachable a second way. Keying off the module closes both.

I also had to fold use_reentrant into it rather than drop it: the config only defaults that flag inside the interval > 0 branch, so reading it directly would be stale for exactly the same reason. The module's activation_checkpoint_func is what use_reentrant ultimately selects, so the helper tests that instead, which is equivalent on the config-driven path and correct on the setter path.

def _reentrant_activation_checkpointing(self):
    if self.module.activation_checkpoint_interval <= 0:
        return False
    return self.module.activation_checkpoint_func is not ds_checkpointing.non_reentrant_checkpoint

Two tests added to TestPipeModuleCheckpointInterval: one starts from a config with checkpointing disabled and asserts the flip happens only after set_checkpoint_interval(), the other pins that use_reentrant: False still does not mark inputs, so the decision tracks the function and not just the interval. yapf clean, and flake8 under the repo's own .flake8 selection is clean.