一、背景

大模型的上下文長度是指我們在使用大模型的時(shí)候，給大模型的輸入加上輸出的字符（Token）總數(shù)，這個(gè)數(shù)字會被限制，如果超過這個(gè)長度的字符會被大模型丟棄。目前開源的大模型上下文長度一般不長，比如 Llama 2 只有 4K，Code-Llama 系列因?yàn)樾枰斎氪a，擴(kuò)展到了 16K。閉源系列模型的提供了更長的上下文長度，比如 OpenAI 在其最新模型 GPT-4 Turbo 中提供了 128K 的上下文長度，Anthropic 的 Claude 2.1 模型提供了 200K 上下文長度。

一些場景需要較長上下文，比如，文檔翻譯需要將整篇文檔輸入給大模型進(jìn)行翻譯，長文檔內(nèi)容抽取需要大模型讀取整篇長文檔進(jìn)行內(nèi)容抽取，會議內(nèi)容總結(jié)則需要給大模型輸入會議聊天記錄進(jìn)行總結(jié)等。

想要得到一個(gè)長上下文的大模型，一般有兩種途徑。一種是大模型在初始階段被設(shè)置為長上下文，然后經(jīng)過預(yù)訓(xùn)練，指令微調(diào)，對齊訓(xùn)練等方式得到一個(gè)長上下文大模型。另外一種方式是選擇已經(jīng)訓(xùn)練好的大模型，通過技術(shù)改造擴(kuò)展其上下文長度，然后再進(jìn)行微調(diào)訓(xùn)練得到長上下文模型。

圖片

本文將基于比較火的 Llama 2 大模型的結(jié)構(gòu)[1]介紹上下文長度的方法與挑戰(zhàn)，然后探討一些業(yè)界流行的上下文長度擴(kuò)展的技術(shù)，最后給大家推薦下 KubeAI 大模型訓(xùn)練推理平臺可以上手實(shí)驗(yàn)。

二、LLAMA的結(jié)構(gòu)

Transformer的結(jié)構(gòu)

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

圖片

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

輸入嵌入層（Input Embedding Layer）：將輸入文本的詞或標(biāo)記轉(zhuǎn)換為向量表示，以便模型能夠理解它們。

多頭自注意力層（Multi-Head Self-Attention Layer）：幫助模型捕捉輸入序列中詞與詞之間的關(guān)系，使模型能夠了解上下文信息。

前饋神經(jīng)網(wǎng)絡(luò)層（Feed-Forward Neural Network Layer）：對多頭自注意力的輸出進(jìn)行進(jìn)一步的特征提取和變換，以增加模型的表示能力。

歸一化層（Layer Normalization Layer）：規(guī)范化每一層的輸出，有助于訓(xùn)練過程的穩(wěn)定性。

總的來說，Transformer 是一種強(qiáng)大的模型，它可以捕捉文本和序列數(shù)據(jù)中的長距離依賴關(guān)系，使其在翻譯、對話、摘要生成等自然語言處理任務(wù)中表現(xiàn)出色。這個(gè)模型已經(jīng)在各種應(yīng)用中取得了顯著的成功。感興趣的同學(xué)可以自行去網(wǎng)上搜索下 Transformer 的結(jié)構(gòu)，深入了解。

LLAMA的結(jié)構(gòu)

目前大多數(shù)生成式語言模型如 Llama 系列，僅僅采用了 Transformer 的 Decoder 模塊結(jié)構(gòu)。在 Huggingface 中,這種結(jié)構(gòu)通常被稱為 CausalLM，即因果語言模型（Causal Language Model）。下面我們來具體看一下 Llama 2 模型系列的結(jié)構(gòu)，Llama 2 相關(guān)的論文[1]。

（下圖是基于 Transfomers 代碼中的 LlamaModel 繪制而成,具體代碼參考 Transfomer 中的 modeling_llama.py[3]）

圖片

我們來解讀下上面 Llama 各層的結(jié)構(gòu)與作用，首先從輸入文本開始。會經(jīng)過下面各層：

