Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

AdaLoRA › Code Implementation

Michael BrenndoerferDecember 5, 202513 min read

Language AI HandbookAll chapters

Part 35 of 67Parameter-Efficient Fine-tuning

Chapter 260AdaLoRA › Code Implementation

Reading section 7 of 11 · Code Implementation · 13 minSection 7 of 11 · 13 min

Section 7 of 11 · 13 min read

AdaLoRA › Code ImplementationOne section · Sections in this chapter7 / 11 ↓

Sections in this chapter11

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Let's implement AdaLoRA using the PEFT library, which gives a clean interface for this technique. We will walk through the key configuration choices, inspect the parameter structure to verify the SVD decomposition, and examine how rank allocation evolves during training.

We start by loading a base model. We use DistilBERT for this demonstration because it is small enough to train quickly while having the same architectural structure as larger BERT-family models, making the AdaLoRA mechanics clearly visible.

In[10]:

Code

from transformers import AutoModelForSequenceClassification, AutoTokenizer

# Load a base model
model_name = "distilbert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(
    model_name, num_labels=2
)

We'll configure AdaLoRA with its key hyperparameters. The most important are the initial rank and target rank, which define the budget reduction during training. Setting init_r=12 and target_r=4 means we start with 12 triplets per matrix and prune down to an average of 4, a 3x reduction in effective adaptation parameters.

In[11]:

Code

from peft import AdaLoraConfig, TaskType, get_peft_model

# AdaLoRA configuration
adalora_config = AdaLoraConfig(
    # SVD decomposition settings
    init_r=12,  # Initial rank for all matrices
    target_r=4,  # Target average rank after pruning
    # Importance scoring
    beta1=0.85,  # EMA smoothing for sensitivity
    beta2=0.85,  # EMA smoothing for importance
    # Regularization
    orth_reg_weight=0.5,  # Weight for orthonormality regularization
    # Training schedule
    total_step=100,  # Total training steps (for pruning schedule)
    # Which modules to adapt
    target_modules=["q_lin", "v_lin", "k_lin", "out_lin"],
    # Task type
    task_type=TaskType.SEQ_CLS,
)

# Create AdaLoRA model
peft_model = get_peft_model(model, adalora_config)

Let's examine the parameter structure that AdaLoRA creates. The trainable parameter count shows the initial rank $r_0$ across all adapted matrices.

In[12]:

Code

def count_parameters(model):
    """Count trainable and total parameters."""
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())
    return trainable, total


trainable_params, total_params = count_parameters(peft_model)

Out[13]:

Console

Trainable parameters: 1,034,786
Total parameters: 67,989,820
Trainable %: 1.52%

The trainable parameter count shows the initial rank $r_0$ across all adapted matrices. This count will effectively decrease during training as singular values are pruned (though the actual parameters remain in memory, just zeroed out through the masking mechanism described earlier).

Inspecting SVD ComponentsLink Copied

Let's look at how AdaLoRA structures its parameters compared to standard LoRA. Understanding the naming convention in the PEFT library is helpful: AdaLoRA uses the names lora_A for the $Q$ matrix (right singular vectors), lora_E for the $\Lambda$ vector (singular values), and lora_B for the $P$ matrix (left singular vectors). The naming follows LoRA's original convention for $A$ and $B$ but adds E for the singular values (from "eigenvalues," a related concept).

In[14]:

Code

# Find an AdaLoRA layer to inspect
def find_adalora_layers(model):
    """Find all AdaLoRA-adapted layers."""
    layers = []
    for name, module in model.named_modules():
        if hasattr(module, "lora_E"):  # AdaLoRA uses E for singular values
            layers.append((name, module))
    return layers


adalora_layers = find_adalora_layers(peft_model)

Out[15]:

Console

Found 24 AdaLoRA layers:
  base_model.model.distilbert.transformer.layer.0.attention.q_lin
  base_model.model.distilbert.transformer.layer.0.attention.k_lin
  base_model.model.distilbert.transformer.layer.0.attention.v_lin
  base_model.model.distilbert.transformer.layer.0.attention.out_lin

Each AdaLoRA layer contains $Q$ (called lora_A in PEFT), $\Lambda$ (called lora_E), and $P$ (called lora_B).

