Upload finetune_scikitllm.py

finetune_scikitllm.py

@@ -0,0 +1,236 @@
+import os
+import torch
+
+#This is the script used to finetune the scikit-llm model.
+#It also contains all the hyperparameters used for training and should be fully reproducible.
+
+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
+
+# We use a previously finetuned model of Mistral, Mistral-Hermes.
+#It already includes many instruction-based features (including the chatml syntax) that makes it easier to finetune. 
+model_name = "mistral-hermes-2.5"
+
+torch.cuda.empty_cache()
+
+# The name of the model.
+new_model_name = "mistral-skikit-reference"
+
+# The output directory where the model predictions and checkpoints will be written
+output_dir = "./mistral-skikit-reference"
+
+# Tensorboard logs
+tb_log_dir = "./mistral-skikit-reference/logs"
+
+# The number of steps. Since we chose a lower learning rate, we took on a long training (8 epochs). Could be lower.
+max_steps = 1200
+
+# Les paramètres importants !!
+per_device_train_batch_size = 4 #Number of batches to send per step. Optimal given our GPU vram.
+learning_rate = 2e-5 #The most important hyperparmater. We take a lower value as mistral-hermes is already finetuned and we want to keep the capacities.
+max_seq_length = 4096 #Context window length. Here we are constrained by Hermes, but Mistral is up to 8128 (32k in the new version)
+save_steps = 1000 # Automated saving of the steps.
+lr_scheduler_type = "linear" #Learning rate scheduler. Better to decrease the learning rate for long training. I prefer linear over to cosine as it is more predictable: easier to restart training if needed.
+
+
+#Other parameters. I don't usually tweak thoses.
+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
+#We go back to 16-bit for inference.
+#Currently this speeds up training significantly we nearly no quality impact.
+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
+#(not used in practice)
+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. Loading the 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.pad_token = tokenizer.eos_token
+
+#3. Preparing the dataset.
+#This is the part most specific to the scikit model.
+#We take an entire conversation, as both the input and the output are part of the same string of texts.
+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
+
+# Loading the data du dataset.
+data_files = {"train": "skikit_administration.json"}
+dataset = load_dataset("json", data_files=data_files, split="train")
+
+# Shuffle the dataset
+dataset_shuffled = dataset.shuffle(seed=42)
+
+#Dataset parsing.
+dataset = dataset.map(template_dataset, remove_columns=list(dataset.features))
+
+print(dataset[40])
+
+#4. Model import
+
+# 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 (actually)
+#We pass all the hyperparmeters, and are ready to go.
+
+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
+)
+
+#Training:
+trainer.train()
+
+#Optionally, if we want to continue training (for instance if there was an issue):
+#trainer.train(resume_from_checkpoint=True)
+
+#6. Export the weights
+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)
+
+#We also save the tokenizer
+tokenizer.save_pretrained(output_merged_dir)