211 lines
6.3 KiB
Python
211 lines
6.3 KiB
Python
import os
|
|
import torch
|
|
from datasets import load_dataset
|
|
from transformers import (
|
|
AutoTokenizer,
|
|
AutoModelForCausalLM,
|
|
TrainingArguments,
|
|
BitsAndBytesConfig,
|
|
)
|
|
from peft import (
|
|
LoraConfig,
|
|
get_peft_model,
|
|
prepare_model_for_kbit_training,
|
|
)
|
|
from trl import SFTTrainer
|
|
|
|
# ----------------------------
|
|
# Environment safety (Windows + AMP fixes)
|
|
# ----------------------------
|
|
os.environ["TORCHDYNAMO_DISABLE"] = "1"
|
|
os.environ["ACCELERATE_MIXED_PRECISION"] = "no" # ✅ disable AMP completely
|
|
os.environ["TORCH_AMP_DISABLE"] = "1" # ✅ disable GradScaler
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = "0" # optional: force first GPU
|
|
|
|
# ----------------------------
|
|
# Global configuration
|
|
# ----------------------------
|
|
MODEL_NAME = "Qwen/Qwen2.5-7B-Instruct"
|
|
OUTPUT_DIR = "./qwen2.5-7b-uk-fr-lora"
|
|
DATA_FILE = "paires_clean.json"
|
|
MAX_SEQ_LENGTH = 512 # Reduce for RTX 4080 SUPER VRAM
|
|
|
|
print(f"\n=== Starting fine-tuning script for {MODEL_NAME} ===\n")
|
|
|
|
# ----------------------------
|
|
# [1/7] Tokenizer
|
|
# ----------------------------
|
|
print(f"{80 * '_'}\n[1/7] Loading tokenizer...")
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
|
MODEL_NAME,
|
|
trust_remote_code=True,
|
|
)
|
|
tokenizer.pad_token = tokenizer.eos_token
|
|
tokenizer.model_max_length = MAX_SEQ_LENGTH
|
|
print("Tokenizer loaded.")
|
|
print(f"Pad token id: {tokenizer.pad_token_id}")
|
|
print(f"Max sequence length: {tokenizer.model_max_length}")
|
|
|
|
# ----------------------------
|
|
# [2/7] Load model in 4-bit (QLoRA)
|
|
# ----------------------------
|
|
print(f"{80 * '_'}\n[2/7] Loading model in 4-bit mode (QLoRA)...")
|
|
assert torch.cuda.is_available(), "CUDA GPU not detected!"
|
|
print(f"Using GPU: {torch.cuda.get_device_name(0)}")
|
|
|
|
bnb_config = BitsAndBytesConfig(
|
|
load_in_4bit=True,
|
|
bnb_4bit_quant_type="nf4",
|
|
bnb_4bit_compute_dtype=torch.float16, # fp16 internally
|
|
bnb_4bit_use_double_quant=True,
|
|
)
|
|
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
MODEL_NAME,
|
|
device_map="auto",
|
|
quantization_config=bnb_config,
|
|
low_cpu_mem_usage=True,
|
|
trust_remote_code=True,
|
|
)
|
|
|
|
# Align model tokens with tokenizer
|
|
model.config.pad_token_id = tokenizer.pad_token_id
|
|
model.config.bos_token_id = tokenizer.bos_token_id
|
|
model.config.eos_token_id = tokenizer.eos_token_id
|
|
print("Model loaded successfully in 4-bit mode on GPU.")
|
|
|
|
# ----------------------------
|
|
# [3/7] Prepare model for k-bit training
|
|
# ----------------------------
|
|
print(f"{80 * '_'}\n[3/7] Preparing model for k-bit training...")
|
|
model = prepare_model_for_kbit_training(model)
|
|
model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
|
|
model.config.use_cache = False # Important with gradient checkpointing + QLoRA
|
|
print("Model prepared for k-bit training.")
|
|
|
|
# ----------------------------
|
|
# [4/7] LoRA configuration
|
|
# ----------------------------
|
|
print(f"{80 * '_'}\n[4/7] Configuring LoRA adapters...")
|
|
lora_config = LoraConfig(
|
|
r=32,
|
|
lora_alpha=64,
|
|
lora_dropout=0.02,
|
|
bias="none",
|
|
task_type="CAUSAL_LM",
|
|
target_modules=[
|
|
"q_proj",
|
|
"k_proj",
|
|
"v_proj",
|
|
"o_proj",
|
|
"gate_proj",
|
|
"up_proj",
|
|
"down_proj",
|
|
],
|
|
)
|
|
model = get_peft_model(model, lora_config)
|
|
model.print_trainable_parameters()
|
|
print("LoRA adapters successfully attached.")
|
|
|
|
# ----------------------------
|
|
# [5/7] Dataset loading & formatting
|
|
# ----------------------------
|
|
print(f"{80 * '_'}\n[5/7] Loading dataset from JSON file...")
|
|
dataset = load_dataset("json", data_files=DATA_FILE)
|
|
print(f"Dataset loaded with {len(dataset['train'])} samples.")
|
|
|
|
print("Formatting dataset for Ukrainian → French translation...")
|
|
def format_prompt(example):
|
|
return {
|
|
"text": (
|
|
"<|im_start|>user\n"
|
|
"Translate the following Ukrainian text into French.\n"
|
|
f"Ukrainian: {example['text']}\n"
|
|
"<|im_end|>\n"
|
|
"<|im_start|>assistant\n"
|
|
f"{example['translation']}"
|
|
"<|im_end|>"
|
|
)
|
|
}
|
|
|
|
dataset = dataset.map(
|
|
format_prompt,
|
|
remove_columns=dataset["train"].column_names
|
|
)
|
|
print("Dataset formatting completed.")
|
|
print("Example prompt:\n")
|
|
print(dataset["train"][0]["text"])
|
|
|
|
# ----------------------------
|
|
# [6/7] Training arguments
|
|
# ----------------------------
|
|
print(f"{80 * '_'}\n[6/7] Initializing training arguments...")
|
|
training_args = TrainingArguments(
|
|
output_dir=OUTPUT_DIR,
|
|
per_device_train_batch_size=1,
|
|
gradient_accumulation_steps=16,
|
|
learning_rate=1e-4,
|
|
num_train_epochs=2,
|
|
max_steps=1000,
|
|
|
|
fp16=False, # ⚠ disable AMP
|
|
bf16=False, # ⚠ disable BF16
|
|
|
|
optim="paged_adamw_32bit",
|
|
logging_steps=10,
|
|
save_steps=500,
|
|
save_total_limit=2,
|
|
report_to="none",
|
|
|
|
dataloader_pin_memory=False,
|
|
max_grad_norm=0.0, # avoid AMP gradient clipping
|
|
)
|
|
print("Training arguments ready.")
|
|
print(f"Output directory: {OUTPUT_DIR}")
|
|
print(f"Epochs: {training_args.num_train_epochs}")
|
|
print(f"Effective batch size: {training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps}")
|
|
|
|
# ----------------------------
|
|
# [7/7] Trainer
|
|
# ----------------------------
|
|
print(f"{80 * '_'}\nInitializing SFTTrainer...")
|
|
trainer = SFTTrainer(
|
|
model=model,
|
|
train_dataset=dataset["train"],
|
|
args=training_args,
|
|
)
|
|
print("Trainer initialized.")
|
|
|
|
# ----------------------------
|
|
# Training
|
|
# ----------------------------
|
|
print(f"{80 * '_'}\nStarting training...")
|
|
checkpoint_exists = False
|
|
if os.path.exists(OUTPUT_DIR):
|
|
checkpoint_exists = any(
|
|
d.startswith("checkpoint-")
|
|
for d in os.listdir(OUTPUT_DIR)
|
|
)
|
|
|
|
if checkpoint_exists:
|
|
print("Checkpoint found → resuming training")
|
|
train_output = trainer.train(resume_from_checkpoint=True)
|
|
else:
|
|
print("No checkpoint found → starting fresh training")
|
|
train_output = trainer.train()
|
|
|
|
print("\n=== Training summary ===")
|
|
print(f"Global steps: {train_output.global_step}")
|
|
print(f"Training loss: {train_output.training_loss}")
|
|
print(f"Metrics: {train_output.metrics}")
|
|
print("Training completed successfully.")
|
|
|
|
# ----------------------------
|
|
# Save LoRA adapter and tokenizer
|
|
# ----------------------------
|
|
print(f"{80 * '_'}\nSaving LoRA adapter and tokenizer...")
|
|
trainer.model.save_pretrained(OUTPUT_DIR)
|
|
tokenizer.save_pretrained(OUTPUT_DIR)
|
|
print("\n=== Fine-tuning finished ===")
|
|
print(f"LoRA adapter saved in: {OUTPUT_DIR}")
|