3.13 Long-Context Pretraining & Context Extension¶

What changes when you train a language model to read 128 thousand tokens at once instead of 2 thousand? The answer turns out to touch nearly every part of the stack: attention arithmetic, positional encoding mathematics, data curation, parallelism strategy, and inference memory management. This chapter develops each of those changes from first principles.

We will see why a naive extension to long contexts fails, how the community learned to patch positional encodings cheaply via RoPE interpolation, how continued pretraining on long documents consolidates those patches, and finally how ring/context parallelism makes the training computationally tractable. We close by connecting long-context training to the inference KV-cache problem you will likely be asked about in a systems design interview.

The Two Fundamental Bottlenecks¶

Quadratic Attention Cost¶

Standard self-attention computes a score matrix \(S \in \mathbb{R}^{T \times T}\) for a sequence of length \(T\):

Both the FLOPs for this product and the memory to store \(S\) grow as \(\mathcal{O}(T^2)\). At \(T = 2048\) this is manageable; at \(T = 128{,}000\) it becomes painful:

A model with 32 heads would need ~2 TB just for score matrices at 128 K. FlashAttention (covered in depth in FlashAttention I: IO-Awareness & The Online Softmax and FlashAttention 2 & 3) eliminates the need to materialise the full \(S\) in HBM by fusing the computation in tiles, reducing the HBM requirement to \(\mathcal{O}(T)\) activations. But the \(\mathcal{O}(T^2)\) FLOPs remain — you simply compute them faster.

Memory From Activations and KV Cache¶

During training, gradient checkpointing can trade memory for compute, but even with it, intermediate activations scale linearly with \(T\) per layer. More acutely, training requires storing the full sequence per sample in the attention buffer. At long sequence lengths the per-token memory cost does not shrink; it just shifts from the score matrix to the gradient and activation buffers.

During inference, all previously computed keys and values must be kept in memory — the KV cache — so the context cost is \(\mathcal{O}(T)\) per forward pass but persists for the entire generation episode. For a 70 B-parameter model with 8 KV heads (GQA), the KV cache for a single 128 K-token context is roughly:

With \(d_{\text{kv}} = 128\) and bf16 (2 bytes): \(2 \times 8 \times 128{,}000 \times 128 \times 2 \approx 524\;\text{MB}\) per layer. With 80 layers that is around 42 GB — comparable to the model weights themselves. Managing this is the focus of PagedAttention & KV-Cache Memory Management.

RoPE and the Extrapolation Problem¶

Why Sinusoidal and Learned Encodings Fail¶

A model trained with absolute positional embeddings at positions \(\{0, \ldots, T_{\text{train}}-1\}\) has never seen the embedding at position \(T_{\text{train}} + k\) for any \(k > 0\). Learned embeddings simply have no weight there. Sinusoidal embeddings produce a valid vector at any position, but the attention patterns were never trained to use those values, so performance degrades sharply. See Positional Encodings: Sinusoidal, Learned, RoPE & ALiBi for the full encoding taxonomy.

RoPE Refresher¶

RoPE (Rotary Position Embedding, Su et al., 2022) applies a position-dependent rotation to query and key vectors before the dot product. For a pair of scalars at dimension index \(2i\) and \(2i+1\):

where \(m\) is the absolute position and \(\theta_i = b^{-2i/d}\) with base \(b = 10{,}000\) by default. The key insight is that the dot product \(q_m \cdot k_n\) depends only on the relative position \(m - n\), because the rotation matrices satisfy \(R_m^\top R_n = R_{m-n}\).

The problem for extrapolation: at training time, \(m\) never exceeds \(T_{\text{train}}\). The angles \(m \theta_i\) stay within \([0, T_{\text{train}} \theta_i]\). For the low-frequency dimensions (large \(i\), tiny \(\theta_i\)) this range can be small — those dimensions barely rotate even within the training window. For the high-frequency dimensions the rotation is well-covered. When you push \(m > T_{\text{train}}\), high-frequency dimensions enter completely unseen phase territory, causing attention to produce pathological scores.

Position Interpolation (PI)¶

Chen et al. (2023, Extending Context Window of Large Language Models via Positional Interpolation) propose a one-line fix: rather than feeding position \(m\) to RoPE, feed the scaled position \(m' = m \cdot \frac{T_{\text{train}}}{T_{\text{target}}}\). This linearly rescales positions so the full target window fits inside the training range.

If the training window is \(T_{\text{train}} = 4{,}096\) and the target is \(T_{\text{target}} = 32{,}768\), every position is divided by \(32{,}768 / 4{,}096 = 8\). All angles stay within the trained range.

The cost: nearby tokens that were originally at relative distance 1 are now at effective distance \(1/8\), so the model initially cannot distinguish adjacent tokens well. A small amount of fine-tuning (on the order of 1 000 steps) on long documents restores performance. The work reports good recovery at 8x extension.

NTK-aware Interpolation¶

NTK-aware interpolation (bloc97, community research 2023) observes that PI uniformly compresses all frequencies, damaging the high-frequency (short-range) components that carry most of the positional signal. Instead, it rescales only the base \(b\):

This increases the base (e.g. from 10 000 to ~83 000 for an 8x extension with \(d = 128\)), making each \(\theta_i\) smaller and thus each dimension rotate more slowly — effectively stretching the rope. High-frequency dimensions retain their resolution for local dependencies while low-frequency dimensions stretch to cover global distances. NTK-aware interpolation often works without any fine-tuning, giving it a zero-shot extension property.

YaRN¶

YaRN (Peng et al., 2023, YaRN: Efficient Context Window Extension of Large Language Models) builds further on NTK-aware scaling by introducing:

1. Frequency ramping. A smooth interpolation factor \(\gamma(i)\) that applies no scaling to the highest-frequency (most local) dimensions and full PI scaling to the lowest-frequency (most global) ones:

where \(\lambda_i = 2\pi / \theta_i\) is the wavelength and \(s = T_{\text{target}} / T_{\text{train}}\) is the scale factor.

2. Attention temperature correction. Spreading the same probability mass over \(s\times\) more keys raises the entropy of the softmax, flattening the attention distribution. YaRN counteracts this by multiplying the pre-softmax logits by \(1/t\), using the empirically fitted rule

At \(s = 8\) this gives \(\sqrt{1/t} \approx 1.21\), i.e. logits sharpened by \(\approx 1.46\times\). Because the factor is a constant, it is free at runtime: implementations (including HuggingFace’s rope_scaling={"rope_type": "yarn", ...}) fold \(\sqrt{1/t}\) into the cached cos/sin tables, which scales both \(q\) and \(k\) and therefore each logit by exactly \(1/t\) — no change to the attention kernel is needed.

YaRN requires only a small amount of continued training (around 400 steps in the original work) and often outperforms PI at the same training budget. It is widely adopted: Mistral, Qwen, and Deepseek use variants of YaRN.

import torch
import math

def build_rope_freqs(dim: int, base: float, seq_len: int) -> torch.Tensor:
    """Build standard RoPE inverse frequencies for a given dimension and base."""
    # Shape: (dim // 2,)
    inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
    # Outer product with positions → (seq_len, dim // 2)
    t = torch.arange(seq_len, dtype=torch.float)
    freqs = torch.outer(t, inv_freq)      # shape: (seq_len, dim // 2)
    return freqs  # each row is the angles for one position

def build_yarn_freqs(
    dim: int,
    base: float,
    orig_max_seq_len: int,
    target_max_seq_len: int,
    beta_fast: float = 32.0,   # high-freq boundary (wavelengths)
    beta_slow: float = 1.0,    # low-freq boundary
) -> torch.Tensor:
    """
    YaRN frequency computation.
    Returns per-dimension scaling factor to apply to the standard freqs.
    """
    scale = target_max_seq_len / orig_max_seq_len

    # Standard inverse frequencies
    inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))

    # Wavelength for each dimension: λ = 2π / θ_i
    wavelengths = 2 * math.pi / inv_freq

    # Ramp function: 0 for high-freq dims (leave unchanged), 1 for low-freq (full scale)
    ramp = (wavelengths - orig_max_seq_len / beta_fast) / (
        orig_max_seq_len / beta_slow - orig_max_seq_len / beta_fast
    )
    ramp = ramp.clamp(0.0, 1.0)

    # Blend: high-freq → no interpolation, low-freq → PI scaling
    # NTK base for fully-scaled dims: b' = b * scale^(dim/(dim-2))
    ntk_base = base * (scale ** (dim / (dim - 2)))
    ntk_inv_freq = 1.0 / (ntk_base ** (torch.arange(0, dim, 2).float() / dim))

    # Interpolate between unscaled (ramp=0) and NTK-scaled (ramp=1) per dimension
    blended_inv_freq = (1 - ramp) * inv_freq + ramp * ntk_inv_freq

    t = torch.arange(target_max_seq_len, dtype=torch.float)
    freqs = torch.outer(t, blended_inv_freq)
    return freqs  # (target_max_seq_len, dim // 2)

def yarn_attention_factor(orig_max_seq_len: int, target_max_seq_len: int) -> float:
    """
    YaRN's constant logit sharpening: sqrt(1/t) = 0.1 * ln(s) + 1.
    Multiply the cached cos/sin by this factor -- it scales q and k, hence
    each logit by 1/t -- instead of touching the attention kernel.
    """
    scale = target_max_seq_len / orig_max_seq_len
    return 0.1 * math.log(scale) + 1.0

# --- demo ---
# GPT-style model: dim=128, base=10000, trained at 4096, extending to 32768
orig_freqs  = build_rope_freqs(128, 10000, 4096)
yarn_freqs  = build_yarn_freqs(128, 10000, 4096, 32768)

print(f"Original freq range at pos 4096: [{orig_freqs[-1].min():.4f}, {orig_freqs[-1].max():.4f}]")
print(f"YaRN freq range at pos 32768:    [{yarn_freqs[-1].min():.4f}, {yarn_freqs[-1].max():.4f}]")
print(f"YaRN attention factor sqrt(1/t): {yarn_attention_factor(4096, 32768):.4f}")  # ~1.2079
# YaRN keeps high-freq dims in known territory; low-freq dims gently scaled.

