workshop/exercises/lora_finetuning/MathEvalLoraLLama.ipynb

Fine-Tuning LLaMA with LoRA for Math Word Problems

This notebook demonstrates how to fine-tune the LLaMA-3.2-1B model using LoRA (Low-Rank Adaptation) on the Orca Math Word Problems dataset. We support multiple modes (Torch, 8-bit, 16-bit) and include evaluation metrics like perplexity.

Setup and Imports

In [1]:

import math
import os
import random
import shutil
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn.functional as F
from datasets import load_dataset
from peft import LoraConfig, get_peft_model
from torch.utils.data import DataLoader
from tqdm import tqdm
from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
from utils_lora import generate_and_print

from torus.ml.torch.lora import LoraTrainer

Configuration Parameters

Define the training mode and other parameters here. Modify these values as needed before running the notebook.

In [2]:

# Training mode: 'torch', '7bit', or '16bit'
MODE = "7bit"

# Number of steps between evaluations
EVAL_STEPS = 100

# Set to True to force CPU
FORCE_CPU = False

# File paths based on mode
mode_str = MODE
EVAL_RESPONSES_FILE = f"eval_generated_responses_{mode_str}.txt"
TRAIN_LOG_FILE = f"training_log_{mode_str}.txt"
SAVE_PATH = Path(f"deployment/llama_lora_finetuned_{mode_str}")

print(f"Mode: {MODE}, Eval Steps: {EVAL_STEPS}, Force CPU: {FORCE_CPU}")

Mode: 7bit, Eval Steps: 100, Force CPU: False

Device and Seed Configuration

Set up the device (CPU/GPU/MPS) and random seeds for reproducibility.

In [3]:

def get_device():
    if FORCE_CPU:
        return "cpu"
    if torch.cuda.is_available():
        return "cuda"
    if torch.backends.mps.is_available() and torch.backends.mps.is_built():
        os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
        return "mps"
    return "cpu"


def set_seed(seed):
    """Set random seeds for reproducibility.

    Args:
        seed (int): The random seed to use
    """
    random.seed(seed)
    np.random.seed(seed)
    torch.use_deterministic_algorithms(True, warn_only=True)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False


SEED = 0
DEVICE = get_device()
set_seed(SEED)
print(f"Using device: {DEVICE}")

Using device: mps

Load Model and Tokenizer

Load the LLaMA model and tokenizer, and test the base model output.

In [4]:

model_name = "meta-llama/Llama-3.2-1B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name).to(DEVICE)

if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
model.config.pad_token_id = model.config.eos_token_id
for p in model.parameters():
    p.requires_grad = False

PROMPT = "When you multiply a number by 7, it becomes 98. What is that number?\n"
_ = generate_and_print(PROMPT, model, tokenizer, seed=SEED)

Prompt: When you multiply a number by 7, it becomes 98. What is that number?

Response: I have a problem with my math. I know that 7 is the answer, but I can't figure out what number it is. I am

LoRA Configuration

Set up LoRA parameters and apply them to the model.

In [5]:

peft_config = LoraConfig(
    r=8,
    lora_alpha=32,
    lora_dropout=0.01,
    bias="none",
    task_type="CAUSAL_LM",
    target_modules="all-linear",
)
peft_model = get_peft_model(model, peft_config)

Dataset Preprocessing

Load the dataset, filter by length, and preprocess for training.

In [6]:

MAX_LENGTH = 64
raw_dataset = load_dataset("microsoft/orca-math-word-problems-200k", split="train")


def length_filter(example):
    q_len = len(tokenizer(example["question"], add_special_tokens=False)["input_ids"])
    a_len = len(tokenizer(example["answer"], add_special_tokens=False)["input_ids"])
    return (q_len + a_len + 1) <= MAX_LENGTH


filtered_dataset = raw_dataset.filter(length_filter)


def get_lengths(example):
    q_len = len(tokenizer(example["question"], add_special_tokens=False)["input_ids"])
    a_len = len(tokenizer(example["answer"], add_special_tokens=False)["input_ids"])
    total_len = q_len + a_len + 1
    return {"q_len": q_len, "a_len": a_len, "total_len": total_len}


lengths = filtered_dataset.map(get_lengths)
q_lengths = [x["q_len"] for x in lengths]
a_lengths = [x["a_len"] for x in lengths]
total_lengths = [x["total_len"] for x in lengths]

