如何擴充套件大模型的上下文長度｜得物技術

大模型的上下文長度是指我們在使用大模型的時候，給大模型的輸入加上輸出的字元（Token）總數，這個數字會被限制，如果超過這個長度的字元會被大模型丟棄。目前開源的大模型上下文長度一般不長，比如 Llama 2 只有 4K，Code-Llama 系列因為需要輸入程式碼，擴充套件到了 16K。閉源系列模型的提供了更長的上下文長度，比如 OpenAI 在其最新模型 GPT-4 Turbo 中提供了 128K 的上下文長度，Anthropic 的 Claude 2.1 模型提供了 200K 上下文長度。

一些場景需要較長上下文，比如，文件翻譯需要將整篇文件輸入給大模型進行翻譯，長文件內容抽取需要大模型讀取整篇長文件進行內容抽取，會議內容總結則需要給大模型輸入會議聊天記錄進行總結等。

想要得到一個長上下文的大模型，一般有兩種途徑。一種是大模型在初始階段被設定為長上下文，然後經過預訓練，指令微調，對齊訓練等方式得到一個長上下文大模型。另外一種方式是選擇已經訓練好的大模型，透過技術改造擴充套件其上下文長度，然後再進行微調訓練得到長上下文模型。

本文將基於比較火的 Llama 2 大模型的結構[1]介紹上下文長度的方法與挑戰，然後探討一些業界流行的上下文長度擴充套件的技術，最後給大家推薦下 KubeAI 大模型訓練推理平臺可以上手實驗。

通常所說的大模型是指大語言模型（Large Language Model，LLM），其模型結構一般基於 Transformer 進行改進而來。Transformer 源自於 2017 年 Google 發表的著名論文"Attention Is All You Need[2]"。論文中的 Transformer 結構如下，下圖源自論文[2]。

它的結構包括兩部分：Encoder（編碼器）和Decoder（解碼器）。Encoder 與 Decoder 大致都包含以下層，每一層都有特定的功能，下面為 Encoder（編碼器）各層的簡單介紹：

• 輸入嵌入層（Input Embedding Layer）：將輸入文字的詞或標記轉換為向量表示，以便模型能夠理解它們。

• 多頭自注意力層（Multi-Head Self-Attention Layer）：幫助模型捕捉輸入序列中詞與詞之間的關係，使模型能夠了解上下文資訊。

• 前饋神經網路層（Feed-Forward Neural Network Layer）：對多頭自注意力的輸出進行進一步的特徵提取和變換，以增加模型的表示能力。

• 歸一化層（Layer Normalization Layer）：規範化每一層的輸出，有助於訓練過程的穩定性。

總的來說，Transformer 是一種強大的模型，它可以捕捉文字和序列資料中的長距離依賴關係，使其在翻譯、對話、摘要生成等自然語言處理任務中表現出色。這個模型已經在各種應用中取得了顯著的成功。感興趣的同學可以自行去網上搜尋下 Transformer 的結構，深入瞭解。

目前大多數生成式語言模型如 Llama 系列，僅僅採用了 Transformer 的 Decoder 模組結構。在 Huggingface 中,這種結構通常被稱為 CausalLM，即因果語言模型（Causal Language Model）。下面我們來具體看一下 Llama 2 模型系列的結構，Llama 2 相關的論文[1]。

（下圖是基於 Transfomers 程式碼中的 LlamaModel 繪製而成,具體程式碼參考 Transfomer 中的 modeling_llama.py[3]）

我們來解讀下上面 Llama 各層的結構與作用，首先從輸入文字開始。會經過下面各層：

• Input Embedding：將 Input 文字轉化為向量表，透過 nn.Embedding 實現。

• Llama Decoder Layer：Decoder 採用多層 Llama Decoder Layer。每一層包括自注意力（Llama Attention）和前饋網路（Llama MLP）。自注意力用於捕捉文字中的長程依賴關係。前饋網路進行非線性對映。

• Llama RMSNorm：一種規範化方式，用於正則化每層的輸，起到預處理的作用。

• lm_head：一個線性層，將 Decoder 最後一層的輸出對映到詞典大小的維，以進行後續的語言模型 Logits 計算。

• Llama Attention：多頭自注意力機制，用於建模文字中的依賴關係。將輸入表示切分為多個頭，然後在每個頭內做點積注意力運算。

• Llama MLP：採用 Gated Linear Units 的多層前饋網路。進行非線性變換來捕捉複雜模式。

總體上，Llama 透過堆疊多層自注意力和前饋網路來表示文字語義，然後預測後續詞元。lm_head 負責將語義表示對映為具體的詞典 Logits。整個模型端到端透過語言模型目標進行訓練。

