Model cards advertise context windows of 128K tokens, a million, ten million. Two questions hide behind each number. Can the model run on that many tokens, with its positions, its memory, and its training budget? And can it use them, finding and combining what it needs from anywhere in the window? Both matter more each year, as agents keep long tool transcripts in context and reasoning models spend their test-time compute on long chains of thought (§18.7). The first question has engineering answers, many of which this chapter has already built. The second is an empirical question, and the honest answer is often "less than advertised".

This section follows one spine, from the simplest wall to the hardest. First we count why long context is hard at all, with real numbers for one model. Then we fix the wall that breaks first, positions, by rescaling the rotary embedding's frequencies three ways (position interpolation, NTK-aware scaling, and YaRN, which is also what Llama 3.1 ships), and see how long-context models are trained. We recall how the sparse and local patterns of §12.6 serve long context. Then we split one sequence across many GPUs with ring attention and its relatives, forward and backward. And finally we measure what a model really does with its context, from needle-in-a-haystack to RULER.

Why long context is hard

Take Llama 3 8B (32 layers, width $d = 4{,}096$, 8 key-value heads of width 128, feed-forward width 14,336) and stretch its context from 8K tokens to a million. Four things go wrong, and each has its own fix.

1. Compute. The weight matrices cost the same per token at any length, but each token's attention looks at every earlier token, so its cost per token grows with $T$. With a causal mask a token sees $T/2$ keys on average, and the two products of the attention core cost $2\cdot 2\cdot(T/2)\cdot d = 2Td$ FLOPs per token per layer (§4.7).
2. Memory. The KV cache of one sequence grows linearly (§6.5). In training, the activations that the backward pass needs grow linearly too, about $34Td$ bytes per layer once FlashAttention has removed the $T\times T$ term, and only $2Td$ if each layer keeps nothing but its input and recomputes the rest (§10.1).
3. Positions. A model trained on 8K tokens has never seen a query 100,000 positions after its key. Its positional encoding must be taught or tricked into handling distances it never saw.
4. Use. Even a model that runs at a length may not read it well, because long documents are rare in training data and long distances are rarely exercised.

At 8K tokens attention is 13% of the forward FLOPs and the activations of one sequence take 34 GiB. At 128K attention is 71% of the FLOPs, the cache is 16 GiB, and the activations are 544 GiB, far more than any GPU holds. Even with every layer checkpointed, the 32 GiB of layer inputs sit on top of the weights, gradients, and optimizer state, about 128 GB for an 8B model trained with Adam in mixed precision (§10.1). At a million tokens attention is 95% of the work and the checkpointed layer inputs alone are 256 GiB.

The chapter has already removed some of this. FlashAttention made memory linear in $T$ (§12.3), sharing and compressing keys and values shrank the cache (§12.5), and sparse patterns cut the compute (§12.6). What remains is positions, splitting the work across devices, and measuring use, in that order.

Positions, the wall that breaks first

A model trained at one length has seen only the relative distances that fit inside it, and RoPE turns each distance into a set of angles. So the first question is which angles a longer context produces that the model has never seen.

Why a 4K model fails at 32K

Recall RoPE from §5.5. Each head's queries and keys are split into $d_k/2$ coordinate pairs, and pair $j$ at position $m$ is rotated by the angle $m\,\omega_j$, with frequencies

for a base $b$ (10,000 in the original RoPE). The score between a query at position $m$ and a key at position $n$ depends on the angles $(m - n)\omega_j$. Each pair has a wavelength $\lambda_j = 2\pi/\omega_j = 2\pi b^{2j/d_k}$, the number of positions it takes to rotate once.

Train on sequences of length $L$ and look at what each pair has seen. A pair whose wavelength is much shorter than $L$ has rotated many times and has seen every angle. A pair whose wavelength is longer than $L$ has seen only an arc of its circle. Run the model at a longer length $L' \gt L$ and those slow pairs produce angles they were never trained on, and the model's attention patterns break. With $d_k = 128$, $b = 10{,}000$, and $L = 4{,}096$, that is 18 of the 64 pairs.

