
MindSpeed LLM長序列并行指南Ring Attention與Ulysses上下文并行詳解【免費(fèi)下載鏈接】MindSpeed-LLM昇騰LLM分布式訓(xùn)練框架項(xiàng)目地址: https://gitcode.com/Ascend/MindSpeed-LLMMindSpeed LLM是昇騰 NPU 上的 LLM 分布式訓(xùn)練框架其上下文并行Context ParallelCP能力專門解決長序列訓(xùn)練難題通過Ring Attention與Ulysses兩大算法把超長序列切分到多卡并行計(jì)算讓 128K 甚至更長的序列訓(xùn)練變得可行。本文面向新手講清兩種長序列并行方案的原理差異、核心參數(shù)與最佳實(shí)踐。為什么長序列訓(xùn)練這么難隨著會(huì)話式 AI、長文檔摘要、代碼庫理解等場景普及訓(xùn)練序列長度從 4K 一路飆升到 128K 以上。麻煩在于顯存爆炸自注意力矩陣隨序列長度 S 呈O(S2)增長單卡根本放不下傳統(tǒng)并行不覆蓋序列維數(shù)據(jù)并行、張量并行、流水線并行都不切分序列維度S 變大時(shí)單步顯存依然失控。因此需要一種沿序列維度切分的并行方式這就是上下文并行CP要解決的問題。兩大上下文并行算法Ring Attention vs UlyssesMindSpeed LLM 提供多種 CP 算法其中最具代表性的是megatron_cp_algoRing Attention和ulysses_cp_algoUlysses通過參數(shù)--context-parallel-algo一鍵切換。 Ring Attention環(huán)狀分塊 KV 通信Ring Attention 借鑒分塊 Softmax 原理把序列切成 N 塊每塊由一個(gè) CP rank 持有本地 QKV。計(jì)算時(shí)各 rank 先做本地 attention再通過環(huán)狀Ring通信把 KV 塊逐段傳給下一個(gè) rank邊傳邊算循環(huán)一輪后得到全局完整結(jié)果。核心特點(diǎn)通信與計(jì)算可互相掩蓋KV 塊在環(huán)上傳遞的同時(shí)本地 attention 繼續(xù)算通信開銷幾乎被藏掉無需數(shù)據(jù)拼接全程分塊計(jì)算支持的理論序列長度近乎無限無頭維整除約束不要求 head_size 能被 cp_size 整除適配性廣要求FlashAttention 必須開啟且序列分塊足夠長才能掩蓋好通信。核心實(shí)現(xiàn)在 AttentionWithCp其中前向過程會(huì)構(gòu)建內(nèi)外兩個(gè) Ring P2P 通信結(jié)構(gòu)按 rank 順序逐塊取 KV 并更新 softmax 歸一化因子。參數(shù)注冊與合法性校驗(yàn)集中在 context_parallel_feature.py例如 CP 不支持 ALiBi 位置編碼、推理 KV Cache 等組合會(huì)被直接攔截。 UlyssesAll-to-All 序列-頭轉(zhuǎn)置Ulysses 的思路完全不同每個(gè) rank 持有完整序列但只負(fù)責(zé) 1/N 的注意力頭。每層 attention 前后各做一次all-to-all通信把切序列的布局轉(zhuǎn)成切頭的布局來計(jì)算算完再轉(zhuǎn)回來。特點(diǎn)通信量與序列長度解耦CP 較小時(shí)經(jīng)驗(yàn)上 CP ≤ 4通信占比低吞吐更高對注意力頭數(shù)量有整除要求頭數(shù)需能被 cp_size 整除框架內(nèi)置--kv-head-repeat-before-uly-alltoall開關(guān)支持 GQA/MQA 模型在 all-to-all 前擴(kuò)展 KV 頭見 ulysses_context_parallel.py。一圖看懂兩種方案怎么選維度Ring Attentionmegatron_cp_algoUlyssesulysses_cp_algo切分方式每 rank 持有序列分塊環(huán)傳 KV每 rank 持有全序列切注意力頭通信模式P2P 環(huán)狀逐塊傳遞可重疊計(jì)算每層兩次 all-to-all頭數(shù)整除要求無需 head 數(shù)被 cp_size 整除適合場景CP 較大、超長序列32K/128KCP 較小≤4、序列中等偏長顯存/延遲特點(diǎn)通信延遲略高但可掩蓋CP 小時(shí)延遲更低、吞吐更高官方長序列微調(diào)文檔中也有實(shí)測佐證Llama2-7B、32K 序列、TP2/CP4 配置下Ulysses 吞吐 192.3 TFLOP/s/GPU高于 Ring 的 102.7 TFLOP/s/GPU詳見 fine-tuning-with-context-parallel.md??焖偕鲜株P(guān)鍵參數(shù)與配置方法一鍵開啟上下文并行步驟設(shè)置并行規(guī)模與算法--context-parallel-sizeCP 卡數(shù)--context-parallel-algo對齊序列長度--seq-length必須能被 cp_size 整除選擇掩碼類型訓(xùn)練建議--attention-mask-type causal開啟通信優(yōu)化--use-cp-send-recv-overlap讓 send/recv 重疊。最快配置方法參數(shù)速查表參數(shù)說明--context-parallel-sizeCP 并行卡數(shù)默認(rèn) 1--context-parallel-algo可選megatron_cp_algoRing/ulysses_cp_algo/hybrid_cp_algo/kvallgather_cp_algo等--seq-length序列總長度需被 cp_size 整除--attention-mask-typecausal倒三角推薦或general全量--use-cp-send-recv-overlap建議開啟掩蓋 CP 通信延遲--cp-window-sizeRing 窗口大小默認(rèn) 1典型超長序列訓(xùn)練配置示例--seq-length 131072 \ --context-parallel-size 8 \ --context-parallel-algo megatron_cp_algo \ --attention-mask-type causal \ --use-cp-send-recv-overlap實(shí)戰(zhàn)避坑讓長序列并行跑出最佳性能分塊長度別太小Ring Attention 的通信掩蓋依賴足夠長的計(jì)算塊。經(jīng)驗(yàn)法則是seq-length / context-parallel-size 8K否則 8K 以下短分塊會(huì)導(dǎo)致通信反而慢于計(jì)算得不償失。Mask 選 causalGPT 類模型訓(xùn)練場景默認(rèn)用causal性能與顯存都優(yōu)于general全量計(jì)算。FlashAttention 是硬依賴開啟 CP 時(shí)必須同時(shí)開啟 Flash Attention 特性否則功能不可用。組合約束要留意CP 與 ALiBi 位置編碼、KV Cache、超長滑窗等互斥DeepSeek V4 類模型需使用kvallgather_cp_algo或deepseek_v4_cp_algo這些都會(huì)在參數(shù)校驗(yàn)階段給出明確報(bào)錯(cuò)見 validate_args。打包微調(diào)場景多樣本 pack 訓(xùn)練時(shí)配合--reset-position-ids與--reset-attention-mask讓 attention mask 按 EOD 分句生成鋸齒狀掩碼與 CP 分塊計(jì)算正確配合方法詳見 長序列微調(diào)文檔。相關(guān)代碼與文檔在哪里官方文檔Ring Attention 原理與使用docs/zh/pytorch/features/mcore/ring-attention-context-parallel.md長序列 CP 微調(diào)指南docs/zh/pytorch/features/mcore/fine-tuning-with-context-parallel.mdCP 特性入口與參數(shù)定義mindspeed_llm/features_manager/context_parallel/context_parallel_feature.pyRing Attention 核心算子mindspeed_llm/core/context_parallel/ring_context_parallel.pyCP 數(shù)據(jù)切批工具mindspeed_llm/core/context_parallel/get_batch_utils.pyUlysses 特性封裝mindspeed_llm/features_manager/context_parallel/ulysses_context_parallel.py總結(jié)MindSpeed LLM 的上下文并行讓長序列訓(xùn)練不再是顯存噩夢Ring Attention以環(huán)狀 KV 通信 計(jì)算掩蓋取勝適合 CP 較大、序列超長的大規(guī)模場景Ulysses以 all-to-all 轉(zhuǎn)置取勝CP 較小時(shí)延遲更低、吞吐更高。記住三條經(jīng)驗(yàn)——序列分塊大于 8K、Mask 用 causal、FlashAttention 必開就能在昇騰集群上穩(wěn)定跑通 128K 級(jí)別的長序列訓(xùn)練?!久赓M(fèi)下載鏈接】MindSpeed-LLM昇騰LLM分布式訓(xùn)練框架項(xiàng)目地址: https://gitcode.com/Ascend/MindSpeed-LLM創(chuàng)作聲明:本文部分內(nèi)容由AI輔助生成(AIGC),僅供參考