上面各層中比較核心的是 Llama Attention 層，該層的結構如下：

論文“Attention Is All You Need”[2]中描述的 Attention 的計算公式如下：

Llama 的 Attention 計算過程如下：

• 輸入會經過線性變換，得到 Query(Q)、Key（K）和 Value（V）矩陣。

• 對 Q 和 K 應用 RoPE 位置編碼。RoPE 包含旋轉的 Sin 和 Cos 編碼，會根據每個 Token 的位置對其表示進行旋轉。

• 用旋轉後的 Q 和 K 計算點積，得到注意力權重 Attention Score，經過 Softmax 計算後得到 Normalized Attention Weight。

• 再把 Attention Weight 與 V 相乘，並進行加權求和，得到 Attention 的輸出。

• 輸出再經過一個線性變換，繼續輸出給下一層作為輸入。

這樣透過 RoPE 位置編碼、加權平均，Llama 的 Attention 可以高效穩定地提取文字序列的上下文語義資訊。

透過上面對 Llama 結構的解析，我們看到在 Llama Attention 中有一個叫做 RoPE（旋轉位置編碼）的層，主要用於對輸入進行位置編碼，讓模型學到輸入文字中每個 Token 的位置關係，從而更好地理解輸入。RoPE 層能處理序列的長度決定了 Llama 的上下文長度，要擴充套件 Llama 的上下文長度，需要對 RoPE 層進行改造和擴充套件。下面我們先簡單介紹下 Llama RoPE 層的工作原理。

RoPE 旋轉位置編碼，最早來自 RoFormer：Enhanced Transformer with Rotary Position Embedding[4]這篇論文。下圖源自論文[4]。

RoPE 層是一種相對位置編碼方法，它給輸入的每個 Token 編碼一個向量，向量中的每個值表示該 Token 與其他 Token 的相對距離。論文以二維向量為例，解釋了這種位置編碼為什麼叫做旋轉位置編碼。如上圖所示，在二維平面，相當於把向量旋轉了一個 Q 的角度。

論文中證明，在進行旋轉位置編碼之後，可以從新的編碼向量中獲取原向量的相對位置資訊，即論文中下面的公式中的 m-（位置 m 減去位置 n）。下圖源自論文[4]。

而我們只需要理解旋轉位置編碼的最終計算公式如下，對於一個輸入向量 X，直接與一個 COS 矩陣和 SIN 矩陣進行內積後求和：下圖源自論文[4]。

其中 WCOS 和 WSIN 分別是一個預先固定的 COS 矩陣和 SIN 矩陣。

下面我們展示下 Huggingface 的 Transformer 對應的計算 COS 和 SIN 矩陣的計算程式碼：

上面介紹了 Llama 結構中的旋轉位置編碼層 RoPE。要擴充套件大模型的上下文長度，就需要擴充套件 RoPE 層，也就是擴充套件其 COS 和 SIN 矩陣，讓RoPE支援更長序列的輸入。

RoPE 中的 COS 和 SIN 矩陣維度（seq_length, embed_dim），其中 seq_length 就是模型支援的最大序列長度，embed_dim 是詞嵌入維度。矩陣中的每個值表示一個位置上的正弦或餘弦編碼。為了支援更長的上下文，需要重新計算更大尺寸的 COS 和 SIN 矩陣。

對於未訓練過的大模型，只需要直接更改其配置檔案中的 max_position_embeddings 即可實現 RoPE 層的擴充套件，然後再進行訓練。但是對於已經訓練過的模型，如果直接修改其配置，會導致模型的效果急劇下降，後面第四部分我們會介紹一些基於已有模型進行改造擴充套件的上下文的方法。max_position_embeddings 的配置如下：

超長上下文的大模型部署推理的時候，往往會面臨如下效能挑戰。

• 推理時間變長

從上面的 Attention 的計算公式可以看出，Attention 進行了點積的運算，其時間複雜度為 L（序列長度）的平方。也就是說大模型在推理的時候，輸入的序列長度越長推理時間越多。所以超長上下文的大模型需要更多的推理時間，這會帶來使用者體驗上的損失。

• 推理視訊記憶體空間變大

大模型在持續推理的過程中，需要快取一個叫做 KV Cache 的資料快，KV Cache 的大小也與序列長度成正比。以 Llama 2 13B 大模型為例，一個 4K 長的序列大約需要 3G 的視訊記憶體去快取 KV Cache，16K 的序列則需要 12G，128K 的序列則需要 100G 視訊記憶體。

超長上下文的大模型需要更多的 KV Cache 儲存空間，但是 GPU 視訊記憶體非常珍貴，比如 A100 也只有 40G 或 80G 視訊記憶體兩個版本，這對本來就比較緊張的 GPU 視訊記憶體來說是一個很大的挑戰。

