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

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

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

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

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

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

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

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

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

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

　　我们相信人工智能为普通人提供了一种“增强工具”，并致力于分享全方位的AI知识。在这里，您可以找到最新的AI科普文章、工具评测、提升效率的秘籍以及行业洞察。 欢迎关注“福大大架构师每日一题”，让AI助力您的未来发展。