Supervised Fine-Tuning (SFT) for PII Masking Using Axolotl

Introduction and Motivation

A vast majority of an organization's data is unstructured (often cited as over 80%). This includes text from sources like emails, internal/external chats, call transcripts, legal contracts, and survey responses. The challenge with this data is its lack of a predefined format. Unlike a structured database with clear labels, identifying sensitive information in unstructured text is a complex task because it requires the ability to interpret context (e.g., telling a zip code from an order ID), resolve ambiguity (distinguishing a sensitive birthday from a simple meeting date), and process varied human expression (like the many different ways a phone number can be written - (555) 123-4567 vs 555.123.4567 vs 555-123-4567). A string of digits, for instance, could be a phone number, a zip code, or an order ID.

This complexity makes it difficult to systematically manage Personally Identifiable Information (PII) within the text, exposing organizations to significant legal and financial risks under regulations like GDPR and CCPA, and compromising user privacy. The only viable solution is to automatically detect and redact this information with a system that understands context.

Simple rule-based systems like regular expressions (regex) are too rigid to handle this nuance, leading to both missed PII and incorrectly flagged data. This necessitates the use of more complex Large Language Models (LLMs), which can understand language context to accurately and reliably identify sensitive information.

This blog post will guide you through the process of fine-tuning a Large Language Model (LLM) for PII redaction using the Axolotl framework, which is designed to simplify this process.

The Landscape of PII Redaction Tools

Before diving into fine-tuning, it's worth noting that a rich ecosystem of tools already exists for PII detection.

• Microsoft Presidio: An open-source framework for detecting and anonymizing PII in text and images. It comes with built-in regex and ML-based recognizers.
• spaCy + Custom NER Models: Allows training domain-specific models for tagging sensitive entities with strong accuracy when paired with quality data.
• Hugging Face Models & Datasets: Open-source models fine-tuned for PII detection, such as BERT- or RoBERTa-based taggers.
• Cloud Provider APIs:
  • Azure AI Language: Detects dozens of PII entity types, including names, emails, and financial identifiers.
  • AWS Comprehend: Offers batch and real-time APIs for PII detection in English and Spanish.
  • Google Cloud DLP: Scans text and structured data for sensitive information using predefined detectors.

Why Fine-Tune Your Own Model?

Fine-tuned models are the right choice when generic tools fail. They solve these key problems:

• Recognize Your Unique Data: They learn to identify your company's domain-specific entities, like internal ticket IDs or custom account formats, which generic models will always miss.
• Understand Complex Context: They grasp long-range context in chats and email threads, correctly identifying PII that isolated, single-text API calls can't.
• Control Cost and Security: Self-hosting eliminates expensive API fees and ensures sensitive data never leaves your infrastructure, meeting strict data residency and privacy rules.
• Adapt Quickly: You can rapidly retrain your model to handle evolving business logic and new PII types, instead of waiting months for a third-party vendor to update their service.
• Data Sovereignty and Security: Many organizations operate under strict data governance policies that prohibit sending sensitive, un-redacted data to third-party APIs.
• Latency and Performance: For real-time applications, such as redacting text in a live chat, network latency from API calls can be a bottleneck. A locally hosted, optimized Small Language Model (SLM) - a sub-1B fine-tuned model can offer significantly lower latency, providing a better user experience.

You likely don't need fine-tuning if:

• Standard Data and Simple Context: Your task is limited to identifying common PII (like names and emails) in short, straightforward texts where conversational history isn't needed.
• Low Volume and Flexible Constraints: Your data volume is small enough that API costs are minimal (more on cost implications in this link), and you don't have strict data residency or privacy rules that prevent using third-party services.
• Minor Customization is Sufficient: Your domain-specific needs are simple enough that you can effectively guide a powerful commercial model (OpenAI, Gemini etc.) using clear, direct instructions in your prompt, without needing to retrain the model itself.

What is Supervised Fine-Tuning (SFT)?

We will make use of Supervised Fine-Tuning (SFT) to teach our model how to redact PII. SFT works by taking a powerful, pre-trained base model and further training it on a specific, high-quality dataset of curated examples.

The model is shown thousands of these input-output pairs. Each pair provides a clear example of the task, showing the model an input (like a sentence with PII) and the exact output it is expected to produce (the same sentence with the PII masked). This labeled dataset provides the "supervision," guiding the model to adjust its internal parameters to accurately replicate this new behavior until it becomes an expert at the task; in this case, PII redaction.

For more details of all the fine-tuning steps in a language model, read this article by Cameron R. Wolfe.

Dataset and Problem Formulation

The project uses the ai4privacy/pii-masking-300k dataset, which is publicly available on Hugging Face. This dataset is specifically designed for PII (Personally Identifiable Information) masking and is structured in the popular "Alpaca" format. This format organizes data into instruction-response pairs, making it ideal for teaching a model how to follow a specific command. The dataset contains data from multiple languages and for brevity, aniket-curlscape/pii-masking-english-1k is created, which is a filtered version containing only data in English language.