綜上所述，擴充套件大模型的上下文長度，一般思路如下：

• 首先透過對位置編碼層進行改造，使其支援更長的上下文。

• 為了取得更好的推理效能，還需要對 Attention 計算進行最佳化。

• 進行微調訓練，讓大模型適應新的模型結構。

上面我們講到，從模型物理結構上擴充套件上下文長度，需要直接修改 RoPE 層，即直接擴充套件其 SIN 和 COS 矩陣。但是大模型都是基於大量短序列資料訓練得到的。如果直接強行擴充套件，會導致模型困惑度提高。所謂困惑度是模型對下一個詞的預測困難程度的量化指標，直觀意義是大模型的輸出是否能夠更容易被人類所理解。

因此，我們需要更好的方法來擴充套件預訓練模型的上下文長度，既要兼顧模型效能，又要控制困惑度。下面我們概括幾種業界常用的上下文長度擴充套件方法。

線性插值法的思想最早來自於這篇文章，Extending Context Window of Large Language Models via Positional Interpolation[5]，該方法已經被 Huggingface 的 Transformer 中 Llama 模型程式碼整合。

下圖源自論文：

思路：透過線性縮小輸入位置索引以匹配原始上下文視窗大小，而不是超出訓練上下文長度進行外推，這樣可以減小注意力機制中相對位置的影響，幫助模型更容易適應擴充套件後的上下文視窗。

效果：在從 LLaMA 7B 到 65B 模型上，透過位置插值擴充套件上下文視窗到 32768（4k擴充套件到32K），僅需微調 1000 步，就能在包括語言建模、長文件摘要撰寫等任務上取得良好效果。

優點：位置插值不僅能有效擴充套件上下文視窗，提高模型在長上下文任務上的效能，還能在原有上下文視窗大小的任務上保持模型質量，且不需要額外的權重或修改模型架構。

缺點：需要重新訓練，有時候擴充後會導致模型困惑度上升。

動態插值法是在位置插值法的基礎上演變而來的，最早提出文章 NTK-Aware Scaled RoPE allows LLaMA models to have extended (8k+) context size without any fine-tuning[6]。現在也被 Huggingface 的 Transformer 中 Llama 模型程式碼整合。

下面是 Chinese Llama 基於位置插值法與動態插值法進行的比較，資料來自 Extend Context Size Without Fine-Tuning[7]。

資料顯示，相比與位置插值法，NTK 動態插值法不會顯著增加大模型的困惑度。

思路：利用神經正切核 (NTK) 理論，設計非線性位置編碼插值方案，改變基數而不是縮放比例，使不同位置可區分，避免線性插值的問題。

效果：與線性插值相比，在大幅擴充套件上下文（如8000+）時，無需微調就可以使困惑度下降極小。

優點：微調成本極低，上下文視窗可以擴充套件很大，困惑度變化小。

總體來說，動態插值法透過考慮模型特性，設計更最佳化的插值方案，能夠在不增加訓練成本的條件下，獲得接近無損的上下文視窗擴充套件效果。這為進一步擴充套件和最佳化大語言模型提供了新的思路。

Yarn 擴充套件上下文的方法來自於文章 YaRN: Efficient Context Window Extension of Large Language Models[8]。程式碼參考Yarn[9]。並給出基於 Llama2 的 128K 上下文擴充套件。

下圖源自論文[8]：

資料顯示，在 128K 的 Proof-Pile 資料集上評測，Yarn-Llama-2-7b-128K/64K 模型的困惑度仍然保持良好下降。

相比於僅進行簡單線性插值或動態插值的方法，Yarn 方法更全面地考慮了不同頻率 RoPE 維度的作用，避免了資訊損失和外推問題。這使得 Yarn 方法在不 Fine-Tuning 的情況下，以及 Fine-Tuning 的資料量很少的情況下，都能更好地擴充套件上下文視窗。

RoPE 中的每個維度對應著不同的正弦波頻率。高頻的正弦波編碼了位置資訊的細微變化，低頻的正弦波編碼了位置資訊的整體趨勢。

如果我們簡單地進行線性插值，會把所有頻率的正弦波都等比例地拉伸。這會導致兩個問題:

• 高頻正弦波被過度拉伸，導致代表細微位置變化的資訊丟失。這個會影響模型區分很接近的詞的能力。

• 低頻正弦波被拉伸，不同位置之間的相對距離變小。這會導致模型判斷近距離詞的先後順序變得困難。

為了解決這個問題，Yarn 方法對不同頻率的正弦波進行不同程度的插值：

• 對高頻正弦波幾乎不進行插值，保留細微位置資訊。

• 對低頻正弦波進行接近線性的插值，保留位置大體資訊。

