[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