The core of the fine-tuning process is teaching the model to map a specific input to a desired output using Supervised Fine-Tuning (SFT) (more on this in section below). In this specific case, it is to train the model to replace specific PII with placeholders:

• Usernames → [USERNAME]
• Emails → [EMAIL]
• Phone numbers → [PHONE]
• Addresses → [ADDRESS]
• Timestamps → [TIME]
• ...

The Process: Fine-Tuning a Llama-3-8B Model with Axolotl

Axolotl is a Python framework built primarily on the Hugging Face ecosystem, used for fine-tuning Large Language Models. It is a powerful framework is known for its simplicity and flexibility. It unifies the entire training process into a single YAML configuration file, removing the need for complex boilerplate code. This allows you to easily experiment with different models, datasets, and advanced, memory-efficient techniques like LoRA and QLoRA just by changing a few lines in a file. Essentially, Axolotl makes the sophisticated process of fine-tuning more accessible, reproducible, and faster to iterate on.

We will use the toolkit available in this GitHub repository: aniket-curlscape/finetune-using-axolotl. This repository provides a comprehensive setup for fine-tuning models on a PII masking task using both local and cloud environments.

For this guide, we'll focus on the cloud-based setup using Modal, which simplifies GPU allocation and environment management.

Step 1: Set Up Your Environment

First, clone the repository and install the required Python packages.

# Clone the project repository
git clone https://github.com/aniket-curlscape/finetune-using-axolotl.git
cd finetune-using-axolotl

# Install the necessary dependencies
pip install -r requirements.txt

Step 2: Configure Modal for Cloud Training

Modal is a serverless platform that lets you run code on cloud GPUs without managing infrastructure. First, you need to set up the Modal client and link your Hugging Face account to access models.

# Configure the Modal client (this will open a browser window for authentication)
modal setup

# Add your Hugging Face token as a secret in Modal so it can download the base model
modal secret create huggingface-secret HUGGING_FACE_TOKEN=<your-hf-token>

Step 3: Benchmark the Base Model (Optional but Recommended)

Before fine-tuning, it's crucial to establish a baseline. This means testing the original Llama-3-8B model to see how it performs on our PII redaction task out-of-the-box. This helps prove that fine-tuning is actually necessary.

modal deploy ./src/modal_base_inference.py

Once deployed, you can send requests to the provided URL. You'll likely observe that the base model fails to identify and mask your specific PII formats, confirming the need for our next step. This serves as a crucial vibe check, giving you a clear performance baseline and setting expectations for what improvement the fine-tuning needs to deliver. It also acts as a valuable control test; if you encounter issues later, you can determine whether an error is due to the fine-tuning process or an inherent limitation of the base model itself.

Step 4: Configure and Launch the Fine-Tuning Job

Axolotl uses a YAML file to configure every aspect of the training process, from the base model to hyperparameters. The repository comes with pre-written configurations. Let's look at a snippet from configs/config-llama-8b-100.yaml:

base_model: NousResearch/Meta-Llama-3-8B
model_type: LlamaForCausalLM
tokenizer_type: AutoTokenizer

load_in_8bit: true
load_in_4bit: false

chat_template: llama3
datasets:
  - path: aniket-curlscape/pii-masking-english-100
    type:
      system_prompt: "You are a helpful assistant that masks all personally identifiable information (PII) in text. Replace each detected PII entity with the correct placeholder from the list below:\nUsernames as [USERNAME]\nGiven names as [GIVENNAME1], [GIVENNAME2]\nLast names as [LASTNAME1], [LASTNAME2]\nTitles (e.g., Mr., Dr., Archduchess) as [TITLE]\nEmails as [EMAIL]\nPhone numbers / Telephones as [TEL]\nAddresses as [BUILDING], [STREET], [CITY], [STATE], [COUNTRY], [POSTCODE], [SECADDRESS]\nDates as [DATE], birthdates as [BOD], times as [TIME]\nIdentity numbers as [SOCIALNUMBER], [PASSPORT], [DRIVERLICENSE], [IDCARD]\nPasswords / secrets as [PASS]\nIP addresses as [IP]\nSex / Gender as [SEX]\nRules:\nAlways use the most specific placeholder available.\nSupport multiple occurrences of the same type (append numbering if needed).\nPreserve all non-PII text exactly.\nApply consistently across structured (CSV, JSON, XML) and unstructured (free text, comments) formats.\nWhen in doubt (e.g., certificate numbers, encoded IDs), map to the closest ID placeholder: [IDCARD] or [DRIVERLICENSE]."
      field_system: system
      field_instruction: source_text
      field_output: target_text
      format: "[INST] {instruction} [/INST]"
      no_input_format: "[INST] {instruction} [/INST]"
