fix tests

Former-commit-id: 23f97bd437424ef43b2b84743d56acc5d1ca70d5

src/llmtuner/data/loader.py

@@ -135,9 +135,9 @@ def merge_dataset(
 
 
 def get_dataset(
+    tokenizer: "PreTrainedTokenizer",
     model_args: "ModelArguments",
     data_args: "DataArguments",
-    tokenizer: "PreTrainedTokenizer",
     training_args: "Seq2SeqTrainingArguments",
     stage: Literal["pt", "sft", "rm", "ppo"],
     # split: Optional[str] = "train", # TODO: add split

src/llmtuner/extras/patches/llama_patch.py

@@ -130,6 +130,20 @@ def llama_flash_attn_forward(
 
     dropout_rate = self.attention_dropout if self.training else 0.0
 
+    input_dtype = query_states.dtype
+    if input_dtype == torch.float32:
+        if torch.is_autocast_enabled():
+            target_dtype = torch.get_autocast_gpu_dtype()
+        elif hasattr(self.config, "_pre_quantization_dtype"):
+            target_dtype = self.config._pre_quantization_dtype
+        else:
+            target_dtype = self.q_proj.weight.dtype
+
+        logger.warning_once("The input hidden states seems to be silently casted in float32.")
+        query_states = query_states.to(target_dtype)
+        key_states = key_states.to(target_dtype)
+        value_states = value_states.to(target_dtype)
+
     if getattr(self.config, "group_size_ratio", None) and self.training: # shift
         groupsz = int(q_len * getattr(self.config, "group_size_ratio"))
         assert q_len % groupsz == 0, "q_len {} should be divisible by group size {}.".format(q_len, groupsz)

src/llmtuner/train/dpo/workflow.py

@@ -25,7 +25,7 @@ def run_dpo(
     callbacks: Optional[List["TrainerCallback"]] = None
 ):
     model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, training_args.do_train)
-    dataset = get_dataset(model_args, data_args, tokenizer, training_args, stage="rm")
+    dataset = get_dataset(tokenizer, model_args, data_args, training_args, stage="rm")
     data_collator = DPODataCollatorWithPadding(
         tokenizer=tokenizer,
         pad_to_multiple_of=8,

src/llmtuner/train/ppo/workflow.py

@@ -29,7 +29,7 @@ def run_ppo(
     callbacks: Optional[List["TrainerCallback"]] = None
 ):
     model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, training_args.do_train, add_valuehead=True)
-    dataset = get_dataset(model_args, data_args, tokenizer, training_args, stage="ppo")
+    dataset = get_dataset(tokenizer, model_args, data_args, training_args, stage="ppo")
 
     tokenizer.padding_side = "left" # use left-padding in generation while using right-padding in training
     data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

src/llmtuner/train/pt/workflow.py

@@ -22,7 +22,7 @@ def run_pt(
     callbacks: Optional[List["TrainerCallback"]] = None
 ):
     model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, training_args.do_train)
-    dataset = get_dataset(model_args, data_args, tokenizer, training_args, stage="pt")
+    dataset = get_dataset(tokenizer, model_args, data_args, training_args, stage="pt")
     data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
 
     # Initialize our Trainer

src/llmtuner/train/rm/workflow.py

@@ -26,7 +26,7 @@ def run_rm(
     callbacks: Optional[List["TrainerCallback"]] = None
 ):
     model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, training_args.do_train, add_valuehead=True)
-    dataset = get_dataset(model_args, data_args, tokenizer, training_args, stage="rm")
+    dataset = get_dataset(tokenizer, model_args, data_args, training_args, stage="rm")
     data_collator = PairwiseDataCollatorWithPadding(tokenizer, pad_to_multiple_of=8)
 
     # Update arguments

src/llmtuner/train/sft/workflow.py

@@ -27,7 +27,7 @@ def run_sft(
     callbacks: Optional[List["TrainerCallback"]] = None
 ):
     model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, training_args.do_train)
-    dataset = get_dataset(model_args, data_args, tokenizer, training_args, stage="sft")
+    dataset = get_dataset(tokenizer, model_args, data_args, training_args, stage="sft")
 
     if training_args.predict_with_generate:
         tokenizer.padding_side = "left" # use left-padding in generation