[Sharding] add new features

[JZ-LIANG]

PR types

New features

PR changes

OPs

Describe

1. fixed bug
   op role
   add_sync_comm_for_test: fixed bug when clone test prog
2. add new features
   comm_analyse: calculating communication volume
   sharding_save_persistables: save sharding model

usage:

import paddle
import paddle.fluid as fluid
import fleetx as X
from paddle.distributed.fleet.meta_optimizers.sharding.utils import add_sync_comm_for_test, sharding_save_persistables, comm_analyse

dist_strategy = fleet.DistributedStrategy()
dist_strategy.sharding = True
dist_strategy.sharding_configs = {
    "fuse_broadcast_MB": 32,
}

model = X.applications.Resnet50(data_layout=args.data_layout)
optimizer = fluid.optimizer.Momentum(
    learning_rate=lr,
    momentum=args.momentum,
    regularization=fluid.regularizer.L2Decay(args.weight_decay))
optimizer = fleet.distributed_optimizer(optimizer, strategy=dist_strategy)
optimizer.minimize(model.loss)

# Clone test program
# when use sharding, test prog clone should be performed after optimizer.minimize(model.loss)
model.test_prog = model.main_prog.clone(for_test=True)
add_sync_comm_for_test(model.test_prog, dist_strategy)

place = fluid.CUDAPlace(int(os.environ.get('FLAGS_selected_gpus', 0)))
exe = fluid.Executor(place)
exe.run(model.startup_prog)

# Analyse COMM 
comm_analyse(fluid.default_main_program())

# Load model
# we could use the orign load_persistables for sharding model
# make sure the dirname contains all the param files for every rank
dirname="/path/to/load_model"  
paddle.fluid.io.load_persistables(exe, dirname, main_program=model.main_prog, filename=None)

# Training
for epoch in range(10):
......

# Save model
# every rank should execute the following function to save a complete sharding model, 
# unlike data parallelism where only the rank0 handles the model saving
dirname="/path/to/save_model"  
sharding_save_persistables(exe, dirname, main_program=model.main_prog, filename=None)

[paddle-bot-old]

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

[fuyinno4]

change the name