Add dygraph sharding stage2

[Baibaifan]

PR types

New features

PR changes

Others

Describe

Add dygraph sharding stage2

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

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

# wrap model & optimizer 
model = model_class(...)
oss_optimizer = ShardingOptimizer(params=model.parameters(), optim=optimizer, group=group)
model = ShardingStage2(model, oss_optimizer, group=group)

# use optimizer as normal
img, label = data
label.stop_gradient = True
img.stop_gradient = True
out = model(img)

loss = paddle.nn.functional.cross_entropy(input=out, label=label)
oss_optimizer.step()
oss_optimizer.clear_grad()

1.Accuracy test
dp2 fp32 - sharding stage2 fp32

dp2 fp32 - sharding stage2 fp16

2.Performance testing
stage2 buffer size=8Mb
0.31B参数 GPT-3 模型
gbs=16，mbs=2：dp(9412 tokens/s) ; stage1(11440 tokens/s) ; stage2(11992 tokens/s)

0.31B参数 GPT-3 模型
gbs=4，mbs=2：dp(9723 tokens/s) ; stage1(10244 tokens/s) ; stage2(10451 tokens/s)

3. GPU memory test
   GPT-3 0.31B sharding stage1 fp32
   
   GPT-3 0.31B sharding stage2 fp32
   
   stage1 memory peak：6899 stage2 memory peak：6353 reduced 546Mb Optimization and promotion 8%.

[paddle-bot-old]

✅ This PR's description meets the template requirements!
Please wait for other CI results.

[paddle-bot-old]

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

[XieYunshen]

LGTM for set_tests_properties(test_dygraph_sharding_stage2 PROPERTIES TIMEOUT 120)

[ForFishes]

LGTM