add Qwen3.5 MTP & CP Training Support and NPU Base Environment Upgrade

docker/Dockerfile.npu

@@ -3,7 +3,7 @@
 ARG HTTP_PROXY
 ARG HTTPS_PROXY
 ARG NO_PROXY
-ARG BASE_IMAGE=quay.io/ascend/cann:8.5.1-a3-ubuntu22.04-py3.11
+ARG BASE_IMAGE=quay.io/ascend/cann:9.0.0-a3-ubuntu22.04-py3.11
 FROM ${BASE_IMAGE} as base
 
 ARG HTTP_PROXY
@@ -45,11 +45,13 @@ RUN ARCH=$(uname -m) && \
         export LD_LIBRARY_PATH=/usr/local/Ascend/ascend-toolkit/latest/x86_64-linux/devlib/linux/x86_64/:$LD_LIBRARY_PATH; \
     fi && \
     source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
-    source /usr/local/Ascend/nnal/atb/set_env.sh
+    source /usr/local/Ascend/nnal/atb/set_env.sh && \
+    source /usr/local/Ascend/cann-9.0.0/share/info/ascendnpu-ir/bin/set_env.sh
 # Configure pip and git.
 ARG PIP_INDEX_URL="https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple"
 RUN pip config set global.index-url ${PIP_INDEX_URL} && \
     git config --global http.sslverify false && \
+    git config --global https.sslverify false && \
     git config --global http.postBuffer 2147483648 && \
     git config --global user.email "temp@example.com" && \
     git config --global user.name "temp"
