ARTICLE DETAIL

资讯详情

深耕商务建站与企业官网运营的一线实战洞察。

Transformers 中的 DeepSpeed Ulysses 序列并行:用多卡切分超长序列进行训练的技术指南

Transformers 中的 DeepSpeed Ulysses 序列并行:用多卡切分超长序列进行训练的技术指南 Transformers 中的 DeepSpeed Ulysses 序列并行用多卡切分超长序列进行训练的技术指南【免费下载链接】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及其配套源码系统讲解 Ulysses 序列并行Sequence Parallelism, SP在 TransformersTrainer中的工作机制、完整配置方式与关键参数含义并结合 Trainer 源码 与 DeepSpeed 集成模块 的调用链帮助你掌握在 2 卡及以上多 GPU 环境下训练超长序列百万 token 级上下文的完整实战方案。Ulysses 序列并行的核心原理超长序列训练的首要瓶颈是显存注意力计算的中间激活随序列长度平方增长单卡往往放不下一条长序列。Ulysses 序列并行通过“切分序列 两次 all-to-all 通信”的方式让多张 GPU 协同处理同一条序列前向开始时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 布局剩余的前向层FFN、LayerNorm 等继续在各自 token 分片上本地执行。官方文档中的布局示意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系统中的Ulysses 序列并行组件。ALST 全家桶还包括 TiledMLP 与激活检查点 offloading这两部分不在Transformers 中提供完整系统请参考 DeepSpeed 官方教程。前置条件与配置方式序列并行的硬性要求Accelerate ≥ v1.12.0至少 2 张 GPU使用Trainer体系并经由accelerate launch启动。SP 配置承载在 Accelerate 的ParallelismConfig中有两种等价的传入途径直接通过TrainingArguments的parallelism_config字段training_args.py 中已声明parallelism_config: ParallelismConfig | None字段或写进 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位于缓存目录。关键内容如下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可作为配置文件格式的对照样本。关键参数逐项解析sp_backend必须为 deepspeed只有sp_backenddeepspeed才启用 Ulysses 序列并行。Trainer源码中多处以此作为开关例如在 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 数加权聚合不能走num_items_in_batch的常规归一化路径。sp_size单条序列被多少张 GPU 并行处理sp_size是处理同一条序列的 GPU 数量。它与张量并行有一个本质区别在 SP 中每个 SP rank 从 DataLoader 接收的是互不相同的数据流各自拿到序列的一段而张量并行中所有 rank 收到的是完全相同的数据。由此带来两个实操约束有效数据并行规模随之缩小。有效dp_world_size world_size / sp_size。4 张 GPU 配sp_size4时dp_world_size1总批量计算中数据并行不再贡献额外样本。这一点可以从 trainer.py 的 get_total_train_batch_size 得到印证其公式注释为dp_world_size world_size // (tp_size * cp_size * sp_size)序列必须填充为 sp_size 的整数倍。需要在数据 collator 中设置pad_to_multiple_offrom transformers import DataCollatorForLanguageModeling data_collator DataCollatorForLanguageModeling( tokenizertokenizer, mlmFalse, pad_to_multiple_ofsp_size, )另外注意力头数必须能被sp_size整除32 头的模型支持sp_size取 1、2、4、8、16 或 32。原因是 all-to-all 后每张卡要拿到整数个头的注意力子集无法切半头。sp_seq_length_is_variable变长序列处理设为True推荐各 batch 之间序列长度可以不同逐批动态处理设为False要求所有序列都 padding 到由sp_seq_length指定的固定长度。训练语料长度波动较大时保持True避免无谓的 padding 开销。sp_attn_implementation注意力后端可选值为sdpa、flash_attention_2、flash_attention_3。选择建议优先 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. 步数计算修正。SP DataLoader 适配器会让每个 rank 只迭代1/sp_size的 batch 数compute_steps 逻辑 因此把 dataloader 长度乘回sp_size保证num_update_steps_per_epoch、max_steps等指标与不开 SP 时语义一致。3. 跨 rank 的加权 loss 聚合。训练时的compute_loss会分流到 deepspeed_sp_compute_loss由 trainer.py 在sp_backend deepspeed且sp_size 1的训练态调用。其核心逻辑利用 DeepSpeed SP 注入的shift_labels预先移位好的标签若 DataLoader 已移除labels键则从shift_labels回填保证模型 forward 能计算 loss通过 DeepSpeed 的 SP 进程组groups._get_sequence_parallel_group()对每个 rank 的 loss 与有效 token 数shift_labels ! -100的计数即good_tokens做all_gather按各 rank 有效 token 数加权求和后归一loss Σ(rank_loss × good_tokens) / Σ(good_tokens)并屏蔽good_tokens 0的 rank例如该分片全是 prompt token 被 -100 掩掉的情况SFT 场景常见。这套机制直接支撑了 SFT 等“prompt 不计入 loss”的场景在 SP 下的正确性——每个 rank 的有效 token 数可能完全不同简单平均会产生偏差加权聚合保证了数学上等价于全局 batch 平均。与数据并行组合使用SP 与 DP 共享同一批 GPU不需要额外硬件。组合规则是一个乘法等式dp_replicate_size × dp_shard_size × sp_size 总 GPU 数例如 8 张 GPU、sp_size4时设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 张卡协同处理一条序列同时保留 2 路数据并行来扩大有效批量。进一步阅读Accelerate 集成文档讲解 Accelerate 与 Trainer 的对接包括parallelism_config的完整用法与config_file加载流程多 GPU 训练并行方法介绍如何将序列并行与 ZeRO 等策略组合Trainer 分布式测试配置仓库内真实使用过的 ZeRO-2 SP 启动配置样例可用于验证自己的config_file写法。如果你编写的是自定义训练循环而非Trainer上述自动化的分片、position_ids、label 移位与 loss 聚合就需要自己实现建议直接参考 Accelerate 官方的 Sequence Parallelism 概念指南而Trainer用户只需按本文配置即可开箱使用。【免费下载链接】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创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表
PREV
查看更多资讯
NEXT
返回资讯列表