Continued Pre-Training: Adapt a Pretrained LM to a New Domain

You've spent months and millions of dollars training a general LM on web text. Now your medical company wants a model that speaks medical, or your code-hosting platform wants one that speaks Python. You don't want to restart from scratch — ~99% of what the model knows is still useful. What you want is Continued Pre-Training (CPT) (also called Domain-Adaptive Pre-Training, DAPT): take the pretrained model and keep training on your domain corpus.

Real examples:

• Code Llama (Roziere et al., 2023) — Meta took Llama 2 and did 500B tokens of continued pretraining on code. Result: a Llama family member that's much stronger at programming while retaining most general ability.
• Don't Stop Pretraining (Gururangan et al., 2020) — the paper that coined DAPT. Showed that even 50–500MB of domain text significantly improves downstream-task accuracy for free.
• BloombergGPT (2023) — a ~50B-param LM with half financial data mixed in from pretraining onward. Their financial benchmark gains come primarily from the domain data, not the size.
• [FinGPT, BioGPT, Galactica, Med-PaLM, StarCoder] — every domain-specialized LM of the last 3 years follows this recipe.

What you'll do

Take GPT-2 (124M) and continue pretraining it on Python source code from codeparrot-clean. Then measure two perplexities:

1. How much better did it get at Python? (should drop substantially)
2. How much worse did it get at English? (catastrophic forgetting check)

That tension — domain gain vs general forgetting — is the central trade-off of CPT, and how you manage it determines how useful your adapted model is.

The three things that make CPT work (or not)

1. Learning rate. This is THE hyperparameter. Use the original pretraining LR and you'll erase everything the model learned. Use zero and nothing changes. The sweet spot is 1–10% of original pretraining LR. GPT-2 was pretrained at ~6e-4; we'll use 5e-5.
2. Data mixing. Pure domain data → catastrophic forgetting. Mixing 10–30% general data back in preserves general ability. We'll run the pure-domain version here to see the forgetting clearly, then discuss the mix.
3. Duration. Too short and nothing sticks. Too long and you overfit. The Code Llama team used 500B tokens on top of Llama 2's 2T — a 25% relative training extension.

---

Before training, establish the baseline. We'll use:

• Python code as the target domain (from codeparrot-clean)
• Wikitext as the held-out general-English reference

On Python code, GPT-2 should struggle — perplexity typically 3–10× higher than on general English. That gap is exactly what we're going to close.

import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForCausalLM
from datasets import load_dataset
import math, time

device = torch.device('cuda')
torch.manual_seed(1337)

print('Loading GPT-2...')
tokenizer = AutoTokenizer.from_pretrained('gpt2')
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained('gpt2').to(device)

print(f'GPT-2: {model.num_parameters()/1e6:.0f}M params')

Build the evaluation corpora — streaming fetch

Two held-out sets: code (CodeParrot, Python source) and general (WikiText-2, English prose). Streaming mode avoids downloading the whole dataset. build_eval_tokens takes the first N docs from each stream, crops each to a max char budget, tokenizes with the GPT-2 BPE, and concatenates with EOS separators.

def build_eval_tokens(n_docs, dataset, key, max_chars_per_doc=3000):
    stream = iter(dataset)
    docs = []
    for _ in range(n_docs):
        try:
            docs.append(next(stream)[key][:max_chars_per_doc])
        except StopIteration:
            break
    ids = []
    for d in docs:
        ids.extend(tokenizer.encode(d) + [tokenizer.eos_token_id])
    return torch.tensor(ids, dtype=torch.long)

print('Fetching eval sets (streaming)...')
code_ds = load_dataset('json', data_files='/workspace/datasets/codeparrot-sample/train.jsonl', split='train')
general_ds = load_dataset('Salesforce/wikitext', 'wikitext-2-raw-v1', split='train')

code_eval_tokens = build_eval_tokens(50, code_ds, key='content')
general_eval_tokens = build_eval_tokens(80, general_ds, key='text')
print(f'Code eval:    {len(code_eval_tokens):,} tokens')
print(f'General eval: {len(general_eval_tokens):,} tokens')

Baseline perplexity — GPT-2 before any pre-training

Perplexity = exp(cross-entropy loss); lower means the model finds the text less surprising. GPT-2 was trained on WebText (~40GB of Reddit-linked pages), which included substantial code. So initial parity between Python and English is normal. What matters is the shift after Step 4's continued pre-training — we want Python ppl to drop while English ppl stays roughly flat (no catastrophic forgetting).

@torch.no_grad()
def perplexity(m, tokens, block_size=256):
    m.eval()
    tokens = tokens.to(device)
    losses = []
    for i in range(0, len(tokens) - block_size, block_size):
        chunk = tokens[i:i + block_size].unsqueeze(0)
        out = m(chunk, labels=chunk)
        losses.append(out.loss.item())
    return math.exp(sum(losses) / len(losses))

baseline_code_ppl = perplexity(model, code_eval_tokens)
baseline_general_ppl = perplexity(model, general_eval_tokens)
print(f'BASELINE  GPT-2 on Python code  : {baseline_code_ppl:.2f}')
print(f'BASELINE  GPT-2 on English text : {baseline_general_ppl:.2f}')
print()
print('Note: GPT-2 was trained on WebText (~40GB of Reddit-linked pages) which includes')
print('substantial code. Initial parity between Python and wikitext perplexity is normal —')
print('what matters is the Python-specific gain in Step 4.')