print("\nLength Distribution Statistics:")
print(f"Original dataset size: {len(raw_dataset):,}")
print(f"Filtered dataset size: {len(filtered_dataset):,}")
print(f"Percentage kept: {100 * len(filtered_dataset)/len(raw_dataset):.1f}%\n")
print("Question lengths: ")
print(f"  Min: {min(q_lengths)}, Max: {max(q_lengths)}")
print(f"  Mean: {sum(q_lengths)/len(q_lengths):.1f}")
print(f"  Median: {sorted(q_lengths)[len(q_lengths)//2]}")
print("\nAnswer lengths:")
print(f"  Min: {min(a_lengths)}, Max: {max(a_lengths)}")
print(f"  Mean: {sum(a_lengths)/len(a_lengths):.1f}")
print(f"  Median: {sorted(a_lengths)[len(a_lengths)//2]}")
print("\nTotal lengths (including newline):")
print(f"  Min: {min(total_lengths)}, Max: {max(total_lengths)}")
print(f"  Mean: {sum(total_lengths)/len(total_lengths):.1f}")
print(f"  Median: {sorted(total_lengths)[len(total_lengths)//2]}\n")


def process_example(example):
    """Tokenize a question-answer pair and prepare labels for training.

    Args:
        example (dict): Dictionary with 'question' and 'answer' strings
    Returns:
        dict: Processed tokens with masked labels for the question portion
    """
    question = example["question"].strip()
    answer = example["answer"].strip()
    tokens = tokenizer(
        question + "\n" + answer,
        padding="max_length",
        truncation=True,
        max_length=MAX_LENGTH,
    )
    question_length = len(tokenizer(question, add_special_tokens=False)["input_ids"]) + 1
    labels = tokens["input_ids"].copy()
    for i in range(question_length):
        if i < len(labels):
            labels[i] = -100
    tokens["labels"] = labels
    return tokens


tokenized_dataset = filtered_dataset.map(
    process_example,
    batched=False,
    remove_columns=filtered_dataset.column_names,
)

tokenized = tokenized_dataset.train_test_split(test_size=0.05, seed=SEED, shuffle=True)
train_dataset, test_dataset = tokenized["train"], tokenized["test"]
print(f"Train samples: {len(train_dataset)}, Test samples: {len(test_dataset)}")

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)

Length Distribution Statistics:
Original dataset size: 200,035
Filtered dataset size: 1,557
Percentage kept: 0.8%

Question lengths: 
  Min: 4, Max: 61
  Mean: 24.7
  Median: 25

Answer lengths:
  Min: 1, Max: 55
  Mean: 23.5
  Median: 23

Total lengths (including newline):
  Min: 6, Max: 64
  Mean: 49.2
  Median: 51

Train samples: 1479, Test samples: 78

Data Collator

Define a custom data collator to handle padding while preserving label masking.

In [7]:

class DataCollator:
    def __init__(self, tokenizer):
        self.pad_id = tokenizer.pad_token_id

    def __call__(self, examples):
        inputs, attention_masks, labels = [], [], []
        max_length = max(len(example["input_ids"]) for example in examples)
        for example in examples:
            inputs.append(example["input_ids"])
            attention_masks.append(example["attention_mask"])
            labels.append(example["labels"])

        def pad(sequences, value):
            return [x + [value] * (max_length - len(x)) for x in sequences]

        return {
            "input_ids": torch.tensor(pad(inputs, self.pad_id)),
            "attention_mask": torch.tensor(pad(attention_masks, 0)),
            "labels": torch.tensor(pad(labels, -100)),
        }


collator = DataCollator(tokenizer)

Training Arguments

Configure the training hyperparameters.

In [8]:

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

Loss and Evaluation Metrics

Define the loss function and evaluation metric (perplexity).

In [9]:

def causal_lm_loss(logits, labels):
    shift_logits = logits[..., :-1, :].contiguous()
    shift_labels = labels[..., 1:].contiguous()
    return F.cross_entropy(
        shift_logits.view(-1, shift_logits.size(-1)),
        shift_labels.view(-1),
        ignore_index=-100,
        reduction="mean",
    )


