AI Engineering Primer | From Basics to Advanced

Chapter: Fine-Tuning Llama 3.2 - 1B Model with LoRA from Scratch for Text Summarization

This chapter provides an end-to-end technical guide to parameter-efficient fine-tuning (PEFT) using Low-Rank Adaptation (LoRA) built from scratch in PyTorch. Using Meta’s Llama 3.2 1B base model and the Tiny Stories dataset, we explore how to adapt a large language model for specialized text summarization while freezing $\sim 99\%$ of its underlying pre-trained parameters.

1. The Resource Challenge of Full Fine-Tuning

Large Language Models (LLMs) are trained on massive multilingual corpora, giving base models strong semantic and grammatical comprehension out of the box. However, adapting these base models to specific downstream tasks (such as structured text summarization) requires targeted instruction tuning.

Updating all weights in a modern language model presents significant technical hurdles:

• Compute Exhaustion: Backpropagating gradients through billions of parameters requires massive floating-point operations (FLOPs) and prolonged execution times.
• VRAM Overhead: Full fine-tuning requires memory not just for the model weights, but also for optimizer states (e.g., Adam's momentum and variance vectors), gradients, and intermediate activation tensors.
• Catastrophic Forgetting: Updating all weights risks overwriting the foundational linguistic knowledge acquired during pre-training.

Parameter-Efficient Fine-Tuning (PEFT) via LoRA

Rather than updating the entire weight matrix $W_0$, Parameter-Efficient Fine-Tuning freezes the pre-trained weights and introduces small, trainable auxiliary matrices. LoRA (Low-Rank Adaptation) targets specific linear projections within the transformer block, drastically reducing the active parameter footprint while preserving the foundational knowledge of the base model.

2. Meta Llama 3.2 1B Model Architecture Breakdown

Meta's Llama 3.2 family consists of multilingual generative models available in 1B and 3B parameter sizes, provided in both instruction-tuned and base auto-complete configurations. For this implementation, we utilize the Llama 3.2 1B Base Model.

3. Mathematical Foundations of Low-Rank Adaptation (LoRA)

LoRA hypothesizes that weight updates during adaptation have a low "intrinsic rank."

For a pre-trained linear weight matrix $W_0 \in \mathbb{R}^{d_{in} \times d_{out}}$, the forward pass normally computes:

LoRA decomposes the weight update matrix $\Delta W$ into two lower-rank matrices $A$ and $B$:

Where:

• $A \in \mathbb{R}^{r \times d_{in}}$
• $B \in \mathbb{R}^{d_{out} \times r}$
• $r \ll \min(d_{in}, d_{out})$ represents the rank hyperparameter.

The modified forward pass becomes:

Where $\alpha$ is a constant scaling hyperparameter that balances the weight of the LoRA update against the pre-trained features.

Parameter Reduction Logic

Consider a linear layer with input dimension $d_{in} = 4096$ and output dimension $d_{out} = 4096$:

• Standard Weight Parameters: $4096 \times 4096 = 16,777,216$ parameters.
• LoRA Parameters ($r = 16$):

This parameter decomposition yields over a 99% reduction in trainable parameters for that layer.

Initialization Dynamics

To ensure that $h = W_0 x$ at the very start of fine-tuning (preserving pre-trained behavior prior to step 1):

1. Matrix $A$ is initialized using a Gaussian or Kaiming Uniform distribution: $A \sim \mathcal{U}(-\sqrt{k}, \sqrt{k})$.
2. Matrix $B$ is initialized to all zeros: $B = 0$.
3. Therefore, $\Delta W = B \cdot A = 0 \cdot A = 0$ at initialization.

4. Dataset Preparation & Custom Prompt Engineering

We utilize a structured subset of the Tiny Stories Instruct Dataset, reduced from $\sim 2\text{ million}$ records down to $\sim 500,000$ high-quality story-summary pairs.

Formatting the Prompt Template

To fine-tune a base language model for instruction-like behavior (auto-completing a summary based on a given text), input texts must follow a consistent, deterministic delimiter template:

<|begin_of_text|>Summarize the following story:

[STORY CONTENT]

Summary:
[SUMMARY CONTENT]<|end_of_text|>

Python Implementation: Prompt Formatting

def generate_prompt_template(example):
    """
    Formats raw story and summary pairs into a structured prompt 
    suitable for auto-regressive text summarization fine-tuning.
    """
    formatted_prompt = (
        f"<|begin_of_text|>Summarize the following story:\n\n"
        f"{example['story']}\n\n"
        f"Summary:\n"
        f"{example['summary']}<|end_of_text|>"
    )
    return {"prompt": formatted_prompt}

5. Tokenization, Custom Padding, and Dynamic Batching

Language models process batch inputs in parallel GPU matrix operations. Because raw stories vary in word count, sequences within a batch must be aligned to uniform token lengths using padding tokens.

Special Tokens and Vocabulary Handling

• Base Vocabulary Size: $128,000$
• Special Character Offsets: $256$
• Total Model Vocabulary Size: $128,256$
• Custom Padding Strategy: The model tokenizer's pad_token_id is assigned to token ID 1131 (representing triple dots ...).

import torch
from transformers import AutoTokenizer

# Initialize tokenizer from Hugging Face Hub
model_id = "meta-llama/Llama-3.2-1B"
tokenizer = AutoTokenizer.from_pretrained(model_id)

# Set custom pad token ID (ID 1131 correspond to '...')
tokenizer.pad_token_id = 1131

Dynamic Batch Collate Function

Instead of statically padding the entire dataset to a globally fixed max length (e.g., 2048 tokens), sequence padding is applied dynamically per batch. Each batch is padded only to the length of its longest constituent sequence, significantly cutting down unnecessary computation.

from torch.utils.data import DataLoader

def dynamic_collate_fn(batch, pad_token_id=1131):
    """
    Pads input sequences dynamically to the maximum sequence length 
    present within the local batch.
    """
    input_ids = [torch.tensor(item['tokenized_prompt']) for item in batch]
    
    # Calculate maximum sequence length in the current batch
    max_len = max(len(ids) for ids in input_ids)
    
    padded_input_ids = []
    attention_masks = []
    labels = []
    
    for ids in input_ids:
        pad_len = max_len - len(ids)
        
        # Pad right side of sequences
        padded_seq = torch.cat([ids, torch.tensor([pad_token_id] * pad_len, dtype=torch.long)])
        
        # Create attention mask (1 for real tokens, 0 for padded tokens)
        attention_mask = torch.cat([torch.ones(len(ids)), torch.zeros(pad_len)])
        
        # Target labels match input IDs for auto-regressive language modeling
        label = padded_seq.clone()
        
        padded_input_ids.append(padded_seq)
        attention_masks.append(attention_mask)
        labels.append(label)
        
    return {
        'input_ids': torch.stack(padded_input_ids),
        'attention_mask': torch.stack(attention_masks),
        'labels': torch.stack(labels)
    }

6. Building LoRA from Scratch in PyTorch

We build LoRA from first principles by creating custom nn.Module subclasses:

1. LoRALayer: Encapsulates matrices $A$ and $B$, handling initialization, forward computation, and parameter scaling.
2. LinearWithLoRA: Wraps an existing base nn.Linear layer alongside a parallel LoRALayer.

import torch
import torch.nn as nn
import math

class LoRALayer(nn.Module):
    def __init__(self, in_features: int, out_features: int, rank: int = 16, alpha: float = 32.0):
        super().__init__()
        self.rank = rank
        self.alpha = alpha
        self.scaling = alpha / rank
        
        # Matrix A: Initialized with Kaiming Uniform
        self.lora_A = nn.Parameter(torch.empty(in_features, rank))
        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        
        # Matrix B: Initialized to Zero
        self.lora_B = nn.Parameter(torch.zeros(rank, out_features))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Output shape: (batch_size, sequence_length, out_features)
        lora_update = (x @ self.lora_A) @ self.lora_B
        return lora_update * self.scaling


class LinearWithLoRA(nn.Module):
    def __init__(self, base_layer: nn.Linear, rank: int = 16, alpha: float = 32.0):
        super().__init__()
        # Retain original pre-trained layer
        self.base_layer = base_layer
        
        # Instantiate parallel LoRA sub-module
        self.lora = LoRALayer(
            in_features=base_layer.in_features,
            out_features=base_layer.out_features,
            rank=rank,
            alpha=alpha
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Combine original output with parallel low-rank adaptation output
        return self.base_layer(x) + self.lora(x)

7. Replacing Base Layers & Freezing Model Parameters

With the LinearWithLoRA module defined, the base model architecture must be transformed:

1. Freeze Base Parameters: Set requires_grad = False across all pre-trained parameters in the original model.
2. Iterative Module Replacement: Traversal functions identify all instances of standard linear projections (nn.Linear) and substitute them with LinearWithLoRA.

def freeze_base_model(model: nn.Module):
    """
    Disables gradient calculations for all pre-trained model parameters.
    """
    for param in model.parameters():
        param.requires_grad = False

def replace_linear_with_lora(model: nn.Module, rank: int = 16, alpha: float = 32.0):
    """
    Recursively replaces standard nn.Linear layers with LinearWithLoRA wrappers.
    """
    for name, child in model.named_children():
        if isinstance(child, nn.Linear):
            # Substitute linear layer with custom LoRA wrapper module
            setattr(model, name, LinearWithLoRA(child, rank=rank, alpha=alpha))
        else:
            # Recurse into sub-modules (attention blocks, MLPs, etc.)
            replace_linear_with_lora(child, rank=rank, alpha=alpha)

Verification of Parameter Scaling

from transformers import AutoModelForCausalLM

# 1. Load Pre-Trained Llama 3.2 1B Model
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-1B",
    torch_dtype=torch.bfloat16,
    device_map="cuda"
)

# 2. Freeze Base Parameters
freeze_base_model(model)

# 3. Inject LoRA Layers
replace_linear_with_lora(model, rank=16, alpha=32.0)

# 4. Count Parameters
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)

print(f"Total Parameters: {total_params:,}")
print(f"Trainable Parameters: {trainable_params:,}")
print(f"Trainable Ratio: {100 * trainable_params / total_params:.2f}%")

Empirical Parameter Results:

• Initial Base Model Parameters: $\sim 1,230,000,000$ ($1.23\text{B}$)
• Parameters Frozen: $\sim 1,217,000,000$
• Active Trainable LoRA Parameters: $\sim 13,000,000$ ($13\text{M}$)
• Active Parameter Ratio: $\mathbf{\sim 1.05\%}$

8. The Training Workflow & Optimization Loop

Training runs on an NVIDIA A100 GPU utilizing cross-entropy loss over auto-regressive targets.

Training Loop Implementation

from torch.optim import AdamW
from transformers import get_linear_schedule_with_warmup

# Hyperparameters
BATCH_SIZE = 8
LEARNING_RATE = 1e-4
EPOCHS = 1

train_dataloader = DataLoader(
    dataset, 
    batch_size=BATCH_SIZE, 
    shuffle=True, 
    collate_fn=dynamic_collate_fn
)

# Filter optimizer targets to strictly include parameters with active gradients
trainable_parameters = [p for p in model.parameters() if p.requires_grad]
optimizer = AdamW(trainable_parameters, lr=LEARNING_RATE)

total_steps = len(train_dataloader) * EPOCHS
scheduler = get_linear_schedule_with_warmup(
    optimizer, 
    num_warmup_steps=100, 
    num_training_steps=total_steps
)

# Training Loop Execution
model.train()
for epoch in range(EPOCHS):
    for step, batch in enumerate(train_dataloader):
        input_ids = batch['input_ids'].to("cuda")
        attention_mask = batch['attention_mask'].to("cuda")
        labels = batch['labels'].to("cuda")
        
        optimizer.zero_grad()
        
        outputs = model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            labels=labels
        )
        
        loss = outputs.loss
        loss.backward()
        
        optimizer.step()
        scheduler.step()
        
        if step % 100 == 0:
            print(f"Epoch: {epoch} | Step: {step}/{total_steps} | Loss: {loss.item():.4f}")

9. Evaluation & Summarization Inference Walkthrough

Once training reaches convergence (typically stabilizing between a Cross-Entropy Loss of $1.5$ to $1.6$), we evaluate the model by passing an un-summarized story prompt and generating tokens auto-regressively.

Story Input Prompt Example

<|begin_of_text|>Summarize the following story:

Once upon a time, there was a little girl named Lily. She loved to play outside in the park. One day, Lily saw a big swing, and she wanted to try it. But she was too scared to swing high. Lily's mom said they had to go home soon... [Story Continues]

Summary:

Generation Function

def generate_summary(prompt_text, model, tokenizer, max_new_tokens=50):
    model.eval()
    inputs = tokenizer(prompt_text, return_tensors="pt").to("cuda")
    
    with torch.no_grad():
        output_tokens = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
            do_sample=False  # Deterministic greedy decoding
        )
        
    generated_text = tokenizer.decode(output_tokens[0], skip_special_tokens=True)
    return generated_text

Generated Inference Output

Base Model (Pre-Fine-Tuning): Auto-completes recursively with coherent but unconstrained stream-of-consciousness text (e.g., "I am doing well, I hope you are doing well too...").

Fine-Tuned Model (Post-LoRA Adaptation): Synthesizes input narratives cleanly into a concise summary: "Lily visits the park, overcomes her fear of swinging high, and enjoys playing before returning home with her mother."

Summary and Key Takeaways

• Parameter-Efficient Fine-Tuning (PEFT): Freezing the base model parameters avoids catastrophic forgetting and slashes computational and memory demands.
• Low-Rank Adaptation Mechanics: Decomposition of updates into low-rank matrices ($A \in \mathbb{R}^{d_{in} \times r}$ and $B \in \mathbb{R}^{r \times d_{out}}$) targets the intrinsic rank of weight adaptation.
• Extreme Parameter Reduction: Injecting LoRA layers with rank $r=16$ reduces total trainable parameters from $\sim 1.2\text{B}$ down to $\sim 13\text{M}$ ($\sim 99\%$ reduction).
• Zero-State Matrix Initialization: Matrix $A$ is initialized with Kaiming Uniform, while Matrix $B$ is set to zero, guaranteeing that $B \cdot A = 0$ at step zero to preserve pre-trained behavior initially.
• Custom Matrix Scaling Factor: Matrix updates are scaled by $\frac{\alpha}{r}$ to adjust the impact of adaptation steps without altering the learning rate.
• Dynamic Batching Efficiencies: Constructing batch-level dynamic collate routines avoids global static padding, cutting unnecessary GPU computation.
• End-to-End Task Adaptation: Custom LoRA integration adapts an unconstrained base model into a focused, production-ready text summarization system.