Update sample_finetune.py

sample_finetune.py

@@ -4,7 +4,9 @@ from trl import SFTTrainer
 from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
 
 """
-Please note that A100 or later generation GPUs are required to finetune Phi-3 models
+A simple example on using SFTTrainer and Accelerate to finetune Phi-3 models. For
+a more advanced example, please follow HF alignment-handbook/scripts/run_sft.py
+
 1. Install accelerate: 
     conda install -c conda-forge accelerate
 2. Setup accelerate config:
@@ -14,19 +16,16 @@ to simply use all the GPUs available:
 check accelerate config:
     accelerate env
 3. Run the code:
-    accelerate launch phi3-mini-sample-ft.py
+    accelerate launch sample_finetune.py
 """
 
 ###################
 # Hyper-parameters
 ###################
-
-
 args = {
     "bf16": True,
     "do_eval": False,
-    "evaluation_strategy": "no",
-    "eval_steps": 100,
+    "eval_strategy": "no",
     "learning_rate": 5.0e-06,
     "log_level": "info",
     "logging_steps": 20,
@@ -34,7 +33,7 @@ args = {
     "lr_scheduler_type": "cosine",
     "num_train_epochs": 1,
     "max_steps": -1,
-    "output_dir": ".",
+    "output_dir": "./checkpoint_dir",
     "overwrite_output_dir": True,
     "per_device_eval_batch_size": 4,
     "per_device_train_batch_size": 8,
@@ -43,51 +42,88 @@ args = {
     "save_total_limit": 1,
     "seed": 0,
     "gradient_checkpointing": True,
+    "gradient_checkpointing_kwargs":{"use_reentrant": False},
     "gradient_accumulation_steps": 1,
-    "warmup_ratio": 0.1,
+    "warmup_ratio": 0.2,
     }
     
 training_args = TrainingArguments(**args)
 
-
 ################
 # Modle Loading
 ################
-
 checkpoint_path = "microsoft/Phi-3-mini-4k-instruct"
 # checkpoint_path = "microsoft/Phi-3-mini-128k-instruct"
 model_kwargs = dict(
+    use_cache=False,
     trust_remote_code=True,
-    attn_implementation="flash_attention_2",  # load the model with flash-attenstion support
+    attn_implementation="flash_attention_2",  # loading the model with flash-attenstion support
     torch_dtype=torch.bfloat16,
     device_map="cuda",
 )
 model = AutoModelForCausalLM.from_pretrained(checkpoint_path, **model_kwargs)
-tokenizer = AutoTokenizer.from_pretrained(checkpoint_path, trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained(checkpoint_path)
+tokenizer.pad_token = tokenizer.unk_token  # use unk rather than eos token to prevent endless generation
+tokenizer.pad_token_id = tokenizer.convert_tokens_to_ids(tokenizer.pad_token)
+tokenizer.padding_side = 'right'
 
-################
-# Data Loading
-################
+##################
+# Data Processing
+##################
+def apply_chat_template(
+    example,
+    tokenizer,
+):
+    messages = example["messages"]
+    # Add an empty system message if there is none
+    if messages[0]["role"] != "system":
+        messages.insert(0, {"role": "system", "content": ""})
+    example["text"] = tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=False)
+    return example
 
-dataset = load_dataset("imdb")
-train_dataset = dataset["train"]
-eval_dataset = dataset["test"]
+raw_dataset = load_dataset("HuggingFaceH4/ultrachat_200k")
+column_names = list(raw_dataset["train_sft"].features)
 
+processed_dataset = raw_dataset.map(
+    apply_chat_template,
+    fn_kwargs={"tokenizer": tokenizer},
+    num_proc=12,
+    remove_columns=column_names,
+    desc="Applying chat template",
+)
+train_dataset = processed_dataset["train_sft"]
+eval_dataset = processed_dataset["test_sft"]
 
-################
+###########
 # Training
-################
-
+###########
 trainer = SFTTrainer(
     model=model,
     args=training_args,
     train_dataset=train_dataset,
+    eval_dataset=eval_dataset,
     max_seq_length=2048,
     dataset_text_field="text",
     tokenizer=tokenizer,
+    packing=True
 )
 train_result = trainer.train()
 metrics = train_result.metrics
 trainer.log_metrics("train", metrics)
 trainer.save_metrics("train", metrics)
 trainer.save_state()
+
+#############
+# Evaluation
+#############
+tokenizer.padding_side = 'left'
+metrics = trainer.evaluate()
+metrics["eval_samples"] = len(eval_dataset)
+trainer.log_metrics("eval", metrics)
+trainer.save_metrics("eval", metrics)
+
+############
+# Save model
+############
+trainer.save_model(training_args.output_dir)