示例工程【免费下载链接】DeepSpeedExamplesExample models using DeepSpeed项目地址https://gitcode.com/gh_mirrors/de/DeepSpeedExamples点击查看免费下载本文以 DeepSpeedExamples 仓库中 training/data_efficiency/variable_batch_size_and_lr/README.md 为骨架结合仓库内的可运行示例 variable_batch_size_and_lr_example.py 与配套插图完整讲解 DeepSpeed 动态批处理Dynamic Batching与学习率缩放LR Scaling的原理、配置与落地代码。读完本文你将掌握如何让训练按总 Token 数而非句子条数打包 batch如何让 batch 维度 B 在训练过程中动态伸缩如何基于全局有效批量大小自动缩放学习率以及如何在 Pipeline Parallelism 下满足微批次同形约束。背景与动机为什么需要按 Token 数打包批次在大语言模型LLM等场景中输入通常是长度不等的句子。一种常见做法是按 Token 数打包token-count based batching而不是固定 batch size即把若干句子拼在一起使它们的某个度量如序列长度之和恰好达到用户给定的目标值。这一思想早在 Transformer 原始论文Attention Is All You Need第 5.1 节中就有体现——训练时句子对按近似序列长度打包每个 batch 约包含 25,000 个源 Token 与 25,000 个目标 Token。动态批量大小这一需求在社区中反复出现如 DeepSpeed issue 1051、3455PyTorch Lightning issue 16914HuggingFace accelerate issue 2647 等且已在 NVIDIA TritonDynamic Batcher与 Meta FairSeq 等库中落地。DeepSpeed 从0.16.5 版本开始正式支持动态批处理。该特性的最直接价值是最大化 GPU 利用率。更重要的是它与**课程学习Curriculum Learning**高度契合理想的BxSxEBatch × Sequence Length × Embedding形状应当是——早期课程步批量大、序列短大量短句打包进一个 batch高 B 低 S后期则批量小、序列长batch 里只有少数长句低 B 高 S。在动态维度上DeepSpeed 此前已支持动态的S例如 Pipeline Engine 的reset_activation_shape()接口课程学习改变每个样本的 seqlen 时需要在 seqlen 即将变化时调用该接口。但动态的B此前并不支持——而B一旦动态变化学习率也必须相应增减这正是本文所讲特性的核心。学习率缩放的两种主流算法README 明确引用线性缩放规则Linear Scaling Rule当 minibatch size 乘以 k 时学习率乘以 k出自Accurate, Large Minibatch SGD: Training ImageNet in 1 HourGoyal et al.。平方根缩放Square Root Scaling当 batch size 乘以 k 时学习率乘以 √k以保持梯度期望中的方差恒定出自One weird trick for parallelizing convolutional neural networksA. Krizhevsky et al.。实际使用中用户把每个 batch 的总 Token 数作为驱动打包的度量而非句子条数。运行时系统计算出动态 batch size再基于配置中给定的 LR 与 batch size 对学习率做相应调整。动态批量、序列长度与学习率的联动示意设每个 micro-batch 的 Token 预算为 30即每次迭代每 GPU 的BxSxE数据会尽可能多地打包序列增加B使该迭代总 Token 数不超过 30。同时在 DeepSpeed 配置中设置参考学习率lr1e-3、参考train_batch_size2。课程学习的打包算法会在早期把数据打包成短句批次左图高 B 低 S后期打包成长句批次右图低 B 高 S。每一次迭代学习率都会根据最终的有效 batch size所有 GPU 上所有 micro-batch 的B之和含梯度累积步进行调整上图中系统持续收集样本直到每个 micro-batch 填满 30 Token。左、右两侧最终的 micro-batch 大小分别为 10 和 4。按线性缩放规则这两个 batch 的学习率分别变为5e-3与2e-3lr × (B/参考B)即1e-3 × 10/2 5e-3、1e-3 × 4/2 2e-3。为便于说明上图只画了 1 张 GPU因此每次迭代的学习率只由单 GPU 的 micro-batch 大小决定多 GPU 的情况在下一节覆盖。注意这里调整学习率用的是全局 batch size而非单 GPU 的 micro-batch 大小详见下文 Pipeline 小节。Pipeline Parallelism 下的微批次约束流水线并行要求一个 batch 内所有 micro-batch 的 micro-batch 大小与序列长度完全相同因为梯度累积步之间的 activation 形状必须固定。而在 batch 之间这些形状可以变化前提是调用engine.reset_activation_shape()让新的形状在 batch 的第一个梯度累积步上被通信同步。在 batch 之间强制相似的BxSxE形状可能导致 micro-batch 变得更小。下图为同一数据集分别面向常规 DDP左与流水线并行右准备数据时的对比——4 个 micro-batch等价于一个 2 节点 × 2 梯度累积步的配置可以看出流水线场景右图中 4 个 micro-batch 的BxSxE形状全部一致但为了满足这一约束它相比常规场景左图打包了更少的样本并增加了 padding。这里有一条重要提醒学习率的调整现在要考虑的是 16左与 12右的样本总数——即LR 依据全局 batch size 缩放而非依据单 GPU 的 micro-batch 大小。这一口径在示例代码中体现为train_batch_size配置全局为 16缩放因子基于所有 GPU 与梯度累积步合并后的有效批量。注意力矩阵的内存分析BxSxS对形状为BxSxE的输入当所有样本长度相同时注意力 mask 形状为SxS当样本长度不同如上文数据集所示时每个样本需要独立的 mask注意力矩阵变为BxSxS3D。以 DDP 场景上图左上角的 micro-batch 14 个句子为例该 3D 注意力矩阵可表示为注意其中的内存收益注意力头的尺寸为BxSxS即对 batch 维度B是线性依赖对 (micro-)batch 中最大序列长度S是二次方依赖。因此支持动态S把 padding 造成的浪费降到最低就能为B腾出空间从而在同样内存预算下塞进更多样本。这也正是动态批处理能提升吞吐量的底层原因——它同时压缩了S的填充浪费并允许B按 Token 预算弹性伸缩。核心实现剖析get_dataloader_and_lr_scheduler_for_variable_batch_size该 PR 实现了动态批处理与 LR 缩放所需的数据加载器与 LR 调度器通过函数get_dataloader_and_lr_scheduler_for_variable_batch_size示例中具体名称为get_dataloader_and_lr_scheduler_for_variable_batch_size_deepspeed从deepspeed.runtime.data_pipeline.data_sampling.variable_batch_size_and_lr导入一次性获取。该函数的核心分工如下scale_lr承载 LR 缩放算法的全部逻辑线性、平方根等规则在此实现。batch_by_seqlen负责把样本划分成 batch。它会返回三组关键信息microbatch_sample_ids每个 micro-batch 的样本 id 列表batch_sizes每个有效 batch 的大小batch_max_seqlens一个 batch 内所有 micro-batch 中的最长序列。dataloader_for_variable_batch_size依赖microbatch_sample_ids对每个 batch 进行迭代/拼接/填充最终返回一个迭代变大小batch 的数据加载器。lr_scheduler_for_variable_batch_size依赖batch_sizes结合配置中的 batch size 与 LR按缩放规则Linear、Square Root 等为每个有效 batch 计算学习率。它返回的lr_scheduler接受两种输入形态用户提供的Optimizer在每个 batch 直接缩放参数组param groups中的学习率用户定义的LRScheduler先由该 scheduler 得到基础学习率再做相应缩放。流水线并行下的强制同形参数针对流水线并行要求一个 pipeline pass 中所有 micro-batch 激活形状相同的约束README 说明了两个可开启的开关required_microbatches_of_same_sizes强制所有数据加载器在一个 batch 内的所有梯度累积步共享相同的B维度。required_microbatches_of_same_lengths强制所有梯度累积步共享相同的S维度。其工作原理是调用用户提供的sample_padding_fn(sentence, len)把给定句子填充到指定长度。完整示例解读variable_batch_size_and_lr_example.py仓库在 variable_batch_size_and_lr_example.py 中提供了一个有/无流水线并行两种用法的端到端示例模型为一个变长BxSxS注意力头 定长前馈网络——这正是 LLM 的主要构成块。跟随注意力的前馈线性层要求输入尺寸恒定等价于整个数据集中最长句子的长度因此注意力输出必须填充到该长度对应代码注释feedforward: needs to convert BxSxE to BxMxE by padding extra tokens。数据集三要素collate / padding / seqlens 三个回调示例中的TestData数据集每条样本是随机长度序列标签即序列长度是接入动态批处理的标准模板它提供了三个必需的回调batch_collate_fn(batch)用nn.utils.rnn.pad_sequence把不同长度的序列拼成BxSxES 取该 batch 内最大 seqlen。sample_padding_fn(sample, size)把形状SxE的序列填充到SxE其中S由参数size指定——这是required_microbatches_of_same_lengths落地时依赖的用户回调。batch_seqlens_fn(batch)给定一个 batch返回其中每条序列的真实长度通过查找 padding 值位置实现。模型与训练循环的关键动作模型AttentionHeadAndFeedForward在 forward 中计算每条序列的真实长度构造BxSxS的 3D 因果掩码torch.tril下三角加逐样本长度屏蔽完成多头注意力后把BxSxE填充成BxMxE再过两层全连接。训练主流程演示了动态批处理的两个关键 API 行为LR 调度器复位每个 epoch 开始时调用lr_scheduler.step(0)把动态 LR 调度器VariableBatchLR复位到数据迭代器起点每个 batch 的形状重置与学习率获取流水线模式下--pipeline-num-stages 1模型先用PipelineModule包装、每个 batch 前调用engine.reset_activation_shape()因为每个 batch 的BxS都不同非流水线模式下在每个梯度累积步内用lr_scheduler.get_lr()读取按当前有效 batch 缩放后的学习率并在engine.step()中生效。示例还演示了如何用dataset_filter_ids模拟课程学习的一步只保留 seqlen 在 5~10 之间的样本子集这与 README 中课程学习早期高 B 低 S、后期低 B 高 S的动机相呼应。DeepSpeed 配置详解动态批处理通过data_efficiency配置节开启。README 给出的完整参考配置如下含逐项注释config { train_batch_size: 16, # train_micro_batch_size_per_gpu tells how many sequence packs of max_tokens each will be collated together. # I.e. the number of tokens per micro batch (ie per gpu iteration) is train_micro_batch_size_per_gpu*max_tokens. train_micro_batch_size_per_gpu: 2, data_efficiency: { enabled: True, # seed to be applied to all data efficiency modules, including dynamic batching seed: 42, data_sampling: { num_workers: 0, # dataloader num_workers argument pin_memory: False, # dataloader pin_memory argument dynamic_batching: { # enables or disables dynamic batching enabled: True, # how many tokens we need to fill a pack of sequences (that will be collated together as a sample) max_tokens: 100, # Input and output write to read from or write the length of every sequence. # Sequence lengths will be loaded from: {metrics_path}/seqlen/seqlen_sample_to_metric.bin and *.idx # If files dont exist, theyll be computed and saved on the first run, and loaded on subsequent runs. metrics_path: ./curriculum_output/, # As batch size increases/decreses, which method to use to scale LR accordingly? # Options: linear, sqrt (square root), or None to disable lr_scaling_method: linear, # how to pick sentences to be packed into samples: # - dataloader: by same order as they come in with the dataloader # - seqlen: by sequence length (shortest to longest) # - random: random order using the seed in config[data_efficiency][seed sentence_picking_order: dataloader, # random / seqlen / dataloader # minimum number of sequences required to reach max_tokens. If sentence pack is smaller, its discarded. min_batch_size: 1, # maximum number of sequences required to reach max_tokens. If sentence pack is larger, its discarded. max_batch_size: 10, # enable the output of microbatching information about sentence packing verbose: True, }, }, }, }配置参数速查表参数默认/示例值含义data_efficiency.enabledTrue开启数据效率模块含动态批处理data_efficiency.seed42应用于所有数据效率模块含动态批处理的随机种子data_sampling.num_workers0传给数据加载器的num_workers参数data_sampling.pin_memoryFalse传给数据加载器的pin_memory参数dynamic_batching.enabledTrue开关动态批处理dynamic_batching.max_tokens100填满一个序列包会被拼接成一个样本所需的 Token 数dynamic_batching.metrics_path./curriculum_output/序列长度的读写路径从{metrics_path}/seqlen/seqlen_sample_to_metric.bin与*.idx加载文件不存在则首次运行计算并保存、后续运行直接加载dynamic_batching.lr_scaling_methodlinear批量变大/变小时的学习率缩放方式可选linear、sqrt平方根或None禁用dynamic_batching.sentence_picking_orderdataloader样本打包顺序dataloader按数据加载器原始顺序/seqlen按序列长度短到长/random用data_efficiency.seed随机dynamic_batching.min_batch_size1达到max_tokens所需的最小序列数序列包小于此值则丢弃dynamic_batching.max_batch_size10达到max_tokens所需的最大序列数序列包大于此值则丢弃dynamic_batching.verboseTrue是否输出关于句子打包的微批次信息关于train_micro_batch_size_per_gpu与max_tokens的乘积关系README 特别注明train_micro_batch_size_per_gpu表示每个 micro-batch 拼接多少个max_tokens的序列包即每 GPU 每次迭代的 Token 数为train_micro_batch_size_per_gpu × max_tokens。示例代码中该项设为 2、max_tokens为 20而 README 的配置中该项同样为 2、max_tokens为 100——真实使用时按显存与数据规模调整即可。需要注意示例代码与 README 的一处细微差异示例配置中打包顺序字段写作sequence_picking_order: seqlenREADME 中写作sentence_picking_order默认dataloader且示例将verbose设为os.environ.get(RANK, 0) 0即只让 rank 0 输出 verbose 信息。这说明该特性仍在快速迭代中字段命名以当前 DeepSpeed 版本的实际实现为准。首次运行时的序列长度统计序列长度度量的存取是动态批处理的前置条件metrics_path下保存的seqlen_sample_to_metric.bin/*.idx记录了每条样本的长度首次运行由数据效率模块内的 DataAnalyzer 计算并落盘后续运行直接加载。这也意味着动态批处理天然依赖一个长度已知的数据管道示例中若dataset_seqlens不显式传入代码中注释if None: use DataAnalyzer to output seqlens and then load them系统会自动走度量分析流程。运行方式与适用前提安装 DeepSpeed需要 DeepSpeed 0.16.5 才支持动态批处理运行示例在 variable_batch_size_and_lr 目录下直接运行python variable_batch_size_and_lr_example.py默认--pipeline-num-stages 0即 DDP 模式需要验证流水线模式时传--pipeline-num-stages 2等大于 1 的数值此时示例会用PipelineModule包装模型并启用engine.reset_activation_shape()路径。示例要求每进程至少 1 张 GPU代码中有dist.get_local_rank() device_count()断言并需以 DeepSpeed 惯用的分布式启动方式如deepspeed/torchrun拉起多进程接入自有训练脚本按上文TestData的模板实现batch_collate_fn、sample_padding_fn、batch_seqlens_fn三个回调用deepspeed.initialize初始化引擎后把get_dataloader_and_lr_scheduler_for_variable_batch_size_deepspeed返回的dataloader与lr_scheduler挂到引擎上示例中分别以engine.lr_scheduler lr_scheduler与可选的engine.training_dataloader dataloader方式接入即可在训练循环中获得按 Token 预算打包的变大小 batch 与自动缩放的学习率。仓库中另有一个相关的学习率实现参考training/data_efficiency/gpt_finetuning/learning_rates.py源自 Megatron-DeepSpeed 的AnnealingLR展示了 Token 感知的 LR 衰减/warmup 写法可作为理解基于 Token 度量而非步数度量调度 LR这一思路的补充。动态批处理正是把这种Token 感知思想扩展到 batch 维度本身以 Token 预算为第一性约束让B、S、LR 三者协同变化在固定显存内换取更高的训练吞吐与更平滑的收敛。赞分享示例工程【免费下载链接】DeepSpeedExamplesExample models using DeepSpeed项目地址https://gitcode.com/gh_mirrors/de/DeepSpeedExamples点击查看免费下载相关推荐DeepSpeed 课程学习Curriculum Learning实战指南以序列长度驱动的稳定高效大规模 GPT 预训练DeepSpeed 课程学习Curriculum Learning实战指南以序列长度驱动的稳定高效大规模 GPT 预训练 本篇教程围绕 DeepSpeed人工智能大模型深度学习分布式训练预训练强化学习模型优化XTuner 序列并行实战指南基于 DeepSpeed Ulysses 思路训练百万级超长序列与百 B 级大模型XTuner 序列并行实战指南基于 DeepSpeed Ulysses 思路训练百万级超长序列与百 B 级大模型 导读 长文档摘要、视频生成等生成式 AI 任大模型模型微调动态学习率调度让DeepSpeed训练效率提升30%的实战指南动态学习率调度让DeepSpeed训练效率提升30%的实战指南 你还在为模型训练时学习率设置不当导致收敛慢或过拟合而烦恼吗本文将带你深入了解DeepSpee人工智能大模型深度学习分布式训练预训练强化学习模型优化上一篇Rendora故障排除手册常见问题与解决方案大全下一篇Signale标签样式设计日志信息层次化展示创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考