Manual rematerialization about alpa

[merrymercy]

@merrymercy Test auto-sharing

[merrymercy]

@zhuohan123 Add the new manual API

[zhuohan123]

@zhuohan123 Add the new manual API

#164
See https://github.com/parax-project/parax/blob/777a3427a2a1323f20077b405aad577d2b7e4b12/tests/test_pipeline_gradient_accumulation.py#L338-L343 for usage

[ZYHowell]

#165
See https://github.com/parax-project/parax/blob/a71e9e203c3a82b8d77a4a4f5d3d5d6b3cdb3551/tests/test_pipeline_gradient_accumulation.py#L133-L141

To use remat without pipeline parallelism, it has an arg use_pipeline. However, remat is used explicitly, then its input should also have pipeline marker: the arg above only control output behavior.

[merrymercy]

#166
We can now use BertConfig to enable remat, which is the most manual way and modifies the model definition

BertConfig(gradient_checkpointing=True, ...)