PR Category

Auto Parallel

PR Types

New features

Description

• support shard optimizer for pir mode.
  • the optimizer op in stage1:

      (%0) = "pd_op.data" () {dtype:(pd_op.DataType)float32,name:"embedding_0.w_0.dist_beta2_pow_acc_0",...} : () -> pd_dist.tensor<1xf32, mesh_shape:[2,2],process_ids:[0,1,2,3],dims_mappings:[-1]>
      (%1) = "pd_op.data"() {dtype:(pd_op.DataType)float32,name:"embedding_0.w_0.dist_beta1_pow_acc_0",...} : () -> pd_dist.tensor<1xf32, mesh_shape:[2,2],process_ids:[0,1,2,3],dims_mappings:[-1]>
      (%2) = "pd_op.data" () {dtype:(pd_op.DataType)float32,name:"embedding_0.w_0.dist_moment2_0",...} : () -> pd_dist.tensor<32000x1024xf32, mesh_shape:[2,2],process_ids:[0,1,2,3],dims_mappings:[0,1]>
      (%3) = "pd_op.data" [id:963] () {dtype:(pd_op.DataType)float32,name:"embedding_0.w_0.dist_moment1_0",...} : () -> pd_dist.tensor<32000x1024xf32, mesh_shape:[2,2],process_ids:[0,1,2,3],dims_mappings:[0,1]>
      (%1147, %1148, %1149, %1150, %1151, %1152) = "pd_op.adamw_"(%171, %906, %907, %3, %2, %1, %0, <<NULL VALUE>>, <<NULL VALUE>>, %1144, %1145, %1146) {...omit attribute} :
           ( # operand type 
                # param
                pd_dist.tensor<32000x1024xf32, mesh_shape:[2,2],process_ids:[0,1,2,3],dims_mappings:[0,1]>, 
                # grad
                pd_dist.tensor<32000x1024xf32, mesh_shape:[2,2],process_ids:[0,1,2,3],dims_mappings:[0,1]>, 
                # learning rate
                pd_dist.tensor<f32, mesh_shape:[2,2],process_ids:[0,1,2,3],dims_mappings:[]>,
                # moment1
                pd_dist.tensor<32000x1024xf32, mesh_shape:[2,2],process_ids:[0,1,2,3],dims_mappings:[0,1]>,
                 # moment2
                pd_dist.tensor<32000x1024xf32, mesh_shape:[2,2],process_ids:[0,1,2,3],dims_mappings:[0,1]>, 
                # beta1_pow
                pd_dist.tensor<1xf32, mesh_shape:[2,2],process_ids:[0,1,2,3],dims_mappings:[-1]>, 
                #beta2_pow
                pd_dist.tensor<1xf32, mesh_shape:[2,2],process_ids:[0,1,2,3],dims_mappings:[-1]>,
                 #master_param
                <<NULL TYPE>>,
                # skip_update
                <<NULL TYPE>>,
                # beta1
                pd_dist.tensor<1xf32, mesh_shape:[2,2],process_ids:[0,1,2,3],dims_mappings:[-1]>,
                # beta2
                pd_dist.tensor<1xf32, mesh_shape:[2,2],process_ids:[0,1,2,3],dims_mappings:[-1]>,
                # epsilon
                pd_dist.tensor<1xf32, mesh_shape:[2,2],process_ids:[0,1,2,3],dims_mappings:[-1]>) -> 
              # ouptuts type
                pd_dist.tensor<32000x1024xf32, mesh_shape:[2,2],process_ids:[0,1,2,3],dims_mappings:[0,1]>, 
                pd_dist.tensor<32000x1024xf32, mesh_shape:[2,2],process_ids:[0,1,2,3],dims_mappings:[0,1]>, 
                pd_dist.tensor<32000x1024xf32, mesh_shape:[2,2],process_ids:[0,1,2,3],dims_mappings:[0,1]>, 
                pd_dist.tensor<1xf32, mesh_shape:[2,2],process_ids:[0,1,2,3],dims_mappings:[-1]>,
                pd_dist.tensor<1xf32, mesh_shape:[2,2],process_ids:[0,1,2,3],dims_mappings:[-1]>, <<NULL TYPE>>
      # reshard parameter to origin dist attr
     (%1153) = "pd_op.c_allgather"  (%1147) {} : 
              (pd_dist.tensor<32000x1024xf32, mesh_shape:[2,2],process_ids:[0,1,2,3],dims_mappings:[0,1]>) -> 
               pd_dist.tensor<32000x1024xf32, mesh_shape:[2,2],process_ids:[0,1,2,3],dims_mappings:[-1,1]>
      # update global parameter
      "pd_op.assign_out_" (%1153, %168) {} :
               (pd_dist.tensor<32000x1024xf32, mesh_shape:[2,2],process_ids:[0,1,2,3],dims_mappings:[-1,1]>, 
                pd_dist.tensor<32000x1024xf32, mesh_shape:[2,2],process_ids:[0,1,2,3],dims_mappings:[-1,1]>)

• support shard startup program.
• support p_to_s reshard function.
• fix the bug in r_to_s reshard function.
• add monkey_patch_value_in_dist called in distribute init function.

Other

pcard-67164