• Input Embedding：將 Input 文本轉(zhuǎn)化為向量表，通過 nn.Embedding 實(shí)現(xiàn)。
• Llama Decoder Layer：Decoder 采用多層 Llama Decoder Layer。每一層包括自注意力（Llama Attention）和前饋網(wǎng)絡(luò)（Llama MLP）。自注意力用于捕捉文本中的長程依賴關(guān)系。前饋網(wǎng)絡(luò)進(jìn)行非線性映射。
• Llama RMSNorm：一種規(guī)范化方式，用于正則化每層的輸，起到預(yù)處理的作用。
• lm_head：一個(gè)線性層，將 Decoder 最后一層的輸出映射到詞典大小的維，以進(jìn)行后續(xù)的語言模型 Logits 計(jì)算。
• Llama Attention：多頭自注意力機(jī)制，用于建模文本中的依賴關(guān)系。將輸入表示切分為多個(gè)頭，然后在每個(gè)頭內(nèi)做點(diǎn)積注意力運(yùn)算。
• Llama MLP：采用 Gated Linear Units 的多層前饋網(wǎng)絡(luò)。進(jìn)行非線性變換來捕捉復(fù)雜模式。

總體上，Llama 通過堆疊多層自注意力和前饋網(wǎng)絡(luò)來表示文本語義，然后預(yù)測后續(xù)詞元。lm_head 負(fù)責(zé)將語義表示映射為具體的詞典 Logits。整個(gè)模型端到端通過語言模型目標(biāo)進(jìn)行訓(xùn)練。

上面各層中比較核心的是 Llama Attention 層，該層的結(jié)構(gòu)如下：

圖片

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

圖片

Llama 的 Attention 計(jì)算過程如下：

• 輸入會經(jīng)過線性變換，得到 Query(Q)、Key（K）和 Value（V）矩陣。
• 對 Q 和 K 應(yīng)用 RoPE 位置編碼。RoPE 包含旋轉(zhuǎn)的 Sin 和 Cos 編碼，會根據(jù)每個(gè) Token 的位置對其表示進(jìn)行旋轉(zhuǎn)。
• 用旋轉(zhuǎn)后的 Q 和 K 計(jì)算點(diǎn)積，得到注意力權(quán)重 Attention Score，經(jīng)過 Softmax 計(jì)算后得到 Normalized Attention Weight。
• 再把 Attention Weight 與 V 相乘，并進(jìn)行加權(quán)求和，得到 Attention 的輸出。
• 輸出再經(jīng)過一個(gè)線性變換，繼續(xù)輸出給下一層作為輸入。

這樣通過 RoPE 位置編碼、加權(quán)平均，Llama 的 Attention 可以高效穩(wěn)定地提取文本序列的上下文語義信息。

三、擴(kuò)展方案與挑戰(zhàn)

位置編碼層(RoPE)

通過上面對 Llama 結(jié)構(gòu)的解析，我們看到在 Llama Attention 中有一個(gè)叫做 RoPE（旋轉(zhuǎn)位置編碼）的層，主要用于對輸入進(jìn)行位置編碼，讓模型學(xué)到輸入文本中每個(gè) Token 的位置關(guān)系，從而更好地理解輸入。RoPE 層能處理序列的長度決定了 Llama 的上下文長度，要擴(kuò)展 Llama 的上下文長度，需要對 RoPE 層進(jìn)行改造和擴(kuò)展。下面我們先簡單介紹下 Llama RoPE 層的工作原理。

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

圖片

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

論文中證明，在進(jìn)行旋轉(zhuǎn)位置編碼之后，可以從新的編碼向量中獲取原向量的相對位置信息，即論文中下面的公式中的 m-（位置 m 減去位置 n）。下圖源自論文[4]。

圖片

而我們只需要理解旋轉(zhuǎn)位置編碼的最終計(jì)算公式如下，對于一個(gè)輸入向量 X，直接與一個(gè) COS 矩陣和 SIN 矩陣進(jìn)行內(nèi)積后求和：下圖源自論文[4]。

圖片

其中 WCOS 和 WSIN 分別是一個(gè)預(yù)先固定的 COS 矩陣和 SIN 矩陣。

下面我們展示下 Huggingface 的 Transformer 對應(yīng)的計(jì)算 COS 和 SIN 矩陣的計(jì)算代碼：

圖片

擴(kuò)展位置編碼層(RoPE)支持長上下文

上面介紹了 Llama 結(jié)構(gòu)中的旋轉(zhuǎn)位置編碼層 RoPE。要擴(kuò)展大模型的上下文長度，就需要擴(kuò)展 RoPE 層，也就是擴(kuò)展其 COS 和 SIN 矩陣，讓RoPE支持更長序列的輸入。

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

