企业官网建设流程全解析

• 示例工程

本文以 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、3455，PyTorch Lightning issue 16914，HuggingFace accelerate issue 2647 等），且已在 NVIDIA Triton（Dynamic Batcher）与 Meta FairSeq 等库中落地。DeepSpeed 从0.16.5 版本开始正式支持动态批处理。

该特性的最直接价值是最大化 GPU 利用率。更重要的是，它与**课程学习（Curriculum Learning）**高度契合：理想的BxSxE（Batch × Sequence Length × Embedding）形状应当是——早期课程步批量大、序列短（大量短句打包进一个 batch，高 B 低 S），后期则批量小、序列长（batch 里只有少数长句，低 B 高 S）。

在动态维度上，DeepSpeed 此前已支持动态的S，例如 Pipeline Engine 的reset_activation_shape()接口：课程学习改变每个样本的 seqlen 时，需要在 seqlen 即将变化时调用该接口。但动态的B此前并不支持——而B一旦动态变化，学习率也必须相应增减，这正是本文所讲特性的核心。

学习率缩放的两种主流算法（README 明确引用）：

1. 线性缩放规则（Linear Scaling Rule）："当 minibatch size 乘以 k 时，学习率乘以 k"，出自Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour（Goyal et al.）。
2. 平方根缩放（Square Root Scaling）："当 batch size 乘以 k 时，学习率乘以 √k，以保持梯度期望中的方差恒定"，出自One weird trick for parallelizing convolutional neural networks（A. Krizhevsky et al.）。

实际使用中，用户把每个 batch 的总 Token 数作为驱动打包的度量（而非句子条数）。运行时，系统计算出动态 batch size，再基于配置中给定的 LR 与 batch size 对学习率做相应调整。

动态批量、序列长度与学习率的联动示意

设每个 micro-batch 的 Token 预算为 30，即每次迭代每 GPU 的BxSxE数据会尽可能多地打包序列（增加B），使该迭代总 Token 数不超过 30。同时，在 DeepSpeed 配置中设置参考学习率lr=1e-3、参考train_batch_size=2。

课程学习的打包算法会在早期把数据打包成短句批次（左图：高 B 低 S），后期打包成长句批次（右图：低 B 高 S）。每一次迭代，学习率都会根据最终的有效 batch size（所有 GPU 上所有 micro-batch 的B之和，含梯度累积步）进行调整：

上图中，系统持续收集样本直到每个 micro-batch 填满 30 Token。左、右两侧最终的 micro-batch 大小分别为 10 和 4。按线性缩放规则，这两个 batch 的学习率分别变为5e-3与2e-3（lr × (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，注意力矩阵变为BxSxS（3D）。以 DDP 场景上图左上角的 micro-batch 1（4 个句子）为例，该 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接受两种输入形态：
  1. 用户提供的Optimizer：在每个 batch 直接缩放参数组（param groups）中的学习率；
  2. 用户定义的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把不同长度的序列拼成BxSxE（S 取该 batch 内最大 seqlen）。
• sample_padding_fn(sample, size)：把形状SxE的序列填充到S'xE，其中S'由参数size指定——这是required_microbatches_of_same_lengths落地时依赖的用户回调。
• batch_seqlens_fn(batch)：给定一个 batch，返回其中每条序列的真实长度（通过查找 padding 值位置实现）。

模型与训练循环的关键动作

模型AttentionHeadAndFeedForward在 forward 中计算每条序列的真实长度，构造BxSxS的 3D 因果掩码（torch.tril下三角加逐样本长度屏蔽），完成多头注意力后把BxSxE填充成BxMxE再过两层全连接。训练主流程演示了动态批处理的两个关键 API 行为：

1. LR 调度器复位：每个 epoch 开始时调用lr_scheduler.step(0)，把动态 LR 调度器（VariableBatchLR）复位到数据迭代器起点；
2. 每个 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, they'll 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, it's discarded. "min_batch_size": 1, # maximum number of sequences required to reach `max_tokens`. If sentence pack is larger, it's discarded. "max_batch_size": 10, # enable the output of microbatching information about sentence packing "verbose": True, }, }, }, }

配置参数速查表

关于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": "seqlen"（README 中写作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），系统会自动走度量分析流程。

运行方式与适用前提

1. 安装 DeepSpeed：需要 DeepSpeed >= 0.16.5 才支持动态批处理；
2. 运行示例：在 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）拉起多进程；
3. 接入自有训练脚本：按上文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 三者协同变化，在固定显存内换取更高的训练吞吐与更平滑的收敛。

• 示例工程