This is a failure of extrapolation. Press et al. (2022) found that rotary models extrapolate only slightly beyond their training length, and Chen et al. (2023) found that a LLaMA model run past its training length without any adjustment reached perplexities above 1,000. Softmax dilution makes it worse (§4.7). With eight times as many keys, every head's weights spread over eight times as many competitors.

All three methods below change the frequencies so that the angles at long positions look like angles the model has seen. They differ in which pairs they change, and by how much.

Position interpolation, squeeze every position

Chen et al. (2023) proposed the simplest possible fix. To extend from $L$ to $L' = sL$ with scale factor $s$, divide every position by $s$ before rotating, which is the same as dividing every frequency by $s$,

In plain words, positions $0, \dots, L'$ are mapped onto the range $0, \dots, L$ that the model saw in training, so no pair ever rotates further than it did then. Nothing is extrapolated, only interpolated, at fractional positions between the integer ones it knew.

Chen et al. also bound how much interpolation can disturb the attention scores, compared with extrapolation (their Theorem 2.1). The interpolated scores are smooth functions of position, and the paper's bound on their deviation is at least about 600 times smaller than the corresponding bound for extrapolation with the original base. The price is resolution. Neighboring tokens are now $s$ times closer in angle, so the fast pairs, which tell adjacent positions apart, lose precision. That is why PI needs fine-tuning. With it, the results were striking. LLaMA models from 7B to 65B were extended to 32,768 tokens within 1,000 fine-tuning steps, and on a passkey-retrieval test the models reached the full extended window after only 200 steps, while direct fine-tuning without interpolation had moved the effective window only from 2,048 to 2,560 tokens after more than 10,000 steps.

NTK-aware scaling, spare the high frequencies

PI compresses every pair by the same factor. But the fast pairs never needed it, since they had already seen every angle. Within days of the PI paper, a user posting as bloc97 on Reddit proposed changing the base instead, so that the slowest pair is compressed by $s$ while the fastest is left almost alone. Require the last pair, $j = d_k/2 - 1$, to slow down by exactly $s$,

With the new base, pair $j$'s frequency is divided by $(b'/b)^{2j/d_k} = s^{2j/(d_k-2)}$, which grows smoothly from 1 for $j = 0$ to $s$ for the last pair. In plain words, raising the base slows the slow pairs a lot and the fast pairs hardly at all, so nearby positions stay as distinguishable as before while distant ones are brought into range. The name comes from its motivation. Neural tangent kernel theory suggests networks have trouble learning high-frequency information, so compressing the high frequencies, as PI does, throws away something hard to relearn. The original post extended LLaMA 7B beyond 8K tokens without any fine-tuning.

Two refinements followed quickly. Dynamic NTK scaling (emozilla, 2023) sets $s$ from the current sequence length instead of fixing it, so short inputs run with the original frequencies and long inputs get only as much stretching as they need. And the YaRN paper pointed out a flaw. Because NTK-aware scaling is not a pure interpolation, some middle pairs are pushed slightly past the angles seen in training, and fine-tuning with it gave worse results than PI.

import numpy as np

def scaled_rope_freqs(d_k, base, s, method):
    return base ** (-np.arange(0, d_k, 2) / d_k) / s      # PI for every method

def rope_rotate(x, positions, freqs):
    return x

YaRN, interpolate by wavelength

YaRN (Peng et al. 2023) makes the per-pair decision explicit. For each pair, count how many full turns it makes within the original context,

A pair with many turns ($r_j$ large) has seen every angle and should be left alone. A pair with less than one turn ($r_j$ small) is exactly the kind that extrapolates badly and should be interpolated fully, as in PI. In between, blend. With two thresholds $\alpha \lt \beta$, define a ramp and the new frequency,

In plain words, pairs that turn fewer than $\alpha$ times in the training window are interpolated by $s$, pairs that turn more than $\beta$ times keep their frequency, and pairs in between are a linear blend. For LLaMA-family models the paper recommends $\alpha = 1$ and $\beta = 32$. The paper calls this part "NTK-by-parts".

