fix partition activations issue when mp=2 and pp=2 (#1589)
* fix partition activations issue when mp=2 and pp=2
* change util function input and fix pre-commit errors
* move print_backward_tensors() to debug.py
Co-authored-by: Olatunji Ruwase <olruwase@microsoft.com>