torus-ml/docs/references/api/torus.ml.torch.lora.md

module torus.ml.torch.lora

This module contains classes for LoRA (Low-Rank Adaptation) FHE training and custom layers.

Global Variables

• LINEAR_LAYERS

---

function setup_logger

setup_logger(log_file: str, level=20)

Set up a logger that logs to both console and a file.

Args:

• log_file (str): The path to the log file.
• level (int): The logging level.

Returns:

• logging.Logger: The logger instance.

---

function grad_to

grad_to(param, device: str) → None

Move parameter gradient to device.

Args:

• param: torch parameter with gradient
• device (str): target device for gradient

---

function optimizer_to

optimizer_to(optim, device)

Move optimizer object to device.

Args:

• optim: torch optimizer
• device (str): target device for gradient

---

function get_remote_names

get_remote_names(
    model: Module,
    include_embedding_layers: bool = False
) → List[str]

Get names of modules to be executed remotely.

Args:

• model (nn.Module): The model to inspect.
• include_embedding_layers (bool): Whether to include embedding layers.

Returns:

• List[str]: List of module names to be executed remotely.

---

class LoraTraining

LoraTraining module for fine-tuning with LoRA in a hybrid model setting.

This class is designed to enable Low-Rank Adaptation (LoRA) fine-tuning in a hybrid model context. It allows selective execution of forward and backward passes in FHE.

The class replaces standard linear layers with custom layers that are compatible with LoRA and FHE operations.

Args:

• model (torch.nn.Module): The base model with LoRA layers to be fine-tuned.
• n_layers_to_skip_for_backprop (int): Number of initial linear layers to keep as standard layers. Since the first layer doesn't need backpropagation (no previous layer to update), we typically skip 1 layer. Defaults to 1.
• loss_fn (callable, optional): Loss function to compute the loss. If None, the model is expected to return a loss.

method __init__

__init__(model, n_layers_to_skip_for_backprop=1, loss_fn=None)

---

method assert_has_lora_layers

assert_has_lora_layers(model)

Assert that the model contains LoRA layers.

Args:

• model (torch.nn.Module): The model to check for LoRA layers.

Raises:

• ValueError: If the model does not contain any LoRA layers.

---

method forward

forward(inputs: Tuple[Tensor, ]) → Tuple[Tensor, Union[Tensor, NoneType]]

Forward pass of the LoRA training module.

Args:

• inputs (tuple): A tuple containing the input tensors.

Returns: A tuple containing the original (unscaled) loss and None.

Raises:

• ValueError: If the model does not return a loss and no loss function is provided.

---

method process_inputs

process_inputs(
    inputs: Any
) → Tuple[Union[Tensor, NoneType], Union[Tensor, NoneType]]

Process training inputs such as labels and attention mask.

Args:

• inputs: a dict, BatchEncoding or tuple containing training data

Returns:

• res: tuple containing the attention mask and the labels

---

method replace_layers_with_custom

replace_layers_with_custom(
    model: Module,
    n_layers_to_skip_for_backprop: int
) → None

Replace linear layers with custom ones.

Args:

• model (nn.Module): The model to replace layers in.
• n_layers_to_skip_for_backprop (int): Number of initial linear layers to keep as standard layers. Since the first layer doesn't need backpropagation (no previous layer to update), we typically skip 1 layer.

---

method set_loss_scaling_factor

set_loss_scaling_factor(loss_scaling_factor: float)

Set a scaling factor for the loss to account for gradient accumulation.

This ensures that gradients are correctly averaged over multiple mini-batches when performing gradient accumulation, preventing them from being scaled up by the number of accumulation steps.

Args:

• loss_scaling_factor (float): The number of gradient accumulation steps. The loss will be divided by this factor before backpropagation.

---

class LoraTrainer

Trainer class for LoRA fine-tuning with FHE support.

This class handles:

• Training loop
• Periodic logging and evaluation
• Loss tracking
• Integration with hybrid FHE model

Args:

• model (nn.Module): The base model with LoRA layers to be fine-tuned.
• optimizer (torch.optim.Optimizer): Optimizer for training.
• loss_fn (callable): Loss function to compute the loss.
• lr_scheduler (optional): Learning rate scheduler.
• training_args (dict): Training arguments.
• n_layers_to_skip_for_backprop (int): Number of initial linear layers to keep as standard layers. Since the first layer doesn't need backpropagation (no previous layer to update), we typically skip 1 layer. Defaults to 1.
• eval_loader (DataLoader, optional): DataLoader for evaluation data.
• eval_metric_fn (callable, optional): Function(model, eval_loader) -> dict of metrics.
• logging_steps (int, optional): Log loss every N training steps. Defaults to 1.
• eval_steps (int, optional): Evaluate on eval set every N training steps. Defaults to 10.
• train_log_path (str, optional): Path to a log file for training. Defaults to "training.log".

method __init__

__init__(
    model,
    optimizer,
    loss_fn=None,
    lr_scheduler=None,
    training_args=None,
    n_layers_to_skip_for_backprop=1,
    eval_loader: Optional[DataLoader] = None,
    eval_metric_fn: Optional[Callable] = None,
    logging_steps: int = 1,
    eval_steps: int = 10,
    train_log_path: str = 'training.log',
    checkpoint_dir: str = None
)

---

method compile

compile(inputset, n_bits=8, use_dynamic_quantization=True)

Compile the hybrid model with the given input set.

Args:

• inputset (tuple): Input set for compilation.
• n_bits (int): Bit width for quantization.
• use_dynamic_quantization (bool): Whether to use dynamic quantization.

Returns:

• Tuple[int, int]: The epoch and global step of the latest checkpoint if found, else (0, 0).

---

method get_gradient_stats

get_gradient_stats()

Return recorded gradient statistics.

Returns:

• Dict[str, List[float]]: Gradient statistics per layer over time.

---

method get_training_losses

get_training_losses()

Return all recorded training losses.

Returns:

• List[float]: All recorded training losses.

---

method load_checkpoint

load_checkpoint(checkpoint_path: str)

Load a training checkpoint and restore model, optimizer, and lr_scheduler.

Args:

• checkpoint_path (str): Path to the checkpoint file.

Returns:

• Tuple[int, int]: The epoch and global step of the checkpoint.

Raises:

• FileNotFoundError: If the checkpoint file is not found.

---

method save_and_clear_private_info

save_and_clear_private_info(path)

Save the model and remove private information.

Args:

• path (str): The path to save the model.

---

method save_checkpoint

save_checkpoint(epoch: int, global_step: int)

Save a training checkpoint.

Args:

• epoch (int): The current epoch number.
• global_step (int): The current global step number.

---

method train

train(
    train_loader: DataLoader,
    num_epochs: int = 1,
    fhe: str = 'simulate',
    device: str = 'cpu'
)

Train the model using the hybrid FHE model.

Args:

• train_loader (DataLoader): DataLoader for training data.
• num_epochs (int): Number of epochs to train.
• fhe (str): FHE mode ('disable', 'simulate', 'execute' or 'torch').
• device (str): A device string that is compatible with PyTorch, used for client-side computations.