def metric_fn(model, dataloader):
    model.eval()
    model.to(DEVICE)
    total_loss, total_tokens, results = 0.0, 0, []
    response = generate_and_print(PROMPT, model, tokenizer, seed=SEED)
    if response:
        results.append({"prompt": PROMPT, "response": response})

    for batch in tqdm(dataloader, desc="Evaluating", leave=False):
        with torch.no_grad():
            input_ids = batch["input_ids"].to(DEVICE)
            batch_labels = batch["labels"].to(DEVICE)
            attention_mask = batch["attention_mask"].to(DEVICE)
            outputs = model(input_ids, attention_mask=attention_mask).logits
            valid = batch_labels[..., 1:] != -100
            loss = F.cross_entropy(
                outputs[..., :-1, :].contiguous().view(-1, outputs.size(-1)),
                batch_labels[..., 1:].contiguous().view(-1),
                ignore_index=-100,
                reduction="sum",
            )
        total_loss += loss.item()
        total_tokens += valid.sum().item()
    perplexity = math.exp(total_loss / total_tokens) if total_tokens > 0 else float("inf")

    with open(EVAL_RESPONSES_FILE, "a", encoding="utf-8") as f:
        f.write(f"Perplexity: {perplexity:.2f}\n")
        for i, r in enumerate(results):
            f.write(
                f"== Generation {i+1} ==\nPrompt:\n{r['prompt']}\n\nResponse:\n{r['response']}\n"
            )
            f.write("=" * 40 + "\n")

    return {"perplexity": perplexity}

Trainer Setup

Initialize the Hugging Face Trainer to set up the optimizer and scheduler.

In [10]:

hf_trainer = Trainer(
    model=peft_model,
    args=training_args,
    train_dataset=train_dataset,
    data_collator=collator,
)
train_dl = hf_trainer.get_train_dataloader()
hf_trainer.create_optimizer_and_scheduler(len(train_dl) * training_args.num_train_epochs)
optimizer, lr_scheduler = hf_trainer.optimizer, hf_trainer.lr_scheduler

Calibration and Evaluation Data

Prepare dummy calibration data and the evaluation DataLoader.

In [11]:

def get_inputset():
    return {
        "input_ids": torch.randint(0, tokenizer.vocab_size, (4, MAX_LENGTH)),
        "attention_mask": torch.ones((4, MAX_LENGTH), dtype=torch.long),
        "labels": torch.randint(0, tokenizer.vocab_size, (4, MAX_LENGTH)),
    }


eval_dl = DataLoader(test_dataset, batch_size=4, shuffle=False, collate_fn=collator)

Model Initialization Function

Add a function to initialize models with consistent weights for our comparison experiments.

In [12]:

def initialize_model(seed=SEED):
    """Initialize the model with consistent weights."""
    set_seed(seed)

    model = AutoModelForCausalLM.from_pretrained(model_name).to(DEVICE)

    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    model.config.pad_token_id = model.config.eos_token_id
    for p in model.parameters():
        p.requires_grad = False

    peft_config = LoraConfig(
        r=8,
        lora_alpha=32,
        lora_dropout=0.01,
        bias="none",
        task_type="CAUSAL_LM",
        target_modules="all-linear",
    )
    peft_model = get_peft_model(model, peft_config)
    peft_model = peft_model.to(DEVICE)
    return peft_model


def setup_trainer(model, steps=EVAL_STEPS, logging_steps=1):
    """Set up the trainer, optimizer and scheduler."""
    hf_trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        data_collator=collator,
    )
    train_dl = hf_trainer.get_train_dataloader()
    hf_trainer.create_optimizer_and_scheduler(len(train_dl) * training_args.num_train_epochs)
    optimizer, lr_scheduler = hf_trainer.optimizer, hf_trainer.lr_scheduler

    lora_trainer = LoraTrainer(
        model=model,
        optimizer=optimizer,
        loss_fn=causal_lm_loss,
        lr_scheduler=lr_scheduler,
        training_args=vars(training_args),
        n_layers_to_skip_for_backprop=3,
        eval_loader=eval_dl,
        eval_metric_fn=metric_fn,
        logging_steps=logging_steps,
        eval_steps=steps,
        train_log_path=TRAIN_LOG_FILE,
    )

    if MODE != "torch":
        bits = 7 if MODE == "7bit" else 16
        lora_trainer.compile(get_inputset(), n_bits=bits)

    return lora_trainer, train_dl