對于未訓(xùn)練過的大模型，只需要直接更改其配置文件中的 max_position_embeddings 即可實(shí)現(xiàn) RoPE 層的擴(kuò)展，然后再進(jìn)行訓(xùn)練。但是對于已經(jīng)訓(xùn)練過的模型，如果直接修改其配置，會導(dǎo)致模型的效果急劇下降，后面第四部分我們會介紹一些基于已有模型進(jìn)行改造擴(kuò)展的上下文的方法。max_position_embeddings 的配置如下：

圖片

超長上下文面臨的挑戰(zhàn)

超長上下文的大模型部署推理的時(shí)候，往往會面臨如下性能挑戰(zhàn)。

• 推理時(shí)間變長

從上面的 Attention 的計(jì)算公式可以看出，Attention 進(jìn)行了點(diǎn)積的運(yùn)算，其時(shí)間復(fù)雜度為 L（序列長度）的平方。也就是說大模型在推理的時(shí)候，輸入的序列長度越長推理時(shí)間越多。所以超長上下文的大模型需要更多的推理時(shí)間，這會帶來用戶體驗(yàn)上的損失。

• 推理顯存空間變大

大模型在持續(xù)推理的過程中，需要緩存一個(gè)叫做 KV Cache 的數(shù)據(jù)快，KV Cache 的大小也與序列長度成正比。以 Llama 2 13B 大模型為例，一個(gè) 4K 長的序列大約需要 3G 的顯存去緩存 KV Cache，16K 的序列則需要 12G，128K 的序列則需要 100G 顯存。

圖片

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

大模型上下文擴(kuò)展的思路

綜上所述，擴(kuò)展大模型的上下文長度，一般思路如下：

圖片

• 首先通過對位置編碼層進(jìn)行改造，使其支持更長的上下文。
• 為了取得更好的推理性能，還需要對 Attention 計(jì)算進(jìn)行優(yōu)化。
• 進(jìn)行微調(diào)訓(xùn)練，讓大模型適應(yīng)新的模型結(jié)構(gòu)。

四、位置編碼層改造擴(kuò)展上下文的案例

上面我們講到，從模型物理結(jié)構(gòu)上擴(kuò)展上下文長度，需要直接修改 RoPE 層，即直接擴(kuò)展其 SIN 和 COS 矩陣。但是大模型都是基于大量短序列數(shù)據(jù)訓(xùn)練得到的。如果直接強(qiáng)行擴(kuò)展，會導(dǎo)致模型困惑度提高。所謂困惑度是模型對下一個(gè)詞的預(yù)測困難程度的量化指標(biāo)，直觀意義是大模型的輸出是否能夠更容易被人類所理解。

因此，我們需要更好的方法來擴(kuò)展預(yù)訓(xùn)練模型的上下文長度，既要兼顧模型性能，又要控制困惑度。下面我們概括幾種業(yè)界常用的上下文長度擴(kuò)展方法。

圖片

線性位置插值法

線性插值法的思想最早來自于這篇文章，Extending Context Window of Large Language Models via Positional Interpolation[5]，該方法已經(jīng)被 Huggingface 的 Transformer 中 Llama 模型代碼集成。

  下圖源自論文：

圖片

思路：通過線性縮小輸入位置索引以匹配原始上下文窗口大小，而不是超出訓(xùn)練上下文長度進(jìn)行外推，這樣可以減小注意力機(jī)制中相對位置的影響，幫助模型更容易適應(yīng)擴(kuò)展后的上下文窗口。

效果：在從 LLaMA 7B 到 65B 模型上，通過位置插值擴(kuò)展上下文窗口到 32768（4k擴(kuò)展到32K），僅需微調(diào) 1000 步，就能在包括語言建模、長文檔摘要撰寫等任務(wù)上取得良好效果。

優(yōu)點(diǎn)：位置插值不僅能有效擴(kuò)展上下文窗口，提高模型在長上下文任務(wù)上的性能，還能在原有上下文窗口大小的任務(wù)上保持模型質(zhì)量，且不需要額外的權(quán)重或修改模型架構(gòu)。

缺點(diǎn)：需要重新訓(xùn)練，有時(shí)候擴(kuò)充后會導(dǎo)致模型困惑度上升。

動態(tài)插值法(NTK-awared)

動態(tài)插值法是在位置插值法的基礎(chǔ)上演變而來的，最早提出文章 NTK-Aware Scaled RoPE allows LLaMA models to have extended (8k+) context size without any fine-tuning[6]?，F(xiàn)在也被 Huggingface 的 Transformer 中 Llama 模型代碼集成。

