[BUG] Partitioned activation checkpointing does not reduce memory usage

Hi, I have an issue with partitioned activation checkpointing.
This is my github repo:
https://github.com/andrasiani/deepspeed_lightning_gpt_partition_activations_checkpointing/tree/master

When I run activation_checkpointing without partitioning I get the same memory usage as in the case when partitioning is enabled.
deepspeed/runtime/activation_checkpoing/checkpointing.py

I checked and partition_activations() method shards the tensor correctly,
the backward all_gather method also seems to work correctly https://github.com/andrasiani/deepspeed_lightning_gpt_partition_activations_checkpointing/blob/master/checkpointing_.py#L663,
the get_partitioned_activations_for_backward() seems to only save partitions.

Somewhere I suspect that the whole checkpointed layer is stored in memory after forward until backward pass, instead of just a partition of it.

This is the checkpointing script I use:
https://github.com/andrasiani/deepspeed_lightning_gpt_partition_activations_checkpointing/blob/master/checkpointing_.py#L521
checkpointing_ custom.txt

and I call it like this:
deepspeed.checkpointing.configure(
mpu_=None,
partition_activations=True,
contiguous_checkpointing=False,
checkpoint_in_cpu=False,
profile=checkpoint_config.get("profile"),
)
mpu object is set to None in pytorch lightning - this is a pytorch lightning bug I think - to address this inside my custom script I initialize an mpu object, so there should be no problem.
https://github.com/andrasiani/deepspeed_lightning_gpt_partition_activations_checkpointing/blob/master/checkpointing_.py#L521

Start by reproducing the reported comparison with the linked checkpointing script and configuration, then inspect deepspeed/runtime/activation_checkpoing/checkpointing.py, especially partition_activations(), get_partitioned_activations_for_backward(), and the backward all_gather path. Done means partitioned activation checkpointing demonstrably uses less memory than the unpartitioned configuration while preserving backward behavior.