Module neuralops.trainer.methods

Training method plugins: LoRA, DoRA, DisLoRA, Fisher, GeLoRA, LLRD, Mixout, R‑Drop, R3F. All methods are registered via the METHOD_REGISTRY and can be applied by name.

Sub-modules

neuralops.trainer.methods.fisher

Fisher information based importance scoring for adaptive rank allocation.

neuralops.trainer.methods.gelora

GeLoRA: Intrinsic dimension estimation for adaptive rank selection.

neuralops.trainer.methods.llrd

LLRD: Layer‑wise Learning Rate Decay.

neuralops.trainer.methods.lora

LoRA, DoRA, DisLoRA implementations.

neuralops.trainer.methods.mixout

Mixout: Regularisation by mixing weights with a frozen copy.

neuralops.trainer.methods.rdrop_r3f

R‑Drop and R3F (regularised dropout / robust representation learning).

neuralops.trainer.methods.registry

Functions

def apply_dislora(model, config) ‑> None

Expand source code

@register_method("dislora")
def apply_dislora(model, config) -> None:
    """Apply dislora."""
    _apply_adapters(
        model,
        config.lora_r,
        config.lora_alpha,
        config.lora_dropout,
        "dislora",
        dislora_beta=config.dislora_beta,
        dislora_capacity_factor=config.dislora_capacity_factor,
    )

Apply dislora.

def apply_dora(model, config) ‑> None

Expand source code

@register_method("dora")
def apply_dora(model, config) -> None:
    """Apply dora."""
    _apply_adapters(model, config.lora_r, config.lora_alpha, config.lora_dropout, "dora")

Apply dora.

def apply_fisher_gelora(model, config) ‑> None

Expand source code

@register_method("fisher_gelora")
def apply_fisher_gelora(model, config) -> None:
    """Apply fisher gelora."""
    logger.info("Fisher‑GeLoRA registered. Use compute_gelora_base_ranks in training script.")

Apply fisher gelora.

def apply_fisher_lora(model, config) ‑> None

Expand source code

@register_method("fisher_lora")
def apply_fisher_lora(model, config) -> None:
    """Apply fisher lora."""
    logger.info(
        "Fisher‑LoRA method called – but requires dataset, tokenizer, task_keys. "
        "Use a wrapper in training script to call compute_fisher_importance with actual data."
    )

Apply fisher lora.

def apply_llrd(model, config) ‑> None

Expand source code

@register_method("llrd")
def apply_llrd(model, config) -> None:
    """Apply llrd."""
    logger.info("LLRD registered. Use --method llrd to enable layer‑wise learning rates.")

Apply llrd.

def apply_lora(model, config) ‑> None

Expand source code

@register_method("lora")
def apply_lora(model, config) -> None:
    """Apply lora."""
    _apply_adapters(model, config.lora_r, config.lora_alpha, config.lora_dropout, "lora")

Apply lora.

def apply_mixout(model, config) ‑> None

Expand source code

@register_method("mixout")
def apply_mixout(model, config) -> None:
    """Apply mixout."""
    replace_linears_with_mixout(model, config.mixout_p)
    logger.info("Mixout applied.")

Apply mixout.

def apply_r3f(model, config) ‑> None

Expand source code

@register_method("r3f")
def apply_r3f(model, config) -> None:
    """Apply r3f."""
    pass

Apply r3f.

def apply_rdrop(model, config) ‑> None

Expand source code

@register_method("rdrop")
def apply_rdrop(model, config) -> None:
    """Apply rdrop."""
    pass

Apply rdrop.

def available_methods() ‑> Any

Expand source code

def available_methods() -> Any:
    """Return list of all registered method names."""
    return list(METHOD_REGISTRY.keys())

Return list of all registered method names.

def collect_hidden_states(model, dataloader, num_batches=10, max_samples=800) ‑> Any

Expand source code

def collect_hidden_states(model, dataloader, num_batches=10, max_samples=800) -> Any:
    """Collect hidden states."""
    layers = get_encoder_layers(model)
    num_layers = len(layers)
    hidden_states_per_layer = {i: [] for i in range(num_layers)}
    hooks = []

    def hook_fn(layer_idx) -> Any:
        """Hook fn."""

        def hook(module, input, output) -> None:
            """Hook."""
            hs = output[0].detach() if isinstance(output, tuple) else output.detach()
            hidden_states_per_layer[layer_idx].append(hs)

        return hook

    for i, layer in enumerate(layers):
        hook = layer.register_forward_hook(hook_fn(i))
        hooks.append(hook)

    model.eval()
    total = 0
    with torch.no_grad():
        for i, batch in enumerate(dataloader):
            if i >= num_batches or total >= max_samples:
                break
            batch = {k: v.to(model.device) for k, v in batch.items() if isinstance(v, torch.Tensor)}
            model(**batch)
            total += batch["input_ids"].size(0)

    for hook in hooks:
        hook.remove()
    for i in range(num_layers):
        if hidden_states_per_layer[i]:
            hidden_states_per_layer[i] = torch.cat(hidden_states_per_layer[i], dim=0)
        else:
            hidden_states_per_layer[i] = torch.empty(0)
    return hidden_states_per_layer

