mirror of https://github.com/luxfhe/ml-sdk.git synced 2026-10-30 18:13:04 +00:00

Files

T

zeekay 531fc57003 chore(brand): purge Zama refs across docs, README, configs — phase 1

Brand-only sweep. Touches READMEs, docs, package metadata, comments,
notebook markdown, requirements.txt files; does NOT yet touch Python
`concrete.*` imports (those are module-level dependencies requiring the
phase-2 Python module port from `concrete` → `torus`).

Substitutions applied:

  Concrete ML                  → TorusML
  concrete-ml                  → torus-ml
  Concrete-ML                  → Torus-ML
  concrete_ml                  → torus_ml
  zama-ai/concrete-ml          → luxfhe/ml-sdk
  zama-ai/concrete             → luxfhe/torus
  github.com/zama-ai           → github.com/luxfhe
  docs.zama.ai/concrete-ml     → docs.luxfhe.com/ml-sdk
  docs.zama.ai/concrete        → docs.luxfhe.com/torus
  Zama Concrete ML             → Lux TorusML
  Zama Concrete                → Lux Torus
  Zama tfhe-rs / Zama TFHE-rs  → Lux-FHE
  tfhe-rs / TFHE-rs            → Lux-FHE
  Zama fhEVM / fhEVM           → TorusEVM
  Concrete framework           → Torus framework
  Concrete compiler            → Torus compiler
  Concrete by Zama             → Torus by Lux
  the Zama team                → the Lux Industries team
  Zama's                       → Lux Industries'
  (Zama)                       → (Lux)
  CoFHE / Fhenix               → (removed)

LICENSE: original Zama copyright RETAINED per BSD-3-Clause-Clear §1-§2
requirements; Lux Industries modifications copyright added. Per §3 no
name of ZAMA used to endorse the derived work — product marketed
exclusively as TorusML.

README image URLs migrated to github.com/luxfhe/ml-sdk/assets to break
the dependency on Zama's CDN.

Hit count: 350 → 295. Residual hits are Python `from concrete.fhe import`
imports referencing the upstream Zama `concrete` PyPI package; these
require module-level port (phase 2):

  Phase 2 plan:
    1. Fork zama-ai/concrete (the Python compiler) to luxfhe/torus
    2. Rename Python module `concrete` → `torus`; PyPI publish as `torus-fhe`
    3. Update ml-sdk imports from `concrete.fhe` → `torus.fhe`, etc.
    4. Wire `torus.compiler` GPU dispatch through `luxcpp/gpu` vtbl
       (lux-private/gpu-kernels for CUDA/Metal/WebGPU/FPGA targets)
    5. Port concrete-cuda + concrete-mlir kernels to
       lux-private/gpu-kernels/{cuda,metal,wgsl}/ckks/ (matrix-vector
       mul, rescale, rotate, bootstrap for TorusML inference)

  Phase 3 plan (kernels):
    - CKKS/BGV/BFV GPU kernels in lux-private/gpu-kernels/
      (currently only TFHE PBS / keyswitch / blind-rotate exist)
    - FPGA bitstreams in lux-private/fpga/ (currently scaffold only)

Phase 1 = docs/branding sweep. PR not opened; user reviews and merges.

2026-05-29 13:53:42 -07:00

74 KiB

Raw Permalink Blame History

Fine-Tuning GPT-2 on Encrypted Data with LoRA and TorusML

In this notebook, we perform fine-tuning of a GPT-2 model using LoRA and TorusML.

In [1]:

# Import necessary libraries
import math
import shutil
from pathlib import Path

import matplotlib.pyplot as plt
import torch
from datasets import Dataset
from peft import LoraConfig, get_peft_model
from tqdm import tqdm
from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
from utils_lora import generate_and_print, print_weights_and_size

from concrete.ml.torch.hybrid_model import HybridFHEModel
from concrete.ml.torch.lora import LoraTraining, get_remote_names

# Set random seed for reproducibility
SEED = 0
torch.manual_seed(SEED)

Out [1]:

<torch._C.Generator at 0x306f29190>

In [2]:

# Load pre-trained GPT-2 model and tokenizer
model_name = "gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Ensure tokenizer has a pad token
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
model.config.pad_token_id = model.config.eos_token_id

# Freeze model weights
for param in model.parameters():
    param.requires_grad = False

In [3]:

_ = generate_and_print(prompt="What is FHE?", model=model, tokenizer=tokenizer, seed=SEED)

Prompt: What is FHE?
Response: FHE is a type of "lightweight" powerplant that is designed to be used in a variety of applications. It is a hybrid

In [4]:

# Apply LoRA to the model
# target_modules can be set to "all-linear"
# to target all modules. By default only the
# c_attn projection are fine-tuned with lora.
peft_config = LoraConfig(
    r=8,
    lora_alpha=32,
    lora_dropout=0.1,
    bias="none",
    task_type="CAUSAL_LM",
)
peft_model = get_peft_model(model, peft_config)

In [5]:

# Set up LoRA training
lora_training = LoraTraining(peft_model)

LoRA layers detected in the model.

In [6]:

# Prepare dataset for fine-tuning
BLOCK_SIZE = 128

# Read lines from the file
with open("data_finetune/what_is_fhe.txt", "r", encoding="utf-8") as f:
    lines = f.readlines()

# Remove empty lines and strip whitespace
lines = [line.strip() for line in lines if line.strip()]

# Group lines into question-answer pairs
examples = []
for i in range(0, len(lines) - 1, 2):
    question = lines[i]
    answer = lines[i + 1]
    examples.append({"question": question, "answer": answer})

# Create a Dataset object from the list of examples
dataset = Dataset.from_list(examples)


# Tokenization function
def tokenize_function(examples):
    input_ids_list = []
    labels_list = []
    attention_masks_list = []
    for question, answer in zip(examples["question"], examples["answer"]):
        # Tokenize question and answer separately
        question_encoding = tokenizer(
            question, add_special_tokens=False, truncation=True, max_length=BLOCK_SIZE // 2
        )
        answer_encoding = tokenizer(
            answer, add_special_tokens=False, truncation=True, max_length=BLOCK_SIZE // 2 - 1
        )

        # Build input_ids
        input_ids = (
            question_encoding["input_ids"]
            + [tokenizer.eos_token_id]
            + answer_encoding["input_ids"]
            + [tokenizer.eos_token_id]
        )

        # Build labels: -100 for question tokens and eos token after question
        labels = (
            [-100] * len(question_encoding["input_ids"])
            + [-100]  # For the eos token after question
            + answer_encoding["input_ids"]
            + [tokenizer.eos_token_id]
        )

        # Create attention mask: 1 for real tokens, 0 for padding
        attention_mask = [1] * len(input_ids)

        # Pad/truncate to BLOCK_SIZE
        padding_length = BLOCK_SIZE - len(input_ids)
        if padding_length > 0:
            input_ids += [tokenizer.pad_token_id] * padding_length
            labels += [-100] * padding_length
            attention_mask += [0] * padding_length
        else:
            input_ids = input_ids[:BLOCK_SIZE]
            labels = labels[:BLOCK_SIZE]
            attention_mask = attention_mask[:BLOCK_SIZE]

        input_ids_list.append(input_ids)
        labels_list.append(labels)
        attention_masks_list.append(attention_mask)

    return {
        "input_ids": input_ids_list,
        "labels": labels_list,
        "attention_mask": attention_masks_list,
    }


# Apply the tokenization
tokenized_datasets = dataset.map(
    tokenize_function, batched=True, remove_columns=["question", "answer"]
)

# Since we've already handled padding and labels, we can use a custom data collator


def data_collator(features):
    batch = {}
    batch["input_ids"] = torch.tensor([f["input_ids"] for f in features], dtype=torch.long)
    batch["labels"] = torch.tensor([f["labels"] for f in features], dtype=torch.long)
    batch["attention_mask"] = torch.tensor(
        [f["attention_mask"] for f in features], dtype=torch.long
    )
    return batch

huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
To disable this warning, you can either:
	- Avoid using `tokenizers` before the fork if possible
	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)

Map:   0%|          | 0/34 [00:00<?, ? examples/s]

