1. 概述

1.1 简介

本提案为 Pipeline Model Parallel Layout 和 Swap-Muon checkpoint 补充单元测试，并优化现有 Swap Optimizer UT 的执行效率。

提案同时修复 Swap-Muon 加载 checkpoint 后 optimizer state 缓存未刷新的问题，保证参数与 momentum 正确恢复。

1.2 动机

当前存在以下问题：

• Pipeline Layout 包含字符串解析、PP/VPP 映射、MTP 布局及参数冲突校验，但缺少专门 UT。
• Swap-Muon 缺少 checkpoint 保存、破坏、恢复的完整闭环测试。
• Optimizer.load_state_dict() 会替换 state 字典，而 _state_map 仍可能指向旧 state，导致 momentum 恢复异常。
• 现有 Swap Optimizer 默认 UT 使用 8 张 NPU、8 层模型和 10 个训练 step，CI 执行成本较高。

若不补充测试，上述功能发生回归时可能只能在训练或 checkpoint 恢复阶段发现，定位和修复成本较高。

1.3 目标

目标：

• 覆盖 Pipeline Layout 的主要正常和异常场景。
• 看护 Swap-Muon checkpoint 参数及 momentum 恢复。
• 修复加载 checkpoint 后 state 映射失效的问题。
• 将默认 Swap Optimizer UT 降低为 2 张 NPU 可执行。
• 保留 8 卡完整拓扑测试能力。

非目标：

• 不增加新的用户配置或公开 API。
• 不修改 Pipeline Layout 的功能语义。
• 不改变 Swap-Muon checkpoint 的使用流程。
• 不在本提案中优化 Swap-Muon 本身的训练性能。

2. 用例分析

主要包含以下用例：

1. Pipeline Layout 解析与校验：
   • 字符串乘法及空 stage。
   • PP/VPP stage 映射。
   • Layer offset 和 layer ID 计算。
   • MTP standalone 布局。
   • 非法字符、层数错误及参数冲突。
2. Swap-Muon checkpoint 恢复：
   • 执行一次优化器更新。
   • 保存 optimizer checkpoint。
   • 主动破坏 CPU mirror 中的参数和 momentum。
   • 加载 checkpoint 并验证恢复结果。
   • 验证恢复后 NPU storage 重新释放。
3. UT 执行效率：
   • 默认 CI 仅使用 2 张 NPU 和小模型。
   • 8 卡 TP/PP 及通信 overlap 组合通过 --run-all 执行。

DFX 要求：

• 测试之间不得残留全局 optimizer 状态。
• 每个 rank 使用独立 checkpoint 临时文件。
• CPU 测试不得初始化 NPU 分布式环境。
• 完整拓扑覆盖不能因默认 UT 提效而被删除。

3. 方案设计

3.1 总体方案

Pipeline Layout 采用单个参数化测试函数覆盖 10 个场景，通过 Mock parallel state 验证 PP/VPP 映射，不依赖 NPU。

Swap-Muon 采用实际生产路径完成 checkpoint 闭环：

初始化 Swap-Muon
    ↓
执行一次 optimizer step
    ↓
save_state_dict_to_file
    ↓
修改 CPU 参数及 momentum mirror
    ↓
load_state_dict_from_file
    ↓
验证参数、momentum 和 NPU storage

加载 optimizer state 后，刷新 SwapOptimizerMixin._state_map 中的 state 引用，再执行 swap-out，防止旧 momentum 覆盖新加载的数据。

3.2 技术选型

考虑过以下方案：

• 仅使用 Mock 测试 Swap-Muon：无需 NPU、执行快，但无法验证异步拷贝、storage 释放及真实 checkpoint 路径，因此不采用。
• 直接调用 state_dict/load_state_dict：多优化器场景中 state 可能为 list，与 LayerWise checkpoint 处理结构不一致，因此采用生产实际使用的文件保存和加载接口。
• 默认执行全部 8 卡场景：覆盖完整但 CI 成本较高，因此默认使用 2 卡，8 卡场景保留为 slow 测试。

3.3 功能与性能设计

Pipeline Layout UT 覆盖 3 类正常布局和 7 类异常或冲突布局，共 10 个 pytest 收集项。

Swap Optimizer 默认 UT 调整如下：

项目	优化前	优化后
NPU 数量	8	2
默认拓扑	TP=2、PP=2	TP=1、PP=1、DP=2
模型层数	8	2
Hidden Size	512	128
Attention Heads	32	8
训练 Step	10	2