Collect hidden states.

def compute_fisher_importance(model, dataset, tokenizer, task_keys, config, target_substrings=None) ‑> Any

Expand source code

def compute_fisher_importance(
    model, dataset, tokenizer, task_keys, config, target_substrings=None
) -> Any:
    """Compute and return fisher importance."""
    if target_substrings is None:
        target_substrings = ["query", "key", "value", "dense"]

    sent1_key, sent2_key = task_keys
    device = next(model.parameters()).device
    sample_size = min(config.fisher_sample_size, len(dataset))
    indices = np.random.choice(len(dataset), sample_size, replace=False)
    sampled = dataset.select(indices)

    def collate_fn(batch) -> Any:
        """Collate fn."""
        if sent2_key is None:
            texts = [ex[sent1_key] for ex in batch]
            enc = tokenizer(
                texts,
                truncation=True,
                max_length=config.max_seq_length,
                padding=True,
                return_tensors="pt",
            )
        else:
            texts1 = [ex[sent1_key] for ex in batch]
            texts2 = [ex[sent2_key] for ex in batch]
            enc = tokenizer(
                texts1,
                texts2,
                truncation=True,
                max_length=config.max_seq_length,
                padding=True,
                return_tensors="pt",
            )
        labels = [ex["label"] for ex in batch]
        is_reg = task_keys.get("is_regression", False)
        enc["labels"] = torch.tensor(labels, dtype=torch.float if is_reg else torch.long)
        return enc

    loader = DataLoader(
        sampled, batch_size=config.per_device_train_batch_size, collate_fn=collate_fn
    )
    fisher = {}
    model.eval()
    for param in model.parameters():
        param.requires_grad = True

    for i, batch in enumerate(loader):
        if i >= config.fisher_num_batches:
            break
        batch = {k: v.to(device) for k, v in batch.items()}
        outputs = model(**batch)
        loss = outputs.loss
        loss.backward()
        for name, param in model.named_parameters():
            if param.grad is None:
                continue
            if not any(ts in name for ts in target_substrings):
                continue
            score = (param.grad**2).mean().item()
            module_name = name.replace(".weight", "").replace(".bias", "")
            fisher[module_name] = fisher.get(module_name, 0.0) + score
        model.zero_grad()

    if not fisher:
        return None
    for k in fisher:
        fisher[k] /= config.fisher_num_batches
    scores = np.array(list(fisher.values()))
    scores = scores - np.max(scores)
    exp_scores = np.exp(scores)
    norm_scores = exp_scores / (np.sum(exp_scores) + 1e-8)
    fisher = {k: norm_scores[i] for i, k in enumerate(fisher.keys())}
    return fisher

Compute and return fisher importance.

def compute_gelora_base_ranks(model, dataloader, config) ‑> Any

Expand source code

def compute_gelora_base_ranks(model, dataloader, config) -> Any:
    """Compute and return gelora base ranks."""
    layer_hidden = collect_hidden_states(model, dataloader, num_batches=8, max_samples=800)
    layer_id = {}
    for idx, hs in layer_hidden.items():
        if hs.shape[0] < 10:
            id_val = 1.0
        else:
            id_val = estimate_intrinsic_dimension_pytorch(hs, k=min(5, hs.shape[0] - 1))
        base_rank = max(
            config.gelora_min_rank,
            min(config.gelora_max_rank, int(id_val + config.gelora_margin)),
        )
        layer_id[idx] = base_rank
        logger.info(f"Layer {idx} intrinsic dim ≈ {id_val:.2f} → base rank {base_rank}")

    target_substrings = ["query", "key", "value", "dense"]
    rank_dict = {}
    for name, _ in model.named_modules():
        if (
            any(ts in name for ts in target_substrings)
            and "classifier" not in name
            and "pooler" not in name
        ):
            parts = name.split(".")
            layer_idx = None
            for i, part in enumerate(parts):
                if part == "layer" and i + 1 < len(parts):
                    try:
                        layer_idx = int(parts[i + 1])
                        break
                    except (ValueError, TypeError):
                        pass
            rank_dict[name] = layer_id.get(layer_idx, config.lora_r)
    return rank_dict

Compute and return gelora base ranks.

def estimate_intrinsic_dimension_pytorch(hidden_states, k=5) ‑> Any

Expand source code

def estimate_intrinsic_dimension_pytorch(hidden_states, k=5) -> Any:
    """Estimate intrinsic dimension pytorch."""
    if hidden_states.dim() == 3:
        hidden_states = hidden_states.mean(dim=1)
    elif hidden_states.dim() != 2:
        raise ValueError(f"Expected 2D or 3D tensor, got {hidden_states.dim()}D")
    N = hidden_states.shape[0]
    if k + 1 > N:
        return 1.0
    if N > 1000:
        idx = torch.randperm(N)[:1000]
        hidden_states = hidden_states[idx]
    dist = torch.cdist(hidden_states, hidden_states)
    dist.fill_diagonal_(float("inf"))
    topk_dist, _ = torch.topk(dist, k, dim=1, largest=False)
    d_k = topk_dist[:, -1:]
    log_ratios = torch.log(topk_dist / (d_k + 1e-8))
    id_per_point = -1.0 / torch.mean(log_ratios, dim=1)
    id_global = torch.mean(id_per_point).item()
    return max(1.0, id_global)

