DeepSpeed
3a47d1dd - activation_checkpointing: default num_layers to None so configure() assert fires (#8041)

Commit
32 days ago
activation_checkpointing: default num_layers to None so configure() assert fires (#8041) ## Summary Calling `deepspeed.checkpointing.configure(contiguous_checkpointing=True, partition_activations=True)` *without* setting `num_checkpoints` (or otherwise providing a layer count) is supposed to fail fast with the assert at `deepspeed/runtime/activation_checkpointing/checkpointing.py:1108`: ```python if CONTIGUOUS_CHECKPOINTING: assert num_layers is not None, "Must specify the number of layers with contiguous memory checkpointing" ``` That assert never fires because `_configure_defaults()` (called inside `configure()` at line 1079) initialized: ```python num_layers = False ``` `False is not None` is `True`, so the assert silently passes. The user instead hits a much later cryptic `IndexError` from `range(num_layers)` (lines 399 / 406) or a 0-element allocation from `numel * num_layers` (lines 457 / 461). The module-level default at line 46 is already `num_layers = None`, and every other path that sets it (`set_num_layers`, `config.number_checkpoints`) assigns an integer — only `_configure_defaults` used `False`, which looks like a copy-paste from the surrounding `PARTITION_ACTIVATIONS = False` etc. block. ## Fix ```diff - num_layers = False + num_layers = None ``` One-character change. No callers compare `num_layers` to `False` (downstream uses are `range(num_layers)` and `numel * num_layers`, both requiring an int), so the only path this changes is the broken one: users now get the documented "Must specify the number of layers" assert instead of a downstream `IndexError`. ## Test Adds `test_configure_with_contiguous_checkpointing_requires_num_checkpoints` to the existing `tests/unit/runtime/activation_checkpointing/test_activation_checkpointing.py` (consolidating per `AGENTS.md`, not a new file). It calls `configure(contiguous_checkpointing=True, partition_activations=True)` and asserts the expected `AssertionError` matches `"number of layers"`. On `main` the assert silently passes and the test fails; with this change it passes. ## CI/lint - `pre-commit run --files <changed files>`: all hooks pass (yapf, flake8, `check-torchdist`, `check-license`, codespell, `check-torchcuda`). - DCO `Signed-off-by` on the commit. - No competing PR (`gh pr list --search "num_layers checkpointing OR _configure_defaults OR contiguous_checkpointing assert"`). --------- Signed-off-by: Kymi808 <zeng.kyle13@gmail.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com> Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Author
Parents
Loading