token_dispatcher support expert_num 64 (#10905)

* token_dispatcher support expert_num 64

* token_dispatcher support expert_num 64

Author: blacksheep-Aristotle

paddlenlp/trainer/trainer_utils.py

@@ -1256,6 +1256,7 @@ def download_recovery_ckpt_from_pdc(recovery_checkpoint_path, timeout):
             f"{PDC_DOWNLOAD_ERROR}; Error occurred when trying to download checkpoint from PDC, recovery_checkpoint_path: {recovery_checkpoint_path}, timeout: {timeout}; error details: {PDCErrorMessageMap[result]}"
         )
 
+
 def parse_nccl_config_file(config_dir):
     json_file = Path(config_dir)
     if json_file.exists():

paddlenlp/trainer/training_args.py

@@ -1223,9 +1223,7 @@ def __post_init__(self):
             if self.sharding_parallel_degree == -1:
                 if len(self.sharding) > 0:
                     self.sharding_parallel_degree = world_size // (
-                        tensor_parallel_degree
-                        * sep_parallel_degree
-                        * pipeline_parallel_degree
+                        tensor_parallel_degree * sep_parallel_degree * pipeline_parallel_degree
                     )
 
             sharding_parallel_degree = max(self.sharding_parallel_degree, 1)
@@ -1234,10 +1232,7 @@ def __post_init__(self):
                 self.sharding = []
 
             self.data_parallel_degree = world_size // (
-                sharding_parallel_degree
-                * tensor_parallel_degree
-                * sep_parallel_degree
-                * pipeline_parallel_degree
+                sharding_parallel_degree * tensor_parallel_degree * sep_parallel_degree * pipeline_parallel_degree
             )
 
             if expert_parallel_degree > 1:
@@ -1513,7 +1508,9 @@ def is_segment_parallel_supported():
                 def is_context_parallel_supported():
                     import inspect
 
-                    members = [name for (name, date) in inspect.getmembers(fleet.base.topology.EPHybridCommunicateGroup)]
+                    members = [
+                        name for (name, date) in inspect.getmembers(fleet.base.topology.EPHybridCommunicateGroup)
+                    ]
                     support_cp = "get_context_parallel_world_size" in members
                     if not support_cp:
                         logger.warning("context parallel is not supported!!! Ignore it.")
@@ -1714,9 +1711,7 @@ def is_context_parallel_supported():
             if self.sharding_parallel_degree == -1:
                 if len(self.sharding) > 0:
                     self.sharding_parallel_degree = world_size // (
-                        self.tensor_parallel_degree
-                        * self.sep_parallel_degree
-                        * self.pipeline_parallel_degree
+                        self.tensor_parallel_degree * self.sep_parallel_degree * self.pipeline_parallel_degree
                     )
 
             self.sharding_parallel_degree = max(self.sharding_parallel_degree, 1)

paddlenlp/trainer/utils/zero_cost_checkpoint.py

@@ -191,10 +191,13 @@ def ema_accumulate(self, global_step, loss, zcc_ema_loss_threshold):
                     _, cpu_buf = self.param_fusion_storage_helper.inited_buffers[index]
                     updated_ema = self.ema_coef * ema_buf + (1 - self.ema_coef) * cpu_buf
                     self.ema_buffer_model_params[index] = updated_ema
-                logger.info(f"[ZCC EMA] accmulating, buffer type:{self.ema_buffer.place} {self.ema_buffer.dtype}, done")
+                logger.info(
+                    f"[ZCC EMA] accmulating, buffer type:{self.ema_buffer.place} {self.ema_buffer.dtype}, done"
+                )
             else:
-                logger.info(f"[ZCC EMA] accmulating SKIP for global_step:{global_step}, because loss:{loss} > threshold:{zcc_ema_loss_threshold}")
-
+                logger.info(
+                    f"[ZCC EMA] accmulating SKIP for global_step:{global_step}, because loss:{loss} > threshold:{zcc_ema_loss_threshold}"
+                )
 
     @imperative_base.no_grad()
     def ema_state_dict(self):
@@ -788,9 +791,9 @@ def process_offload_task(self, dump, global_step):
 
             if self.ema_coef is not None:
                 self.zcc_ema_processor.ema_accumulate(
-                    self.trainer_state.global_step, 
+                    self.trainer_state.global_step,
                     self.trainer_state.loss,
-                    self.training_args_content.zcc_ema_loss_threshold
+                    self.training_args_content.zcc_ema_loss_threshold,
                 )
 
         # continue to process dumping task at the last chunk

slm/model_zoo/gpt-3/external_ops/token_dispatcher_utils/utils.h

@@ -124,5 +124,6 @@ __device__ __forceinline__ void vectorized_memcpy(const T* src,
     PD_SWITCH_NUM_EXPERTS_IMPL(__num_expert, 8, __VA_ARGS__);    \
     PD_SWITCH_NUM_EXPERTS_IMPL(__num_expert, 16, __VA_ARGS__);   \
     PD_SWITCH_NUM_EXPERTS_IMPL(__num_expert, 32, __VA_ARGS__);   \
+    PD_SWITCH_NUM_EXPERTS_IMPL(__num_expert, 64, __VA_ARGS__);   \
     PD_THROW("Unsupported expert number %d", int(__num_expert)); \
   } while (0)