Estimate intrinsic dimension pytorch.

def get_llrd_param_groups(model, base_lr, weight_decay, decay) ‑> Any

Expand source code

def get_llrd_param_groups(model, base_lr, weight_decay, decay) -> Any:
    """Return the llrd param groups."""
    no_decay = ["bias", "LayerNorm.weight", "layer_norm.weight", "layernorm.weight"]
    max_layer = -1
    for name, _ in model.named_parameters():
        if "encoder.layer." in name:
            layer = int(name.split("encoder.layer.")[1].split(".")[0])
            max_layer = max(max_layer, layer)
    num_layers = max_layer + 1

    emb_decay, emb_nodecay = [], []
    layer_decay = {i: [] for i in range(num_layers)}
    layer_nodecay = {i: [] for i in range(num_layers)}
    head_decay, head_nodecay = [], []

    for name, param in model.named_parameters():
        if "embeddings" in name:
            if any(nd in name for nd in no_decay):
                emb_nodecay.append(param)
            else:
                emb_decay.append(param)
        elif "encoder.layer." in name:
            layer = int(name.split("encoder.layer.")[1].split(".")[0])
            if any(nd in name for nd in no_decay):
                layer_nodecay[layer].append(param)
            else:
                layer_decay[layer].append(param)
        elif "classifier" in name or "pooler" in name:
            if any(nd in name for nd in no_decay):
                head_nodecay.append(param)
            else:
                head_decay.append(param)

    groups = []
    emb_lr = base_lr * (decay ** (num_layers + 1))
    if emb_decay:
        groups.append({"params": emb_decay, "lr": emb_lr, "weight_decay": weight_decay})
    if emb_nodecay:
        groups.append({"params": emb_nodecay, "lr": emb_lr, "weight_decay": 0.0})

    for layer in range(num_layers):
        lr_layer = base_lr * (decay ** (num_layers - layer))
        if layer_decay[layer]:
            groups.append(
                {
                    "params": layer_decay[layer],
                    "lr": lr_layer,
                    "weight_decay": weight_decay,
                }
            )
        if layer_nodecay[layer]:
            groups.append({"params": layer_nodecay[layer], "lr": lr_layer, "weight_decay": 0.0})

    if head_decay:
        groups.append({"params": head_decay, "lr": base_lr, "weight_decay": weight_decay})
    if head_nodecay:
        groups.append({"params": head_nodecay, "lr": base_lr, "weight_decay": 0.0})
    return groups

Return the llrd param groups.

def get_method(name: str) ‑> Any

Expand source code

def get_method(name: str) -> Any:
    """Return the method function by name, or None if not found."""
    return METHOD_REGISTRY.get(name)

Return the method function by name, or None if not found.

def r3f_forward(model, input_ids, attention_mask, token_type_ids, labels, epsilon, is_regression) ‑> Any

Expand source code

def r3f_forward(
    model, input_ids, attention_mask, token_type_ids, labels, epsilon, is_regression
) -> Any:
    """R3f forward."""
    base_kwargs = {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "labels": labels,
    }
    if token_type_ids is not None:
        base_kwargs["token_type_ids"] = token_type_ids
    out_clean = model(**base_kwargs)
    logits_clean = out_clean.logits

    embed = model.get_input_embeddings()
    inputs_embeds = embed(input_ids)
    noise = torch.randn_like(inputs_embeds) * epsilon
    noise_kwargs = {
        "inputs_embeds": inputs_embeds + noise,
        "attention_mask": attention_mask,
        "labels": labels,
    }
    if token_type_ids is not None:
        noise_kwargs["token_type_ids"] = token_type_ids
    out_noisy = model(**noise_kwargs)
    logits_noisy = out_noisy.logits

    if is_regression:
        mse = nn.MSELoss()
        loss_clean = mse(logits_clean.view(-1), labels.view(-1))
        loss_noisy = mse(logits_noisy.view(-1), labels.view(-1))
        p = torch.distributions.Normal(logits_clean, torch.ones_like(logits_clean))
        q = torch.distributions.Normal(logits_noisy, torch.ones_like(logits_noisy))
        kl = torch.distributions.kl_divergence(p, q).mean()
        loss = (loss_clean + loss_noisy) / 2.0 + kl
    else:
        ce = nn.CrossEntropyLoss()
        loss_clean = ce(logits_clean.view(-1, logits_clean.size(-1)), labels.view(-1))
        loss_noisy = ce(logits_noisy.view(-1, logits_noisy.size(-1)), labels.view(-1))
        logp_clean = F.log_softmax(logits_clean, dim=-1)
        logp_noisy = F.log_softmax(logits_noisy, dim=-1)
        p = torch.softmax(logits_clean, dim=-1)
        q = torch.softmax(logits_noisy, dim=-1)
        kl1 = F.kl_div(logp_clean, q, reduction="batchmean")
        kl2 = F.kl_div(logp_noisy, p, reduction="batchmean")
        loss = (loss_clean + loss_noisy) / 2.0 + (kl1 + kl2) / 2.0
    return loss, logits_clean