In[16]:

Code

# Examine one layer's structure
layer_name, adalora_layer = adalora_layers[0]

Out[17]:

Console

Layer: base_model.model.distilbert.transformer.layer.0.attention.q_lin
  Q (lora_A) shape: torch.Size([12, 768])
  Lambda (lora_E) shape: torch.Size([12, 1])
  P (lora_B) shape: torch.Size([768, 12])
  Effective rank: 12

The shapes confirm the SVD structure: $P$ has dimensions matching the output dimension, $Q$ matches the input dimension, and $\Lambda$ (stored as a vector $E$ ) has length equal to the rank. Notice that $\Lambda$ is stored as a 1D vector rather than a diagonal matrix: this saves memory (storing $r$ values instead of $r^2$ ) and makes the elementwise multiplication with the mask straightforward.

Training with Dynamic PruningLink Copied

To see AdaLoRA's pruning in action, let's set up a simple training loop that tracks rank changes. We use the SST-2 sentiment analysis dataset from the GLUE benchmark as a representative fine-tuning task.

In[18]:

Code

from datasets import load_dataset
from torch.utils.data import DataLoader

# Load a small dataset
dataset = load_dataset("nyu-mll/glue", "sst2", split="train[:1000]")


def tokenize_function(examples):
    return tokenizer(
        examples["sentence"],
        padding="max_length",
        truncation=True,
        max_length=128,
    )


tokenized_dataset = dataset.map(tokenize_function, batched=True)
tokenized_dataset.set_format(
    type="torch", columns=["input_ids", "attention_mask", "label"]
)

train_dataloader = DataLoader(tokenized_dataset, batch_size=16, shuffle=True)

The RankAllocator is the key object in PEFT's AdaLoRA implementation. It encapsulates the importance scoring and global ranking logic, and exposes the update_and_allocate method that should be called after each backward pass.

In[19]:

Code

from peft.tuners.adalora import RankAllocator
from torch.optim import AdamW

# Set up optimizer and training
optimizer = AdamW(peft_model.parameters(), lr=1e-4)

# Create the rank allocator that handles pruning
rankallocator = RankAllocator(peft_model, adalora_config, "default")

# Track rank changes during training
rank_history = []
latest_rank_pattern = None

In[20]:

Code

import torch

# Training loop with rank allocation
peft_model.train()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
peft_model.to(device)

global_step = 0
for epoch in range(2):
    for batch in train_dataloader:
        # Move batch to device
        batch = {k: v.to(device) for k, v in batch.items()}

        # Forward pass
        outputs = peft_model(
            input_ids=batch["input_ids"],
            attention_mask=batch["attention_mask"],
            labels=batch["label"],
        )
        loss = outputs.loss

        # Add orthogonality regularization (API varies by PEFT version)
        if hasattr(rankallocator, "compute_orth_regu"):
            orth_loss = rankallocator.compute_orth_regu(
                peft_model, regu_weight=adalora_config.orth_reg_weight
            )
            total_loss = loss + orth_loss
        else:
            total_loss = loss  # Skip orth reg if not available

        # Backward pass
        total_loss.backward()

        optimizer.step()

        # Mask after the optimizer step so AdamW cannot immediately move
        # pruned singular values away from zero again.
        _, current_rank_pattern = rankallocator.update_and_allocate(
            peft_model, global_step
        )
        if current_rank_pattern is not None:
            latest_rank_pattern = current_rank_pattern

        optimizer.zero_grad()

        global_step += 1

        # Record effective ranks periodically
        if global_step % 10 == 0 and latest_rank_pattern is not None:
            current_ranks = {
                name: int(sum(mask))
                for name, mask in latest_rank_pattern.items()
            }
            rank_history.append((global_step, current_ranks))

        if global_step >= adalora_config.total_step:
            break
    if global_step >= adalora_config.total_step:
        break

Out[21]:

Console

Training completed at step 100
Final loss: 1.1802

The final loss confirms that the model has successfully adapted to the training data. The loss reduction shows that the rank allocation strategy, despite progressively pruning components, does not prevent the model from fitting the task. In fact, the adaptive allocation typically produces lower loss at the same parameter budget compared to uniform LoRA, because it concentrates parameters in the most important directions.

