Gradient Checkpointing, Frozen Embeddings, and LoRA Adapters

Fine-tuning large models like LLAMA on specific tasks or domains can be a resource-intensive process. To address these challenges, techniques such as LoRA (Low-Rank Adaptation), frozen embeddings with hooks, and gradient checkpointing can be applied. This article provides a detailed exploration of these techniques, how they interact, and how to optimize them when combined. Whether you’re dealing with large models on limited resources or trying to fine-tune specific model parts, this guide offers practical insights for efficient training.

Background on LoRA and Frozen Embeddings Why Use LoRA?Freezing Embeddings Using Frozen Embeddings with Forward Hooks Implementing Forward Hooks to Enable Gradient Flow Gradient Checkpointing Explained How Gradient Checkpointing Works When to Use Gradient Checkpointing Challenges with Gradient Checkpointing and Forward Hooks Interaction Between Hooks and Checkpointing Practical Impact Practical Implementation Strategies 1. Avoid Relying on Hooks with Checkpointing 2. Use Functional API for Embeddings 3. Custom Checkpointing Function 4. Verify Correct Gradient Flow Key Takeaways

Background on LoRA and Frozen Embeddings

When fine-tuning large models, we often want to avoid updating the entire model to save memory and computation time. LoRA (Low-Rank Adaptation) is a technique that inserts low-rank matrices into specific parts of the model (usually linear layers) to learn task-specific adjustments without modifying the original weights.

Why Use LoRA?

Low Memory Usage: LoRA introduces only a few additional parameters, leaving the main model weights frozen.

Efficient Training: Fine-tuning only these low-rank adapters significantly reduces the computation required, making it feasible to use large models in resource-constrained settings.

Freezing Embeddings