YaRN adds one more ingredient. Peng et al. found that dividing the attention logits by a temperature $t$ lowered perplexity uniformly, across samples and positions, and fitted

which is about 1.21 at $s = 8$. Implementations multiply both the rotated queries and the rotated keys by $\sqrt{1/t}$, so the logits are multiplied by $1/t$, which costs nothing at inference because it can be folded into the cached cosines and sines. A reasonable reading of why it helps is the softmax dilution of §4.7. Longer contexts flatten attention, and a slightly lower temperature sharpens it back.

The results made YaRN the default way to extend RoPE models. Starting from Llama 2 7B and 13B (4K context), 400 fine-tuning steps on 64K-token sequences gave $s = 16$ models, and 200 more steps from there gave $s = 32$ models that worked at 128K tokens, even though they were still trained on 64K-token sequences, with over 99% passkey retrieval across the full 128K window. The paper reports needing 10 times fewer tokens and 2.5 times fewer training steps than earlier methods, under 0.1% of the original pretraining data. DeepSeek-V2 and V3, gpt-oss, Kimi K2, and Qwen's long-context variants all ship with YaRN scaling in their configurations.

Llama 3.1 uses the same ramp

Llama 3.1's configuration lists rope_scaling as {"rope_type": "llama3", "factor": 8.0, "low_freq_factor": 1.0, "high_freq_factor": 4.0, "original_max_position_embeddings": 8192} with base 500,000, and the Llama 3 paper does not describe it. The code (in Meta's reference implementation and in Hugging Face Transformers) keeps a frequency unchanged if its wavelength is below $8192/4 = 2048$, divides it by 8 if its wavelength is above $8192/1 = 8192$, and in between blends with weight

Since $8192/\lambda_j$ is exactly $r_j$, this is YaRN's ramp with $\alpha = 1$ and $\beta = 4$, and with no temperature. The cell confirms it numerically.

import numpy as np

def yarn(d_k, base, s, L_orig, alpha=1.0, beta=32.0):
    w = base ** (-np.arange(0, d_k, 2) / d_k)
    return w / s, 0.1 * np.log(s) + 1                    # PI frequencies, and the square root of the scale

Training for long context

Frequency scaling makes a model able to attend far back. Training is what makes it good at it, and the recipes have converged.

Raise the base. Xiong et al. (2023) continued pretraining Llama 2 on 400B tokens with the RoPE base raised from 10,000 to 500,000 ("adjusted base frequency"), reaching 32,768 tokens for the 7B and 13B models and 16,384 for the larger ones. Code Llama (Rozière et al. 2023) used a base of 1,000,000, fine-tuned on 16,384-token sequences, and stayed stable up to about 100,000 tokens. Llama 3 used a base of 500,000 from the start of pretraining. Raising the base is the same move as NTK-aware scaling, made for training rather than patched in at inference. Xiong et al. motivate it by RoPE's decay of attention with distance (§5.5), which a larger base weakens for distant tokens. Note that it does not by itself remove the extrapolation problem. A larger base slows every pair, so more pairs fail to complete a turn within a short training window (32 of 64 at 4K tokens with base 500,000, by the wavelength cell above), which is why it is paired with training at the long length.

Extend in stages. Llama 3 pretrained at 8K tokens and then increased the context in six stages up to 128K, using about 800B tokens of long-context training, moving to the next stage only once short-context performance had recovered and long-context retrieval worked at the current length. Shorter recipes work for smaller budgets. Fu et al. (2024) found that 500 million to 5 billion tokens of continued pretraining suffice to extend a 7B model to 128K, provided each source is upsampled toward long documents while keeping the original mix of domains, rather than, say, training mostly on books.

Sparse and local patterns at long context

