support offload in sharding stage2

[haohongxiang]

PR types

Function optimization

PR changes

Others

Describe

Support offload, grad_clip and loss_scaler in dygraph sharding stage2
Optimize the performance of offload in PR-38064

1. 用户手册

import paddle
from paddle.distributed.fleet.meta_optimizers.dygraph_optimizer.sharding_optimizer_stage2 import ShardingOptimizerStage2
from paddle.distributed.fleet.meta_parallel.sharding.sharding_stage2 import ShardingStage2
from paddle.distributed.fleet.meta_parallel.sharding.sharding_utils import ShardingScaler

fleet.init(is_collective=True)
group = paddle.distributed.new_group([0, 1])

# state model and optimizer
model = model_class(...)
clip = paddle.nn.ClipGradByGlobalNorm(clip_norm=1.0)
optimizer = paddle.optimizer.AdamW(
        parameters=model.parameters(),
        learning_rate=0.001,
        weight_decay=0.00001,
        grad_clip=clip,
        multi_precision=True)

# convert to pfp16 model
model = paddle.amp.decorate(models=model, level='O2', save_dtype='float32')
scaler = paddle.amp.GradScaler(init_loss_scaling=32768)
scaler = ShardingScaler(scaler, group)

# convert to sharding_stage2 model and optimizer
optimizer = ShardingOptimizerStage2(params=model.parameters(), optim=optimizer, group=group, offload=True)
model = ShardingStage2(model, optimizer, group=group)

# forward, backward and optimization
img, label = data
label.stop_gradient = True
img.stop_gradient = True
out = model(img)
loss = paddle.nn.functional.cross_entropy(input=out, label=label)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
model.clear_gradients()

2. 精度验证
   PaddleNLP GPT-3模型，sharding stage2+pfp16 with/without offload：

3. 显存优化
   1> PaddleNLP GPT-3模型 0.31B参数量
   单机两卡，sharding stage2+pfp16 without offload，峰值显存为 5319 MiB，显存变化曲线为：

单机两卡，sharding stage2+pfp16 with offload，峰值显存为 3137 MiB（节省 2182 MiB，约 41%），显存变化曲线为：

2> PaddleNLP GPT-3模型 1.02B参数量
单机两卡，sharding stage2+pfp16 without offload，峰值显存为 11941 MiB
单机两卡，sharding stage2+pfp16 with offload，峰值显存为 5369 MiB（节省 6572 MiB，约 55%）

[paddle-bot-old]

Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

[Baibaifan]

这个也改成.value().get_tensor()吧

[haohongxiang]

好的

[ForFishes]

这个地方会增加显存，需要先释放param，在shareddata master参数。

[haohongxiang]

好的

[ForFishes]

LGTM