Walkthrough: Question Answering - Privacy Aware Fine Tuning

In this tutorial, we will showcase using DynamoFL's ML SDK to:

• Fine-Tune LLMs for Q/A with differential privacy.

Prerequisites:

• DynamoFL API token
• DynamoFL ML SDK Runtime License Key

Fine-tuning (with and without DP)

Environment Setup

Let's import DynamoFL's SDKs and other libraries needed for fine-tuning.

Install DynamoFL's ML SDK packages

from dynamofl.privacy import (
    get_dp_causal_lm_trainer,
    PrivacyArguments,
)

Install other necessary packages

import transformers
import os
import time
from transformers import (
    set_seed,
    Trainer,
    TrainingArguments,
    default_data_collator
)
from huggingface_hub import HfApi
from peft import LoraConfig, TaskType, get_peft_model
import datasets
import yaml

Load Fine-Tuning Helper Functions

The DynamoFL team has prepared helper files to simplify the process of fine-tuning with and without Differential Privacy and LoRA.

These include:

• hf_utils.ipynb include helper functions for DP fine-tuning
• configs/ directory include yaml files for fine-tuning with and without DP

'''
    DynamoFL tutorial helper functions
'''
from hf_utils import (
    DataTrainingArguments,
    setup_logging,
    model_loading,
    ModelArguments,
    prepare_data_for_training,
    DynamoFLCallback,
    EvaluateFirstStepCallback
)

Training

We wrap the training process in a function to make it easier and more readable. The arguments train_args, privacy_args, model_args, data_args are parsed from the config file.

def train(train_args, privacy_args, model_args, data_args, resume_from_checkpoint=False):
    setup_logging(train_args)
    # Set seed before initializing model.
    set_seed(train_args.seed)


    # Load in model and tokenizer
    model, tokenizer = model_loading(model_args)
    tokenizer.pad_token = tokenizer.eos_token
    model = model.cuda()
    model.train()

    # Load in dataset and parse it for the correct format
    train_dataset, test_dataset, raw_train_dataset = prepare_data_for_training(data_args, model_args, train_args, model, tokenizer)

    # LoRA configs
    if USE_LORA:
        lora_config = LoraConfig(
            task_type=TaskType.CAUSAL_LM,
            inference_mode=False,
            r=train_config_yaml['lora_rank'],
            lora_alpha=train_config_yaml['lora_alpha'],
            lora_dropout=train_config_yaml['lora_dropout'],
            target_modules=train_config_yaml['lora_target_modules'],
        )

        model = get_peft_model(model, lora_config)
        model.print_trainable_parameters()

    # [DynamoFL] Easily set your trainer for differentially private training with one line of code

    if privacy_args.disable_dp:
        # Normal training, no DP (for comparison)
        trainer = Trainer(
            model=model,
            tokenizer=tokenizer,
            args=train_args,
            train_dataset=train_dataset,
            eval_dataset=test_dataset,
            data_collator=default_data_collator,
        )
    else:
        trainer = get_dp_causal_lm_trainer(
            model=model,
            tokenizer=tokenizer,
            train_args=train_args,
            privacy_args=privacy_args,
            train_dataset=train_dataset,
            eval_dataset=test_dataset,
        )

    trainer.add_callback(EvaluateFirstStepCallback())

    # train
    trainer.train(resume_from_checkpoint=resume_from_checkpoint)

    if not privacy_args.disable_dp:
        '''
            Log differential privacy metrics
        '''
        eps_prv = trainer.get_prv_epsilon()
        eps_rdp = trainer.get_rdp_epsilon()
        trainer.log({
            "final_epsilon_prv": eps_prv, # estimated upper bound of true epsilon
            "final_epsilon_rdp": eps_rdp, # privacy budget (epsilon) expended
        })

    return trainer

# Path to your yaml file containing training configs
train_config = "configs/train_config_dp.yml"

# fine-tune from pre-trained model
resume_from_checkpoint = False

arg_parser = transformers.HfArgumentParser((
    TrainingArguments,
    PrivacyArguments,
    ModelArguments,
    DataTrainingArguments
))

train_args, privacy_args, model_args, data_args = arg_parser.parse_yaml_file(
    yaml_file=os.path.abspath(train_config),
    allow_extra_keys=True
)

trainer = train(
    train_args=train_args,
    privacy_args=privacy_args,
    model_args=model_args,
    data_args=data_args,
    resume_from_checkpoint=resume_from_checkpoint
)

---

Copyright 2023 DynamoFL, Inc. All rights reserved.

All use of the software is subject to the license terms and conditions in the Terms of Use between you or your employer and DynamoFL. Without limiting the terms and conditions of the Terms of Use, you may not:

• Use the software in any manner that isn't directly tied to your use of the DynamoFL platform, including any use that attempts to remove or circumvent the dependency of the software on the DynamoFL platform;
• Reverse engineer or otherwise attempt to discover underlying structure, ideas, or algorithms of the software;
• Copy, distribute, publish, or otherwise make available the software.

You may make minor edits to the software solely as necessary to train your models on the DynamoFL platform.