Full Fine-tuning › Implementation Deep Dive

Let's walk through a complete fine-tuning implementation, examining each component in detail. Understanding the underlying mechanics of the training loop gives you the ability to debug problems, implement custom modifications, and reason about what is happening during training. High-level APIs are convenient, but knowing what they are doing under the hood is needed for serious practitioners.

The training loop has a deceptively simple structure: iterate over batches, compute loss, compute gradients, update parameters. But each of these steps involves important decisions that affect training stability and final performance. The details of gradient clipping, optimizer parameter grouping, and evaluation frequency all matter for achieving good results.

While high-level APIs handle many details, understanding the underlying training loop gives insight into what is happening during fine-tuning. The key components are the optimizer, the learning rate scheduler, and the parameter grouping strategy.

Notice the parameter grouping: we apply weight decay to most parameters but exclude biases and layer normalization weights. This follows the original BERT fine-tuning recipe and prevents these parameters from being pushed toward zero unnecessarily. The no_decay list identifies parameters that should remain free from regularization pressure. Biases are typically small in magnitude and contribute little to overfitting. Layer normalization parameters are scaling and shifting factors that should be free to take whatever values the model needs for proper normalization.

The AdamW optimizer is the standard choice for fine-tuning transformers. It implements decoupled weight decay, maintaining the beneficial adaptive learning rate behavior of Adam while correctly separating the regularization effect of weight decay. Using standard Adam with L2 regularization added to the loss produces suboptimal results for transformers because the adaptive learning rates interact with the gradient in a way that diminishes the regularization effect. AdamW avoids this issue by applying weight decay directly to parameters after the gradient update.

Each training step involves a forward pass, loss computation, backward pass, and parameter update. The implementation below shows the needed operations, including gradient clipping, which is a necessary stability measure for fine-tuning.

Gradient clipping prevents any single batch from causing excessively large updates. The maximum gradient norm of 1.0 is a common choice for fine-tuning, constraining the total magnitude of gradients across all parameters to prevent training instability. The clip_grad_norm_ function computes the global gradient norm (the L2 norm across all parameter gradients), and if it exceeds the specified maximum, scales all gradients down proportionally so that the total norm equals exactly 1.0. This operation does not change the direction of the update, only its magnitude.

Gradient clipping is particularly important in the early epochs of fine-tuning, when the classification head is still learning rapidly and may produce large gradients. Without clipping, a single bad batch could push the backbone parameters far from their pre-trained values. The clipping threshold of 1.0 was chosen empirically and has become a community standard for BERT-style fine-tuning. For larger models, you may occasionally need to adjust this threshold if you observe that gradients are being clipped too aggressively.

Evaluating during fine-tuning helps monitor for overfitting and select the best checkpoint. Evaluation should be done frequently enough to detect overfitting early, but not so frequently that it materially slows down training. Evaluating once per epoch is usually sufficient for fine-tuning runs of 2-5 epochs.

The torch.no_grad() context manager is necessary during evaluation: it disables gradient computation, reducing memory usage and speeding up the forward pass. During evaluation, we don't need gradients because we are not updating parameters. Forgetting to use torch.no_grad() during evaluation can cause memory exhaustion and materially slow down your training loop, especially when evaluating on large validation sets.

The model.eval() call is equally important: it switches certain layers from their training behavior to their evaluation behavior. Dropout layers, for example, randomly zero out activations during training to prevent co-adaptation of neurons, but should be disabled during evaluation to produce consistent predictions. Batch normalization layers similarly behave differently during training and evaluation. Failing to call model.eval() before evaluation will produce inconsistent results and may underestimate the model's true performance.

Putting it all together into a complete training run:

The key parameters for full fine-tuning are:

• learning_rate: Controls how much weights change per update. Typically 1e-5 to 5e-5 for transformer fine-tuning.
• num_train_epochs: Number of passes through the training data. Usually 2-5 epochs for fine-tuning.
• per_device_train_batch_size: Samples per batch on each device. Common values are 16-32 for BERT-sized models.
• warmup_ratio: Fraction of training for learning rate warmup. Typically 0.06-0.1 (6-10%).
• weight_decay: L2 regularization strength. Common value is 0.01.
• gradient_accumulation_steps: Number of batches to accumulate before updating. Use to simulate larger batch sizes.

The training progresses through multiple epochs, with both training and validation metrics tracked at each step. The best model checkpoint is saved based on validation accuracy, which helps prevent overfitting by selecting the model state that generalizes best to unseen data. The final validation accuracy shows how well the fine-tuned model performs on held-out examples.

Tracking loss and accuracy across epochs helps diagnose training issues. The training curves encode a wealth of information about what is happening during optimization, and learning to read them fluently is one of the most useful skills for a practitioner.