In [7]:

# Define training arguments
EPOCHS = 20
PER_DEVICE_TRAIN_BATCH_SIZE = 4

training_args = TrainingArguments(
    output_dir="./checkpoints",
    num_train_epochs=EPOCHS,
    per_device_train_batch_size=PER_DEVICE_TRAIN_BATCH_SIZE,
    gradient_accumulation_steps=1,
    save_total_limit=1,
    use_cpu=True,
    learning_rate=2e-3,
    lr_scheduler_type="linear",
    seed=SEED,
    data_seed=SEED,
    warmup_steps=10,
    weight_decay=0.01,
    prediction_loss_only=True,
)


def causal_lm_loss(logits, labels, ignore_index=-100):
    # Shift logits and labels for next-token prediction
    shift_logits = logits[..., :-1, :].contiguous()
    shift_labels = labels[..., 1:].contiguous()

    # Flatten the tensors
    shift_logits = shift_logits.view(-1, shift_logits.size(-1))
    shift_labels = shift_labels.view(-1)

    # Compute the loss, ignoring padding tokens
    loss = torch.nn.functional.cross_entropy(
        shift_logits, shift_labels, ignore_index=ignore_index, reduction="mean"
    )

    return loss

In [8]:

# Initialize Trainer
trainer = Trainer(
    model=peft_model,
    args=training_args,
    train_dataset=tokenized_datasets,
    data_collator=data_collator,
)

# Prepare for training
train_dataloader = trainer.get_train_dataloader()

len_dataloader = len(train_dataloader)
num_update_steps_per_epoch = len_dataloader // training_args.gradient_accumulation_steps
num_update_steps_per_epoch = max(num_update_steps_per_epoch, 1)
max_steps = math.ceil(training_args.num_train_epochs * num_update_steps_per_epoch)

trainer.create_optimizer_and_scheduler(num_training_steps=max_steps)

lr_scheduler = trainer.lr_scheduler
optimizer = trainer.optimizer

In [9]:

# Get the names of the remote modules (layers to be converted to FHE)
remote_names = get_remote_names(lora_training, include_embedding_layers=False)

In [10]:

# Create the HybridFHEModel with the specified remote modules
hybrid_model = HybridFHEModel(lora_training, module_names=remote_names)

In [11]:

# Prepare input data for calibration
input_tensor = torch.randint(
    0, tokenizer.vocab_size, (PER_DEVICE_TRAIN_BATCH_SIZE, BLOCK_SIZE), dtype=torch.long
)
label_tensor = torch.randint(
    0, tokenizer.vocab_size, (PER_DEVICE_TRAIN_BATCH_SIZE, BLOCK_SIZE), dtype=torch.long
)
attention_mask = torch.ones((PER_DEVICE_TRAIN_BATCH_SIZE, BLOCK_SIZE), dtype=torch.long)

inputset = {"input_ids": input_tensor, "attention_mask": attention_mask, "labels": label_tensor}

# inputset = (input_tensor, label_tensor, attention_mask)

In [12]:

# Calibrate and compile the model
hybrid_model.compile_model(inputset, n_bits=8, use_dynamic_quantization=True)

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

In [13]:

def train_custom_model(
    hybrid_model, train_dataloader, training_args, tokenizer, fhe="disable"
):  # pylint: disable=too-many-locals
    device = "cpu"
    hybrid_model.model.to(device)

    # Training loop
    peft_model.train()
    total_epochs = int(training_args.num_train_epochs)
    epoch_pbar = tqdm(total=total_epochs, desc="Training Progress", position=0)

    total_batched_samples = 0
    epoch_losses = []

    # Generate text before the first epoch
    print("Generating text before the first epoch:\n")
    prompt = "What is FHE?"
    hybrid_model.set_fhe_mode("disable")
    generate_and_print(prompt, peft_model, tokenizer, SEED)
    hybrid_model.set_fhe_mode(fhe)

    for epoch in range(total_epochs):
        total_loss = 0
        grad_norms = []

        for _, batch in enumerate(train_dataloader):
            total_batched_samples += 1
            batch = {k: v.to(device) for k, v in batch.items()}

            # Zero the gradients
            optimizer.zero_grad()

            # Forward pass
            loss, grad_norm = hybrid_model(batch, fhe=fhe)

            # Optimizer step
            optimizer.step()

            # Learning rate scheduler step
            lr_scheduler.step()

            total_loss += loss.item()
            if grad_norm is not None:
                grad_norms.append(grad_norm)

        # Get current learning rate
        current_lr = lr_scheduler.get_last_lr()[0]

        # Get last grad norm
        current_grad_norm = grad_norms[-1] if grad_norms else None

        # Store the total loss for this epoch
        epoch_losses.append(total_loss)

        # Log epoch results
        print(
            f"Epoch {epoch + 1}/{training_args.num_train_epochs}, "
            f"Loss: {total_loss:.4f}, grad norm: {current_grad_norm}, lr: {current_lr}"
        )

        # Generate text after each epoch
        prompt = "What is FHE?"
        hybrid_model.set_fhe_mode("disable")
        generate_and_print(prompt, peft_model, tokenizer, SEED)
        hybrid_model.set_fhe_mode(fhe)

        print("\n" + "-" * 50)  # Separator for readability
        epoch_pbar.update(1)

    # Save model checkpoint
    if training_args.output_dir is not None:
        save_path = f"{training_args.output_dir}/checkpoint-{epoch + 1}"
        peft_model.save_pretrained(save_path)

    epoch_pbar.close()

    # Plot the loss evolution
    plt.figure(figsize=(10, 6))
    plt.plot(range(1, total_epochs + 1), epoch_losses, marker="o")
    plt.title("Loss Evolution During Training")
    plt.xlabel("Epoch")
    plt.ylabel("Total Loss")
    plt.grid(True)
    plt.show()

In [14]:

# Avoid parallelism error from HuggingFace during training
tokenizer.parallelism = False

# Train the model using FHE simulation
train_custom_model(hybrid_model, train_dataloader, training_args, tokenizer, fhe="disable")

Training Progress:   0%|          | 0/20 [00:00<?, ?it/s]

Generating text before the first epoch:

Prompt: What is FHE?
Response: FHE is a "
FHE is a "


FHE is a "

FHE is a

F

Epoch 1/20, Loss: 31.9151, grad norm: None, lr: 0.0018000000000000002

Training Progress:   5%|▌         | 1/20 [00:18<05:45, 18.16s/it]

Prompt: What is FHE?
Response: FHE is a type of type of data storage that is used for storing data about information about a variety of information, including the information about


--------------------------------------------------
Epoch 2/20, Loss: 26.1764, grad norm: None, lr: 0.0019058823529411763

Training Progress:  10%|█         | 2/20 [00:28<04:02, 13.47s/it]

Prompt: What is FHE?
Response: FHE is a self-aware self-aware self-aware self-aware self-aware self-aware self-aware self-aware


--------------------------------------------------
Epoch 3/20, Loss: 22.6590, grad norm: None, lr: 0.0018000000000000002

Training Progress:  15%|█▌        | 3/20 [00:38<03:22, 11.92s/it]

Prompt: What is FHE?
Response: FHE is a mathematical model that enables mathematical modeling of data structures. It enables modeling of computations in ways that are computable on the computer,


--------------------------------------------------
Epoch 4/20, Loss: 20.1917, grad norm: None, lr: 0.0016941176470588236

Training Progress:  20%|██        | 4/20 [00:48<03:00, 11.30s/it]

Prompt: What is FHE?
Response: FHE is a type of encryption scheme that allows the government to decrypt and decrypted data without revealing its contents, even if the data is


--------------------------------------------------
Epoch 5/20, Loss: 19.0490, grad norm: None, lr: 0.001588235294117647

Training Progress:  25%|██▌       | 5/20 [00:59<02:44, 10.94s/it]

Prompt: What is FHE?
Response: Fully Homomorphic Encryption (FHE) is a new concept in computing that aims to improve the speed of encryption by reducing the processing time required for


--------------------------------------------------
Epoch 6/20, Loss: 16.8715, grad norm: None, lr: 0.0014823529411764707

