workshop/notebooks/LoraMLP.ipynb

Low-Rank Approximation Fine-Tuning

This notebook demonstrates encrypted fine-tuning of a small MLP model with LoRA. A model trained on an initial dataset is adapted to a second dataset using LoRA fine-tuning.

The fine-tuning dataset and the trained LoRA weights are protected using encryption. Thus, training can be securely outsourced to a remote server without compromising any sensitive data.

The hybrid approach is applied to fine-tuning: only the linear layers of the original model are outsourced to the server. The forward and backward passes on these layers are performed using encrypted activations and gradients. Meanwhile, the LoRA weights are kept by the client, which performs locally the forward and backward passes on the LoRA weights.

In [1]:

import shutil
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import torch
from peft import LoraConfig, get_peft_model
from sklearn.datasets import make_circles, make_moons
from torch import nn, optim
from torch.utils.data import DataLoader, TensorDataset

from torus.ml.torch.lora import LoraTrainer

# Set random seed for reproducibility
SEED = 42
np.random.seed(SEED)
torch.manual_seed(SEED);

Data preparation

Two datasets are generated: one for the original training, and a second one on which LoRA fine-tuning is performed.

In [2]:

# Task 1: Two interleaving half circles
X_task1, y_task1 = make_moons(n_samples=500, noise=0.1)
# Task 2: Two concentric circles
X_task2, y_task2 = make_circles(n_samples=500, noise=0.2, factor=0.5)


def plot_datasets_and_boundaries(X_task1, y_task1, X_task2, y_task2, model=None, titles=None):
    _, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6))

    if titles is None:
        titles = ["Task 1 Dataset", "Task 2 Dataset"]

    for ax, X, y, title in zip([ax1, ax2], [X_task1, X_task2], [y_task1, y_task2], titles):
        ax.scatter(X[:, 0], X[:, 1], c=y, cmap="viridis", edgecolor="k")
        ax.set_title(title)

        if model is not None:
            x_min, x_max = X[:, 0].min() - 0.5, X[:, 0].max() + 0.5
            y_min, y_max = X[:, 1].min() - 0.5, X[:, 1].max() + 0.5
            h = 0.1  # step size in the mesh
            xx, yy = np.meshgrid(np.arange(x_min, x_max, h), np.arange(y_min, y_max, h))
            grid = torch.FloatTensor(np.c_[xx.ravel(), yy.ravel()])

            with torch.no_grad():
                Z = model(grid)
                probabilities = torch.softmax(Z, dim=1)
                Z = probabilities[:, 1].numpy().reshape(xx.shape)

            ax.contourf(xx, yy, Z, cmap="viridis", alpha=0.3)

    plt.tight_layout()
    plt.show()


# Plot datasets
plot_datasets_and_boundaries(X_task1, y_task1, X_task2, y_task2)

Create the original MLP and test it

In [3]:

# Convert datasets to PyTorch tensors
X_task1 = torch.FloatTensor(X_task1)
y_task1 = torch.LongTensor(y_task1)
X_task2 = torch.FloatTensor(X_task2)
y_task2 = torch.LongTensor(y_task2)

# Create DataLoaders
batch_size = 32
train_loader_task1 = DataLoader(
    TensorDataset(X_task1, y_task1), batch_size=batch_size, shuffle=True
)
train_loader_task2 = DataLoader(
    TensorDataset(X_task2, y_task2), batch_size=batch_size, shuffle=True
)

# Define an MLP model without LoRA layers


class SimpleMLP(nn.Module):
    """Simple MLP model without LoRA layers."""

    def __init__(self, input_size=2, hidden_size=128, num_classes=2):
        super().__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, num_classes)

    def forward(self, x, labels=None):  # pylint: disable=unused-argument
        """Forward pass of the MLP."""
        out = self.fc1(x)
        out = self.relu(out)
        out = self.fc2(out)
        return out


# Instantiate the model
model = SimpleMLP()

# Training loop for Task 1


def train_model(model, train_loader, num_epochs=100):
    """Train the model.

    Args:
        model (nn.Module): The model to train.
        train_loader (DataLoader): DataLoader for training data.
        num_epochs (int): Number of epochs to train.
    """
    device = torch.device("cpu")
    model.to(device)
    model.train()

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.01)

    for epoch in range(num_epochs):
        total_loss = 0
        for x_batch, y_batch in train_loader:
            x_batch = x_batch.to(device)
            y_batch = y_batch.to(device)

            optimizer.zero_grad()
            outputs = model(x_batch)
            loss = criterion(outputs, y_batch)
            loss.backward()
            optimizer.step()

            total_loss += loss.item()

        # Print loss every 20 epochs
        if (epoch + 1) % 20 == 0:
            print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {total_loss/len(train_loader):.4f}")