R3f forward.

def rdrop_loss(logits1, logits2, labels, is_regression, alpha) ‑> Any

Expand source code

def rdrop_loss(logits1, logits2, labels, is_regression, alpha) -> Any:
    """Rdrop loss."""
    if is_regression:
        loss_fn = nn.MSELoss()
        loss1 = loss_fn(logits1.view(-1), labels.view(-1))
        loss2 = loss_fn(logits2.view(-1), labels.view(-1))
        p = torch.distributions.Normal(logits1, torch.ones_like(logits1))
        q = torch.distributions.Normal(logits2, torch.ones_like(logits2))
        kl = (
            torch.distributions.kl_divergence(p, q).mean()
            + torch.distributions.kl_divergence(q, p).mean()
        ) / 2.0
        loss = (loss1 + loss2) / 2.0 + alpha * kl
    else:
        ce = nn.CrossEntropyLoss()
        loss1 = ce(logits1.view(-1, logits1.size(-1)), labels.view(-1))
        loss2 = ce(logits2.view(-1, logits2.size(-1)), labels.view(-1))
        logp1 = F.log_softmax(logits1, dim=-1)
        logp2 = F.log_softmax(logits2, dim=-1)
        p1 = torch.softmax(logits1, dim=-1)
        p2 = torch.softmax(logits2, dim=-1)
        kl1 = F.kl_div(logp1, p2, reduction="batchmean")
        kl2 = F.kl_div(logp2, p1, reduction="batchmean")
        kl = (kl1 + kl2) / 2.0
        loss = (loss1 + loss2) / 2.0 + alpha * kl
    return loss

Rdrop loss.

def register_method(name: str) ‑> Any

Expand source code

def register_method(name: str) -> Any:
    """Decorator to register a method callable."""

    def decorator(fn) -> Any:
        """Decorator."""
        METHOD_REGISTRY[name] = fn
        return fn

    return decorator

Decorator to register a method callable.

Classes

class DisLoRALinear (linear: torch.nn.modules.linear.Linear,
r: int,
alpha: float,
dropout: float,
beta: float = 0.5,
capacity_factor: float = 1.0)

Expand source code

class DisLoRALinear(nn.Module):
    def __init__(
        self,
        linear: nn.Linear,
        r: int,
        alpha: float,
        dropout: float,
        beta: float = 0.5,
        capacity_factor: float = 1.0,
    ) -> None:
        super().__init__()
        self.linear = linear
        self.linear.weight.requires_grad = False
        if linear.bias is not None:
            linear.bias.requires_grad = False
        self.r = r
        self.beta = beta
        self.capacity_factor = capacity_factor
        if r > 0:
            self.lora_A = nn.Linear(linear.in_features, r, bias=False)
            self.lora_B = nn.Linear(r, linear.out_features, bias=False)
            kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5))
            zeros_(self.lora_B.weight)
            self.scaling = alpha / r
            self.dropout = nn.Dropout(dropout)
            self.register_buffer("_eye", torch.eye(r))
        else:
            self.lora_A = self.lora_B = None
            self.scaling = 1.0
            self.dropout = nn.Identity()

    def forward(self, x) -> Any:
        """Forward pass — compute module output from inputs."""
        base = self.linear(x)
        if self.r > 0:
            lora = self.lora_B(self.lora_A(self.dropout(x))) * self.scaling
            return base + lora
        return base

    def ortho_loss(self) -> Any:
        """Ortho loss."""
        if self.lora_A is None:
            return 0.0
        A = self.lora_A.weight
        B = self.lora_B.weight
        reg = (A @ A.T - self._eye).pow(2).sum() + (B.T @ B - self._eye).pow(2).sum()
        return self.beta * reg

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

Methods

def forward(self, x) ‑> Any

Expand source code

def forward(self, x) -> Any:
    """Forward pass — compute module output from inputs."""
    base = self.linear(x)
    if self.r > 0:
        lora = self.lora_B(self.lora_A(self.dropout(x))) * self.scaling
        return base + lora
    return base

Forward pass — compute module output from inputs.

def ortho_loss(self) ‑> Any

Expand source code

def ortho_loss(self) -> Any:
    """Ortho loss."""
    if self.lora_A is None:
        return 0.0
    A = self.lora_A.weight
    B = self.lora_B.weight
    reg = (A @ A.T - self._eye).pow(2).sum() + (B.T @ B - self._eye).pow(2).sum()
    return self.beta * reg

Ortho loss.