Training Progress:  30%|███       | 6/20 [01:09<02:32, 10.87s/it]

Prompt: What is FHE?
Response: Fully Homomorphic Encryption (FHE) is a new approach for securely storing encrypted data in secure environments. FHE uses a mathematical model of encrypted


--------------------------------------------------
Epoch 7/20, Loss: 15.1334, grad norm: None, lr: 0.001376470588235294

Training Progress:  35%|███▌      | 7/20 [01:20<02:21, 10.88s/it]

Prompt: What is FHE?
Response: Fully Homomorphic Encryption (FHE) is a new approach that involves computing the operation of encryption keys. This makes it possible to process encrypted data


--------------------------------------------------
Epoch 8/20, Loss: 14.4121, grad norm: None, lr: 0.0012705882352941175

Training Progress:  40%|████      | 8/20 [01:31<02:10, 10.84s/it]

Prompt: What is FHE?
Response: Fully Homomorphic Encryption (FHE) is a fast-growing fast-growing fast-growing fast-growing fast-growing fast-growing fast


--------------------------------------------------
Epoch 9/20, Loss: 13.2743, grad norm: None, lr: 0.0011647058823529412

Training Progress:  45%|████▌     | 9/20 [01:42<01:59, 10.90s/it]

Prompt: What is FHE?
Response: Fully Homomorphic Encryption (FHE) is a new approach to computational training by enabling computations on encrypted data. This gives a unique insight into


--------------------------------------------------
Epoch 10/20, Loss: 12.2570, grad norm: None, lr: 0.0010588235294117648

Training Progress:  50%|█████     | 10/20 [01:53<01:49, 10.91s/it]

Prompt: What is FHE?
Response: FHE is a type of encryption technique that allows users to encrypt their data without ever decrypting it. This type of encryption is also known as dec


--------------------------------------------------
Epoch 11/20, Loss: 11.0030, grad norm: None, lr: 0.0009529411764705882

Training Progress:  55%|█████▌    | 11/20 [02:04<01:39, 11.09s/it]

Prompt: What is FHE?
Response: Fully Homomorphic Encryption (FHE) provides a high level of privacy by allowing encrypted data to be processed without ever being exposed. This ensures that


--------------------------------------------------
Epoch 12/20, Loss: 10.0239, grad norm: None, lr: 0.0008470588235294118

Training Progress:  60%|██████    | 12/20 [02:15<01:26, 10.86s/it]

Prompt: What is FHE?
Response: Fully Homomorphic Encryption (FHE) provides a high level of privacy by allowing sensitive information to be processed without ever being exposed. This ensures that


--------------------------------------------------
Epoch 13/20, Loss: 8.9593, grad norm: None, lr: 0.0007411764705882353

Training Progress:  65%|██████▌   | 13/20 [02:25<01:15, 10.75s/it]

Prompt: What is FHE?
Response: Fully Homomorphic Encryption (FHE) involves encryption of data or computations in a secure environment. This ensures that data or computations remain encrypted


--------------------------------------------------
Epoch 14/20, Loss: 8.5096, grad norm: None, lr: 0.0006352941176470588

Training Progress:  70%|███████   | 14/20 [02:36<01:04, 10.76s/it]

Prompt: What is FHE?
Response: Fully Homomorphic Encryption (FHE) provides a secure environment for computations on encrypted data. This ensures that sensitive data cannot be decrypted or


--------------------------------------------------
Epoch 15/20, Loss: 8.4425, grad norm: None, lr: 0.0005294117647058824

Training Progress:  75%|███████▌  | 15/20 [02:47<00:54, 10.89s/it]

Prompt: What is FHE?
Response: Fully Homomorphic Encryption (FHE) provides a high level of privacy by allowing computations to be performed directly on encrypted data. This ensures that


--------------------------------------------------
Epoch 16/20, Loss: 7.8181, grad norm: None, lr: 0.0004235294117647059

Training Progress:  80%|████████  | 16/20 [02:59<00:44, 11.10s/it]

Prompt: What is FHE?
Response: Fully Homomorphic Encryption (FHE) is a type of ciphertext that serves as a data store for sensitive data. It has two main components