@@ -71,13 +73,13 @@ RUN pip_install() { pip install "$@" || pip install --index-url https://pypi.org
     pip_install pyyaml && \
     git clone https://gitcode.com/Ascend/pytorch.git /root/pytorch && \
     cd /root/pytorch && \
-    git checkout v2.9.0-7.3.0 && \
+    git checkout v26.0.1-pytorch2.9.0 && \
     git cherry-pick -n f495de675bce38a2fa21edbf067b73d2a5f26733 && \
     bash ci/build.sh --python=3.11 && \
-    pip install dist/torch_npu-2.9.0*.whl
+    pip install dist/torch_npu*.whl
+RUN pip install triton-ascend==3.2.1 --extra-index-url=https://triton-ascend.osinfra.cn/pypi/simple && \
+    cd /root && rm -rf /root/pytorch
 RUN pip_install() { pip install "$@" || pip install --index-url https://pypi.org/simple "$@"; }; \
-    cd /root && rm -rf /root/pytorch && \
-    pip_install triton-ascend==3.2.0 && \
     pip_install tensordict==0.10.0 pyvers==0.1.0 --no-deps
 RUN git clone https://github.com/redai-studio/TransferQueue.git /root/TransferQueue && \
     cd /root/TransferQueue && \
@@ -89,57 +91,77 @@ RUN git clone https://github.com/redai-studio/TransferQueue.git /root/TransferQu
 # Install the MindSpeed/Megatron stack.
 RUN git clone https://gitcode.com/ascend/MindSpeed.git /root/MindSpeed && \
     git clone https://github.com/NVIDIA/Megatron-LM.git /root/Megatron-LM && \
+    git clone https://gitcode.com/ascend/MindSpeed-Ops.git /root/MindSpeed-Ops && \
     git clone https://gitcode.com/ascend/MindSpeed-Bridge.git /root/MindSpeed-Bridge && \
     git clone https://github.com/NVIDIA-NeMo/Megatron-Bridge.git /root/Megatron-Bridge
 
 RUN pip_install() { pip install "$@" || pip install --index-url https://pypi.org/simple "$@"; }; \
-    cd /root/MindSpeed && git checkout core_r0.16.0 && pip_install -r requirements.txt && pip_install -e . --no-build-isolation && \
-    cd /root/Megatron-LM && git checkout core_v0.16.1 && pip_install -e . --no-build-isolation && \
+    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+    source /usr/local/Ascend/nnal/atb/set_env.sh && \
+    source /usr/local/Ascend/cann-9.0.0/share/info/ascendnpu-ir/bin/set_env.sh && \
+    cd /root/MindSpeed && git checkout core_r0.16.0 && pip install -r requirements.txt && pip install -e . && git checkout e4772499 && \
+    cd /root/Megatron-LM && git checkout core_v0.16.1 && pip install -e . --no-build-isolation && \
     cd /root/Megatron-Bridge && git checkout v0.3.1 && \
-    cd /root/MindSpeed-Bridge && git checkout 3655c07cbcc9 && pip_install -r requirements.txt && bash tools/install_auto.sh
+    cd /root/MindSpeed-Ops && git checkout 33ac80f7 && pip install -e . --no-build-isolation --no-deps && \
+    cd /root/MindSpeed-Bridge && git checkout v0.3.1 && pip install -r requirements.txt && pip install -e . --no-deps
 
 # Keep Relax code changes out of the dependency cache.
 COPY docker/npu_patch /root/Relax/docker/npu_patch
 
 # Apply NPU patches.
-RUN cd /root/MindSpeed && \
-    patch -p1 < /root/Relax/docker/npu_patch/mindspeed.patch && \
-    cd /root/Megatron-Bridge && \
+RUN cd /root/Megatron-Bridge && \
     patch -p1 < /root/Relax/docker/npu_patch/megatron-bridge.patch && \
+    git add . && git commit -m "base line" && \
     cd /root/Megatron-LM && \
     patch -p1 < /root/Relax/docker/npu_patch/megatron.patch && \
+    git add . && git commit -m "base line" && \
     cd /root/MindSpeed-Bridge && \
-    patch -p1 < /root/Relax/docker/npu_patch/mindspeed-bridge.patch
+    patch -p1 < /root/Relax/docker/npu_patch/mindspeed-bridge.patch && \
+    git add . && git commit -m "base line" && \
+    cd /root/MindSpeed-Ops && \
+    patch -p1 < /root/Relax/docker/npu_patch/mindspeed-ops.patch && \
+    git add . && git commit -m "base line" && \
+    cd /root/MindSpeed && \
+    patch -p1 < /root/Relax/docker/npu_patch/mindspeed.patch && \
+    git add . && git commit -m "base line"
 
-# Merge bridge packages into Megatron-LM.
-RUN cp -r /root/MindSpeed-Bridge/mindspeed_bridge /root/Megatron-LM/ && \
-    cp -r /root/Megatron-Bridge/src/megatron/bridge /root/Megatron-LM/megatron/ && \
-    cd /root && rm -rf /root/MindSpeed-Bridge && rm -rf /root/Megatron-Bridge
 
-# Install SGLang.
+# Install sglang
 RUN pip_install() { pip install "$@" || pip install --index-url https://pypi.org/simple "$@"; }; \
     git clone https://github.com/sgl-project/sglang.git /root/sglang && \
-    cd /root/sglang && git checkout v0.5.10 && \
+    cd /root/sglang && git checkout v0.5.15 && \
     mv python/pyproject.toml python/pyproject.toml.backup && \
     mv python/pyproject_npu.toml python/pyproject.toml && \
-    echo "torch==2.9.0" > /tmp/sglang-constraints.txt && \
-    pip_install -e "python[srt_npu]" --constraint /tmp/sglang-constraints.txt && \
-    python -c 'import importlib.metadata as m; v = m.version("torch"); assert v.startswith("2.9.0"), v' && \
-    rm /tmp/sglang-constraints.txt && \
-    git stash  && \
-    git fetch origin pull/23815/head:pr-23815 && \
-    git checkout pr-23815 && \
-    patch -p1 < /root/Relax/docker/npu_patch/sglang-npu.patch
+    pip install -e "python[srt_npu]" --constraint <(echo "torch==2.9.0") && \
+    git add . && git commit -m "install info" && \
+    git fetch && \
+    git cherry-pick ece02ffc9cc32e94382d4f1b553b2c755f83f722 && \
+    patch -p1 < /root/Relax/docker/npu_patch/sglang-npu.patch && \
+    git add . && git commit -m "sglang-npu.patch"
 
 # Install SGLang NPU kernels.
 RUN git clone https://github.com/sgl-project/sgl-kernel-npu /root/sgl-kernel-npu && \
-    cd /root/sgl-kernel-npu && git checkout 2026.04.15.rc3 && \
-    git cherry-pick -n 23519771d347 --no-gpg-sign && \
-    patch -p1 < /root/Relax/docker/npu_patch/torch-memory-saver.patch  && \
-    bash build.sh -a kernels && bash build.sh -a memory-saver && \
+    cd /root/sgl-kernel-npu && git checkout 2026.7.2 && \
+    patch -p1 < /root/Relax/docker/npu_patch/sgl-kernel-npu.patch && \
+    git add . && git commit -m "sgl-kernel-npu.patch" && \
+    bash build.sh && \
     pip install output/*.whl && \
     cd /root && rm -rf /root/sgl-kernel-npu
 
+# Install AscendC FLA
+RUN git clone https://github.com/flashserve/flash-linear-attention-npu.git /root/flash-linear-attention-npu && \
+    cd /root/flash-linear-attention-npu && git checkout v26.1.0 && \
+    apt update && apt-get install -y --no-install-recommends gawk && \
+    # 编译命令，注意--soc=${soc_version}需要指定为当前机器的芯片类型{ascend910b/ascend910_93/ascend950}
+    bash build.sh --soc=ascend910_93 --pkg --ops=causal_conv1d,chunk_bwd_dv_local,chunk_bwd_dqkwg,chunk_gated_delta_rule_bwd_dhu,prepare_wy_repr_bwd_da,prepare_wy_repr_bwd_full,chunk_fwd_o,chunk_gated_delta_rule_fwd_h,recurrent_gated_delta_rule,recompute_wu_fwd && \
+    # 安装run包
+    ./build_out/cann-*.run && \
+    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+    source /usr/local/Ascend/nnal/atb/set_env.sh && \
+    source /usr/local/Ascend/cann-9.0.0/share/info/ascendnpu-ir/bin/set_env.sh && \
+    # 一键编译安装脚本，先调用torchnpugen自动接入算子，再运行setup编whl包，最后安装whl包
+    cd torch_custom/fla_npu && bash build.sh
+
 FROM train as relax
 
 WORKDIR /root

docker/npu-training.md

@@ -12,14 +12,22 @@
 | Qwen3.5-9B      | DAPO     | √    | √     | 910C 2卡         | `scripts/training/text/run-qwen35-9B-4xnpu-colocate.sh`       |
 | Qwen3.5-35B-A3B | DAPO     | √    | √     | 910C 8卡         | `scripts/training/text/run-qwen35-35B-A3B-16xnpu-colocate.sh` |
 
+## 特性支持
+
+| 模型            | 训练场景 | 多模态 | MTP | CP  | 训练所需最小卡数 | 参考脚本                                                       |
+| --------------- | -------- | ------ | --- | --- | ---------------- | -------------------------------------------------------------- |
+| Qwen3.5-9B      | DAPO     | -      | -   | √   | 910C 8卡         | `scripts/training/text/run-qwen35-9B-16xnpu-cp.sh`             |
+| Qwen3.5-9B      | DAPO     | -      | √   | -   | 910C 4卡         | `scripts/training/text/run_qwen35_9B_mtp_8xnpu_thd.sh`         |
+| Qwen3.5-35B-A3B | SFT      | √      | √   | -   | 910C 4卡         | `scripts/training/sft/run_qwen35-35B-pokemon-sft-mtp-8xnpu.sh` |
+
 ## 环境准备
 
 ### 前置准备
 
 - 资源类型：`Ascend910 Snt9b23`
 - 驱动版本：`Software Version 25.5.1`
 - 固件版本：`Firmware Version 7.8.0.6.201`
-- 基础镜像：`quay.io/ascend/cann:8.5.1-a3-ubuntu22.04-py3.11`
+- 基础镜像：`quay.io/ascend/cann:9.0.0-a3-ubuntu22.04-py3.11`
 
 ### 环境检查
 
@@ -49,7 +57,7 @@ REGISTRY=<registry> make docker-ascend
 REGISTRY=<registry> ASCEND_DOCKER_BUILDKIT=0 make docker-ascend
 ```
 
-可配置变量：`BASE_IMAGE`（默认 `quay.io/ascend/cann:8.5.1-a3-ubuntu22.04-py3.11`）、`SOC_VERSION`（默认 `ascend910_9391`）、`REGISTRY`、`DO_PUSH`（默认 `1`，设为 `0` 时不推送并检查本地镜像）、`ASCEND_DOCKER_BUILDKIT`（默认 `1`）。远端已存在同名镜像时会跳过构建。
+可配置变量：`BASE_IMAGE`（默认 `quay.io/ascend/cann:9.0.0-a3-ubuntu22.04-py3.11`）、`SOC_VERSION`（默认 `ascend910_9391`）、`REGISTRY`、`DO_PUSH`（默认 `1`，设为 `0` 时不推送并检查本地镜像）、`ASCEND_DOCKER_BUILDKIT`（默认 `1`）。远端已存在同名镜像时会跳过构建。
 
 （可选，内部 QS 镜像）复用 `ml-engine/tools/relax-ci` 的 `docker/Dockerfile.qs`（纯 Python 依赖，架构无关，ARM64 可直接构建）。由 CI 先 checkout relax-ci，再指向其 `Dockerfile.qs`：
 
@@ -146,6 +154,10 @@ bash scripts/training/text/run-qwen3-4B-8xgpu-async-npu.sh
 
   > MISC_ARGS，显示启用FlashAttention实现 `--use-flash-attn`
 
+### MTP 特性说明
+
+开启 MTP 训练时，`MTP_NUM_LAYERS`（MTP 头层数，脚本透传给训练参数 `--mtp-num-layers`）只能为 `1`：由于 Qwen3.5 原始 checkpoint 中仅包含 1 层 MTP 权重（`mtp_num_hidden_layers=1`），MTP 相关启动脚本（`scripts/training/text/run_qwen35_9B_mtp_8xnpu_thd.sh`、`scripts/training/sft/run_qwen35-35B-pokemon-sft-mtp-8xnpu.sh`）已加入校验，当该参数被设置为非 `1` 的值时脚本会报错退出
+
 ## 下一步
 
-- [ ] 特性支持：多模态、CP长序列、MTP 等
+- [ ] 性能优化：Qwen3.5-35B-A3B 等

docker/npu_patch/megatron-bridge.patch

@@ -1,11 +1,32 @@
 diff --git a/src/megatron/bridge/models/conversion/utils.py b/src/megatron/bridge/models/conversion/utils.py
-index 5a66e719..3d411c17 100644
+index 86ddf8661..115624b16 100644
 --- a/src/megatron/bridge/models/conversion/utils.py
 +++ b/src/megatron/bridge/models/conversion/utils.py
-@@ -203,6 +203,15 @@ def remove_non_pickleables(obj, max_depth: int = 3, current_depth: int = 0):
+@@ -203,6 +203,17 @@ def remove_non_pickleables(obj, max_depth: int = 3, current_depth: int = 0):
+         ):  # bound methods
+             return None
+
++    # Convert OmegaConf containers to plain dict/list to avoid in-place
++    # mutation triggering "dictionary changed size during iteration" errors
++    # inside OmegaConf's internal _flags_cache handling.
++    try:
++        from omegaconf import DictConfig, ListConfig, OmegaConf as _OmegaConf
++
++        if isinstance(obj, (DictConfig, ListConfig)):
++            obj = _OmegaConf.to_container(obj, resolve=True)
++    except ImportError:
++        pass
++
+     # Handle dataclass/object with attributes
+     if hasattr(obj, "__dict__"):
+         # Create a copy to avoid modifying the original
+@@ -213,9 +224,18 @@ def remove_non_pickleables(obj, max_depth: int = 3, current_depth: int = 0):
 
              # Recursively clean attribute
              cleaned_value = remove_non_pickleables(attr_value, max_depth, current_depth + 1)
+-
+-            # Set the cleaned value (or None if it was removed)
+-            setattr(cleaned_obj, attr_name, cleaned_value)
 +            if hasattr(cleaned_obj, '__setattr__'):
 +                try:
 +                    setattr(cleaned_obj, attr_name, cleaned_value)
@@ -15,6 +36,33 @@ index 5a66e719..3d411c17 100644
 +                        print(f"Skipping attribute '{attr_name}' due to Union type")
 +                        continue
 +                    raise
++            else:
++                # Fallback for objects without __setattr__ override
++                setattr(cleaned_obj, attr_name, cleaned_value)
+
+         return cleaned_obj
+
+diff --git a/src/megatron/bridge/peft/utils.py b/src/megatron/bridge/peft/utils.py
+index 1ca5b18bd..4797e2e42 100644
+--- a/src/megatron/bridge/peft/utils.py
++++ b/src/megatron/bridge/peft/utils.py
+@@ -33,6 +33,7 @@ from megatron.core.transformer.moe.router import TopKRouter
+
+ from megatron.bridge.utils.import_utils import safe_import_from
+
++from relax.utils.device import is_npu_available
+
+ TEColumnParallelLinear, HAVE_TE_COL_LINEAR = safe_import_from(
+     "megatron.core.extensions.transformer_engine", "TEColumnParallelLinear"
+@@ -62,7 +63,10 @@ HAVE_TE = all(
+     )
+ )
+
+-MixedFusedLayerNorm, HAVE_APEX = safe_import_from("apex.normalization.fused_layer_norm", "MixedFusedLayerNorm")
++if is_npu_available:
++    MixedFusedLayerNorm, HAVE_APEX = None, False
++else:
++    MixedFusedLayerNorm, HAVE_APEX = safe_import_from("apex.normalization.fused_layer_norm", "MixedFusedLayerNorm")
 
-             # Set the cleaned value (or None if it was removed)
-             setattr(cleaned_obj, attr_name, cleaned_value)
+ TECL = (TEColumnParallelLinear, TELayerNormColumnParallelLinear, TEColumnParallelGroupedLinear)
+ TERL = (TERowParallelLinear, TERowParallelGroupedLinear)

docker/npu_patch/megatron.patch

@@ -655,7 +655,7 @@ index 4be974017..0dfdb0928 100644
          """
          Maintain the expert bias in float32.
 diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
-index 2edb652bf..58fb89c4e 100755
+index 2edb652bf..7a068271b 100755
 --- a/megatron/core/transformer/multi_token_prediction.py
 +++ b/megatron/core/transformer/multi_token_prediction.py
 @@ -586,6 +586,102 @@ class MTPLossAutoScaler(torch.autograd.Function):
@@ -857,6 +857,42 @@ index 2edb652bf..58fb89c4e 100755
                  )
 
          if self.config.recompute_method == 'uniform':
+@@ -1053,7 +1192,7 @@ class MultiTokenPredictionBlock(MegatronModule):
+         self._build_layers(pg_collection)
+         assert len(self.layers) > 0, "MultiTokenPredictionBlock must have at least one layer."
+         self.cp_group = pg_collection.cp
+-
++        self._register_mtp_grad_hooks()
+     def _build_layers(self, pg_collection):
+         def build_layer(layer_spec, layer_number):
+             fp8_init_context = get_fp8_context(self.config, is_init=True)
+@@ -1074,6 +1213,26 @@ class MultiTokenPredictionBlock(MegatronModule):
+             ]
+         )
+
++    def _register_mtp_grad_hooks(self):
++        if len(self.layers) == 0:
++            return
++
++        def _make_sync_hook():
++            def _hook(grad):
++                torch.npu.synchronize()
++            return _hook
++
++        sync_keys = [
++            "self_attention.linear_qkv.weight",
++            "eh_proj.weight",                    
++        ]
++
++        for name, param in self.layers[0].named_parameters():
++            for key in sync_keys:
++                if key in name:
++                    param.register_hook(_make_sync_hook())
++                    break
++
+     def forward(
+         self,
+         input_ids: Tensor,
 diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
 index eaae58590..0f01f6bc0 100644
 --- a/megatron/core/transformer/transformer_config.py