原 8 卡拓扑矩阵继续通过 slow/--run-all 执行。默认 UT 耗时目标为降低 50% 以上，最终数据以相同 NPU 环境的 CI 结果为准。

3.4 安全隐私与 DFX 设计

• 不读取用户数据或外部模型文件。
• Checkpoint 写入 pytest 临时目录，并按 rank 区分文件名。
• 测试结束后重置 stream、event、state map 和计数器。
• 不新增网络访问、序列化格式或权限要求。
• 保持 MindSpeed master 与 Megatron-LM core_v0.12.1 兼容。

3.5 编程与调用设计

3.5.1 编程模型基本设计

开发环境：

• Python 3.10 及以上。
• PyTorch NPU 及 CANN 环境。
• MindSpeed master。
• Megatron-LM core_v0.12.1。
• Swap-Muon 测试需要至少 2 张 Ascend NPU。

验收标准：

• Pipeline Layout 的 10 个参数场景全部通过。
• Swap-Muon checkpoint round-trip 测试通过。
• 完整默认 UT 和 ST 通过。
• 8 卡 slow 场景可通过 --run-all 运行。
• pre-commit 检查通过。

3.5.2 接口定义与设计

本提案不增加或修改公开 API，测试使用现有接口：

PipelineParallelLayerLayout(layout, pipeline_model_parallel_size)
layout.validate_layer_layout(num_layers, mtp_num_layers)
optimizer.save_state_dict_to_file(checkpoint_path)
optimizer.load_state_dict_from_file(checkpoint_path)

Swap-Muon 仅在 load_state_dict 内部刷新 state 缓存引用，不改变接口参数、返回值或调用方式。

3.5.3 编程手册设计

本提案不需要新增用户编程手册。测试设计、执行命令、环境约束和性能对比记录在 RFC、PR 描述及关联 issue 中。

4. 测试设计

Pipeline Layout：

python -m pytest -vv \
  tests_extend/unit_tests/features/pipeline_parallel/test_pipeline_model_parallel_layout.py

Swap-Muon checkpoint：

python -m pytest -vv -s \
  tests_extend/unit_tests/features/optimizer/test_swap_optimizer.py::TestDistributedOptimizer::test_swap_muon_checkpoint_round_trip

完整默认 optimizer UT：

python -m pytest -vv --durations=20 \
  tests_extend/unit_tests/features/optimizer/test_swap_optimizer.py

完整 8 卡场景：

python -m pytest -vv --run-all \
  tests_extend/unit_tests/features/optimizer/test_swap_optimizer.py

最终执行 MindSpeed 完整 UT、pretrain_base.sh ST 及 pre-commit 门禁。

5. 缺点和风险

• 小模型默认 UT 可能无法发现大规模拓扑问题。应对措施是保留 8 卡 slow 测试矩阵。
• Swap-Muon 测试依赖 2 张 NPU。通过分离 CPU 逻辑与 NPU 闭环，避免扩大其他测试的设备需求。
• 测试涉及 optimizer 私有 state 和 storage 状态，相关断言仅保留在测试代码中，生产接口保持不变。
• Checkpoint 修复可能影响加载后的 swap 流程，需要通过参数、momentum、storage 断言及完整训练 ST 共同验证。

6. 现有技术

本提案复用以下现有机制：

• Pytest 参数化测试。
• MindSpeed DistributedTest 多进程测试框架。
• Megatron Chained/LayerWise Optimizer checkpoint 机制。
• PyTorch optimizer state_dict 加载语义。
• Pytest 临时目录和 slow marker 机制。

与端到端训练测试相比，本方案使用最小模型和定向断言，在保留关键功能覆盖的同时降低 CI 成本。

7. 未解决问题

• 默认 optimizer UT 优化前后的实际耗时数据需在相同 NPU CI 节点补充。
• Swap-Muon checkpoint 修复是否需要回合到其他维护分支，需由社区确认。
• 8 卡 slow 场景的执行周期由 CI 资源策略决定。

附录

• 相关代码
  • mindspeed/core/optimizer/swap_muon/swap_muon.py
  • tests_extend/unit_tests/features/optimizer/test_swap_optimizer.py
  • tests_extend/unit_tests/features/pipeline_parallel/test_pipeline_model_parallel_layout.py
• 术语表
  • PP：Pipeline Parallel
  • VPP：Virtual Pipeline Parallel
  • MTP：Multi-Token Prediction
  • UT：Unit Test
  • CPU Mirror：NPU 参数或优化器状态在 CPU 侧的镜像

欢迎加入社区，感谢您对社区的贡献 🎉！