mirror of
https://github.com/luxfi/torus-ml.git
synced 2026-10-30 15:15:23 +00:00
zeekay
e037e8f7ea
The distribution called itself torus-ml while shipping a package directory named concrete, so the name a user installs and the name they import disagreed — and most of the tree already said torus.ml, which could not resolve against src/concrete. Moved src/concrete to src/torus and pointed the include at it. The dependency was Zama's concrete-python from their index. It is torus-fhe from ours, which is what luxfhe/torus already ships: that frontend declares torus-fhe and ships a torus/fhe module, so every 'from torus.fhe' in this tree was already naming our compiler and only the dependency lagged. Two distributions share one namespace, which is the shape the upstream this forked from used and the reason the split works: torus-fhe gives torus.fhe, torus-ml gives torus.ml. Swept every other spelling — identifiers, enum members, workflow variables, filenames, notebooks. LICENSE and NOTICE keep the upstream attribution a fork owes; the figures carry the old name in JPEG metadata a text sweep cannot reach.
74 KiB
74 KiB
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 torus.ml.torch.hybrid_model import HybridFHEModel
from torus.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 = FalseIn [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 batchhuggingface/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 lossIn [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.optimizerIn [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]
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.