def extract_lora_weights(model):
    """Extract LoRA weights from model for comparison."""
    weights = {}
    for name, param in model.named_parameters():
        if "lora_" in name and param.requires_grad:
            weights[name] = param.detach().cpu().numpy().copy()
    return weights


# Select only the first 5 batches for initial comparison
def get_limited_batches(dataloader, num_batches=5):
    """Get a limited number of batches from dataloader."""
    limited_batches = []
    for i, batch in enumerate(dataloader):
        if i < num_batches:
            limited_batches.append(batch)
        else:
            break
    return limited_batches

Experiment: Comparing fhe="disable" vs fhe="execute"

Train two identical models for 5 steps using the quantized clear model (disable mode) and compare against FHE fine-tuning (execute mode)

In [13]:

# Original model trainer for reference
lora_trainer = LoraTrainer(
    model=peft_model,
    optimizer=optimizer,
    loss_fn=causal_lm_loss,
    lr_scheduler=lr_scheduler,
    training_args=vars(training_args),
    n_layers_to_skip_for_backprop=3,
    eval_loader=eval_dl,
    eval_metric_fn=None,
    logging_steps=1,
    eval_steps=EVAL_STEPS,
    train_log_path=TRAIN_LOG_FILE,
)

# Clear the existing file for our experiments
with open(EVAL_RESPONSES_FILE, "w", encoding="utf-8") as f:
    f.write("=== Experiment: disable vs simulate ===\n")

2025-03-25 15:18:36,880 - INFO - === Starting new training session ===

LoRA layers detected in the model.

In [14]:

print("Initializing first model (fhe=disable)...")
model_disable = initialize_model()
initial_weights_disable = extract_lora_weights(model_disable)

trainer_disable, train_dl = setup_trainer(model_disable)

# Evaluate pre-training
print("Evaluating pre-training model...")
model_disable.eval()
pre_training_metrics = metric_fn(model_disable, eval_dl)
print(f"Pre-training perplexity: {pre_training_metrics['perplexity']:.2f}")

# Get limited batches for consistent training
limited_batches = get_limited_batches(train_dl, 5)

# Train with fhe=disable for 5 steps
print("\nTraining with fhe=disable for 5 steps...")
trainer_disable.train(limited_batches, fhe="disable", device=DEVICE)
losses_disable = trainer_disable.get_training_losses()

# Store LoRA weights
weights_disable = extract_lora_weights(model_disable)

Initializing first model (fhe=disable)...

2025-03-25 15:18:39,194 - INFO - === Starting new training session ===

LoRA layers detected in the model.

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

2025-03-25 15:20:24,975 - INFO - Compilation complete.

Evaluating pre-training model...
Prompt: When you multiply a number by 7, it becomes 98. What is that number?

Response: I have a problem with this question. I know that 7 is a factor of 98, but what is the number that is a factor of

Pre-training perplexity: 37.06

Training with fhe=disable for 5 steps...

Epoch 1:   0%|          | 0/5 [00:00<?, ?it/s]

2025-03-25 15:21:27,313 - INFO - Step 1: loss=2.163410, avg_loss=2.163410
2025-03-25 15:21:27,374 - INFO - Average gradient magnitude: 0.764078
2025-03-25 15:21:29,799 - INFO - Step 2: loss=3.122139, avg_loss=2.642774
2025-03-25 15:21:29,842 - INFO - Average gradient magnitude: 1.383245
2025-03-25 15:21:32,195 - INFO - Step 3: loss=4.424260, avg_loss=3.236603
2025-03-25 15:21:32,240 - INFO - Average gradient magnitude: 2.932004
2025-03-25 15:21:34,564 - INFO - Step 4: loss=0.981824, avg_loss=2.672908
2025-03-25 15:21:34,607 - INFO - Average gradient magnitude: 0.219038
2025-03-25 15:21:37,091 - INFO - Step 5: loss=0.893339, avg_loss=2.316994
2025-03-25 15:21:37,138 - INFO - Average gradient magnitude: 0.286395
2025-03-25 15:21:37,204 - INFO - Epoch 1 completed. Avg Loss: 2.316994, FHE Mode: disable
2025-03-25 15:21:37,204 - INFO - Training completed. Final Avg Loss: 2.316994, FHE Mode: disable