# Train the model on Task 1
print("Training on Task 1 without LoRA:")
train_model(model, train_loader_task1, num_epochs=20)

# Plot datasets with decision boundaries
plot_datasets_and_boundaries(
    X_task1.numpy(),
    y_task1.numpy(),
    X_task2.numpy(),
    y_task2.numpy(),
    model=model,
    titles=["Task 1 after Training", "Task 2 after Training"],
)

Training on Task 1 without LoRA:

Epoch [20/20], Loss: 0.0036

Apply LoRA to the model using peft

In [4]:

# Apply LoRA to the model using peft
lora_config = LoraConfig(
    r=1, lora_alpha=1, lora_dropout=0.01, target_modules=["fc1", "fc2"], bias="none"
)

peft_model = get_peft_model(model, lora_config)

Setup FHE fine-tuning with LoraTrainer

In [5]:

# Update training parameters, including loss function
optimizer = optim.Adam(filter(lambda p: p.requires_grad, peft_model.parameters()), lr=0.01)
loss_fn = nn.CrossEntropyLoss()
training_args = {"gradient_accumulation_steps": 1}

# Set up LoRA training
lora_trainer = LoraTrainer(
    peft_model, optimizer=optimizer, loss_fn=loss_fn, training_args=training_args
)

# Prepare input data for calibration
batch_size_per_task = batch_size // 2
inputset = (
    torch.cat([X_task1[:batch_size_per_task], X_task2[:batch_size_per_task]]),
    torch.cat([y_task1[:batch_size_per_task], y_task2[:batch_size_per_task]]),
)

# Compile the model
lora_trainer.compile(inputset, n_bits=8)

LoRA layers detected in the model.

Compiling FHE layers:   0%|          | 0/3 [00:00<?, ?it/s]

In [6]:

# Fine-tune the model on Task 2 using LoRA
lora_trainer.train(train_loader_task2, num_epochs=10, fhe="execute")

Training:   0%|          | 0/10 [00:00<?, ?epoch/s]

Training completed. Final Avg Loss: 0.2978, FHE Mode: execute

Visualize fine-tuned model

In [7]:

# Enable LoRA adapters (already enabled by default)
peft_model.enable_adapter_layers()

# Plot datasets with decision boundaries after fine-tuning
plot_datasets_and_boundaries(
    X_task1.numpy(),
    y_task1.numpy(),
    X_task2.numpy(),
    y_task2.numpy(),
    model=peft_model,
    titles=["Task 1 after Fine-tuning", "Task 2 after Fine-tuning"],
)

Disable LoRA adapters to see the original model

In [8]:

# Disable LoRA adapters
peft_model.disable_adapter_layers()

# Plot datasets with decision boundaries after fine-tuning
plot_datasets_and_boundaries(
    X_task1.numpy(),
    y_task1.numpy(),
    X_task2.numpy(),
    y_task2.numpy(),
    model=peft_model,
    titles=["Task 1 after Fine-tuning", "Task 2 after Fine-tuning"],
)

Print statistics and save the model

In [9]:

# Enable LoRA adapters (already enabled by default)
peft_model.enable_adapter_layers()

# Print trainable (lora) parameters
peft_model.print_trainable_parameters()

# Save the model and remove all layers that will be done on the server
path = Path("lora_mlp")

if path.is_dir() and any(path.iterdir()):
    shutil.rmtree(path)

lora_trainer.save_and_clear_private_info(path)

# At this point, the client's model only contains the trainable parameters of the LoRA layers.
peft_model.print_trainable_parameters()

trainable params: 260 || all params: 902 || trainable%: 28.8248
trainable params: 260 || all params: 260 || trainable%: 100.0000

Conclusion

This notebook demonstrates how Low-Rank Adaptation (LoRA) facilitates parameter-efficient fine-tuning of models on private data. By leveraging FHE, the training process ensures that sensitive data and private model updates remains secure.

Key Takeaways:

• Efficiency with LoRA: While this example utilizes an MLP model with a relatively high proportion of LoRA weights due to its simplicity, the approach scales effectively to larger models like large language models (LLMs). In such cases, LoRA typically accounts for less than one percent of the total model parameters, ensuring minimal memory and computational overhead on the client side.
• Scalability and Practicality: The hybrid approach demonstrated here is particularly beneficial for scenarios where client devices have limited resources. Memory heavy computations are offloaded to a secure server and the client handles only the lightweight LoRA adjustments locally.