下面是 Chinese Llama 基于位置插值法與動態(tài)插值法進(jìn)行的比較，數(shù)據(jù)來自 Extend Context Size Without Fine-Tuning[7]。

圖片

數(shù)據(jù)顯示，相比與位置插值法，NTK 動態(tài)插值法不會顯著增加大模型的困惑度。

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

效果：與線性插值相比，在大幅擴(kuò)展上下文（如8000+）時(shí)，無需微調(diào)就可以使困惑度下降極小。

優(yōu)點(diǎn)：微調(diào)成本極低，上下文窗口可以擴(kuò)展很大，困惑度變化小。

總體來說，動態(tài)插值法通過考慮模型特性，設(shè)計(jì)更優(yōu)化的插值方案，能夠在不增加訓(xùn)練成本的條件下，獲得接近無損的上下文窗口擴(kuò)展效果。這為進(jìn)一步擴(kuò)展和優(yōu)化大語言模型提供了新的思路。

Yarn(NTK升級)

Yarn 擴(kuò)展上下文的方法來自于文章 YaRN: Efficient Context Window Extension of Large Language Models[8]。代碼參考Yarn[9]。并給出基于 Llama2 的 128K 上下文擴(kuò)展。

下圖源自論文[8]：

圖片

數(shù)據(jù)顯示，在 128K 的 Proof-Pile 數(shù)據(jù)集上評測，Yarn-Llama-2-7b-128K/64K 模型的困惑度仍然保持良好下降。

相比于僅進(jìn)行簡單線性插值或動態(tài)插值的方法，Yarn 方法更全面地考慮了不同頻率 RoPE 維度的作用，避免了信息損失和外推問題。這使得 Yarn 方法在不 Fine-Tuning 的情況下，以及 Fine-Tuning 的數(shù)據(jù)量很少的情況下，都能更好地?cái)U(kuò)展上下文窗口。

RoPE 中的每個(gè)維度對應(yīng)著不同的正弦波頻率。高頻的正弦波編碼了位置信息的細(xì)微變化，低頻的正弦波編碼了位置信息的整體趨勢。

如果我們簡單地進(jìn)行線性插值，會把所有頻率的正弦波都等比例地拉伸。這會導(dǎo)致兩個(gè)問題:

• 高頻正弦波被過度拉伸，導(dǎo)致代表細(xì)微位置變化的信息丟失。這個(gè)會影響模型區(qū)分很接近的詞的能力。
• 低頻正弦波被拉伸，不同位置之間的相對距離變小。這會導(dǎo)致模型判斷近距離詞的先后順序變得困難。

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

• 對高頻正弦波幾乎不進(jìn)行插值，保留細(xì)微位置信息。
• 對低頻正弦波進(jìn)行接近線性的插值，保留位置大體信息。
• 中頻正弦波進(jìn)行漸變的插值。

這樣既保留了高頻表示細(xì)微位置變化的信息，也保留了低頻表示位置整體關(guān)系的信息，避免了簡單線性插值的問題。

五、優(yōu)化Attention擴(kuò)展上下文的案例

上面我們提到，長上下文對大模型的正向與反向傳播的性能來說是個(gè)挑戰(zhàn)。其主要原因是 Attention（注意力）的計(jì)算復(fù)雜度比較高，為了解決這個(gè)問題，業(yè)界提出了很多優(yōu)化 Attention 計(jì)算的方法。

LongLoRA方法

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

圖片

LongLoRA 提出了一種移位稀疏注意力（Shifted Sparse Attention，S2-Attn）來近似標(biāo)準(zhǔn)的自注意力。

在傳統(tǒng)的自注意力（self-attention）中，模型需要計(jì)算輸入序列中所有元素對之間的注意力權(quán)重，這在處理長序列時(shí)會導(dǎo)致計(jì)算復(fù)雜度呈二次方增長。S2-Attn 在訓(xùn)練時(shí)，將輸入序列劃分成若干個(gè)組，在每個(gè)組內(nèi)部進(jìn)行自注意力計(jì)算。為了使不同組之間有信息流通，在一半的注意力頭內(nèi)，向其中一個(gè)組的 Tokens 做平（Shift）操作，平移的長度為組的一半。這樣就引入了不同組之間的信息交換，又不增加計(jì)算量。

S2-Attn 的設(shè)計(jì)使得大型語言模型能夠在處理長序列時(shí)保持較高的性能，同時(shí)顯著降低了訓(xùn)練和推理時(shí)的計(jì)算資源需求。