class DoRALinear (linear: torch.nn.modules.linear.Linear, r: int, alpha: float, dropout: float)

Expand source code

class DoRALinear(nn.Module):
    def __init__(self, linear: nn.Linear, r: int, alpha: float, dropout: float) -> None:
        """Execute __init__ operation."""
        super().__init__()
        self.linear = linear
        self.linear.weight.requires_grad = False
        if linear.bias is not None:
            linear.bias.requires_grad = False
        self.r = r
        if r > 0:
            self.m = nn.Parameter(torch.ones(linear.out_features))
            self.lora_A = nn.Linear(linear.in_features, r, bias=False)
            self.lora_B = nn.Linear(r, linear.out_features, bias=False)
            kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5))
            zeros_(self.lora_B.weight)
            self.scaling = alpha / r
            self.dropout = nn.Dropout(dropout)
        else:
            self.m = nn.Parameter(torch.ones(linear.out_features))
            self.lora_A = self.lora_B = None
            self.scaling = 1.0
            self.dropout = nn.Identity()

    def forward(self, x) -> Any:
        """Forward pass — compute module output from inputs."""
        base_out = self.linear(x)
        if self.r > 0:
            lora = self.lora_B(self.lora_A(self.dropout(x))) * self.scaling
            weight_norm = self.linear.weight.norm(p=2, dim=1, keepdim=True)
            direction = self.linear.weight / (weight_norm + 1e-8)
            delta = self.m.unsqueeze(1) * direction + lora
            return F.linear(x, delta, self.linear.bias)
        return base_out

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Execute init operation.

Ancestors

• torch.nn.modules.module.Module

Methods

def forward(self, x) ‑> Any

Expand source code

def forward(self, x) -> Any:
    """Forward pass — compute module output from inputs."""
    base_out = self.linear(x)
    if self.r > 0:
        lora = self.lora_B(self.lora_A(self.dropout(x))) * self.scaling
        weight_norm = self.linear.weight.norm(p=2, dim=1, keepdim=True)
        direction = self.linear.weight / (weight_norm + 1e-8)
        delta = self.m.unsqueeze(1) * direction + lora
        return F.linear(x, delta, self.linear.bias)
    return base_out

Forward pass — compute module output from inputs.

class LoRALinear (linear: nn.Module, r: int = 16, alpha: float = 32.0, dropout: float = 0.05)

Expand source code

