如何幫你節(jié)省顯存)
MHA還是GQAStableLM-3B-4E1T的num_key_value_heads參數(shù)如何幫你節(jié)省顯存【免費(fèi)下載鏈接】stablelm-3b-4e1t項(xiàng)目地址: https://ai.gitcode.com/hf_mirrors/ai-gitcode/stablelm-3b-4e1tStableLM-3B-4E1T 是 Stability AI 開(kāi)源的 30 億參數(shù)語(yǔ)言模型其config.json中的num_key_value_heads參數(shù)決定了模型使用 MHA 還是 GQA 注意力機(jī)制直接影響推理時(shí)的 KV Cache 顯存占用。這篇文章帶你快速看懂這個(gè)參數(shù)并告訴你如何用它省下顯存。一句話說(shuō)清MHA、GQA、MQA 到底是什么在 Transformer 解碼器中每個(gè)注意力頭都會(huì)計(jì)算 QueryQ、KeyK、ValueV。區(qū)別在于 K 和 V 各有幾套注意力機(jī)制KV 頭數(shù)KV Cache 大小MHA多頭注意力與 Q 頭數(shù)相同最大GQA分組查詢注意力介于兩者之間適中MQA多查詢注意力只有 1 個(gè)最小判斷規(guī)則就藏在 configuration_stablelm.py 的參數(shù)文檔里num_key_value_headsnum_attention_heads→ MHAnum_key_value_heads1→ MQA其他情況 → GQA。查看 StableLM-3B-4E1T 的默認(rèn)配置打開(kāi)倉(cāng)庫(kù)根目錄的 config.json可以看到關(guān)鍵參數(shù)參數(shù)值含義hidden_size2560隱藏層維度num_hidden_layers32解碼層數(shù)num_attention_heads32Q 注意力頭數(shù)num_key_value_heads32K/V 頭數(shù)max_position_embeddings4096最大序列長(zhǎng)度num_key_value_heads32與num_attention_heads32相等所以StableLM-3B-4E1T 默認(rèn)就是標(biāo)準(zhǔn) MHA——K/V 各 32 個(gè)頭KV Cache 開(kāi)銷(xiāo)是全頭配置中最大的。KV Cache 為什么是顯存大頭推理時(shí)權(quán)重只占一部分顯存bf16 下約 5.6GB真正隨序列長(zhǎng)度增長(zhǎng)的是 KV Cache估算公式為顯存 ≈ 2K和V× 層數(shù) × KV頭數(shù) × 頭維度 × 序列長(zhǎng)度 × 每元素字節(jié)數(shù) × 批大小以 StableLM-3B-4E1T 為例頭維度 2560 ÷ 32 80bf16 即 2 字節(jié)單條 4096 長(zhǎng)度序列KV 頭數(shù)注意力類(lèi)型單序列 KV Cache4096 長(zhǎng)度32默認(rèn)MHA約 1.3GB8GQA約 0.33GB1MQA約 40MB可以看到把 KV 頭從 32 降到 8KV Cache 直接縮小到 1/4。對(duì)長(zhǎng)文本生成或大批量并發(fā)服務(wù)來(lái)說(shuō)這就是實(shí)打?qū)嵉娘@存節(jié)省。在源碼里找到它repeat_kv 的分組擴(kuò)展邏輯這個(gè)參數(shù)的運(yùn)行時(shí)作用體現(xiàn)在 modeling_stablelm.py 中StableLmAttention類(lèi)里計(jì)算num_key_value_groups num_heads // num_key_value_heads即每個(gè) KV 頭要服務(wù)的 Q 頭組數(shù)k_proj/v_proj的投影維度是num_key_value_heads × head_dim而不是全部 Q 頭數(shù)repeat_kv函數(shù)在注意力計(jì)算前把每個(gè) KV 頭沿分組方向復(fù)制擴(kuò)展成與 Q 頭對(duì)齊。當(dāng) MHA組數(shù)為 1時(shí)它不做任何復(fù)制GQA 時(shí)才真正一魚(yú)多吃。這正是 GQA 省顯存的本質(zhì)權(quán)重和 KV Cache 都按少量 KV 頭存儲(chǔ)計(jì)算時(shí)再按需復(fù)用。如何正確修改 num_key_value_heads 來(lái)省顯存??重要提醒不能只改配置文件里的數(shù)字。k_proj、v_proj的權(quán)重形狀由 KV 頭數(shù)決定直接改小會(huì)導(dǎo)致權(quán)重?zé)o法加載。正確做法分兩步第一步MHA → GQA 權(quán)重轉(zhuǎn)換按官方配置文檔的說(shuō)明把每組內(nèi)的原始 K/V 頭做平均池化mean pooling得到新的分組頭權(quán)重再按新num_key_value_heads保存 checkpoint。第二步部署時(shí)啟用更多省顯存手段使用 Flash Attention 2加載模型時(shí)設(shè)置attn_implementationflash_attention_2官方 README 已驗(yàn)證支持??結(jié)合量化如 8bit / 4bit 加載進(jìn)一步壓縮權(quán)重與緩存。一張表總結(jié)新手省顯存速查你的場(chǎng)景建議短對(duì)話、單用戶、12GB 顯卡默認(rèn) MHA 配置即可權(quán)重 5.6GB KV 約 1.3GB 完全夠用長(zhǎng)上下文接近 4096、多用戶并發(fā)轉(zhuǎn)換為 8 頭 GQAKV Cache 縮小 75%極致顯存優(yōu)先、可接受精度微損進(jìn)一步轉(zhuǎn)為 MQAKV 頭數(shù) 1結(jié)論StableLM-3B-4E1T 默認(rèn)num_key_value_heads32是標(biāo)準(zhǔn) MHA 配置理解這個(gè)參數(shù)后你可以通過(guò)MHA→GQA 平均池化轉(zhuǎn)換 Flash Attention 2的組合在不改變模型主體結(jié)構(gòu)的前提下把推理時(shí)的 KV Cache 顯存最多節(jié)省 97%?！久赓M(fèi)下載鏈接】stablelm-3b-4e1t項(xiàng)目地址: https://ai.gitcode.com/hf_mirrors/ai-gitcode/stablelm-3b-4e1t創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考 MHA還是GQA？StableLM-3B-4E1T的num_key_value_heads參數(shù)如何幫你節(jié)省顯存
MHA還是GQAStableLM-3B-4E1T的num_key_value_heads參數(shù)如何幫你節(jié)省顯存【免費(fèi)下載鏈接】stablelm-3b-4e1t項(xiàng)目地址: https://ai.gitcode.com/hf_mirrors/ai-gitcode/stablelm-3b-4e1tStableLM-3B-4E1T 是 Stability AI 開(kāi)源的 30 億參數(shù)語(yǔ)言模型其config.json中的num_key_value_heads參數(shù)決定了模型使用 MHA 還是 GQA 注意力機(jī)制直接影響推理時(shí)的 KV Cache 顯存占用。這篇文章帶你快速看懂這個(gè)參數(shù)并告訴你如何用它省下顯存。一句話說(shuō)清MHA、GQA、MQA 到底是什么在 Transformer 解碼器中每個(gè)注意力頭都會(huì)計(jì)算 QueryQ、KeyK、ValueV。區(qū)別在于 K 和 V 各有幾套注意力機(jī)制KV 頭數(shù)KV Cache 大小MHA多頭注意力與 Q 頭數(shù)相同最大GQA分組查詢注意力介于兩者之間適中MQA多查詢注意力只有 1 個(gè)最小判斷規(guī)則就藏在 configuration_stablelm.py 的參數(shù)文檔里num_key_value_headsnum_attention_heads→ MHAnum_key_value_heads1→ MQA其他情況 → GQA。查看 StableLM-3B-4E1T 的默認(rèn)配置打開(kāi)倉(cāng)庫(kù)根目錄的 config.json可以看到關(guān)鍵參數(shù)參數(shù)值含義hidden_size2560隱藏層維度num_hidden_layers32解碼層數(shù)num_attention_heads32Q 注意力頭數(shù)num_key_value_heads32K/V 頭數(shù)max_position_embeddings4096最大序列長(zhǎng)度num_key_value_heads32與num_attention_heads32相等所以StableLM-3B-4E1T 默認(rèn)就是標(biāo)準(zhǔn) MHA——K/V 各 32 個(gè)頭KV Cache 開(kāi)銷(xiāo)是全頭配置中最大的。KV Cache 為什么是顯存大頭推理時(shí)權(quán)重只占一部分顯存bf16 下約 5.6GB真正隨序列長(zhǎng)度增長(zhǎng)的是 KV Cache估算公式為顯存 ≈ 2K和V× 層數(shù) × KV頭數(shù) × 頭維度 × 序列長(zhǎng)度 × 每元素字節(jié)數(shù) × 批大小以 StableLM-3B-4E1T 為例頭維度 2560 ÷ 32 80bf16 即 2 字節(jié)單條 4096 長(zhǎng)度序列KV 頭數(shù)注意力類(lèi)型單序列 KV Cache4096 長(zhǎng)度32默認(rèn)MHA約 1.3GB8GQA約 0.33GB1MQA約 40MB可以看到把 KV 頭從 32 降到 8KV Cache 直接縮小到 1/4。對(duì)長(zhǎng)文本生成或大批量并發(fā)服務(wù)來(lái)說(shuō)這就是實(shí)打?qū)嵉娘@存節(jié)省。在源碼里找到它repeat_kv 的分組擴(kuò)展邏輯這個(gè)參數(shù)的運(yùn)行時(shí)作用體現(xiàn)在 modeling_stablelm.py 中StableLmAttention類(lèi)里計(jì)算num_key_value_groups num_heads // num_key_value_heads即每個(gè) KV 頭要服務(wù)的 Q 頭組數(shù)k_proj/v_proj的投影維度是num_key_value_heads × head_dim而不是全部 Q 頭數(shù)repeat_kv函數(shù)在注意力計(jì)算前把每個(gè) KV 頭沿分組方向復(fù)制擴(kuò)展成與 Q 頭對(duì)齊。當(dāng) MHA組數(shù)為 1時(shí)它不做任何復(fù)制GQA 時(shí)才真正一魚(yú)多吃。這正是 GQA 省顯存的本質(zhì)權(quán)重和 KV Cache 都按少量 KV 頭存儲(chǔ)計(jì)算時(shí)再按需復(fù)用。如何正確修改 num_key_value_heads 來(lái)省顯存??重要提醒不能只改配置文件里的數(shù)字。k_proj、v_proj的權(quán)重形狀由 KV 頭數(shù)決定直接改小會(huì)導(dǎo)致權(quán)重?zé)o法加載。正確做法分兩步第一步MHA → GQA 權(quán)重轉(zhuǎn)換按官方配置文檔的說(shuō)明把每組內(nèi)的原始 K/V 頭做平均池化mean pooling得到新的分組頭權(quán)重再按新num_key_value_heads保存 checkpoint。第二步部署時(shí)啟用更多省顯存手段使用 Flash Attention 2加載模型時(shí)設(shè)置attn_implementationflash_attention_2官方 README 已驗(yàn)證支持??結(jié)合量化如 8bit / 4bit 加載進(jìn)一步壓縮權(quán)重與緩存。一張表總結(jié)新手省顯存速查你的場(chǎng)景建議短對(duì)話、單用戶、12GB 顯卡默認(rèn) MHA 配置即可權(quán)重 5.6GB KV 約 1.3GB 完全夠用長(zhǎng)上下文接近 4096、多用戶并發(fā)轉(zhuǎn)換為 8 頭 GQAKV Cache 縮小 75%極致顯存優(yōu)先、可接受精度微損進(jìn)一步轉(zhuǎn)為 MQAKV 頭數(shù) 1結(jié)論StableLM-3B-4E1T 默認(rèn)num_key_value_heads32是標(biāo)準(zhǔn) MHA 配置理解這個(gè)參數(shù)后你可以通過(guò)MHA→GQA 平均池化轉(zhuǎn)換 Flash Attention 2的組合在不改變模型主體結(jié)構(gòu)的前提下把推理時(shí)的 KV Cache 顯存最多節(jié)省 97%?！久赓M(fèi)下載鏈接】stablelm-3b-4e1t項(xiàng)目地址: https://ai.gitcode.com/hf_mirrors/ai-gitcode/stablelm-3b-4e1t創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考