per-token loss calculation

Adjust loss logging based on per-token loss calculation.

Author: aigaolc

mcore_adapter/src/mcore_adapter/trainer/trainer.py

@@ -927,7 +927,12 @@ def _maybe_log_save_evaluate(
             # reset tr_loss to zero
             tr_loss -= tr_loss
 
-            logs["loss"] = round(tr_loss_scalar / (self.state.global_step - self._globalstep_last_logged), 4)
+            # logs["loss"] = round(tr_loss_scalar / (self.state.global_step - self._globalstep_last_logged), 4)
+            if self.args.calculate_per_token_loss:
+                logs["loss"] = round(tr_loss_scalar, 4)
+            else:
+                logs["loss"] = round(tr_loss_scalar / (self.state.global_step - self._globalstep_last_logged), 4)
+            
             if grad_norm is not None:
                 logs["grad_norm"] = grad_norm.detach().item() if isinstance(grad_norm, torch.Tensor) else grad_norm
             logs["learning_rate"] = self._get_learning_rate()