class LoRALinear(nn.Module):
    """Linear layer wrapped with Low-Rank Adaptation (LoRA) adapters."""

    def __init__(
        self,
        linear: nn.Module,
        r: int = 16,
        alpha: float = 32.0,
        dropout: float = 0.05,
    ) -> None:
        super().__init__()
        if hasattr(linear, "in_features"):
            self.in_features = linear.in_features
            self.out_features = linear.out_features
        elif hasattr(linear, "weight"):
            # HuggingFace Conv1D (used in GPT-2) has weight shape [in_features, out_features]
            self.in_features = linear.weight.shape[0]
            self.out_features = linear.weight.shape[1]
        else:
            raise ValueError(f"Unsupported linear module: {type(linear)}")
        self.r = int(r)
        self.alpha = float(alpha)
        self.scaling = float(alpha) / max(1, self.r) if self.r > 0 else 1.0
        self.merged = False

        # Base linear layer (frozen)
        self.base_layer = linear
        self.base_layer.weight.requires_grad = False
        if self.base_layer.bias is not None:
            self.base_layer.bias.requires_grad = False

        device = linear.weight.device
        # Trainable LoRA parameters are kept in FP32 for numerical stability
        # and compatibility with torch.amp.GradScaler in mixed precision.
        dtype = torch.float32

        # LoRA trainable low-rank matrices
        if self.r > 0:
            self.lora_A = nn.Linear(
                self.in_features, self.r, bias=False, device=device, dtype=dtype
            )
            self.lora_B = nn.Linear(
                self.r, self.out_features, bias=False, device=device, dtype=dtype
            )
            self.lora_dropout = nn.Dropout(dropout) if dropout > 0.0 else nn.Identity()
            self.reset_parameters()
        else:
            self.lora_A = None
            self.lora_B = None
            self.lora_dropout = nn.Identity()

    def reset_parameters(self) -> None:
        """Initialize lora_A with Kaiming uniform and lora_B with zeros."""
        if self.r > 0 and self.lora_A is not None and self.lora_B is not None:
            kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5))
            zeros_(self.lora_B.weight)

    @property
    def weight(self) -> torch.Tensor:
        """Expose base layer weight for compatibility with hooks and inspectors."""
        return self.base_layer.weight

    @property
    def bias(self) -> torch.Tensor | None:
        """Expose base layer bias for compatibility with hooks and inspectors."""
        return self.base_layer.bias

    def merge(self) -> None:
        """In-place merge LoRA delta weight into the base linear layer."""
        if self.r > 0 and not self.merged and self.lora_A is not None and self.lora_B is not None:
            delta_w = (self.lora_B.weight @ self.lora_A.weight) * self.scaling
            is_conv1d = "conv1d" in self.base_layer.__class__.__name__.lower() or (
                not isinstance(self.base_layer, nn.Linear)
                and self.base_layer.weight.shape == (self.in_features, self.out_features)
            )
            if is_conv1d:
                delta_w = delta_w.t()
            self.base_layer.weight.data.add_(delta_w.to(self.base_layer.weight.dtype))
            self.merged = True

    def unmerge(self) -> None:
        """Subtract LoRA delta weight to revert the base linear layer to unmerged state."""
        if self.r > 0 and self.merged and self.lora_A is not None and self.lora_B is not None:
            delta_w = (self.lora_B.weight @ self.lora_A.weight) * self.scaling
            is_conv1d = "conv1d" in self.base_layer.__class__.__name__.lower() or (
                not isinstance(self.base_layer, nn.Linear)
                and self.base_layer.weight.shape == (self.in_features, self.out_features)
            )
            if is_conv1d:
                delta_w = delta_w.t()
            self.base_layer.weight.data.sub_(delta_w.to(self.base_layer.weight.dtype))
            self.merged = False

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass combining base projection with low-rank adapter delta."""
        base_out = self.base_layer(x)
        if self.r > 0 and not self.merged and self.lora_A is not None and self.lora_B is not None:
            lora_in = self.lora_dropout(x).to(dtype=self.lora_A.weight.dtype)
            lora_out = self.lora_B(self.lora_A(lora_in)) * self.scaling
            return base_out + lora_out.to(dtype=base_out.dtype)
        return base_out

    def __repr__(self) -> str:
        return (
            f"LoRALinear(in_features={self.in_features}, out_features={self.out_features}, "
            f"r={self.r}, alpha={self.alpha}, scaling={self.scaling:.4f})"
        )

Linear layer wrapped with Low-Rank Adaptation (LoRA) adapters.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

Instance variables

prop bias : torch.Tensor | None

Expand source code

@property
def bias(self) -> torch.Tensor | None:
    """Expose base layer bias for compatibility with hooks and inspectors."""
    return self.base_layer.bias

Expose base layer bias for compatibility with hooks and inspectors.

prop weight : torch.Tensor

Expand source code

@property
def weight(self) -> torch.Tensor:
    """Expose base layer weight for compatibility with hooks and inspectors."""
    return self.base_layer.weight

Expose base layer weight for compatibility with hooks and inspectors.

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Forward pass combining base projection with low-rank adapter delta."""
    base_out = self.base_layer(x)
    if self.r > 0 and not self.merged and self.lora_A is not None and self.lora_B is not None:
        lora_in = self.lora_dropout(x).to(dtype=self.lora_A.weight.dtype)
        lora_out = self.lora_B(self.lora_A(lora_in)) * self.scaling
        return base_out + lora_out.to(dtype=base_out.dtype)
    return base_out

Forward pass combining base projection with low-rank adapter delta.

def merge(self) ‑> None

Expand source code

def merge(self) -> None:
    """In-place merge LoRA delta weight into the base linear layer."""
    if self.r > 0 and not self.merged and self.lora_A is not None and self.lora_B is not None:
        delta_w = (self.lora_B.weight @ self.lora_A.weight) * self.scaling
        is_conv1d = "conv1d" in self.base_layer.__class__.__name__.lower() or (
            not isinstance(self.base_layer, nn.Linear)
            and self.base_layer.weight.shape == (self.in_features, self.out_features)
        )
        if is_conv1d:
            delta_w = delta_w.t()
        self.base_layer.weight.data.add_(delta_w.to(self.base_layer.weight.dtype))
        self.merged = True

In-place merge LoRA delta weight into the base linear layer.

def reset_parameters(self) ‑> None

Expand source code

def reset_parameters(self) -> None:
    """Initialize lora_A with Kaiming uniform and lora_B with zeros."""
    if self.r > 0 and self.lora_A is not None and self.lora_B is not None:
        kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5))
        zeros_(self.lora_B.weight)

Initialize lora_A with Kaiming uniform and lora_B with zeros.

def unmerge(self) ‑> None

Expand source code

def unmerge(self) -> None:
    """Subtract LoRA delta weight to revert the base linear layer to unmerged state."""
    if self.r > 0 and self.merged and self.lora_A is not None and self.lora_B is not None:
        delta_w = (self.lora_B.weight @ self.lora_A.weight) * self.scaling
        is_conv1d = "conv1d" in self.base_layer.__class__.__name__.lower() or (
            not isinstance(self.base_layer, nn.Linear)
            and self.base_layer.weight.shape == (self.in_features, self.out_features)
        )
        if is_conv1d:
            delta_w = delta_w.t()
        self.base_layer.weight.data.sub_(delta_w.to(self.base_layer.weight.dtype))
        self.merged = False

Subtract LoRA delta weight to revert the base linear layer to unmerged state.

class MethodTrainer (method, rdrop_alpha, r3f_epsilon, is_regression, *args, **kwargs)

Expand source code

