I believe there is a bug in deepspeed activation checkpointing as partition activations feature does not work.