CANN 9.0.0 NPU 官方镜像中 MindSpeed 与 Megatron 版本冲突，导致 MoE 模型训练无法启动

Checklist / 检查清单

• I have searched existing issues, and this is a new bug report. / 我已经搜索过现有的 issues，确认这是一个新的 bug report。

Bug Description / Bug 描述

在使用官方 Docker 镜像docker pull quay.io/ascend/ms-swift:v4.3.0-A2-py311-CANN9.0.0-ubuntu22.04在华为昇腾 NPU 910b集群上进行 MoE 模型的 Megatron SFT 训练时，由于镜像内部同时存在两个不兼容的 megatron 版本，导致训练因模块缺失或导入失败而中断。

环境变量 PYTHONPATH 中同时包含 /Megatron-LM（旧版）和 site-packages 中的 megatron_core-0.16.1（新版）。这导致了的依赖冲突：优先导入旧版会导致 mcore-bridge 报错；强制导入新版则会导致 MindSpeed 报错。

export PATH=/usr/local/python3.11.15/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/python3.11.15/lib:/usr/local/Ascend/cann-9.0.0/lib64:/usr/local/Ascend/cann-9.0.0/lib64/plugin/opskernel:/usr/local/Ascend/cann-9.0.0/lib64/plugin/nnengine

export PYTHONPATH=/usr/local/Ascend/cann-9.0.0/python/site-packages:/usr/local/Ascend/cann-9.0.0/opp/built-in/op_impl/ai_core/tbe:/ms-swift:/mcore-bridge/src:/Megatron-LM:$PYTHONPATH

ASCEND 环境变量（容器内路径）

export ASCEND_OPP_PATH=/usr/local/Ascend/cann-9.0.0/opp
export ASCEND_AICPU_PATH=/usr/local/Ascend/cann-9.0.0
export ASCEND_HOME_PATH=/usr/local/Ascend/cann-9.0.0
export ASCEND_TOOLKIT_HOME=/usr/local/Ascend/cann-9.0.0
export ASCEND_TOOLKIT_LATEST_HOME=/usr/local/Ascend/cann-9.0.0

Megatron-LM

export MEGATRON_LM_PATH=/Megatron-LM
export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
export USE_MCORE_GDN=1

Triton 缓存

rm -rf /tmp/triton_cache_*
export TRITON_CACHE_DIR=/tmp/triton_cache_${USER}local_rank${LOCAL_RANK}

export HCCL_SOCKET_IFNAME=eno1
export GLOO_SOCKET_IFNAME=eno1

export MODELSCOPE_CACHE="..."

export HCCL_CONNECT_TIMEOUT=7200
export HCCL_EXEC_TIMEOUT=7200

nproc_per_node=8
DISTRIBUTED_ARGS="--master_addr=$MASTER_ADDR --master_port=$MASTER_PORT --nnodes=$NNODES --node_rank=$RANK --nproc_per_node=$nproc_per_node"

torchrun $DISTRIBUTED_ARGS /ms-swift/swift/cli/_megatron/sft.py
--model qwen3_6 35b-a3
--save_safetensors true
--dataset
--sequence_parallel true
--micro_batch_size 1
--tensor_model_parallel_size 2
--pipeline_model_parallel_size 4
--expert_model_parallel_size 8
--packing true
--packing_length 4096
--global_batch_size 64
--logging_steps 5
--recompute_granularity full
--recompute_method uniform
--recompute_num_layers 10
--finetune true
--cross_entropy_loss_fusion true
--lr 1e-5
--lr_warmup_fraction 0.05
--model_type qwen3_5_moe
--lr_decay_style cosine
--num_train_epochs 3
--output_dir
--save_steps 100
--max_length 4096
--dataloader_num_workers 8
--no_save_optim true
--no_save_rng true
--attn_impl sdpa
--use_distributed_optimizer true
--dataset_num_proc 8
--use_precision_aware_optimizer true
--optimizer_cpu_offload true
--optimizer_offload_fraction 0.7

How to Reproduce / 如何复现

见bug描述

Additional Information / 补充信息

默认镜像配置下的最初报错
直接运行上述脚本，系统报错如下：
ModuleNotFoundError: No module named 'megatron.core.models.gpt.experimental_attention_variant_module_specs'

根本原因：镜像内同时存在两个 megatron 包。由于 PYTHONPATH 中 /Megatron-LM 排在前面，Python 优先找到了旧版。而旧版中缺少 experimental_attention_variant_module_specs.py 文件，导致 mcore-bridge 导入失败。

路径与版本对比：

路径：/Megatron-LM/megatron/ | 版本：旧版 | experimental_attention_variant_module_specs.py：缺失 | expert_dist_ckpt_decorator：有

路径：site-packages/megatron/ (megatron_core-0.16.1) | 版本：新版 | experimental_attention_variant_module_specs.py：有 | expert_dist_ckpt_decorator：缺失

禁用旧版后的第二次报错
为了让 Python 使用 site-packages 中的新版，将 /Megatron-LM/megatron 重命名禁用。此时 mcore-bridge 导入成功，但紧接着触发了第二个错误：
ImportError: cannot import name 'expert_dist_ckpt_decorator' from 'megatron.core.transformer.moe.experts'

根本原因：MindSpeed 的 grouped_linear.py 强依赖 expert_dist_ckpt_decorator 函数，但该函数在 megatron_core-0.16.1（新版）中已被移除。由于 MindSpeed 是 ms-swift 在检测到 NPU 环境时自动导入的（swift/megatron/init.py），无法通过修改训练参数来绕过。

该版本的官方镜像此前是否针对 NPU 环境下的 MoE 模型（如 --model_type qwen3_5_moe）进行过完整的端到端训练测试

Reproduce the failure with the provided Docker image and torchrun command, then inspect swift/megatron/init.py, MindSpeed's grouped_linear.py, /Megatron-LM/megatron, and the installed megatron_core package. Compare the two reported import errors and dependency versions. Done means the image has a compatible Megatron/MindSpeed setup and the NPU MoE training command passes initialization.