diff --git a/docs/source/BestPractices/NPU-support.md b/docs/source/BestPractices/NPU-support.md index 0c49fbfe76..967778691b 100644 --- a/docs/source/BestPractices/NPU-support.md +++ b/docs/source/BestPractices/NPU-support.md @@ -7,10 +7,10 @@ 如果你是第一次在 NPU 上使用 ms-swift，推荐按以下顺序阅读： 1. 先查看“支持范围速览”，确认模型、算法和后端是否已验证。 -2. 根据“选择你的使用路线”决定只装基础环境，还是额外安装 MindSpeed/Megatron-SWIFT。 +2. 根据“选择你的使用路线”决定只装基础环境，还是额外安装 Megatron/MegatronAdaptor/TransformerEngineNPU。 3. 根据自己的环境管理习惯选择“本地环境安装”或“镜像/容器环境安装”，然后执行“NPU 可用性检查”。 4. 使用“快速跑通”完成一次 ModelScope 模型 LoRA 训练、合并、推理和部署。 -5. 需要更大规模训练时，再阅读 DDP、DeepSpeed 和 MindSpeed/Megatron-SWIFT 相关章节。 +5. 需要更大规模训练时，再阅读 DDP、DeepSpeed 和 Megatron 相关章节。 ## 硬件配套和支持的操作系统 @@ -58,7 +58,7 @@ vLLM-Ascend 需要与 CANN、torch 和 torch_npu 按整套兼容矩阵安装， | | FSDP | 已支持 | | | FSDP2 | 已支持 | | | DeepSpeed | 已支持 | -| | MindSpeed(Megatron) | 已支持 | +| | Megatron (MCore 0.18 + MegatronAdaptor) | 已支持 | | 低参微调 | FULL | 已支持 | | | LoRA | 已支持 | | | QLoRA | 暂不支持 | @@ -125,12 +125,12 @@ vLLM-Ascend 需要与 CANN、torch 和 torch_npu 按整套兼容矩阵安装， ## 选择你的使用路线 -| 场景 | 推荐路线 | 是否需要 MindSpeed | -| ---------------------------- | --------------------------------------------- | ------------------ | -| 只做普通 SFT/LoRA/推理 | 本地环境安装或镜像/容器环境安装 | 不需要 | -| 需要 Megatron-SWIFT 大模型训练 | 先装基础环境，再装 MindSpeed/Megatron/mcore-bridge | 需要 | -| 需要 GRPO/PPO/DPO 等 RLHF | 基础训练环境 + vLLM-Ascend rollout/deploy | 通常不需要 | -| 只是验证 NPU 是否可用 | 跑 NPU 可用性检查脚本 | 不需要 | +| 场景 | 推荐路线 | 是否需要 MegatronAdaptor | +| ---------------------------- | --------------------------------------------- | ------------------------ | +| 只做普通 SFT/LoRA/推理 | 本地环境安装或镜像/容器环境安装 | 不需要 | +| 需要 Megatron 大模型训练 | 先装基础环境，再装 MCore 0.18、MegatronAdaptor、TransformerEngineNPU 和 mcore-bridge | 需要 | +| 需要 GRPO/PPO/DPO 等 RLHF | 基础训练环境 + vLLM-Ascend rollout/deploy | 通常不需要 | +| 只是验证 NPU 是否可用 | 跑 NPU 可用性检查脚本 | 不需要 | ## 环境准备 @@ -240,74 +240,67 @@ print(torch.npu.device_count()) # 8 print(torch.randn(10, device='npu:0')) ``` -### MindSpeed/Megatron-SWIFT 可选安装 +### MegatronAdaptor + TransformerEngineNPU 可选安装 -如果需要使用 MindSpeed(Megatron-LM)，请按照下面引导安装必要依赖。 +NPU Megatron 路径固定使用 MCore 0.18、MegatronAdaptor、TransformerEngineNPU 和 mcore-bridge，不安装或导入 MindSpeed。下面的命令对应当前首个准入环境：Python 3.12、CANN 9.1 beta、torch/torch_npu 2.10。 ```shell -# 1. 获取并切换 Megatron-LM 至 v0.16.0 版本 -git clone https://github.com/NVIDIA/Megatron-LM.git -cd Megatron-LM -git checkout core_v0.16.0 -cd .. +export SWIFT_ROOT=/home/zyh/code +export MEGATRON_LM_PATH=$SWIFT_ROOT/Megatron-LM +export PYTHONPATH=$SWIFT_ROOT/Megatron-LM:$SWIFT_ROOT/MegatronAdaptor:$SWIFT_ROOT/TransformerEngineNPU:$SWIFT_ROOT/mcore-bridge/src:$SWIFT_ROOT/ms-swift:$PYTHONPATH +# 可选：设置为 0 使用 transformers GDN；默认值 1 使用 MCore GDN。 +# export USE_MCORE_GDN=0 -# 2. 获取并安装 MindSpeed -git clone https://gitcode.com/Ascend/MindSpeed.git -cd MindSpeed -git checkout core_r0.16.0 -pip install -e . -cd .. - -# 3. 获取并安装 mcore-bridge -git clone https://github.com/modelscope/mcore-bridge.git -cd mcore-bridge -pip install -e . -cd .. - -# 4. 获取并安装 triton-ascend -pip install triton-ascend==3.2.1 --extra-index-url=https://triton-ascend.osinfra.cn/pypi/simple - -# 5. 设置环境变量 -export PYTHONPATH=$PYTHONPATH: -export MEGATRON_LM_PATH= - -# 6. 如需回退到 transformers 的 GatedDeltaNet 实现，可关闭 Megatron GDN -export USE_MCORE_GDN=0 -``` +# 固定 checkout：Megatron-LM core_v0.18.0 ba7b5ebce12a， +# MegatronAdaptor core_r0.18.0 2a164d4548e2， +# TransformerEngineNPU main 0a389a677586，mcore-bridge adapt_018_main c820ed47960e +cd $SWIFT_ROOT/Megatron-LM && git checkout ba7b5ebce12a +cd $SWIFT_ROOT/MegatronAdaptor && git checkout 2a164d4548e2 +cd $SWIFT_ROOT/TransformerEngineNPU && git checkout 0a389a677586 +cd $SWIFT_ROOT/mcore-bridge && git checkout c820ed47960e -执行如下命令验证 MindSpeed(Megatron-LM) 是否配置成功： -```shell -python -c "import mindspeed.megatron_adaptor; from swift.megatron.init import init_megatron_env; init_megatron_env(); print('✓ NPU环境下的Megatron-SWIFT配置验证成功！')" +# 先安装与 torch 2.10 配套的 Python 依赖，再对本地 checkout 使用 no-build-isolation/no-deps。 +# Qwen3.5 的 Transformers loader 需要 FLA 的导出符号；其 NPU backend 也来自该 checkout。 +python -m pip install 'packaging>=26.2' transformers==5.12.1 torchvision==0.25.0 torchaudio==2.10.0 nvdlfw-inspect 'qwen-vl-utils>=0.0.14' +python -m pip install -e $SWIFT_ROOT/flash-linear-attention --no-build-isolation --no-deps +python -m pip install -e $SWIFT_ROOT/Megatron-LM --no-build-isolation --no-deps +python -m pip install -e $SWIFT_ROOT/MegatronAdaptor --no-build-isolation --no-deps +python -m pip install -e $SWIFT_ROOT/TransformerEngineNPU --no-build-isolation --no-deps +python -m pip install -e $SWIFT_ROOT/mcore-bridge --no-build-isolation --no-deps +python -m pip install -e $SWIFT_ROOT/ms-swift --no-build-isolation --no-deps ``` -### Qwen3.5 FLA补丁说明 - -Qwen3.5 在昇腾 NPU 上直接使用 `flash-linear-attention`（FLA）的原生 Triton-Ascend GDN backend。ms-swift 不再内置或维护一份 MindSpeed `chunk_gated_delta_rule` 副本。 +固定的 `mcore-bridge` 提交已包含 TransformerEngineNPU main 的 NPU LoRA `LayerNormLinear`、GroupedLinear/expert 和 checkpoint 映射适配，不依赖未提交的 working-tree patch。 -ms-swift 只保留两处 NPU 兼容处理： +配套 bridge 修复同时删除了 MindSpeed defaults 探测。准入条件应同时检查 `find_spec('mindspeed') is None`、`mindspeed` 不在 `sys.modules`，并确认 bridge 运行时代码中不存在 MindSpeed 引用。 -1. Transformers 的 FLA 可用性检查包含 CUDA 条件，ms-swift 会在 NPU 上按 FLA 的公共 causal-conv 与 GDN 入口是否可导入补充判断。具体 NPU backend 的选择与校验由 FLA 自己负责。 -2. Transformers 当前使用 `torch.cuda.current_device()` 初始化 `FusedRMSNormGated`，因此 NPU 上仍保留 Qwen3.5 原生 torch norm；这不影响 GDN 使用 FLA。 +验证时应按代表性 dense、GDN、MoE 拓扑分别记录完整命令、源码 SHA、版本、实际拓扑、日志、退出码以及 pass/framework failure/environment failure/capacity failure 分类；不要把两步 smoke 或单元测试结果当成长跑验收。 -可以将这条调用链理解为： +验证时要同时检查源码归属和 MindSpeed 缺失；只检查 `import` 成功不足以证明运行时绑定正确： -```text -Qwen3.5 modeling.chunk_gated_delta_rule - -> fla.ops.gated_delta_rule.chunk_gated_delta_rule - -> fla.ops.gated_delta_rule.backends.triton_ascend +```shell +python - <<'PY' +import importlib.util +import megatron_adaptor, megatron.core, transformer_engine, mcore_bridge, fla +from swift.megatron.init import init_megatron_env +init_megatron_env() +for module in (megatron_adaptor, megatron.core, transformer_engine, mcore_bridge): + print(module.__name__, module.__file__) +assert '/home/zyh/code/' in megatron_adaptor.__file__ +assert '/home/zyh/code/' in megatron.core.__file__ +assert '/home/zyh/code/' in transformer_engine.__file__ +assert '/home/zyh/code/' in mcore_bridge.__file__ +assert '/home/zyh/code/' in fla.__file__ +assert importlib.util.find_spec('mindspeed') is None +print('NPU MegatronAdaptor/MCore 0.18 environment verified') +PY ``` -- 当前请安装 FLA main 分支最新版本以获得 NPU GDN 支持。 -- 当前验证环境：torch 2.9.0+cpu、torch-npu 2.9.0.post2、FLA 0.5.2 main、triton-ascend 3.2.1、transformers 5.12.1。 - - -当前 Qwen3.5 在 NPU 上如果走 transformers 后端或 Megatron-SWIFT 后端训练，还需要额外注意版本和功能约束： - -1. 当前 NPU 文档中约定的 MindSpeed 训练组合是 `Megatron-LM v0.16.0 + MindSpeed core_r0.16.0`。在这个组合下，`mcore-bridge` 默认按 `USE_MCORE_GDN=1` 走 Megatron-Core/MindSpeed GDN；若显式设置 `USE_MCORE_GDN=0`，则走 transformers 版 GDN。该路径会优先尝试使用上述 FLA Triton-Ascend backend；FLA 不可用时仅记录 warning 并保留当前 MindSpeed/Megatron 选择的 GDN 实现，不主动替换为 Torch 实现，因此 FLA 不是普通 Megatron 训练的强依赖。 +### Qwen3.5 MCore GDN 说明 -2. MindSpeed 初始化以及训练参数 `repatch()` 可能会替换 FLA 的公共 GDN 入口。ms-swift 会在构建 `mcore-bridge` 前和 `repatch()` 后，优先从 MindSpeed patch manager 保存的 `orig_func` 恢复 FLA 原始 callable；仅在成功恢复 FLA 时才同步刷新 `mcore-bridge` 已缓存的 callable。如果 FLA 缺失或恢复失败，则不修改当前 GDN callable，并记录 warning。packed/varlen 能力取决于保留的实现；如果它不支持非空 `cu_seqlens`，实际 GDN 调用会直接失败，此时需要安装可用的 FLA。 +Qwen3.5 的 Megatron NPU 路径默认使用 MCore 0.18 的 `GatedDeltaNet`（GDN，门控 DeltaNet 状态空间层），由 MegatronAdaptor 提供 NPU 所需的 Torch-native causal-conv、GDN 和归一化替换。`USE_MCORE_GDN` 仍然有效：设为 `1` 使用 MCore GDN，设为 `0` 选择 transformers GDN；它是模型实现选择项，与本次 0.18 依赖迁移无关。 -3. 已在 8 卡 Atlas 900 A2 上完成 Qwen3.5-4B 的 Transformers/FSDP LoRA 验证：BF16、`alpaca-gpt4-data-zh`、`packing=true`、`max_length=512`、每卡 batch size 1、梯度累积 1，共训练 300 steps；全程 loss/grad_norm 为有限值，并成功保存 checkpoint。相同配置下，NPU 与 GPU 对照实验的 loss 变化趋势对齐。 +MCore-bridge 负责把 Qwen3.5/Omni 的模型参数和 packed/padding-free 元数据接到 MCore；CP（context parallel，沿序列维度切分的并行）使用 MCore 0.18 的实现。需要分别验证 dense、GDN、MoE 的 forward/backward、optimizer、packing、CP 和 save/resume；这里的环境安装成功不等于长训练或性能已经验收。 ### 环境查看 @@ -754,7 +747,7 @@ ASCEND_RT_VISIBLE_DEVICES=0 swift deploy \ ### Q2: 训练时应该选择 FSDP、DeepSpeed 还是 Megatron-SWIFT？ -普通 SFT 优先参考本文兼容性表中的 `FSDP1/FSDP2/deepspeed` 组合；如果模型规模较大、需要更高并行能力，再使用 Megatron-SWIFT，并按安装章节额外安装 MindSpeed、Megatron-LM 和 mcore-bridge。DeepSpeed 可以降低显存压力，但速度可能下降，遇到性能问题时可以对比 FSDP 方案。 +普通 SFT 优先参考本文兼容性表中的 `FSDP1/FSDP2/deepspeed` 组合；如果模型规模较大、需要更高并行能力，再使用 Megatron，并按安装章节额外安装 MCore 0.18、MegatronAdaptor、TransformerEngineNPU 和 mcore-bridge。DeepSpeed 可以降低显存压力，但速度可能下降，遇到性能问题时可以对比 FSDP 方案。 ### Q3: NPU 模型 Patch 需要手动关闭吗？ @@ -824,14 +817,14 @@ python -c "import torch, vllm_ascend; print(torch.__version__); print(vllm_ascen 不要默认可以。下载或加载大模型前，先检查 `config.json` 是否包含 `quantization_config`，再检查 safetensors 的真实 dtype。当前 NPU 支持范围中量化/QLoRA 仍属于暂不支持或未完全验证能力；如果模型权重是 FP8 block quantized，而当前 NPU 软件栈不支持对应 FP8 路径，应先换用 BF16 权重，或离线转换为 BF16 后再训练/加载。 -### Q14: Megatron-SWIFT 导入到错误的 Megatron/MindSpeed 怎么排查？ +### Q14: Megatron 导入到错误的 MCore/Adaptor 怎么排查？ -跑 Megatron-SWIFT 前，`PYTHONPATH` 和 `MEGATRON_LM_PATH` 必须指向同一份 Megatron-LM 源码树。否则 Python 可能能启动，但实际导入到的是另一套 Megatron/MindSpeed 组合，后续报错会很像模型或参数问题。 +跑 Megatron 前，`PYTHONPATH`、`MEGATRON_LM_PATH` 和本地 MegatronAdaptor/TransformerEngineNPU 路径必须指向同一套 `/home/zyh/code` checkout。否则 Python 可能能启动，但实际导入到的是另一套 MCore/Adaptor 组合，后续报错会很像模型或参数问题。 ```shell -export PYTHONPATH=$PYTHONPATH: -export MEGATRON_LM_PATH= -python -c "import megatron, os; print(megatron.__file__); print(os.environ.get('MEGATRON_LM_PATH'))" +export MEGATRON_LM_PATH=/home/zyh/code/Megatron-LM +export PYTHONPATH=/home/zyh/code/Megatron-LM:/home/zyh/code/MegatronAdaptor:/home/zyh/code/TransformerEngineNPU:/home/zyh/code/mcore-bridge/src:/home/zyh/code/ms-swift:$PYTHONPATH +python -c "import importlib.util, megatron_adaptor, megatron.core, transformer_engine, os; print(megatron_adaptor.__file__); print(megatron.core.__file__); print(transformer_engine.__file__); print(os.environ.get('MEGATRON_LM_PATH')); assert importlib.util.find_spec('mindspeed') is None" ``` 如果二者不一致，先修环境变量，再继续排查模型构建、权重加载或并行配置。 diff --git a/docs/source_en/BestPractices/NPU-support.md b/docs/source_en/BestPractices/NPU-support.md index 4382e2e6ae..911c06874c 100644 --- a/docs/source_en/BestPractices/NPU-support.md +++ b/docs/source_en/BestPractices/NPU-support.md @@ -7,10 +7,10 @@ This document describes how to prepare the environment, train models, save and m If this is your first time using ms-swift on NPUs, we recommend reading this document in the following order: 1. Check "Support Scope at a Glance" first to confirm whether your model, algorithm, and backend have been verified. -2. Use "Choose Your Usage Path" to decide whether you only need the base environment or also need MindSpeed/Megatron-SWIFT. +2. Use "Choose Your Usage Path" to decide whether you only need the base environment or also need Megatron/MegatronAdaptor/TransformerEngineNPU. 3. Choose "Local Environment Installation" or "Image/Container Environment Installation" according to your own environment management preference, then run "NPU Availability Check". 4. Use "Quick Start" to complete one ModelScope model LoRA training, merge, inference, and deployment flow. -5. For larger-scale training, continue reading the DDP, DeepSpeed, and MindSpeed/Megatron-SWIFT sections. +5. For larger-scale training, continue reading the DDP, DeepSpeed, and Megatron sections. ## Hardware and Supported Operating Systems @@ -44,7 +44,6 @@ Recommended base environment versions: | CANN | >= 8.5.1 | | torch | >= 2.7.1 | | torch_npu | >= 2.7.1.post4 | - Install vLLM-Ascend together with CANN, torch, and torch_npu as one compatibility set. For details, see the [vLLM-Ascend 0.23.0 installation guide](https://docs.vllm.ai/projects/ascend/en/v0.23.0/installation.html). For base environment setup, see the [Ascend PyTorch installation guide](https://gitcode.com/Ascend/pytorch). The examples in this document were verified on 8 * Ascend 910B3 64G. @@ -59,7 +58,7 @@ For base environment setup, see the [Ascend PyTorch installation guide](https:// | | FSDP | Supported | | | FSDP2 | Supported | | | DeepSpeed | Supported | -| | MindSpeed(Megatron) | Supported | +| | Megatron (MCore 0.18 + MegatronAdaptor) | Supported | | PEFT | FULL | Supported | | | LoRA | Supported | | | QLoRA | Not supported | @@ -126,12 +125,12 @@ For regular Qwen3.5 GRPO/SFT LoRA training, avoid explicitly passing `--model_ty ## Choose Your Usage Path -| Scenario | Recommended path | Need MindSpeed | -| -------------------------------------- | ----------------------------------------------------- | -------------- | +| Scenario | Recommended path | Need MegatronAdaptor | +| -------------------------------------- | ----------------------------------------------------- | -------------------- | | Ordinary SFT/LoRA/inference | Local environment installation or image/container installation | No | -| Megatron-SWIFT large-model training | Install the base environment, then install MindSpeed/Megatron/mcore-bridge | Yes | -| GRPO/PPO/DPO and other RLHF workflows | Base training environment + vLLM-Ascend rollout/deploy | Usually no | -| Only verifying whether NPUs are usable | Run the NPU availability check script | No | +| Megatron large-model training | Install the base environment, then install MCore 0.18, MegatronAdaptor, TransformerEngineNPU, and mcore-bridge | Yes | +| GRPO/PPO/DPO and other RLHF workflows | Base training environment + vLLM-Ascend rollout/deploy | Usually no | +| Only verifying whether NPUs are usable | Run the NPU availability check script | No | ## Environment Preparation @@ -243,74 +242,68 @@ print(torch.npu.device_count()) # 8 print(torch.randn(10, device='npu:0')) ``` -### Optional MindSpeed/Megatron-SWIFT Installation +### Optional MegatronAdaptor + TransformerEngineNPU Installation -If you need MindSpeed(Megatron-LM), install the required dependencies as follows. +The NPU Megatron path uses MCore 0.18, MegatronAdaptor, TransformerEngineNPU, and mcore-bridge. MindSpeed must not be installed or imported. The commands below describe the first admission environment: Python 3.12, CANN 9.1 beta, and torch/torch_npu 2.10. ```shell -# 1. Clone Megatron-LM and switch to v0.16.0 -git clone https://github.com/NVIDIA/Megatron-LM.git -cd Megatron-LM -git checkout core_v0.16.0 -cd .. - -# 2. Clone and install MindSpeed -git clone https://gitcode.com/Ascend/MindSpeed.git -cd MindSpeed -git checkout core_r0.16.0 -pip install -e . -cd .. - -# 3. Clone and install mcore-bridge -git clone https://github.com/modelscope/mcore-bridge.git -cd mcore-bridge -pip install -e . -cd .. - -# 4. Download and install triton-ascend -pip install triton-ascend==3.2.1 --extra-index-url=https://triton-ascend.osinfra.cn/pypi/simple +export SWIFT_ROOT=/home/zyh/code +export MEGATRON_LM_PATH=$SWIFT_ROOT/Megatron-LM +export PYTHONPATH=$SWIFT_ROOT/Megatron-LM:$SWIFT_ROOT/MegatronAdaptor:$SWIFT_ROOT/TransformerEngineNPU:$SWIFT_ROOT/mcore-bridge/src:$SWIFT_ROOT/ms-swift:$PYTHONPATH +# Optional: set this to 0 for Transformers GDN; the default value 1 uses MCore GDN. +# export USE_MCORE_GDN=0 -# 5. Set environment variables -export PYTHONPATH=$PYTHONPATH: -export MEGATRON_LM_PATH= +# Pinned checkouts: Megatron-LM core_v0.18.0 ba7b5ebce12a, +# MegatronAdaptor core_r0.18.0 2a164d4548e2, +# TransformerEngineNPU main 0a389a677586, mcore-bridge adapt_018_main c820ed47960e +cd $SWIFT_ROOT/Megatron-LM && git checkout ba7b5ebce12a +cd $SWIFT_ROOT/MegatronAdaptor && git checkout 2a164d4548e2 +cd $SWIFT_ROOT/TransformerEngineNPU && git checkout 0a389a677586 +cd $SWIFT_ROOT/mcore-bridge && git checkout c820ed47960e -# 6. Disable Megatron GDN if you need to fall back to the transformers GatedDeltaNet implementation -export USE_MCORE_GDN=0 +# Install dependencies compatible with torch 2.10, then install each local checkout +# without build isolation or dependency resolution. Qwen3.5's Transformers loader +# needs FLA exports; its NPU backend is also provided by this checkout. +python -m pip install 'packaging>=26.2' transformers==5.12.1 torchvision==0.25.0 torchaudio==2.10.0 nvdlfw-inspect 'qwen-vl-utils>=0.0.14' +python -m pip install -e $SWIFT_ROOT/flash-linear-attention --no-build-isolation --no-deps +python -m pip install -e $SWIFT_ROOT/Megatron-LM --no-build-isolation --no-deps +python -m pip install -e $SWIFT_ROOT/MegatronAdaptor --no-build-isolation --no-deps +python -m pip install -e $SWIFT_ROOT/TransformerEngineNPU --no-build-isolation --no-deps +python -m pip install -e $SWIFT_ROOT/mcore-bridge --no-build-isolation --no-deps +python -m pip install -e $SWIFT_ROOT/ms-swift --no-build-isolation --no-deps ``` -Run the following command to verify that MindSpeed(Megatron-LM) is configured correctly: +The pinned `mcore-bridge` commit contains the TransformerEngineNPU main adaptations for NPU LoRA `LayerNormLinear`, GroupedLinear/expert handling, and checkpoint mapping; it does not rely on an uncommitted working-tree patch. -```shell -python -c "import mindspeed.megatron_adaptor; from swift.megatron.init import init_megatron_env; init_megatron_env(); print('✓ Megatron-SWIFT configuration verified successfully in the NPU environment!')" -``` +The companion bridge fix also removes the MindSpeed-defaults probe. Admission should check that `find_spec('mindspeed') is None`, `mindspeed` is absent from `sys.modules`, and the bridge runtime source contains no MindSpeed references. -### Qwen3.5 FLA Patch Notes +For validation, record the complete command, source SHAs, versions, runtime topology, log, exit code, and pass/framework failure/environment failure/capacity failure classification for representative dense, GDN, and MoE topologies. Do not treat a two-step smoke or a unit-test result as long-run acceptance. -Qwen3.5 on Ascend NPU now uses the native Triton-Ascend GDN backend from `flash-linear-attention` (FLA). ms-swift no longer embeds or maintains a copy of MindSpeed's `chunk_gated_delta_rule`. +Check both source ownership and the absence of MindSpeed; a successful import alone does not prove the runtime bindings: -ms-swift retains only two NPU compatibility adjustments: - -1. Transformers' FLA availability check includes a CUDA condition. On NPU, ms-swift supplements it by checking FLA's public causal-conv and GDN entry points. FLA remains responsible for selecting and validating the concrete NPU backend. -2. Transformers currently initializes `FusedRMSNormGated` with `torch.cuda.current_device()`, so NPU keeps Qwen3.5's native torch norm. This does not affect the FLA GDN path. - -The call chain can be understood as: - -```text -Qwen3.5 modeling.chunk_gated_delta_rule - -> fla.ops.gated_delta_rule.chunk_gated_delta_rule - -> fla.ops.gated_delta_rule.backends.triton_ascend +```shell +python - <<'PY' +import importlib.util +import megatron_adaptor, megatron.core, transformer_engine, mcore_bridge, fla +from swift.megatron.init import init_megatron_env +init_megatron_env() +for module in (megatron_adaptor, megatron.core, transformer_engine, mcore_bridge): + print(module.__name__, module.__file__) +assert '/home/zyh/code/' in megatron_adaptor.__file__ +assert '/home/zyh/code/' in megatron.core.__file__ +assert '/home/zyh/code/' in transformer_engine.__file__ +assert '/home/zyh/code/' in mcore_bridge.__file__ +assert '/home/zyh/code/' in fla.__file__ +assert importlib.util.find_spec('mindspeed') is None +print('NPU MegatronAdaptor/MCore 0.18 environment verified') +PY ``` -- Install the latest FLA main branch to obtain NPU GDN support. -- Current verification environment: torch 2.9.0+cpu, torch-npu 2.9.0.post2, FLA 0.5.2 main, triton-ascend 3.2.1, and transformers 5.12.1. - -When running Qwen3.5 on NPU with either the transformers backend or the Megatron-SWIFT backend, note the following version and feature constraints: - -1. The MindSpeed training combination currently pinned by the NPU documentation is `Megatron-LM v0.16.0 + MindSpeed core_r0.16.0`. `mcore-bridge` defaults to the Megatron-Core/MindSpeed GDN path with `USE_MCORE_GDN=1`. With `USE_MCORE_GDN=0`, it uses the transformers GDN. That path first attempts to use the FLA Triton-Ascend backend described above; when FLA is unavailable, it only logs a warning and preserves the currently selected MindSpeed/Megatron GDN implementation rather than proactively replacing it with Torch, so FLA is not a hard dependency for ordinary Megatron training. +### Qwen3.5 MCore GDN Notes -2. MindSpeed initialization and training-argument `repatch()` may replace FLA's public GDN entry point. Before building `mcore-bridge` and again after `repatch()`, ms-swift first tries to restore FLA's original callable from the `orig_func` saved by the MindSpeed patch manager. The `mcore-bridge` cached callable is refreshed only after FLA is restored successfully. If FLA is missing or restoration fails, the current GDN callable is not modified and a warning is recorded. Packed/varlen support depends on the retained implementation; if it does not support non-empty `cu_seqlens`, the actual GDN call fails and a working FLA installation is required. +The Qwen3.5 Megatron NPU path uses MCore 0.18's `GatedDeltaNet` (GDN, the gated DeltaNet state-space layer) by default. MegatronAdaptor provides the NPU Torch-native replacements for causal-conv, GDN, and normalization. `USE_MCORE_GDN` remains supported: `1` selects MCore GDN and `0` selects the Transformers GDN; it is a model-implementation choice independent of the 0.18 dependency migration. -3. Qwen3.5-4B has completed a Transformers/FSDP LoRA validation on 8 Atlas 900 A2 cards: BF16, `alpaca-gpt4-data-zh`, `packing=true`, `max_length=512`, per-device batch size 1, gradient accumulation 1, and 300 training steps. Loss and grad norm remained finite throughout, and a checkpoint was saved successfully. Under the same configuration, the NPU loss trend aligned with the GPU reference. +mcore-bridge connects Qwen3.5/Omni model parameters and packed/padding-free metadata to MCore; CP (context parallel, sequence-dimension parallelism) uses the MCore 0.18 implementation. Dense, GDN, and MoE forward/backward, optimizer, packing, CP, and save/resume must be validated independently; a successful environment install is not evidence of long-run stability or performance. ### Environment Viewing Check the P2P connections of the NPU, where we can see that each NPU is interconnected through 7 HCCS links with other NPUs. @@ -758,7 +751,7 @@ First confirm that you have run `source /usr/local/Ascend/ascend-toolkit/set_env ### Q2: How should I choose between FSDP, DeepSpeed, and Megatron-SWIFT? -For ordinary SFT, first refer to the verified `FSDP1/FSDP2/deepspeed` combinations in this document. For larger models or higher parallelism requirements, use Megatron-SWIFT and install MindSpeed, Megatron-LM, and mcore-bridge as described in the installation section. DeepSpeed can reduce memory pressure but may reduce speed, so compare it with FSDP when tuning performance. +For ordinary SFT, first refer to the verified `FSDP1/FSDP2/deepspeed` combinations in this document. For larger models or higher parallelism requirements, use Megatron and install MCore 0.18, MegatronAdaptor, TransformerEngineNPU, and mcore-bridge as described in the installation section. DeepSpeed can reduce memory pressure but may reduce speed, so compare it with FSDP when tuning performance. ### Q3: Do I need to manually disable the NPU model patch? @@ -828,14 +821,14 @@ If the error message contains `Current device type ... does not match the instal Do not assume they can. Before downloading or loading a large model, check whether `config.json` contains `quantization_config`, and check the actual dtype in safetensors. Quantization/QLoRA is still listed as not supported or not fully verified in the NPU support scope. If model weights are FP8 block quantized and your NPU software stack does not support that FP8 path, use BF16 weights first, or convert the model offline to BF16 before training/loading. -### Q14: How do I troubleshoot Megatron-SWIFT importing the wrong Megatron/MindSpeed? +### Q14: How do I troubleshoot Megatron importing the wrong MCore/Adaptor? -Before running Megatron-SWIFT, `PYTHONPATH` and `MEGATRON_LM_PATH` must point to the same Megatron-LM source tree. Otherwise Python may start successfully while importing a different Megatron/MindSpeed combination, which can make later errors look like model or argument issues. +Before running Megatron, `PYTHONPATH`, `MEGATRON_LM_PATH`, and the local MegatronAdaptor/TransformerEngineNPU paths must point to the same `/home/zyh/code` checkouts. Otherwise Python may start successfully while importing a different MCore/Adaptor combination, which can make later errors look like model or argument issues. ```shell -export PYTHONPATH=$PYTHONPATH: -export MEGATRON_LM_PATH= -python -c "import megatron, os; print(megatron.__file__); print(os.environ.get('MEGATRON_LM_PATH'))" +export MEGATRON_LM_PATH=/home/zyh/code/Megatron-LM +export PYTHONPATH=/home/zyh/code/Megatron-LM:/home/zyh/code/MegatronAdaptor:/home/zyh/code/TransformerEngineNPU:/home/zyh/code/mcore-bridge/src:/home/zyh/code/ms-swift:$PYTHONPATH +python -c "import importlib.util, megatron_adaptor, megatron.core, transformer_engine, os; print(megatron_adaptor.__file__); print(megatron.core.__file__); print(transformer_engine.__file__); print(os.environ.get('MEGATRON_LM_PATH')); assert importlib.util.find_spec('mindspeed') is None" ``` If they do not match, fix the environment variables before continuing with model construction, weight loading, or parallel configuration troubleshooting. diff --git a/swift/megatron/__init__.py b/swift/megatron/__init__.py index bfc2bc35ab..ffbe91274c 100644 --- a/swift/megatron/__init__.py +++ b/swift/megatron/__init__.py @@ -4,10 +4,15 @@ from transformers.utils import is_torch_npu_available if is_torch_npu_available(): - from swift.model.npu_patch.mindspeed import prepare_mindspeed_gdn_import - prepare_mindspeed_gdn_import() - # Enable Megatron on Ascend NPU - import mindspeed.megatron_adaptor # F401 + try: + # MegatronAdaptor must patch MCore before mcore-bridge or any TE + # wrapper caches the original callables. + import megatron_adaptor # F401 + except ModuleNotFoundError as exc: + if exc.name == 'megatron_adaptor': + raise ImportError('Megatron on Ascend NPU requires MegatronAdaptor. Install the local ' + 'MegatronAdaptor checkout before importing swift.megatron.') from exc + raise from .init import init_megatron_env init_megatron_env() except Exception: diff --git a/swift/megatron/arguments/megatron_args.py b/swift/megatron/arguments/megatron_args.py index b1417c6c8a..c584359b7a 100644 --- a/swift/megatron/arguments/megatron_args.py +++ b/swift/megatron/arguments/megatron_args.py @@ -917,6 +917,10 @@ def __post_init__(self): def _check_megatron_fsdp(self): if not self.use_megatron_fsdp: return + if is_torch_npu_available() and self.virtual_pipeline_model_parallel_size is not None: + raise ValueError('NPU Megatron-FSDP does not support virtual pipeline parallelism yet. ' + 'Its DCP preprocessing currently handles one model chunk per pipeline rank; ' + 'set virtual_pipeline_model_parallel_size=None and use ordinary PP instead.') # Megatron-FSDP is only compatible with the distributed optimizer. if not self.use_distributed_optimizer: logger.info('Megatron-FSDP is only compatible with use_distributed_optimizer=True; setting it to True.') diff --git a/swift/megatron/init.py b/swift/megatron/init.py index e85a5be3c3..d064a8f9d1 100644 --- a/swift/megatron/init.py +++ b/swift/megatron/init.py @@ -351,9 +351,6 @@ def init_megatron_env(): os.environ.pop('VLLM_USE_MODELSCOPE', None) logging_level = logging.root.level _patch_unified_memory() - if is_torch_npu_available(): - from swift.model.npu_patcher import patch_mindspeed_fla_gdn_implementation - patch_mindspeed_fla_gdn_implementation() _patch__batched_p2p_ops() logging.root.setLevel(logging_level) # revert logger level try: diff --git a/swift/megatron/model/utils.py b/swift/megatron/model/utils.py index 0e822a5f1c..fa3264aa5a 100644 --- a/swift/megatron/model/utils.py +++ b/swift/megatron/model/utils.py @@ -3,7 +3,6 @@ from mcore_bridge import ModelConfig from mcore_bridge import get_mcore_model as _get_mcore_model from mcore_bridge import hf_to_mcore_config -from transformers.utils import is_torch_npu_available from typing import Any, Generator, Optional, Tuple from swift.utils import get_logger @@ -72,8 +71,6 @@ def get_mcore_model_config(args, hf_config): if args.megatron_extra_kwargs: kwargs.update(args.megatron_extra_kwargs) config = ModelConfig(**kwargs) - if is_torch_npu_available() and getattr(args, 'attention_backend', 'flash') != 'local': - setattr(config, 'use_flash_attn', True) _check_attention_backend(args, config) _check_padding_free(args, config) return config diff --git a/swift/megatron/pipelines/train/sft.py b/swift/megatron/pipelines/train/sft.py index e3fe6a1ffb..500bc2babf 100644 --- a/swift/megatron/pipelines/train/sft.py +++ b/swift/megatron/pipelines/train/sft.py @@ -2,8 +2,6 @@ import os import torch import torch.distributed as dist -from dataclasses import asdict -from transformers.utils import is_torch_npu_available from typing import List, Optional, Union from swift.megatron.arguments import MegatronSftArguments @@ -11,12 +9,6 @@ from swift.pipelines import SwiftSft from swift.utils import append_to_jsonl, get_logger, is_last_rank, plot_images -if is_torch_npu_available(): - # Enable Megatron on Ascend NPU - from swift.model.npu_patcher import apply_mindspeed_patches -else: - apply_mindspeed_patches = None - logger = get_logger() @@ -40,14 +32,6 @@ def __init__(self, args: Optional[Union[List[str], MegatronSftArguments]] = None self.train_msg = {} super(SwiftSft, self).__init__(args) args = self.args - if apply_mindspeed_patches is not None: - megatron_args = asdict(self.args) - if args.attention_backend != 'local': - # MindSpeed requires passing `use_flash_attn` to Megatron - # to enable flash attention on Ascend NPU. - args.use_flash_attn = True - megatron_args['use_flash_attn'] = True - apply_mindspeed_patches(megatron_args) template_cls = args.template_meta.template_cls if args.model_meta.is_multimodal and template_cls and template_cls.use_model: kwargs = {'return_dummy_model': True} diff --git a/swift/megatron/trainers/utils.py b/swift/megatron/trainers/utils.py index 36a1004df2..a13b6e1adb 100644 --- a/swift/megatron/trainers/utils.py +++ b/swift/megatron/trainers/utils.py @@ -361,9 +361,10 @@ def _should_use_npu_generated_attention_mask(args) -> bool: return False if args.task_type != 'causal_lm' or args.padding_free: return False - if getattr(args, 'attention_backend', None) == 'local': + attention_backend = getattr(args, 'attention_backend', None) + if getattr(attention_backend, 'name', attention_backend) == 'local': return False - return bool(getattr(args, 'use_flash_attn', False)) + return True def _prepare_npu_generated_attention_mask(batch, *, keep_attention_mask_2d: bool) -> None: diff --git a/swift/megatron/utils/megatron_fsdp_checkpoint.py b/swift/megatron/utils/megatron_fsdp_checkpoint.py index 7e0cb3a86a..5402bdce54 100644 --- a/swift/megatron/utils/megatron_fsdp_checkpoint.py +++ b/swift/megatron/utils/megatron_fsdp_checkpoint.py @@ -86,8 +86,8 @@ def _validate_optimizer_state(state_dict): def _prepare_state_dict(args, state_dict, model, preserve_raw_state: bool = False): _validate_optimizer_state(state_dict) if is_torch_npu_available(): - from swift.model.npu_patch.mindspeed import complete_mindspeed_fsdp_dtensor_optimizer_state - complete_mindspeed_fsdp_dtensor_optimizer_state(state_dict, model) + from swift.model.npu_patch.megatron_fsdp import complete_npu_fsdp_dtensor_optimizer_state + complete_npu_fsdp_dtensor_optimizer_state(state_dict, model) # Preprocessing rewrites the model and optimizer containers. Keep their original structure # for the wrapper and optimizer load_state_dict calls after DCP has populated the tensors. diff --git a/swift/megatron/utils/megatron_lm_utils.py b/swift/megatron/utils/megatron_lm_utils.py index 8d641bb02c..8dceb04c3e 100644 --- a/swift/megatron/utils/megatron_lm_utils.py +++ b/swift/megatron/utils/megatron_lm_utils.py @@ -184,10 +184,6 @@ def _optimizer_sharded_state_dict(optimizer, state_dict, optim_sd_kwargs): def _load_optimizer_state_dict(optimizer, state_dict): - if is_torch_npu_available(): - from swift.model.npu_patch.megatron_checkpoint import load_optimizer_state_dict - load_optimizer_state_dict(optimizer, state_dict) - return optimizer.load_state_dict(state_dict) @@ -309,7 +305,7 @@ def save_mcore_checkpoint( models = unwrap_model(models) fsdp_dtensor = bool(models) and getattr(args, 'use_megatron_fsdp', False) if fsdp_dtensor and args.async_save: - raise ValueError('Megatron-FSDP fsdp_dtensor checkpoint does not support async_save in Megatron-Core 0.16.') + raise ValueError('Megatron-FSDP fsdp_dtensor checkpoint does not support async_save.') rng_state = ( _get_rng_state( fsdp_dtensor=fsdp_dtensor, @@ -370,6 +366,24 @@ def save_mcore_checkpoint( if torch.distributed.is_initialized(): torch.distributed.barrier() + if getattr(args, 'use_megatron_fsdp', False) and is_torch_npu_available(): + # ``state_dict_for_save_checkpoint`` temporarily exposes the + # FSDP-managed DTensor parameters on the wrapped module. The + # regular Megatron-FSDP loop restores raw module parameters before + # its next forward, but Swift immediately calls mcore-bridge's + # HF exporter after this function. First complete the same + # parameter all-gather used by the GPU FSDP forward boundary, then + # expose those full raw parameters to bridge. Restoring raw + # parameters without this gather would leave empty DP shards on + # non-zero ranks and can crash during bridge's reshape/copy path. + for model in models: + start_param_sync = getattr(model, 'start_param_sync', None) + if start_param_sync is not None: + start_param_sync(force_sync=True) + restore_raw_parameters = getattr(model, '_replace_param_with_raw_if_needed', None) + if restore_raw_parameters is not None: + restore_raw_parameters() + if is_master(): tracker_path = os.path.join(output_dir, 'latest_checkpointed_iteration.txt') @@ -641,6 +655,11 @@ def _ensure_fsdp_tensor_parallel_attributes(args, model) -> None: def wrap_model(args, models, wrap_with_ddp: bool = True): + if getattr(args, 'use_megatron_fsdp', False) and is_torch_npu_available(): + from swift.model.npu_patch.megatron_fsdp import patch_megatron_fsdp_optimizer + patch_megatron_fsdp_optimizer( + use_precision_aware_optimizer=getattr(args, 'use_precision_aware_optimizer', False)) + # Set tensor model parallel attributes if not set. # Only parameters that are already tensor model parallel have these # attributes set for them. We should make sure the default attributes diff --git a/swift/model/npu_patch/__init__.py b/swift/model/npu_patch/__init__.py index e5df1a7199..0a8495a595 100644 --- a/swift/model/npu_patch/__init__.py +++ b/swift/model/npu_patch/__init__.py @@ -5,7 +5,6 @@ from transformers.utils import strtobool from .fsdp import NPUCastError -from .mindspeed import apply_mindspeed_patches, patch_mindspeed_fla_gdn_implementation _APPLIED = False _ENABLE_NPU_MODEL_PATCH_ARGS = ('--enable_npu_model_patch', '--enable-npu-model-patch') @@ -46,4 +45,4 @@ def apply_all_patches() -> None: _APPLIED = True -__all__ = ['NPUCastError', 'apply_all_patches', 'apply_mindspeed_patches', 'patch_mindspeed_fla_gdn_implementation'] +__all__ = ['NPUCastError', 'apply_all_patches'] diff --git a/swift/model/npu_patch/megatron_checkpoint.py b/swift/model/npu_patch/megatron_checkpoint.py index 8744f37f48..30d4fe7806 100644 --- a/swift/model/npu_patch/megatron_checkpoint.py +++ b/swift/model/npu_patch/megatron_checkpoint.py @@ -1,8 +1,9 @@ # Copyright (c) ModelScope Contributors. All rights reserved. """NPU-only Megatron checkpoint compatibility helpers. -MindSpeed patches Megatron's distributed optimizer on NPU, but some Megatron-Core -checkpoint formats still need the native Megatron param_state loaders. +Megatron Core builds optimizer checkpoint state from scalar step values. On +Ascend, those values may be NPU tensors with object identity based hashing, so +normalize them temporarily while MCore builds the sharded state dict. """ from __future__ import annotations @@ -98,129 +99,4 @@ def optimizer_sharded_state_dict(optimizer, state_dict, **optim_sd_kwargs): return optimizer.sharded_state_dict(state_dict, **optim_sd_kwargs) -def _iter_distributed_optimizers(optimizer): - visited = set() - - def visit(obj): - if obj is None or id(obj) in visited: - return - visited.add(id(obj)) - - if hasattr(obj, 'load_parameter_state_from_dp_reshardable') or hasattr( - obj, 'load_parameter_state_from_fully_reshardable'): - yield obj - return - - for child in getattr(obj, 'chained_optimizers', []) or []: - yield from visit(child) - for child in getattr(obj, 'sub_optimizers', []) or []: - yield from visit(child) - - yield from visit(optimizer) - - -def _has_mindspeed_patched_load_state_dict(distributed_optimizer): - load_state_dict = getattr(type(distributed_optimizer), 'load_state_dict', None) - return getattr(load_state_dict, '__module__', '').startswith('mindspeed.') - - -_MEGATRON_RESHARDABLE_PARAM_STATE_LOADERS = { - 'dp_reshardable': 'load_parameter_state_from_dp_reshardable', - 'fully_reshardable': 'load_parameter_state_from_fully_reshardable', -} - - -def _current_npu_device(): - if hasattr(torch, 'npu'): - return torch.device('npu', torch.npu.current_device()) - return torch.cuda.current_device() - - -def _restore_mindspeed_optimizer_step_tensors(optimizer): - restored_count = 0 - for param_groups in _iter_optimizer_param_groups(optimizer): - for param_group in param_groups: - step = param_group.get('step') - if isinstance(step, torch.Tensor): - continue - if isinstance(step, (int, float)): - param_group['step'] = torch.tensor(int(step), dtype=torch.int64, device=_current_npu_device()) - restored_count += 1 - if restored_count: - logger.warning(f'Restored {restored_count} MindSpeed optimizer param-group step values to NPU tensors.') - - -def _split_chained_optimizer_state_dict(chained_optimizers, state_dict): - if isinstance(state_dict, dict): - state_dicts = [v for _k, v in sorted(state_dict.items())] - else: - state_dicts = list(state_dict) - if len(chained_optimizers) != len(state_dicts): - raise RuntimeError( - f'Expected {len(chained_optimizers)} entries in optimizer state dict, but got {len(state_dicts)}.') - return state_dicts - - -def _load_chained_optimizer_state_dict(optimizer, state_dict): - chained_optimizers = getattr(optimizer, 'chained_optimizers', None) - if not chained_optimizers or len(chained_optimizers) <= 1: - return False - - state_dicts = _split_chained_optimizer_state_dict(chained_optimizers, state_dict) - for child_optimizer, child_state_dict in zip(chained_optimizers, state_dicts): - load_optimizer_state_dict(child_optimizer, child_state_dict) - synchronize_steps = getattr(optimizer, '_synchronize_steps', None) - if synchronize_steps is not None: - synchronize_steps() - return True - - -def load_optimizer_state_dict(optimizer, state_dict): - if _load_chained_optimizer_state_dict(optimizer, state_dict): - return - - distributed_optimizers = list(_iter_distributed_optimizers(optimizer)) - mindspeed_patched = any( - _has_mindspeed_patched_load_state_dict(distributed_optimizer) - for distributed_optimizer in distributed_optimizers) - if mindspeed_patched: - from .mindspeed import load_mindspeed_fsdp_dtensor_optimizer_state_dict - if load_mindspeed_fsdp_dtensor_optimizer_state_dict(distributed_optimizers, state_dict): - _restore_mindspeed_optimizer_step_tensors(optimizer) - return - - sharding_type = state_dict.get('param_state_sharding_type') if isinstance(state_dict, dict) else None - native_loader_name = _MEGATRON_RESHARDABLE_PARAM_STATE_LOADERS.get(sharding_type) - if native_loader_name is None: - optimizer.load_state_dict(state_dict) - if mindspeed_patched: - _restore_mindspeed_optimizer_step_tensors(optimizer) - return - - if not mindspeed_patched: - optimizer.load_state_dict(state_dict) - return - - if len(distributed_optimizers) != 1: - raise RuntimeError(f'MindSpeed optimizer checkpoint compatibility supports exactly one distributed optimizer, ' - f'got {len(distributed_optimizers)}.') - distributed_optimizer = distributed_optimizers[0] - if not hasattr(distributed_optimizer, native_loader_name): - raise RuntimeError(f'Distributed optimizer does not support sharding type {sharding_type}.') - - state_dict_without_param_state = dict(state_dict) - param_state = state_dict_without_param_state.pop('param_state', None) - state_dict_without_param_state.pop('param_state_sharding_type', None) - if param_state is None: - raise RuntimeError(f'Optimizer checkpoint missing param_state for sharding type {sharding_type}.') - - logger.warning(f'Loading optimizer param_state with ms-swift compatibility path because MindSpeed ' - f'DistributedOptimizer.load_state_dict does not support {sharding_type}.') - # Let MindSpeed restore the generic optimizer state; load the missing - # reshardable param_state with Megatron-Core's native implementation. - optimizer.load_state_dict(state_dict_without_param_state) - _restore_mindspeed_optimizer_step_tensors(optimizer) - getattr(distributed_optimizer, native_loader_name)(param_state) - - -__all__ = ['load_optimizer_state_dict', 'optimizer_sharded_state_dict'] +__all__ = ['optimizer_sharded_state_dict'] diff --git a/swift/model/npu_patch/megatron_fsdp.py b/swift/model/npu_patch/megatron_fsdp.py new file mode 100644 index 0000000000..138e3bfc1f --- /dev/null +++ b/swift/model/npu_patch/megatron_fsdp.py @@ -0,0 +1,294 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""NPU-only Megatron-FSDP optimizer compatibility. + +MCore's GPU implementation has two valid optimizer contracts for +Megatron-FSDP: Transformer Engine's precision-aware FusedAdam, or the native +PyTorch AdamW fallback when the fused optimizer is unavailable. The +TransformerEngineNPU FusedAdam currently rejects DTensor parameters, so the +NPU path implements the latter contract locally. The adapter below accepts +the small TE/Apex constructor surface that MCore uses, then delegates the +actual update and state handling to ``torch.optim.AdamW``. + +The patch is requested only by NPU Megatron-FSDP, but changes process-global +MCore optimizer aliases. Mixing FSDP and non-FSDP optimizers in the same +process is not supported by this adapter. GPU-only processes are unchanged. +""" + +from __future__ import annotations + +import copy +import torch + +from swift.utils import get_logger + +logger = get_logger() + +# 1. NPU Megatron-FSDP optimizer adaptation. + +_APPLIED = False + + +class NPUFSDPAdamW(torch.optim.AdamW): + """Native AdamW with MCore's FusedAdam-compatible constructor surface. + + MCore selects the optimizer class through module-level aliases. When TE + is installed, its normal FusedAdam alias receives arguments that native + AdamW does not understand (``adam_w_mode``, ``bias_correction`` and the + precision-aware state options). This class makes that boundary explicit + instead of relying on a bare alias to ``torch.optim.AdamW``. + + Only the standard, non-precision-aware FSDP contract is supported. A + precision-aware request must fail in ``patch_megatron_fsdp_optimizer``; + silently dropping master weights or decoupled gradients would change the + optimizer semantics. + """ + + def __init__( + self, + params, + lr=1e-3, + betas=(0.9, 0.999), + eps=1e-8, + weight_decay=1e-2, + amsgrad=False, + *, + adam_w_mode=True, + bias_correction=True, + capturable=False, + use_decoupled_grad=False, + master_weights=False, + master_weight_dtype=torch.float32, + exp_avg_dtype=None, + exp_avg_sq_dtype=None, + store_param_remainders=False, + fused=None, + set_grad_none=None, + **kwargs, + ): + if not adam_w_mode: + raise ValueError('NPU Megatron-FSDP requires AdamW decoupled weight decay.') + if not bias_correction: + raise ValueError('NPU Megatron-FSDP native AdamW requires bias correction.') + if use_decoupled_grad: + raise ValueError('NPU Megatron-FSDP native AdamW does not support decoupled_grad; ' + 'disable precision-aware optimizer.') + if master_weights or store_param_remainders: + raise ValueError('NPU Megatron-FSDP native AdamW does not own TE master weights or parameter remainders.') + if exp_avg_dtype is not None or exp_avg_sq_dtype is not None: + raise ValueError('NPU Megatron-FSDP native AdamW does not support precision-aware optimizer state dtypes.') + if set_grad_none is not None: + raise ValueError('NPU Megatron-FSDP native AdamW uses zero_grad(set_to_none=...); ' + 'set_grad_none is not a constructor option.') + if kwargs: + unknown = ', '.join(sorted(kwargs)) + raise TypeError(f'Unsupported NPU Megatron-FSDP AdamW arguments: {unknown}') + + # ``fused`` is a TE/Apex hint. The NPU implementation intentionally + # uses torch's ordinary AdamW kernels; dropping this hint is safe for + # the standard FSDP path and avoids asking torch_npu for CUDA fusion. + del fused, master_weight_dtype + super().__init__( + params, + lr=lr, + betas=betas, + eps=eps, + weight_decay=weight_decay, + amsgrad=amsgrad, + capturable=capturable, + ) + + +def patch_megatron_fsdp_optimizer(*, use_precision_aware_optimizer: bool = False) -> None: + """Use native Torch AdamW for NPU Megatron-FSDP DTensor parameters. + + MCore keeps the optimizer class and backend-selection flags in both + ``megatron.core.optimizer`` and ``megatron.core.optimizer.distrib_optimizer`` + module globals. Updating both modules is necessary because the latter + performs its own ``Adam`` import and type checks. + + Precision-aware Adam is intentionally not silently downgraded: the NPU + FSDP compatibility path currently targets the standard BF16/FP32 AdamW + contract. Callers using precision-aware optimizer settings receive a + clear error instead of a partially compatible optimizer. + """ + + global _APPLIED + if use_precision_aware_optimizer: + raise ValueError('NPU Megatron-FSDP currently requires use_precision_aware_optimizer=False; ' + 'the NPU FSDP compatibility path uses native torch.optim.AdamW for DTensor parameters.') + if _APPLIED: + return + + import megatron.core.optimizer as mcore_optimizer + import megatron.core.optimizer.distrib_optimizer as distributed_optimizer + + # MCore's default import prefers TE whenever it is installed. Switch only + # the already-loaded optimizer aliases used to build the FSDP optimizer; + # TE model/attention/linear callables remain untouched. Mark the alias as + # a non-TE optimizer so MCore passes the FusedAdam-compatible constructor + # arguments to our explicit adapter and uses its native DTensor state + # handling in DistributedOptimizer. + mcore_optimizer.Adam = NPUFSDPAdamW + mcore_optimizer.USING_PYTORCH_OPTIMIZER = False + distributed_optimizer.Adam = NPUFSDPAdamW + distributed_optimizer.HAVE_APEX_OR_TE = False + distributed_optimizer.USING_TE_OPTIMIZER = False + distributed_optimizer.USING_APEX_OPTIMIZER = False + _APPLIED = True + logger.info('NPU Megatron-FSDP: using Swift NPUFSDPAdamW (native torch.optim.AdamW) for DTensor parameters.') + + +# 2. Refresh stale optimizer steps on empty local shards. + + +def _optimizer_step_value(step, param_name): + if isinstance(step, torch.Tensor): + if step.numel() != 1: + raise RuntimeError( + f'Megatron-FSDP optimizer step for `{param_name}` must be scalar, got shape {tuple(step.shape)}.') + return step.detach().item() + if isinstance(step, (int, float)): + return step + raise RuntimeError(f'Megatron-FSDP optimizer step for `{param_name}` has unsupported type {type(step).__name__}.') + + +def _clone_optimizer_step(step, existing_step=None): + if isinstance(step, torch.Tensor): + cloned_step = step.detach().clone() + if isinstance(existing_step, torch.Tensor): + cloned_step = cloned_step.to(device=existing_step.device, dtype=existing_step.dtype) + return cloned_step + return copy.deepcopy(step) + + +def _refresh_empty_optimizer_step(param_state, param_name, canonical_step, canonical_step_value): + existing_step = param_state.get('step') + if existing_step is None or _optimizer_step_value(existing_step, param_name) != canonical_step_value: + param_state['step'] = _clone_optimizer_step(canonical_step, existing_step) + return 1 + return 0 + + +# 3. Complete empty optimizer shards for DCP save/load. + + +def _get_fsdp_model_parameter(model, param_name): + model_param_name = param_name[len('module.'):] if param_name.startswith('module.') else param_name + num_experts = getattr(getattr(model, 'config', None), 'num_moe_experts', None) + if num_experts: + from megatron.core.transformer.fsdp_dtensor_checkpoint import expert_param_local_key + model_param_name = expert_param_local_key(model_param_name, num_experts) + return model.get_parameter(model_param_name) + + +def _build_empty_optimizer_state(state_template, template_param, dist_param, param_name): + missing_state = {} + for state_key, template_value in state_template.items(): + if not isinstance(template_value, torch.Tensor): + missing_state[state_key] = copy.deepcopy(template_value) + continue + + if state_key == 'step': + missing_state[state_key] = template_value.detach().clone() + continue + + if tuple(template_value.shape) != tuple(template_param.shape): + raise RuntimeError( + f'Cannot infer the empty Megatron-FSDP optimizer state `{state_key}` for `{param_name}`: ' + f'template state shape {tuple(template_value.shape)} does not match its parameter shape ' + f'{tuple(template_param.shape)}.') + missing_state[state_key] = torch.zeros_like(dist_param, dtype=template_value.dtype) + return missing_state + + +def complete_npu_fsdp_dtensor_optimizer_state(state_dict, model) -> None: + """Complete and refresh empty AdamW shards for consistent DCP state. + + MCore skips the dummy optimizer step for empty local DTensor shards. Native + ``torch.optim.AdamW`` therefore has no state entry for those parameters, + while Torch DCP requires a consistent optimizer key set across DP ranks. + + An empty shard created during checkpoint load remains in the live optimizer + state afterwards, but AdamW never visits it and therefore never increments + its scalar ``step``. Refresh those placeholders from a non-empty local + shard before every save so a save-resume-save chain does not persist a + stale optimizer step for the corresponding global parameter. + + This assumes synchronized updates across all optimized parameters and a + uniform AdamW state layout/dtype. Conditional updates with unequal parameter + steps or heterogeneous optimizer state layouts are not supported here. + """ + optimizer_state_dict = state_dict.get('optimizer') + if not optimizer_state_dict: + return + optimizer_state = optimizer_state_dict.get('state', {}) + param_to_group_meta = optimizer_state_dict.get('param_to_group_meta', {}) + if not param_to_group_meta: + return + if not optimizer_state: + raise RuntimeError('Cannot infer Megatron-FSDP optimizer state fields from an empty local state dict.') + + parameter_info = {} + active_steps = [] + template_param_name = None + for param_name in param_to_group_meta: + dist_param = _get_fsdp_model_parameter(model, param_name) + local_param = dist_param.to_local() if hasattr(dist_param, 'to_local') else dist_param + parameter_info[param_name] = (dist_param, local_param) + param_state = optimizer_state.get(param_name) + if local_param.numel() == 0 or not param_state: + continue + if 'step' not in param_state: + raise RuntimeError(f'Megatron-FSDP AdamW state for `{param_name}` is missing scalar `step`.') + active_steps.append((param_name, param_state['step'], _optimizer_step_value(param_state['step'], param_name))) + if template_param_name is None: + template_param_name = param_name + + if template_param_name is None: + raise RuntimeError('Cannot match a non-empty Megatron-FSDP optimizer state template to a local parameter.') + + canonical_step_name, canonical_step, canonical_step_value = active_steps[0] + inconsistent_steps = [(name, value) for name, _step, value in active_steps if value != canonical_step_value] + if inconsistent_steps: + details = ', '.join(f'{name}={value}' for name, value in inconsistent_steps[:8]) + raise RuntimeError(f'Inconsistent Megatron-FSDP optimizer steps on non-empty local shards: ' + f'{canonical_step_name}={canonical_step_value}; {details}.') + + state_template = optimizer_state[template_param_name] + template_param = _get_fsdp_model_parameter(model, template_param_name) + completed_state = {} + added_count = 0 + refreshed_count = 0 + for param_name in param_to_group_meta: + dist_param, local_param = parameter_info[param_name] + if optimizer_state.get(param_name): + param_state = optimizer_state[param_name] + if local_param.numel() == 0: + refreshed_count += _refresh_empty_optimizer_step(param_state, param_name, canonical_step, + canonical_step_value) + completed_state[param_name] = param_state + continue + + if local_param.numel() != 0: + raise RuntimeError( + f'Megatron-FSDP optimizer state is missing or empty for `{param_name}`, but its local parameter shard ' + f'is not empty (numel={local_param.numel()}).') + completed_state[param_name] = _build_empty_optimizer_state( + state_template, + template_param, + dist_param, + param_name, + ) + added_count += 1 + + completed_state.update({key: value for key, value in optimizer_state.items() if key not in completed_state}) + optimizer_state_dict['state'] = completed_state + if added_count or refreshed_count: + logger.info( + 'Completed NPU Megatron-FSDP optimizer state: added %d empty local shards and refreshed %d stale steps.', + added_count, + refreshed_count, + ) + + +__all__ = ['NPUFSDPAdamW', 'complete_npu_fsdp_dtensor_optimizer_state', 'patch_megatron_fsdp_optimizer'] diff --git a/swift/model/npu_patch/mindspeed.py b/swift/model/npu_patch/mindspeed.py deleted file mode 100644 index 2c6b73e0eb..0000000000 --- a/swift/model/npu_patch/mindspeed.py +++ /dev/null @@ -1,537 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -from __future__ import annotations - -import copy -import importlib -import inspect -import sys -import torch -from functools import wraps -from types import ModuleType -from typing import Any - -from swift.utils.logger import get_logger - -logger = get_logger() - -_ORIGINAL_MINDSPEED_TE_CP_CLASS = None -_ORIGINAL_MINDSPEED_GDN = None -_FLA_GDN_PATCH_TARGET = 'fla.ops.gated_delta_rule.chunk_gated_delta_rule' -_MINDSPEED_FSDP_GRAD_REDUCE_MODULE = 'mindspeed.core.distributed.custom_fsdp.param_and_grad_buffer' - - -def _wrap_mindspeed_fsdp_gradient_reduce(original_gradient_reduce): - """Treat a missing Megatron-FSDP gradient scale as the multiplicative identity.""" - - @wraps(original_gradient_reduce) - def gradient_reduce_with_none_scaling(self, bucket_group, *args, **kwargs): - restored_buffers = [] - seen_buffers = set() - for bucket_id in bucket_group: - grad_buffer = self.get_fsdp_buffer(bucket_id) - buffer_id = id(grad_buffer) - if buffer_id in seen_buffers or grad_buffer.gradient_scaling_factor is not None: - continue - seen_buffers.add(buffer_id) - original_ddp_config = grad_buffer.ddp_config - restored_buffers.append((grad_buffer, original_ddp_config)) - grad_buffer.gradient_scaling_factor = 1.0 - - # `None` means SUM without any pre-scaling. Keep that semantic even if - # average-in-collective is enabled, without changing the shared config - # subsequently used by the outer-HSDP reduction. - if original_ddp_config.average_in_collective: - grad_buffer.ddp_config = copy.copy(original_ddp_config) - grad_buffer.ddp_config.average_in_collective = False - - try: - return original_gradient_reduce(self, bucket_group, *args, **kwargs) - finally: - for grad_buffer, original_ddp_config in restored_buffers: - grad_buffer.gradient_scaling_factor = None - grad_buffer.ddp_config = original_ddp_config - - gradient_reduce_with_none_scaling._swift_handles_none_gradient_scaling = True - return gradient_reduce_with_none_scaling - - -def patch_mindspeed_megatron_fsdp_gradient_scaling(megatron_args: dict[str, Any]) -> None: - """Patch the MindSpeed 0.16 Megatron-FSDP reducer that multiplies by ``None``.""" - if not megatron_args.get('use_megatron_fsdp', False): - return - - grad_buffer_module = importlib.import_module( - 'megatron.core.distributed.fsdp.src.megatron_fsdp.param_and_grad_buffer') - grad_reduce = grad_buffer_module.GradReducePipeline._bucket_group_gradient_reduce - if getattr(grad_reduce, '_swift_handles_none_gradient_scaling', False): - return - if grad_reduce.__module__ != _MINDSPEED_FSDP_GRAD_REDUCE_MODULE: - logger.info('MindSpeed Megatron-FSDP gradient reducer is not active; skip the None-scaling patch.') - return - - try: - source = inspect.getsource(grad_reduce) - except (OSError, TypeError) as e: - logger.warning('Cannot inspect the active MindSpeed Megatron-FSDP gradient reducer; skip patch: %s', e) - return - has_unconditional_scale = 'bucket.data.mul_(scaling_factor)' in source - has_none_guard = 'if scaling_factor is None' in source or 'if scaling_factor is not None' in source - if has_none_guard: - logger.info('MindSpeed Megatron-FSDP gradient reducer already handles None scaling; skip patch.') - return - if not has_unconditional_scale: - logger.warning('The active MindSpeed Megatron-FSDP gradient reducer no longer matches the expected ' - 'None-scaling implementation; skip patch and verify the reducer compatibility.') - return - - grad_buffer_module.GradReducePipeline._bucket_group_gradient_reduce = _wrap_mindspeed_fsdp_gradient_reduce( - grad_reduce) - logger.info('Patched MindSpeed Megatron-FSDP gradient reduction to handle scaling_factor=None.') - - -def _get_mindspeed_fsdp_model_parameter(model, param_name): - model_param_name = param_name[len('module.'):] if param_name.startswith('module.') else param_name - num_experts = getattr(getattr(model, 'config', None), 'num_moe_experts', None) - if num_experts: - from megatron.core.transformer.fsdp_dtensor_checkpoint import expert_param_local_key - model_param_name = expert_param_local_key(model_param_name, num_experts) - return model.get_parameter(model_param_name) - - -def _build_empty_optimizer_state(state_template, template_param, dist_param, param_name): - missing_state = {} - for state_key, template_value in state_template.items(): - if not isinstance(template_value, torch.Tensor): - missing_state[state_key] = copy.deepcopy(template_value) - continue - - if template_value.ndim == 0 or state_key == 'step': - missing_state[state_key] = template_value.detach().clone() - continue - - if tuple(template_value.shape) != tuple(template_param.shape): - raise RuntimeError( - f'Cannot infer the empty Megatron-FSDP optimizer state `{state_key}` for `{param_name}`: ' - f'template state shape {tuple(template_value.shape)} does not match its parameter shape ' - f'{tuple(template_param.shape)}.') - missing_state[state_key] = torch.zeros_like(dist_param, dtype=template_value.dtype) - return missing_state - - -def complete_mindspeed_fsdp_dtensor_optimizer_state(state_dict, model) -> None: - """Add optimizer shards omitted by MindSpeed so all DP ranks use the same state keys.""" - optimizer_state_dict = state_dict.get('optimizer') - if not optimizer_state_dict: - return - optimizer_state = optimizer_state_dict.get('state', {}) - param_to_group_meta = optimizer_state_dict.get('param_to_group_meta', {}) - if not param_to_group_meta: - return - if set(param_to_group_meta).issubset(optimizer_state): - optimizer_state_dict['state'] = { - **{ - key: optimizer_state[key] - for key in param_to_group_meta - }, - **{ - key: value - for key, value in optimizer_state.items() if key not in param_to_group_meta - }, - } - return - if not optimizer_state: - raise RuntimeError('Cannot infer Megatron-FSDP optimizer state fields from an empty local state dict.') - - template_param_name = next( - (name for name in param_to_group_meta if name in optimizer_state and optimizer_state[name]), None) - if template_param_name is None: - raise RuntimeError('Cannot match a Megatron-FSDP optimizer state template to a model parameter.') - state_template = optimizer_state[template_param_name] - template_param = _get_mindspeed_fsdp_model_parameter(model, template_param_name) - completed_state = {} - added_count = 0 - for param_name in param_to_group_meta: - if param_name in optimizer_state: - completed_state[param_name] = optimizer_state[param_name] - continue - - dist_param = _get_mindspeed_fsdp_model_parameter(model, param_name) - local_param = dist_param.to_local() if hasattr(dist_param, 'to_local') else dist_param - if local_param.numel() != 0: - raise RuntimeError( - f'MindSpeed omitted Megatron-FSDP optimizer state for `{param_name}`, but its local parameter shard ' - f'is not empty (numel={local_param.numel()}). Refusing to replace a non-empty optimizer shard with ' - f'zeros.') - completed_state[param_name] = _build_empty_optimizer_state( - state_template, - template_param, - dist_param, - param_name, - ) - added_count += 1 - - completed_state.update({key: value for key, value in optimizer_state.items() if key not in completed_state}) - optimizer_state_dict['state'] = completed_state - logger.info('Added %d empty local optimizer shards for MindSpeed Megatron-FSDP checkpointing.', added_count) - - -def load_mindspeed_fsdp_dtensor_optimizer_state_dict(distributed_optimizers, state_dict) -> bool: - """Restore an FSDP DTensor optimizer state without MindSpeed's legacy checkpoint loader.""" - is_fsdp_dtensor_state = isinstance(state_dict, - dict) and 'state' in state_dict and 'param_to_group_meta' in state_dict - if not is_fsdp_dtensor_state: - return False - if len(distributed_optimizers) != 1: - raise RuntimeError(f'MindSpeed FSDP DTensor optimizer compatibility supports exactly one distributed ' - f'optimizer, got {len(distributed_optimizers)}.') - - distributed_optimizer = distributed_optimizers[0] - logger.warning('Loading FSDP DTensor optimizer state with the Megatron-Core compatibility path because ' - 'MindSpeed DistributedOptimizer.load_state_dict expects the legacy checkpoint structure.') - inner_state_dict = dict(state_dict) - inner_state_dict['param_groups'] = distributed_optimizer._param2group_meta_to_param_groups( - inner_state_dict.pop('param_to_group_meta'), distributed_optimizer.optimizer.param_groups) - distributed_optimizer.optimizer.load_state_dict(inner_state_dict) - return True - - -def _mindspeed_gdn_with_safe_varlen(q, - k, - v, - g, - beta, - scale=None, - initial_state=None, - output_final_state=False, - use_qk_l2norm_in_kernel=False, - cu_seqlens=None, - chunk_size=64, - head_first=False): - kwargs = { - 'scale': scale, - 'output_final_state': output_final_state, - 'use_qk_l2norm_in_kernel': use_qk_l2norm_in_kernel, - 'chunk_size': chunk_size, - 'head_first': head_first, - } - if cu_seqlens is None: - return _ORIGINAL_MINDSPEED_GDN(q, k, v, g, beta, initial_state=initial_state, **kwargs) - - # MindSpeed's arch35 varlen backward uses the local sequence length as the packed gate stride. - # Keep the same implementation but run each sequence independently to avoid the invalid indexing. - import torch - sequence_dim = 2 if head_first else 1 - offsets = cu_seqlens.detach().cpu().tolist() - outputs, final_states = [], [] - for i, (start, end) in enumerate(zip(offsets, offsets[1:])): - length = end - start - inputs = [x.narrow(sequence_dim, start, length) for x in (q, k, v, g, beta)] - state = None if initial_state is None else initial_state[i:i + 1] - output, final_state = _ORIGINAL_MINDSPEED_GDN(*inputs, initial_state=state, **kwargs) - outputs.append(output) - if output_final_state: - final_states.append(final_state) - output = torch.cat(outputs, dim=sequence_dim) - final_state = torch.cat(final_states) if output_final_state else None - return output, final_state - - -def prepare_mindspeed_gdn_import() -> None: - try: - import fla.utils - except ModuleNotFoundError as e: - if e.name not in {'fla', 'fla.utils'}: - raise - gdn_module = ModuleType('mindspeed.core.ssm.chunk_gated_delta_rule') - - def torch_chunk_gated_delta_rule(q, - k, - v, - g, - beta, - scale=None, - initial_state=None, - output_final_state=False, - use_qk_l2norm_in_kernel=False, - cu_seqlens=None, - chunk_size=64, - head_first=False, - **kwargs): - if cu_seqlens is not None: - raise ValueError('Torch GDN fallback does not support cu_seqlens.') - from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import torch_chunk_gated_delta_rule as torch_gdn - return torch_gdn( - q, - k, - v, - g=g, - beta=beta, - chunk_size=chunk_size, - initial_state=initial_state, - output_final_state=output_final_state, - use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, - ) - - gdn_module.chunk_gated_delta_rule = torch_chunk_gated_delta_rule - gdn_module._ms_swift_torch_fallback = True - sys.modules[gdn_module.__name__] = gdn_module - else: - import torch_npu - device_name = torch_npu.npu.get_device_name() - # MindSpeed still imports this flag after it was removed from upstream FLA. - if not hasattr(fla.utils, 'USE_CUDA_GRAPH'): - if 'Ascend910_95' in device_name or 'Ascend950' in device_name: - fla.utils.USE_CUDA_GRAPH = False - - -def _apply_gdn_patch(MindSpeedPatchesManager, patch, implementation) -> None: - if patch is not None: - MindSpeedPatchesManager.register_patch( - _FLA_GDN_PATCH_TARGET, - implementation, - force_patch=True, - ) - MindSpeedPatchesManager.apply_patches() - else: - try: - fla_gated_delta_rule = importlib.import_module('fla.ops.gated_delta_rule') - except Exception: - pass - else: - fla_gated_delta_rule.chunk_gated_delta_rule = implementation - - # mcore-bridge may have cached the callable before a runtime repatch. - bridge_gdn = sys.modules.get('mcore_bridge.model.modules.gated_delta_net') - if bridge_gdn is not None: - bridge_gdn.chunk_gated_delta_rule = implementation - - if patch is not None: - fla_gated_delta_rule = importlib.import_module('fla.ops.gated_delta_rule') - if fla_gated_delta_rule.chunk_gated_delta_rule is not implementation: - raise RuntimeError('MindSpeed did not install the selected Megatron GDN implementation.') - if bridge_gdn is not None and bridge_gdn.chunk_gated_delta_rule is not implementation: - raise RuntimeError('Failed to refresh the mcore-bridge cached GDN implementation.') - - -def _patch_mindspeed_fla_gdn_implementation(MindSpeedPatchesManager) -> None: - patch = MindSpeedPatchesManager.patches_info.get(_FLA_GDN_PATCH_TARGET) - - mindspeed_gdn_module = sys.modules.get('mindspeed.core.ssm.chunk_gated_delta_rule') - if getattr(mindspeed_gdn_module, '_ms_swift_torch_fallback', False): - torch_gdn = mindspeed_gdn_module.chunk_gated_delta_rule - _apply_gdn_patch(MindSpeedPatchesManager, patch, torch_gdn) - logger.info('Using torch chunk_gated_delta_rule for Megatron GDN because FLA is unavailable.') - return - - import torch_npu - device_name = torch_npu.npu.get_device_name() - if 'Ascend910_95' in device_name or 'Ascend950' in device_name: - from mindspeed.core.ssm.chunk_gated_delta_rule import chunk_gated_delta_rule as mindspeed_gdn - global _ORIGINAL_MINDSPEED_GDN - if _ORIGINAL_MINDSPEED_GDN is None: - _ORIGINAL_MINDSPEED_GDN = mindspeed_gdn - _apply_gdn_patch(MindSpeedPatchesManager, patch, _mindspeed_gdn_with_safe_varlen) - logger.info( - 'Using MindSpeed chunk_gated_delta_rule with safe varlen fallback for Megatron GDN on Ascend arch35.') - return - - fla_error = None - if (patch is not None and patch.orig_func is not None and patch.orig_func.__module__.startswith('fla.')): - # MindSpeed propagates its replacement into already imported submodules, - # so importing from ``fla.ops.gated_delta_rule.chunk`` again is not enough. - fla_chunk_gated_delta_rule = patch.orig_func - else: - try: - from fla.ops.gated_delta_rule.chunk import chunk_gated_delta_rule as fla_chunk_gated_delta_rule - except Exception as e: - fla_chunk_gated_delta_rule = None - fla_error = e - - if fla_chunk_gated_delta_rule is not None: - try: - if not fla_chunk_gated_delta_rule.__module__.startswith('fla.'): - raise RuntimeError('resolved a non-FLA callable: ' - f'{fla_chunk_gated_delta_rule.__module__}.' - f'{fla_chunk_gated_delta_rule.__name__}') - _apply_gdn_patch(MindSpeedPatchesManager, patch, fla_chunk_gated_delta_rule) - logger.info( - 'Using upstream FLA chunk_gated_delta_rule for Megatron GDN: module=%s, source=%s.', - fla_chunk_gated_delta_rule.__module__, - inspect.getsourcefile(inspect.unwrap(fla_chunk_gated_delta_rule)), - ) - return - except Exception as e: - fla_error = e - - logger.warning( - 'FLA GDN is unavailable (%s); keep the current MindSpeed/Megatron GDN implementation unchanged. ' - 'If it does not support packed cu_seqlens, the GDN call will fail at runtime.', - fla_error, - ) - - -def patch_mindspeed_fla_gdn_implementation() -> None: - """Use torch GDN without FLA, MindSpeed GDN on arch35, and upstream FLA elsewhere.""" - from mindspeed.patch_utils import MindSpeedPatchesManager - - try: - _patch_mindspeed_fla_gdn_implementation(MindSpeedPatchesManager) - except Exception as e: - logger.warning('Failed to apply the optional FLA GDN patch; keep the current implementation: %s', e) - - -def patch_mindspeed_te_cp_implementation(megatron_args: dict[str, Any]) -> None: - """ - Route NPU CP to the legacy MindSpeed TE adaptor when the new strategy factory - only supports kvallgather. - """ - # MindSpeed 0.15.3 replaced the TE context-parallel attention class with a - # new implementation. That new class does not yet cover all CP algorithms, - # so the default non-kvallgather path can fail during Megatron training. - # For those algorithms, temporarily route TE attention back to the legacy - # MindSpeedCPDotProductAttention adaptor. Once MindSpeed's new CP class has - # feature parity, this compatibility patch can be removed. - try: - import mindspeed.te.pytorch.attention.dot_product_attention.dot_product_attention as ms_te_dpa - from mindspeed.core.context_parallel.adaptor import MindSpeedCPDotProductAttention - except ImportError as e: - logger.warning(f'Failed to import MindSpeed CP modules before repatch: {e}') - return - - global _ORIGINAL_MINDSPEED_TE_CP_CLASS - if _ORIGINAL_MINDSPEED_TE_CP_CLASS is None: - _ORIGINAL_MINDSPEED_TE_CP_CLASS = getattr(ms_te_dpa, 'MindSpeedTEDotProductAttention', None) - - if _ORIGINAL_MINDSPEED_TE_CP_CLASS is None: - logger.warning('MindSpeedTEDotProductAttention is unavailable before repatch; skip CP workaround.') - return - - cp_algo = megatron_args.get('context_parallel_algo', 'megatron_cp_algo') - use_legacy_cp_te = int(megatron_args.get('context_parallel_size', 1)) > 1 and cp_algo != 'kvallgather_cp_algo' - target_cls = MindSpeedCPDotProductAttention if use_legacy_cp_te else _ORIGINAL_MINDSPEED_TE_CP_CLASS - - if getattr(ms_te_dpa, 'MindSpeedTEDotProductAttention', None) is target_cls: - return - - ms_te_dpa.MindSpeedTEDotProductAttention = target_cls - logger.info( - 'Patched MindSpeedTEDotProductAttention to %s for context_parallel_size=%s, context_parallel_algo=%s.', - target_cls.__name__, - megatron_args.get('context_parallel_size', 1), - cp_algo, - ) - - -def patch_mindspeed_te_layernorm_linear_frozen_weight() -> None: - """Route frozen MindSpeed TE LayerNormLinear weights through Megatron's frozen-weight path.""" - try: - ms_te_layernorm_linear = importlib.import_module('mindspeed.te.pytorch.module.layernorm_column_parallel_linear') - from megatron.core.tensor_parallel.layers import linear_with_frozen_weight - except ImportError as e: - logger.warning('Failed to import MindSpeed TE LayerNormLinear modules: %s', e) - return - - linear_impl_name = 'linear_with_grad_accumulation_and_async_allreduce' - trainable_weight_impl = getattr(ms_te_layernorm_linear, linear_impl_name, None) - if trainable_weight_impl is None: - logger.warning('MindSpeed TE LayerNormLinear does not expose %s; skip frozen-weight patch.', linear_impl_name) - return - if getattr(trainable_weight_impl, '_swift_supports_frozen_weight', False): - return - - @wraps(trainable_weight_impl) - def linear_with_frozen_weight_dispatch( - input, - weight, - bias, - gradient_accumulation_fusion, - allreduce_dgrad, - sequence_parallel, - grad_output_buffer=None, - wgrad_deferral_limit=0, - async_grad_allreduce=None, - tp_group=None, - ): - if weight.requires_grad: - return trainable_weight_impl( - input=input, - weight=weight, - bias=bias, - gradient_accumulation_fusion=gradient_accumulation_fusion, - allreduce_dgrad=allreduce_dgrad, - sequence_parallel=sequence_parallel, - grad_output_buffer=grad_output_buffer, - wgrad_deferral_limit=wgrad_deferral_limit, - async_grad_allreduce=async_grad_allreduce, - tp_group=tp_group, - ) - return linear_with_frozen_weight( - input=input, - weight=weight, - bias=bias, - gradient_accumulation_fusion=gradient_accumulation_fusion, - allreduce_dgrad=allreduce_dgrad, - sequence_parallel=sequence_parallel, - async_grad_allreduce=async_grad_allreduce, - tp_group=tp_group, - ) - - linear_with_frozen_weight_dispatch._swift_supports_frozen_weight = True - setattr(ms_te_layernorm_linear, linear_impl_name, linear_with_frozen_weight_dispatch) - logger.info('Patched MindSpeed TE LayerNormLinear to use Megatron frozen-weight backward for frozen weights.') - - -def patch_mindspeed_te_grouped_linear_save_original_input() -> None: - """Allow BF16 MindSpeed grouped linear layers to use Megatron activation offloading.""" - te_extension = importlib.import_module('megatron.core.extensions.transformer_engine') - set_save_original_input = te_extension.set_save_original_input - if getattr(set_save_original_input, '_swift_supports_mindspeed_bf16_grouped_linear', False): - return - - @wraps(set_save_original_input) - def set_save_original_input_dispatch(module): - config = getattr(module, 'config', None) - is_mindspeed_grouped_linear = type(module).__module__ == 'mindspeed.te.pytorch.module.grouped_linear' - uses_quantized_tensors = bool(getattr(config, 'fp8', None) or getattr(config, 'fp4', None)) - if is_mindspeed_grouped_linear and not uses_quantized_tensors: - # MindSpeed's BF16 TEGroupedLinearGMM saves the original input tensor directly. - return - return set_save_original_input(module) - - set_save_original_input_dispatch._swift_supports_mindspeed_bf16_grouped_linear = True - te_extension.set_save_original_input = set_save_original_input_dispatch - logger.info('Patched Megatron set_save_original_input for MindSpeed BF16 grouped linear layers.') - - -def patch_mindspeed_gdn_cp_helpers(megatron_args: dict[str, Any]) -> None: - """Expose MindSpeed's backported GDN CP helpers to Megatron Core versions before 0.18.""" - if int(megatron_args.get('context_parallel_size', 1)) <= 1: - return - - megatron_gdn = importlib.import_module('megatron.core.ssm.gated_delta_net') - mindspeed_gdn = importlib.import_module('mindspeed.core.ssm.gated_delta_net') - patched_helpers = [] - for helper_name in ('tensor_a2a_cp2hp', 'tensor_a2a_hp2cp'): - if hasattr(megatron_gdn, helper_name): - continue - helper = getattr(mindspeed_gdn, helper_name, None) - if helper is None: - raise RuntimeError(f'MindSpeed does not provide the required GDN CP helper: {helper_name}.') - setattr(megatron_gdn, helper_name, helper) - patched_helpers.append(helper_name) - - if patched_helpers: - logger.info('Patched Megatron GDN CP helpers from MindSpeed: %s.', ', '.join(patched_helpers)) - - -def apply_mindspeed_patches(megatron_args: dict[str, Any]) -> None: - """Apply MindSpeed compatibility patches around its runtime repatch in the required order.""" - from mindspeed.megatron_adaptor import repatch - - patch_mindspeed_te_cp_implementation(megatron_args) - repatch(megatron_args) - patch_mindspeed_megatron_fsdp_gradient_scaling(megatron_args) - patch_mindspeed_gdn_cp_helpers(megatron_args) - patch_mindspeed_te_layernorm_linear_frozen_weight() - patch_mindspeed_te_grouped_linear_save_original_input() - patch_mindspeed_fla_gdn_implementation() diff --git a/swift/model/npu_patcher.py b/swift/model/npu_patcher.py index 4ea78eaf6d..4914c28883 100644 --- a/swift/model/npu_patcher.py +++ b/swift/model/npu_patcher.py @@ -1,8 +1,8 @@ # Copyright (c) ModelScope Contributors. All rights reserved. from __future__ import annotations -from .npu_patch import NPUCastError, apply_all_patches, apply_mindspeed_patches, patch_mindspeed_fla_gdn_implementation +from .npu_patch import NPUCastError, apply_all_patches apply_all_patches() -__all__ = ['NPUCastError', 'apply_all_patches', 'apply_mindspeed_patches', 'patch_mindspeed_fla_gdn_implementation'] +__all__ = ['NPUCastError', 'apply_all_patches']