Efficient Feature Extraction with LoRA-Enhanced LLaMA: Optimizing Memory for Large-Scale Fine-Tuning

Introduction

Fine-tuning large language models (LLMs) has become a powerful approach for adapting pre-trained models to specific tasks, such as text classification, named entity recognition (NER), and feature extraction. However, full fine-tuning is often expensive and resource-intensive, requiring substantial GPU memory and compute power.

To address this challenge, LoRA (Low-Rank Adaptation) has emerged as a lightweight fine-tuning technique that significantly reduces the number of trainable parameters while maintaining competitive performance. In this article, we explore how LoRA can be applied to feature extraction from documents.

The Challenges of Fine-Tuning Large Models

Fine-tuning involves training an existing model on a smaller dataset to adapt it to a specific task. While this technique has proven effective, it presents several key challenges:

1. High Memory Consumption: Large models like LLaMA, Gemma, and GPT-3 require enormous memory for fine-tuning, making them impractical for many users.
2. Compute Costs: Fine-tuning requires significant GPU resources, making it expensive for companies and researchers.
3. Overfitting Risks: Training all model parameters on a small dataset can lead to overfitting, reducing the model’s ability to generalize.
4. Catastrophic Forgetting: Standard fine-tuning can overwrite the model’s pre-trained knowledge, leading to degraded performance on other tasks.

These challenges call for a more efficient fine-tuning approach, which is where LoRA comes into play.

What is LoRA?

LoRA (Low-Rank Adaptation) is a parameter-efficient fine-tuning method that modifies only a small subset of model parameters while keeping the rest frozen. Instead of updating all weights, LoRA introduces low-rank matrices that are trained alongside the model, significantly reducing memory requirements.

How LoRA Works

LoRA applies the following key techniques:

• Freezing the Pre-Trained Model: Instead of modifying the entire model, LoRA freezes the main weights and injects trainable low-rank matrices into specific layers.
• Targeting Key Layers: LoRA modifies only the query (q_proj) and value (v_proj) layers in attention mechanisms, which are the most crucial for adapting to new tasks.
• Reducing Trainable Parameters: By keeping rank r small (e.g., 4 or 8), LoRA reduces the number of learnable parameters by orders of magnitude.

This approach allows LoRA to achieve comparable fine-tuning results while using much less memory and compute power.

My Use Case: Feature Extraction from Documents

Extracting Named Entities

For this case study, I focused on using LoRA fine-tuning to extract structured features from unstructured text. Specifically, I wanted to identify named entities like people, organizations, locations, and miscellaneous names in documents.

To train the model, I used the WNUT-17 dataset, which is designed for recognizing emerging and rare named entities in user-generated content like social media posts.

What is the WNUT-17 Dataset?

The WNUT-17 dataset includes:

• Tokens: Individual words from informal text sources.
• NER Tags: Labels for each token showing whether it’s part of an entity.
• Entity Types: Six main categories:
• person
• location
• corporation
• product
• creative_work
• group

A sample record from WNUT-17 looks like this:

{
  "tokens": ["@paulwalk", "It", "'s", "the", "view", "from", "where", "I", "'m", "living", "for", "two", "weeks", ".", "Empire", "State", "Building", "=", "ESB", "."],
  "ner_tags": ["B-person", "O", "O", "O", "O", "O", "O", "O", "O", "O", "O", "O", "O", "O", "B-location", "I-location", "I-location", "O", "B-location", "O"]
}

Here, “Empire State Building” is labeled as a Location, and “@paulwalk” is identified as a Person.

!pip install --upgrade datasets
!pip install --upgrade transformers
!pip install --upgrade peft
!pip install --upgrade trl
!pip install bitsandbytes
!pip install accelerate
!pip install tensorboard
!pip install jsonlines
!pip install evaluate
!pip install seqeval

from huggingface_hub import notebook_login
from google.colab import userdata

hf_token = userdata.get('hf_token_mz')
notebook_login()

This code sets up the environment for fine-tuning large language models using LoRA (Low-Rank Adaptation). It installs all necessary libraries, including Hugging Face’s datasets and transformers, LoRA (via peft), and evaluation tools like seqeval. The script also logs into the Hugging Face Hub using a stored token (hf_token) for accessing pre-trained models and datasets.

Preprocessing and Tokenization

Before fine-tuning, I needed to preprocess the dataset and ensure that entity labels were properly aligned with tokenized text. Since transformer models break words into subwords, each entity label had to be correctly assigned to all corresponding subword tokens. This process involved:

