Skip to content
Ashish.
All posts
Diagram illustrating the relationship between learning rate, batch size, and epochs in LLM fine-tuning.
6 min readDevelopmentML engineersFeatured#llm#fine-tuning#machine-learning#hyperparameters#deep-learning#nlp

Hyperparameters That Actually Matter

A practical examination of key hyperparameters like learning rate, batch size, and epochs for ML engineers fine-tuning LLMs.

By Ashish KumarPart 6 of Fine-Tuning LLMs Hands-On #6 of 8

When fine-tuning a Large Language Model (LLM), engineers often treat hyperparameters as dial knobs to be turned until the validation loss flattens. This approach is inefficient and theoretically opaque. To understand what actually matters, we must look at the mechanism of optimization. Fine-tuning is a process of navigating a high-dimensional loss landscape using Stochastic Gradient Descent (SGD) or its variants (like Adam). The hyperparameters that drive performance are not arbitrary settings; they are the physical constraints of your search algorithm.

This article is Part 6 of the Fine-Tuning LLMs Hands-On series.

The three hyperparameters that dictate the trajectory of your model are Learning Rate (how far you step), Batch Size (how noisy your steps are), and Epochs (how long you walk). Everything else—optimizer choice, weight decay, dropout—is secondary to the geometry defined by these three.

The Learning Rate: Step Size and Stability

The learning rate (η\eta) determines the magnitude of the update to the model's weights after each gradient computation. In the context of LLMs, which have billions of parameters, the loss surface is rugged and non-convex. The learning rate is the single most critical parameter for convergence.

If the learning rate is too high, the optimization steps overshoot the minimum, causing the loss to diverge or oscillate wildly. If it is too low, the model converges slowly and may get stuck in sharp local minima that generalize poorly.

The Warmup Mechanism

In pre-trained models, the initial loss is often high because the model's weights are initialized for a different task distribution. Starting with a full learning rate in this region is dangerous; the gradients can be large and unstable, leading to immediate divergence.

This is why warmup is essential. During the warmup phase, the learning rate increases linearly from zero to a peak value over the first NN steps. This allows the model to stabilize its representations before taking large steps. Without warmup, fine-tuning often fails silently, producing garbage output with no obvious error signal.

# Typical warmup schedule implementation
def get_cosine_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps):
    def lr_lambda(current_step):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
        return max(0.0, 0.5 * (1.0 + math.cos(math.pi * progress)))
    return LambdaLR(optimizer, lr_lambda)

Typical Values

For full fine-tuning, learning rates are typically in the range of $1e-5$ to $5e-5$. For parameter-efficient fine-tuning methods like LoRA (Low-Rank Adaptation), where only a small subset of weights are updated, the learning rate can be higher, often $1e-4$ to $2e-4$, because the updates are more stable and localized.

Batch Size: Noise and Generalization

Batch size is often misunderstood as merely a memory constraint. While it does determine GPU memory usage, its deeper role is in regulating the variance of the gradient estimate.

In Stochastic Gradient Descent, we compute the gradient on a mini-batch of data. This gradient is an unbiased estimate of the true gradient over the entire dataset, but it has variance. The size of the batch controls this variance:

  • Small Batch Size: High variance (noisy gradients). This noise acts as a regularizer, helping the model escape sharp minima and find broader, flatter minima that generalize better. However, it leads to slower convergence and requires more careful tuning of the learning rate.
  • Large Batch Size: Low variance (stable gradients). This allows for faster convergence per epoch and better utilization of parallel hardware. However, large batches tend to converge to sharp minima, which can hurt generalization performance.

The Linear Scaling Rule and Its Limits

Traditionally, it was believed that if you increase the batch size by a factor of kk, you should increase the learning rate by the same factor kk to maintain the same effective step size in weight space. This is known as the linear scaling rule.

However, in the context of LLMs, this rule breaks down. Research shows that simply scaling the learning rate with batch size does not always improve performance. Instead, a square-root scaling rule or even no scaling at all is often more effective for fine-tuning. This is because the noise introduced by smaller batches is beneficial for generalization in deep learning.

# Example: Adjusting batch size based on available memory
# If you have limited VRAM, you cannot use large batches.
# Use gradient accumulation to simulate larger batches.
 
accumulation_steps = 4
global_effective_batch_size = actual_batch_size * accumulation_steps
 
# global_effective_batch_size represents the total batch size
# across accumulation steps, distinguishing it from the per-forward-pass batch size.
# The learning rate should remain relatively stable,
# but you may need to adjust slightly based on empirical testing.

