CANN / torchtitan-npu：SwapOptimizer 特性

SwapOptimizer 特性【免费下载链接】torchtitan-npuAscend Extension for torchtitan项目地址: https://gitcode.com/cann/torchtitan-npu在分布式训练任务中随着模型参数量的剧增除了模型权重和激活值优化器状态如 Adam/AdamW 的动量和方差通常会占用极大的显存空间。现有的常规训练方案在处理优化器状态时暴露出以下局限性优化器状态在整个前向和反向传播过程中静态驻留在设备显存中但实际上仅在最后的参数更新阶段才会被使用造成了计算过程中的显存严重浪费。传统的单纯卸载Offload方案通常会引入巨大的 CPU 与加速卡之间的传输延迟阻塞主计算流导致训练吞吐量大幅下降。实现原理参考MindSpeed框架设计的SwapOptimizer特性我们在torchtitan框架的优化器构建逻辑中引入了自定义的SwapOptimizersContainer对原生的Adam和AdamW优化器的step方法进行了无缝拦截与替换从而使能了优化器状态的动态显存管理。相关代码定义在torchtitan_npu/patches/optimizer/swap_optimizer.py。本特性主要遵循“按块加载 - 异步更新 - 及时卸载”的流水线机制在确保训练精度的前提下用多流通算重叠换取巨大的显存空间收益。具体拆解如下状态初始化与零显存占位在初始化阶段系统会在host内存中为每个参数分配固定内存pin_memoryTrue以存放一阶动量exp_avg、二阶动量exp_avg_sq等状态。而在device端通过底层接口untyped_storage().resize_(0)巧妙地将这些状态的底层物理显存清空仅保留 Tensor 的元数据。这使得在耗时较长的前向和反向传播期间优化器状态完全不占用宝贵的设备显存。异步多流通信与流水线更新重写后的swap_optimizer_step引入了两个独立的底层数据流swap_to_device_stream和swap_to_host_stream并且将整个参数更新过程切分为多个小块通过swap_optimizer_times控制单次处理的容量阈值swap_numel以实现高效的流水线更新Load加载在设备流中异步将下一块所需参数的优化器状态从 CPU 拷贝到 Device并恢复其显存大小。Update更新主计算流通过事件Event阻塞等待状态加载完成随后立即调用底层 Fused 算子进行参数更新。Offload卸载更新完成后记录事件并在主机流中将更新后的最新状态异步写回 CPU同时再次清空设备侧显存以复用给下一个参数块。配置选项在训练任务的 TOML 配置文件例如torchtitan_npu/models/deepseek_v32/train_configs/deepseek_v32_671b_debug.toml或实际启动训练时--job.config_file所指向的路径中找到对应的[optimizer]节并添加以下配置以启用 Swap Optimizer配置项类型默认值说明namestrAdamW使用的优化器类型。当前 Swap 特性拦截并支持Adam和AdamW。swap_optimizerboolfalse是否启用 Swap Optimizer 特性以进行显存流水线卸载。swap_optimizer_timesint16切片划分次数。用于决定优化器参数卸载与加载时的分块大小。值越大单次峰值显存占用越小但可能增加系统的流调度开销。配置示例首先在配置文件中使能本代码仓的自定义配置随后在[optimizer]节中添加以下配置为AdamW优化器开启SwapOptimizer特性并设置流水线切片参数[job] custom_config_module torchtitan_npu.config.custom_config # 使能本代码仓的自定义配置 [optimizer] name AdamW lr 3e-4 weight_decay 0.01 swap_optimizer true # 启用 Swap Optimizer swap_optimizer_times 16 # 设置流水切分次数【免费下载链接】torchtitan-npuAscend Extension for torchtitan项目地址: https://gitcode.com/cann/torchtitan-npu创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考