In [15]:

print("\nInitializing second model (fhe=execute)...")
model_execute = initialize_model()
trainer_execute, _ = setup_trainer(model_execute)

# Verify weights are the same before training
print("Verifying initial weights are identical...")
initial_weights_execute = extract_lora_weights(model_execute)
weights_match = True
for name, weight in initial_weights_disable.items():
    if not np.allclose(weight, initial_weights_execute[name]):
        print(f"Weight mismatch in {name}")
        weights_match = False
if weights_match:
    print("Initial weights match between models.")

# Train with fhe=execute for 5 steps
print("\nTraining with fhe=execute for 5 steps...")
trainer_execute.train(limited_batches, fhe="execute", device=DEVICE)
losses_execute = trainer_execute.get_training_losses()

# Compare LoRA weights after training
weights_execute = extract_lora_weights(model_execute)

Initializing second model (fhe=execute)...

2025-03-25 15:24:17,754 - INFO - === Starting new training session ===

LoRA layers detected in the model.

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

2025-03-25 15:26:01,179 - INFO - Compilation complete.

Verifying initial weights are identical...
Initial weights match between models.

Training with fhe=execute for 5 steps...

Epoch 1:   0%|          | 0/5 [00:00<?, ?it/s]

2025-03-25 08:47:17,002 - INFO - Step 1: loss=2.198043, avg_loss=2.198043
2025-03-25 08:47:17,018 - INFO - Average gradient magnitude: 10.796187
2025-03-25 09:45:10,070 - INFO - Step 2: loss=3.195699, avg_loss=2.696871
2025-03-25 09:45:10,087 - INFO - Average gradient magnitude: 14.245789
2025-03-25 10:43:01,591 - INFO - Step 3: loss=5.219449, avg_loss=3.537730
2025-03-25 10:43:01,608 - INFO - Average gradient magnitude: 22.919750
2025-03-25 11:40:54,081 - INFO - Step 4: loss=0.987054, avg_loss=2.900061
2025-03-25 11:40:54,097 - INFO - Average gradient magnitude: 3.477989
2025-03-25 12:38:47,204 - INFO - Step 5: loss=0.822733, avg_loss=2.484596
2025-03-25 12:38:47,222 - INFO - Average gradient magnitude: 7.610127
2025-03-25 12:38:47,240 - INFO - Epoch 1 completed. Avg Loss: 2.484596, FHE Mode: execute
2025-03-25 12:38:47,241 - INFO - Training completed. Final Avg Loss: 2.484596, FHE Mode: execute

In [16]:

plt.figure(figsize=(10, 6))
steps = np.arange(len(losses_disable))

plt.plot(steps, losses_disable, label="Disable Loss", color="blue")
plt.plot(steps, losses_execute, label="Execute Loss", color="red", linestyle="--")

plt.scatter(steps, losses_disable, color="blue", s=100)
plt.scatter(steps, losses_execute, color="red", marker="s", s=100)

plt.xlabel("Step")
plt.ylabel("Loss")
plt.title("Loss Comparison")
plt.legend()
plt.grid()
plt.xticks(steps)

plt.tight_layout()
plt.show()

Full Training

After validating the model's FHE behavior, the quantized model is trained in clear on the filtered dataset to validate convergence.

In [17]:

# Fully train the model in disable mode
model_final = initialize_model()
trainer_final, full_train_dl = setup_trainer(model_final, logging_steps=100)

trainer_final.train(full_train_dl, fhe="disable", device=DEVICE)

# Final evaluation
print("\nEvaluating final model after full training...")
model_final.eval()
metrics_final = metric_fn(model_final, eval_dl)
print(f"Final perplexity after extended training: {metrics_final['perplexity']:.2f}")
print(
    f"Improvement from initial model: "
    f"{pre_training_metrics['perplexity'] - metrics_final['perplexity']:.2f}"
)

2025-03-26 10:29:58,595 - INFO - === Starting new training session ===

LoRA layers detected in the model.

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

2025-03-26 10:31:43,299 - INFO - Compilation complete.

Epoch 1:   0%|          | 0/1479 [00:00<?, ?it/s]