The patterns of §12.6 are the other half of most long-context designs, because they attack the compute and cache walls directly. Mistral 7B's 4,096-token window keeps the cache of a 32K-token sequence at an eighth of its full size. Gemma 3 reaches 128K tokens with 5 local layers of window 1,024 for every global one, which cuts its long-context cache to about a sixth, and it treats positions differently in the two kinds of layer, keeping the RoPE base at 10,000 in the local layers and raising it to 1,000,000 in the global ones (Gemma Team 2025). DeepSeek-V3.2-Exp's sparse attention reads 2,048 selected tokens per query at 128K context, and attention sinks keep windowed generation stable past any training length.

Two limits follow from what those patterns are. They do not remove the position problem, because whatever layers still look far back (the global layers, or the selected tokens) see distances they may never have been trained on, so they need the frequency fixes above. And they do not guarantee use. A local layer can only relay distant information through intermediate positions, and a selection method can only read what its scorer ranks highly, so a sparse long-context model must be measured as carefully as a dense one, which the last part of this section does.

One sequence on many GPUs

Even with FlashAttention, the activations of one 128K-token sequence do not fit on one GPU, as the first cell showed, so the sequence itself must be split. Almost every layer of a Transformer is per token (the projections, the MLP, the norms), so if device $p$ of $P$ devices holds a slice of $c = T/P$ tokens, it can run those layers on its slice with no communication at all. Attention is the one layer that objects, because every query needs every earlier key. Splitting the sequence this way is context parallelism, and how it sits on a cluster beside data, tensor, and pipeline parallelism is the subject of §10.6. Here we build the attention that makes it work.

There are three ways for a device to get the keys and values it lacks. It can gather every other device's keys and values. It can receive them one block at a time around a ring. Or the devices can swap the split, trading a slice of tokens for a slice of heads. We start with the ring, the most interesting of the three, and come back to the others at the end.

Ring attention, the forward pass

Ring attention (Liu, Zaharia, and Abbeel 2023) arranges the $P$ devices in a ring. Device $p$ keeps its own queries and starts with its own key-value block. At each of $P$ steps, every device computes attention between its queries and the key-value block it currently holds, merges the result into its running statistics with the online softmax of §12.2, sends the block to the next device, and receives one from the previous device. So at step $k$ device $p$ holds the block that started on device $(p - k) \bmod P$. After $P$ steps every query block has seen every key block, and because the merge is exact and order-free, the output is exactly attention. No device ever holds more than two key-value blocks, the one it computes with and the one arriving.

In plain words, ring attention is FlashAttention's inner loop with the key-value blocks arriving over a network link instead of from HBM. The running maximum, sum, and unnormalized output per query row are the same, and each device ends with its rows of $\mathbf{O}$ and their log-sum-exps.

When the ring's communication is free

The transfer to the next device can run while the current block is computed, so the question is whether it finishes in time. Count one step of one layer, with $h$ query heads, $n_{kv}$ key-value heads of width $d_k$, and bf16 numbers. Without a mask, the attention of $c$ queries against $c$ keys costs $4c^2d_k$ FLOPs per query head, $4c^2hd_k$ in all. The key and value block is $2cn_{kv}d_k$ numbers, $4cn_{kv}d_k$ bytes. With a device doing $F$ FLOP/s and a link moving $\beta_{\text{net}}$ bytes per second, the transfer is hidden when

In plain words, each device must hold at least as many tokens as the ratio of its compute speed to its link bandwidth, and grouped-query attention lowers the bar by the group size $h/n_{kv}$, because the blocks it sends are smaller while the work is not. Compute grows like $c^2$ and communication like $c$, so a long enough block always wins. Liu et al. state the condition for multi-head attention, $c \ge F/\beta_{\text{net}}$, about 1,000 tokens on the hardware they used, and report over a million tokens of context for a 7B model on 32 A100s and over 30 million on a TPU v4-512 slice. With a causal mask a balanced device does about half the work per step, so the bar roughly doubles.

A 128K-token sequence over 16 GPUs gives 8,192 tokens per device, well above both NVLink bars. Over InfiniBand only grouped-query attention without a causal mask clears it, which is why rings are usually laid out within a node's NVLink domain, or given longer blocks when they cross nodes. These are peak rates. A kernel running at half of peak halves the bar, which is one reason ring attention works better in practice than the worst case suggests.

