Created
February 19, 2024 01:47
-
-
Save celsowm/d158950ef8815f4b42f6681476a64ae6 to your computer and use it in GitHub Desktop.
fine_tunning_llama2_bode_alpaca.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
from datasets import load_dataset, Dataset | |
from trl import SFTTrainer | |
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TrainingArguments | |
from peft import LoraConfig | |
import torch, sys | |
dataset = load_dataset("celsowm/auryn", split="train" | |
#, download_mode="force_redownload" | |
) | |
# Specify the model | |
model_name = 'recogna-nlp/bode-7b-alpaca-pt-br-no-peft' | |
# Load the tokenizer and the model | |
tokenizer = AutoTokenizer.from_pretrained(model_name) | |
# BitsAndBytesConfig int-4 config | |
bnb_config = BitsAndBytesConfig( | |
load_in_4bit=True, | |
bnb_4bit_use_double_quant=True, | |
bnb_4bit_quant_type="nf4", | |
bnb_4bit_compute_dtype=torch.bfloat16, | |
) | |
model = AutoModelForCausalLM.from_pretrained( | |
model_name, | |
quantization_config=bnb_config, | |
use_cache=False, | |
device_map="auto", | |
torch_dtype=torch.float16 | |
) | |
peft_config = LoraConfig( | |
lora_alpha=64, | |
lora_dropout=0.05, | |
r=128, | |
bias="none", | |
target_modules="all-linear", | |
task_type="CAUSAL_LM", | |
) | |
training_args = TrainingArguments( | |
output_dir='outputs/auryn', | |
num_train_epochs=3, # number of training epochs | |
per_device_train_batch_size=3, # batch size per device during training | |
gradient_accumulation_steps=2, # number of steps before performing a backward/update pass | |
gradient_checkpointing=True, # use gradient checkpointing to save memory | |
optim="adamw_torch_fused", # use fused adamw optimizer | |
logging_steps=10, # log every 10 steps | |
save_strategy="epoch", # save checkpoint every epoch | |
learning_rate=2e-4, # learning rate, based on QLoRA paper | |
bf16=True, # use fploat16 precision | |
tf32=True, # use tf32 precision | |
max_grad_norm=0.3, # max gradient norm based on QLoRA paper | |
warmup_ratio=0.03, # warmup ratio based on QLoRA paper | |
lr_scheduler_type="constant" | |
) | |
trainer = SFTTrainer( | |
model=model, | |
peft_config=peft_config, | |
args=training_args, | |
train_dataset=dataset, | |
packing=True | |
) | |
# Train the model | |
trainer.train() | |
# Save the trained model and tokenizer | |
trainer.model.save_pretrained(training_args.output_dir) | |
tokenizer.save_pretrained(training_args.output_dir) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment