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>
Signed-off-by: Guokai Ma <guokai.ma@intel.com>

Author: 3 people

deepspeed/module_inject/auto_tp.py

@@ -571,7 +571,7 @@ def _replace_module(self, r_module, prev_name='', prev_class_name=''):
             # When using partition_config (custom patterns/presets), use pattern-based routing
             # instead of linear_policies. This keeps all pattern logic centralized here.
             if self.partition_config is not None:
-                full_name = prev_name + '.' + name if prev_name else name
+                full_name = class_name + '.' + name if class_name else name
                 if isinstance(child, nn.Embedding):
                     # Check if embedding matches any pattern
                     param_name = full_name + ".weight"
@@ -588,7 +588,7 @@ def _replace_module(self, r_module, prev_name='', prev_class_name=''):
                         setattr(r_module, name, new_child)
                 else:
                     self.update_mp_params(child)
-                    self._replace_module(child, full_name, class_name)
+                    self._replace_module(child, name, class_name)
             # Traditional path: use linear_policies for type-based routing
             elif child.__class__ in self.linear_policies:
                 setattr(r_module, name, self.linear_policies[child.__class__](child, prev_name + '.' + name,

deepspeed/runtime/activation_checkpointing/checkpointing.py

@@ -1019,7 +1019,7 @@ def _configure_defaults():
 
     PARTITION_ACTIVATIONS = False
     CONTIGUOUS_CHECKPOINTING = False
-    num_layers = False
+    num_layers = None
     CPU_CHECKPOINT = False
     SYNCHRONIZE = False
     PROFILE_TIME = False

tests/unit/module_inject/test_tp_partition_config_path.py

@@ -0,0 +1,128 @@
+# SPDX-License-Identifier: Apache-2.0
+# DeepSpeed Team
+"""Test that partition_config receives correct full hierarchical module paths.
+
+The bug: AutoTP._replace_module built ``full_name`` from ``prev_name`` (the
+immediate parent only) instead of ``class_name`` (the accumulated hierarchical
+path).  Patterns like ``model.layers.0.self_attn.q_proj`` never matched
+because the name was just ``0.self_attn.q_proj``.
+"""
+
+import pytest
+import torch.nn as nn
+
+from deepspeed.module_inject.auto_tp import AutoTP, AutoTPConfig, PartitionType, TPLayerSpec
+
+
+class SubAttn(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.q_proj = nn.Linear(32, 32, bias=False)
+        self.k_proj = nn.Linear(32, 32, bias=False)
+        self.v_proj = nn.Linear(32, 32, bias=False)
+        self.o_proj = nn.Linear(32, 32, bias=False)
+
+
+class DecoderLayer(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.self_attn = SubAttn()
+        self.mlp = nn.Sequential(nn.Linear(32, 64), nn.GELU(), nn.Linear(64, 32))
+
+
+class DummyModel(nn.Module):
+
+    def __init__(self, num_layers=2):
+        super().__init__()
+        self.embed = nn.Embedding(100, 32)
+        self.layers = nn.ModuleList([DecoderLayer() for _ in range(num_layers)])
+        self.head = nn.Linear(32, 100, bias=False)
+
+
+def _build_config():
+    """Partition config that matches q_proj and o_proj via regex."""
+    return AutoTPConfig(layer_specs=[
+        TPLayerSpec(patterns=[r".*\.self_attn\.q_proj"], partition_type=PartitionType.COLUMN),
+        TPLayerSpec(patterns=[r".*\.self_attn\.o_proj"], partition_type=PartitionType.ROW),
+    ])
+
+
+def _capture_matched_names(model, config):
+    """Run _replace_module and capture full_name values that match a spec."""
+    matched_names = []
+    original = AutoTP._replace_with_config
+
+    def capture(self, child, full_name):
+        # Only capture if a spec actually matches
+        param_name = full_name + ".weight"
+        model_type = self._get_model_type() if hasattr(self, '_get_model_type') else None
+        spec = config.find_matching_spec(param_name, model_type)
+        if spec is not None:
+            matched_names.append(full_name)
+        return None
+
+    AutoTP._replace_with_config = capture
+    try:
+        autotp = AutoTP(
+            module=model,
+            all_reduce_linears=[],
+            prefix="model",
+            state_dict=None,
+            linear_layer_setting=None,
+            orig_layer_impl=None,
+            partition_config=config,
+        )
+        autotp._replace_module(model)
+    finally:
+        AutoTP._replace_with_config = original
+    return matched_names
+
+
+def test_partition_config_receives_full_path():
+    """Verify that pattern matching sees the full hierarchical path."""
+    model = DummyModel(num_layers=2)
+    config = _build_config()
+    matched_names = _capture_matched_names(model, config)
+
+    for layer_idx in range(2):
+        assert f"layers.{layer_idx}.self_attn.q_proj" in matched_names, \
+            f"Expected 'layers.{layer_idx}.self_attn.q_proj', got: {matched_names}"
+        assert f"layers.{layer_idx}.self_attn.o_proj" in matched_names, \
+            f"Expected 'layers.{layer_idx}.self_attn.o_proj', got: {matched_names}"
+
+
+def test_no_truncated_paths():
+    """Ensure paths are never truncated to just the immediate parent prefix."""
+    model = DummyModel(num_layers=3)
+    config = _build_config()
+    matched_names = _capture_matched_names(model, config)
+
+    for name in matched_names:
+        assert name.startswith("layers."), \
+            f"Path should start with 'layers.', got: {name}"
+        assert ".self_attn." in name, \
+            f"Path should contain '.self_attn.', got: {name}"
+        # With the bug, paths would be '0.self_attn.q_proj' (only layer index as prefix)
+        assert name.count(".") >= 3, \
+            f"Path should have at least 3 dots (layers.N.self_attn.X_proj), got: {name}"
+
+
+def test_nested_depth_correct():
+    """Verify correct count and paths at 3 layers deep."""
+    model = DummyModel(num_layers=3)
+    config = _build_config()
+    matched_names = _capture_matched_names(model, config)
+
+    expected_count = 3 * 2  # 3 layers × (q_proj + o_proj)
+    assert len(matched_names) == expected_count, \
+        f"Expected {expected_count} matches, got {len(matched_names)}: {matched_names}"
+
+    for layer_idx in range(3):
+        assert f"layers.{layer_idx}.self_attn.q_proj" in matched_names
+        assert f"layers.{layer_idx}.self_attn.o_proj" in matched_names
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

tests/unit/runtime/activation_checkpointing/test_activation_checkpointing.py

@@ -309,3 +309,45 @@ def __init__(self):
         assert model._is_checkpointable([layers[0]]) == True  # ParallelTransformerLayerPipe
         assert model._is_checkpointable([layers[1]]) == True  # GMLPBlock
         assert model._is_checkpointable([layers[2]]) == False  # Linear layer
+
+
+def test_configure_with_contiguous_checkpointing_requires_num_checkpoints():
+    # Regression: ``_configure_defaults`` previously initialized ``num_layers``
+    # to ``False`` while the assert below uses ``is not None``; ``False is not
+    # None`` is True, so the missing-config assert silently passed and a
+    # cryptic ``IndexError`` surfaced later from ``range(num_layers)``. With
+    # the default switched to ``None`` (matching the module-level default),
+    # the helpful assert message fires at the configure() call site.
+    #
+    # ``configure()`` mutates module globals before raising, so snapshot and
+    # restore them around the call to avoid order-dependent failures in other
+    # activation-checkpointing tests sharing the same pytest worker.
+    cp = deepspeed.checkpointing
+    saved = (
+        cp.PARTITION_ACTIVATIONS,
+        cp.CONTIGUOUS_CHECKPOINTING,
+        cp.num_layers,
+        cp.CPU_CHECKPOINT,
+        cp.SYNCHRONIZE,
+        cp.PROFILE_TIME,
+        cp.mpu,
+        cp.deepspeed_checkpointing_enabled,
+    )
+    try:
+        with pytest.raises(AssertionError, match="number of layers"):
+            deepspeed.checkpointing.configure(
+                mpu_=None,
+                partition_activations=True,
+                contiguous_checkpointing=True,
+            )
+    finally:
+        (
+            cp.PARTITION_ACTIVATIONS,
+            cp.CONTIGUOUS_CHECKPOINTING,
+            cp.num_layers,
+            cp.CPU_CHECKPOINT,
+            cp.SYNCHRONIZE,
+            cp.PROFILE_TIME,
+            cp.mpu,
+            cp.deepspeed_checkpointing_enabled,
+        ) = saved