DeepSpeed

87e4c8c8 - Make PipelineModule.set_checkpoint_interval actually change the interval (#8178)

Make PipelineModule.set_checkpoint_interval actually change the interval (#8178) ## Problem `PipelineModule.set_checkpoint_interval()` does not change the activation checkpoint interval. ```python 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: ```python 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: ```python 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: ```python self.is_checkpointable_results = [] ``` ```python 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. --------- Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>