Continued Pretraining on Long Documents¶

Why Fine-Tuning Alone Is Insufficient¶

RoPE scaling patches the positional encoding but does not teach the model new attention patterns or update MLP weights to reason over document-scale structure. A 4 K-trained model re-encodes positions at 32 K but the attention heads have never learned which distant tokens to attend to: they were never penalised for ignoring tokens at offset >4 096. Continued pretraining on genuinely long documents rectifies this.

Data Curation for Long-Context Training¶

The key challenge is obtaining sufficiently long training documents. The pretraining corpus — described in Pretraining Data: Sources, Crawling & The Data Pipeline — mostly consists of web paragraphs shorter than the target context. Long-context continued pretraining typically draws from:

• Books (Project Gutenberg / PG-19, and the books subsets of open corpora such as AI2’s Dolma): naturally tens of thousands of tokens. Note that Books3, ubiquitous in 2022–2023 papers, was withdrawn over copyright claims and should not be used in new work.
• Scientific papers (arXiv LaTeX source, peS2o from Semantic Scholar): 6–20 K tokens each.
• Code repositories (The Stack v2, GitHub): single files are short — the length comes from repo-level concatenation, ordering files by import/dependency so the concatenation is semantically coherent.
• Legal and financial documents: contracts, 10-K filings.
• Long web articles: Wikipedia pages, journalism.

Filtering and deduplicating these follows the same machinery as the base corpus (Data Cleaning, Deduplication & Quality Filtering); in practice you run the same datatrove pipeline with an extra length filter (doc_len >= target_len / 2) and re-pack the shards at the new sequence length.

A useful rule of thumb from community experimentation: 5–20 billion tokens of continued pretraining at the target context length is often sufficient to unlock the new length. The ratio of long-to-short documents matters — too many short documents and the model never practices long-range attention. Two refinements from frontier recipes:

• Extend in stages, not one jump. The Llama 3 report describes raising the window from 8 K to 128 K in six successive stages over a large token budget (on the order of hundreds of billions of tokens), advancing to the next stage only once short-context evals had recovered and needle retrieval at the current length was solved. Staging keeps each RoPE rescale small, so less has to be repaired by training.
• Keep a short-context anchor in the mix. A mix that is 100% long documents reliably regresses short-context benchmarks; the 60/40 long/original-mix split in the recipe below exists for exactly that reason.

Learning Rate and Scheduler Strategy¶

Continued pretraining for context extension typically uses:

• A learning rate significantly lower than the original pretraining peak (on the order of 1–5% of it), to avoid catastrophic forgetting of short-context capabilities.
• A cosine decay or flat schedule, often without warmup since the model already has stable weights.
• Gradient clipping set conservatively (e.g., \(\ell_2\) norm of 1.0).

See Learning Rate Schedules, Warmup, Batch Size & Hyperparameters for the general scheduling machinery.

Document Packing for Long-Context Training¶

The Problem With Naïve Packing¶

Training efficiency requires that every GPU sees a tensor of exactly \(T_{\text{seq}}\) tokens per step. Short documents must be concatenated to fill the window, and long documents must be split across windows (or within one window). The naïve approach simply concatenates documents end-to-end and slices at \(T_{\text{seq}}\) boundaries.

This creates a cross-document attention problem: the model computes attention between the end of document \(A\) and the beginning of document \(B\), even though they are unrelated. At short context lengths this is a minor source of noise; at 128 K context it can mean most of the window is cross-document garbage, actively harming training signal.

Document Masking in Attention¶

The principled fix is a document attention mask: each token attends only to tokens that belong to the same document (and obeys causal ordering within that document). This is equivalent to block-diagonal causal masking over the packed sequence.

import torch

def build_packed_attention_mask(
    doc_lengths: list[int],
    device: str = "cpu",
) -> torch.Tensor:
    """
    Build a boolean causal attention mask for a packed sequence of multiple docs.

    doc_lengths: list of token counts per document, must sum to seq_len.
    Returns: (seq_len, seq_len) bool tensor where True = can attend.
    """
    seq_len = sum(doc_lengths)
    # Start with no attention allowed
    mask = torch.zeros(seq_len, seq_len, dtype=torch.bool, device=device)

    offset = 0
    for length in doc_lengths:
        # Within this document's block, allow lower-triangular (causal) attention
        block_end = offset + length
        for row in range(offset, block_end):
            # Token at 'row' can attend to [offset .. row] (same doc, causal)
            mask[row, offset:row + 1] = True
        offset = block_end

    return mask  # (seq_len, seq_len)

# Example: two docs packed into a single context of length 8 (4 + 4)
mask = build_packed_attention_mask([4, 4])
print(mask.int())
# tensor([[1, 0, 0, 0, 0, 0, 0, 0],
#         [1, 1, 0, 0, 0, 0, 0, 0],
#         [1, 1, 1, 0, 0, 0, 0, 0],
#         [1, 1, 1, 1, 0, 0, 0, 0],
#         [0, 0, 0, 0, 1, 0, 0, 0],
#         [0, 0, 0, 0, 1, 1, 0, 0],
#         [0, 0, 0, 0, 1, 1, 1, 0],
#         [0, 0, 0, 0, 1, 1, 1, 1]], dtype=torch.int32)

