Skip to content

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

Open
Lw135 wants to merge 17 commits into
redai-studio:mainfrom
Lw135:ascend-dev-0821
Open

add Qwen3.5 MTP & CP Training Support and NPU Base Environment Upgrade#288
Lw135 wants to merge 17 commits into
redai-studio:mainfrom
Lw135:ascend-dev-0821

Conversation

@Lw135

@Lw135 Lw135 commented Aug 26, 2026

Copy link
Copy Markdown
Contributor

This batch of merges focuses on NPU (Ascend) support upgrades and Qwen3.5 training capabilities, including a CANN 9.0.0 base-image upgrade, an NPU patch-set refresh, and new Qwen3.5 MTP / CP training scripts along with related fixes.

Key Changes

1. NPU Base Environment Upgrade

  • Upgraded the NPU Docker image to CANN 9.0.0 and updated Dockerfile.npu
  • Refreshed the NPU patch set (megatron, mindspeed, mindspeed-bridge, sglang, sgl-kernel) to sync with the latest Relax
  • Fixed a DDR issue in sgl-kernel for MTP scenarios

2. Qwen3.5 MTP Training Support

  • Added Qwen3.5 MTP training scripts:
    • run_qwen35_9B_mtp_8xnpu_thd.sh (9B, 8-NPU THD)
image
  • run_qwen35-35B-pokemon-sft-mtp-8xnpu.sh (SFT example)
image
  • Fixed MTP bridge weight conversion: _hf_config attachment when pp>1, and updated the Qwen3.5 MTP weight mapping in the mindspeed-bridge patch

3. Qwen3.5-9B CP (Context Parallel) Training

  • Added a Qwen3.5-9B CP training script (16-NPU)
image

4. Documentation

  • Added an NPU feature support table to docker/npu-training.md

Major Commits

Commit Description
25a42cc build(docker): upgrade NPU image to CANN 9.0.0
3f51dcc chore(docker): refresh NPU patch set
9d3ac5e [NPU] fix DDR issue for MTP
700b291 / e5f12e5 feat(npu): add Qwen3.5 MTP SFT/training scripts
0127ec4 feat(npu): add Qwen3.5-9B CP training script
57809b4 fix(megatron): attach _hf_config for MTP bridge when pp>1
4160612 fix: update Qwen3.5 MTP weight mapping (mindspeed-bridge patch)
9c39605 docs(npu): add feature support table

Lw135 and others added 12 commits August 21, 2026 11:50
@Lw135
Merge branch 'wqw_base_lw_dev_0808' into ascend-dev-0808
6c91af9 
feat(docker): upgrade NPU env to CANN 9.0.0 with FLA and MindSpeed-Ops

# ⭐ Feature

## Upgrade CANN base image and toolchain

- Upgrade CANN base image from 8.5.1-a3 to 9.0.0-a3
- Upgrade torch_npu from v2.9.0-7.3.0 to v26.0.1-pytorch2.9.0
- Upgrade triton-ascend from 3.2.0 to 3.2.1

## Add AscendC Flash Linear Attention (FLA) support

- Clone and build fla_npu from flash-linear-attention-npu v26.1.0
- Compile causal_conv1d and gated_delta_rule ops for ascend910_93

## Add MindSpeed-Ops support

- Clone, checkout and install MindSpeed-Ops
- Add mindspeed-ops.patch for gated_delta_rule autotune key fix

---
# ♻️ Refactor

## Restructure Dockerfile build order

- Move torch/torch_npu install before repo clone
- Delay COPY . /root/Relax to just before patching
- Consolidate Megatron-Bridge into Megatron-LM via cp instead of separate path
- Install MindSpeed-Bridge with --no-deps to avoid circular dependency

## Migrate mindspeed-bridge to FLA ops

- Replace causal_conv1d from causal_conv1d with FLA implementation
- Replace mindspeed_ops l2norm with naive_l2norm fallback
- Switch to flash_gated_delta_rule when FLA is available

---
# 🐛 Bug Fix

## Fix various compatibility issues

- Fix autotune key in gated_delta_rule kernel by removing USE_G and IS_VARLEN
- Fix lambda closure bug in transformer_config_init_subclass (use default arg binding)
- Fix attention_mask dtype from int32 to bool in model preprocess
- Fix vision model config with MoE permute fusion disabled
- Handle OmegaConf DictConfig/ListConfig in remove_non_pickleables
- Guard apex MixedFusedLayerNorm import with is_npu_available check

## Update sgl-kernel-npu

- Upgrade sgl-kernel-npu checkout from 2026.04.15.rc3 to 2026.7.2
- Remove obsolete cherry-pick workaround