Examining Final Rank DistributionLink Copied

After training, we can inspect how AdaLoRA distributed ranks across different layers. This distribution is the key output of the algorithm: a data-driven answer to the question of which weight matrices deserve more adaptation capacity.

In[22]:

Code

def get_effective_ranks(model):
    """Get the effective rank (non-zero singular values) for each layer."""
    ranks = {}
    for name, module in model.named_modules():
        if hasattr(module, "lora_E"):
            for adapter_name in module.lora_E.keys():
                E = module.lora_E[adapter_name]
                # Count non-zero singular values
                effective_rank = (E.abs() > 1e-6).sum().item()
                ranks[name] = effective_rank
    return ranks


final_ranks = get_effective_ranks(peft_model)

Out[23]:

Console

Final rank distribution:
  base_model.model.distilbert.transformer.layer.0.attention.k_lin: rank 4
  base_model.model.distilbert.transformer.layer.0.attention.out_lin: rank 4
  base_model.model.distilbert.transformer.layer.0.attention.q_lin: rank 4
  base_model.model.distilbert.transformer.layer.0.attention.v_lin: rank 2
  base_model.model.distilbert.transformer.layer.1.attention.k_lin: rank 4
  base_model.model.distilbert.transformer.layer.1.attention.out_lin: rank 5
  base_model.model.distilbert.transformer.layer.1.attention.q_lin: rank 4
  base_model.model.distilbert.transformer.layer.1.attention.v_lin: rank 5
  base_model.model.distilbert.transformer.layer.2.attention.k_lin: rank 3
  base_model.model.distilbert.transformer.layer.2.attention.out_lin: rank 4
  base_model.model.distilbert.transformer.layer.2.attention.q_lin: rank 3
  base_model.model.distilbert.transformer.layer.2.attention.v_lin: rank 5
  base_model.model.distilbert.transformer.layer.3.attention.k_lin: rank 3
  base_model.model.distilbert.transformer.layer.3.attention.out_lin: rank 3
  base_model.model.distilbert.transformer.layer.3.attention.q_lin: rank 3
  base_model.model.distilbert.transformer.layer.3.attention.v_lin: rank 4
  base_model.model.distilbert.transformer.layer.4.attention.k_lin: rank 4
  base_model.model.distilbert.transformer.layer.4.attention.out_lin: rank 4
  base_model.model.distilbert.transformer.layer.4.attention.q_lin: rank 2
  base_model.model.distilbert.transformer.layer.4.attention.v_lin: rank 7
  base_model.model.distilbert.transformer.layer.5.attention.k_lin: rank 4
  base_model.model.distilbert.transformer.layer.5.attention.out_lin: rank 4
  base_model.model.distilbert.transformer.layer.5.attention.q_lin: rank 4
  base_model.model.distilbert.transformer.layer.5.attention.v_lin: rank 7

Total effective rank budget: 96
Average rank per layer: 4.0

The distribution reveals which layers AdaLoRA determined were most important for this task. Layers with higher retained ranks contribute more to the adaptation. Notice that the distribution is heterogeneous even within a single transformer block: different projection types receive different ranks based on their actual importance for sentiment analysis, rather than a uniform assumption about all projections being equally important.

Out[24]:

Visualization

Bar chart showing the final effective rank per adapted layer after AdaLoRA training, with taller bars for higher-importance layers and a horizontal dashed line showing the target average rank. — Effective rank distribution across adapted layers after AdaLoRA training. The algorithm allocates higher ranks to layers with greater impact on task performance. Color-coding separates the four projection types: query, value, key, and output projections. The dashed line shows the target average rank, showing how the learned allocation departs from a uniform distribution.

Visualizing Rank EvolutionLink Copied

In[25]:

Code

import matplotlib.pyplot as plt

use_book_style()
plt.rcParams["figure.figsize"] = (6.0, 4.0)

# Extract training steps and calculated ranks
steps = [h[0] for h in rank_history]
total_ranks = [sum(h[1].values()) for h in rank_history]
target_budget = adalora_config.target_r * len(adalora_layers)

