update, test=develop

paddle/fluid/framework/distributed_strategy.proto

@@ -32,6 +32,9 @@ message ShardingConfig {
   optional float fuse_broadcast_MB = 1 [ default = 32.0 ];
   optional bool hybrid_dp = 2 [ default = false ];
   optional int32 sharding_group_size = 3 [ default = 8 ];
+  optional bool as_outer_parallelism = 4 [ default = false ];
+  optional int32 inner_parallelism_size = 5 [ default = 8 ];
+  optional bool use_pipeline = 6 [ default = false ];
 }
 
 message AMPConfig {
@@ -117,6 +120,8 @@ message AsyncConfig {
 
 message PipelineConfig { optional int32 micro_batch = 1 [ default = 1 ]; }
 
+message ModelParallelConfig { optional int32 parallelism = 1 [ default = 1 ]; }
+
 message DistributedStrategy {
   // bool options
   optional Mode mode = 1 [ default = COLLECTIVE ];
@@ -146,6 +151,7 @@ message DistributedStrategy {
   optional bool fp16_allreduce = 25 [ default = false ];
   optional bool sharding = 26 [ default = false ];
   optional float last_comm_group_size_MB = 27 [ default = 1 ];
+  optional bool model_parallel = 28 [ default = false ];
 
   optional RecomputeConfig recompute_configs = 101;
   optional AMPConfig amp_configs = 102;
@@ -158,6 +164,7 @@ message DistributedStrategy {
   optional LambConfig lamb_configs = 109;
   optional AdaptiveLocalSGDConfig adaptive_localsgd_configs = 110;
   optional ShardingConfig sharding_configs = 111;
+  optional ModelParallelConfig model_parallel_configs = 112;
   optional BuildStrategy build_strategy = 201;
   optional ExecutionStrategy execution_strategy = 202;
 }

paddle/fluid/framework/section_worker.cc

@@ -107,7 +107,7 @@ void SectionWorker::TrainFiles() {
     int op_role = op->Attr<int>(std::string("op_role"));
     if (op_role == static_cast<int>(OpRole::kOptimize)) {
       VLOG(3) << "Update: running op " << op->Type();
-      op->Run(*microbatch_scopes_[0], place_);
+      op->Run(*microbatch_scopes_[num_microbatches_ - 1], place_);
       if (gc) {
         DeleteUnusedTensors(*microbatch_scopes_[0], op.get(), unused_vars_,
                             gc.get());

python/paddle/distributed/fleet/meta_optimizers/amp_optimizer.py

@@ -56,9 +56,10 @@ class AMPOptimizer(MetaOptimizerBase):
         # add is_distributed to optimize amp, overlap communication and
         # computation by split the check_finite_and_unscale op.
         is_distributed = self.role_maker._worker_num() > 1
-        if self.user_defined_strategy.sharding:
-            # FIXME(wangxi). sharding failed when split check_finite_and_unscale
-            is_distributed = False
+        #if self.user_defined_strategy.sharding or self.user_defined_strategy.model_parallel:
+        #    # FIXME(wangxi). sharding failed when split check_finite_and_unscale
+        #    # FIXME(JZ-LIANG). To support Sharding-Megatron-AMP, Megatron should follow Sharding's behavior
+        #    is_distributed = False
         self.wrapped_opt._set_distributed(is_distributed)
 
     def _can_apply(self):

python/paddle/distributed/fleet/meta_optimizers/pipeline_optimizer.py

@@ -154,8 +154,10 @@ class PipelineOptimizer(MetaOptimizerBase):
     def __init__(self, optimizer):
         super(PipelineOptimizer, self).__init__(optimizer)
         self.inner_opt = optimizer
-        # we do not allow meta optimizer to be inner optimizer currently
-        self.meta_optimizers_white_list = []
+        self.meta_optimizers_white_list = [
+            "RecomputeOptimizer",
+            "AMPOptimizer",
+        ]
         self.meta_optimizers_black_list = ["GraphExecutionOptimizer", ]
 
     def _set_basic_info(self, loss, role_maker, user_defined_optimizer,