special_tokens:
   pad_token: <|end_of_text|>
dataset_prepared_path:
val_set_size: 0.05
output_dir: ./outputs/pii

sequence_len: 4096
sample_packing: false

adapter: lora
lora_model_dir:
lora_r: 16
lora_alpha: 32
lora_dropout: 0.05
lora_target_linear: true

gradient_accumulation_steps: 4
micro_batch_size: 2
num_epochs: 1
optimizer: adamw_bnb_8bit
lr_scheduler: cosine
learning_rate: 0.0002

bf16: auto
tf32: false

gradient_checkpointing: true
resume_from_checkpoint:
logging_steps: 1
flash_attention: true

warmup_ratio: 0.1
evals_per_epoch: 5
saves_per_epoch: 3
weight_decay: 0.0

use_tensorboard: true

This configuration specifies that we are fine-tuning Meta-Llama-3-8B using the efficient QLoRA method on the pii-masking-english-1k dataset.

With the configuration ready, you can launch the fine-tuning job on Modal with a single command:

modal run src/modal_train.py --config configs/config-llama-8b-100.yaml

Modal will spin up a container with the specified GPU, download the dataset and base model, run the Axolotl training process, and save the resulting fine-tuned model adapters.

The src.modal_train.py script automates the entire fine-tuning job. Here’s a step-by-step breakdown of how it works:

• Environment and Storage Setup: First, the script defines the cloud environment in the axolotl_image object, packaging all necessary software. It also prepares persistent storage volumes (VOLUME_CONFIG) to save base models and training results so they are not lost between runs.
• Kicking Off the Job: The process starts locally in the main function, which reads your config file and calls the launch function in the cloud. launch then creates a unique run folder, downloads the base model if needed, and kicks off the main training job.
• The Fine-Tuning Process: Executing on a GPU, the run_axolotl function runs Axolotl's core training command (axolotl.cli.train). After successfully creating the LoRA adapters, it automatically triggers the final merging step.
• Merging the Model: In the final merge step, a function combines the lightweight LoRA adapters with the original base model using the axolotl.cli.merge_lora command. This produces a single, standalone fine-tuned model that is saved permanently.

Step 5: Test Your Fine-Tuned Model

Once the training is complete, you can immediately run inference to test its performance:

modal run --quiet -m src.modal_infer --prompt "I would like to meet John at 9am."

The model, now aware of PII context, should correctly identify "9am" as a timestamp, "John" as given name and mask them accordingly.

The Importance of Observability and Evaluation

Fine-tuning isn't a "set it and forget it" process. To build a robust and trustworthy AI system, you must stay in the loop. This is where observability and evaluation come in.

• Observability: Tools like TensorBoard and Weights & Biases are crucial for monitoring training runs. They help you track metrics like loss, learning rate, and GPU utilization in real-time, allowing you to catch issues early and understand your model's behavior. Axolotl can be configured to use any of these tools.
• Evaluation: How do you know if your model is actually working well? Simply observing its output isn't enough. A systematic evaluation process is needed like calculating Precision, Recall, and F1-score for each PII entity type to get a quantitative measure of performance. This can range from simple string matching against a test dataset to more advanced techniques like using another powerful LLM as a judge to score the quality of the redactions. Consistent evaluation ensures your model's performance doesn't degrade over time and that it meets the required accuracy standards for your use case.

From Training to Production

A successful fine-tuning run is just the beginning. To deploy this model in a production environment, consider these next steps:

• Merge the Adapter Weights and Push: For inference, loading the base model and the LoRA adapter separately adds overhead. You should merge the adapter weights into the base model to create a single, unified model and push to a central model repository such as Huggingface.
• Model Quantization: To further reduce the model's memory footprint and speed up inference, you can quantize the merged model to a format like GGUF (for CPU inference) or use techniques like AWQ/GPTQ for GPU deployment.
• High-Performance Serving: Don't run inference using a basic Hugging Face pipeline in production. Use a dedicated serving engine like vLLM, Text Generation Inference (TGI), or SGLang. These tools are optimized for high-throughput, low-latency LLM serving and include features like continuous batching. The GitHub repository already includes vLLM integration for testing the model without fine-tuning here.

Fine-tuning a large language model for a specialized task like PII redaction is more accessible than ever. As we've seen, while off-the-shelf APIs have their place, they often struggle with unique, domain-specific data. By leveraging a powerful framework like Axolotl and an automation platform like Modal, you can transform a generalist base model into a specialist that understands your specific context—leading to higher accuracy, better security, and more predictable costs.

The power to create custom, production-grade AI solutions is no longer confined to a few large labs. The complete code for this guide is available on GitHub, and the best way to learn is by building. We encourage you to clone the repository, run the training yourself, and start experimenting with your own datasets.

Happy fine-tuning! 🚀