plt.figure()
plt.plot(
    steps,
    total_ranks,
    color=PALETTE["primary"],
    label="Total effective rank",
)
plt.axhline(
    y=target_budget,
    color=PALETTE["coral"],
    linestyle="--",
    label="Target budget",
)
plt.fill_between(steps, total_ranks, color=PALETTE["primary"], alpha=0.2)

plt.xlabel("Training Step")
plt.ylabel("Total Effective Rank")
plt.title("AdaLoRA Rank Pruning Schedule")
plt.legend(
    loc="upper center",
    bbox_to_anchor=(0.5, -0.16),
    ncol=2,
    borderaxespad=0,
)
plt.subplots_adjust(bottom=0.23)
plt.show()

Out[25]:

Visualization

Line plot showing rank decreasing over training steps with shaded area, from initial maximum down toward the target budget dashed line. — Temporal evolution of the total effective rank during training. The shaded area is the active parameter budget, which starts at the maximum capacity and decreases following the cubic schedule until reaching the target budget (dashed red line). The rapid initial decline shows aggressive early pruning, while the flatter tail shows gentler pruning as the budget approaches its target.

The cubic pruning schedule is evident in the rank evolution: aggressive rank reduction early in training (when many clearly unimportant components exist and the budget is far from its target) followed by gentler pruning as the budget approaches its target. The shaded area is the total "active capacity" of the adaptation: the area under the curve corresponds to the total computational work performed by the adaptation components. AdaLoRA concentrates this work in the early phases of training where it has the most impact on learning.

Comparing with Standard LoRALink Copied

To appreciate AdaLoRA's benefits, let's compare the parameter efficiency at equivalent performance levels. The comparison highlights the qualitative difference in how parameters are allocated.

In[26]:

Code

from peft import LoraConfig

# Equivalent LoRA configuration
# To match AdaLoRA's average target rank, we'd use r=4
lora_config = LoraConfig(
    r=4,
    lora_alpha=16,
    target_modules=["q_lin", "v_lin", "k_lin", "out_lin"],
    task_type=TaskType.SEQ_CLS,
)

# Count parameters for comparison
lora_model = AutoModelForSequenceClassification.from_pretrained(
    model_name, num_labels=2
)
lora_peft_model = get_peft_model(lora_model, lora_config)
lora_trainable, _ = count_parameters(lora_peft_model)

# Calculate AdaLoRA effective parameters
# Each rank unit has: 1 column P (d), 1 row Q (d), 1 singular value (1)
hidden_size = model.config.dim
adalora_effective_params = sum(final_ranks.values()) * (2 * hidden_size + 1)

Out[27]:

Console

Parameter comparison (target rank 4):
  Standard LoRA: 739,586 trainable parameters
  AdaLoRA (final): 147,552 effective parameters (approximate)

Key difference: AdaLoRA allocates these parameters adaptively
across layers based on importance, rather than uniformly.

The parameter counts are similar in total, but AdaLoRA's adaptive allocation concentrates parameters where they matter most. Standard LoRA's uniform allocation means that less important layers receive the same rank as necessary ones, effectively wasting parameters on layers that need less adaptation while leaving necessary layers underserved.

Out[28]:

Visualization

Bar chart showing uniform rank assigned to each layer in standard LoRA, with a horizontal red dashed line at the constant rank value. — Standard LoRA applying a uniform rank to every layer, with the fixed rank shown as a red dashed reference line. Every adapted layer receives identical modeling capacity regardless of its actual importance to the task.

Bar chart showing variable rank per layer in AdaLoRA after adaptive pruning, with a horizontal red dashed line at the average rank. — AdaLoRA's adaptive rank allocation across layers after training. The varying bar heights reflect the data-driven discovery of which layers need more adaptation capacity. The red dashed line shows the average rank, revealing how individual allocations depart from uniformity.

The visual contrast makes AdaLoRA's core advantage concrete. Standard LoRA's flat allocation treats every layer identically, regardless of whether the task requires large adaptation in that layer or almost none. AdaLoRA's allocation tells a differentiated story: some layers receive substantially more capacity, others receive far less. This shows the actual demands of the fine-tuning task.