fix act recomp test

microsoft · Dec 18, 2024 · b212177 · b212177
1 parent 6334229
commit b212177
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/unit/runtime/activation_checkpointing/test_activation_checkpointing.py b/tests/unit/runtime/activation_checkpointing/test_activation_checkpointing.py
@@ -271,7 +271,7 @@ def test_gpt2_checkpointable_layers(self):
 
         # Create a simple topology for testing
         from deepspeed.runtime.pipe.topology import PipeModelDataParallelTopology
-        topo = PipeModelDataParallelTopology(num_pp=1, num_mp=1)
+        topo = PipeModelDataParallelTopology(num_pp=1, num_mp=1, num_dp=1)
 
         # Create test classes that we want to checkpoint
         class TestTransformerLayer(torch.nn.Module):