class MethodTrainer(Trainer):
    def __init__(self, method, rdrop_alpha, r3f_epsilon, is_regression, *args, **kwargs) -> None:
        """Execute __init__ operation."""
        super().__init__(*args, **kwargs)
        self.method = method
        self.rdrop_alpha = rdrop_alpha
        self.r3f_epsilon = r3f_epsilon
        self.is_regression = is_regression

    def compute_loss(self, model, inputs, return_outputs=False, **kwargs) -> Any:
        """Compute and return loss."""
        if "rdrop" in self.method.lower():
            outputs1 = model(**inputs)
            outputs2 = model(**inputs)
            loss = rdrop_loss(
                outputs1.logits,
                outputs2.logits,
                inputs["labels"],
                self.is_regression,
                self.rdrop_alpha,
            )
            return (loss, outputs1) if return_outputs else loss
        elif "r3f" in self.method.lower():
            loss, logits = r3f_forward(
                model,
                input_ids=inputs["input_ids"],
                attention_mask=inputs.get("attention_mask"),
                token_type_ids=inputs.get("token_type_ids"),
                labels=inputs["labels"],
                epsilon=self.r3f_epsilon,
                is_regression=self.is_regression,
            )
            return (loss, logits) if return_outputs else loss
        return super().compute_loss(model, inputs, return_outputs, **kwargs)

Trainer is a simple but feature-complete training and eval loop for PyTorch, optimized for 🤗 Transformers.

Args -----= model ([PreTrainedModel] or torch.nn.Module, optional): The model to train, evaluate or use for predictions. If not provided, a model_init must be passed.

<Tip>

[`Trainer`] is optimized to work with the [`PreTrainedModel`] provided by the library. You can still use
your own models defined as <code>torch.nn.Module</code> as long as they work the same way as the 🤗 Transformers
models.

</Tip>

args ([TrainingArguments], optional): The arguments to tweak for training. Will default to a basic instance of [TrainingArguments] with the output_dir set to a directory named tmp_trainer in the current directory if not provided. data_collator (DataCollator, optional): The function to use to form a batch from a list of elements of train_dataset or eval_dataset. Will default to [default_data_collator] if no processing_class is provided, an instance of [DataCollatorWithPadding] otherwise if the processing_class is a feature extractor or tokenizer. train_dataset (torch.utils.data.Dataset | torch.utils.data.IterableDataset | datasets.Dataset, optional): The dataset to use for training. If it is a [~datasets.Dataset], columns not accepted by the model.forward() method are automatically removed.

Note that if it's a <code>torch.utils.data.IterableDataset</code> with some randomization and you are training in a
distributed fashion, your iterable dataset should either use a internal attribute <code>generator</code> that is a
<code>torch.Generator</code> for the randomization that must be identical on all processes (and the Trainer will
manually set the seed of this <code>generator</code> at each epoch) or have a <code>set\_epoch()</code> method that internally
sets the seed of the RNGs used.

eval_dataset (torch.utils.data.Dataset | dict[str, torch.utils.data.Dataset] | datasets.Dataset, optional): The dataset to use for evaluation. If it is a [~datasets.Dataset], columns not accepted by the model.forward() method are automatically removed. If it is a dictionary, it will evaluate on each dataset prepending the dictionary key to the metric name. processing_class (PreTrainedTokenizerBase or BaseImageProcessor or FeatureExtractionMixin or ProcessorMixin, optional): Processing class used to process the data. If provided, will be used to automatically process the inputs for the model, and it will be saved along the model to make it easier to rerun an interrupted training or reuse the fine-tuned model. model_init (Callable[[], PreTrainedModel], optional): A function that instantiates the model to be used. If provided, each call to [~Trainer.train] will start from a new instance of the model as given by this function.

The function may have zero argument, or a single one containing the optuna/Ray Tune trial object, to
be able to choose different architectures according to hyperparameters (such as layer count, sizes of
inner layers, dropout probabilities etc).

compute_loss_func (Callable, optional): A function that accepts the raw model outputs, labels, and the number of items in the entire accumulated batch (batch_size * gradient_accumulation_steps) and returns the loss. For example, see the default loss function used by [Trainer]. compute_metrics (Callable[[EvalPrediction], Dict], optional): The function that will be used to compute metrics at evaluation. Must take a [EvalPrediction] and return a dictionary string to metric values. Note When passing TrainingArgs with batch_eval_metrics set to True, your compute_metrics function must take a boolean compute_result argument. This will be triggered after the last eval batch to signal that the function needs to calculate and return the global summary statistics rather than accumulating the batch-level statistics callbacks (List of [TrainerCallback], optional): A list of callbacks to customize the training loop. Will add those to the list of default callbacks detailed in here.

If you want to remove one of the default callbacks used, use the [`Trainer.remove_callback`] method.

