| |
| |
|
|
| |
|
|
| from typing import List, Optional |
|
|
| import torch |
| from torch.nn import Module |
| from torch.nn.modules.loss import _Loss |
| from torch.optim.lr_scheduler import _LRScheduler |
|
|
| from internlm.core.gradient_handler import BaseGradientHandler |
| from internlm.solver.beta2_scheduler import Beta2Scheduler |
| from internlm.solver.optimizer.hybrid_zero_optim import BaseOptimizer |
| from internlm.utils.common import get_batch_size, move_to_device |
|
|
|
|
| class Engine: |
| """ |
| The Engine class is responsible for managing the training and evaluation process of a neural network model. |
| It handles the forward and backward passes, parameter updates, gradient handling, and mode switching between |
| training and evaluation. |
| |
| Args: |
| model (torch.nn.Module): The neural network model to be trained or evaluated. |
| optimizer (BaseOptimizer): The optimizer used for updating the parameters of the model. |
| lr_scheduler (torch.optim.lr_scheduler._LRScheduler, optional): The learning rate scheduler for the optimizer. |
| Default is None. |
| beta2_scheduler (internlm.solver.beta2_scheduler.Beta2Scheduler, optional): The beta2 scheduler for the |
| optimizer. Default is None. |
| criterion (torch.nn.modules.loss._Loss, optional): The loss function used for calculating the loss during |
| training. Default is None. |
| gradient_handlers (List[BaseGradientHandler], optional): A list of gradient handlers used in the backward pass. |
| Default is None. |
| clip_grad_norm (float, optional): The norm value for gradient clipping. Default is 0.0. |
| |
| Examples: |
| >>> # define model, criterion, optimizer, lr_scheduler, train_dataloader for your training |
| >>> model = ... |
| >>> criterion = ... |
| >>> optimizer = ... |
| >>> train_dataloader = ... |
| >>> engine, _, _, _ = internlm.initialize_engine(model, optimizer, criterion) |
| >>> engine.train() |
| >>> for inputs, labels in train_dataloader |
| >>> # set gradients to zero |
| >>> engine.zero_grad() |
| >>> # run forward pass |
| >>> outputs = engine(inputs) |
| >>> # compute loss value and run backward pass |
| >>> loss = engine.criterion(outputs, labels) |
| >>> engine.backward(loss) |
| >>> # update parameters |
| >>> engine.step() |
| """ |
|
|
| def __init__( |
| self, |
| model: Module, |
| optimizer: BaseOptimizer, |
| lr_scheduler: Optional[_LRScheduler] = None, |
| beta2_scheduler: Optional[Beta2Scheduler] = None, |
| criterion: Optional[_Loss] = None, |
| gradient_handlers: Optional[List[BaseGradientHandler]] = None, |
| clip_grad_norm: float = 0.0, |
| ): |
| self._model = model |
| self._optimizer = optimizer |
| self._lr_scheduler = lr_scheduler |
| self._beta2_scheduler = beta2_scheduler |
| self._criterion = criterion |
| self._clip_grad_norm = clip_grad_norm |
|
|
| |
| self.training = True |
|
|
| |
| self._gradient_handlers = gradient_handlers if gradient_handlers else [] |
|
|
| @property |
| def model(self): |
| """Returns the model attached to the engine.""" |
| return self._model |
|
|
| @property |
| def optimizer(self): |
| """Returns the optimizer attached to the engine.""" |
| return self._optimizer |
|
|
| @property |
| def criterion(self): |
| """Returns the criterion (loss function) attached to the engine.""" |
| return self._criterion |
|
|
| def _all_reduce_gradients(self): |
| """Handles all-reduce operations of gradients across different parallel groups.""" |
| for handler in self._gradient_handlers: |
| handler.handle_gradient() |
|
|
| def zero_grad(self): |
| """Sets the gradient of all parameters in the model to zero.""" |
| self.optimizer.zero_grad() |
|
|
| def step(self): |
| """ |
| Executes the parameter update step. This includes all-reduce operations of gradients, gradient clipping, |
| and parameter update. If successful, it also steps the learning rate scheduler and beta2 scheduler |
| if they exist. |
| |
| Returns: |
| success (bool): Whether the parameter update was successful. |
| grad_norm (float): The norm of the gradient after clipping. |
| """ |
| self._all_reduce_gradients() |
| self.optimizer.clip_grad_norm(self.model, self._clip_grad_norm) |
|
|
| success, grad_norm = self.optimizer.step() |
|
|
| if success and self._lr_scheduler is not None: |
| self._lr_scheduler.step() |
|
|
| if success and self._beta2_scheduler is not None: |
| self._beta2_scheduler.step() |
|
|
| return success, grad_norm |
|
|
| def train(self): |
| """Sets the model to training mode.""" |
| self.training = True |
| self._model.train() |
|
|
| def eval(self): |
| """Sets the model to evaluation mode.""" |
| self.training = False |
| self._model.eval() |
|
|
| def backward(self, loss: torch.Tensor): |
| """ |
| Starts the backward propagation given the loss value computed by a loss function. |
| |
| Args: |
| loss (torch.Tensor): The loss value computed by a loss function. |
| """ |
| return self.optimizer.backward(loss) |
|
|
| def backward_by_grad(self, tensor, grad): |
| """ |
| Starts the backward propagation given the gradient of the output tensor. |
| |
| Args: |
| tensor (torch.Tensor): The output tensor. |
| grad (torch.Tensor): The gradient passed back to the output tensor. |
| """ |
| return self.optimizer.backward_by_grad(tensor, grad) |
|
|
| def __call__(self, *args, **kwargs): |
| """ |
| Runs the forward step for the model. |
| |
| Returns: |
| torch.Tensor: The output of the model. |
| """ |
| return self.model(*args, **kwargs) |
|
|
| def load_batch(self, data_iter, to_gpu=True): |
| """ |
| Loads a batch from the data iterator. It returns the data and labels which are |
| already in the same GPU as where the model is. |
| |
| Args: |
| data_iter (Iterable): The data iterator from which to get a batch of data, obtained by calling |
| iter(dataloader). |
| to_gpu (bool, optional): Whether the data should be moved to the GPU. Default is True. |
| |
| Returns: |
| Tuple (torch.Tensor, torch.Tensor): A tuple of (data, label). |
| """ |
| if data_iter is None: |
| raise RuntimeError("Dataloader is not defined.") |
| try: |
| batch_data = next(data_iter) |
| except TypeError: |
| batch_data = data_iter |
|
|
| if to_gpu: |
| batch_data = move_to_device(batch_data) |
| batch_size = get_batch_size(batch_data) |
|
|
| return batch_data, batch_size |
|
|
|
|
| class KDEngine(Engine): |
| def __init__( |
| self, |
| model: Module, |
| teacher: Module, |
| optimizer: BaseOptimizer, |
| lr_scheduler: Optional[_LRScheduler] = None, |
| beta2_scheduler: Optional[Beta2Scheduler] = None, |
| criterion: Optional[_Loss] = None, |
| kd_criterion: Optional[_Loss] = None, |
| gradient_handlers: Optional[List[BaseGradientHandler]] = None, |
| clip_grad_norm: float = 0.0, |
| ): |
| self._teacher = teacher |
| self._kd_criterion = kd_criterion |
|
|
| super().__init__( |
| model=model, |
| optimizer=optimizer, |
| lr_scheduler=lr_scheduler, |
| beta2_scheduler=beta2_scheduler, |
| criterion=criterion, |
| gradient_handlers=gradient_handlers, |
| clip_grad_norm=clip_grad_norm, |
| ) |
|
|
| @property |
| def teacher(self): |
| """Returns the model attached to the engine.""" |
| return self._teacher |
|
|
| @property |
| def kd_criterion(self): |
| """Returns the model attached to the engine.""" |
| return self._kd_criterion |
|
|