feat(docker): upgrade NPU env to CANN 9.0.0 with FLA and MindSpeed-Ops

feat(docker): upgrade NPU env to CANN 9.0.0 with FLA and MindSpeed-Ops

# ⭐ Feature

## Upgrade CANN base image and toolchain

- Upgrade CANN base image from 8.5.1-a3 to 9.0.0-a3
- Upgrade torch_npu from v2.9.0-7.3.0 to v26.0.1-pytorch2.9.0
- Upgrade triton-ascend from 3.2.0 to 3.2.1

## Add AscendC Flash Linear Attention (FLA) support

- Clone and build fla_npu from flash-linear-attention-npu v26.1.0
- Compile causal_conv1d and gated_delta_rule ops for ascend910_93

## Add MindSpeed-Ops support

- Clone, checkout and install MindSpeed-Ops
- Add mindspeed-ops.patch for gated_delta_rule autotune key fix

---
# ♻️ Refactor

## Restructure Dockerfile build order

- Move torch/torch_npu install before repo clone
- Delay COPY . /root/Relax to just before patching
- Consolidate Megatron-Bridge into Megatron-LM via cp instead of separate path
- Install MindSpeed-Bridge with --no-deps to avoid circular dependency

## Migrate mindspeed-bridge to FLA ops

- Replace causal_conv1d from causal_conv1d with FLA implementation
- Replace mindspeed_ops l2norm with naive_l2norm fallback
- Switch to flash_gated_delta_rule when FLA is available

---
# 🐛 Bug Fix

## Fix various compatibility issues

- Fix autotune key in gated_delta_rule kernel by removing USE_G and IS_VARLEN
- Fix lambda closure bug in transformer_config_init_subclass (use default arg binding)
- Fix attention_mask dtype from int32 to bool in model preprocess
- Fix vision model config with MoE permute fusion disabled
- Handle OmegaConf DictConfig/ListConfig in remove_non_pickleables
- Guard apex MixedFusedLayerNorm import with is_npu_available check

## Update sgl-kernel-npu

- Upgrade sgl-kernel-npu checkout from 2026.04.15.rc3 to 2026.7.2
- Remove obsolete cherry-pick workaround


[NPU] update sglang
@Lw135
!57 merge ascend-dev-0808 into ascend-dev-0821
7d33717 
Merge branch 'wqw_base_lw_dev_0808' into ascend-dev-0808

Created-by: dabuliu123
Commit-by: wuqiwei;lixionglong;dabuliu123;@dabuliu123
Merged-by: dabuliu123
Description: ## What

<!-- What changes does this PR introduce? -->

## Why

<!-- Why are these changes needed? Link related issues with "Fixes redai-studio#123" or "Relates to #456". -->

## How

<!-- How do the changes work? Describe the technical approach. -->

## Testing

<!-- How were the changes tested? Include commands, test results, or screenshots. -->

- [ ] `pre-commit run --all-files` passes
- [ ] Tests pass (`pytest tests/`)
- [ ] New tests added (if applicable)
- [ ] Documentation updated (if applicable)

## Type of Change

- [ ] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to change)
- [ ] Documentation update
- [ ] Refactoring (no functional changes)
- [ ] Performance improvement
- [ ] CI/CD or build changes

## Screenshots / Logs

<!-- If applicable, add screenshots or log output to help explain the changes. -->


See merge request: hw-pbclouds/Relax!57
lixionglong
fix(megatron): attach _hf_config for MTP bridge when pp>1
57809b4 
# 🐛 Bug Fix

## Fix bridge mapping registry failure with PP > 1

- Set model_bridge._hf_config from bridge.hf_pretrained.config when missing
@Weird733
chore(docker): refresh NPU patch set
3f51dcc 
# 🔩 Chore

## Sync NPU patches

- Add sgl-kernel-npu.patch (1847 lines)
- Update sglang-npu.patch (+1070 lines)
- Update mindspeed.patch / mindspeed-bridge.patch / megatron patches
@Weird733
build(docker): upgrade NPU image to CANN 9.0.0
25a42cc 
# 🔩 Chore

## Upgrade NPU Dockerfile dependencies

- Bump base image CANN 8.5.1 → 9.0.0
- Upgrade torch_npu to v26.0.1 (PyTorch 2.9.0)
- Upgrade triton-ascend 3.2.0 → 3.2.1
@Weird733
Merge branch 'wqw_xianwang_0821' into remotes/hw/ascend-dev-0821
a6f3205
@tang27-1
feat(npu): add Qwen3.5 MTP SFT/training scripts
700b291 
feat(npu): add Qwen3.5 MTP training scripts for 0821

