use non_reentrant_checkpoint fix requires_grad of input must be true for activation checkpoint layer in pipeline train.

[inkcherry]

from original PR
#4128 rebase on feature proposed in
#4118. should be reviewed after #4118 merged.

• set config ['pipeline']['use_reentrant']=False to open it, it won't affect the original stable workload.
• add cifar10 train ut and passed, it seems that there is no megatron-deepspeed pipeline train ut for reference. I have verified it locally（model parameters update process is consistent), and can still reduce a small memory on rank0（stage0, also the most memory occupied rank) on some 3D configurations.
  @tohtana @tjruwase @hughpu

[tohtana]

@inkcherry Thank you for submitting this as a new PR! Sorry for the delay of my review.
It is great that the change looks much simpler now. But the new tests failed. Can you check it?

[inkcherry]

Hi ,@tohtana
Thanks for you replay! But I couldn't reproduce the errors that occurred in the ci workflow, launching
the test_pipe.py file or test_pipe_use_reentrant function separately in my cuda device is always passed.
could this ut can be passed separately on nv-torch-latest-v100?
hope to get some suggestions. Thanks!

[tohtana]

Hi @inkcherry,

I could reproduced the same error (AssertionError on line 109). How did you run the test? I ran the following:

pytest test_pipe.py::TestPipeCifar10::test_pipe_use_reentrant[topo_config0]

[inkcherry]

Hi @inkcherry,

I could reproduced the same error (AssertionError on line 109). How did you run the test? I ran the following:

pytest test_pipe.py::TestPipeCifar10::test_pipe_use_reentrant[topo_config0]

Hi, @tohtana, very Thanks for helping me find the reason~ seems my testing environment version is a bit higher than ci.
It can be reproduced through the following two dockers+install deepspeed

pytorch/pytorch:2.0.1-cuda11.7-cudnn8-devel
(same version to ci, always failed)
nvcr.io/nvidia/pytorch:23.07-py3
(higher version than ci, always passed)

I have reduced the check level from params equal to loss convergence and can be passed currently. params equal check could be reused when the version is updated.

[tohtana]

Thank you @inkcherry, the changes look good to me and I approved this PR.
It is not surprising that the updated parameters do not exactly match without setting options for reproducibility.
Let's merge this PR after the CI tests completed.