Fine-Tune LLM on Text Summarization Task from Scratch | PyTorch | Python
This chapter provides a comprehensive, step-by-step guide to fine-tuning a small language model (SLM) for text summarization using PyTorch. We begin with a pre-trained 19.3-million-parameter base model ("Tiny Stories") and transition its behavior from simple auto-completion to instruction-driven summarization. Through custom dataset preprocessing, tokenization, dynamic batch collation, and a PyTorch training loop, you will learn how to adapt language models to specialized downstream tasks.
1. The Paradigm Shift: Auto-Completion vs. Instruction-Following
Base large language models (LLMs) are trained primarily on the objective of causal language modeling: predicting the next token given a sequence of preceding tokens. While this enables them to generate grammatically correct and coherent text, they do not inherently understand instructions or tasks like summarization.
To transform a base model into an assistant capable of executing specific tasks, we must perform Supervised Fine-Tuning (SFT).
Base Model (Auto-complete)
|
| Input: "Once upon a time there was a little girl..."
v
| Output: "...named Lily. She loved to play outside in the park."
Fine-Tuned Model (Instruction-Following)
|
| Input: "Story: [Full Narrative] \n Summary:"
v
| Output: "[Concise high-level summary of the narrative] <|endoftext|>"
During SFT, the model is trained on structured prompt templates containing both the input task (the story) and the target output (the summary). This teaches the model to recognize the structural boundary where its generation should transition from reading the context to producing the summary.
2. Base Model Architecture & Specifications
The model used in this guide is a reproduction of the architecture introduced in the research paper "Tiny Stories: How Small Language Models Can Be and Still Speak Coherent English?"
Real-Life Analogy: The Children's Book Vocabulary
The core thesis of the Tiny Stories paper is that language models do not need billions of parameters to speak coherent English if their training domain is highly constrained. By training on a dataset containing simple vocabulary, basic grammatical structures, and themes understandable to a 5- to 7-year-old child, a 19.3M parameter model can achieve remarkable linguistic coherence.
The structural specifications of our base model are detailed below:
3. Base Model Inference Walkthrough
Before fine-tuning, we load the base model to observe its default auto-complete behavior. The model is hosted on Hugging Face as 997/TinyStories-19M (or your respective user path).
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
# 1. Load Tokenizer and Base Model
model_id = "997/TinyStories-19M"
tokenizer = AutoTokenizer.from_pretrained(model_id)
base_model = AutoModelForCausalLM.from_pretrained(model_id)
# Set model to evaluation mode
base_model.eval()
# 2. Prepare the Input Prompt
prompt = "Once upon a time there was a little girl"
inputs = tokenizer(prompt, return_tensors="pt")
# 3. Stream Generation to Console
streamer = TextStreamer(tokenizer, skip_prompt=False)
print("--- Base Model Generation ---")
with torch.no_grad():
base_model.generate(
**inputs,
max_new_tokens=150,
temperature=0.7,
top_k=50,
pad_token_id=tokenizer.eos_token_id,
streamer=streamer
)
Expected Output Behavior
The base model completes the story naturally, generating a narrative about a girl named Lily playing in a park, climbing trees, and watching ducks. It does not summarize; it simply continues the story because it has only been trained on the next-token prediction task.
4. The Fine-Tuning Data Format & Autoregressive Alignment
To train the model to summarize, we must structure our training data using a consistent prompt template.
The Prompt Template
Story: <STORY_TEXT>
Summary: <SUMMARY_TEXT> <|endoftext|>
During training, the model processes the entire sequence. However, because it is an autoregressive causal language model, every token in the sequence acts as a target for the preceding token.
Mathematical Formulation of Causal Loss
For a sequence of tokens , the model minimizes the negative log-likelihood of the target tokens:
During fine-tuning, the model learns the conditional probability of generating the summary tokens given the story context tokens :
5. Dataset Ingestion & Custom Preprocessing Pipeline
We use the Tiny Stories Instruct dataset from Hugging Face. The raw dataset contains unstructured text lines where stories and summaries are mixed together without explicit boundaries. We must parse this raw text into distinct story-summary pairs.
Preprocessing Code
The raw dataset contains millions of lines, but because our model is already pre-trained, we only need a fraction of this data (e.g., parsing 100,000 raw lines to extract roughly 11,000 clean story-summary pairs).
from datasets import load_dataset
# Load raw dataset
raw_dataset = load_dataset("roneneldan/TinyStoriesInstruct")
def process_story_summary_pairs(raw_data, limit=100000):
"""
Parses unstructured text lines into structured story-summary pairs.
"""
pairs = []
current_story = []
current_summary = []
in_story = True
for i, line in enumerate(raw_data['train']['text'][:limit]):
line_str = line.strip()
if not line_str:
continue
# Custom parsing logic based on dataset markers
if "Summary:" in line_str or "summary:" in line_str:
in_story = False
current_summary.append(line_str.replace("Summary:", "").replace("summary:", "").strip())
elif in_story:
current_story.append(line_str)
else:
# Once a new story block starts, save the previous pair
story_txt = " ".join(current_story).strip()
summary_txt = " ".join(current_summary).strip()
if story_txt and summary_txt:
pairs.append({"story": story_txt, "summary": summary_txt})
current_story = [line_str]
current_summary = []
in_story = True
if i % 10000 == 0:
print(f"Processed {i} lines ({(i/limit)*100:.1f}% complete)")
return pairs
processed_pairs = process_story_summary_pairs(raw_dataset, limit=100000)
print(f"Extracted {len(processed_pairs)} clean story-summary pairs.")
6. Tokenization & Prompt Templating Strategy
Once we have extracted the story-summary pairs, we format them using our template and tokenize them. To optimize performance, we run the tokenization process in parallel.
import os
from datasets import Dataset
def tokenize_pair(example):
"""
Formats the story and summary into the target prompt template and tokenizes.
"""
prompt_text = f"Story: {example['story']}\nSummary: {example['summary']}{tokenizer.eos_token}"
tokenized = tokenizer(
prompt_text,
truncation=True,
max_length=512, # Context window limit
add_special_tokens=True
)
return {
"input_ids": tokenized["input_ids"],
"attention_mask": tokenized["attention_mask"]
}
# Convert processed list to Hugging Face Dataset object
hf_dataset = Dataset.from_list(processed_pairs)
# Tokenize dataset in parallel
tokenized_dataset = hf_dataset.map(
tokenize_pair,
batched=False,
num_proc=os.cpu_count(),
remove_columns=["story", "summary"]
)
# Save tokenized dataset to disk to avoid re-running preprocessing
tokenized_dataset.save_to_disk("./tokenized_stories_instruct")
print("Tokenized dataset saved successfully.")
Note on Saving Checkpoints: Saving the tokenized dataset to disk is a standard production practice. If your training run is interrupted, you can reload the tokenized data instantly using
Dataset.load_from_disk(), bypassing the expensive parsing and tokenization steps.
7. PyTorch Dataset & Dynamic Collation with Custom Padding
When training in batches, all sequences within a single batch must have the same length. We use a custom padding token and a dynamic collator to pad shorter sequences to match the longest sequence in the batch.
Special Tokens Configuration
We define a distinct pad token to ensure the model does not confuse padding with the end of a sentence.
- End of Sentence (EOS) Token:
<|endoftext|>(Tells the model when to stop generating). - Pad Token:
...(ID:986) (Used to fill empty space in batches; ignored during loss calculation).
from torch.utils.data import Dataset as PyTorchDataset, DataLoader
import torch
# Configure Pad Token
tokenizer.pad_token = "..."
pad_token_id = tokenizer.convert_tokens_to_ids(tokenizer.pad_token)
print(f"Pad Token: {tokenizer.pad_token} | ID: {pad_token_id}")
class TinyStoriesDataset(PyTorchDataset):
def __init__(self, tokenized_ds, device):
self.dataset = tokenized_ds
self.device = device
def __len__(self):
return len(self.dataset)
def __getitem__(self, idx):
item = self.dataset[idx]
return {
"input_ids": torch.tensor(item["input_ids"], dtype=torch.long, device=self.device),
"attention_mask": torch.tensor(item["attention_mask"], dtype=torch.long, device=self.device)
}
def collate_fn(batch):
"""
Dynamically pads sequences in the batch to the maximum sequence length.
"""
input_ids = [item["input_ids"] for item in batch]
attention_masks = [item["attention_mask"] for item in batch]
# Pad input IDs and attention masks
padded_inputs = torch.nn.utils.rnn.pad_sequence(
input_ids,
batch_first=True,
padding_value=pad_token_id
)
padded_masks = torch.nn.utils.rnn.pad_sequence(
attention_masks,
batch_first=True,
padding_value=0 # 0 tells attention layers to ignore these tokens
)
# In causal language modeling, labels are shifted versions of inputs
labels = padded_inputs.clone()
labels[labels == pad_token_id] = -100 # -100 is ignored by PyTorch CrossEntropyLoss
return {
"input_ids": padded_inputs,
"attention_mask": padded_masks,
"labels": labels
}
8. The Fine-Tuning Training Loop
We set up a native PyTorch training loop using the Adam optimizer and a learning rate scheduler. The training runs on an NVIDIA A100 GPU to ensure fast convergence.
from torch.optim import AdamW
from transformers import get_scheduler
from tqdm.auto import tqdm
# 1. Setup Device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
# 2. Move Model to Device
base_model.to(device)
# 3. Instantiate Dataset and DataLoader
train_dataset = TinyStoriesDataset(tokenized_dataset, device)
train_dataloader = DataLoader(train_dataset, batch_size=16, shuffle=True, collate_fn=collate_fn)
# 4. Optimizer and Scheduler Setup
optimizer = AdamW(base_model.parameters(), lr=5e-5)
num_epochs = 1
num_training_steps = num_epochs * len(train_dataloader)
lr_scheduler = get_scheduler(
"linear",
optimizer=optimizer,
num_warmup_steps=100,
num_training_steps=num_training_steps
)
# 5. Training Loop
base_model.train()
progress_bar = tqdm(range(num_training_steps))
print("Starting Fine-Tuning...")
for epoch in range(num_epochs):
for step, batch in enumerate(train_dataloader):
# Skip sequences that exceed the context window limit
if batch["input_ids"].size(1) >= 512:
continue
outputs = base_model(
input_ids=batch["input_ids"],
attention_mask=batch["attention_mask"],
labels=batch["labels"]
)
loss = outputs.loss
loss.backward()
optimizer.step()
lr_scheduler.step()
optimizer.zero_grad()
progress_bar.update(1)
progress_bar.set_postfix({"loss": loss.item()})
# Optional: Save intermediate checkpoints
if step % 500 == 0:
torch.save(base_model.state_dict(), f"./checkpoint-step-{step}.pt")
Loss Curve Expectation
During fine-tuning, the loss curve should show a steady downward trend. Because the model is already pre-trained on English grammar, the loss will start relatively low and decrease as the model adapts to the structured summarization format.
9. Inference & Evaluation of the Fine-Tuned Model
To evaluate the fine-tuned model, we run inference on stories it did not see during training. We construct the prompt using the story text and the Summary: label, leaving the model to generate the summary and stop when it hits the End of Sentence (EOS) token.
def summarize_story(model, tokenizer, story_text, device):
"""
Generates a summary for a given story using the fine-tuned model.
"""
model.eval()
# Construct prompt without the target summary
prompt = f"Story: {story_text}\nSummary:"
inputs = tokenizer(prompt, return_tensors="pt").to(device)
with torch.no_grad():
output_tokens = model.generate(
**inputs,
max_new_tokens=60,
temperature=0.3, # Lower temperature for more focused, deterministic summaries
eos_token_id=tokenizer.eos_token_id,
pad_token_id=pad_token_id
)
# Decode only the newly generated tokens
input_len = inputs["input_ids"].shape[1]
generated_summary = tokenizer.decode(output_tokens[0][input_len:], skip_special_tokens=True)
return generated_summary.strip()
# Example Evaluation
unseen_story = (
"Once upon a time, there was a tiny bird named Pip. Pip lived in a big oak tree. "
"One day, a strong wind blew Pip out of his nest. He fell all the way to the ground "
"and could not fly back up. A kind boy named Leo saw Pip. Leo gently picked Pip up, "
"climbed the tree, and placed him safely back in his nest. Pip chirped happily."
)
summary = summarize_story(base_model, tokenizer, unseen_story, device)
print("--- Original Story ---")
print(unseen_story)
print("\n--- Generated Summary ---")
print(summary)
Comparison of Outputs
Even when trained on only a fraction of the dataset (e.g., 9% of the target steps), the model successfully shifts from continuing the story to generating a concise, accurate summary.
10. Summary and Key Takeaways
- Supervised Fine-Tuning (SFT) changes a model's behavior from simple next-token prediction to structured task execution (such as text summarization).
- Small Language Models (SLMs) like the 19.3M parameter Tiny Stories model can achieve high coherence when trained on simplified, domain-specific vocabularies.
- Structured Prompt Templates are essential for SFT. They establish clear boundaries between the input context (the story) and the target output (the summary).
- Dynamic Padding and Collation ensure efficient batch training by padding sequences to the maximum length in each batch while ignoring pad tokens during loss calculation.
- Special Token Separation is critical. Keeping the pad token (
...) distinct from the EOS token (<|endoftext|>) prevents the model from getting confused about when to stop generating. - Causal Masking and Labels: Setting pad labels to
-100tells PyTorch's cross-entropy loss function to ignore those tokens, focusing the model's learning entirely on active text. - Hugging Face Integration: Uploading fine-tuned weights to Hugging Face allows you to easily share your models and load them for downstream applications.