Practical Implication

For fine-tuning, a batch size that fits within your GPU memory, combined with gradient accumulation to achieve an effective batch size of 16–64, is often sufficient. There is rarely a need for massive batch sizes (e.g., 1024+) unless you are pre-training from scratch.

Epochs: The Boundary of Overfitting

Epochs represent the number of times the entire training dataset is passed through the model. Unlike learning rate and batch size, which affect the shape of the optimization path, epochs affect the duration of the search.

The Overfitting Threshold

LLMs are powerful enough to memorize small datasets. If you train for too many epochs, the model will start to memorize the specific examples in your training set rather than learning the underlying patterns. This is evident in two ways:

  1. Training Loss Continues to Decrease: The model gets better at predicting the training data.
  2. Validation Loss Increases: The model performs worse on unseen data.

The point where validation loss starts to increase is the early stopping point. Ignoring this and training for a fixed number of epochs (e.g., 3 or 10) is a common mistake. The optimal number of epochs depends entirely on the dataset size and complexity.

Monitoring Validation Loss

Instead of fixing the number of epochs, you should monitor the validation loss. Stop training when the validation loss plateaus or begins to rise. This is often called "early stopping."

# Pseudocode for early stopping
best_val_loss = float('inf')
patience = 2
no_improve_count = 0
 
for epoch in range(max_epochs):
    train_loss = train_epoch(model, train_loader)
    val_loss = validate(model, val_loader)
    
    if val_loss < best_val_loss:
        best_val_loss = val_loss
        save_model(model, "best_checkpoint.pt")
        no_improve_count = 0
    else:
        no_improve_count += 1
        
    if no_improve_count >= patience:
        print("Early stopping triggered")
        break

Dataset Size vs. Epochs

A useful heuristic is that the number of epochs should be inversely proportional to the size of the dataset. If you have a small dataset (e.g., 1000 examples), you might need 3–5 epochs to ensure the model sees the data enough times. If you have a large dataset (e.g., 100,000 examples), 1 epoch might be sufficient because the model is exposed to a vast variety of examples in a single pass.

Conclusion

Fine-tuning an LLM is not about guessing hyperparameters. It is about understanding the mechanics of optimization:

  1. Learning Rate: Controls the step size. Use warmup to stabilize early training, and cosine decay to fine-tune convergence.
  2. Batch Size: Controls gradient noise. Use small-to-medium batches with gradient accumulation to balance stability and generalization.
  3. Epochs: Controls data exposure. Monitor validation loss and stop early to prevent overfitting.

By focusing on these three parameters, you can build a reliable fine-tuning pipeline that is both efficient and effective. The other hyperparameters, such as weight decay and optimizer settings, are fine-tuning knobs that can be adjusted later, but they are secondary to the fundamental geometry defined by LR, batch size, and epochs.

FAQ

What is the ideal learning rate for LoRA fine-tuning? A common starting point is $1e-4$ to $2e-4$, though this varies by model and dataset. Always use a warmup phase.

Does larger batch size always mean better performance? No. While larger batches speed up training, they can lead to poorer generalization (sharp minima). Smaller batches with gradient accumulation often yield better results for fine-tuning.

How do I know when to stop training (early stopping)? Monitor the validation loss. Stop training when the validation loss stops decreasing and begins to plateau or increase, indicating overfitting.

Why is warmup important? Warmup prevents instability in the early stages of training by gradually increasing the learning rate, allowing the model to adapt to the new task distribution without diverging.

Common Pitfalls

  1. Ignoring Warmup: Starting with a high learning rate immediately can cause divergence, especially in pre-trained models.
  2. Using Large Batches Without Scaling: Simply increasing batch size without adjusting the learning rate or understanding its impact on generalization can hurt performance.
  3. Not Monitoring Validation Loss: Training for a fixed number of epochs often leads to overfitting, as the model may memorize the training data.
  4. Assuming Linear Scaling Works: The linear scaling rule for learning rate and batch size is not universally applicable, especially in fine-tuning scenarios.

Practical Takeaways

  • Learning Rate: Start with $1e-5$ to $5e-5$ for full fine-tuning, $1e-4$ to $2e-4$ for LoRA. Always use warmup.
  • Batch Size: Use gradient accumulation to achieve an effective batch size of 16–64. Avoid very large batches unless pre-training.
  • Epochs: Do not fix the number of epochs. Use early stopping based on validation loss.
  • Monitoring: Always track both training and validation loss to detect overfitting early.

Related posts