# ⭐ Feature

## Add Qwen3.5 MTP training launch scripts

- run_qwen35-35B-pokemon-sft-mtp-8xnpu.sh: Qwen3.5-35B-A3B MTP SFT on pokemon-gpt4o-captions, 8xNPU single-node, ray-submit launch
- run_qwen35-35B-A3B-16xnpu-colocate-thd.sh: Qwen3.5-35B-A3B colocate THD training on 16xNPU
- run_qwen35_9B_mtp_8xnpu_thd.sh: Qwen3.5-9B MTP THD training on 8xNPU

---

# 🐛 Bug Fix

## Fix EXP_DIR silently overridden by MODEL_DIR default

- EXP_DIR now uses its own default with MODEL_DIR following EXP_DIR, matching the 9B THD script pattern

---

# 🔩 Chore

## Align script naming and comments

- Rename scripts to the 8xnpu naming convention and drop duplicated .sh suffix
- Remove commented-out --qkv-format bshd / --micro-batch-size 1 lines
- Sync Usage comments with actual script names
@tang27-1
!58 merge feat/npu-mtp-scripts-0824 into ascend-dev-0821
2e18d5a 
feat(npu): add Qwen3.5 MTP SFT/training scripts

Created-by: Tgz27
Commit-by: 唐二十七
Merged-by: Tgz27
Description: 新增 3 个 NPU MTP 训练脚本：

- run_qwen35-35B-pokemon-sft-mtp.sh: 35B-A3B MTP SFT on pokemon dataset, 8xNPU
- run_qwen35-35B-A3B-16xnpu-colocate-thd.sh: 35B-A3B colocate THD training on 16xNPU
- run_qwen35_9B_mtp_8npu_thd.sh.sh: 9B MTP THD training on 8xNPU

See merge request: hw-pbclouds/Relax!58
lixionglong
Revert "fix(megatron): attach _hf_config for MTP bridge when pp>1"
bb7dda7 
This reverts commit 57809b4.
lixionglong
fix: update qwen35 mtp mapping in mindspeed-bridge.patch for relax br…
4160612 
…idge_converter.py
@Lw135
fix(npu): rename Qwen3.5-9B CP script to 16xnpu
905857a 
feat(npu): add Qwen3.5-9B CP training script

# ⭐ Feature

## Add Qwen3.5-9B CP colocate training script

- Add run-qwen35-9B-8xnpu-cp.sh for DAPO math training with TP4/CP4 on 16 NPUs
- Set MINDSPEED_BRIDGE_GDN_USE_TORCH_CONV=1 to use torch conv fallback for GDN

## Support GDN causal conv on NPU via fla_npu

- Replace causal_conv1d with fla_npu in mindspeed-bridge.patch GDN layer
- Add MINDSPEED_BRIDGE_GDN_USE_TORCH_CONV env switch to fall back to torch conv

---

# 🐛 Bug Fix

## Fix repatch ordering before Megatron init

- Move repatch(args) before init(args) in MegatronTrainRayActor so bridge patches apply during initialization

fix(npu): rename Qwen3.5-9B CP script to 16xnpu

# 🐛 Bug Fix

## Fix Qwen3.5-9B CP script naming for 16 NPUs

- Rename run-qwen35-9B-8xnpu-cp.sh to run-qwen35-9B-16xnpu-cp.sh to match the actual 16-NPU setup
- Update header comment to 16xNPU and log filename to qwen35-9B-GRPO-npu16

docs(npu): add feature support table

# 📝 Documentation

## Add feature support table to NPU training guide

- Add 特性支持 table covering Qwen3.5-9B CP, Qwen3.5-9B MTP, and Qwen3.5-35B-A3B SFT MTP with reference scripts
- Update 下一步 checklist from feature support to performance optimization
- Fix MTP row minimum card count to 4卡 to match the script's 8 NPUs
@Lw135
!59 merge ascend-dev-0821-cp into ascend-dev-0821
d5aa205 
fix(npu): rename Qwen3.5-9B CP script to 16xnpu

Created-by: dabuliu123
Commit-by: dabuliu123
Merged-by: dabuliu123
Description: ## What

<!-- What changes does this PR introduce? -->

## Why

<!-- Why are these changes needed? Link related issues with "Fixes redai-studio#123" or "Relates to #456". -->

## How

<!-- How do the changes work? Describe the technical approach. -->

## Testing

<!-- How were the changes tested? Include commands, test results, or screenshots. -->

- [ ] `pre-commit run --all-files` passes
- [ ] Tests pass (`pytest tests/`)
- [ ] New tests added (if applicable)
- [ ] Documentation updated (if applicable)

