guillaumetell-7b / finetuning.py

Update finetuning.py

0b60c78 verified 4 months ago

No virus

6.32 kB

	import os
	import torch

	device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

	print(device)


	from datasets import load_dataset
	from transformers import (
	AutoModelForCausalLM,
	AutoTokenizer,
	BitsAndBytesConfig,
	HfArgumentParser,
	TrainingArguments,
	pipeline,
	logging,
	LlamaTokenizerFast
	)
	from peft import LoraConfig, PeftModel, get_peft_model
	from trl import SFTTrainer

	# Le modèle que nous allons utiliser dans le Hugging Face hub
	model_name = "mistral-hermes-2.5"

	torch.cuda.empty_cache()

	#project_directory = "~/finetuning/sigmund-spplus"

	# Le nom du nouveau modèle
	new_model_name = "mistral-mfs-reference-2"

	# The output directory where the model predictions and checkpoints will be written
	output_dir = "./mistral-mfs-reference-2"

	# Tensorboard logs
	tb_log_dir = "./mistral-mfs-reference-2/logs"

	# Nombre de steps : à ajuster selon la taille du corpus et le nombre d'epochs à faire tourner.
	max_steps = 2000

	# Les paramètres importants !!
	per_device_train_batch_size = 4 #Nombre d'exemples envoyés par batch. En mettre plus pour aller plus vite.
	learning_rate = 2e-5 #De préférence un taux d'apprentissage bas comme Mistral-Hermes se débrouille bien en français
	max_seq_length = 4096 #C'est la fenêtre contextuelle. Elle peut être portée jusqu'à 4096 tokens (mais attention à la mémoire disponible !)
	save_steps = 1000 # Sauvegarde des steps (permet de faire redémarrer l'entraînement si le fine-tuning ne fonctionne pas)
	# Learning rate schedule (constant a bit better than cosine, and has advantage for analysis)
	lr_scheduler_type = "linear"


	#Les autres paramètres
	local_rank = -1
	per_device_eval_batch_size = 1
	gradient_accumulation_steps = 4
	max_grad_norm = 0.3
	weight_decay = 0.001
	lora_alpha = 16
	lora_dropout = 0.1
	lora_r = 64
	# Group sequences into batches with same length (saves memory and speeds up training considerably)
	group_by_length = True

	# Activate 4-bit precision base model loading
	use_4bit = True

	# Activate nested quantization for 4-bit base models
	use_nested_quant = False

	# Compute dtype for 4-bit base models
	bnb_4bit_compute_dtype = "float16"

	# Quantization type (fp4 or nf4=
	bnb_4bit_quant_type = "nf4"

	# Number of training epochs
	num_train_epochs = 1

	# Enable fp16 training
	fp16 = True

	# Enable bf16 training
	bf16 = False

	# Use packing dataset creating
	packing = False

	# Enable gradient checkpointing
	gradient_checkpointing = True

	# Optimizer to use, original is paged_adamw_32bit
	optim = "paged_adamw_32bit"

	# Fraction of steps to do a warmup for
	warmup_ratio = 0.03

	# Log every X updates steps
	logging_steps = 1

	# Load the entire model on the GPU 0
	device_map = {"": 0}

	# Visualize training
	report_to = "tensorboard"


	#2. Import du tokenizer.
	peft_config = LoraConfig(
	lora_alpha=lora_alpha,
	lora_dropout=lora_dropout,
	r=lora_r,
	inference_mode=False,
	task_type="CAUSAL_LM",
	target_modules = ["q_proj", "v_proj"]
	)

	tokenizer = AutoTokenizer.from_pretrained(model_name)

	# This is the fix for fp16 training
	#tokenizer.padding_side = "right"
	tokenizer.pad_token = tokenizer.eos_token

	#3. Préparation de la base de données

	from datasets import load_dataset

	def format_alpaca(sample):
	prompt = f"{sample['conversation']}"
	return prompt

	# template dataset to add prompt to each sample
	def template_dataset(sample):
	sample["text"] = f"{format_alpaca(sample)}{tokenizer.eos_token}"
	return sample

	# Chargement du dataset.
	#dataset = load_dataset("databricks/databricks-dolly-15k", split="train")
	data_files = {"train": "corpus_guillaume_tell_2.json"}
	dataset = load_dataset("json", data_files=data_files, split="train")

	# Shuffle the dataset
	dataset_shuffled = dataset.shuffle(seed=42)

	# Select the first 250 rows from the shuffled dataset, comment if you want 15k
	#dataset = dataset_shuffled.select(range(512))

	#Transformation du dataset pour utiliser le format guanaco
	dataset = dataset.map(template_dataset, remove_columns=list(dataset.features))

	print(dataset[40])

	#4. Import du modèle

	# Load tokenizer and model with QLoRA configuration
	compute_dtype = getattr(torch, bnb_4bit_compute_dtype)

	bnb_config = BitsAndBytesConfig(
	load_in_4bit=use_4bit,
	bnb_4bit_quant_type=bnb_4bit_quant_type,
	bnb_4bit_compute_dtype=compute_dtype,
	bnb_4bit_use_double_quant=use_nested_quant,
	)

	if compute_dtype == torch.float16 and use_4bit:
	major, _ = torch.cuda.get_device_capability()
	if major >= 8:
	print("=" * 80)
	print("Your GPU supports bfloat16, you can accelerate training with the argument --bf16")
	print("=" * 80)

	model = AutoModelForCausalLM.from_pretrained(
	model_name,
	device_map=device_map,
	quantization_config=bnb_config
	)

	model.config.use_cache = False
	model.config.pretraining_tp = 1

	#5. Fine-tuning

	torch.cuda.empty_cache()

	training_arguments = TrainingArguments(
	output_dir=output_dir,
	per_device_train_batch_size=per_device_train_batch_size,
	gradient_accumulation_steps=gradient_accumulation_steps,
	gradient_checkpointing=True,
	optim=optim,
	save_steps=save_steps,
	logging_steps=logging_steps,
	learning_rate=learning_rate,
	fp16=fp16,
	bf16=bf16,
	max_grad_norm=max_grad_norm,
	max_steps=max_steps,
	warmup_ratio=warmup_ratio,
	group_by_length=group_by_length,
	lr_scheduler_type=lr_scheduler_type,
	report_to="tensorboard"
	)

	trainer = SFTTrainer(
	model=model,
	train_dataset=dataset,
	peft_config=peft_config,
	dataset_text_field="text",
	max_seq_length=max_seq_length,
	tokenizer=tokenizer,
	args=training_arguments,
	packing=packing
	)

	trainer.train()
	#trainer.train(resume_from_checkpoint=True)

	#6. Sauvegarde

	model_to_save = trainer.model.module if hasattr(trainer.model, 'module') else trainer.model # Take care of distributed/parallel training
	model_to_save.save_pretrained(new_model_name)

	torch.cuda.empty_cache()

	from peft import AutoPeftModelForCausalLM

	model = AutoPeftModelForCausalLM.from_pretrained(new_model_name, device_map="auto", torch_dtype=torch.bfloat16)
	model = model.merge_and_unload()

	output_merged_dir = os.path.join(new_model_name, new_model_name)
	model.save_pretrained(output_merged_dir, safe_serialization=True)

	tokenizer.save_pretrained(output_merged_dir)