1. Loading and inspecting the dataset to extract entity labels and convert them into ID-based mappings for model training.
2. Tokenizing the text using a pre-trained tokenizer while maintaining word boundaries.
3. Aligning entity labels with tokenized subwords so that labels remain meaningful even when words are split into multiple tokens.
4. Handling special tokens and padding, ensuring that [CLS], [SEP], and padding tokens don’t interfere with training by assigning them a default ignored label (100).

This preprocessing step ensures that the model receives properly structured inputs, making it ready for fine-tuning with LoRA.

dataset = load_dataset("wnut_17")
print(dataset)

labels = dataset["train"].features["ner_tags"].feature.names
label2id = {label: i for i, label in enumerate(labels)}
id2label = {i: label for label, i in label2id.items()}
num_labels = len(labels)
print("Labels:", labels)
print("Number of labels:", num_labels)

def tokenize_and_align_labels(examples):
    """
    Tokenizes input 'tokens' while aligning the word-level ner_tags
    to the resulting subword tokens.
    """
    tokenized_inputs = tokenizer(
        examples["tokens"],
        truncation=True,
        is_split_into_words=True,
        return_offsets_mapping=True  # needed for alignment
    )
    offsets = tokenized_inputs.pop("offset_mapping")

    labels_aligned = []
    word_ids = tokenized_inputs.word_ids()
    for word_id in word_ids:
        if word_id is None:
            # Special tokens -> label = -100 (ignored in training)
            labels_aligned.append(-100)
        else:
            # Use the original word-level label
            labels_aligned.append(examples["ner_tags"][word_id])

    tokenized_inputs["labels"] = labels_aligned
    return tokenized_inputs

model_name = "meta-llama/Llama-2-7b-chat-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token)
# Set EOS token as PAD token
tokenizer.pad_token = tokenizer.eos_token
encoded_dataset = dataset.map(
    tokenize_and_align_labels,
    batched=False,   # WNUT data is already splitted into "tokens", so we can do batched=False or True
    remove_columns=dataset["train"].column_names
)

print(encoded_dataset)

# We'll name them to match your variables:
encoded_dataset_train = encoded_dataset["train"]
encoded_dataset_val   = encoded_dataset["validation"]
encoded_dataset_test  = encoded_dataset["test"]  

Optimizing Model with 4-bit Quantization

Before fine-tuning with LoRA, I optimized the model for efficient memory usage by loading it in 4-bit precision using BitsAndBytesConfig. This significantly reduces GPU memory requirements while maintaining model performance. The model was then prepared for LoRA training using prepare_model_for_kbit_training, allowing for parameter-efficient fine-tuning without modifying the entire model.

import torch
from datasets import load_dataset
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    TrainingArguments,
    Trainer,
    BitsAndBytesConfig
)
from peft import (
    LoraConfig,
    TaskType,
    get_peft_model,
    prepare_model_for_kbit_training
)

bnb_4bit_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4"  # NormalFloat4
)

#  Load base model in 4-bit + tokenizer
# -------------------------------------------------------------------

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_4bit_config,  # pass your 4-bit config
    device_map="auto",                    # let accelerate handle device placement
    use_auth_token=hf_token
)

# Prepare the model for 4-bit training
model = prepare_model_for_kbit_training(model)

Configuring LoRA for Memory-Efficient Fine-Tuning

After preparing the model, I configured LoRA to fine-tune only the query (q_proj) and value (v_proj) layers, making the process highly memory-efficient. The LoRA configuration included:

• Rank (r=4): Defines the low-rank decomposition of weight updates.
• Scaling factor (lora_alpha=16): Controls the influence of LoRA updates.
• Dropout (lora_dropout=0.05): Helps prevent overfitting.
• Bias strategy (bias="none"): Ensures only LoRA parameters are updated.
• Task type (TaskType.CAUSAL_LM): Specifies that this is for causal language modeling.

I also enabled gradient checkpointing to further reduce memory usage during training.

#  LoRA config (memory-friendly)
# -------------------------------------------------------------------
peft_config = LoraConfig(
    r=4,
    lora_alpha=16,
    target_modules=["q_proj", "v_proj"],  # only Q and V
    lora_dropout=0.05,
    bias="none",
    task_type=TaskType.CAUSAL_LM  # required by PEFT for a causal LM approach
)

lora_model = get_peft_model(model, peft_config)
print(lora_model)

# Optionally enable gradient checkpointing to reduce memory:
lora_model.gradient_checkpointing_enable()

