Spaces:
Sleeping
Sleeping
File size: 5,950 Bytes
bfe5d0e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 |
import os
from datasets import load_dataset
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
BitsAndBytesConfig,
TrainingArguments,
pipeline,
logging,
DataCollatorForLanguageModeling,
)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from trl import SFTTrainer
import torch
import logging
from torch.utils.data import DataLoader
import multiprocessing
# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
def preprocess_function(examples, tokenizer):
# Format the text
texts = []
for i in range(len(examples["text"])):
text = examples["text"][i]
texts.append(text)
# Tokenize the texts with shorter max length
tokenized = tokenizer(
texts,
padding=True,
truncation=True,
max_length=512, # Reduced from 1024 to 512
return_tensors="pt"
)
return tokenized
def main():
try:
# Load dataset
logger.info("Loading dataset...")
dataset = load_dataset("OpenAssistant/oasst1")
# Use a smaller subset for faster training
logger.info("Selecting smaller dataset subset...")
dataset["train"] = dataset["train"].select(range(2000)) # Reduced to 2k examples
# Model and tokenizer setup
logger.info("Setting up model and tokenizer...")
model_name = "microsoft/phi-2"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
# Preprocess dataset
logger.info("Preprocessing dataset...")
tokenized_dataset = dataset.map(
lambda x: preprocess_function(x, tokenizer),
batched=True,
remove_columns=dataset["train"].column_names,
num_proc=4 # Parallel processing for faster preprocessing
)
# Split dataset into train and eval
logger.info("Splitting dataset into train and eval sets...")
split_dataset = tokenized_dataset["train"].train_test_split(test_size=0.1, seed=42)
train_dataset = split_dataset["train"]
eval_dataset = split_dataset["test"]
# Configure 4-bit quantization with memory optimizations
logger.info("Configuring 4-bit quantization...")
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.float16,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_storage=torch.float16
)
# Load model with quantization and memory optimizations
logger.info("Loading model...")
model = AutoModelForCausalLM.from_pretrained(
model_name,
quantization_config=bnb_config,
device_map="auto",
trust_remote_code=True,
torch_dtype=torch.float16,
low_cpu_mem_usage=True
)
# Enable gradient checkpointing for memory efficiency
model.gradient_checkpointing_enable()
model.enable_input_require_grads()
# Prepare model for k-bit training
logger.info("Preparing model for k-bit training...")
model = prepare_model_for_kbit_training(model)
# LoRA configuration with optimized parameters
logger.info("Configuring LoRA...")
lora_config = LoraConfig(
r=8, # Reduced from 16 to 8
lora_alpha=16, # Reduced from 32 to 16
target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM"
)
# Get PEFT model
logger.info("Getting PEFT model...")
model = get_peft_model(model, lora_config)
# Create data collator
data_collator = DataCollatorForLanguageModeling(
tokenizer=tokenizer,
mlm=False
)
# Training arguments with memory-optimized settings
logger.info("Setting up training arguments...")
training_args = TrainingArguments(
output_dir="./phi2-qlora",
num_train_epochs=2,
per_device_train_batch_size=4, # Reduced from 16 to 4
per_device_eval_batch_size=4,
gradient_accumulation_steps=4, # Increased from 1 to 4
learning_rate=2e-4,
fp16=True,
logging_steps=5,
save_strategy="epoch",
evaluation_strategy="epoch",
# Additional optimizations
dataloader_num_workers=2, # Reduced from 4 to 2
dataloader_pin_memory=True,
warmup_ratio=0.05,
lr_scheduler_type="cosine",
optim="adamw_torch",
max_grad_norm=1.0,
group_by_length=True,
)
# Create trainer
logger.info("Creating trainer...")
trainer = SFTTrainer(
model=model,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
tokenizer=tokenizer,
args=training_args,
data_collator=data_collator,
)
# Train the model
logger.info("Starting training...")
trainer.train()
# Save the model
logger.info("Saving model...")
trainer.save_model("./phi2-qlora-final")
# Save tokenizer
logger.info("Saving tokenizer...")
tokenizer.save_pretrained("./phi2-qlora-final")
logger.info("Training completed successfully!")
except Exception as e:
logger.error(f"An error occurred: {str(e)}")
raise
if __name__ == "__main__":
multiprocessing.set_start_method('spawn')
main() |