Fix moe optimizer broadcast (#8813)

Author: DesmonDay

paddlenlp/trainer/trainer.py

@@ -2512,7 +2512,9 @@ def _load_optimizer_and_scheduler(self, checkpoint):
  dist.barrier()
  if self.args.use_expert_parallel:
  opt_state_dict = broadcast_moe_optimizer(
- opt_state_dict, broadcast_dp=not self.args.should_load_sharding_stage1_model
+ opt_state_dict,
+ model_state_dict=self.model.state_dict(),
+ broadcast_dp=not self.args.should_load_sharding_stage1_model,
  )
  else:
  if not self.args.should_load_sharding_stage1_model:

paddlenlp/trainer/utils/helper.py

@@ -229,7 +229,7 @@ def broadcast_dp_optimizer(state_dict):
  return state_dict
 
 
-def broadcast_moe_optimizer(state_dict, broadcast_dp=True):
+def broadcast_moe_optimizer(state_dict, model_state_dict=None, broadcast_dp=True):
 
  try:
  hcg = fleet.get_hybrid_communicate_group()
@@ -242,7 +242,29 @@ def broadcast_moe_optimizer(state_dict, broadcast_dp=True):
  except:
  dp_group = None
  src_rank = 0
- data_parallel_rank = 0
+ data_parallel_rank = dist.get_rank()
+
+ def _filter_sync_optimizer_state(model_state_dict, opt_state_dict):
+ # get sync name
+ sync_vname = []
+ for k, v in model_state_dict.items():
+ if not getattr(v, "no_sync", False):
+ sync_vname.append(v.name)
+
+ filter_opt_state_dict = {"master_weights": {}}
+ filter_opt_state_dict["LR_Scheduler"] = opt_state_dict.get("LR_Scheduler", {})
+ for op_k, op_v in opt_state_dict.items():
+ if op_k not in ["master_weights", "LR_Scheduler"]:
+ for sync_v in sync_vname:
+ if op_k.startswith(sync_v):
+ filter_opt_state_dict[op_k] = op_v
+ break
+ elif op_k == "master_weights":
+ for k, v in op_v.items():
+ for sync_v in sync_vname:
+ if k.startswith(sync_v):
+ filter_opt_state_dict["master_weights"][k] = v
+ return filter_opt_state_dict
 
  def _broadcast_moe_optimizer_state(state_dict):
  # boardcast_keys
@@ -272,9 +294,11 @@ def _broadcast_moe_optimizer_state(state_dict):
  return base_state_dict
 
  if broadcast_dp:
- base_state_dict = broadcast_dp_optimizer(state_dict)
+ filter_opt_state_dict = _filter_sync_optimizer_state(model_state_dict, state_dict)
+ base_state_dict = broadcast_dp_optimizer(filter_opt_state_dict)
  else:
  base_state_dict = _broadcast_moe_optimizer_state(state_dict)
+
  if data_parallel_rank > 0:
  master_weight = state_dict.pop("master_weights", {})
  base_state_dict.update(state_dict)

paddlenlp/transformers/conversion_utils.py

@@ -1284,7 +1284,7 @@ def merge_tensor_parallel(cls, state_dict, config) -> None:
 
  if len(name_action_mappings) > 0:
  for x in name_action_mappings.keys():
- logger.warning(f"key <{x}> need to merge tensor parallel but we can't find in model state.")
+ logger.debug(f"key <{x}> need to merge tensor parallel but we can't find in model state.")
 
  return state_dict_to_save
 
@@ -1318,7 +1318,7 @@ def _resolve_prefix_keys(state_keys_base, state_keys_real, ignore_error=False):
  break
  if key not in state_keys_map:
  if not ignore_error:
- logger.error(f"tensor parallel conversion: could not find name {key} in loaded state dict!")
+ logger.debug(f"tensor parallel conversion: could not find name {key} in loaded state dict!")
  else:
  state_keys_real.remove(state_keys_map[key])