序列進(jìn)行訓(xùn)練的技術(shù)指南)
Transformers 中的 DeepSpeed Ulysses 序列并行用多卡切分超長(zhǎng)序列進(jìn)行訓(xùn)練的技術(shù)指南【免費(fèi)下載鏈接】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.項(xiàng)目地址: https://gitcode.com/GitHub_Trending/tra/transformers本篇基于 Transformers 官方文檔deepspeed_alst.md及其配套源碼系統(tǒng)講解 Ulysses 序列并行Sequence Parallelism, SP在 TransformersTrainer中的工作機(jī)制、完整配置方式與關(guān)鍵參數(shù)含義并結(jié)合 Trainer 源碼 與 DeepSpeed 集成模塊 的調(diào)用鏈幫助你掌握在 2 卡及以上多 GPU 環(huán)境下訓(xùn)練超長(zhǎng)序列百萬(wàn) token 級(jí)上下文的完整實(shí)戰(zhàn)方案。Ulysses 序列并行的核心原理超長(zhǎng)序列訓(xùn)練的首要瓶頸是顯存注意力計(jì)算的中間激活隨序列長(zhǎng)度平方增長(zhǎng)單卡往往放不下一條長(zhǎng)序列。Ulysses 序列并行通過(guò)“切分序列 兩次 all-to-all 通信”的方式讓多張 GPU 協(xié)同處理同一條序列前向開始時(shí)seq-sharded 布局整條長(zhǎng)度為 N 的序列被均勻切分每張 GPU 只持有N/sp_size個(gè) token但持有全部 H 個(gè)注意力頭第一次 all-to-all 集合通信將分片維度從“序列”換到“注意力頭”。交換后每張 GPU 持有完整序列但只負(fù)責(zé)H/sp_size個(gè)頭注意力計(jì)算head-sharded 布局每張 GPU 對(duì)自己負(fù)責(zé)的頭部子集、在完整序列上本地計(jì)算注意力無(wú)需逐 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 │ ← 此時(shí)每卡擁有完整序列 (head-sharded)│ heads 0..H/2 │ │ heads H/2..H │ ← 但只有半個(gè)頭的注意力 └───────┬───────┘ └───────┬───────┘ └───────── 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 與激活檢查點(diǎn) offloading這兩部分不在Transformers 中提供完整系統(tǒng)請(qǐng)參考 DeepSpeed 官方教程。前置條件與配置方式序列并行的硬性要求Accelerate ≥ v1.12.0至少 2 張 GPU使用Trainer體系并經(jīng)由accelerate launch啟動(dòng)。SP 配置承載在 Accelerate 的ParallelismConfig中有兩種等價(jià)的傳入途徑直接通過(guò)TrainingArguments的parallelism_config字段training_args.py 中已聲明parallelism_config: ParallelismConfig | None字段或?qū)戇M(jìn) 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啟動(dòng)基于Trainer的腳本accelerate launch --num_processes 4 train.py \ --output_dir output_dir \ --per_device_train_batch_size 1 \ --gradient_accumulation_steps 1方式二Accelerate 配置文件運(yùn)行accelerate config命令按提示回答硬件與訓(xùn)練拓?fù)鋯?wèn)題后生成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啟動(dòng)命令accelerate launch --config_file alst_config.yaml train.py \ --output_dir output_dir \ --per_device_train_batch_size 1 \ --gradient_accumulation_steps 1倉(cāng)庫(kù)測(cè)試目錄中同樣存在一份最小可參考的 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可作為配置文件格式的對(duì)照樣本。關(guān)鍵參數(shù)逐項(xiàng)解析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單條序列被多少?gòu)?GPU 并行處理sp_size是處理同一條序列的 GPU 數(shù)量。它與張量并行有一個(gè)本質(zhì)區(qū)別在 SP 中每個(gè) SP rank 從 DataLoader 接收的是互不相同的數(shù)據(jù)流各自拿到序列的一段而張量并行中所有 rank 收到的是完全相同的數(shù)據(jù)。由此帶來(lái)兩個(gè)實(shí)操約束有效數(shù)據(jù)并行規(guī)模隨之縮小。有效dp_world_size world_size / sp_size。4 張 GPU 配sp_size4時(shí)dp_world_size1總批量計(jì)算中數(shù)據(jù)并行不再貢獻(xiàn)額外樣本。這一點(diǎn)可以從 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ù)個(gè)頭的注意力子集無(wú)法切半頭。sp_seq_length_is_variable變長(zhǎng)序列處理設(shè)為True推薦各 batch 之間序列長(zhǎng)度可以不同逐批動(dòng)態(tài)處理設(shè)為False要求所有序列都 padding 到由sp_seq_length指定的固定長(zhǎng)度。訓(xùn)練語(yǔ)料長(zhǎng)度波動(dòng)較大時(shí)保持True避免無(wú)謂的 padding 開銷。sp_attn_implementation注意力后端可選值為sdpa、flash_attention_2、flash_attention_3。選擇建議優(yōu)先 FlashAttention尤其是 batch 中 packing 了多條樣本時(shí)。打包packing場(chǎng)景下SDPA 可能錯(cuò)誤地跨樣本邊界計(jì)算注意力——因?yàn)榇虬鼧颖局g缺乏顯式的塊狀注意力掩碼隔離不支持 Eager 注意力其 4Dattention_mask出于顯存與可擴(kuò)展性考慮會(huì)被直接丟棄導(dǎo)致掩碼失效、注意力計(jì)算錯(cuò)誤。Trainer 自動(dòng)處理的隱藏工作官方文檔提示Trainer會(huì)自動(dòng)處理 SP 下的 DataLoader 分片、position_ids生成、label 移位與跨 SP rank 的 loss 聚合。從源碼看這些并非空話而是有明確落點(diǎn)1. DataLoader 適配。在模型 prepare 之后trainer.py 會(huì)補(bǔ)一次 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ù)計(jì)算修正。SP DataLoader 適配器會(huì)讓每個(gè) rank 只迭代1/sp_size的 batch 數(shù)compute_steps 邏輯 因此把 dataloader 長(zhǎng)度乘回sp_size保證num_update_steps_per_epoch、max_steps等指標(biāo)與不開 SP 時(shí)語(yǔ)義一致。3. 跨 rank 的加權(quán) loss 聚合。訓(xùn)練時(shí)的compute_loss會(huì)分流到 deepspeed_sp_compute_loss由 trainer.py 在sp_backend deepspeed且sp_size 1的訓(xùn)練態(tài)調(diào)用。其核心邏輯利用 DeepSpeed SP 注入的shift_labels預(yù)先移位好的標(biāo)簽若 DataLoader 已移除labels鍵則從shift_labels回填保證模型 forward 能計(jì)算 loss通過(guò) DeepSpeed 的 SP 進(jìn)程組groups._get_sequence_parallel_group()對(duì)每個(gè) rank 的 loss 與有效 token 數(shù)shift_labels ! -100的計(jì)數(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 場(chǎng)景常見。這套機(jī)制直接支撐了 SFT 等“prompt 不計(jì)入 loss”的場(chǎng)景在 SP 下的正確性——每個(gè) rank 的有效 token 數(shù)可能完全不同簡(jiǎn)單平均會(huì)產(chǎn)生偏差加權(quán)聚合保證了數(shù)學(xué)上等價(jià)于全局 batch 平均。與數(shù)據(jù)并行組合使用SP 與 DP 共享同一批 GPU不需要額外硬件。組合規(guī)則是一個(gè)乘法等式dp_replicate_size × dp_shard_size × sp_size 總 GPU 數(shù)例如 8 張 GPU、sp_size4時(shí)設(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é)同處理一條序列同時(shí)保留 2 路數(shù)據(jù)并行來(lái)擴(kuò)大有效批量。進(jìn)一步閱讀Accelerate 集成文檔講解 Accelerate 與 Trainer 的對(duì)接包括parallelism_config的完整用法與config_file加載流程多 GPU 訓(xùn)練并行方法介紹如何將序列并行與 ZeRO 等策略組合Trainer 分布式測(cè)試配置倉(cāng)庫(kù)內(nèi)真實(shí)使用過(guò)的 ZeRO-2 SP 啟動(dòng)配置樣例可用于驗(yàn)證自己的config_file寫法。如果你編寫的是自定義訓(xùn)練循環(huán)而非Trainer上述自動(dòng)化的分片、position_ids、label 移位與 loss 聚合就需要自己實(shí)現(xiàn)建議直接參考 Accelerate 官方的 Sequence Parallelism 概念指南而Trainer用戶只需按本文配置即可開箱使用?!久赓M(fèi)下載鏈接】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.項(xiàng)目地址: https://gitcode.com/GitHub_Trending/tra/transformers創(chuàng)作聲明:本文部分內(nèi)容由AI輔助生成(AIGC),僅供參考