Balancing the causal mask

With a causal mask the obvious layout is unfair. If device 0 holds the first $T/P$ tokens, its queries can see only its own block, so it skips $P - 1$ of its $P$ steps, while the last device's queries see everything. Every device waits for the slowest, so the ring runs at the speed of the last device. Two layouts fix it by giving every device a mix of early and late tokens.

• Striped attention (Brandon et al. 2023) assigns tokens round-robin, token $i$ to device $i \bmod P$. Every block now holds early and late tokens, every step has work, and the work is nearly equal. They measured up to 1.45 times end-to-end speedups at 256K tokens on 8 A100s.
• Zigzag cuts the sequence into $2P$ chunks and gives device $p$ chunks $p$ and $2P - 1 - p$, one from the front and its mirror from the back, so the cheap and the expensive chunk add up to the same work on every device. Llama 3's context parallelism uses this layout, and it keeps each device's tokens in two contiguous runs, which suits kernels and position ids better than striping.

The cell runs the forward pass on a simulated ring for all three layouts, checks it against causal attention, and counts the unmasked scores each device computes.

All three are exact. With contiguous blocks the busiest device does 1.87 times the average work (the ratio approaches 2 as $P$ grows), striping brings it to 1.03, and zigzag to exactly 1.

import numpy as np

