Fine-tuning Learning Rates › Complete Implementation
Michael BrenndoerferNovember 27, 20259 min read
Language AI HandbookAll chapters Part 34 of 67Fine-tuning Fundamentals
Fine-tuning Learning Rates › Complete ImplementationOne section · Sections in this chapter8 / 12 ↓
Sections in this chapter12
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 combine these techniques into a complete fine-tuning setup. We will create a training configuration with layer-wise learning rates and a warmup-plus-cosine schedule. This implementation shows how the individual components work together in practice rather than as a collection of independent tricks.
In[17]:
Codeimport math
from torch.optim import AdamW
from torch.optim.lr_scheduler import LambdaLR
class FineTuningConfig:
"""Configuration for fine-tuning with advanced learning rate strategies."""
def __init__(
self,
base_lr=2e-5,
llrd_decay=0.95,
weight_decay=0.01,
warmup_ratio=0.1,
num_training_steps=10000,
schedule="cosine",
):
self.base_lr = base_lr
self.llrd_decay = llrd_decay
self.weight_decay = weight_decay
self.warmup_ratio = warmup_ratio
self.num_training_steps = num_training_steps
self.warmup_steps = int(warmup_ratio * num_training_steps)
self.schedule = schedule
def create_optimizer_and_scheduler(model, task_head, config):
"""
Create optimizer with LLRD and scheduler with warmup + decay.
Args:
model: Pre-trained transformer model
task_head: Task-specific classification head
config: FineTuningConfig instance
Returns:
Tuple of (optimizer, scheduler)
"""
# Get encoder layers for LLRD
if hasattr(model, "encoder"):
encoder_layers = list(model.encoder.layer)
elif hasattr(model, "transformer"):
encoder_layers = list(model.transformer.h)
else:
encoder_layers = []
num_layers = len(encoder_layers)
parameter_groups = []
# Embeddings - lowest learning rate
embedding_params = []
for name, param in model.named_parameters():
if "embed" in name.lower() and param.requires_grad:
embedding_params.append(param)
if embedding_params:
embedding_lr = config.base_lr * (config.llrd_decay**num_layers)
parameter_groups.append(
{
"params": embedding_params,
"lr": embedding_lr,
"weight_decay": config.weight_decay,
}
)
# Encoder layers with LLRD
processed_params = set(id(p) for p in embedding_params)
for layer_idx, layer in enumerate(encoder_layers):
layer_lr = config.base_lr * (
config.llrd_decay ** (num_layers - 1 - layer_idx)
)
layer_params = [
p
for p in layer.parameters()
if p.requires_grad and id(p) not in processed_params
]
if layer_params:
parameter_groups.append(
{
"params": layer_params,
"lr": layer_lr,
"weight_decay": config.weight_decay,
}
)
processed_params.update(id(p) for p in layer_params)
# Remaining model parameters (pooler, etc.)
remaining_params = [
p
for name, p in model.named_parameters()
if p.requires_grad and id(p) not in processed_params
]
if remaining_params:
parameter_groups.append(
{
"params": remaining_params,
"lr": config.base_lr,
"weight_decay": config.weight_decay,
}
)
# Task head - highest learning rate (often 10x base)
task_head_lr = config.base_lr * 10
parameter_groups.append(
{
"params": list(task_head.parameters()),
"lr": task_head_lr,
"weight_decay": config.weight_decay,
}
)
# Create optimizer
optimizer = AdamW(parameter_groups, betas=(0.9, 0.999), eps=1e-8)
# Create scheduler
def lr_lambda(current_step):
if current_step < config.warmup_steps:
return float(current_step) / float(max(1, config.warmup_steps))
if config.schedule == "linear":
progress = float(current_step - config.warmup_steps) / float(
max(1, config.num_training_steps - config.warmup_steps)
)
return max(0.0, 1.0 - progress)
elif config.schedule == "cosine":
progress = float(current_step - config.warmup_steps) / float(
max(1, config.num_training_steps - config.warmup_steps)
)
return max(0.0, 0.5 * (1.0 + math.cos(math.pi * progress)))
return 1.0
scheduler = LambdaLR(optimizer, lr_lambda)
return optimizer, schedulerLet's see this in action with a simple classification task:
In[18]:
Codeimport torch
from transformers import BertConfig, BertModel
# Create a small model for demonstration (avoids downloading weights)
bert_config = BertConfig(
hidden_size=128,
num_hidden_layers=6,
num_attention_heads=2,
intermediate_size=512,
)
model = BertModel(bert_config)
# Create a simple task head
class ClassificationHead(torch.nn.Module):
def __init__(self, hidden_size, num_classes):
super().__init__()
self.dropout = torch.nn.Dropout(0.1)
self.classifier = torch.nn.Linear(hidden_size, num_classes)
def forward(self, pooled_output):
x = self.dropout(pooled_output)
return self.classifier(x)
task_head = ClassificationHead(model.config.hidden_size, num_classes=2)
# Create fine-tuning configuration
config = FineTuningConfig(
base_lr=2e-5,
llrd_decay=0.95,
warmup_ratio=0.1,
num_training_steps=1000,
schedule="cosine",
)
# Create optimizer and scheduler
optimizer, scheduler = create_optimizer_and_scheduler(model, task_head, config)
# Calculate parameter stats for display
param_group_stats = []
for i, group in enumerate(optimizer.param_groups):
num_params = sum(p.numel() for p in group["params"])
param_group_stats.append(
{"group_idx": i, "lr": group["lr"], "num_params": num_params}
)Out[19]:
ConsoleFine-tuning Configuration ================================================== Base learning rate: 2.00e-05 LLRD decay factor: 0.95 Warmup steps: 100 Total training steps: 1,000 Schedule: cosine Learning Rates per Parameter Group: -------------------------------------------------- Group 0: LR = 0.00e+00, Params = 3,972,864 Group 1: LR = 0.00e+00, Params = 198,272 Group 2: LR = 0.00e+00, Params = 198,272 Group 3: LR = 0.00e+00, Params = 198,272 Group 4: LR = 0.00e+00, Params = 198,272 Group 5: LR = 0.00e+00, Params = 198,272 Group 6: LR = 0.00e+00, Params = 198,272 Group 7: LR = 0.00e+00, Params = 16,512 Group 8: LR = 0.00e+00, Params = 258
In this configuration, the task head has a higher learning rate () than the embeddings. This hierarchy ensures that the randomly initialized task head can adapt quickly while the pre-trained backbone layers are updated more cautiously. Notice that the absolute difference between groups is largest at the top of the network and compresses toward the bottom. This shows the geometric nature of LLRD.
Now let's simulate a training run and visualize how learning rates evolve:
In[20]:
Codeimport matplotlib.pyplot as plt
import numpy as np
use_book_style()
plt.rcParams["figure.figsize"] = (6.0, 4.0)
# Reset optimizer state and track learning rates over training
optimizer, scheduler = create_optimizer_and_scheduler(model, task_head, config)
# Track LRs for selected groups
lr_history = {f"Group {i}": [] for i in range(len(optimizer.param_groups))}
for step in range(config.num_training_steps):
# Record current learning rates
for i, group in enumerate(optimizer.param_groups):
lr_history[f"Group {i}"].append(group["lr"])
# Step the scheduler (in real training, this happens after optimizer.step())
scheduler.step()
# Plot selected groups for clarity
fig, ax = plt.subplots()
steps = np.arange(config.num_training_steps)
# Plot embedding layer (lowest LR)
ax.plot(
steps,
[lr * 1e5 for lr in lr_history["Group 0"]],
label="Embeddings",
linewidth=1,
alpha=0.8,
)
# Plot a middle transformer layer
mid_group = len(optimizer.param_groups) // 2
ax.plot(
steps,
[lr * 1e5 for lr in lr_history[f"Group {mid_group}"]],
label="Middle layers",
linewidth=1,
alpha=0.8,
)
# Plot top transformer layer
ax.plot(
steps,
[lr * 1e5 for lr in lr_history[f"Group {len(optimizer.param_groups) - 2}"]],
label="Top layers",
linewidth=1,
alpha=0.8,
)
# Plot task head (highest LR)
ax.plot(
steps,
[lr * 1e5 for lr in lr_history[f"Group {len(optimizer.param_groups) - 1}"]],
label="Task head",
linewidth=1,
alpha=0.8,
)
ax.axvline(
x=config.warmup_steps, color=PALETTE["muted"], linestyle="--", alpha=0.5
)
ax.annotate(
"End warmup",
xy=(config.warmup_steps, 0),
xytext=(config.warmup_steps + 50, 3.0),
fontsize=8,
color=PALETTE["ink"],
bbox=dict(
boxstyle="round,pad=0.3",
facecolor=PALETTE["paper"],
edgecolor="none",
alpha=0.92,
),
zorder=4,
)
ax.set_xlabel("Training Step")
ax.set_ylabel(r"Learning Rate ($\times 10^{-5}$)")
ax.set_title(
"Learning Rate Evolution During Fine-tuning\n(LLRD + Warmup + Cosine Decay)"
)
ax.legend(loc="upper right")
ax.set_xlim(0, config.num_training_steps)
polish_axes(ax)The separate trajectories show the combined effect of our strategies: all layers warm up together, but they peak at different values due to LLRD, and then decay in sync according to the cosine schedule. The task head's learning rate is an order of magnitude larger than the embeddings' rate throughout training. This shows our confidence that the task head needs rapid adaptation while the embeddings should remain relatively stable.
Out[22]:
Visualization
Complete Training Loop Example
Here is how these components integrate into a complete fine-tuning loop. Notice how the optimizer step, gradient clipping, and scheduler step must occur in the correct order: compute gradients, clip them, apply the optimizer update, then advance the scheduler. Reversing any of these steps would break the intended behavior.
In[12]:
Codedef fine_tune_model(
model, task_head, train_dataloader, val_dataloader, config, device="cuda"
):
"""
Complete fine-tuning loop with LLRD, warmup, and decay.
"""
model.to(device)
task_head.to(device)
optimizer, scheduler = create_optimizer_and_scheduler(
model, task_head, config
)
criterion = torch.nn.CrossEntropyLoss()
model.train()
task_head.train()
global_step = 0
best_val_accuracy = 0
for epoch in range(num_epochs):
epoch_loss = 0
for batch in train_dataloader:
input_ids = batch["input_ids"].to(device)
attention_mask = batch["attention_mask"].to(device)
labels = batch["labels"].to(device)
# Forward pass
outputs = model(input_ids, attention_mask=attention_mask)
pooled_output = outputs.last_hidden_state[:, 0] # CLS token
logits = task_head(pooled_output)
loss = criterion(logits, labels)
# Backward pass
optimizer.zero_grad()
loss.backward()
# Gradient clipping (important for stability)
torch.nn.utils.clip_grad_norm_(
list(model.parameters()) + list(task_head.parameters()),
max_norm=1.0,
)
optimizer.step()
scheduler.step()
global_step += 1
epoch_loss += loss.item()
# Log learning rates periodically
if global_step % 100 == 0:
current_lr = scheduler.get_last_lr()[0]
print(
f"Step {global_step}: Loss = {loss.item():.4f}, LR = {current_lr:.2e}"
)
# Validation
val_accuracy = evaluate(model, task_head, val_dataloader, device)
print(f"Epoch {epoch + 1}: Val Accuracy = {val_accuracy:.4f}")
if val_accuracy > best_val_accuracy:
best_val_accuracy = val_accuracy
# Save best model checkpoint
torch.save(
{
"model_state_dict": model.state_dict(),
"task_head_state_dict": task_head.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
},
"best_model.pt",
)
return best_val_accuracy