[megatron] Fix router replay CP partition and expert index overflow

swift/megatron/trainers/utils.py

@@ -422,8 +422,11 @@ def compute_per_token_logps_fn(model, args, data_iterator, temperature=1.0, no_g
     global_topk_idx = data.pop('routed_experts', None)
     if enable_routing_replay and RouterReplayHelper.is_replay_forward_action(model.config):
         assert global_topk_idx is not None, 'When router_replay_mode = R3, routed_experts must be in data'
-        routing_topk_idx = get_local_topk_idx_for_current_rank(global_topk_idx, model.config,
-                                                               data.get('packed_seq_params'))
+        routing_topk_idx = get_local_topk_idx_for_current_rank(
+            global_topk_idx,
+            model.config,
+            data.get('packed_seq_params'),
+            cp_partition_mode=getattr(args, 'cp_partition_mode', 'zigzag'))
         set_router_replay_data(routing_topk_idx, model.config)
 
     data_for_forward = {k: v for k, v in data.items() if k != 'labels'}

swift/megatron/utils/router_replay_utils.py

@@ -79,7 +79,7 @@ def get_local_layer_range(tf_config, vp_rank=None, only_moe_layer=True):
     return offset, count
 
 
-def get_local_topk_idx_for_current_rank(global_topk_idx, tf_config, packed_seq_params=None):
+def get_local_topk_idx_for_current_rank(global_topk_idx, tf_config, packed_seq_params=None, cp_partition_mode='zigzag'):
     if global_topk_idx is None:
         return None
     # 1. pp slice
@@ -100,17 +100,29 @@ def get_local_topk_idx_for_current_rank(global_topk_idx, tf_config, packed_seq_p
     # 2. cp slice
     cp_size = mpu.get_context_parallel_world_size()
     if cp_size > 1:
-        local_topk_idx = split_cp_inputs(local_topk_idx, getattr(packed_seq_params, 'cu_seqlens_q', None), 1)
+        kwargs = {}
+        if cp_partition_mode == 'contiguous':
+            kwargs['cp_partition_mode'] = 'contiguous'
+        local_topk_idx = split_cp_inputs(local_topk_idx, getattr(packed_seq_params, 'cu_seqlens_q', None), 1, **kwargs)
     # 3. sp slice
     local_topk_idx = scatter_to_sequence_parallel_region(local_topk_idx.transpose(0, 1)).transpose(0, 1)
     return local_topk_idx
 
 
 def get_router_replay_data(tf_config, vp_rank=None):
     router_instances_list = RouterReplayHelper.get_micro_batch_router_list(tf_config, vp_rank)
-    layers_topk_idx = []
-    for router in router_instances_list:
-        layers_topk_idx.append(router.recorded_topk_idx.to(torch.uint8))
+    recorded_topk_idx = [router.recorded_topk_idx for router in router_instances_list]
+    nonempty_topk_idx = [topk_idx for topk_idx in recorded_topk_idx if topk_idx.numel()]
+    num_experts = tf_config.num_moe_experts
+    if nonempty_topk_idx:
+        bounds = torch.stack([torch.stack(torch.aminmax(topk_idx)) for topk_idx in nonempty_topk_idx]).cpu()
+        min_idx = bounds[:, 0].min().item()
+        max_idx = bounds[:, 1].max().item()
+        if min_idx < 0 or max_idx >= num_experts:
+            raise ValueError(
+                f'Recorded expert indices must be in [0, {num_experts}), but found range [{min_idx}, {max_idx}].')
+    index_dtype = torch.uint8 if num_experts <= 256 else torch.int32
+    layers_topk_idx = [topk_idx.to(index_dtype) for topk_idx in recorded_topk_idx]
     # layer_num, seq_len, topk -> 1, seq_len, layer_num, topk
     layers_topk_idx = torch.stack(layers_topk_idx).transpose(0, 1).unsqueeze(0).to(device_name)
     return layers_topk_idx