refactor: move class function inside Trainer

Author: Tridu33

mindnlp/engine/trainer/base.py

@@ -90,7 +90,6 @@
     TrainerControl,
     TrainerState,
 )
-from ..utils import _get_learning_rate
 
 
 logger = logging.get_logger(__name__)
@@ -126,7 +125,7 @@ def _is_peft_model(model):
 class Trainer:
     """
     Trainer is a simple but feature-complete training and eval loop for MindSpore, optimized for 🤗 Transformers.
-    """
+    """ 
     def __init__(
         self,
         model: Union[PreTrainedModel, nn.Module] = None,
@@ -288,6 +287,29 @@ def __init__(
         self._created_lr_scheduler = False
         self.actual_distributed_type = accelerate_distributed_type
 
+
+    def _get_learning_rate(self):
+        r"""
+        This function retrieves the learning rate used by the optimizer.
+        
+        Args:
+            self: An instance of the class containing the optimizer and learning rate scheduler.
+        
+        Returns:
+            The learning rate value (float) used by the optimizer.
+        
+        Raises:
+            None.
+        """
+        if isinstance(self.lr_scheduler, optim.lr_scheduler.ReduceLROnPlateau):
+            last_lr = self.optimizer.param_groups[0]["lr"]
+        else:
+            last_lr = self.lr_scheduler.get_last_lr()[0]
+        if ops.is_tensor(last_lr):
+            last_lr = last_lr.item()
+        return last_lr
+
+
     def _activate_neftune(self, model):
         r"""
         Activates the neftune as presented in this code: https://github.com/neelsjain/NEFTune and paper:
@@ -1136,6 +1158,7 @@ def _inner_training_loop(
                             model.parameters(),
                             args.max_grad_norm,
                         )
+                    
                     # Optimizer step
                     self.optimizer.step()
 
@@ -1376,7 +1399,7 @@ def training_step(self, model: nn.Module, inputs: Dict[str, Union[mindspore.Tens
         inputs = self._prepare_inputs(inputs)
 
         def forward(inputs):
-            if accelerate_distributed_type == DistributedType.MULTI_NPU_DP:
+            if accelerate_distributed_type == DistributedType.MULTI_NPU:
                 from mindspore.communication import get_group_size
                 import mindspore.ops as msops
                 rank_size = get_group_size()

mindnlp/engine/utils.py

@@ -505,28 +505,6 @@ def speed_metrics(split, start_time, num_samples=None, num_steps=None, num_token
         result[f"{split}_tokens_per_second"] = round(tokens_per_second, 3)
     return result
 
-def _get_learning_rate(self):
-    r"""
-    This function retrieves the learning rate used by the optimizer.
-    
-    Args:
-        self: An instance of the class containing the optimizer and learning rate scheduler.
-    
-    Returns:
-        The learning rate value (float) used by the optimizer.
-    
-    Raises:
-        None.
-    """
-    if isinstance(self.lr_scheduler, optim.lr_scheduler.ReduceLROnPlateau):
-        last_lr = self.optimizer.param_groups[0]["lr"]
-    else:
-        last_lr = self.lr_scheduler.get_last_lr()[0]
-    if ops.is_tensor(last_lr):
-        last_lr = last_lr.item()
-    return last_lr
-
-
 def find_batch_size(tensors):
     """
     Find the first dimension of a tensor in a nested list/tuple/dict of tensors.