The next step is to define the optimizer to ensure efficient training for LLaMA-2 (7B) with LoRA and 4-bit quantization. Instead of using AdamW, the setup utilizes Adam8bit from bitsandbytes, which is specifically optimized for low-bit precision models. The Trainer API orchestrates training, handling gradient updates, evaluation, and checkpointing. By fine-tuning only the LoRA adapter parameters, the setup ensures efficient training on commodity GPUs while achieving high performance.

from transformers import AdamW, get_scheduler
import bitsandbytes as bnb
import evaluate
import numpy as np

optimizer = bnb.optim.Adam8bit(
    lora_model.parameters(),  # Use LoRA model parameters
    lr=2e-4,
)

training_args = TrainingArguments(
    output_dir="my-wnut17-lora-4bit",
    evaluation_strategy="epoch",
    save_strategy="epoch",
    learning_rate=3e-5,
    num_train_epochs=15,
    per_device_train_batch_size=2,
    per_device_eval_batch_size=2,
    logging_dir="logs",
    logging_steps=50,
    report_to="tensorboard",
    fp16=True,     # often still beneficial for matmul ops
    push_to_hub=False
)

# 
seqeval_metric = evaluate.load("seqeval")
def compute_metrics(eval_preds):
    logits, labels = eval_preds
    predictions = np.argmax(logits, axis=-1)

    true_labels_list = []
    true_preds_list = []

    for pred_seq, label_seq in zip(predictions, labels):
        temp_labels = []
        temp_preds = []
        for p, l in zip(pred_seq, label_seq):
            if l == -100:  # Ignore padding tokens
                continue
            temp_labels.append(id2label.get(l, "O"))  # ✅ Prevent KeyError
            temp_preds.append(id2label.get(p, "O"))  # ✅ Prevent KeyError

        true_labels_list.append(temp_labels)
        true_preds_list.append(temp_preds)

    # 🔍 Debugging: Print the first few label pairs
    print("\n=== Debug: True vs. Predicted Labels ===")
    for i in range(min(3, len(true_labels_list))):
        print(f"True: {true_labels_list[i]}")
        print(f"Pred: {true_preds_list[i]}")

    # Compute NER metrics
    results = seqeval_metric.compute(
        predictions=true_preds_list,
        references=true_labels_list
    )

    # 🔍 Debugging: Print computed metrics
    print("\n=== Debug: Computed Metrics ===")
    print(results)

    # ✅ Ensure required keys exist, return defaults if missing
    return {
        "precision": results.get("overall_precision", 0.0),
        "recall": results.get("overall_recall", 0.0),
        "f1": results.get("overall_f1", 0.0),
        "accuracy": results.get("overall_accuracy", 0.0)
    }

from transformers import DataCollatorForTokenClassification
data_collator = DataCollatorForTokenClassification(tokenizer=tokenizer, return_tensors="pt")
lora_model.gradient_checkpointing_enable()
small_train_dataset = encoded_dataset_train.select(range(1000))
small_val_dataset = encoded_dataset_val.select(range(100))
trainer = Trainer(
    model=lora_model,
    args=training_args,
    train_dataset=small_train_dataset,
    eval_dataset=small_val_dataset,
    tokenizer=tokenizer,
    data_collator=data_collator,
    optimizers=(optimizer, None), 
    compute_metrics=compute_metrics
)

trainer.train()

The last step is to initiate the training process using trainer.train(). During training, the model iteratively optimizes its parameters across multiple epochs. As the training progresses, the loss decreases steadily, indicating that the model is learning to better align its predictions with the true labels.

This study employs several key techniques to reduce memory consumption while fine-tuning LLaMA-2 (7B) for NER. 4-bit quantization minimizes memory usage by storing model weights in lower precision, significantly reducing storage and computational overhead. LoRA (Low-Rank Adaptation) optimizes only a small subset of trainable parameters, avoiding the need to update the full model, which drastically cuts memory requirements. Adam8bit optimizer further reduces memory consumption by shrinking optimizer states compared to traditional AdamW. Mixed-precision training (fp16=True) enables computations in lower precision where possible, reducing memory load while maintaining stability. Gradient checkpointing optimizes memory usage by recomputing intermediate activations during backpropagation instead of storing them, lowering memory overhead at the cost of slightly increased computation. Finally, dataset subsampling ensures efficient experimentation by limiting the number of training examples, making fine-tuning feasible on commodity GPUs. These techniques collectively enable efficient training without sacrificing model performance.