In practice, FlashAttention 2 supports “variable-length” (varlen) sequences via the cu_seqlens API, which achieves exactly this masking without materialising the full \(T \times T\) mask tensor:

# FlashAttention varlen API (flash_attn library)
from flash_attn import flash_attn_varlen_func
import torch

# Suppose we have a batch of two sequences packed: lengths 512 and 1024
batch_q = torch.randn(512 + 1024, 16, 64, device="cuda", dtype=torch.bfloat16)  # (total_tokens, heads, d_head)
batch_k = torch.randn_like(batch_q)
batch_v = torch.randn_like(batch_q)

# Cumulative sequence lengths (must start at 0 and end at total_tokens)
cu_seqlens = torch.tensor([0, 512, 512 + 1024], dtype=torch.int32, device="cuda")
max_seqlen = 1024

out = flash_attn_varlen_func(
    batch_q, batch_k, batch_v,
    cu_seqlens_q=cu_seqlens,
    cu_seqlens_k=cu_seqlens,
    max_seqlen_q=max_seqlen,
    max_seqlen_k=max_seqlen,
    causal=True,             # causal within each segment
)
# out.shape == (total_tokens, heads, d_head)

The varlen interface ensures that attention computations never cross document boundaries, at zero mask-storage cost. In HuggingFace transformers the same effect is available without calling flash_attn directly: pass position_ids that restart at 0 per document and the FlashAttention-2 path unpads and builds cu_seqlens for you (TRL exposes this as padding_free=True in SFTConfig).

Position IDs Must Reset Too¶

Masking is only half of document-aware packing. If the second document in a packed window still receives positions \(5000, 5001, \ldots\), then RoPE encodes it as though it began five thousand tokens into a document — and worse, a packed window at 32 K exercises positions up to 32 K only if some single document is that long. With per-document position resets, the largest position the model ever sees is the length of the longest document, not the length of the packed window. This is the single most common way a context-extension run silently trains nothing: you rescale RoPE for 32 K, spend 10 B tokens, and every position the model saw was below 4 K. Always assert on your repacked shards that position_ids.max() actually reaches into the newly extended range before launching.

The two conventions therefore pair up: reset positions per document and mask attention per document (what varlen does), or keep global positions and allow cross-document attention (the naive scheme). Mixing them — global positions with document masking — wastes the long-range positions you paid to train.

Loss Masking¶

The cross-entropy loss should also exclude the first token of each document’s continuation from predicting the last token of the previous document. A simple loss_mask tensor of 0/1 per token accomplishes this — set the first token of each document to 0 (do not compute loss for it). See also Chat Templates, Data Formatting & Sequence Packing for how the same principle applies to fine-tuning packing.

Evaluating Long-Context Models: Needle in a Haystack¶

The Test¶

Needle in a Haystack (NIAH) is the canonical evaluation for long-context retrieval fidelity. A single “needle” sentence containing a unique fact (e.g., “The magic phrase is pineapple-strawberry-7.”) is inserted at a controlled depth (e.g., 25%, 50%, 75%) inside a long “haystack” of unrelated text (often Paul Graham essays or Wikipedia). The model must retrieve the needle after reading the full context.

The result is plotted as a 2D heatmap: the x-axis is context length, the y-axis is needle depth, and color indicates whether the model retrieved it correctly. A well-trained model shows uniform success across the entire grid; a poorly extended model shows failures at long contexts and specific depths.

import random

