import torch from datasets import load_dataset 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 1. Install accelerate: conda install -c conda-forge accelerate 2. Setup accelerate config: accelerate config to simply use all the GPUs available: python -c "from accelerate.utils import write_basic_config; write_basic_config(mixed_precision='bf16')" check accelerate config: accelerate env 3. Run the code: accelerate launch phi3-mini-sample-ft.py """ ################### # Hyper-parameters ################### args = { "bf16": True, "do_eval": False, "evaluation_strategy": "no", "eval_steps": 100, "learning_rate": 5.0e-06, "log_level": "info", "logging_steps": 20, "logging_strategy": "steps", "lr_scheduler_type": "cosine", "num_train_epochs": 1, "max_steps": -1, "output_dir": ".", "overwrite_output_dir": True, "per_device_eval_batch_size": 4, "per_device_train_batch_size": 8, "remove_unused_columns": True, "save_steps": 100, "save_total_limit": 1, "seed": 0, "gradient_checkpointing": True, "gradient_accumulation_steps": 1, "warmup_ratio": 0.1, } 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( trust_remote_code=True, attn_implementation="flash_attention_2", # load 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) ################ # Data Loading ################ dataset = load_dataset("imdb") train_dataset = dataset["train"] eval_dataset = dataset["test"] ################ # Training ################ trainer = SFTTrainer( model=model, args=training_args, train_dataset=train_dataset, max_seq_length=2048, dataset_text_field="text", tokenizer=tokenizer, ) train_result = trainer.train() metrics = train_result.metrics trainer.log_metrics("train", metrics) trainer.save_metrics("train", metrics) trainer.save_state()