[Dygraph]Add sharding stage3 offload

[Baibaifan]

PR types

New features

PR changes

Others

Describe

Add sharding stage3 offload.

import paddle
from paddle.distributed.fleet.meta_parallel.sharding.sharding_stage3 import ShardingStage3

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

# wrap model
model = model_class(...)
model = ShardingStage3(model, optimizer=optimizer, group=group, offload=True)

# 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)
loss.backward()
optimizer.step()
optimizer.clear_grad()
# Get all parameter from parameter slice
# If parameter need to convert to cpu, please add convert2cpu=True
model.get_all_parameters(convert2cpu=True)

stage3 and stage3 offload fp32 GPT 117M

[paddle-bot-old]

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

[ForFishes]

LGTM