def ring_attention(Q, K, V, P, layout="contiguous"):
    T = len(Q)
    S = np.where(np.tril(np.ones((T, T), bool)), Q @ K.T / np.sqrt(Q.shape[1]), -np.inf)
    E = np.exp(S - S.max(1, keepdims=True))
    return E @ V / E.sum(1, keepdims=True), np.full(P, P), np.full(P, T * (T + 1) // (2 * P))   # global softmax, no ring

The backward pass around the ring

Training needs gradients, and the backward pass has the same problem as the forward pass, every query's gradient involves every key. The tile formulas of §12.3 carry over unchanged. Device $p$ keeps, for its own rows, $\mathbf{Q}_p$, the output $\mathbf{O}_p$, the upstream gradient $\mathrm{d}\mathbf{O}_p$, the log-sum-exps $\operatorname{lse}_p$ that the forward pass left behind, and the row terms $\mathbf{D}_p = \operatorname{rowsum}(\mathrm{d}\mathbf{O}_p\odot\mathbf{O}_p)$. When the key-value block $j$ arrives, it computes

then adds $\mathrm{d}\mathbf{S}_{pj}\mathbf{K}_j/\sqrt{d_k}$ into its own $\mathrm{d}\mathbf{Q}_p$, and adds $\mathrm{d}\mathbf{S}_{pj}^\top\mathbf{Q}_p/\sqrt{d_k}$ and $\mathbf{P}_{pj}^\top\mathrm{d}\mathbf{O}_p$ into $\mathrm{d}\mathbf{K}_j$ and $\mathrm{d}\mathbf{V}_j$.

Two observations make this work. First, the backward pass needs no online softmax at all. The forward pass ended with each row's log-sum-exp over all keys, so $\mathbf{P}_{pj}$ computed from it is already the final, globally normalized weight, and each block's contribution is added as it is. Second, $\mathrm{d}\mathbf{Q}_p$ is a sum over key blocks, which all visit device $p$, so it accumulates locally. But $\mathrm{d}\mathbf{K}_j$ and $\mathrm{d}\mathbf{V}_j$ are sums over query blocks, which live on different devices. So the gradient accumulators travel with their block. Block $j$ goes around the ring carrying $(\mathbf{K}_j, \mathbf{V}_j, \mathrm{d}\mathbf{K}_j, \mathrm{d}\mathbf{V}_j)$, each device adds its contribution as the block passes through, and one final hop brings the finished gradients back to the block's owner.

In plain words, queries stay home and collect their own gradients, while keys and values travel and collect theirs. The price is communication. Each step now sends the gradient accumulators as well as the block, about twice the forward traffic in numbers, and more in bytes when, as usual, the accumulators are kept in float32. Compute per step is about 2.5 times the forward pass's (five tile products instead of two), so the overlap condition is no harder to meet.

import numpy as np

def ring_attention_backward(Q, K, V, O, dO, lse, P):
    raise NotImplementedError

Gathering, and splitting by heads instead

The ring is not the only way, and not always the best.

Gather the keys and values. Each device all-gathers every other device's key-value block, then computes attention for its own queries against the full $\mathbf{K}$ and $\mathbf{V}$. Llama 3's context parallelism does this, with the zigzag layout (Dubey et al. 2024). Its reasons are worth knowing. With grouped-query attention the keys and values are much smaller than the queries, so the all-gather is cheap, and holding the full $\mathbf{K}$ and $\mathbf{V}$ makes any mask, such as the document masks of packed sequences, easy to apply. The cost is that the all-gather is exposed rather than overlapped, and each device briefly holds the whole sequence's keys and values.

Swap the split. DeepSpeed-Ulysses (Jacobs et al. 2023) uses the fact that heads are independent. If each device held all tokens but only some heads, it could compute exact attention locally, with any kernel and any mask. So it switches layouts with an all-to-all. Each device starts with $c = T/P$ tokens of $\mathbf{Q}$, $\mathbf{K}$, $\mathbf{V}$ for all heads, and sends device $q$ the heads of group $q$, so afterwards device $q$ holds all $T$ tokens of $h/P$ heads. After ordinary attention, a second all-to-all sends the outputs back to the token split.

The counts explain the trade-off. Each all-to-all keeps $1/P$ of a device's slice and sends the rest, so Ulysses moves about $(2h + 2n_{kv})d_kT/P$ numbers per device per layer, which falls as $P$ grows. The ring and the all-gather deliver every other device's keys and values, about $2Tn_{kv}d_k$ numbers, which does not fall with $P$. But the ring hides its traffic behind compute, while Ulysses's all-to-alls sit on the critical path, and Ulysses cannot split further than the number of heads. With grouped-query attention it is the key-value heads that run out first, 8 for most Llama models, after which they must be replicated. That is why large systems often combine the two, Ulysses across the GPUs of a node, where all-to-alls are fast, and the ring across nodes (Fang and Zhao 2024).

A ring attention step in real code overlaps the send and receive of key-value blocks with the local attention computation. The sketch below uses torch.distributed point-to-point operations and the log-sum-exp merge of §12.2. It needs several GPUs, and production versions (such as the ring and zigzag implementations used with FlashAttention kernels) call a fused kernel for each block and run the backward pass around the ring as derived above.

Does the model use its context?

A model that runs at 128K tokens is not necessarily a model that reads 128K tokens. Measuring the difference has become its own small field.

Needle in a haystack. Kamradt (2023) hid one sentence, "The best thing to do in San Francisco is eat a sandwich and sit in Dolores Park on a sunny day", at a chosen depth inside a long stretch of Paul Graham's essays, asked the model what the most fun thing to do in San Francisco is, and plotted accuracy over a grid of context lengths and depths. The heatmaps exposed real failures at the time and became a standard figure in model announcements. Passkey retrieval, used to evaluate PI and YaRN, is the same idea with a random number hidden in filler text.

The trouble is that these tests are too easy. Finding one sentence with an unusual phrase is a lexical match that a model can solve with a single induction-style lookup (§4.7). Most frontier models now score near 100%, which says little about tasks that need several facts or reasoning over the whole context. A related failure is positional. Liu et al. (2023) placed the one relevant document among 10 to 30 retrieved ones and found accuracy highest when it came first or last and lowest in the middle, sometimes below the model's accuracy with no documents at all (§7.6).

RULER. Hsieh et al. (2024) built a synthetic benchmark of 13 tasks in four categories, with lengths that can be set freely. Retrieval extends the needle test with multiple keys, multiple values, multiple queries, and distracting near-duplicates. Multi-hop tracing asks the model to follow chains of variable assignments through the text. Aggregation asks for the most common words in a long list. Question answering hides the answer document among many distractors. They defined a model's effective context length as the longest tested length at which its average score stays above a fixed threshold, Llama 2 7B's score at 4K tokens (85.6%), and evaluated 17 models.

Almost every model's effective context was a fraction of its claim, and the authors found that only about half of the models kept satisfactory performance at 32K tokens, though nearly all passed the plain needle test.

Even RULER's retrieval tasks let the question and the hidden fact share words, which an attention head can match directly. NoLiMa (Modarressi et al. 2025) removes that overlap, so that connecting question and needle takes one step of inference (the question names a city, the needle a landmark in it). Of 12 models that claim at least 128K tokens, 10 fell below half of their short-context score at 32K, and even GPT-4o dropped from 99.3% to 69.7%. HELMET (Yen et al. 2024) makes the complementary point from the other side, that synthetic recall scores predict poorly how a model does on long-context applications, and evaluates seven categories of application-like tasks, among them retrieval-augmented generation, reranking, generation with citations, summarization, and many-shot in-context learning. The lesson for practice is the same in both. Test with tasks that look like yours, at the lengths you will use.

Why effective context falls short

An et al. (2024) offered a simple explanation, and it comes from counting. In a training sequence of length $L$, a relative distance $r$ between a query and an earlier key occurs $L - r$ times, so the distribution of training distances is skewed toward small ones,

In plain words, only about a quarter of the query-key pairs a model trains on are more than half a window apart, and the longest distances are seen almost never, since distance $L - 1$ occurs once per sequence. Long-distance positions are undertrained, so it is not surprising that effective lengths often fall to half the training length or less, as An et al. report. Their fix, STRING, shifts the well-trained position offsets to replace the undertrained ones at inference time, with no training. On RULER at 128K it raised Llama 3.1 70B from 66.6 to 81.7, slightly above GPT-4's 81.2.

def effective_length(lengths, scores, threshold):
    return max(L for L, s in zip(lengths, scores) if s >= threshold)   # ignores dips at shorter lengths

def far_pair_fraction(L, frac=0.5):
    return 1 - frac

In PyTorch

The frequency schedules are a few lines each, and the property that matters, that the score depends only on the offset between positions, survives every one of them, because each is still a rotation by position times frequency. The block below builds PI, NTK, YaRN, and Llama 3.1 frequencies in torch, applies them to random queries and keys, and checks the relative property and the YaRN temperature.

In a real model these frequencies replace the cached cosines and sines (and, for YaRN, are multiplied by the temperature factor), and the rest of the attention layer is unchanged. In Hugging Face Transformers, all of this is selected by the rope_scaling entry of a model's configuration, such as {"rope_type": "yarn", "factor": 4.0, "original_max_position_embeddings": 32768}.

Read more

• Chen et al. (2023). Extending Context Window of Large Language Models via Positional Interpolation. Position interpolation, its stability bound, and the passkey results.
• Peng et al. (2023). YaRN: Efficient Context Window Extension of Large Language Models. The clearest account of PI, NTK-aware, NTK-by-parts, and dynamic scaling, with the temperature.
• Xiong et al. (2023). Effective Long-Context Scaling of Foundation Models. Raising the RoPE base and continued pretraining, with careful ablations.
• Liu, Zaharia & Abbeel (2023). Ring Attention with Blockwise Transformers for Near-Infinite Context. Ring attention and the condition for overlapping communication.
• Brandon et al. (2023). Striped Attention: Faster Ring Attention for Causal Transformers. Why contiguous blocks unbalance the causal ring, and the round-robin fix.
• Jacobs et al. (2023). DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models. Sequence parallelism by all-to-all into a head split, with its communication analysis.
• Hsieh et al. (2024). RULER: What's the Real Context Size of Your Long-Context Language Models?. The benchmark and the effective-length tables.
• An et al. (2024). Why Does the Effective Context Length of LLMs Fall Short?. The position-frequency explanation and the STRING fix.

Watch next

▶

▶