Llama/GPTNeoX: add RoPE scaling

[gante]

What does this PR do?

This is an experimental PR for discussion, so we can decide whether to add this pattern.

Context

In the past week, there have been several developments about scaling RoPE (Rotary Position Embeddings, i.e. Llama's position embeddings) so as to be able to extrapolate beyond 2048 tokens. Without any scaling and/or finetuning, the perplexity quickly explodes when we go beyond 2048 tokens. Here's the sequence of RoPE scaling improvements, announced mostly on Reddit:

1. Linear scaling -- Simply divide the position index by a scaling factor. Needs fine-tuning to observe the best results. Discussed in this lmsys blog post. Credits to the reddit user /u/kaiokendev.
2. NTK-aware scaling -- proposed in this reddit thread. Scaling the RoPE Fourier space linearly is not optimal to evenly distribute information, so this can be seen as a improved linear scaling. Works okay without fine-tuning, but seems to benefit from it. Credits to the reddit user /u/bloc97. EDIT: following the comments in this thread, this technique will not be added!
3. Dynamic NTK scaling -- proposed in this reddit thread. It's a form of NTK-aware scaling that a) works quite well without fine-tuning; b) doesn't degrade the performance if the model is used with short sequences; c) gracefully scales to long sequences, under a fixed parameterization. Credits to the reddit user /u/emozilla.

Changes in the PR

The goal of this PR is to debate whether we want to include RoPE scaling support, with working code as reference. The field is evolving quite fast, so I've added it in a way that we can quicky add to new scaling strategies and keep surfing the wave 🏄 Of course, the implementation itself is up for discussion! (An alternative implementation would be to have separate classes for the scalable RoPEs)

Pros:

• Flexible implementation that allows adding new scaling methods in minutes;
• Works quite well with pre-trained models (see example below), through dynamic NTK scaling;
• Supports strategies that are compatible with fine-tuning (it is unclear whether dynamic NTK works well with fine-tuning, and it seems like Linear scaling is better after fine-tuning)

Cons:

• rope_scaling is a dictionary input, which is somewhat undesirable;
• additional if/else branches in RoPE

Example

Consider the following prompt from a paper transcript, containing ~6k tokens:

prompt built from the transcript of https://arxiv.org/abs/2306.15595

