TOWARDSAI.NET
How to Fine-Tune Language Models: First Principles to Scalable Performance
Author(s): Ehssan Originally published on Towards AI. Image by AuthorIn this article, well explore the process of fine-tuning language models for text classification. Well do so in three levels: first, by manually adding a classification head in PyTorch* and training the model so you can see the full process; second, by using the Hugging Face* Transformers library to streamline the process; and third, by leveraging PyTorch Lightning* and accelerators to optimize training performance. By the end of this guide, youll have a well-rounded understanding of the fine-tuning workflow.IntroductionThe idea behind using fine-tuning in Natural Language Processing (NLP) was borrowed from Computer Vision (CV). The CV models were first trained on large datasets such as ImageNet to teach them the basic features of images such as edges or colors. These pretrained models were then fine-tuned on a downstream task such as classifying birds with a relatively small number of labeled examples.Fine-tuned models typically achieved a higher accuracy than supervised models trained from scratch on the same amount of labeled data.Despite the popularity and success of transfer learning in CV, for many years it wasnt clear what the analogous pretraining process was for NLP. Consequently, NLP applications required large amounts of labeled data to achieve high performance.How is Fine-tuning Different from Pretraining?With pretraining, language models gain a general understanding of languages. During this process, they learn language patterns but typically are not capable of following instructions or answering questions. In the case of GPT models, this self-supervised learning includes predicting the next word (unidirectional) based on their training data, which is often webpages. In the case of BERT (Bidirectional Encoder Representations from Transformers), learning involves predicting randomly masked words (bidirectional) and sentence-order prediction. But how can we adapt language models for our own data or our own tasks?Fine-tuning continues training a pretrained model to increase its performance on specific tasks. For instance, through instruction fine-tuning you can teach a model to behave more like a chatbot. This is the process for specializing a general purpose model like OpenAI* GPT-4 into an application like ChatGPT* or GitHub* Copilot. By fine-tuning your own language model, you can increase the reliability, performance, and privacy of your model while reducing the associated inference costs compared to subscription-based services, especially if you have a large volume of data or frequent requests.Fine-Tuning a Language Model for Text ClassificationPreprocessing and Preparing DataLoaderFeel free to skip this section if youre comfortable with preprocessing data. Throughout we assume that we have our labeled data saved in train, validation, and test csv files each with a text and a label column. For training, the labels should be numeric, so if thats not the case, youcan use a label_to_id dictionary such as {"negative": 0, "positive": 1} and do a mapping to get the desired format.For concreteness, we will use BERT as the base model and set the number of classification labels to 4. After running the code below, you are encouraged to swap BERT for DistilBERT which reduces the size of the BERT model by 40%, speeding inference by 60%, while retaining 97% of BERTs language understanding capabilities.A Quick Look at BERTBERT was introduced by Google in 2018 and has since revolutionized the field of NLP. Unlike traditional models that process text in a unidirectional manner, BERT is designed to understand the context of a word in a sentence by looking at both its left and right surroundings. This bidirectional approach allows BERT to capture the nuances of language more effectively.Key Features of BERTPretraining: BERT is pretrained on a massive corpus of text, including the entire Wikipedia and BookCorpus. The pretraining involves two tasks: Masked Language Modeling (MLM) and Next Sentence Prediction (NSP).Architecture: BERT_BASE has 12 layers (transformer blocks), 768 hidden units, and 12 attention heads, totaling 110 million parameters.You can run this tutorial on Intel Tiber AI Cloud, using an Intel Xeon CPU instance. This platform provides ample computing resources for smooth execution of our code.import osimport torchfrom torch.utils.data import DataLoader, Datasetfrom transformers import AutoTokenizerimport pandas as pd# Parametersmodel_ckpt = "bert-base-uncased"num_labels = 4batch_size = 16num_workers = 6# Load the tokenizertokenizer = AutoTokenizer.from_pretrained(model_ckpt)# Custom Dataset classclass TextDataset(Dataset): def __init__(self, dataframe, tokenizer, max_length=512): self.data = dataframe self.tokenizer = tokenizer self.max_length = max_length def __len__(self): return len(self.data) def __getitem__(self, idx): row = self.data.iloc[idx] text = row["text"] # Replace "text" with your actual column name for text label = row["label"] # Replace "label" with your actual column name for labels # Tokenize the input text encoding = self.tokenizer( text, max_length=self.max_length, padding="max_length", truncation=True, return_tensors="pt", ) return { "input_ids": encoding["input_ids"].squeeze(0), # Remove batch dimension with squeeze "attention_mask": encoding["attention_mask"].squeeze(0), "label": torch.tensor(label, dtype=torch.long), }os.environ["TOKENIZERS_PARALLELISM"] = "false"# Load csv filestrain_df = pd.read_csv("train.csv")val_df = pd.read_csv("val.csv")test_df = pd.read_csv("test.csv")# Create Dataset objectstrain_dataset = TextDataset(train_df, tokenizer)val_dataset = TextDataset(val_df, tokenizer)test_dataset = TextDataset(test_df, tokenizer)# Create DataLoaderstrain_loader = DataLoader( train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers)val_loader = DataLoader(val_dataset, batch_size=batch_size, num_workers=num_workers)test_loader = DataLoader(test_dataset, batch_size=batch_size, num_workers=num_workers)The Classification Token [CLS]The [CLS] token is typically added at the beginning of the input sequence in transformer-based models such as BERT and its variants. During fine-tuning, the model learns to assign meaningful information to the [CLS] token, which aggregates the input sequences context. The last hidden state corresponding to the [CLS] token is then used as a representation of the entire input, which can be passed through a classifier layer for downstream tasks like sentiment analysis, topic categorization, or any task requiring a decision based on the entire sequence. This mechanism allows the model to focus on both the global understanding of the text and task-specific features for accurate predictions.Unlike traditional models that may rely on static embeddings (like word2vec), transformers generate contextualized embeddings, so that the meaning of a token depends on the tokens around it. The [CLS] token, as it passes through the layers, becomes increasingly aware of the entire sequences meaning, which makes it a good summary representation for downstream tasks. For some tasks, especially those requiring finer-grained understanding, other strategies might be employed. For instance, for document classification, where every word contributes equally, some models use mean pooling over all token embeddings.Level 1: PyTorchIn this section, we manually add a classification head to the base model and do the fine-tuning. We achieve this using the AutoModelclass which converts the tokens (or rather token encodings) to embeddings and then then feeds them through the encoder stack to return the hidden states. While AutoModel is helpful for understanding the idea behind what were doing, to fine-tune for text classification its better practice to work withAutoModelForSequenceClassification instead, as we discuss below.import torchfrom torch import nnfrom transformers import AutoModel# Load the base model with AutoModel and add a classifierclass CustomModel(nn.Module): def __init__(self, model_ckpt, num_labels): super(CustomModel, self).__init__() self.model = AutoModel.from_pretrained(model_ckpt) # Base transformer model self.classifier = nn.Linear( self.model.config.hidden_size, num_labels ) # Classification head. The 1st parameter equals 768 for BERT as discussed above def forward(self, input_ids, attention_mask): # Forward pass through the transformer model outputs = self.model(input_ids=input_ids, attention_mask=attention_mask) # Use the [CLS] token (0-th token in the sequence) for classification cls_output = outputs.last_hidden_state[ :, 0, : ] # Shape: (batch_size, hidden_size) # Pass through the classifier head logits = self.classifier(cls_output) return logits# Initialize the modelmodel = CustomModel(model_ckpt, num_labels)# Loss function and optimizerloss_fn = nn.CrossEntropyLoss()optimizer = torch.optim.Adam(model.parameters(), lr=5e-5)# Training functiondef train(model, optimizer, train_loader, loss_fn): model.train() total_loss = 0 for batch in train_loader: optimizer.zero_grad() # Unpack the batch data input_ids = batch["input_ids"] attention_mask = batch["attention_mask"] label = batch["label"] # Forward pass output = model(input_ids, attention_mask) # Compute loss loss = loss_fn(output, label) loss.backward() # Update the model parameters optimizer.step() total_loss += loss.item() print(f"Train loss: {total_loss / len(train_loader):.2f}")import torchdef evaluate(model, test_loader, loss_fn): model.eval() # Set model to evaluation mode total_loss = 0 total_acc = 0 total_samples = 0 with torch.no_grad(): # No gradient computation needed during evaluation for batch in test_loader: input_ids = batch["input_ids"] attention_mask = batch["attention_mask"] labels = batch["label"] # Forward pass output = model(input_ids, attention_mask) # Compute loss loss = loss_fn(output, labels) total_loss += loss.item() # Compute accuracy predictions = torch.argmax(output, dim=1) total_acc += torch.sum(predictions == labels).item() total_samples += labels.size(0) # Calculate average loss and accuracy avg_loss = total_loss / len(test_loader) avg_acc = total_acc / total_samples * 100 print(f"Test loss: {avg_loss:.2f}, Test acc: {avg_acc:.2f}%")Finally, we can train, evaluate, and save the model.num_epochs = 3for epoch in range(num_epochs): train(model, optimizer, train_loader, loss_fn)evaluate(model, test_loader, loss_fn)torch.save(model.state_dict(), "./fine-tuned-model.pt")Level 2: Hugging Face TransformersNow, we use the convenience of AutoModelForSequenceClassification class that will add the classification head to the base model automatically. Compare this against what we did with the AutoModel class in the previous section!Also note that the Trainer class from Hugging Faces Transformerslibrary can directly handle Dataset objects without needing aDataLoader, as it automatically handles batching and shuffling foryou.from transformers import AutoModelForSequenceClassification, Trainer, TrainingArgumentsmodel = AutoModelForSequenceClassification.from_pretrained( model_ckpt, num_labels=num_labels)training_args = TrainingArguments( output_dir="./results", num_train_epochs=3, warmup_steps=500, weight_decay=0.01, logging_dir="./logs", logging_steps=10, # Log every 10 steps evaluation_strategy="steps", save_steps=500, # Save model checkpoint every 500 steps load_best_model_at_end=True, # Load the best model at the end of training metric_for_best_model="accuracy",)# Train the modeltrainer = Trainer( model=model, args=training_args, train_dataset=train_dataset, eval_dataset=test_dataset,)trainer.train()trainer.evaluate(test_dataset)Level 3: PyTorch LightningLightning is, in the words of its documentation, the deep learning framework with batteries included for professional AI researchers and machine learning engineers who need maximal flexibility while super-charging performance at scale.As we shall see, with a bit of additional organizational code, the Lightning Trainer automates the following:Epoch and batch iterationoptimizer.step(), loss.backward(), optimizer.zero_grad() callsCalling of model.eval(), enabling and disabling grads during evaluationCheckpoint Saving and LoadingLoggingAccelerator, Multi-GPU, and TPU Support (No .to(device) calls required.)Mixed-precision trainingYou can accelerate training with Intel Gaudi processors, which allow you to conduct more deep learning training at a lower expense. You can try an Intel Gaudi instance for free on Intel Tiber AI Cloud.import torchmetricsimport lightning as Lfrom lightning.pytorch.callbacks import ModelCheckpointfrom lightning.pytorch.loggers import TensorBoardLoggerfrom transformers import AutoModelForSequenceClassification# A LightningModule is a torch.nn.Module with added functionality.# It wraps around a regular PyTorch model.class LightningModel(L.LightningModule): def __init__(self, model, learning_rate=5e-5): super().__init__() self.learning_rate = learning_rate self.model = model self.val_acc = torchmetrics.Accuracy(task="multiclass", num_classes=num_labels) self.test_acc = torchmetrics.Accuracy(task="multiclass", num_classes=num_labels) def forward(self, input_ids, attention_mask, labels): return self.model(input_ids, attention_mask=attention_mask, labels=labels) def _shared_step(self, batch, batch_idx): outputs = self( batch["input_ids"], attention_mask=batch["attention_mask"], labels=batch["label"], ) return outputs def training_step(self, batch, batch_idx): outputs = self._shared_step(batch, batch_idx) self.log("train_loss", outputs["loss"]) return outputs["loss"] def validation_step(self, batch, batch_idx): outputs = self._shared_step(batch, batch_idx) self.log("val_loss", outputs["loss"], prog_bar=True) logits = outputs["logits"] self.val_acc(logits, batch["label"]) self.log("val_acc", self.val_acc, prog_bar=True) def test_step(self, batch, batch_idx): outputs = self._shared_step(batch, batch_idx) logits = outputs["logits"] self.test_acc(logits, batch["label"]) self.log("accuracy", self.test_acc, prog_bar=True) def configure_optimizers(self): optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate) return optimizermodel = AutoModelForSequenceClassification.from_pretrained( model_ckpt, num_labels=num_labels)lightning_model = LightningModel(model)callbacks = [ ModelCheckpoint(save_top_k=1, mode="max", monitor="val_acc") # Save top 1 model]logger = TensorBoardLogger(save_dir="./logs", name="fine-tuned-model")trainer = L.Trainer( max_epochs=3, callbacks=callbacks, accelerator="hpu", precision="bf16-mixed", # By default, HPU training uses 32-bit precision. To enable mixed precision, set the precision flag. devices="auto", logger=logger, log_every_n_steps=10,)trainer.fit(lightning_model, train_dataloaders=train_loader, val_dataloaders=val_loader)trainer.test(lightning_model, train_loader, ckpt_path="best")trainer.test(lightning_model, val_loader, ckpt_path="best")trainer.test(lightning_model, test_loader, ckpt_path="best")While the Transformers Trainer class supports distributed training, it doesnt offer the same level of integration and flexibility as Lightning when it comes to advanced features like custom callbacks, logging, and seamless scaling across multiple GPUs or nodes.Practical AdviceNow that youre familiar with the fine-tuning process, you mightwonder how you can apply it to your specific task. Heres somepractical advice:Collect real data for your task, or generate synthetic data. See, for instance, Synthetic Data Generation with Language Models: A Practical Guide.Fine-tune a relatively small model.Evaluate your language model on your test set, and on a benchmark if available for your task.Increase the training dataset size, base model size, and, if necessary, task complexity.Keep in mind that the standard or conventional fine-tuning of languagemodels as described in this writing can be expensive. Rather thanupdating all the weights and biases, we could update only the lastlayer as follows:# Freeze all layersfor param in model.parameters(): param.requires_grad = False# Unfreeze last layerfor param in model.pre_classifier.parameters(): param.requires_grad = Truefor param in model.classifier.parameters(): param.requires_grad = TrueIn future articles, we shall discuss more efficient fine-tuning techniques, so stay tuned!For more AI development how-to content, visit Intel AI Development.ResourcesAcknowledgmentsThe author thanks Jack Erickson for providing detailed feedback on an earlier draft of this work.Suggested Reading*Other names and brands may be claimed as the property of others.Join thousands of data leaders on the AI newsletter. Join over 80,000 subscribers and keep up to date with the latest developments in AI. From research to projects and ideas. If you are building an AI startup, an AI-related product, or a service, we invite you to consider becoming asponsor. Published via Towards AI
0 Commenti 0 condivisioni 48 Views