Update train_seallm_khm_sum.py

train_seallm_khm_sum.py

@@ -5,8 +5,9 @@ from transformers import (
     AutoTokenizer,
     AutoModelForCausalLM,
     BitsAndBytesConfig,
+    TrainingArguments,
 )
-from trl import SFTTrainer, SFTConfig
+from trl import SFTTrainer
 from peft import LoraConfig
 
 MODEL_NAME = "SeaLLMs/SeaLLMs-v3-1.5B"
@@ -101,8 +102,8 @@ def main():
         task_type="CAUSAL_LM",
     )
 
-    # NOTE: no max_seq_length here
-    sft_config = SFTConfig(
+    # Use standard TrainingArguments instead of SFTConfig
+    training_args = TrainingArguments(
         output_dir="seallm-khm-sum-lora",
         num_train_epochs=2,
         per_device_train_batch_size=2,
@@ -110,16 +111,14 @@ def main():
         gradient_accumulation_steps=8,
         learning_rate=2e-4,
         logging_steps=10,
-        evaluation_strategy="steps",  # <- was eval_strategy
+        evaluation_strategy="steps",  # eval every eval_steps
         eval_steps=200,
         save_steps=200,
         save_total_limit=2,
-        packing=True,
         lr_scheduler_type="cosine",
         warmup_ratio=0.03,
-        bf16=True,
-        gradient_checkpointing=True,
-        report_to="none",  # or "wandb"
+        bf16=True,              # ok on modern GPUs; set False if it crashes
+        report_to="none",       # or "wandb"
     )
 
     trainer = SFTTrainer(
@@ -128,9 +127,10 @@ def main():
         train_dataset=train_ds,
         eval_dataset=eval_ds,
         peft_config=lora_config,
-        args=sft_config,
+        args=training_args,
         dataset_text_field="text",
-        max_seq_length=1024,  # <- moved here
+        max_seq_length=1024,    # set here instead of in config
+        # packing=False,        # keep off for compatibility
     )
 
     trainer.train()