DeepSpeed v0.17.4发布：优化序列并行与内存效率的关键更新

DeepSpeed团队近日发布了v0.17.4版本，这是继v0.17.3之后的一个重要补丁更新。本次更新主要围绕序列并行训练和内存效率优化展开，引入了多项关键改进，包括修复了维度变量错误、新增分片融合对数损失函数以及相关bug修复。本文将详细解析这些更新的技术细节、实现原理以及对大规模语言模型训练的实际影响。

DeepSpeed v0.17.4版本包含以下主要变更：

1. 版本号更新：从v0.17.3升级到v0.17.4

2. 错误修复：解决了'dim'变量的UnboundLocalError问题

3. 新功能引入：添加了TiledFusedLogitsLoss功能

4. Bug修复：修正了TiledFusedLogitsLoss中的问题

5. 代码变更：共4个提交，修改了4个文件，涉及5位贡献者

这些更新主要集中在序列并行训练领域，特别是针对超长序列训练场景的优化，相关技术源自论文《Arctic Long Sequence Training: Scalable And Efficient Training For Multi-Million Token Sequences》。

DeepSpeed在序列并行方面的核心组件包括：

• UlyssesSPDataLoaderAdapter：数据加载器适配器，用于将普通数据批次分片以供UlyssesSPAttentionHF使用

• SequenceTiledCompute：通用自动微分函数，用于在序列维度分片后执行计算

• TiledMLP：特定的自动微分函数，用于执行分片MLP计算

• TiledFusedLogitsLoss：新增的自动微分函数，无需完整生成对数张量即可计算损失

在v0.17.4中，这些组件得到了进一步优化，特别是在梯度计算和内存管理方面。

TiledFusedLogitsLoss是本次更新的核心功能之一，它通过分片计算损失函数，避免了生成完整的对数(logits)张量，从而显著降低了内存使用。其主要特点包括：

1. 内存效率：不生成完整的对数张量，而是分片计算损失

2. 灵活配置：支持"mean"和"sum"两种输出缩减方式

3. 自动微分支持：完整集成到PyTorch自动微分系统中

4. ZeRO兼容：专门设计支持DeepSpeed的ZeRO优化器

该功能的典型使用场景如下：

def loss_fn(self, x, y): logits = self.lm_head(x) return self.cross_entropy_loss(logits.view(-1, self.vocab_size), y.view(-1))loss = TiledFusedLogitsLoss.apply( loss_fn, self, x, y, mask, shards, compute_params, output_reduction,)

TiledFusedLogitsLoss的实现基于PyTorch的autograd.Function，其核心思想是将输入张量在序列维度上进行分片，然后分别计算每个分片的损失，最后合并结果。关键技术点包括：

1. 输入验证：确保输入张量的维度、形状和掩码(如果存在)符合要求

2. 分片处理：使用torch.chunk将输入张量分割为多个分片

3. 梯度计算：通过调整传入梯度来处理输出缩减的影响

4. 内存优化：避免完整对数张量的实例化，直接在分片上计算损失

特别值得注意的是，由于该自动微分函数通常位于调用堆栈的最后，它在forward方法内部执行backward，并人工补偿output_reduction的影响，这消除了在backward中重新运行forward的需要。

本次更新修复了关于变量'dim'的UnboundLocalError问题。该问题出现在反向传播过程中，当某些输入不需要梯度时，'dim'变量可能未被定义。修复方案是在使用前明确定义'dim'变量：

dim = grad_output.dimif ctx.needs_input_grad[0]: grad_input = grad_output.matmul(weight)if ctx.needs_input_grad[1]: if dim > 2: grad_weight = grad_output.reshape(-1, grad_output.shape[-1]).t.matmul(input.reshape(-1, input.shape[-1]))

在初步实现中发现并修复了TiledFusedLogitsLoss的几个问题，包括：

1. 输出缩减处理：确保在"mean"缩减模式下正确调整梯度

2. 梯度计算：修复了分片间梯度累积的问题

3. 内存布局：解决了某些情况下的张量步幅(stride)问题

DeepSpeed团队为这些更新添加了全面的单元测试，验证内容包括：

1. 数值正确性：确保分片计算结果与未分片版本一致

2. 梯度准确性：验证参数梯度和输入梯度的正确性

3. 内存节省：确认内存使用量的降低

4. ZeRO兼容性：测试与ZeRO各阶段的兼容情况

测试用例覆盖了不同批次大小(batch size)和ZeRO阶段(1和3)的组合，包括边界情况如非2^n长度的序列。

在实际应用中，这些更新带来了以下优势：

1. 内存效率提升：对于大词汇表模型，可显著减少峰值内存使用

2. 长序列支持：使训练超长序列(百万token级别)更加可行

3. 计算效率：分片计算可以更好地利用现代GPU的并行能力

测试数据显示，在保持数值精度的前提下，内存使用量可降低30%-50%(取决于模型配置和序列长度)。

对于现有DeepSpeed用户，建议通过以下方式升级：

pip install deepspeed==0.17.4

1. 分片数量选择：根据GPU内存和序列长度合理选择分片数量

2. 缩减模式选择：对于不均匀分片(序列长度不能被分片数整除)，建议使用"sum"模式

3. 混合精度训练：与FP16/BF16配合使用时需注意梯度缩放

4. 调试技巧：可以先使用全精度(torch.float)验证数值正确性，再切换到混合精度

1. 序列长度限制：虽然支持长序列，但极端长度(如超过100万token)可能需要额外优化

2. 批处理维度：当前实现假设批处理维度是连续的，某些特殊数据布局可能需要调整

3. 分布式训练：完全支持DDP/FSDP需要额外配置

六、未来展望

基于v0.17.4的更新，DeepSpeed在序列并行方面的路线图可能包括：

1. 更灵活的分片策略：支持非均匀分片和动态分片

2. 更多算子优化：将分片计算模式扩展到其他常见算子

3. 自动分片配置：根据硬件特性自动选择最优分片参数

4. 更紧密的PyTorch集成：利用PyTorch的新特性进一步优化性能

DeepSpeed v0.17.4虽然是一个小版本更新，但在序列并行和内存效率方面带来了重要改进。特别是TiledFusedLogitsLoss的引入，为训练大规模语言模型提供了新的内存优化手段。这些更新使得在有限硬件资源下训练更长序列的模型成为可能，为自然语言处理和多模态研究提供了有力支持。

·

我们相信人工智能为普通人提供了一种“增强工具”，并致力于分享全方位的AI知识。在这里，您可以找到最新的AI科普文章、工具评测、提升效率的秘籍以及行业洞察。

·