Spaces:
Running
Running
| # run as a module using: python3 -m scripts.finetune | |
| # Using: https://huggingface.co/blog/mlabonne/sft-llama3 | |
| import torch | |
| from trl import SFTTrainer | |
| from datasets import load_dataset | |
| from transformers import TrainingArguments, TextStreamer | |
| from unsloth.chat_templates import get_chat_template | |
| from unsloth import FastLanguageModel, is_bfloat16_supported | |
| from data.fine_tune_dataset import load_data | |
| def finetune(model="unsloth/Meta-Llama-3.1-8B-bnb-4bit", dataset="mlabonne/FineTome-100k"): | |
| hf_token = "" | |
| # Loading the model and restricting context window | |
| max_seq_length = 2048 | |
| model, tokenizer = FastLanguageModel.from_pretrained( | |
| model_name=model, | |
| max_seq_length=max_seq_length, | |
| load_in_4bit=True, | |
| dtype=None, | |
| ) | |
| # Loading prepared dataset | |
| dataset = load_data(dataset, tokenizer) | |
| # Loading the model for fine tuning - only set to FT 42million/8billion parameters | |
| model = FastLanguageModel.get_peft_model( | |
| model, | |
| r=16, # rank determines LoRA (Low rank adaptation - freezing much of the model for fine tuning) matrix size, higher increases memory and compute cost | |
| lora_alpha=16, # scaling factor for updates | |
| lora_dropout=0, # not used for speedup | |
| target_modules=["q_proj", "k_proj", "v_proj", "up_proj", "down_proj", "o_proj", "gate_proj"], # where LoRA targets | |
| use_rslora=True, # rank stabilised | |
| use_gradient_checkpointing="unsloth" | |
| ) | |
| # Saving the untrained model, save_method can be lora to only save adapters or merged (16 or 4 bit) | |
| model.save_pretrained_merged("models/PreFineLlama-3.1-8B", tokenizer, save_method="merged_16bit") # save to models directory locally | |
| model.push_to_hub_merged("thebigoed/PreFineLlama-3.1-8B", tokenizer, token=hf_token, save_method="merged_16bit") | |
| trainer=SFTTrainer( | |
| model=model, | |
| tokenizer=tokenizer, | |
| train_dataset=dataset, | |
| dataset_text_field="text", | |
| max_seq_length=max_seq_length, | |
| dataset_num_proc=2, | |
| packing=True, | |
| args=TrainingArguments( | |
| learning_rate=3e-4, # to low = slow and local minima, too high = unstable | |
| lr_scheduler_type="linear", # adjusts the learning rate (linear and cosine are most popular) | |
| per_device_train_batch_size=8, | |
| gradient_accumulation_steps=2, | |
| num_train_epochs=1, | |
| fp16=not is_bfloat16_supported(), | |
| bf16=is_bfloat16_supported(), | |
| logging_steps=1, | |
| optim="adamw_8bit", | |
| weight_decay=0.01, | |
| warmup_steps=10, | |
| output_dir="output", | |
| seed=0, | |
| ), | |
| ) | |
| trainer.train() | |
| # Saving the model, save_method can be lora to only save adapters or merged (16 or 4 bit) | |
| model.save_pretrained_merged("models/FineLlama-3.1-8B", tokenizer, save_method="merged_16bit") # save to models directory locally | |
| model.push_to_hub_merged("thebigoed/FineLlama-3.1-8B", tokenizer, token=hf_token, save_method="merged_16bit") | |
| # Use to save in GGUF quantised format | |
| # quant_methods = ["q2_k", "q3_k_m", "q4_k_m", "q5_k_m", "q6_k", "q8_0"] | |
| # for quant in quant_methods: | |
| # model.push_to_hub_gguf("", tokenizer, quant) | |
| return | |
| if __name__ == "__main__": | |
| finetune() | |