2025-03-26 10:35:42,918 - INFO - Step 100: loss=0.335604, avg_loss=0.592850
2025-03-26 10:35:42,962 - INFO - Average gradient magnitude: 0.244334
2025-03-26 10:35:42,963 - INFO - Running evaluation at step 100...

Prompt: When you multiply a number by 7, it becomes 98. What is that number?

Response: If you multiply a number by 7, it becomes 98. So, the number you're asking about is 98.

2025-03-26 10:36:44,321 - INFO - [Evaluation at step 100] perplexity: 1.4078
2025-03-26 10:40:45,460 - INFO - Step 200: loss=0.247193, avg_loss=0.450941
2025-03-26 10:40:45,501 - INFO - Average gradient magnitude: 0.335301
2025-03-26 10:40:45,502 - INFO - Running evaluation at step 200...

Prompt: When you multiply a number by 7, it becomes 98. What is that number?

Response: The number that 98 is multiplied by is 7. So, the number is 98 * 7 = 702.

2025-03-26 10:41:44,560 - INFO - [Evaluation at step 200] perplexity: 1.4067
2025-03-26 10:45:56,970 - INFO - Step 300: loss=0.079890, avg_loss=0.444334
2025-03-26 10:45:57,016 - INFO - Average gradient magnitude: 0.091397
2025-03-26 10:45:57,017 - INFO - Running evaluation at step 300...

Prompt: When you multiply a number by 7, it becomes 98. What is that number?

Response: If you multiply 7 by 98, you get 98 * 7 = 702.

2025-03-26 10:46:49,066 - INFO - [Evaluation at step 300] perplexity: 1.3947
2025-03-26 10:51:04,550 - INFO - Step 400: loss=0.016669, avg_loss=0.420956
2025-03-26 10:51:04,598 - INFO - Average gradient magnitude: 0.034923
2025-03-26 10:51:04,598 - INFO - Running evaluation at step 400...

Prompt: When you multiply a number by 7, it becomes 98. What is that number?

Response: The number 98 is the answer to the question.

2025-03-26 10:51:50,182 - INFO - [Evaluation at step 400] perplexity: 1.3715
2025-03-26 10:55:48,938 - INFO - Step 500: loss=0.268339, avg_loss=0.405399
2025-03-26 10:55:48,981 - INFO - Average gradient magnitude: 0.268719
2025-03-26 10:55:48,982 - INFO - Running evaluation at step 500...

Prompt: When you multiply a number by 7, it becomes 98. What is that number?

Response: If you multiply 7 by 98, you get 98 * 7 = 672.

2025-03-26 10:56:41,760 - INFO - [Evaluation at step 500] perplexity: 1.3689
2025-03-26 11:00:47,376 - INFO - Step 600: loss=0.061164, avg_loss=0.389228
2025-03-26 11:00:47,424 - INFO - Average gradient magnitude: 0.094080
2025-03-26 11:00:47,424 - INFO - Running evaluation at step 600...

Prompt: When you multiply a number by 7, it becomes 98. What is that number?

Response: The number that becomes 98 when you multiply 7 by itself is 49.

2025-03-26 11:01:37,540 - INFO - [Evaluation at step 600] perplexity: 1.3429
2025-03-26 11:05:34,106 - INFO - Step 700: loss=0.049174, avg_loss=0.372756
2025-03-26 11:05:34,148 - INFO - Average gradient magnitude: 0.076843
2025-03-26 11:05:34,149 - INFO - Running evaluation at step 700...

Prompt: When you multiply a number by 7, it becomes 98. What is that number?

Response: The number that is being multiplied by 7 is 98.

2025-03-26 11:06:18,797 - INFO - [Evaluation at step 700] perplexity: 1.3076
2025-03-26 11:10:16,397 - INFO - Step 800: loss=0.114913, avg_loss=0.366715
2025-03-26 11:10:16,440 - INFO - Average gradient magnitude: 0.201189
2025-03-26 11:10:16,441 - INFO - Running evaluation at step 800...

Prompt: When you multiply a number by 7, it becomes 98. What is that number?

Response: To find the number, you need to divide 98 by 7. 

98 ÷ 7 = 14

So the number is

2025-03-26 11:11:14,821 - INFO - [Evaluation at step 800] perplexity: 1.2973
2025-03-26 11:15:14,813 - INFO - Step 900: loss=0.126561, avg_loss=0.354234
2025-03-26 11:15:14,867 - INFO - Average gradient magnitude: 0.104785
2025-03-26 11:15:14,867 - INFO - Running evaluation at step 900...