In many cases, embedding layers (which represent the model's initial input layer) are also frozen. By freezing embeddings, we avoid modifying the model's representation of its input space, which is often already well-trained. However, in some fine-tuning tasks, it can be beneficial to allow gradient flow through these embeddings without actually updating them.

Using Frozen Embeddings with Forward Hooks

When we freeze embeddings, their parameters have requires_grad=False, which stops the model from updating them during backpropagation. But in cases like LoRA, where we add task-specific adapters, it's crucial to allow gradient flow through the embeddings to the adapters for them to update correctly. This is where forward hooks come into play.

Implementing Forward Hooks to Enable Gradient Flow

Consider the following Python snippet that demonstrates how a forward hook can enable gradient computation on the outputs of a frozen embedding layer:


def enable_input_require_grads(self):
    """
    Enables gradients for the input embeddings.
    Useful for fine-tuning adapters like LoRA without updating the main model's weights.
    """
    def make_inputs_require_grads(module, input, output):
        output.requires_grad_(True)

    self._require_grads_hook = self.get_input_embeddings().register_forward_hook(make_inputs_require_grads)

Here’s how this setup works:

The Hook: Sets requires_grad_(True) on the outputs of the frozen embedding layer, allowing gradients to flow through these outputs even though the embeddings themselves remain frozen.

Effect: This modification lets adapters like LoRA receive gradients from the embedding layer’s outputs, ensuring that they learn effectively during backpropagation without altering the embeddings.

Gradient Checkpointing Explained

Gradient checkpointing is a memory optimization strategy that trades extra computation for reduced memory usage, which is beneficial when training large models with limited resources.

How Gradient Checkpointing Works

Standard Forward Pass: Normally, all activations are stored during the forward pass for use in backpropagation.

With Checkpointing: Instead of storing all activations, only specific “checkpoints” are stored. During backpropagation, the non-stored activations are recomputed on-the-fly as needed.

This approach saves memory since fewer activations are stored, but increases computation time due to the need for recomputation.

When to Use Gradient Checkpointing

Gradient checkpointing is especially helpful when working with models that have a large number of layers or a high memory footprint. By using checkpointing, you can train models that would otherwise exceed memory limits, although with the trade-off of increased computation time.

Challenges with Gradient Checkpointing and Forward Hooks

When using gradient checkpointing along with frozen embeddings and forward hooks, issues can arise in backpropagation, specifically due to inconsistent activation of hooks during recomputation.

Interaction Between Hooks and Checkpointing

In frameworks like PyTorch, the forward hook might not be called during the recomputation phase of checkpointing. If the hook does not execute, the embeddings’ outputs may not have requires_grad=True during recomputation, preventing gradient flow to the adapters and causing incorrect or zero gradients for LoRA layers.

Practical Impact

Without the hook being applied during recomputation:

Adapter Layers Fail to Train: LoRA adapters or other components relying on gradients from the embedding outputs will not receive proper updates.

Loss and Gradient Issues: Model training becomes ineffective as gradient flow is interrupted.

Practical Implementation Strategies

1. Avoid Relying on Hooks with Checkpointing

Instead of using hooks, directly set requires_grad=True on embedding outputs within the model's forward method. This approach guarantees that requires_grad will be applied correctly during both the initial forward pass and any recomputation steps.


class ModifiedModel(nn.Module):
    def __init__(self, ...):
        super().__init__()
        self.embeddings = EmbeddingLayer(...)
        for param in self.embeddings.parameters():
            param.requires_grad = False

    def forward(self, input_ids, ...):
        embeddings = self.embeddings(input_ids)
        embeddings.requires_grad_(True)  # Directly set requires_grad on embeddings' outputs
        output = checkpoint(self.compute_rest, embeddings, ...)
        return output

    def compute_rest(self, embeddings, ...):
        # Continue forward computation with adapters and other layers
        ...

2. Use Functional API for Embeddings

Another effective method is to compute the embeddings in a standalone function that sets requires_grad=True, ensuring this setting is applied both initially and during recomputation.


def compute_embeddings(input_ids):
    embeddings = embedding_layer(input_ids)
    embeddings.requires_grad_(True)
    return embeddings

# In the model's forward method
embeddings = checkpoint(compute_embeddings, input_ids)

3. Custom Checkpointing Function

If you need more control over checkpointing behavior, create a custom checkpointing function to guarantee correct gradient settings across both forward and backward passes.


from torch.utils.checkpoint import CheckpointFunction

class CustomCheckpointFunction(CheckpointFunction):
    @staticmethod
    def forward(ctx, run_function, *args):
        ctx.run_function = run_function
        ctx.save_for_backward(*args)
        with torch.no_grad():
            outputs = run_function(*args)
        return outputs

    @staticmethod
    def backward(ctx, *grad_outputs):
        args = ctx.saved_tensors
        with torch.enable_grad():
            outputs = ctx.run_function(*args)
        torch.autograd.backward(outputs, grad_outputs)
        return (None, *args)

# Usage
embeddings = CustomCheckpointFunction.apply(compute_embeddings, input_ids)

4. Verify Correct Gradient Flow

After implementing gradient checkpointing, inspect the gradients to ensure correct flow through the model:


loss.backward()
for name, param in model.named_parameters():
    if param.grad is None:
        print(f"No gradient for {name}")
    else:
        print(f"Gradient for {name}: {param.grad.norm()}")

Check that:

Frozen embedding parameters should have grad=None.

Adapter parameters have non-zero gradients if they are receiving the correct backpropagation signal.

Key Takeaways

Using gradient checkpointing with frozen embeddings and adapters like LoRA requires a precise setup to ensure gradient flow is maintained. Here are the main considerations:

Directly Set requires_grad on Embeddings: Avoid using forward hooks with checkpointing. Instead, set requires_grad=True on embeddings directly within the forward method to ensure consistency during recomputation.

Trade-Offs of Gradient Checkpointing: While checkpointing reduces memory usage, it increases computation time, so balance your memory and time requirements accordingly.

Monitor Gradient Flow: Always verify that gradients are flowing correctly, especially for models with frozen layers and adapters.

By following these steps, you can optimize the training of large models with LoRA adapters, gradient checkpointing, and frozen embeddings, achieving efficient, task-specific fine-tuning without unnecessary computational overhead. This setup enables fine-tuning of complex models like LLAMA even with limited resources, making it a powerful tool for practical deep learning applications.