Compute Loss inside the training step.

optimum/onnxruntime/trainer.py

@@ -23,6 +23,12 @@
 import warnings
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Type, Union
+from tqdm.auto import tqdm
+
+
+# Integrations must be imported before ML frameworks:
+from transformers.integrations import hp_params, is_fairscale_available # isort: split
+
 
 
 # Integrations must be imported before ML frameworks:
@@ -45,6 +51,7 @@
 from transformers.debug_utils import DebugOption, DebugUnderflowOverflow
 from transformers.deepspeed import deepspeed_init, is_deepspeed_zero3_enabled
 from transformers.dependency_versions_check import dep_version_check
+from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
 from transformers.file_utils import (
  is_apex_available,
  is_sagemaker_dp_enabled,
@@ -134,6 +141,28 @@
 SCALER_NAME = "scaler.pt"
 
 
+class ModuleWithLoss(nn.Module):
+ def __init__(self, model, args) -> None:
+ super().__init__()
+ self._original_model = model
+ self.args = args
+ self.hf_trainer = Trainer(model)
+ # Label smoothing
+ if self.args.label_smoothing_factor != 0:
+ from transformers.trainer_pt_utils import LabelSmoother
+
+ self.label_smoother = LabelSmoother(epsilon=self.args.label_smoothing_factor)
+ else:
+ self.label_smoother = None
+
+ def forward(self, inputs: Dict[str, Union[torch.Tensor, Any]], return_outputs):
+ return self.hf_trainer.compute_loss(self._original_model, inputs, return_outputs=False)
+
+ @property
+ def config(self):
+ return self._original_model.config
+
+
 class ORTFeaturesManager:
  _TASKS_TO_ORTMODELS = {
  "default": ORTModelForFeatureExtraction,
@@ -279,12 +308,48 @@ def __init__(
  preprocess_logits_for_metrics=preprocess_logits_for_metrics,
  )
 
+ # We leverage both training_model and inference_model in conjunction with model.
+ # _training_model will be wrapped so it will use ORT and will use the overriden functions in ModuleWithLoss.
+ # _inferencing_model will be storing the default version of the model and we will switch to it in case of eval/test.
+
+ # Only Wrap the model if we pass --loss_in_train flag.
+ if args.loss_in_train:
+ self._training_model = ModuleWithLoss(model, args)
+ else:
+ self._training_model = model
+
+ self.model = model
+ self._inferencing_model = model
  self.feature = feature
  self.onnx_model_path = onnx_model_path
  self.exported_with_loss = False
  if self.args.local_rank:
  torch.cuda.set_device(self.args.local_rank)
 
+ # we assume that training_model and inference_model have the same forward signature column.
+ # self._signature_columns attribute only stores the first-time parsed signature
+ def _set_signature_columns_if_needed(self):
+ if self._signature_columns is None:
+ # Inspect model forward signature to keep only the arguments it accepts.
+ import inspect
+
+ if isinstance(self.model, ModuleWithLoss):
+ signature = inspect.signature(self.model._original_model.forward)
+ else:
+ signature = inspect.signature(self.model.forward)
+
+ self._signature_columns = list(signature.parameters.keys())
+ # Labels may be named label or label_ids, the default data collator handles that.
+ self._signature_columns += list(set(["label", "label_ids"] + self.label_names))
+
+ def compute_loss(self, model_with_loss, inputs, return_outputs=False):
+ # Run model forward + loss compute.
+ if self.args.loss_in_train and self.model == self._training_model:
+ outputs = model_with_loss(inputs, return_outputs)
+ return outputs
+ else:
+ return super().compute_loss(self.model, inputs, return_outputs)
+
  def train(
  self,
  resume_from_checkpoint: Optional[Union[str, bool]] = None,
@@ -314,6 +379,8 @@ def train(
  "https://huggingface.co/docs/optimum/onnxruntime/usage_guides/trainer#install-onnx-runtime."
  )
 
+ self.model = self._training_model
+
  if resume_from_checkpoint is False:
  resume_from_checkpoint = None
 
@@ -592,7 +659,6 @@ def _inner_training_loop(
  # Otherwise we need to call the whooooole sampler cause there is some random operation added
  # AT THE VERY END!
  _ = list(train_dataloader.sampler)
-
  for epoch in range(epochs_trained, num_train_epochs):
  if isinstance(train_dataloader, DataLoader) and isinstance(train_dataloader.sampler, DistributedSampler):
  train_dataloader.sampler.set_epoch(epoch)
@@ -801,6 +867,8 @@ def evaluate(
  dictionary also contains the epoch number which comes from the training state.
  """
  # memory metrics - must set up as early as possible
+ # TODO: We need to enable evaluation using ORT backend.
+ self.model = self._inferencing_model
  self._memory_tracker.start()
 
  eval_dataloader = self.get_eval_dataloader(eval_dataset)
@@ -892,6 +960,9 @@ def predict(
  - metrics (`Dict[str, float]`, *optional*): The potential dictionary of metrics (if the dataset contained
  labels).
  """
+ # TODO: We need to enable evaluation using ORT backend.
+ self.model = self._inferencing_model
+
  # memory metrics - must set up as early as possible
  self._memory_tracker.start()
 
@@ -909,10 +980,7 @@ def predict(
 
  try:
  output = eval_loop(
- test_dataloader,
- description="Prediction",
- ignore_keys=ignore_keys,
- metric_key_prefix=metric_key_prefix,
+ test_dataloader, description="Prediction", ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix
  )
  except Exception as error:
  logger.error(error)
@@ -1697,11 +1765,7 @@ def create_optimizer(self):
  optimizer_cls, optimizer_kwargs = Trainer.get_optimizer_cls_and_kwargs(self.args)
 
  if self.sharded_ddp == ShardedDDPOption.SIMPLE:
- self.optimizer = OSS(
- params=optimizer_grouped_parameters,
- optim=optimizer_cls,
- **optimizer_kwargs,
- )
+ self.optimizer = OSS(params=optimizer_grouped_parameters, optim=optimizer_cls, **optimizer_kwargs)
  else:
  self.optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)
  if optimizer_cls.__name__ == "Adam8bit":
@@ -1731,10 +1795,7 @@ def get_ort_optimizer_cls_and_kwargs(args: ORTTrainingArguments) -> Tuple[Any, A
  The training arguments for the training session.
  """
  optimizer_kwargs = {"lr": args.learning_rate}
- adam_kwargs = {
- "betas": (args.adam_beta1, args.adam_beta2),
- "eps": args.adam_epsilon,
- }
+ adam_kwargs = {"betas": (args.adam_beta1, args.adam_beta2), "eps": args.adam_epsilon}
  if args.optim == ORTOptimizerNames.ADAMW_ORT_FUSED:
  try:
  from onnxruntime.training.optim import FusedAdam

optimum/onnxruntime/training_args.py

@@ -65,6 +65,11 @@ class ORTTrainingArguments(TrainingArguments):
  metadata={"help": "The optimizer to use."},
  )
 
+ loss_in_train: Optional[bool] = field(
+ default=False,
+ metadata={"help": "Use ModuleWithLoss Wrapper to compute loss inside the training loop."},
+ )
+
  # This method will not need to be overriden after the deprecation of `--adafactor` in version 5 of 🤗 Transformers.
  def __post_init__(self):
  # Handle --use_env option in torch.distributed.launch (local_rank not passed as an arg then).
@@ -336,3 +341,12 @@ def __post_init__(self):
  f"{self.hub_model_id}).",
  FutureWarning,
  )
+ if self.loss_in_train is True:
+ logger.info(
+ "Using ModuleWithLoss Wrapper."
+ "loss will be computed during training loop and it will save memory peak "
+ )
+ else:
+ logger.info(
+ "Not Using ModuleWithLoss Wrapper."
+ )