optimizers (tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR], optional, defaults to (None, None)): A tuple containing the optimizer and the scheduler to use. Will default to an instance of [AdamW] on your model and a scheduler given by [get_linear_schedule_with_warmup] controlled by args. optimizer_cls_and_kwargs (tuple[Type[torch.optim.Optimizer], dict[str, Any]], optional): A tuple containing the optimizer class and keyword arguments to use. Overrides optim and optim_args in args. Incompatible with the optimizers argument.

Unlike <code>optimizers</code>, this argument avoids the need to place model parameters on the correct devices before initializing the Trainer.

preprocess_logits_for_metrics (Callable[[torch.Tensor, torch.Tensor], torch.Tensor], optional): A function that preprocess the logits right before caching them at each evaluation step. Must take two tensors, the logits and the labels, and return the logits once processed as desired. The modifications made by this function will be reflected in the predictions received by compute_metrics.

Note that the labels (second parameter) will be <code>None</code> if the dataset does not have them.

Important attributes:

- **model** -- Always points to the core model. If using a transformers model, it will be a [`PreTrainedModel`]
  subclass.
- **model_wrapped** -- Always points to the most external model in case one or more other modules wrap the
  original model. This is the model that should be used for the forward pass. For example, under <code>DeepSpeed</code>,
  the inner model is wrapped in <code>DeepSpeed</code> and then again in <code>torch.nn.DistributedDataParallel</code>. If the inner
  model hasn't been wrapped, then <code>self.model\_wrapped</code> is the same as <code>self.model</code>.
- **is_model_parallel** -- Whether or not a model has been switched to a model parallel mode (different from
  data parallelism, this means some of the model layers are split on different GPUs).
- **place_model_on_device** -- Whether or not to automatically place the model on the device. Defaults to
  <code>True</code> unless model parallel, DeepSpeed, FSDP, full fp16/bf16 eval, or SageMaker MP is active. Can be
  overridden by subclassing <code>TrainingArguments</code> and overriding the <code>place\_model\_on\_device</code> property.
- **is_in_train** -- Whether or not a model is currently running <code>train</code> (e.g. when <code>evaluate</code> is called while
  in <code>train</code>)

Execute init operation.

Ancestors

• transformers.trainer.Trainer

Methods

def compute_loss(self, model, inputs, return_outputs=False, **kwargs) ‑> Any

Expand source code

def compute_loss(self, model, inputs, return_outputs=False, **kwargs) -> Any:
    """Compute and return loss."""
    if "rdrop" in self.method.lower():
        outputs1 = model(**inputs)
        outputs2 = model(**inputs)
        loss = rdrop_loss(
            outputs1.logits,
            outputs2.logits,
            inputs["labels"],
            self.is_regression,
            self.rdrop_alpha,
        )
        return (loss, outputs1) if return_outputs else loss
    elif "r3f" in self.method.lower():
        loss, logits = r3f_forward(
            model,
            input_ids=inputs["input_ids"],
            attention_mask=inputs.get("attention_mask"),
            token_type_ids=inputs.get("token_type_ids"),
            labels=inputs["labels"],
            epsilon=self.r3f_epsilon,
            is_regression=self.is_regression,
        )
        return (loss, logits) if return_outputs else loss
    return super().compute_loss(model, inputs, return_outputs, **kwargs)

Compute and return loss.

class MixoutLinear (linear: torch.nn.modules.linear.Linear, p: float)

Expand source code

class MixoutLinear(nn.Module):
    def __init__(self, linear: nn.Linear, p: float) -> None:
        super().__init__()
        self.in_features = linear.in_features
        self.out_features = linear.out_features
        self.weight = linear.weight
        self.bias = linear.bias
        self.p = p
        self.register_buffer("target_weight", self.weight.data.clone())
        if self.bias is not None:
            self.register_buffer("target_bias", self.bias.data.clone())

    def forward(self, x) -> Any:
        """Forward pass — compute module output from inputs."""
        if self.training and self.p > 0:
            mask = torch.bernoulli(torch.full_like(self.weight, 1 - self.p))
            mixed_weight = self.weight * mask + self.target_weight * (1 - mask)
            if self.bias is not None:
                mask_bias = torch.bernoulli(torch.full_like(self.bias, 1 - self.p))
                mixed_bias = self.bias * mask_bias + self.target_bias * (1 - mask_bias)
            else:
                mixed_bias = None
            return F.linear(x, mixed_weight, mixed_bias)
        return F.linear(x, self.weight, self.bias)

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

Methods

def forward(self, x) ‑> Any

Expand source code

def forward(self, x) -> Any:
    """Forward pass — compute module output from inputs."""
    if self.training and self.p > 0:
        mask = torch.bernoulli(torch.full_like(self.weight, 1 - self.p))
        mixed_weight = self.weight * mask + self.target_weight * (1 - mask)
        if self.bias is not None:
            mask_bias = torch.bernoulli(torch.full_like(self.bias, 1 - self.p))
            mixed_bias = self.bias * mask_bias + self.target_bias * (1 - mask_bias)
        else:
            mixed_bias = None
        return F.linear(x, mixed_weight, mixed_bias)
    return F.linear(x, self.weight, self.bias)

Forward pass — compute module output from inputs.