prompt = '''
You are given this machine learning research paper, please read it carefully and answer the follow up question.

=== BEGIN ===

2306.15595v2 [cs.CL] 28 Jun 2023

arXiv

EXTENDING CONTEXT WINDOW OF LARGE LAN-
GUAGE MODELS VIA POSITION INTERPOLATION

Shouyuan Chen Sherman Wong Liangjian Chen  Yuandong Tian
Meta Platforms Inc.
{chenshouyuan, shermanwong, cli, yuandong}@meta . com

1 INTRODUCTION

Large language models (LLMs) typically come with a pre-defined context window size. For exam-
ple, inputs to LLaMA models (Touvron et al., 2023) must be fewer than 2048 tokens. This pre-set
context window limit is frequently exceeded in applications such as conducting long conversations,
summarizing long documents, or executing long-term planning. For these applications, LLMs with
longer context windows are preferred. However, training an LLM from scratch with long context
windows requires significant investments. This naturally leads to a question: Can we extend the
context window of an existing pre-trained LLM?

One straightforward approach is to fine-tune an existing pre-trained Transformer with a longer con-
text window. However, empirically, we found that models trained this way adapt to long context
windows very slowly. After training for more than 10000 batches, the effective context window
saw a minimal increase, moving from 2048 to 2560 (Table 4). This suggests that such method is
inefficient for extending to substantially longer context windows.

While certain techniques such as ALiBi (Press et al., 2022) and LeX (Sun et al., 2022) enable length
extrapolation of Transformers, i.e. train on short context windows and inference on longer ones,
many existing pre-trained LLMs, including LLaMA (Touvron et al., 2023), use positional encodings
that have weak extrapolation properties (e.g., RoPE (Su et al., 2021)). Therefore, the applicability
of these techniques for extending the context window sizes of such LLMs remains limited.

In this work, we introduce Position Interpolation to enable context window extensions for certain
existing pre-trained LLMs, including LLaMA. The key idea is, instead of extrapolation, we directly
down-scale the position indices so that the maximum position index matches the previous context
window limit in the pre-training stage. See Figure 1 for an illustration. In other words, to accom-
modate more input tokens, we interpolate the position encodings at neighboring integer positions,
utilizing the fact that position encodings can be applied on non-integer positions, as opposed to
extrapolating outside the trained positions, which may lead to catastrophic values. We verify our
approach theoretically, by showing that the interpolated attention score has a much smaller upper

bound (~ 600x smaller in LLaMA 7B setting) than the extrapolated one, and is thus much more
stable. Therefore, interpolated position encodings are easier for the model to adapt.

Empirically, we found that Position Interpolation is highly effective and efficient, requiring only a
very short period of fine-tuning for the model to fully adapt to greatly extended context windows.
We present experimental results for extending the context window to up to 32768 from the initial
2048 across 7B to 65B LLaMA models using Position Interpolation. Our results show that

1. Position Interpolation can easily enable very long context windows (e.g. 32768), requiring
only fine-tuning for 1000 steps on the Pile (Gao et al., 2020) to achieve a good quality.
The cost of fine-tuning is negligible compared to the pre-training costs. This confirms
our hypothesis that it is relatively easy for the models to adapt to interpolated position
encodings.

2. Position Interpolation generates strong models that can effectively make use of much ex-
tended context window. We show that models extended by Position Interpolation enjoy
significant perplexity gains from greatly extended context windows for text modeling, and
we show that the perplexity reduces graceful with the enlargement of context windows.
We also applied Position Interpolation in a long text summarization task, and demonstrate
competitive performances.

3. Position Interpolation preserves model quality relatively well for tasks within its original
context window sizes. We present a variety of evaluation results for the extended LLaMA
models on the original LLaMA benchmark. Compared with original LLaMA models, the
extended LLLaM A models saw a minor degradation on several standard benchmarks within
a 2048 token limit.

Our results highlight the innate ability of Transformer models to “extrapolate to sequence lengths
longer than the ones encountered during training” as hypothesized in the seminal work of Vaswani
et al. (2017). We reaffirm this hypothesis and suggest that the previously known weakness of ex-
trapolating to longer sequences for language modeling (Press et al., 2022) may be due to direct

extrapolation of positional encodings and it can be largely mitigated by interpolating position en-
codings instead.

Concurrent work. Right before our release, we are informed with a concurrent blogpost (Super-
HOT kaiokendev (2023)) that also interpolates positional encoding in RoPE to extend the context
window from 2K to 8K. Recently, open source community picks it up in Reddit post ! and Github
Issues 2, which shows that fine-tuning with LoRA (Hu et al., 2021) also seems to work well. Our
paper shows a full fine-tuning with up to 65B model work well with Position Interpolation, and we
also give theoretical explanations why interpolation achieves much more stable results than extrap-
olation, by showing that the upper bound of interplated attention score is much lower than that of
extrapolated ones.

2 METHOD

2.1 BACKGROUND: ROTARY POSITION EMBEDDING (ROPE)

Transformer models require explicit positional information to be injected, typically in the form of
positional encodings, to represent the order of inputs. We consider Rotary Position Embedding
(ROPE) (Su et al., 2021), which is the position encoding used in the LLLaMA model (Touvron et al.,
2023). Given a position index m € [0, ¢) and an embedding vector x := [zg, 71,..., 241], Where
d is the dimension of the attention head, RoPE defines a vector-valued complex function f{x, m) as
follows

Using RoPE, the self-attention score
is only dependent on relative position m — 7 through trigonometric functions. Here q and k are the
query and key vector for a specific attention head. At each layer, RoPE is applied on both query and
key embeddings for computing attention scores.

2.2 DIRECT EXTRAPOLATION

While the attention score in RoPE only depends on the relative positions, which is what we want,
its extrapolation performance is not great . In particular, when directly extending to larger context
windows unseen in the training, the perplexity may shoot up to very high numbers (i.e., > 10%),
comparable to untrained models.

Ideally, we want to see the model trained on a context window of size L = 2048 to still work
reasonably well on longer context window, but may not have the capability to leverage information
that appears beyond L. For example, to answer a question located at 3000, the model trained on
maximal window size of I = 2048 cannot leverage evidences provided at location 0, but still
can leverage the evidences provided at location 2900. In contrast, in reality we see catastrophic
behaviors, i.e., question at location 3000 cannot be answered correctly, even if the evidences are
located at location 2900.

What is the reason behind? How could this happen if the attention score a,,,—,, decays as the relative
distance |m — n/| increases, according to Section 3.4.3 of (Su et al., 2021), and content from very
far distances should not matter that much? It turns out that the upper bound derived in Section 3.4.3
of (Su et al., 2021) may be too loose: while it indeed decays with respect to |m — nl, the bound
can still be quite large (i.e., the bound can be critically depends on the magnitude of v;) and thus
vacuous. In fact, if we treat all trigonometric functions as basis functions (i.e, ¢;(s) := #93), and
think about Eqn. 2 as basis expansion as the following:

where s is the positional span between a query and a key and h; := (ga; + igaj+1){k2j — tk2j+1)
are complex coefficients depending on q and k (here the definition of h; is exactly the same as the
definition of k; in Sec 3.4.3 in RoPE (Su et al., 2021)). Now the the issue becomes clear: as shown
in Fig. 2, a, can be small in magnitude in the range of [0, 2048], but gives huge values out of the
region. The underlying reason is that the trigonometric family {¢;} (with sufficiently large d) is
a universal approximator and can fit any arbitrary functions. Therefore, for a, there always exist
coefficients {h;} (i.e. key and query) that corresponds to small function values in [0, 2048] but

much larger in regions beyond.

2.3 PROPOSED APPROACH: POSITION INTERPOLATION (PI)

In Fig. 2, thanks to the smoothness of bases functions ¢; interpolation is much more stable and will
not lead to wild values. Therefore, instead of extrapolate the attention score in Eqn. 3 to s > L,
how about we define an attention score a{s) = a(Ls/L’) where L’ is the longer context window?
Formally, we replace RoPE f by {’ defined as follows

We call this transformation on the position encoding Position Interpolation. In this step, we reduce
position indices from [0, L') to [0, L) to match the original range of indices before computing RoPE.
Consequently, as inputs to RoPE, the maximum relative distance between any two tokens has been
reduced from I’ to L. Since we align the ranges of position indices and relative distances before
and after extension, we mitigate the effect on attention score computation due to context window
extensions, which can allow the model easier to adapt. To further demonstrate this is the case, in the
following theorem, we show that the interpolated attention score is well-behaved:

While there is no close form for B(s) := 4/21 |Ag41(s)|, numerically it is at least larger than d, and for many positional difference s, B(s) is much larger than d
(check Appendix B for the plot). Therefore, the interpolation bound is at least 2 - 294.73 ~ 600 x
smaller than the extrapolation bound, and thus the interpolated attention score is much more stable
than extrapolated one.

Notably, our method of rescaling of position indices does not introduce extra weight, or modify
the model architecture in any way. This makes it attractive in practical applications, since most
infrastructure and optimization for the original model can be reused after the extension.

Fine-tuning. We can further fine-tune the interpolated model using the next token prediction task
with interpolated position encodings on the extended context window size using a pre-training cor-
pus such as the Pile (Gao et al., 2020). In the next section, we show that our fine-tuning process
only needs tens to hundreds thousands of examples. We also find that the result of the fine-tuning
is not sensitive to the choice of examples. The reason may be that the model is only adapting to the
new context window during the fine-tuning phase, starting from a good initialization, as opposed to
acquiring new knowledge.

Other ways to reduce interpolation/extrapolation bound. From the expression of the interpola-
tion (Eqn. 5) and extrapolation bound (Eqn. 8), a common term is max; ||, which is the maximal
magnitude of query/key products. If we enforce a regularization on || during LLM training, it is
possible that the catastrophic extrapolation error can be mitigated or even resolved. In fact, if we
apply ridge regression with proper regularization to fit a curve in Fig. 2, the magnitude of extrapo-
lated a(s) when s > L can be comparable to that within [0, L]. To our knowledge, we are not aware
of existing LLM pre-training techniques that leverage this regularization and will leave it for future
work.

3 EXPERIMENTS

We show Position Interpolation can effectively extend context window up to 32 times of the original
size, and such extension can be done with only several hundreds of training steps. We show the
resulting models are strong LLMs with fully effective long context windows. We demonstrate its
performance in a number of tasks including language modeling, passkey retrieval, and long doc-
ument summarization. We also present benchmark results of the extended models on the original
LLaMA evaluation benchmarks.
3.1 SETUP

Model Variants. We extended the pre-trained 7B, 13B, 33B and 65B LLaMA models (Touvron
et al., 2023) to various context window of sizes up to 32768, using either direct fine-tuning or
Position Interpoloation method. Except for rescaling the position indices for models extended with
Position Interpolation, we did not modify LLaMA model architectures (Touvron et al., 2023) in any
ways.

Training Procedure. We fine-tune all model variants using the next token prediction objective. We
use AdamW (Loshchilov & Hutter, 2019) with 5; = 0.9 and 2 = 0.95. We use a linear learning
rate warmup of 20 steps starting from 10% of the maximum learning rate. For 7B and 13B models,
we set the learning rate to 2 x 1075 and for 33B and 65B models we set the learning rate to 1072. We
set the weight decay to zero. For extending 7B, 13B and 33B models to the 8192 context window
size, we use 32 A100 GPUs and 64 global batch size. For all other cases we use 128 A100 GPUs and
128 global batch size. We note that the main need of using more GPUs is memory limitation during
fine-tuning, and it is possible to use fewer GPUs in certain cases. We train all models using PyTorch
(Paszke et al., 2019) with Fully Sharded Data Parallel (Zhao et al., 2023) and Flash Attention (Dao
et al., 2022).

If not specified otherwise, for the Position Interpolation method, we fine-tune the models for 1000
steps. For the direct fine-tuning method, we use 10000 steps. We primarily fine-tune using the Pile
training dataset (Gao et al., 2020). In Section 3.4 we also compared fine-tuning performance on the
RedPajama dataset (Computer, 2023).

3.2 LONG SEQUENCE LANGUAGE MODELING

We evaluate the long sequence language modeling performance of our extended models and base-
lines on two datasets: book corpus (PG-19) (Rae et al., 2020) and cleaned Arxiv Math proof-pile
dataset (Azerbayev et al., 2022).

We use the test splits of PG19 (Rae et al., 2020) and proof-pile (Azerbayev et al., 2022). For PG19,
we use the whole test split consisting of 100 documents. For the proof-pile dataset, we use a random
subsample of 128 documents with at least 32768 SentencePiece (Kudo & Richardson, 2018) tokens
and truncate to the first 32768 tokens for each test document. We evaluate perplexity at various
context window size by using a sliding window approach following Press et al. (2022) with stride
S = 256.

In Table 1 and Table 2, we report the perplexity results for our models and baselines on the datasets.
From the results, we found that models extended with our method enjoy a significantly improved
perplexity from longer context window sizes. By increasing the context window size from 2048 to
16384, we observed -0.28 and -0.5 reductions of perplexity for extending LLaMA 7B models on
both datasets, -0.27 and -0.48 reductions for extending LL.aMA 13B models, and -0.14 and -0.42
reductions for extending LLaMA 33B models. For LLaMA 65B models, we observed -0.12 and
-0.3 reductions of perplexity by extending to the 8192 context window size.

In general, we observed a consistent trend of our models achieving better perplexity with longer
context windows. This indicates our models can effectively make use of the longer context windows
to better predict next tokens in language modeling tasks. Moreover, we found this trend extends to
32768 window size without diminishing on the PG19 dataset for LLaMA 7B and 13B models. This
indicates that our method may enable extension to even longer context windows.

In contrast, we observed that models extended via the direct fine-tuning method has shown regres-
sion (up to +0.48) or minor improvement (up to -0.12) on the perplexity at longer context windows.
This indicates that models extended this way have limited capability of making use of context win-
dows longer than their pre-trained settings.

We saw a minor degradation of the perplexity on the original context window of 2048 for our ex-
tended models in some cases. For example, on the Proof-pile dataset, we saw a degradation ranging
from 0.01 to 0.05 across all models with extended with Position Interpolation. A small degradation
of performance within original evaluation context window is expected since Position Interpolation
forces position encodings in original context window to reside in a much narrower region, which
may negatively affect the language model’s performance. We present more benchmark results on
the original context window size in Section 3.4.

In Table 3 we report the relationship between perplexity and the number of fine-tuning steps for
LLaMA 7B model extending to 8192 and 16384 context window sizes using Position Interpolation
evaluated on the PG19 dataset. We can see without fine-tuning (at step 0) the model can exhibit
certain language modeling capability, as indicated by < 20 perplexity for extending to 8192 context
window (in contrast, the direct extrapolation method leads to > 10% perplexity). With fine-tuning,
we observed that the perplexity improves quickly. At 200 steps the models surpassed the original
model’s perplexity on 2048 context window size, indicating the models gaining ability of effectively
using sequences longer than the pre-training settings for language modeling. At 1000 steps, we can
see the models have improved steadily and achieve a significantly better perplexity.

3.3 MEASURING EFFECTIVE CONTEXT WINDOW SIZE THROUGH PASSKEY RETRIEVAL

We study the effective context window size, i.e. the maximum distance of a token can effectively
attend to during inference, of our models after extension. To measure this, we follow a synthetic
evaluation task of passkey retrieval proposed by Mohtashami & Jaggi (2023). In this task, the models
are asked to recover a random passkey hidden in a long document. See Figure 3 for the format of
the document.

Given a language model, we estimate the upper and lower bounds of effective context windows as
follows. Suppose the random passkey is k tokens away from the end of the input. When a model
persistently fails to retrieve the correct passkey value across several independent attempts, it suggests
that the effective context window size of the model is less than k. Conversely, if a model consistently
succeeds in retrieving the correct passkey value, we deduce that the effective context window size
of the model is at least k.

We evaluate the 7B and 33B LLaMA model variants that are extended via Position Interpolation or
direct fine-tuning. For each model, we use 32 different &£ uniformly spaced in the targeted context
window L’ and run the above tests for 10 times for each k, where each time a random passkey of 5
random digits is used. In Table 4, we report kyax as a function of the number of fine-tuning steps,

We can see that models extended via Position Interpolation all successfully attain their desired ex-
tension objectives in terms of effective context window sizes, indicating by the effective context
window size reaching maximum kp, = L/, after merely fine-tuning for 200 steps, consistently
across both 7B and 33B model sizes and up to 32768 context windows. In contrast, LLLaMA models
that are extended via direct fine-tuning only saw a minimal increase of the effective context win-
dow size kay from 2048 to 2560, even after fine-tuning for more than 10000 steps, with no clear
indication of an acceleration in the increase of window size.

3.4 BENCHMARKS ON ORIGINAL CONTEXT WINDOW SIZE

We evaluate the models extended by Position Interpolation on several standard benchmark tasks
within the original context window size of 2048. The evaluation results are listed in Table 5. From
the results, we saw that models extended to 8192 produce comparable results on the original bench-
mark which is designed for a much smaller context window, with a degradation of up to 2% on
the benchmark tasks, for both 7B and 33B model sizes. Models extended to longer context win-
dows regressed more on the benchmarks, but still in reasonable ranges for most tasks. We also note
that the choice of fine-tuning datasets does not seem to lead significant difference in the benchmark
performances, which may be due to the limited number of fine-tuning steps used in our method.
The regression on benchmark tasks is consistent with our observation on perplexity regression in
Section 3.2.

3.5 LONG DOCUMENT SUMMARIZATION

In this task, we evaluate our models’ performance on the long document summarization task. In
particular, we consider the GovReport (Huang et al., 2021) dataset, which contains 17457 documents
for training and 972 documents for evaluation. Each document comes with a human generated
summary. We truncate all input documents to their first 15000 tokens.

We fine-tune the LL.aMA models extended with Position Interpolation with a context window of
16384. Note the rescaling of position indices are still required during this fine-tuning step. We first
Model Size Context Window Fine-tune on  BoolQ PIQA Race-M Race-H WinoGrande

format the raw document using the prompt template in Figure 4, and then concatenate the prompt
with the ground-truth summary (truncate to 1000 tokens) associated with each document. We fine-
tune the model using the next token prediction task with the above setup for 10 epochs. The losses
from the input prompt proportion of training examples are excluded during our fine-tuning.

We use a generation temperature of 0.5 and top, = 0.95 as our inference parameter to generate a
summarization of each document in the test set. The final output is truncated at 1000 tokens. We
used the ROUGE-1/ROUGE-2/ROUGE-L scores (Lin, 2004) as the evaluation metrics to evaluate
the models’ outputs vs the ground-truth summaries.

In Table 6 we report our evaluation results. We have also included results from two baselines in
existing SCROLLS Leaderboard (Shaham et al., 2022; Ainslie et al., 2023). In general, we have
obtained competitive R1 score among other models with minimal tuning of hyper-parameters. This
result suggests our models with 16384 context window can effectively handle the long document
summarization task.

=== END OF FILE ===

'''