## Type of Change

- [ ] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to change)
- [ ] Documentation update
- [ ] Refactoring (no functional changes)
- [ ] Performance improvement
- [ ] CI/CD or build changes

## Screenshots / Logs

<!-- If applicable, add screenshots or log output to help explain the changes. -->


See merge request: hw-pbclouds/Relax!59
@Lw135 @claude
fix(gitleaks): drop /mnt/ defaults in NPU scripts
f49de70 
# 🔒 Security

## Remove hardcoded internal paths from Qwen3.5 NPU scripts

- Replace /mnt/tidalfs-hwwh01 EXP_DIR defaults with the repo-standard `${SCRIPT_DIR}/../../../../exps` in run_qwen35-35B-pokemon-sft-mtp-8xnpu.sh, run_qwen35-35B-A3B-16xnpu-colocate-thd.sh and run_qwen35_9B_mtp_8xnpu_thd.sh
- Replace bare /mnt/ placeholder in run-qwen35-9B-16xnpu-cp.sh with the same repo-standard default
- Switch the SFT script DATA_DIR default to the `${DATA_DIR:-${EXP_DIR}}` convention used by all other training scripts

---

# 🎨 Style

## End-of-file fixer

- Remove trailing blank line from docker/Dockerfile.npu

Co-Authored-By: Claude <noreply@anthropic.com>
@Michael-Salon

Michael-Salon commented Aug 28, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

建议修改：
P1 — 在线更新权重后，Qwen3.5 卷积层继续使用旧权重。
docker/npu_patch/sglang-npu.patch:43-48
_get_conv_weights_t() 永久缓存 transpose().contiguous() 生成的独立副本。Actor→SGLang 在线同步只更新原始 conv1d.weight，缓存不会刷新；首次 rollout 后再更新权重，后续 decode/extend/verify 会形成新旧权重混合模型。建议在每次权重更新后失效缓存，或避免持久缓存该副本。

P1 — 指定 token-ID logprobs 的请求会直接报 AttributeError。
docker/npu_patch/sglang-npu.patch:495-509
前面赋值的是 next_token_token_ids_logprobs_val，随后却遍历不存在的 token_ids_logprobs_val。应改用正确字段并补一条 specific-token-logprob 测试。

P1 — Qwen2-MoE 的 NPU multi-stream 路径必崩。
docker/npu_patch/sglang-npu.patch:696-708
新代码调用 self.alt_stream.wait_stream()，但 Qwen2-MoE 在 NPU 上仍将 alt_stream 初始化为 None。启用 SGLANG_NPU_USE_MULTI_STREAM=1、DeepEP 和 graph forward 后会在首个非空 batch 崩溃。需要像本 PR 的 Qwen3.5 路径一样创建 NPU stream，或保留原 helper。

P1 — 当前 NPU Dockerfile 无法从声明的干净 base 完成构建。
docker/Dockerfile.npu:135-140
apt install gawk 缺少 -y。精确 base 镜像未预装 gawk，非交互 Docker build 会在确认提示处 abort。建议改成 apt-get install -y --no-install-recommends gawk。

P1 — 新 CP 脚本会把 WANDB_API_KEY 明文写入日志。
scripts/training/text/run-qwen35-9B-16xnpu-cp.sh:28-32
脚本启用了 set -x，随后将 key 作为 jq --arg 和 Ray CLI JSON 参数传递，xtrace 会输出完整密钥。应在处理和提交 secret 时关闭 xtrace，最好不要经命令行 JSON 传递。

P2 — MTP_NUM_LAYERS > 1 时 bridge 权重转换缺少 layer 1+ 映射。
docker/npu_patch/mindspeed-bridge.patch:195-241
新映射硬编码 layers.0，但启动脚本允许覆盖 MTP_NUM_LAYERS，Megatron 也会实际创建多层。首次 Actor→SGLang 权重转换会因 layer 1 缺少 registry entry 失败。应恢复按层构建映射，或显式限制该 recipe 只能为 1。

P2 — top-logprob 路径对全词表执行了两次 topk。
docker/npu_patch/sglang-npu.patch:475-493
get_top_logprobs() 已完成一次 topk，新增代码立即重复计算并覆盖结果。任何 top_logprobs_num > 0 的请求都会多执行一次昂贵 kernel，抵消本段的性能优化。应删除第一次 helper 调用。

