From 14d9499d3fbbc77cf7f14098b4fad6fa346b20d0 Mon Sep 17 00:00:00 2001 From: thisisatharva-rh Date: Thu, 29 May 2025 09:39:38 -0400 Subject: [PATCH 1/3] use LigerCrossEntropy with reduction set to sum --- src/instructlab/training/main_ds.py | 6 ++++++ src/instructlab/training/utils.py | 3 +-- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/src/instructlab/training/main_ds.py b/src/instructlab/training/main_ds.py index 19ee76f3..a28b1b53 100644 --- a/src/instructlab/training/main_ds.py +++ b/src/instructlab/training/main_ds.py @@ -17,6 +17,7 @@ # Third Party from accelerate import Accelerator + try: # Third Party from deepspeed.ops.adam import DeepSpeedCPUAdam @@ -145,6 +146,7 @@ def setup_model( base_model_args["attn_implementation"] = "flash_attention_2" if args.use_dolomite: + from torch.nn import CrossEntropyLoss with ensure_loadable_dolomite_checkpoint( args.model_name_or_path, args.output_dir ) as path: @@ -153,6 +155,7 @@ def setup_model( model = GPTDolomiteForCausalLM.from_pretrained( **base_model_args, ) + model.loss_fct = CrossEntropyLoss(reduction="sum") elif args.use_liger: # TODO(osilkin): we duplicate some checks here because someone may run this script through # torchrun directly and not `run_training`. To fix this, we should eventually move everything @@ -164,6 +167,7 @@ def setup_model( try: # Third Party from liger_kernel.transformers import AutoLigerKernelForCausalLM + from liger_kernel.transformers.cross_entropy import LigerCrossEntropyLoss except ImportError as e: raise ValueError( "Liger kernels are not installed. Please install Liger kernels using the following command: pip install liger-kernel" @@ -175,8 +179,10 @@ def setup_model( model = AutoLigerKernelForCausalLM.from_pretrained( **base_model_args, cross_entropy=True, fused_linear_cross_entropy=False ) + model.loss_fct = LigerCrossEntropyLoss(reduction="sum") else: model = AutoModelForCausalLM.from_pretrained(**base_model_args) + model.loss_fct = CrossEntropyLoss(reduction="sum") # store the base model args so we can recall them later if saving a LoRA model args.base_model_args = base_model_args diff --git a/src/instructlab/training/utils.py b/src/instructlab/training/utils.py index 0c3644a4..a81346fa 100644 --- a/src/instructlab/training/utils.py +++ b/src/instructlab/training/utils.py @@ -425,8 +425,7 @@ def reduce_sum_forward( shift_labels = shift_labels.view(-1) # Ensure tensors are on the same device shift_labels = shift_labels.to(shift_logits.device) - loss_fct = torch.nn.CrossEntropyLoss(reduction="sum") - loss = loss_fct(shift_logits, shift_labels) + loss = model.loss_fct(shift_logits, shift_labels) if not return_dict: return ((loss,) + output) if loss is not None else output From 9dca5724d0b21cbf96de151cb2a2d3de91caaaed Mon Sep 17 00:00:00 2001 From: thisisatharva-rh Date: Thu, 29 May 2025 09:47:25 -0400 Subject: [PATCH 2/3] fix linting issue --- src/instructlab/training/main_ds.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/instructlab/training/main_ds.py b/src/instructlab/training/main_ds.py index a28b1b53..7e1c7eb0 100644 --- a/src/instructlab/training/main_ds.py +++ b/src/instructlab/training/main_ds.py @@ -17,7 +17,6 @@ # Third Party from accelerate import Accelerator - try: # Third Party from deepspeed.ops.adam import DeepSpeedCPUAdam From 0442be5e8e5c07b31905350e995dc095b2038738 Mon Sep 17 00:00:00 2001 From: thisisatharva-rh Date: Thu, 29 May 2025 09:50:09 -0400 Subject: [PATCH 3/3] fix import issue --- src/instructlab/training/main_ds.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/instructlab/training/main_ds.py b/src/instructlab/training/main_ds.py index 7e1c7eb0..bf3ac587 100644 --- a/src/instructlab/training/main_ds.py +++ b/src/instructlab/training/main_ds.py @@ -55,6 +55,7 @@ ) import torch import torch.distributed +from torch.nn import CrossEntropyLoss # First Party from instructlab.training import config @@ -145,7 +146,6 @@ def setup_model( base_model_args["attn_implementation"] = "flash_attention_2" if args.use_dolomite: - from torch.nn import CrossEntropyLoss with ensure_loadable_dolomite_checkpoint( args.model_name_or_path, args.output_dir ) as path: