3.16 Continual & Domain-Adaptive Pretraining¶

The most expensive thing in this book is a from-scratch pretraining run. A frontier base model represents tens of millions of dollars of compute and months of wall-clock time. So the question that dominates real-world LLM engineering is rarely “how do I train a model from scratch?” — it is “I have a good base model and a new requirement (a new domain, a new language, fresher data, a bigger model); how do I get there without paying the full pretraining bill again?”

That is continual pretraining (CPT) — also called continued pretraining or continued training (CT), or, when the new data is concentrated in one domain, domain-adaptive pretraining (DAPT). You take a converged checkpoint and keep optimizing the language-modeling objective on a new data distribution. It sounds trivial — “just keep training” — and that is exactly why so many teams get it wrong. Naively resuming training on new data triggers two opposing failure modes at once: if you use a small learning rate, the model barely learns the new domain (the plasticity problem); if you use a large one, it forgets everything it knew (the catastrophic forgetting problem, also called the stability problem). The entire craft of CPT is navigating this stability–plasticity dilemma at a cost that is a single-digit percentage of the original run.

This chapter covers the full toolkit. We start with the learning-rate dynamics that make or break a CPT run (re-warming and re-decay), then the data-side defense against forgetting (replay), then the specific recipes for domain, language, and streaming adaptation. The second half is about a more aggressive idea: not just continuing to train the same weights, but growing the model — adding depth and width, upcycling a dense model into a Mixture-of-Experts (MoE), and transplanting an old model’s knowledge into a new, larger vocabulary. We close with how to plan a CPT run: predicting the loss trajectory and budgeting compute before you spend it.

Read this chapter alongside Scaling Laws: Kaplan, Chinchilla & Beyond, which sets the compute-vs-data tradeoffs CPT exploits; Learning Rate Schedules, Warmup, Batch Size & Hyperparameters, whose schedule machinery we reuse; Data Mixing, Domain Weighting & Curriculum for the data-mix knobs; and Mixture-of-Experts (MoE) Architectures for the upcycling target. The capstone applies this chapter’s machinery at 100M scale: Mid-Training: Quality Annealing, Long-Context Extension & Capability Injection is literally a CPT pass — resume the pre-decay WSD checkpoint, swap the data mix, re-decay — and is the cheapest place to practice every knob below.

The Stability–Plasticity Dilemma¶

Let \(\theta_0\) be the weights of a base model converged on data distribution \(\mathcal{D}_{\text{base}}\) (say, general web text). We want a model that performs well on a new distribution \(\mathcal{D}_{\text{new}}\) (say, medical literature) without materially regressing on \(\mathcal{D}_{\text{base}}\). Formally we are minimizing a two-term objective,

but with two complications: we usually do not have \(\mathcal{D}_{\text{base}}\) exactly (the original mix is proprietary or lost), and we want the new term to dominate the compute. Catastrophic forgetting is the empirical fact that gradient descent on the first term alone drives the second term up sharply — the parameters that encoded general knowledge are overwritten by the new-domain gradient signal.

Why does forgetting happen so violently in transformers? The loss landscape near a converged minimum is not flat in the directions that matter for old tasks. SGD on new data moves \(\theta\) along whatever directions reduce the new loss fastest; those directions are generally not orthogonal to the directions that the old task cares about. The Fisher information matrix \(F\) of the base task quantifies this — high-curvature directions (large \(F_{ii}\)) are precisely the parameters whose perturbation most damages old performance. Elastic Weight Consolidation (EWC, Kirkpatrick et al., 2017) added an explicit penalty \(\sum_i F_{ii}(\theta_i - \theta_{0,i})^2\), but for LLM-scale CPT the dominant practical tools turn out to be learning-rate control and replay, not curvature penalties — they are cheaper and more robust. We will see why.

A related question every team asks: why not just use LoRA? A low-rank adapter (PEFT I: LoRA, QLoRA, DoRA & The Adapter Family) does bound forgetting almost for free — you can always drop the adapter to recover the base model — but it also bounds learning. CPT’s job is to inject genuinely new knowledge (a new language’s morphology, a new domain’s facts), which is diffuse and high-rank; the empirical pattern is that LoRA is competitive for style/format adaptation and clearly weaker than full-parameter training for knowledge acquisition at fixed token budget. Use full-parameter CPT to add knowledge, and reserve PEFT for behavior. Everything below assumes full-parameter training.

The good news, established by Gupta et al. (2023) and the Ibrahim et al. (2024) “Simple and Scalable Strategies to Continually Pre-train LLMs” study, is that a simple recipe — re-warm the learning rate, re-decay it, and mix in a small fraction of replay data — recovers nearly the performance of an idealized from-scratch run on the union of the data, at a fraction of the cost. The rest of this chapter unpacks each ingredient.

Learning-Rate Re-Warming and Re-Decaying¶

The single most important CPT knob is the learning-rate schedule. The base model finished training at its minimum LR \(\eta_{\min}\) (the floor of a cosine decay, see Learning Rate Schedules, Warmup, Batch Size & Hyperparameters). If you simply resume at \(\eta_{\min}\), the model is in “polish” mode — gradients are tiny, and it will absorb almost nothing from \(\mathcal{D}_{\text{new}}\). To actually learn, you must raise the LR back up. This is re-warming.

The mechanics¶

A CPT schedule has three phases, mirroring a from-scratch run but compressed:

1. Re-warmup (\(T_w\) steps): linearly ramp from a low value (often \(\eta_{\min}\) of the base run, or even lower) up to a re-warm peak \(\eta_{\text{cpt}}\).
2. Stable/decay body: hold or cosine-decay over the CPT token budget.
3. Re-decay (\(T_d\) steps): cosine or linear decay down to a new floor.

Two decisions dominate outcomes: how high to re-warm (\(\eta_{\text{cpt}}\)) and how to decay.

How high? The peak controls the stability–plasticity tradeoff directly. Empirically, re-warming to the original pretraining peak learns the new domain fastest but forgets the most. Re-warming to roughly \(10\%\)–\(30\%\) of the original peak is the common sweet spot for domain adaptation; for a large distribution shift (e.g., a brand-new language) you push higher, toward \(50\%\)+; for a gentle freshness update (new web crawl, same distribution) you stay lower. A crucial subtlety from the Ibrahim et al. study: the transient loss spike caused by re-warming is temporary — the loss on both old and new data jumps up when you raise the LR, then recovers below the starting point. Do not panic at the spike; judge the run by where it settles after re-decay.

How to decay? You must re-decay to a low floor at the end, because — exactly as in the Warmup-Stable-Decay (WSD) schedule — most of the committed loss reduction happens during the decay phase. A CPT run held at high LR and never decayed looks unconverged. If you plan to do several CPT rounds (a streaming scenario), the WSD-style “decay only at the end” or “infinite LR schedule” of Ibrahim et al. is attractive: hold a constant high LR across rounds and only spend a short decay when you need a deployable checkpoint.

Schedule code¶

Here is a from-scratch CPT scheduler with explicit re-warm and re-decay, written as a plain multiplier on the optimizer’s base LR.

import math
from torch.optim.lr_scheduler import LambdaLR

def make_cpt_schedule(
    optimizer,
    base_peak_lr: float,        # the ORIGINAL pretraining peak LR
    rewarm_fraction: float,     # eta_cpt = rewarm_fraction * base_peak_lr
    num_rewarm_steps: int,      # T_w: linear re-warmup
    num_total_steps: int,       # full CPT step budget
    num_decay_steps: int,       # T_d: final cosine re-decay
    start_lr_frac: float = 0.0, # LR at step 0 as a fraction of eta_cpt
    final_lr_frac: float = 0.05 # floor as a fraction of eta_cpt
):
    """
    Three-phase CPT LR schedule (re-warm -> stable -> re-decay).
    Returns a multiplier in [0, 1] RELATIVE to base_peak_lr, so set the
    optimizer's base lr == base_peak_lr.

        |       ____________________
        |      /                    \\
        |     /                      \\___
        |    /                           (floor)
        +---/------------------------------------>
          T_w        stable           T_d
    """
    eta_cpt_mult = rewarm_fraction          # peak as fraction of base peak
    stable_end   = num_total_steps - num_decay_steps

    def lr_lambda(step: int) -> float:
        if step < num_rewarm_steps:
            # Phase 1: linear re-warm from start_lr_frac*eta_cpt -> eta_cpt
            frac = step / max(1, num_rewarm_steps)
            mult = start_lr_frac + frac * (1.0 - start_lr_frac)
            return eta_cpt_mult * mult
        elif step < stable_end:
            # Phase 2: hold at the re-warm peak
            return eta_cpt_mult
        else:
            # Phase 3: cosine re-decay from eta_cpt -> final_lr_frac*eta_cpt
            prog = (step - stable_end) / max(1, num_decay_steps)
            prog = min(prog, 1.0)
            cos = 0.5 * (1.0 + math.cos(math.pi * prog))   # 1 -> 0
            mult = final_lr_frac + (1.0 - final_lr_frac) * cos
            return eta_cpt_mult * mult

    return LambdaLR(optimizer, lr_lambda)

# Example: base model trained at peak 3e-4, finished at 3e-5.
# We re-warm to 20% of peak (6e-5), over 2% of the CPT budget, then
# decay over the final 20%.
# optimizer base lr is set to 3e-4 (base_peak_lr).
# sched = make_cpt_schedule(opt, base_peak_lr=3e-4, rewarm_fraction=0.20,
#             num_rewarm_steps=400, num_total_steps=20_000,
#             num_decay_steps=4_000, final_lr_frac=0.05)

In real training stacks you rarely write this from scratch — you configure it. With Hugging Face transformers, the cleanest route is to build the LambdaLR above and hand it to Trainer(optimizers=(opt, sched)), since the built-in get_scheduler("cosine", ...) has no re-warm/re-decay concept. With Megatron-LM / Megatron-Core the whole recipe is command-line flags: --finetune loads the model weights but resets the optimizer state, RNG state, and iteration counter (exactly the “reset and re-warm” move); --override-opt_param-scheduler makes your new --lr, --lr-warmup-iters, and --lr-decay-style win over the schedule saved in the checkpoint — forget that flag and Megatron will silently resume the base run’s schedule at its floor LR, which is the single most common CPT configuration bug. --no-load-optim / --no-load-rng give finer-grained control if you want weights-only loading without --finetune’s counter reset. MosaicML llm-foundry exposes the same idea as load_path plus load_weights_only: true in YAML, and torchtitan and nanotron both support weights-only warm starts. See Megatron-LM, DeepSpeed & Parallelism in Practice for the surrounding launch configuration.

What about the optimizer state?¶

A subtle decision: do you resume the Adam optimizer state (\(m_t\), \(v_t\)) from the base checkpoint, or reset it? Both are defensible. Resetting forces a fresh re-warmup (the \(v_t\) estimates rebuild from scratch, so warmup is mandatory). Resuming preserves per-parameter learning-rate adaptation but the stale \(v_t\) can mis-scale the first new gradients. The common practical choice is to reset the optimizer state and re-warm, which is also what you must do anyway if the model architecture changed (depth/width growth, new vocab) — the optimizer state has the wrong shape. We discuss precision and checkpoint-format issues in Checkpointing, Fault Tolerance & Long-Running Jobs.

Replay: Fighting Catastrophic Forgetting With Data¶

Re-warming makes the model plastic; replay keeps it stable. The idea is dead simple and remarkably effective: mix a fraction of old-distribution data back into the CPT data stream so that the gradient never sees pure \(\mathcal{D}_{\text{new}}\). The replay ratio \(r\) is the fraction of tokens drawn from \(\mathcal{D}_{\text{base}}\) (or a proxy for it):

Typical values are \(r \in [0.01, 0.30]\). The Ibrahim et al. study found that even 5% replay dramatically reduces forgetting on the original distribution while costing almost nothing in new-domain learning, and that the marginal benefit of replay above ~25–50% is small. The intuition: a small but steady gradient signal pointing back toward \(\mathcal{D}_{\text{base}}\) is enough to keep the model anchored, because forgetting is driven by unopposed drift away from the old minimum.

What if you don’t have the original data?¶

You usually don’t have the exact \(\mathcal{D}_{\text{base}}\). Three options, in rough order of preference:

1. Proxy replay: use a high-quality public general corpus (web text, books, Wikipedia, code) as a stand-in. It need not match the original mix exactly — it just needs to broadly cover the capabilities you want to preserve.
2. Self-generated replay (pseudo-rehearsal): sample text from the base model itself and replay that. The base model’s own distribution is, by definition, the thing you’re trying to preserve. This is the LLM analogue of generative replay and connects to Synthetic Data for Pre- and Post-Training.
3. No replay, low LR: if you genuinely cannot replay, lean harder on a low re-warm peak to limit drift — but accept more forgetting.

A replay-aware data sampler¶

import random
from itertools import cycle

def replay_mixed_stream(new_iter, replay_iter, replay_ratio: float, seed: int = 0):
    """
    Yield documents from a CPT stream that is `replay_ratio` old-distribution
    and the rest new-distribution. Token-level ratios are approximated at the
    document level here; for exact token budgets, weight by document length.

    new_iter    : iterator over NEW-domain documents (the target)
    replay_iter : iterator over BASE/REPLAY documents (anti-forgetting)
    replay_ratio: probability a given document comes from replay (e.g. 0.05)
    """
    rng = random.Random(seed)
    new_pool    = cycle(new_iter)      # in practice these are sharded streams,
    replay_pool = cycle(replay_iter)   # not in-memory; cycle is illustrative
    while True:
        if rng.random() < replay_ratio:
            yield next(replay_pool)
        else:
            yield next(new_pool)

# For *exact* token accounting you instead track running token counts and
# pull from whichever source is behind its target share -- this matters when
# documents vary wildly in length (code files vs. tweets).
def token_balanced_stream(new_iter, replay_iter, replay_ratio, len_fn=len):
    new_pool, replay_pool = cycle(new_iter), cycle(replay_iter)
    seen_new = seen_replay = 0
    while True:
        total = seen_new + seen_replay
        # current replay share; pull replay if we are below target, else new
        cur_replay_share = seen_replay / total if total else 0.0
        if cur_replay_share < replay_ratio:
            doc = next(replay_pool); seen_replay += len_fn(doc)
        else:
            doc = next(new_pool); seen_new += len_fn(doc)
        yield doc

In practice you get this for free from the data libraries. Hugging Face datasets interleaves streaming corpora with explicit sampling probabilities, which is exactly a replay ratio:

from datasets import load_dataset, interleave_datasets

# new domain: your own cleaned corpus, streamed from local JSONL shards
new    = load_dataset("json", data_files="medical/*.jsonl",
                      split="train", streaming=True)
# proxy replay: a general-web corpus standing in for the lost base mix
replay = load_dataset("HuggingFaceFW/fineweb-edu", "sample-10BT",
                      split="train", streaming=True)

# 5% replay; "all_exhausted" oversamples the smaller stream instead of stopping
cpt_stream = interleave_datasets([new, replay], probabilities=[0.95, 0.05],
                                 seed=42, stopping_strategy="all_exhausted")

Note that probabilities are document-level, so — as in the code above — token-level shares only match if the two sources have similar mean document length; pre-tokenize and pack into fixed-length shards (Data Mixing, Domain Weighting & Curriculum) if you need exact token accounting. If you are running on Megatron-LM, the equivalent is a weighted blend passed to --data-path as alternating weight/prefix pairs (--data-path 0.95 /shards/medical_text_document 0.05 /shards/general_text_document), which Megatron samples at the sample level from pre-packed binary shards — that is the exact-token-accounting version. Build those replay shards with the same datatrove / datasets pipeline you used for pretraining data (Pretraining Data: Sources, Crawling & The Data Pipeline) so filtering and dedup are identical on both sides.

Merging back toward the base model¶

Replay fights forgetting during training; weight interpolation fights it after. Because a CPT run stays in the same loss basin as its initialization, the base and adapted checkpoints are usually linearly mode-connected, so you can simply average them:

Sweeping \(\alpha\) traces out an explicit stability–plasticity curve — a cheap post-hoc dial you can tune against the two-curve diagnostic above, with no extra training. This is the same machinery as model soups (Wortsman et al., 2022) and task arithmetic (Ilharco et al., 2022), and mergekit implements it (linear, ties, dare_ties) for real checkpoints; see PEFT II: Prompt/Prefix Tuning, IA3, Model Merging & Soups. Two caveats: merging only works when both checkpoints share an architecture and a tokenizer (so it is unavailable after vocabulary extension), and it can only trade off capabilities the two endpoints already have — it cannot recover knowledge neither model retains.

Domain, Language & Streaming Adaptation¶

The re-warm + replay recipe is the engine. The application — which domain, how big the shift, how the data arrives — sets the dials. Below are the canonical regimes.

Domain adaptation (code, medical, legal, finance)¶

The classic motivation (Gururangan et al., Don’t Stop Pretraining, 2020) is that even a strong general model benefits from a continued pass over in-domain text before any task fine-tuning. Modern recipes (StarCoder-style code adaptation, the BloombergGPT finance model, PMC/Meditron-style medical models, SaulLM legal models) follow the same arc:

• Token budget: a domain pass is typically 10–100 B tokens — single-digit percent of the base run. Below ~5 B tokens you mostly polish; above ~100 B you risk drifting into a new base model with all the forgetting that implies.
• Re-warm peak: \(\sim 10\)–\(30\%\) of base peak. Code adaptation often goes higher because code is structurally far from prose.
• Replay: \(5\)–\(30\%\) general text, plus — importantly — keep some of the base domain’s neighbors. For medical, replay general English and general science so you don’t forget how to write fluently.
• Tokenizer: domain text may tokenize poorly (chemical formulae, legal citations, code symbols). You may need vocabulary extension — covered in the tokenizer-transfer section below.

The order matters too: domain CPT comes before instruction tuning and alignment (Supervised Fine-Tuning & Instruction Tuning). CPT injects knowledge; SFT/RLHF inject behavior. Doing CPT after alignment tends to wash out the alignment.

Language adaptation (adding or strengthening a language)¶

Extending a predominantly-English model to, say, Hindi or Swahili is the hardest CPT regime because the shift is large and the tokenizer is usually a bottleneck: an English-centric BPE tokenizer fragments other scripts into many tokens (high fertility), so the model wastes context and compute, and the new-language embeddings barely exist. The recipe:

1. Extend the tokenizer with new-language merges/tokens (see vocabulary transfer below), initializing new embeddings sensibly.
2. Re-warm higher (\(30\)–\(60\%\) of base peak) — you are teaching genuinely new statistics.
3. Heavy replay of the strong language(s) (\(\geq 25\%\)) so the model retains English reasoning, which often transfers to the new language.
4. Larger budget (\(50\)–\(300\) B tokens) than a same-language domain pass.

A counterintuitive but well-documented effect: cross-lingual transfer means a model that is strong in English and given even modest target-language CPT often reasons better in the target language than a model trained on the target language alone, because abstract reasoning circuits are largely language-agnostic.

Streaming / temporal adaptation (fresh data over time)¶

Models go stale: a 2024-trained model does not know 2026 events. The streaming regime does CPT repeatedly as new data arrives. Key differences from a one-shot pass:

• Schedule: a one-shot cosine that decays to a floor “ends” the model. For repeated rounds, use a WSD/infinite schedule — hold a constant LR across rounds, decay only to produce a release checkpoint, then resume from the pre-decay weights (Ibrahim et al. show the decayed checkpoint is a worse starting point than the high-LR one).
• Replay across all past rounds, not just the original base, to avoid forgetting intermediate knowledge.
• Distribution-shift detection: monitor per-domain perplexity to decide when a refresh is warranted.

This is the pretraining-side analogue of the production Data Flywheels & Continuous Improvement loop.

Model Growth & Weight Reuse¶

So far we kept the architecture fixed. The more ambitious form of “don’t retrain from scratch” changes the architecture but reuses the trained weights as initialization. The premise — validated by Net2Net (Chen et al., 2015), bert2BERT, Gong et al.’s progressive stacking, and the dense-to-MoE sparse upcycling of Komatsuzaki et al. (2022) — is that a smaller trained model is a far better initialization for a larger one than random init, so you can “grow” into a bigger model and finish with a short CPT pass.

Depth growth (layer stacking)¶

To go from \(L\) to \(2L\) layers, duplicate existing layers (interleaved or stacked) so the initial deeper network computes nearly the same function as the shallow one. The cleanest trick is function-preserving growth: add new layers initialized so their contribution is zero at step 0. Because transformer blocks are residual, \(x \mapsto x + f(x)\), you can make a new block an identity by zeroing the output projection of its attention and MLP sublayers — then \(f(x)=0\) and the block passes its input through untouched. Training then “wakes up” the new layers gradually.

import torch, torch.nn as nn, copy

def grow_depth_identity(model, insert_after: list[int]):
    """
    Insert function-preserving (identity-at-init) transformer blocks.
    `model.layers` is a ModuleList of residual transformer blocks, each of
    the form: x = x + attn(ln1(x)); x = x + mlp(ln2(x)).
    We clone an existing block and ZERO its residual-output projections so the
    new block is the identity map at initialization -> the grown model
    computes the SAME function as the original at step 0 (loss is preserved).
    """
    new_layers = []
    for i, layer in enumerate(model.layers):
        new_layers.append(layer)
        if i in insert_after:
            twin = copy.deepcopy(layer)
            with torch.no_grad():
                # zero the output projection of attention (o_proj) and MLP
                # (down_proj). Names depend on your block; adapt accordingly.
                twin.attn.o_proj.weight.zero_()
                if twin.attn.o_proj.bias is not None:
                    twin.attn.o_proj.bias.zero_()
                twin.mlp.down_proj.weight.zero_()
                if getattr(twin.mlp.down_proj, "bias", None) is not None:
                    twin.mlp.down_proj.bias.zero_()
            new_layers.append(twin)
    model.layers = nn.ModuleList(new_layers)
    model.config.num_layers = len(new_layers)
    return model

The alternative, progressive stacking (train shallow, copy-stack to deeper, continue), was shown by Gong et al. to reach a target deep model faster in total FLOPs than training the deep model from scratch — the shallow phase is cheap and the warm start is good.

Width growth¶

Widening a hidden dimension \(d \to d'\) reuses Net2Net’s net2wider: copy existing neurons and split their outgoing weights so the function is preserved. If neuron \(j\) is duplicated \(k\) times, divide each copy’s outgoing weights by \(k\) so the summed contribution is unchanged. To break the symmetry (otherwise the copies receive identical gradients forever and never diverge), add tiny noise. bert2BERT applies exactly this “function-preserving” widening to initialize a wide BERT from a narrow one and recovers most of the from-scratch performance with far less compute.

import torch

@torch.no_grad()
def net2wider_linear(W_in: torch.Tensor, W_out: torch.Tensor, new_width: int,
                     noise: float = 1e-3):
    """
    Function-preserving width expansion of one hidden layer.
      W_in : (hidden, in)   -> produces the hidden activations
      W_out: (out, hidden)  -> consumes them
    Returns expanded (W_in', W_out') with hidden -> new_width such that the
    composed map is unchanged (up to small symmetry-breaking noise).
    """
    hidden, _in = W_in.shape
    assert new_width >= hidden
    # choose which existing neurons to duplicate (uniform with replacement)
    idx = torch.randint(0, hidden, (new_width - hidden,))
    pick = torch.cat([torch.arange(hidden), idx])         # (new_width,)
    # count copies of each original neuron, for the divide-by-k correction
    counts = torch.bincount(pick, minlength=hidden).float()
    W_in_new  = W_in[pick].clone()                        # replicate rows
    W_in_new += noise * torch.randn_like(W_in_new)        # break symmetry
    # outgoing weights: divide each replica by the # of copies of its source
    W_out_new = (W_out[:, pick] / counts[pick].unsqueeze(0)).clone()
    return W_in_new, W_out_new

Dense-to-MoE upcycling¶

Sparse upcycling turns a trained dense model into a sparse Mixture-of-Experts model, reusing the dense weights. The recipe (Komatsuzaki et al., 2022; used at scale for several production MoEs): replace each (or every other) dense MLP block with an MoE layer whose \(E\) experts are each initialized as a copy of the original dense MLP. The router is added fresh (small random init). At step 0, every expert is identical, so — if the router is roughly uniform — the MoE layer computes approximately the original dense MLP’s output, preserving the function. Continued pretraining then differentiates the experts. This buys you a high-capacity MoE (more parameters, same or modestly higher FLOPs per token) for the cost of a CPT pass, rather than training an MoE from scratch.

import torch, torch.nn as nn, copy

class UpcycledMoE(nn.Module):
    """
    Replace a dense MLP with an MoE whose experts are clones of that MLP.
    At init all experts are identical, so with a near-uniform router the layer
    approximately reproduces the dense MLP -> function-preserving upcycle.
    See chapter 2.9 for routing, load balancing, and capacity factors.
    """
    def __init__(self, dense_mlp, num_experts=8, top_k=2):
        super().__init__()
        d_model = dense_mlp.up_proj.in_features
        # one router (gate) added fresh; small init so logits start ~uniform
        self.gate = nn.Linear(d_model, num_experts, bias=False)
        nn.init.normal_(self.gate.weight, std=1e-3)
        # each expert is a deep copy of the trained dense MLP
        self.experts = nn.ModuleList(copy.deepcopy(dense_mlp)
                                     for _ in range(num_experts))
        self.top_k = top_k

    def forward(self, x):                       # x: (tokens, d_model)
        logits = self.gate(x)                   # (tokens, E)
        w, idx = torch.topk(logits.softmax(-1), self.top_k, dim=-1)
        w = w / w.sum(-1, keepdim=True)         # renormalize top-k weights
        out = torch.zeros_like(x)
        for slot in range(self.top_k):          # gather-scatter over experts
            for e, expert in enumerate(self.experts):
                mask = idx[:, slot] == e
                if mask.any():
                    out[mask] += w[mask, slot:slot+1] * expert(x[mask])
        return out

See Mixture-of-Experts (MoE) Architectures for routing, auxiliary load-balancing losses, and capacity factors — all of which you must tune during the upcycling CPT, because freshly-added routers are prone to collapse (routing all tokens to one expert). The loop above is deliberately written for clarity, not speed: a real implementation sorts tokens by expert and does one grouped GEMM per expert (see the MoE chapter, or the fused kernels in megablocks / grouped_gemm). You do not have to write the checkpoint surgery yourself either — Megatron-Core ships a dense-to-MoE upcycling path (enabled with --moe-use-upcycling, which converts a loaded dense checkpoint into the MoE architecture described by your --num-experts / --moe-router-topk flags on the first save), and mergekit’s passthrough merge method is the off-the-shelf way to do the depth-stacking growth of the previous subsection on Hugging Face checkpoints.

Vocabulary & tokenizer transfer¶

When you change the tokenizer — extending the vocab for a new domain/language, or swapping to a different tokenizer entirely — the embedding matrix \(E \in \mathbb{R}^{|V|\times d}\) and the output (unembedding) matrix must change shape, and the new rows must be initialized. Random init for new tokens is wasteful; the trained model already “knows” the meaning of the pieces of a new token. Two standard moves:

• Mean-of-subtokens init: a new token (e.g., a whole word that the old tokenizer split into pieces) gets its embedding initialized to the mean of its old sub-token embeddings. This is the heuristic behind FOCUS and the widely-used embedding-init utilities, and it dramatically shortens the CPT needed to make the new tokens useful.
• Shared-token preservation: tokens present in both vocabularies keep their trained embeddings exactly. Only genuinely new tokens are initialized.

import torch

@torch.no_grad()
def init_new_embeddings(old_emb, new_vocab, old_tokenizer, mean_init=True):
    """
    Build a new embedding matrix for an extended/changed vocabulary.
      old_emb       : (|V_old|, d) trained embedding matrix
      new_vocab     : dict {new_token_str -> new_id}
      old_tokenizer : can encode a string into OLD ids
    Shared tokens copy their trained vector; new tokens are initialized to the
    mean of the OLD sub-token embeddings of their surface string (FOCUS-style).
    Falls back to the overall mean (a safe centroid) when no sub-tokens exist.
    """
    d = old_emb.shape[1]
    new_emb = torch.empty(len(new_vocab), d)
    overall_mean = old_emb.mean(0)
    old_vocab = old_tokenizer.get_vocab()       # {token_str -> old_id}
    for tok, new_id in new_vocab.items():
        if tok in old_vocab:                     # shared: copy trained vector
            new_emb[new_id] = old_emb[old_vocab[tok]]
        elif mean_init:                          # new: mean of OLD sub-tokens
            sub_ids = old_tokenizer.encode(tok, add_special_tokens=False)
            if sub_ids:
                new_emb[new_id] = old_emb[torch.tensor(sub_ids)].mean(0)
            else:
                new_emb[new_id] = overall_mean
        else:
            new_emb[new_id] = overall_mean
    return new_emb

Three details that bite in practice. (1) The output head. If the model unties input and output embeddings, you must initialize the new lm_head rows too — run the same sub-token-mean procedure on the unembedding matrix (and on the output bias, if any). If weights are tied, resizing the input embedding is enough, but then a mean-initialized row is simultaneously an input vector and a logit direction, which is one reason tied models sometimes need a few hundred extra steps to stop over-predicting new tokens. (2) Padding. Round the new vocabulary size up to a multiple of 64 (128 with tensor parallelism) with dummy tokens; unaligned vocab sizes cost real throughput in the final GEMM, and Megatron’s --make-vocab-size-divisible-by does this for you. (3) Order of operations. After re-initializing embeddings you almost always CPT the whole model (not just the embeddings) so the body learns to use the new vocabulary; a common warm-up is to first train only the new rows for a few hundred steps with the rest frozen, then unfreeze everything.

The library path mirrors the code above. Train the new merges with Hugging Face tokenizers (or SentencePiece) on target-domain text, splice them in with tokenizer.add_tokens([...]) — or, for a language pass, train a fresh tokenizer and take the union of vocabularies — then call model.resize_token_embeddings(len(tokenizer), pad_to_multiple_of=64). Recent transformers versions default to mean_resizing=True, which draws new rows from a distribution fitted to the existing embeddings rather than from a plain \(\mathcal{N}(0,\sigma^2)\); that is strictly better than random but is not sub-token-aware, so for a serious vocabulary transfer overwrite those rows with the FOCUS-style means computed above. See Tokenization: BPE, WordPiece, Unigram & Byte-Level for how the merges that define new tokens are produced, Embeddings & The Input Pipeline for embedding/unembedding tying, and A Byte-Level BPE Tokenizer From Scratch for the capstone’s own vocabulary, whose 100M-scale embedding table makes these tradeoffs cheap to measure.

Planning CPT Compute With Loss-Trajectory Models¶

The final discipline is planning: before spending the compute, predict where the loss will land and how to split the budget. Three questions recur.

1. How many tokens? CPT obeys a scaling-law-like diminishing-returns curve. The new-domain loss after \(D\) CPT tokens is well-modeled by a shifted power law,

where \(D_0\) encodes the “head start” the base model already has on the new domain, and \(A,\alpha,\mathcal{L}_\infty\) are fit from a short pilot. Run a small pilot at, say, 1 B, 2 B, 4 B tokens, fit the curve, and extrapolate to find the knee where additional tokens stop paying. The same machinery is developed in Scaling Laws: Kaplan, Chinchilla & Beyond; CPT just adds the \(D_0\) offset that represents transferred knowledge.

2. How much forgetting? Model the base-domain loss as rising with CPT tokens at a rate damped by the replay ratio \(r\). A serviceable empirical form is

so larger replay \(r\) flattens forgetting. Fitting both curves from the pilot lets you pick the \((D, r, \eta_{\text{cpt}})\) that hits a target — e.g., “maximize medical gain subject to base-loss regression \(\le 2\%\).”

3. Grow now or later? If you also plan model growth, the loss-trajectory model tells you whether to grow first (cheaper if the growth’s function-preserving init starts near the base loss) or finish CPT on the small model first.

Here is a tiny fitter you can run on pilot points.

import numpy as np
from scipy.optimize import curve_fit

def fit_cpt_trajectory(tokens, losses):
    """
    Fit L(D) = L_inf + A / (D0 + D)**alpha to pilot (tokens, loss) points,
    then return a predictor and the token count to hit a target loss.
    tokens : array of CPT token counts (e.g. [1e9, 2e9, 4e9])
    losses : measured new-domain loss at each.
    """
    def law(D, L_inf, A, D0, alpha):
        return L_inf + A / np.power(D0 + D, alpha)
    p0 = [min(losses) * 0.9, 1.0, 1e9, 0.3]            # rough initial guess
    popt, _ = curve_fit(law, np.asarray(tokens), np.asarray(losses),
                        p0=p0, maxfev=20000)
    L_inf, A, D0, alpha = popt

    def predict(D):                                    # loss at D tokens
        return L_inf + A / (D0 + D) ** alpha

    def tokens_for_target(target_loss):                # invert the law
        if target_loss <= L_inf:
            return float("inf")                        # unreachable: below L_inf
        return (A / (target_loss - L_inf)) ** (1.0 / alpha) - D0

    return predict, tokens_for_target, popt

# Example usage:
# predict, tokens_for, params = fit_cpt_trajectory(
#     [1e9, 2e9, 4e9], [2.41, 2.30, 2.22])
# print(predict(40e9))           # extrapolated loss at the full 40B budget
# print(tokens_for(2.10))        # tokens needed to reach loss 2.10

Key Takeaways¶

Further Reading¶

• Gururangan et al., “Don’t Stop Pretraining: Adapt Language Models to Domains and Tasks” (2020) — the foundational case for domain-adaptive and task-adaptive continued pretraining (DAPT/TAPT).
• Ibrahim et al., “Simple and Scalable Strategies to Continually Pre-train Large Language Models” (2024) — the definitive modern recipe: re-warming, re-decaying, replay ratios, and infinite LR schedules, with large-scale ablations.
• Gupta et al., “Continual Pre-Training of Large Language Models: How to (re)warm your model?” (2023) — focused study of LR re-warming dynamics and the transient spike.
• Kirkpatrick et al., “Overcoming Catastrophic Forgetting in Neural Networks” (EWC, 2017) — the Fisher-information curvature penalty; the classic framing of forgetting that CPT replay sidesteps cheaply.
• Chen et al., “Net2Net: Accelerating Learning via Knowledge Transfer” (2015) — function-preserving net2wider/net2deeper transformations underlying model growth.
• Gong et al., “Efficient Training of BERT by Progressively Stacking” (2019) and Chen et al., “bert2BERT: Towards Reusable Pretrained Language Models” (2021) — depth/width growth with weight reuse for transformers.
• Komatsuzaki et al., “Sparse Upcycling: Training Mixture-of-Experts from Dense Checkpoints” (2022) — the dense-to-MoE upcycling recipe.
• Dobler & de Melo, “FOCUS: Effective Embedding Initialization for Monolingual Specialization of Multilingual Models” (EMNLP 2023) — sub-token-mean embedding initialization for vocabulary/tokenizer transfer.
• Wu et al., “BloombergGPT” (2023) and Chen et al., “Meditron” / SaulLM legal-LM reports (2023–2024) — real domain-adaptive pretraining recipes for finance, medicine, and law.

Exercises¶

1. (Conceptual.) A colleague wants to specialize your base model on legal text. The base model finished its cosine schedule at its floor LR \(\eta_{\min} = 3\times10^{-5}\), and your colleague proposes simply resuming training on the legal corpus at exactly \(\eta_{\min}\), “since that’s where the model left off.” (a) Which horn of the stability–plasticity dilemma does this proposal fall on, and what will the legal-domain loss curve look like? (b) Your colleague, told to “raise the LR,” instead jumps straight to \(3\times10^{-4}\) with no warmup. What goes wrong mechanically? © After you finally run a proper re-warm, the loss on both the legal set and the held-out base set spikes upward in the first few hundred steps. Should you abort? Why or why not?

2. (Quantitative — replay accounting.) You must ensure the model actually sees at least \(50\) B tokens of new-domain (biomedical) text, and you want a replay ratio of \(r = 0.2\) general-text tokens to fight forgetting. (a) What is the total CPT token budget \(D\), and how many replay tokens does that imply? (b) Relative to a hypothetical replay-free run that also sees \(50\) B new tokens, what percentage of extra compute does the replay cost? © Ibrahim et al. found the marginal benefit of replay above ~25–50% is small, and even 5% helps a lot. Given that, is \(r=0.2\) a defensible choice here?

3. (Quantitative — compute budgeting.) Your base model has \(N = 3\times10^{9}\) parameters and was pretrained on \(1.5\) T tokens. You plan a \(30\) B-token domain-adaptive pass on a cluster of \(16\) H100s, each sustaining an achieved \(4\times10^{14}\) FLOP/s. Using the \(C \approx 6ND\) estimate: (a) What is the CPT FLOP cost and the wall-clock time? (b) What fraction of the original pretraining token budget is this pass?

4. (Conceptual — why upcycling preserves the function.) Look at the UpcycledMoE.forward code. At initialization every expert is a deep copy of the same dense MLP. (a) Show that at step 0 the MoE layer reproduces the dense MLP’s output exactly, and note precisely which line makes this independent of the router’s logits. (b) The chapter’s prose says the layer computes the dense output “approximately, if the router is roughly uniform.” Reconcile that hedge with your exact result in (a). © Name the failure mode the chapter warns about for freshly-added routers, and explain why it becomes a real risk only after step 0.

5. (Quantitative — loss-trajectory planning.) You ran a short pilot and fit_cpt_trajectory returned the shifted power law \(\mathcal{L}_{\text{new}}(D) = \mathcal{L}_\infty + A/(D_0 + D)^{\alpha}\) with \(\mathcal{L}_\infty = 2.00\), \(A = 200\), \(D_0 = 1\times10^{9}\), \(\alpha = 0.30\) (tokens in absolute counts). (a) Predict the new-domain loss at the planned full budget of \(D = 40\) B tokens. (b) Using tokens_for_target, how many tokens would you need to reach a loss of \(2.10\)? © What does the answer to (b), compared to the \(40\) B plan, tell you about where the “knee” of this curve is?

6. (Implementation — verify function-preserving depth growth.) The grow_depth_identity function claims the grown model computes the same function as the original at step 0. Write a small runnable harness that builds a toy residual transformer, records its output on a batch, grows it, and asserts the output is unchanged. Then explain in one sentence which lines of grow_depth_identity are responsible for the invariance.

import torch, torch.nn as nn

class ToyAttn(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.qkv    = nn.Linear(d, d)
        self.o_proj = nn.Linear(d, d)      # residual output projection
    def forward(self, x):
        return self.o_proj(torch.relu(self.qkv(x)))

class ToyMLP(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.up_proj   = nn.Linear(d, 4 * d)
        self.down_proj = nn.Linear(4 * d, d)   # residual output projection
    def forward(self, x):
        return self.down_proj(torch.relu(self.up_proj(x)))

class ToyBlock(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.attn, self.mlp = ToyAttn(d), ToyMLP(d)
    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.mlp(self.ln2(x))
        return x

class ToyModel(nn.Module):
    def __init__(self, d, L):
        super().__init__()
        self.layers = nn.ModuleList(ToyBlock(d) for _ in range(L))
        self.config = type("cfg", (), {"num_layers": L})()
    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

torch.manual_seed(0)
model = ToyModel(d=16, L=4)
x = torch.randn(2, 5, 16)

with torch.no_grad():
    before = model(x)
grow_depth_identity(model, insert_after=[1, 3])   # from the chapter
with torch.no_grad():
    after = model(x)

print("layers:", model.config.num_layers)                 # 4 -> 6
print("max abs diff:", (before - after).abs().max().item())
assert torch.allclose(before, after, atol=1e-6)
print("function preserved at init: OK")