lixionglong and others added 4 commits September 1, 2026 15:43
lixionglong
fix: enforce MTP_NUM_LAYERS=1 for Qwen3.5 NPU scripts
abf8a43 
add validation in the *sft.sh script to ensure that MTP_NUM_LAYERS can only be 1
lixionglong
!60 merge lxl-mtp-num-layers-recipe-0821 into ascend-dev-0821
c3adb7b 
fix: enforce MTP_NUM_LAYERS=1 for Qwen3.5 NPU scripts

Created-by: lixionglong
Commit-by: lixionglong
Merged-by: lixionglong
Description: ## What

<!-- What changes does this PR introduce? -->

在 Qwen3.5 的 NPU MTP 启动脚本（`scripts/training/text/run_qwen35_9B_mtp_8xnpu_thd.sh`、`scripts/training/sft/run_qwen35-35B-pokemon-sft-mtp-8xnpu.sh`）中加入 `MTP_NUM_LAYERS` 校验：开启 MTP 训练时该参数只能为 `1`，否则脚本报错退出；同时在 `docker/npu-training.md` 中补充了该约束的说明。

## Why

<!-- Why are these changes needed? Link related issues with "Fixes redai-studio#123" or "Relates to #456". -->

Qwen3.5 原始 checkpoint 中只包含 1 层 MTP 权重（`mtp_num_hidden_layers=1`），若用户将 `MTP_NUM_LAYERS` 设为非 `1` 的值，训练会在模型加载/初始化阶段失败且报错不直观。增加快速校验可以让非法配置在脚本启动阶段就给出明确报错（fail fast）。

## How

<!-- How do the changes work? Describe the technical approach. -->

两个脚本在组装 `MTP_ARGS` 前检查 `${MTP_NUM_LAYERS:-1}`（未设置时默认为 `1`）：若值非 `1`，向 stderr 输出 `ERROR: MTP_NUM_LAYERS must be 1 for Qwen3.5 (checkpoint has mtp_num_hidden_layers=1).` 并以退出码 `1` 终止。

## Testing

<!-- How were the changes tested? Include commands, test results, or screenshots. -->

纯 shell 启动脚本改动，不涉及 Python 代码，`pytest tests/` 不受影响。可通过如下命令快速验证校验逻辑：

```bash
MTP_NUM_LAYERS=2 bash scripts/training/text/run_qwen35_9B_mtp_8xnpu_thd.sh   # 应报错退出
bash scripts/training/text/run_qwen35_9B_mtp_8xnpu_thd.sh                     # 未设置，默认 1，校验通过


See merge request: hw-pbclouds/Relax!60
@Lw135 @tang27-1
refactor(npu): drop wandb wiring from CP script
ebfc364 
fix(docker): install gawk non-interactively

# 🐛 Bug Fix

## Make gawk install non-interactive in NPU Dockerfile

- Replace `apt install gawk` with `apt-get install -y --no-install-recommends gawk`
- `apt install` prompts for confirmation and aborts in unattended Docker builds

refactor(npu): drop wandb wiring from CP script

# ♻️ Refactor

## Drop script-local wandb wiring from Qwen3.5-9B CP script

- Remove WANDB_API_KEY injection into Ray runtime_env via jq
- Remove WANDB_API_KEY-gated --use-wandb/--wandb-project/--wandb-group args
- Wandb is configured centrally via --use-wandb/--wandb-key CLI args
@tang27-1
!61 merge fix-docker-ascend-dev-0821 into ascend-dev-0821
b600c8e 
refactor(npu): drop wandb wiring from CP script

Created-by: dabuliu123
Commit-by: dabuliu123
Merged-by: Tgz27
Description: ## What

<!-- What changes does this PR introduce? -->

## Why

<!-- Why are these changes needed? Link related issues with "Fixes redai-studio#123" or "Relates to #456". -->

## How

<!-- How do the changes work? Describe the technical approach. -->

## Testing

<!-- How were the changes tested? Include commands, test results, or screenshots. -->

- [ ] `pre-commit run --all-files` passes
- [ ] Tests pass (`pytest tests/`)
- [ ] New tests added (if applicable)
- [ ] Documentation updated (if applicable)

## Type of Change

- [ ] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to change)
- [ ] Documentation update
- [ ] Refactoring (no functional changes)
- [ ] Performance improvement
- [ ] CI/CD or build changes

## Screenshots / Logs

<!-- If applicable, add screenshots or log output to help explain the changes. -->


See merge request: hw-pbclouds/Relax!61
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@Yangruipis Yangruipis Awaiting requested review from Yangruipis Yangruipis is a code owner
@NINGBENZHE NINGBENZHE Awaiting requested review from NINGBENZHE NINGBENZHE is a code owner

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@Lw135 @Michael-Salon @Weird733 @tang27-1