If we place it in the following example

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("huggyllama/llama-7b")
model = AutoModelForCausalLM.from_pretrained(
    "huggyllama/llama-7b",
    load_in_8bit=True,
    device_map="auto",
)

prompt = ...
question = "Question: What is the paper about?"

inputs = tokenizer(prompt + question, return_tensors="pt").to("cuda")

print(inputs.input_ids.shape)
gen_out = model.generate(**inputs, max_new_tokens=100)
print(tokenizer.batch_decode(gen_out)[0])

we get:

Question: What is the paper about? a a a a a a a a a a a a b: a a a a a a: a in a a a a a a [(b a b. a [b [b [b. [b [b [( [( [( [( [( [( [b [(b [b [b
[b [(( [((: [(: [: [: [((((((0:(((((al:

However, if we add rope_scaling={"type": "dynamic", "factor": 2.0} in from_pretrained, we now get:

Question: What is the paper about?
Answer: The paper is about extending the context window of Transformer models.
Answer: The paper is about extending the context window of Transformer models.
Answer: The paper is about extending the context window of Transformer models.
Answer: The paper is about extending the context window of Transformer models.
Answer: The paper is about extending the context window of Transformer models.
Answer: The paper is about extending the context window of Transformer models.
Answer: The

Better generation parameterization can definitely be selected, but you get the idea -- with these changes, models with RoPE can handle much larger contexts right out of the box 🔥

[gante]

(Of course, tests are missing. Proper validation of whether the feature is working as expected is also missing. I'll add them if we decide to move forward with this feature!)

[gante]

(changes in open_llama are copy/paste)

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[TheBloke]

Having this in transformers would be excellent!

I've uploaded a bunch of fp16 and GPTQ repos to HF using @jquesnelle 's trust_remote_code Llama modelling patch that implements RoPE using @kaiokendev's method, and I know there are quite a number of people using those already, and I've had a few requests to put out more. And even more are using RoPE outside of transformers via the ExLlama GPTQ implementation.

So there's a great deal of appetite for this feature amongst users, understandably.

[versae]

Could this also be applied to GPT-J models?

[gante]

@versae yes, it can :) The code needs to be modified there as well, but the concept can be applied to any model with rotary position embeddings

[jquesnelle]

For dynamic, we want to do the scaling only when seq_len > max_position_embeddings (i.e. when we're going past the model's pre-trained length). My original code did this by just having the scaling in the forward() code that re-calculated the frequency cache when seq_len > self.max_seq_len_cached but not in the __init__. Since this code has now been deduplicated (makes sense!), I think this needs to be

if self.scaling_name == "dynamic" and seq_len > self.max_position_embeddings:

[gante]

@jquesnelle that's a great catch, and completely missed in my attempts to unify and refactor the cos/sin cache!

[bloc97]

Thank you for your work! Just letting you know that I've improved the NTK-aware method in this PR. jquesnelle/yarn#1 It decreases non-finetuned PPL even further (preliminary testing shows 4.9 -> 4.6 PPL at 8192 context size) and theoretically will significantly improve a finetune's convergence/stability compared to previous NTK-aware method.

Also because the alpha hyperparameter was difficult to use when predicting effective context size (alpha=4 was something close to ~6400 context size instead of 8192), that problem was fixed and it is now changed to a "scale" factor, which can be used the same way to the "scale" in linear RoPE scaling. (eg. for LLaMA scale=2 is 4096 and scale=4 is 8192)

I hope this improved method might be also considered one day as it is one more step towards extending context size for all LLMs! 🚀

[sgugger]

As long as this can be used with all existing checkpoints, I have no problem adding this as a new feature. Not too sure about the API you picked as a dictionary though. Since there are two arguments, I would just have put two arguments as a dict can be complicated. Or if we expect more arguments to arrive, having some dataclass to hold them and accept both dict/dataclass at init (though this can also be done later since we are warning this is an experimental feature subject to changes).

And as was mentioned in the comments, this should also be added to other models with Rotary Embeddings :-)

[sgugger]

I would put "Experimental feature" as a warning at the end of the docstring.

This is an experimental feature, subject to breaking API changes in future versions.

[gante]

Hey @bloc97 @jquesnelle 👋

Looking at your recent PR (this one) -- am I right in saying that

1. There is no way to parameterize the new class such that it is equivalent to the original NTK-aware scaling?
2. @bloc97's PR and @jquesnelle's dynamic implementation are slightly different, in the sense that @bloc97's targets a specific length (but can extrapolate) and @jquesnelle's dynamically adjusts to the maximum observed length?
3. Because @jquesnelle's implementation base may suddenly change due to a longer sequence, it is less friendly to fine-tune?

I'm trying to determine how to integrate and document the goodies, while keeping the diff size manageable 🤗

[gante]

The technique also seems to work out-of-the-box with GPTNeoX models 🔥 With the latest commit, running the script below

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

tokenizer = AutoTokenizer.from_pretrained("EleutherAI/pythia-1.4b-deduped")
model = AutoModelForCausalLM.from_pretrained(
    "EleutherAI/pythia-1.4b-deduped",
    torch_dtype=torch.bfloat16,
    device_map="auto",
    rope_scaling={"type": "dynamic", "factor": 2.0},
)

prompt = ...  # see PR header for the prompt, >5k tokens
question = "Question: What is the paper about?"

inputs = tokenizer(prompt + question, return_tensors="pt").to("cuda")

print(inputs.input_ids.shape)
gen_out = model.generate(**inputs, max_new_tokens=50)
print(tokenizer.batch_decode(gen_out)[0])

gets us

Question: What is the paper about?

3.6 CONCLUSION

We have shown that Position Interpolation can extend the context window of pre-trained models to substantially longer context windows. We have
demonstrated that the models can be effectively extended to longer context windows, and

Without the rope_scaling argument, we get

Question: What is the paper about? The. The.
. The.
.
.
.
.
4.
. The. The. The. The. The. The. The.al.s. The. The.... a. The.

This is particularly amazing, since we're talking about a 1.4B model 👀

(cc @bloc97 @jquesnelle, you may be interested in this finding)

[gante]

@amyeroberts @sgugger I'd like to request the review of you two, since this is an unusual post hoc modeling change PR.

A few key points:

1. The technique also works well with gptneox, see this comment for a cool example on a 1.4B model
2. Adding the functionality to gptneox implied a minor modeling change -- the causal mask was limited to the original maximum sequence size, but there is no reason for that limitation. It's just a triangular matrix with ones.
3. Decided NOT to implement on gptneox-japanese and esm, the two other models with rotary embeddings. I'm not sure if their usage justifies the implementation cost (it takes some time to validate everything is working correctly, as there are variations in the expected usage), so I'd suggest letting demand speak for itself :)
4. RoPE scaling is parameterized by a dict, and not a dataclass. A dataclass would be better, as @sgugger suggested, but it complicates (de)serialization, needing extra code. I'd like to first work on the config file base class I've mentioned on slack, if you're okay with it -- it would make the new dataclass a ~50 line change, as opposed to a >200 one!
5. There are new scaling strategies in the works, as mentioned in the comments above, so we can quickly add them in follow up PRs if their results are superior. As it stands, we can already hack llama and gptneox beyond their original maximum length without fine-tuning 🔥

[sgugger]

LGTM! As said before, still not a fan of the API but since you are planning to address this in a follow-up PR, that works for me.

[jquesnelle]

(For brevity, I'll refer to the new NTK-By-Parts method as NTKv2)

NTKv2 is an improved version of NTK. We found that NTK did not perform well when fine-tuned; the reason for this was that the resulting embedding matrix still contained some extrapolated out-of-population values that model had not seen during training. Dynamic NTK hid this by continually scaling base so that you never actually got to this part of the embedding values.

NTKv2 is parameterized by scale, which has the same meaning as linear interpolation, e.g. you set it to 4 to target 8K context length. We've found that this method, when fine-tuned, beats fine-tuned linear interpolation, which is to say it gives even better results than the recent Meta paper.

In the repository there is also now a Dynamic NTKv2, which is the same idea as the previous dynamic method, i.e. scale the hyperparamter relative to the ratio between the current context length and the model's original trained context length, while using the original embedding values when under the native pre-trained length. This also beats Dynamic NTK in the no-fine-tuning scenario.

In the above graph, LLongMA are the fine-tuned OpenLLaMA models we've released, trained on 1B extra tokens (v2 still in the process of training)

1. There is no way to parameterize the new class such that it is equivalent to the original NTK-aware scaling?

Unfortunately no. I understand these different methods can get unwieldly quickly, but NTKv2 appears to be strictly better than original NTK -- I would potentially just advocate replacing the original NTK with this, but that could also be done in a follow-up PR too; the results that this gives you is already Very Good (TM).

FWIW the LLongMA models use the exact modeling code here to maintain compatibility without needing trust_remote_code if/when this PR gets merged 🙂

[bloc97]

Hey @bloc97 @jquesnelle 👋

Looking at your recent PR (this one) -- am I right in saying that

1. There is no way to parameterize the new class such that it is equivalent to the original NTK-aware scaling?
2. @bloc97's PR and @jquesnelle's dynamic implementation are slightly different, in the sense that @bloc97's targets a specific length (but can extrapolate) and @jquesnelle's dynamically adjusts to the maximum observed length?
3. Because @jquesnelle's implementation base may suddenly change due to a longer sequence, it is less friendly to fine-tune?

I'm trying to determine how to integrate and document the goodies, while keeping the diff size manageable 🤗

1. Unfortunately "NTK v1" was just not good for fine-tuning unless alpha is set correctly, so I think going forward people should strictly use "v2" for fine-tuning, and consider v1 to be only for inference. However it is possible for me to parameterize the "v2" class so that you can make it equivalent to original NTK scaling, but it will take additional effort that is probably best used elsewhere. There are only few "NTK v1" finetunes are out there.
2. For points 2 and 3, finetuning with Dynamic method will need additional consideration in the code on the training side, because training happens on all the tokens at once, dynamic implemented as is (for inference) will probably not be applied correctly. We are still working on the theoretical side of potentially training a dynamic model.

[amyeroberts]

Really nice! 🤩 Thanks for adding ❤️

+1 on @sgugger's comment on using a dataclass, also happy for it to be done in a follow-up.

[gante]

@bloc97 @jquesnelle thank you for your input -- and excited to hear about the performance of NTK-By-Parts!

Based on your comments, I will:
1 - Delete the ntk approach, as NTK-By-Parts is superior;
2 - Merge what I have now -- we are going to have a release early next week, so this would already be included in v4.31;
3 - Open a follow-up PR with NTK-By-Parts 🤗 Or, if you're interested in contributing with the technique, we'd highly appreciate it! Just let me know over the next days.

⚠️ Note -- the latest commits have changed the structure of the modeling code from overloading the existing RoPE class to inheriting from the original implementation, so we don't risk ending up with a Frankenstein class as we add more strategies. The parameterization stayed nearly the same, so you probably only need to make minor adjustments to the model config files to load without trust_remote_code! (changed from {"name": scaling type, "factor": scaling factor} to {"type": scaling type, "factor": scaling factor}, as name is often attributed to an instance name in transformers)

[fahadh4ilyas]

Hi, I'm very glad to see that transformers supports RoPE scaling! Experiments show low ppl on long input sequences.

But in the current implementation, would there be a mismatch in generation? Here are my thoughts.

Since the seq_len increases during the generation, the base is scaled in every generation step with different scaling factor. Since the history key_states are store in the kv_cache , they are not scaled with the new base. The scaling only affects the state of the current token.

For example, if the input sequence is of length 2048, after generating the first token, the new input length is 2049, and we scale the base with seq_len=2049. After generating the second token, the new input length is 2050, and we scale the base with seq_len=2050. But during the generation, the kv_cache is used and thus the key_states before position 2049 are not scaled according to the new length.

Should all the key_states be scaled with the same base? Would it be a problem?

I have question similar to this. The graph showing dynamic scaling in this reddit post showing that the perplexity of the model with dynamic scaling are same with model without scaling until 2048 tokens length (Of course this must be because the base value did not change before 2048 tokens).

This got me thinking, If I first generate with long context (say 4096 tokens), the base value would change accordingly (which is around 35000). Then, if I next generate with short context like 1024 context, the sin_cache and cos_cache will not be reverted back when the base value still 10000 hence the perplexity is raised. Should there be changed to forward call especially for dynamic scaled embeddings?

[airaria]

This got me thinking, If I first generate with long context (say 4096 tokens), the base value would change accordingly (which is around 35000). Then, if I next generate with short context like 1024 context, the sin_cache and cos_cache will not be reverted back when the base value still 10000 hence the perplexity is raised. Should there be changed to forward call especially for dynamic scaled embeddings?

I have the same concern. In the dynamic scaling, the sin and os may should not be cached

[guozhiyao]

Hi
I try to test ntk effect on my trained neox model. Using dynamic ntk(https://www.reddit.com/r/LocalLLaMA/comments/14mrgpr/dynamically_scaled_rope_further_increases/). However, it is found that the ppl will oscillate. What is the reason for this?

here is my test code. I modified from https://huggingface.co/docs/transformers/perplexity .

import json
from transformers import AutoModel, AutoTokenizer, AutoConfig
import torch
from tqdm import tqdm
import traceback

device = "cpu"
if torch.cuda.is_available():
    device = "cuda"

config = AutoConfig.from_pretrained(model_dir, trust_remote_code=True)
config.rope_scaling = {
    "type": "dynamic",
    "factor": 2,
}

model = AutoModel.from_pretrained(model_dir, config=config, trust_remote_code=True, torch_dtype=torch.float16)
model.eval()
model = model.to(device)
tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)

with torch.inference_mode():
    kv = {}
    try:
        for value in tqdm(range(32, 12000, 32)):
            max_length = stride = value

            with open("gov_report_test.json") as f:
                data = json.load(f)

            ppls = []
            for idx, line in enumerate(data):
                if idx >= 1:
                    break
                encodings = tokenizer(line, return_tensors="pt")
                seq_len = encodings.input_ids.size(1)

                nlls = []
                prev_end_loc = 0
                for begin_loc in range(0, seq_len, stride):
                    end_loc = min(begin_loc + max_length, seq_len)
                    trg_len = end_loc - prev_end_loc  # may be different from stride on last loop
                    input_ids = encodings.input_ids[:, begin_loc:end_loc].to(device)
                    target_ids = input_ids.clone()
                    target_ids[:, :-trg_len] = -100

                    outputs = model(input_ids, labels=target_ids)

                    # loss is calculated using CrossEntropyLoss which averages over valid labels
                    # N.B. the model only calculates loss over trg_len - 1 labels, because it internally shifts the labels
                    # to the left by 1.
                    neg_log_likelihood = outputs.loss

                    nlls.append(neg_log_likelihood)

                    prev_end_loc = end_loc
                    if end_loc == seq_len:
                        break

                ppl = torch.exp(torch.stack(nlls).mean())
                ppls.append(ppl)

            total_ppl = torch.stack(ppls).mean()
            kv[value] = total_ppl.item()
            print(value, total_ppl.item())
    except Exception as e:
        print(e)
        print(value, seq_len)
        print(traceback.format_exc())

[gante]

@guozhiyao Nothing immediately comes to mind, it could be even a model "feature" (looking at the plot for the original model, which also has the periodicity).

Would you be able to a) run the same script for LLaMA and b) repeat your experiment using the script @jquesnelle used (this one)? a) should rule out model-specific issues and b) should rule out code-specific issues.

[guozhiyao]

@guozhiyao Nothing immediately comes to mind, it could be even a model "feature" (looking at the plot for the original model, which also has the periodicity).

Would you be able to a) run the same script for LLaMA and b) repeat your experiment using the script @jquesnelle used (this one)? a) should rule out model-specific issues and b) should rule out code-specific issues.

@gante Thanks a lot. It is solved by using the code.

[guozhiyao]

This got me thinking, If I first generate with long context (say 4096 tokens), the base value would change accordingly (which is around 35000). Then, if I next generate with short context like 1024 context, the sin_cache and cos_cache will not be reverted back when the base value still 10000 hence the perplexity is raised. Should there be changed to forward call especially for dynamic scaled embeddings?

I have the same concern. In the dynamic scaling, the sin and os may should not be cached

@airaria I had the same problem, not only cos and sin, inv_freq also don't cache. The _set_cos_sin_cache of GPTNeoXDynamicNTKScalingRotaryEmbedding can be changed to the following form, but the efficiency is not optimized.

    def _set_cos_sin_cache(self, seq_len, device):
        self.max_seq_len_cached = 0

        base = self.base
        if seq_len > self.max_position_embeddings:
            base = self.base * (
                (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)
            ) ** (self.dim / (self.dim - 2))

        inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
        self.register_buffer("inv_freq", inv_freq)

        t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)

        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
        # Different from paper, but it uses a different permutation in order to obtain the same calculation
        emb = torch.cat((freqs, freqs), dim=-1)
        self.cos_cached = emb.cos()[None, None, :, :]
        self.sin_cached = emb.sin()[None, None, :, :]

[guozhiyao]

This got me thinking, If I first generate with long context (say 4096 tokens), the base value would change accordingly (which is around 35000). Then, if I next generate with short context like 1024 context, the sin_cache and cos_cache will not be reverted back when the base value still 10000 hence the perplexity is raised. Should there be changed to forward call especially for dynamic scaled embeddings?

I have the same concern. In the dynamic scaling, the sin and os may should not be cached

@airaria I had the same problem, not only cos and sin, inv_freq also don't cache. The _set_cos_sin_cache of GPTNeoXDynamicNTKScalingRotaryEmbedding can be changed to the following form, but the efficiency is not optimized.

    def _set_cos_sin_cache(self, seq_len, device):
        self.max_seq_len_cached = 0

        base = self.base
        if seq_len > self.max_position_embeddings:
            base = self.base * (
                (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)
            ) ** (self.dim / (self.dim - 2))

        inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
        self.register_buffer("inv_freq", inv_freq)

        t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)

        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
        # Different from paper, but it uses a different permutation in order to obtain the same calculation
        emb = torch.cat((freqs, freqs), dim=-1)
        self.cos_cached = emb.cos()[None, None, :, :]
        self.sin_cached = emb.sin()[None, None, :, :]

There is a precision difference between the inv_freq here and the inv_freq defined in __init__, and the reason is not found. In order to ensure the same performance as the original when seq_len <= self.max_position_embeddings, it can only be modified to this form.

    def _set_cos_sin_cache(self, seq_len, device):
        self.max_seq_len_cached = 0

        if seq_len > self.max_position_embeddings:
            base = self.base * (
                    (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)
            ) ** (self.dim / (self.dim - 2))
            inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
        else:
            inv_freq = self.inv_freq

        t = torch.arange(max(seq_len, self.max_position_embeddings), device=device, dtype=inv_freq.dtype)

        freqs = torch.einsum("i,j->ij", t, inv_freq)
        # Different from paper, but it uses a different permutation in order to obtain the same calculation
        emb = torch.cat((freqs, freqs), dim=-1)
        self.cos_cached = emb.cos()[None, None, :, :]
        self.sin_cached = emb.sin()[None, None, :, :]