Move the gradients of torch based params to CPU before Adam

[zhuzilin]

The CPUAdam requires the torch based params to be on CPU before updation, otherwise we need to move the momentum and variance of the tensor to GPU. There are 2 design I could think of:

1. Load the torch based params to CPU the entire forward and backward computation and offload them in the loop of adam.
2. Load them to GPU before modules and offload after modules for both forward and backward.

This PR implements the latter design, which may introduce more offloading overhead, but more consistent with other parts of the current design.

[feifeibear]

The CPUAdam requires the torch based params to be on CPU before updation, otherwise we need to move the momentum and variance of the tensor to GPU. There are 2 design I could think of:

1. Load the torch based params to CPU the entire forward and backward computation and offload them in the loop of adam.
2. Load them to GPU before modules and offload after modules for both forward and backward.

This PR implements the latter design, which may introduce more offloading overhead, but more consistent with other parts of the current design.

I currently have no specific idea about managing torch-based tensors in heterogenous memory (HM) space.
If torch and chunk-based tensor exist in HM concurrently, we have to reconfigure the class Manager and take into consideration the footprint of torch-based model data.

[feifeibear]

The CPUAdam requires the torch based params to be on CPU before updation, otherwise we need to move the momentum and variance of the tensor to GPU. There are 2 design I could think of:

1. Load the torch based params to CPU the entire forward and backward computation and offload them in the loop of adam.
2. Load them to GPU before modules and offload after modules for both forward and backward.

This PR implements the latter design, which may introduce more offloading overhead, but more consistent with other parts of the current design.

Maybe we can still use chunk to manage torch based tensors. Instead the communication logic should be different for torch-based and patrickstar-based tensors?

[zhuzilin]

After offline discussion, we'll probably remove the torch_scope, which is mainly for introducing MoE layers. Instead, we 'll try to separate the communication and storage and let user customize the communication setting for certain layers. Therefore, close this PR.