Prompt: When you multiply a number by 7, it becomes 98. What is that number?

Response: To find the number, you need to divide 98 by 7. 

98 ÷ 7 = 14

So, the number is

2025-03-26 11:16:16,442 - INFO - [Evaluation at step 900] perplexity: 1.2808
2025-03-26 11:20:13,539 - INFO - Step 1000: loss=0.344479, avg_loss=0.344842
2025-03-26 11:20:13,582 - INFO - Average gradient magnitude: 0.255919
2025-03-26 11:20:13,583 - INFO - Running evaluation at step 1000...

Prompt: When you multiply a number by 7, it becomes 98. What is that number?

Response: To find the number, you need to divide 98 by 7.

98 ÷ 7 = 14

So the number is 14

2025-03-26 11:21:16,761 - INFO - [Evaluation at step 1000] perplexity: 1.2511
2025-03-26 11:25:15,729 - INFO - Step 1100: loss=0.354381, avg_loss=0.335532
2025-03-26 11:25:15,771 - INFO - Average gradient magnitude: 0.290001
2025-03-26 11:25:15,771 - INFO - Running evaluation at step 1100...

Prompt: When you multiply a number by 7, it becomes 98. What is that number?

Response: To find the number, you need to divide 98 by 7.

98 ÷ 7 = 14

So the number is 14

2025-03-26 11:26:15,104 - INFO - [Evaluation at step 1100] perplexity: 1.2496
2025-03-26 11:30:04,476 - INFO - Step 1200: loss=0.130505, avg_loss=0.326929
2025-03-26 11:30:04,518 - INFO - Average gradient magnitude: 0.197031
2025-03-26 11:30:04,518 - INFO - Running evaluation at step 1200...

Prompt: When you multiply a number by 7, it becomes 98. What is that number?

Response: To find the number, you need to divide 98 by 7.

98 ÷ 7 = 14

So the number is 14

2025-03-26 11:31:01,618 - INFO - [Evaluation at step 1200] perplexity: 1.2336
2025-03-26 11:34:59,935 - INFO - Step 1300: loss=0.782768, avg_loss=0.323717
2025-03-26 11:34:59,978 - INFO - Average gradient magnitude: 0.438046
2025-03-26 11:34:59,979 - INFO - Running evaluation at step 1300...

Prompt: When you multiply a number by 7, it becomes 98. What is that number?

Response: To find the number, you need to divide 98 by 7. 

98 ÷ 7 = 14

So the number is

2025-03-26 11:35:59,463 - INFO - [Evaluation at step 1300] perplexity: 1.2246
2025-03-26 11:39:52,631 - INFO - Step 1400: loss=0.323173, avg_loss=0.316051
2025-03-26 11:39:52,673 - INFO - Average gradient magnitude: 0.343263
2025-03-26 11:39:52,673 - INFO - Running evaluation at step 1400...

Prompt: When you multiply a number by 7, it becomes 98. What is that number?

Response: To find the number, you need to divide 98 by 7. 

98 ÷ 7 = 14

So the number is

2025-03-26 11:40:52,313 - INFO - [Evaluation at step 1400] perplexity: 1.2274
2025-03-26 11:43:57,682 - INFO - Epoch 1 completed. Avg Loss: 0.311628, FHE Mode: disable
2025-03-26 11:43:57,682 - INFO - Training completed. Final Avg Loss: 0.311628, FHE Mode: disable

Evaluating final model after full training...
Prompt: When you multiply a number by 7, it becomes 98. What is that number?

Response: To find the number, you need to divide 98 by 7. 

98 ÷ 7 = 14

So the number is

Final perplexity after extended training: 1.22
Improvement from initial model: 35.84

Save the Model

Save the fine-tuned model to disk.

In [18]:

if SAVE_PATH.is_dir() and any(SAVE_PATH.iterdir()):
    shutil.rmtree(SAVE_PATH)
trainer_final.save_and_clear_private_info(SAVE_PATH)
print("Model saved to:", SAVE_PATH)

2025-03-26 11:54:06,388 - INFO - Model saved at deployment/llama_lora_finetuned_7bit

Model saved to: deployment/llama_lora_finetuned_7bit