def build_haystack_with_needle(
    haystack_text: str,
    needle: str,
    context_length: int,           # target total tokens (approx by words)
    depth_fraction: float,         # 0.0 = beginning, 1.0 = end
    tokenizer_approx_ratio: float = 1.3,  # tokens per word (rough)
) -> tuple[str, int]:
    """
    Insert a needle sentence at a given depth in the haystack.
    Returns (full_text, needle_start_word_index).
    """
    target_words = int(context_length / tokenizer_approx_ratio)

    # Tile haystack words to fill the context
    words = haystack_text.split()
    if len(words) < target_words:
        multiplier = (target_words // len(words)) + 1
        words = words * multiplier
    words = words[:target_words]

    # Determine insertion point
    insert_idx = int(len(words) * depth_fraction)

    # Insert needle words at that point
    needle_words = needle.split()
    words = words[:insert_idx] + needle_words + words[insert_idx:]

    return " ".join(words), insert_idx

def evaluate_niah(
    model_fn,           # callable(text: str) -> str
    needle: str,
    answer_keyword: str,
    context_lengths: list[int],
    depths: list[float],
    haystack: str,
) -> dict:
    """
    Run a full NIAH evaluation grid.
    Returns dict[(context_len, depth)] -> bool
    """
    results = {}
    for ctx_len in context_lengths:
        for depth in depths:
            text, _ = build_haystack_with_needle(haystack, needle, ctx_len, depth)
            prompt = text + "\n\nWhat is the magic phrase?"
            response = model_fn(prompt)
            results[(ctx_len, depth)] = answer_keyword.lower() in response.lower()
    return results

NIAH is a necessary but not sufficient evaluation: a model can pass NIAH yet still fail at multi-hop reasoning over long contexts because NIAH only requires retrieval of a single salient fact. More challenging variants scatter multiple needles or require combining information from two distant locations. The original public implementation is gkamradt/LLMTest_NeedleInAHaystack; the heatmap above is its signature output.

Beyond NIAH: Effective Context Length¶

The benchmark that superseded plain NIAH is RULER (Hsieh et al., NVIDIA, 2024; repo NVIDIA/RULER). RULER synthesises tasks at any configurable length, so there is no test-set contamination and no length ceiling, and it goes well past retrieval:

• multi-key / multi-value / multi-query needle retrieval (distractor needles present),
• variable tracking (follow a chain of X = Y assignments scattered through the context),
• frequent-word extraction (aggregation, not retrieval — the answer is not any single span),
• long-context QA over injected paragraphs.

RULER’s key contribution is the notion of effective context length: the longest input length at which a model still clears a fixed quality bar (set by a strong short-context baseline). Its headline result is uncomfortable and still holds — most models advertising 32 K–128 K windows have effective lengths well below their claimed window, because aggregation and tracking degrade long before single-needle retrieval does. Report your model’s effective length, not its config’s max_position_embeddings.

Complements worth running: LongBench v2 (realistic bilingual tasks, multiple-choice so it is cheap to score), HELMET (a 2024–2025 suite that argues synthetic retrieval correlates poorly with downstream long-context ability and spans RAG, summarisation, and many-shot ICL), and simple per-position perplexity — plot loss as a function of token index within the window. A healthy extended model shows loss still decreasing with position out to the target length; a flat or rising tail past the old window means the extension did not take. That last check needs no benchmark at all and should be the first thing you run.

Ring Attention and Context Parallelism for Training¶

Why We Need a New Parallelism Dimension¶

Even with FlashAttention, training a sequence of 128 K tokens on a single GPU is impossible: the tiled attention still needs \(\mathcal{O}(T)\) intermediate buffers per layer. At 128 K tokens in bf16, the KV activations for one layer with \(d_{\text{model}} = 4096\) are:

For 80 layers that is 160 GB, exceeding a single A100 80 GB GPU. We need to shard the sequence across devices — a different axis than the data, tensor, or pipeline parallelism of Distributed Training II.

Ring Attention¶

Ring Attention (Liu et al., 2023, Ring Attention with Blockwise Transformers) achieves sequence parallelism by distributing non-overlapping chunks of the sequence across \(N\) GPUs arranged in a logical ring. Each GPU holds \(T/N\) tokens’ keys and values locally. Attention is computed as follows:

Ring of 4 GPUs, each holding T/4 tokens:

Step 0: GPU-i computes attention of its Q chunk against its own KV chunk.
Step 1: Each GPU sends its KV chunk to the next GPU in the ring.
        While the data is in flight, each GPU continues computing
        attention of its Q chunk against the KV chunk it is about to
        receive (overlap communication with computation).
Step 2: Repeat for N-1 total rounds.
After N rounds, each GPU has accumulated the full softmax-normalised
attention output for its Q chunk.

The key algorithmic insight is that FlashAttention’s online softmax accumulation is associative: you can accumulate the numerator sums and the log-sum-exp denominators independently across rounds and produce the exact same result as if all KV blocks were available at once. No approximation is involved.

Communication volume per round: \(2 \times (T/N) \times d \times 2\;\text{bytes}\) per direction (sending K and V). With \(N = 8\) GPUs and \(T = 128{,}000\), \(d = 4096\): each round moves \(128{,}000/8 \times 4096 \times 2 \times 2 = 256\;\text{MB}\) per GPU. On NVLink at 400 GB/s this takes ~0.6 ms; a single attention kernel at this scale takes much longer, so the overlap is effective.

# Pseudocode sketch of ring attention forward pass
# (actual implementation uses NCCL send/recv and a FlashAttention tile kernel)

import torch
import torch.distributed as dist

def ring_attention_forward(
    q_local: torch.Tensor,     # (T/N, H, D) — this GPU's query chunk
    k_local: torch.Tensor,     # (T/N, H, D) — this GPU's key chunk (starting KV)
    v_local: torch.Tensor,     # (T/N, H, D) — this GPU's value chunk
    causal: bool,
    rank: int,
    world_size: int,
) -> torch.Tensor:
    """
    Simplified ring attention. Real implementations fuse this with FlashAttention.
    """
    T_local = q_local.shape[0]
    d_model = q_local.shape[-1]

    # Accumulators for online softmax (see FlashAttention chapter for derivation)
    output_acc = torch.zeros_like(q_local)        # O accumulator
    lse_acc = torch.full((q_local.shape[0], q_local.shape[1]), float('-inf'))  # log-sum-exp

    k_chunk = k_local.clone()
    v_chunk = v_local.clone()
    # We track which absolute position offset these k/v come from
    kv_rank = rank

    for step in range(world_size):
        # Determine the global position ranges of this Q chunk and the current KV chunk.
        q_global_start  = rank * T_local
        kv_global_start = kv_rank * T_local

        # CAUSAL BLOCK SKIP (critical for correctness):
        # If this entire KV chunk lies in the future of every local query, no query
        # may attend to any of it. Building an all -inf score row and calling softmax
        # returns NaN (exp(-inf)/sum(exp(-inf)) = 0/0), and NaN * 0 in the merge step
        # (exp_new = 0) then poisons output_acc. So we skip the whole block. Because
        # ring rotations are block-aligned, the only partially masked block is the
        # diagonal (own) block, which is always processed first; every other processed
        # block is strictly in the past, so no score row is ever fully -inf once
        # future blocks are skipped. Production kernels (zigzag / striped ring
        # attention) skip these blocks the same way -- which is exactly what creates
        # the causal load imbalance across the ring (rank 0 processes 1 block,
        # rank N-1 processes N) discussed in Distributed Training II.
        fully_future = causal and (kv_global_start > q_global_start + T_local - 1)

        if not fully_future:
            # Compute attention of q_local against k_chunk / v_chunk
            # (in reality this calls a FlashAttention tile kernel)
            scale = d_model ** -0.5
            scores = torch.einsum("thd,khd->thk", q_local, k_chunk) * scale  # (T, H, T_kv)

            if causal:
                # Mask out future KV positions relative to each Q position. Since
                # fully-future blocks are skipped above, cutoff > 0 always holds
                # here, so no score row is ever entirely -inf.
                for t in range(T_local):
                    q_pos = q_global_start + t
                    cutoff = q_pos - kv_global_start + 1   # keep kv positions [0, cutoff)
                    if cutoff < T_local:
                        scores[t, :, cutoff:] = float('-inf')

            # Online softmax accumulation (simplified; real code tracks m and l)
            new_lse = torch.logsumexp(scores, dim=-1)                    # (T, H)
            new_out = torch.einsum("thk,khd->thd",
                                    torch.softmax(scores, dim=-1), v_chunk)  # (T, H, D)

            # Merge with running accumulator
            m = torch.maximum(lse_acc, new_lse)
            exp_old = torch.exp(lse_acc - m)
            exp_new = torch.exp(new_lse - m)
            lse_acc = m + torch.log(exp_old + exp_new)
            output_acc = (output_acc * exp_old.unsqueeze(-1) + new_out * exp_new.unsqueeze(-1)) / (exp_old + exp_new).unsqueeze(-1)

        # Rotate KV to next GPU in ring (overlap with next step's compute in real code)
        next_rank = (rank + 1) % world_size
        prev_rank = (rank - 1) % world_size
        dist.send(k_chunk, dst=next_rank)
        dist.recv(k_chunk, src=prev_rank)
        dist.send(v_chunk, dst=next_rank)
        dist.recv(v_chunk, src=prev_rank)
        kv_rank = (kv_rank - 1) % world_size  # the received chunk came from one step earlier

    return output_acc

Verifying correctness: run this ring implementation on a toy problem (e.g. world_size=4, T_local=8, H=2, D=16, random q/k/v in fp32) under torchrun --nproc_per_node=4, gather the per-rank output_acc back into a full (T, H, D) tensor in rank order, and compare it against a single-device reference over the un-sharded sequence: torch.nn.functional.scaled_dot_product_attention(q, k, v, is_causal=True). The two should agree to within ~1e-5 in fp32 (~1e-2 in bf16). This is also the fastest way to catch the causal-masking bug: without the fully_future skip, rank 0 – whose queries may attend only to their own block – still receives every other rank’s future KV chunk, softmaxes an all -inf row into NaN, and returns NaN for its entire output, so the reference comparison fails immediately with nan.

The Libraries That Implement This¶

Nobody ships the loop above by hand. In practice you pick one of four stacks:

Ring vs Ulysses. DeepSpeed-Ulysses attacks the same problem from the other side: instead of rotating KV around a ring, it performs an all-to-all just before attention that converts a sequence-sharded tensor into a head-sharded one, runs ordinary full-sequence attention on a subset of heads, then all-to-alls back. Its advantages are that it needs no custom attention kernel (any FlashAttention build works unmodified) and its communication per layer is \(\mathcal{O}(T d)\) regardless of world size. Its constraint is severe: the parallel degree must divide the number of KV heads, which with GQA (often 8) caps you at 8 devices. Ring attention has no head constraint and scales to arbitrary sequence length, at the cost of \(N-1\) communication rounds that must be overlapped with compute. Hybrid 2D schemes that put Ulysses inside a node and ring across nodes — “unified sequence parallelism”, implemented in libraries such as yunchang — get the best of both and are what large 1M-token runs typically use.

Sequence Parallelism in Megatron-LM¶

Megatron-LM (see Megatron-LM, DeepSpeed & Parallelism in Practice) has a separate “sequence parallelism” mode that shards the LayerNorm and Dropout computations along the sequence dimension while using tensor parallelism for the attention projections. This is distinct from ring attention: it reduces activation memory for non-attention operations by splitting the sequence across tensor-parallel ranks, using all-gather and reduce-scatter at the sequence axis.

In practice, the two approaches compose: one might use tensor parallelism for weight matrices, ring/context attention for the attention itself, and pipeline parallelism for layers — the “4D parallelism” configuration used by some large-scale long-context runs.

Connecting Long-Context Training to the Inference KV Cache¶

Training and inference share the same fundamental tension: both need to store or compute over \(T\) keys and values per layer. The training solution (ring attention, gradient checkpointing) addresses the activation memory during the forward-backward pass. The inference solution (paged/chunked KV cache) addresses the persistent memory across generation steps.

The positional encoding choices made during training directly constrain what the inference engine can do:

• RoPE base and scaling method: The inference engine must apply the same RoPE modification as training. A mismatch (e.g., serving with standard RoPE a model trained with YaRN) produces garbage scores.
• Maximum context length: The inference engine’s KV cache pre-allocates memory up to max_model_len. Extending this at inference time without retraining is limited by the interpolation range of the chosen scaling method.
• GQA vs MHA: GQA (8 KV heads instead of 64) cuts the KV cache footprint 8x and is a prerequisite for long contexts at scale. See Multi-Head Attention, MQA, GQA & MLA.

For chunked prefill — processing a 128 K context in chunks rather than all at once at inference startup — see Disaggregated Prefill/Decode & Chunked Prefill.

Practical Training Recipe¶

Putting the above together, here is a concrete recipe to extend a pretrained 7 B RoPE model from 4 K to 32 K:

# 1. Curate a long-context continued-pretraining dataset.
#    Mix: 60% books/papers/repo-concatenated code >16K tokens,
#         40% original pretraining mix (short docs) to prevent regression.
#    Target total: ~10B tokens, RE-PACKED at seq_len = 32768 with
#    per-document position_ids and a document attention mask.

# 2. Declare the RoPE scaling in the model config (HuggingFace transformers).
#    In config.json:
#      "rope_scaling": {"rope_type": "yarn", "factor": 8.0,
#                       "original_max_position_embeddings": 4096},
#      "max_position_embeddings": 32768
#    ("rope_type" is the current key; older configs use "type". Other accepted
#     values: "linear" (= PI), "dynamic" (NTK), "longrope", "llama3".)

# 3. Launch continued pretraining with a reduced LR.
#    learning_rate 2e-5 is ~2% of a typical ~1e-3 pretraining peak LR.
#    Packing and the per-document loss mask live in your collator (see the
#    "Document Packing" section); they are not Trainer flags.
torchrun --nproc_per_node=8 train.py \
  --model_name_or_path meta-llama/Llama-2-7b-hf \
  --max_seq_length 32768 \
  --per_device_train_batch_size 1 \
  --gradient_accumulation_steps 4 \
  --learning_rate 2e-5 \
  --lr_scheduler_type cosine \
  --num_train_epochs 1 \
  --warmup_steps 20 \
  --bf16 \
  --gradient_checkpointing \
  --attn_implementation flash_attention_2

# Beyond ~32K, add a context-parallel degree rather than shrinking the batch
# further -- e.g. under Megatron-LM:
#   --context-parallel-size 4 --tensor-model-parallel-size 2

# 4. Validate with NIAH sweep before serving
import matplotlib.pyplot as plt
import numpy as np

def plot_niah_heatmap(results: dict, context_lengths: list, depths: list):
    """
    Visualise needle-in-haystack results as a 2D heatmap.
    results: dict[(ctx_len, depth)] -> bool
    """
    grid = np.array([
        [float(results.get((c, d), False)) for c in context_lengths]
        for d in depths
    ])
    fig, ax = plt.subplots(figsize=(12, 5))
    im = ax.imshow(grid, vmin=0, vmax=1, cmap="RdYlGn", aspect="auto")
    ax.set_xticks(range(len(context_lengths)))
    ax.set_xticklabels([f"{c//1000}K" for c in context_lengths], rotation=45)
    ax.set_yticks(range(len(depths)))
    ax.set_yticklabels([f"{int(d*100)}%" for d in depths])
    ax.set_xlabel("Context Length")
    ax.set_ylabel("Needle Depth")
    ax.set_title("Needle-in-Haystack Retrieval Accuracy")
    plt.colorbar(im, ax=ax)
    plt.tight_layout()
    plt.savefig("niah_heatmap.png", dpi=150)
    print("Saved niah_heatmap.png")
    return grid.mean()  # return overall accuracy

Key Takeaways¶

Further Reading¶

• Su et al., “RoFormer: Enhanced Transformer with Rotary Position Embedding” (2022) — original RoPE paper; introduces the rotation formulation and relative-position property.
• Chen et al., “Extending Context Window of Large Language Models via Positional Interpolation” (2023) — the PI approach; simple and highly cited.
• Peng et al., “YaRN: Efficient Context Window Extension of Large Language Models” (2023) — introduces frequency-ramped interpolation and temperature correction; practical state of the art.
• Liu et al., “Ring Attention with Blockwise Transformers for Near-Infinite Context” (2023) — ring attention paper; the foundational sequence-parallelism approach.
• Dao et al., “FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning” (2023) — the varlen API for document packing is documented here.
• Anthropic, “Claude’s Long Context” (technical reports, 2024) — high-level discussion of training recipes for very long contexts (100 K+).
• LongBench (Bai et al., 2023) and LongBench v2 (2024) — comprehensive long-context evaluation benchmarks covering summarisation, QA, few-shot learning, and code; more demanding than needle-in-haystack alone.
• Hsieh et al., “RULER: What’s the Real Context Size of Your Long-Context Language Models?” (NVIDIA, 2024) — synthetic, contamination-free, length-configurable tasks; introduces effective context length.
• HELMET (Princeton, 2024–2025) — argues synthetic retrieval correlates poorly with downstream long-context ability; evaluates RAG, many-shot ICL, summarisation and citation tasks instead.
• Jacobs et al., “DeepSpeed-Ulysses” (2023) — all-to-all sequence parallelism, the main alternative to ring attention.
• LLaMA-3 technical report (Meta, 2024) — describes the multi-stage context extension recipe used in a public large model, including data mix and positional encoding choices.

Exercises¶

1. The chapter’s table lists a 16 MB (fp16) score matrix per head at \(T = 2{,}048\). Using the \(\mathcal{O}(T^2)\) scaling of the attention score matrix, (a) compute the per-head score-matrix size at \(T = 16{,}384\), (b) compute the total across a 32-head layer, and © state in one sentence why FlashAttention removes this memory cost even though the attention FLOPs remain \(\mathcal{O}(T^2)\).

2. Explain, grounded in the chapter, why (a) learned absolute positional embeddings and (b) sinusoidal embeddings both fail when a model trained to length \(T_{\text{train}}\) is run at a position \(T_{\text{train}} + k\). Then © explain why standard RoPE also degrades past \(T_{\text{train}}\), even though RoPE produces a mathematically valid rotation at any position, and (d) identify which specific mechanism Position Interpolation exploits to fix this.

3. A model has \(d = 128\), RoPE base \(b = 10{,}000\), and was trained at \(T_{\text{train}} = 4{,}096\). You want to extend to \(T_{\text{target}} = 32{,}768\) (an 8x extension).

(a) Under Position Interpolation, what effective position does the last token \(m = 32{,}768\) map to, and what effective positions do the adjacent early tokens \(m = 1\) and \(m = 2\) map to? What practical problem does the latter reveal?

(b) Under NTK-aware interpolation, compute the rescaled base \(b'\) using the chapter’s formula \(b' = b \cdot \left(T_{\text{target}}/T_{\text{train}}\right)^{d/(d-2)}\).

© In one sentence, contrast what PI and NTK-aware scaling each do to the highest-frequency dimension (\(i = 0\)).

4. The chapter describes a per-token loss_mask that sets the first token of each packed document to 0 (so the loss never rewards predicting a document’s opening token from the previous document’s final token), but gives no code for it. Implement build_packed_loss_mask(doc_lengths) in the same style as build_packed_attention_mask, returning a 1-D float tensor of length sum(doc_lengths) where the first token of each document is 0.0 and all other tokens are 1.0. Show the output for doc_lengths = [4, 4].

import torch

def build_packed_loss_mask(
    doc_lengths: list[int],
    device: str = "cpu",
) -> torch.Tensor:
    """
    Per-token loss weights for a packed multi-document sequence.
    The first token of every document is masked out (0.0); every other
    token contributes to the loss (1.0). Prevents rewarding the model for
    'predicting' a new document's first token from the previous document.
    """
    seq_len = sum(doc_lengths)
    mask = torch.ones(seq_len, dtype=torch.float, device=device)

    offset = 0
    for length in doc_lengths:
        mask[offset] = 0.0     # first token of this document
        offset += length

    return mask  # (seq_len,)

# Example: two docs of length 4 packed into one context of length 8
print(build_packed_loss_mask([4, 4]))
# tensor([0., 1., 1., 1., 0., 1., 1., 1.])

5. Ring attention shards a \(T = 128{,}000\)-token sequence across \(N = 8\) GPUs (\(d = 4096\), bf16), each holding \(T/N\) tokens.

(a) Reproduce the chapter’s per-round, per-GPU communication volume (each round sends one K chunk and one V chunk).

(b) With causal masking, the pseudocode skips any KV block that is entirely in the future of a rank’s local queries (fully_future). Show that rank \(r\) ends up actually processing exactly \(r+1\) of the \(N\) blocks, and use this to give the total number of blocks computed across the whole ring versus the naive \(N^2\). Then write a one-line function blocks_processed(rank) and state the resulting load-imbalance ratio between the busiest and idlest rank.

def blocks_processed(rank: int) -> int:
    return rank + 1