指南)
Transformers 中的 DeepSpeed Ulysses 序列并行用多卡切分超長序列進行訓練的技術(shù)指南【免費下載鏈接】transformers Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.項目地址: https://gitcode.com/GitHub_Trending/tra/transformers本篇基于 Transformers 官方文檔deepspeed_alst.md及其配套源碼系統(tǒng)講解 Ulysses 序列并行Sequence Parallelism, SP在 TransformersTrainer中的工作機制、完整配置方式與關(guān)鍵參數(shù)含義并結(jié)合 Trainer 源碼 與 DeepSpeed 集成模塊 的調(diào)用鏈幫助你掌握在 2 卡及以上多 GPU 環(huán)境下訓練超長序列百萬 token 級上下文的完整實戰(zhàn)方案。Ulysses 序列并行的核心原理超長序列訓練的首要瓶頸是顯存注意力計算的中間激活隨序列長度平方增長單卡往往放不下一條長序列。Ulysses 序列并行通過“切分序列 兩次 all-to-all 通信”的方式讓多張 GPU 協(xié)同處理同一條序列前向開始時seq-sharded 布局整條長度為 N 的序列被均勻切分每張 GPU 只持有N/sp_size個 token但持有全部 H 個注意力頭第一次 all-to-all 集合通信將分片維度從“序列”換到“注意力頭”。交換后每張 GPU 持有完整序列但只負責H/sp_size個頭注意力計算head-sharded 布局每張 GPU 對自己負責的頭部子集、在完整序列上本地計算注意力無需逐 token 的跨卡通信第二次 all-to-all換回 seq-sharded 布局剩余的前向?qū)覨FN、LayerNorm 等繼續(xù)在各自 token 分片上本地執(zhí)行。官方文檔中的布局示意2 GPU 示例GPU 0 GPU 1 ┌───────────────┐ ┌───────────────┐ forward │ tokens 0..N/2 │ │ tokens N/2..N │ ← 每卡持有半段序列 (seq-sharded) │ all H heads │ │ all H heads │ └───────┬───────┘ └───────┬───────┘ └───────── all-to-all ──────┘ ┌───────────────┐ ┌───────────────┐ attention │ all N tokens │ │ all N tokens │ ← 此時每卡擁有完整序列 (head-sharded)│ heads 0..H/2 │ │ heads H/2..H │ ← 但只有半個頭的注意力 └───────┬───────┘ └───────┬───────┘ └───────── all-to-all ──────┘ ┌───────────────┐ ┌───────────────┐ forward │ tokens 0..N/2 │ │ tokens N/2..N │ ← 回到 seq-sharded (seq-sharded) │ all H heads │ │ all H heads │ └───────────────┘ └──────────────┘需要明確的邊界Transformers 目前集成的是完整 ALSTArctic Long Sequence Training系統(tǒng)中的Ulysses 序列并行組件。ALST 全家桶還包括 TiledMLP 與激活檢查點 offloading這兩部分不在Transformers 中提供完整系統(tǒng)請參考 DeepSpeed 官方教程。前置條件與配置方式序列并行的硬性要求Accelerate ≥ v1.12.0至少 2 張 GPU使用Trainer體系并經(jīng)由accelerate launch啟動。SP 配置承載在 Accelerate 的ParallelismConfig中有兩種等價的傳入途徑直接通過TrainingArguments的parallelism_config字段training_args.py 中已聲明parallelism_config: ParallelismConfig | None字段或?qū)戇M Accelerate 的config_file。方式一代碼中直接傳入 ParallelismConfigfrom accelerate.utils import ParallelismConfig, DeepSpeedSequenceParallelConfig parallelism_config ParallelismConfig( sp_backenddeepspeed, sp_size4, dp_replicate_size1, sp_handlerDeepSpeedSequenceParallelConfig( sp_seq_length_is_variableTrue, sp_attn_implementationflash_attention_2, ), ) training_args TrainingArguments( ..., deepspeedpath/to/deepspeed_config.json, parallelism_configparallelism_config, )然后用accelerate launch啟動基于Trainer的腳本accelerate launch --num_processes 4 train.py \ --output_dir output_dir \ --per_device_train_batch_size 1 \ --gradient_accumulation_steps 1方式二Accelerate 配置文件運行accelerate config命令按提示回答硬件與訓練拓撲問題后生成default_config.yaml位于緩存目錄。關(guān)鍵內(nèi)容如下distributed_type: DEEPSPEED deepspeed_config: deepspeed_config_file: path/to/ds_config.json machine_rank: 0 num_machines: 1 num_processes: 4 parallelism_config: parallelism_config_sp_size: 4 parallelism_config_dp_replicate_size: 1 parallelism_config_sp_backend: deepspeed parallelism_config_sp_seq_length_is_variable: true parallelism_config_sp_attn_implementation: flash_attention_2啟動命令accelerate launch --config_file alst_config.yaml train.py \ --output_dir output_dir \ --per_device_train_batch_size 1 \ --gradient_accumulation_steps 1倉庫測試目錄中同樣存在一份最小可參考的 SP 配置 deepspeed_zero2_sp.yaml展示了 ZeRO-2 2 卡 SP 的組合sp_size: 2、sp_backend: deepspeed、sp_seq_length_is_variable: true、sp_attn_implementation: sdpa可作為配置文件格式的對照樣本。關(guān)鍵參數(shù)逐項解析sp_backend必須為 deepspeed只有sp_backenddeepspeed才啟用 Ulysses 序列并行。Trainer源碼中多處以此作為開關(guān)例如在 trainer.py 的初始化階段pc getattr(self.accelerator, parallelism_config, None) if pc is not None and pc.sp_backend deepspeed and pc.sp_enabled: self.model_accepts_loss_kwargs False注釋寫明原因“Sequence Parallelism computes its own good_tokens count”——SP 模式下 loss 需要按各 SP rank 的有效 token 數(shù)加權(quán)聚合不能走num_items_in_batch的常規(guī)歸一化路徑。sp_size單條序列被多少張 GPU 并行處理sp_size是處理同一條序列的 GPU 數(shù)量。它與張量并行有一個本質(zhì)區(qū)別在 SP 中每個 SP rank 從 DataLoader 接收的是互不相同的數(shù)據(jù)流各自拿到序列的一段而張量并行中所有 rank 收到的是完全相同的數(shù)據(jù)。由此帶來兩個實操約束有效數(shù)據(jù)并行規(guī)模隨之縮小。有效dp_world_size world_size / sp_size。4 張 GPU 配sp_size4時dp_world_size1總批量計算中數(shù)據(jù)并行不再貢獻額外樣本。這一點可以從 trainer.py 的 get_total_train_batch_size 得到印證其公式注釋為dp_world_size world_size // (tp_size * cp_size * sp_size)序列必須填充為 sp_size 的整數(shù)倍。需要在數(shù)據(jù) collator 中設(shè)置pad_to_multiple_offrom transformers import DataCollatorForLanguageModeling data_collator DataCollatorForLanguageModeling( tokenizertokenizer, mlmFalse, pad_to_multiple_ofsp_size, )另外注意力頭數(shù)必須能被sp_size整除32 頭的模型支持sp_size取 1、2、4、8、16 或 32。原因是 all-to-all 后每張卡要拿到整數(shù)個頭的注意力子集無法切半頭。sp_seq_length_is_variable變長序列處理設(shè)為True推薦各 batch 之間序列長度可以不同逐批動態(tài)處理設(shè)為False要求所有序列都 padding 到由sp_seq_length指定的固定長度。訓練語料長度波動較大時保持True避免無謂的 padding 開銷。sp_attn_implementation注意力后端可選值為sdpa、flash_attention_2、flash_attention_3。選擇建議優(yōu)先 FlashAttention尤其是 batch 中 packing 了多條樣本時。打包packing場景下SDPA 可能錯誤地跨樣本邊界計算注意力——因為打包樣本之間缺乏顯式的塊狀注意力掩碼隔離不支持 Eager 注意力其 4Dattention_mask出于顯存與可擴展性考慮會被直接丟棄導致掩碼失效、注意力計算錯誤。Trainer 自動處理的隱藏工作官方文檔提示Trainer會自動處理 SP 下的 DataLoader 分片、position_ids生成、label 移位與跨 SP rank 的 loss 聚合。從源碼看這些并非空話而是有明確落點1. DataLoader 適配。在模型 prepare 之后trainer.py 會補一次 Ulysses 專用 DataLoader 包裝# since DataLoader was Accelerate prepared w/o a model arg in the same call, we now # have to complete the DL wrapping for ALST/UlyssesSP, after model has been prepared pc getattr(self.accelerator, parallelism_config, None) if pc is not None and pc.sp_backend deepspeed and pc.sp_enabled: train_dataloader self.accelerator.deepspeed_ulysses_dl_adapter(train_dataloader, model)2. 步數(shù)計算修正。SP DataLoader 適配器會讓每個 rank 只迭代1/sp_size的 batch 數(shù)compute_steps 邏輯 因此把 dataloader 長度乘回sp_size保證num_update_steps_per_epoch、max_steps等指標與不開 SP 時語義一致。3. 跨 rank 的加權(quán) loss 聚合。訓練時的compute_loss會分流到 deepspeed_sp_compute_loss由 trainer.py 在sp_backend deepspeed且sp_size 1的訓練態(tài)調(diào)用。其核心邏輯利用 DeepSpeed SP 注入的shift_labels預(yù)先移位好的標簽若 DataLoader 已移除labels鍵則從shift_labels回填保證模型 forward 能計算 loss通過 DeepSpeed 的 SP 進程組groups._get_sequence_parallel_group()對每個 rank 的 loss 與有效 token 數(shù)shift_labels ! -100的計數(shù)即good_tokens做all_gather按各 rank 有效 token 數(shù)加權(quán)求和后歸一loss Σ(rank_loss × good_tokens) / Σ(good_tokens)并屏蔽good_tokens 0的 rank例如該分片全是 prompt token 被 -100 掩掉的情況SFT 場景常見。這套機制直接支撐了 SFT 等“prompt 不計入 loss”的場景在 SP 下的正確性——每個 rank 的有效 token 數(shù)可能完全不同簡單平均會產(chǎn)生偏差加權(quán)聚合保證了數(shù)學上等價于全局 batch 平均。與數(shù)據(jù)并行組合使用SP 與 DP 共享同一批 GPU不需要額外硬件。組合規(guī)則是一個乘法等式dp_replicate_size × dp_shard_size × sp_size 總 GPU 數(shù)例如 8 張 GPU、sp_size4時設(shè)dp_replicate_size22 × 1 × 4 8parallelism_config ParallelismConfig( sp_backenddeepspeed, sp_size4, dp_replicate_size2, sp_handlerDeepSpeedSequenceParallelConfig( sp_seq_length_is_variableTrue, sp_attn_implementationflash_attention_2, ), )這樣每 4 張卡協(xié)同處理一條序列同時保留 2 路數(shù)據(jù)并行來擴大有效批量。進一步閱讀Accelerate 集成文檔講解 Accelerate 與 Trainer 的對接包括parallelism_config的完整用法與config_file加載流程多 GPU 訓練并行方法介紹如何將序列并行與 ZeRO 等策略組合Trainer 分布式測試配置倉庫內(nèi)真實使用過的 ZeRO-2 SP 啟動配置樣例可用于驗證自己的config_file寫法。如果你編寫的是自定義訓練循環(huán)而非Trainer上述自動化的分片、position_ids、label 移位與 loss 聚合就需要自己實現(xiàn)建議直接參考 Accelerate 官方的 Sequence Parallelism 概念指南而Trainer用戶只需按本文配置即可開箱使用?!久赓M下載鏈接】transformers Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.項目地址: https://gitcode.com/GitHub_Trending/tra/transformers創(chuàng)作聲明:本文部分內(nèi)容由AI輔助生成(AIGC),僅供參考