fix overflow of efficient_attn_ratio (#1436)

Author: hhaAndroid

xtuner/v1/engine/train_engine.py

@@ -281,8 +281,8 @@ def train_step(self, data_batches: list[ModelItem]) -> tuple[LossLog, OtherLog]:
                 step_consumed_tokens += seq_ctx.mask.sum()
 
                 num_tokens = seq_ctx.cu_seq_lens_k[1:] - seq_ctx.cu_seq_lens_k[:-1]
-                efficient_forward_tokens += (num_tokens**2).sum()
-                total_forward_tokens += (num_tokens.sum()) ** 2
+                efficient_forward_tokens += (num_tokens.long() ** 2).sum()
+                total_forward_tokens += (num_tokens.long().sum()) ** 2
 
             if self.intra_layer_micro_batch == 1:
                 output = self.model(seq_ctx=seq_ctx_list[0], loss_ctx=loss_ctx_list[0])

xtuner/v1/engine/vision_compose_train_engine.py

@@ -167,8 +167,8 @@ def train_step(self, data_batches: List[ModelItem]) -> tuple[LossLog, OtherLog]:
                         step_consumed_img_tokens /= seq_ctx.sequence_parallel_mesh.size()
 
                 num_tokens = seq_ctx.cu_seq_lens_k[1:] - seq_ctx.cu_seq_lens_k[:-1]
-                efficient_forward_tokens += (num_tokens**2).sum()
-                total_forward_tokens += (num_tokens.sum()) ** 2
+                efficient_forward_tokens += (num_tokens.long() ** 2).sum()
+                total_forward_tokens += (num_tokens.long().sum()) ** 2
 
             # todo: support intra_layer_micro_batch
             output = self.model(seq_ctx=seq_ctx_list[0], loss_ctx=loss_ctx_list[0])