• 中頻正弦波進行漸變的插值。

這樣既保留了高頻表示細微位置變化的資訊，也保留了低頻表示位置整體關係的資訊，避免了簡單線性插值的問題。

上面我們提到，長上下文對大模型的正向與反向傳播的效能來說是個挑戰。其主要原因是 Attention（注意力）的計算複雜度比較高，為了解決這個問題，業界提出了很多最佳化 Attention 計算的方法。

LongLoRA 是香港中文大學聯合 MIT 提出的一種模型微調方法，其論文 LONGLORA: EFFICIENT FINE-TUNING OF LONG-CONTEXT LARGE LANGUAGE MODELS[10]。下圖是論文[10]中描述的方法：

LongLoRA 提出了一種移位稀疏注意力（Shifted Sparse Attention，S2-Attn）來近似標準的自注意力。

在傳統的自注意力（self-attention）中，模型需要計算輸入序列中所有元素對之間的注意力權重，這在處理長序列時會導致計算複雜度呈二次方增長。S2-Attn 在訓練時，將輸入序列劃分成若干個組，在每個組內部進行自注意力計算。為了使不同組之間有資訊流通，在一半的注意力頭內，向其中一個組的 Tokens 做平（Shift）操作，平移的長度為組的一半。這樣就引入了不同組之間的資訊交換，又不增加計算量。

S2-Attn 的設計使得大型語言模型能夠在處理長序列時保持較高的效能，同時顯著降低了訓練和推理時的計算資源需求。

實驗結果如下，用了一個 8 個 A100 的機器微調，將 Llama2 7B/13B/70B 模型分別擴充套件到 100K，64K, 32K 長度，而大模型的困惑度並沒有明顯變化。

最近 Technology Innovation Institute(TII)發表了一篇論文綜述The What, Why, and How of Context Length Extension Techniques in Large Language Models – A Detailed Survey[11]，調研了業界的擴充套件大模型的技術，下圖源自論文。

論文中介紹了更多的業界上下文擴充套件的方法，大致可以簡單分為一下幾大主要方式。

上面我們分別講解了 Llama 的結構，然後基於 Llama 的結構去講解了業界最新擴充套件大模型上下文長度的方法與效果。

在 KubeAI 訓練推理平臺上，使用者只需要上傳資料、選擇大模型，就可以完成一次訓練和推理部署。如果想了解詳細使用方法，可以參考我們之前發表的關於得物大模型平臺的系列介紹的文章。

KubeAI 平臺為使用者提供了非常便捷的大模型訓練和部署功能。使用者無需關注底層基礎設施，就可以透過簡單的步驟上傳資料、配置引數、選擇模型，從而獲得針對自己業務自定義的大模型。

本文從 Llama 大模型的結構入手，介紹了其模組結構，重點解析了 Attention 機制中的 RoPE 層。要實現 Llama 模型上下文長度的擴充套件，需要對應擴充套件 RoPE 位置編碼層。但是直接擴充套件會導致模型困惑度上升，針對這個問題，我們介紹了業界常見的幾種上下文擴充套件方法，包括位置查詢法、動態插值法和 Yarn 方法等。

長下文推理對效能要求比較高，為此我們也介紹了一些為了提升效能而最佳化 Attention 的方法，比如 LongLoRA[10] 這篇論文的 S2-Atten 的方法。有興趣的同學可以閱讀相關論文了解細節。

本文透過剖析 Llama 模型結構，解析上下文擴充套件的關鍵層 RoPE，並概述各種擴充套件方法的原理，希望能夠幫助大家對大模型上下文擴充套件有一個系統的瞭解。後續如果有機會，我們會繼續分享更多大模型的核心技術，讓更多人對大模型的內在機制有更深的認識。歡迎持續關注我們的內容和分享！

參考資料

• [1] Llama 2: Open Foundation and Fine-Tuned Chat Models()
• [2] Attention Is All You Need()
• [3]modeling_llama.py()
• [4] RoFormer: Enhanced Transformer with Rotary Position Embedding()
• [5] Extending Context Window of Large Language Models via Positional Interpolation()
• [6] NTK-Aware Scaled RoPE allows LLaMA models to have extended (8k+) context size without any fine-tuning()
• [7] Extend context size without fine-tuning(https://github.com/ymcui/Chinese-LLaMA-Alpaca/pull/705)
• [8] YaRN: Efficient Context Window Extension of Large Language Models()
• [9] Yarn()
• [10] LONGLORA: EFFICIENT FINE-TUNING OF LONG-CONTEXT LARGE LANGUAGE MODELS()
• [11] The What, Why, and How of Context Length Extension Techniques in Large Language Models – A Detailed
  Survey()