實(shí)驗(yàn)結(jié)果如下，用了一個(gè) 8 個(gè) A100 的機(jī)器微調(diào)，將 Llama2 7B/13B/70B 模型分別擴(kuò)展到 100K，64K, 32K 長度，而大模型的困惑度并沒有明顯變化。

圖片

六、業(yè)界更多擴(kuò)展上下文的方法

最近 Technology Innovation Institute(TII)發(fā)表了一篇論文綜述The What, Why, and How of Context Length Extension Techniques in Large Language Models – A Detailed Survey[11]，調(diào)研了業(yè)界的擴(kuò)展大模型的技術(shù)，下圖源自論文。

圖片

論文中介紹了更多的業(yè)界上下文擴(kuò)展的方法，大致可以簡單分為一下幾大主要方式。

圖片

七、Kubeai大模型訓(xùn)練推理平臺

上面我們分別講解了 Llama 的結(jié)構(gòu)，然后基于 Llama 的結(jié)構(gòu)去講解了業(yè)界最新擴(kuò)展大模型上下文長度的方法與效果。

圖片

在 KubeAI 訓(xùn)練推理平臺上，用戶只需要上傳數(shù)據(jù)、選擇大模型，就可以完成一次訓(xùn)練和推理部署。如果想了解詳細(xì)使用方法，可以參考我們之前發(fā)表的關(guān)于得物大模型平臺的系列介紹的文章。

KubeAI 平臺為用戶提供了非常便捷的大模型訓(xùn)練和部署功能。用戶無需關(guān)注底層基礎(chǔ)設(shè)施，就可以通過簡單的步驟上傳數(shù)據(jù)、配置參數(shù)、選擇模型，從而獲得針對自己業(yè)務(wù)自定義的大模型。

八、總結(jié)與展望

本文從 Llama 大模型的結(jié)構(gòu)入手，介紹了其模塊結(jié)構(gòu)，重點(diǎn)解析了 Attention 機(jī)制中的 RoPE 層。要實(shí)現(xiàn) Llama 模型上下文長度的擴(kuò)展，需要對應(yīng)擴(kuò)展 RoPE 位置編碼層。但是直接擴(kuò)展會導(dǎo)致模型困惑度上升，針對這個(gè)問題，我們介紹了業(yè)界常見的幾種上下文擴(kuò)展方法，包括位置查找法、動態(tài)插值法和 Yarn 方法等。

長下文推理對性能要求比較高，為此我們也介紹了一些為了提升性能而優(yōu)化 Attention 的方法，比如 LongLoRA[10] 這篇論文的 S2-Atten 的方法。有興趣的同學(xué)可以閱讀相關(guān)論文了解細(xì)節(jié)。

本文通過剖析 Llama 模型結(jié)構(gòu)，解析上下文擴(kuò)展的關(guān)鍵層 RoPE，并概述各種擴(kuò)展方法的原理，希望能夠幫助大家對大模型上下文擴(kuò)展有一個(gè)系統(tǒng)的了解。后續(xù)如果有機(jī)會，我們會繼續(xù)分享更多大模型的核心技術(shù)，讓更多人對大模型的內(nèi)在機(jī)制有更深的認(rèn)識。歡迎持續(xù)關(guān)注我們的內(nèi)容和分享！

參考資料

• [1] Llama 2: Open Foundation and Fine-Tuned Chat Models(https://arxiv.org/abs/2307.09288)
• [2] Attention Is All You Need(https://arxiv.org/pdf/1706.03762.pdf)
• [3]modeling_llama.py(https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py)
• [4] RoFormer: Enhanced Transformer with Rotary Position Embedding(https://arxiv.org/abs/2104.09864)
• [5] Extending Context Window of Large Language Models via Positional Interpolation(https://arxiv.org/abs/2306.15595)
• [6] NTK-Aware Scaled RoPE allows LLaMA models to have extended (8k+) context size without any fine-tuning(https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/)
• [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(https://arxiv.org/abs/2309.00071)
• [9] Yarn(https://github.com/jquesnelle/yarn)
• [10] LONGLORA: EFFICIENT FINE-TUNING OF LONG-CONTEXT LARGE LANGUAGE MODELS(https://arxiv.org/pdf/2309.12307.pdf)
• [11] The What, Why, and How of Context Length Extension Techniques in Large Language Models – A DetailedSurvey(https://arxiv.org/pdf/2401.07872.pdf)