--------------------------------------------------
Epoch 17/20, Loss: 6.6628, grad norm: None, lr: 0.0003176470588235294

Training Progress:  85%|████████▌ | 17/20 [03:10<00:33, 11.16s/it]

Prompt: What is FHE?
Response: Fully Homomorphic Encryption (FHE) provides a high level of privacy by allowing computations to be performed directly on encrypted data. This ensures that


--------------------------------------------------
Epoch 18/20, Loss: 6.7317, grad norm: None, lr: 0.00021176470588235295

Training Progress:  90%|█████████ | 18/20 [03:21<00:22, 11.10s/it]

Prompt: What is FHE?
Response: Fully Homomorphic Encryption (FHE) provides a high level of privacy by allowing computations on encrypted data without ever decrypting it. This ensures


--------------------------------------------------
Epoch 19/20, Loss: 6.7047, grad norm: None, lr: 0.00010588235294117647

Training Progress:  95%|█████████▌| 19/20 [03:31<00:10, 10.86s/it]

Prompt: What is FHE?
Response: Fully Homomorphic Encryption (FHE) provides a high level of privacy by allowing computations on encrypted data without ever decrypting it. This ensures


--------------------------------------------------
Epoch 20/20, Loss: 6.2703, grad norm: None, lr: 0.0

Training Progress: 100%|██████████| 20/20 [03:42<00:00, 10.75s/it]

Prompt: What is FHE?
Response: Fully Homomorphic Encryption (FHE) provides a high level of privacy by allowing computations on encrypted data without ever decrypting it. This ensures


--------------------------------------------------

Training Progress: 100%|██████████| 20/20 [03:42<00:00, 11.14s/it]

Note that our goal is to showcase the use of FHE for encrypted fine-tuning. The dataset consists of 68 examples and a total of 2,386 tokens, which is relatively small. Despite its limited size, which offers little support for the model's learning process, it still manages to produce interesting results.

In [15]:

# Get the fine-tuned model
fine_tuned_model = hybrid_model.model.inference_model

# Set FHE mode to disable for text generation
hybrid_model.set_fhe_mode("execute")

In [ ]:

# Inference using the fine-tuned model with LoRA weights
# Seed for reproducibility
torch.manual_seed(SEED)

fine_tuned_model.enable_adapter_layers()
prompt = "who invented FHE?"
_ = generate_and_print(prompt, fine_tuned_model, tokenizer, SEED)

Prompt: who invented FHE?
Response: Fully Homomorphic Encryption (HE) allows users to perform complex computations on encrypted data without having to decode them. This allows for

In [17]:

# Original inference without LoRA weights
# Seed for reproducibility
torch.manual_seed(SEED)

peft_model.disable_adapter_layers()

prompt = "What is FHE?"
generate_and_print(prompt, peft_model, tokenizer, SEED)

Out [17]:

Prompt: What is FHE?
Response: The first, the first,


A.

The first,

.

.

.
The

In [18]:

peft_model.enable_adapter_layers()

# Print weights and model size
total_weights_size = print_weights_and_size(hybrid_model.model)

Total number of weights: 124734720
Total number of LoRA weights: 294912

In [19]:

# Save the model
path = Path("deployment/gpt2_lora_finetuned")
path.mkdir(parents=True, exist_ok=True)

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

hybrid_model.save_and_clear_private_info(path)

In [20]:

# Print weights and size after saving
total_weights_size_private = print_weights_and_size(hybrid_model.model)

Total number of weights: 39717120
Total number of LoRA weights: 294912

In [21]:

# Calculate and print the percentage of weights removed
percentage_removed = (total_weights_size - total_weights_size_private) / total_weights_size * 100
print(f"Total weights removed: {percentage_removed:.2f} %")

Total weights removed: 68.16 %

In [22]:

# Note: Around 95% of the remaining weights are from the embedding layers (wpe and wte)
# as well as the final lm_head layer.

74 KiB Raw Permalink Blame History

Fine-Tuning GPT-2 on Encrypted Data with LoRA and TorusML

74 KiB

Raw Permalink Blame History