Apply-MegatronAdaptor-NPU-migration-to-clean-branch by momo609 · Pull Request #385 · vllm-project/vime

momo609 · 2026-08-14T02:56:31Z

Summary

Remove MindSpeed-specific dependencies and imports.
Switch NPU backend integration to MegatronAdaptor and TransformerEngineNPU.
Align Megatron-related patches with the mainline implementation.
Remove obsolete MindSpeed and Megatron communication patches.
Update Megatron, vLLM, and bridge patch application order.
Fix MegatronAdaptor initialization and feature re-patching in the actor.
Update NPU attention patch and weight-transfer compatibility logic.
Fix nested NPU memory-pool allocation during training offload.
Update the related unit test and deployment scripts.

Files Changed

vime/backends/megatron_utils/
vime/utils/
docker/
scripts/
examples/
tests/
train.py
train_async.py

gemini-code-assist

Code Review

This pull request transitions the NPU training stack from MindSpeed to MegatronAdaptor and TransformerEngineNPU, updates the Megatron-LM commit, and applies several patches for optimizer steps, memory pool regions, and vLLM speculative draft updates. The reviewer's feedback is highly actionable and should be addressed: the new megatron_bridge.patch is orphaned in the Dockerfile and lacks path resolution in the environment update script; cloning dependencies from main without pinning commits makes the Docker build fragile; and the run-qwen3-4B-npu.sh script contains several hardcoded absolute paths and NPU device IDs that break portability.

gemini-code-assist · 2026-08-14T02:58:39Z

+RUN git clone https://gitcode.com/Ascend/MegatronAdaptor.git /root/MegatronAdaptor && \
+    git clone https://gitcode.com/Ascend/TransformerEngineNPU.git /root/TransformerEngineNPU


Cloning MegatronAdaptor and TransformerEngineNPU from the main branch without pinning a specific commit or tag makes the Docker build non-reproducible and fragile. Any breaking change in their upstream repositories will break the build. Please define build arguments for their commits (e.g., MEGATRON_ADAPTOR_COMMIT and TRANSFORMER_ENGINE_NPU_COMMIT) and check them out, similar to Megatron-LM.

read-the-docs-community · 2026-08-14T03:04:36Z

Documentation build overview

📚 vime | 🛠️ Build #34327151 | 📁 Comparing 5611e37 against latest (8144096)

🔍 Preview build

48 files changed · + 3 added · ± 27 modified · - 18 deleted

+ Added

± Modified

- Deleted

Meihan-chen · 2026-08-18T06:35:10Z

    done && \
    git -C /root/Megatron-LM apply --whitespace=nowarn \
-      /tmp/npu_patch/megatron_comm.patch && \
+      /tmp/npu_patch/megatron_mainline.patch && \


Could we simplify the Megatron patch layering here? The two patches modify several of the same core files, making future mainline updates prone to apply conflicts or duplicated behavior; please consider a single composite patch or a minimal non-overlapping NPU overlay.

Meihan-chen · 2026-08-18T06:35:47Z

+
+        self._select_weight_update_target(is_draft)
+         self.weight_transfer_engine.start_weight_update()
+         self._weight_update_active = True


Could you clarify why thist file are being synchronized here？It seems not used by Dockerfile.npu，so it may be better to remove these unrelated changes or update these in a separate PR.

Meihan-chen · 2026-08-18T06:35:47Z

+    return True
+
+
+_register_auto_classes()


floatlibai · 2026-08-20T08:08:31Z


-git -C "${WORKSPACE}/MindSpeed" apply --whitespace=nowarn "${PATCH_DIR}/mindspeed.patch"
+pip install --no-deps --no-build-isolation -e ${WORKSPACE}/MegatronAdaptor
+pip install --no-deps --no-build-isolation -e ${WORKSPACE}/TransformerEngineNPU


Do we need to clone these two repos or not? I think this step could use more detail.

need install

floatlibai · 2026-08-20T08:13:46Z

+    # MegatronAdaptor must run before Megatron imports so its dummy NPU modules
+    # are bound in Megatron tensor-parallel modules.
+    import megatron_adaptor  # noqa: F401
+    from . import npu_attention_patch  # noqa: F401


Why is this line needed here? Isn't the import in actor.py already sufficient?

Meihan-chen · 2026-08-25T02:59:09Z

+ARG MEGATRON_COMMIT=1dcf0dafa884ad52ffb243625717a3471643e087
 ARG MEGATRON_BRIDGE_COMMIT=3fd3768045422d0aa5c97e90a4e6c659aea9acb9
-ARG MINDSPEED_COMMIT=fc63de5c48426dd019c3b3f39e65f5bdf56e4086
 ARG MBRIDGE_COMMIT=89eb10887887bc74853f89a4de258c0702932a1c


Could you please remove the megatron_comm.patch and mindspeed.patch entries from docker/npu_patch/series.conf? This keeps the declared patch stack consistent with the Dockerfile and available patch files.

CalvinXKY · 2026-08-28T01:24:29Z

    pip install --no-deps --no-build-isolation -e /root/Megatron-LM && \
+    pip install --no-deps --no-build-isolation -e /root/TransformerEngineNPU && \
+    pip install --no-deps --no-build-isolation -e /root/MegatronAdaptor && \
    pip install --no-deps --no-build-isolation -e /root/MindSpeed


NPU CI (#193) likely fails during image build: Dockerfile.npu removes the MindSpeed clone but still runs pip install -e /root/MindSpeed.
Drop that line？ This migration also requires image-build (not smk-only on the legacy MindSpeed image), and the common docker/patch/latest/megatron.patch is Dockerfile-only and not in series.conf.

Offline weight updates depend on MindSpeed, so the MindSpeed code has been re-added.

Signed-off-by: wangxiaoxin-sherie <wangxiaoxin7@huawei.com>

This reverts commit 6d7e969. Signed-off-by: wangxiaoxin-sherie <wangxiaoxin7@huawei.com>

This reverts commit 5e53e29. Signed-off-by: wangxiaoxin-sherie <wangxiaoxin7@huawei.com>

This reverts commit f5ce534. Signed-off-by: wangxiaoxin-sherie <wangxiaoxin7@huawei.com>

This reverts commit 70d763c. Signed-off-by: wangxiaoxin-sherie <wangxiaoxin7@huawei.com>

Revert f5b8491 (vllm-project#409) and the associated NPU adaptations in 7a27808 and 81ae79d. Defer Qwen3.5 validation to the next stage with a fresh, matched serving and training environment. Remove the NPU-only GDN interfaces, FLA bootstrap and serving isolation, forced spawn, dedicated scripts/tests/CI entry, and FLA build recipe. Restore the pre-409 Bridge pin and patch while preserving the common TransformerLayer implementation. Retain S7 vllm-project#385 training-stack migration, vllm-project#396 torch_dist/ref-load, Qwen3-VL fixes, and main Qwen3.5 model code. Serving patches and latest patches are unchanged. No installed environment is rolled back. Validation: 183 grouped CPU tests passed; Ruff and runner shell syntax passed; common-to-NPU Megatron and reverted Bridge patches pass clean-base apply checks. Existing basic E2E and torch_dist PASS logs are retained; no fresh post-revert E2E was run. Signed-off-by: Meihan-chen <zr010426ztt@outlook.com>

gemini-code-assist Bot reviewed Aug 14, 2026

View reviewed changes

Meihan-chen reviewed Aug 18, 2026

View reviewed changes

momo609 force-pushed the ascendupdate2 branch from db2468f to 24fab47 Compare August 20, 2026 06:22

floatlibai reviewed Aug 20, 2026

View reviewed changes

Meihan-chen mentioned this pull request Aug 21, 2026

[NPU][CI] Refactor patch reconciliation #390

Merged

momo609 force-pushed the ascendupdate2 branch from 24fab47 to 9339c10 Compare August 24, 2026 12:10

Meihan-chen reviewed Aug 25, 2026

View reviewed changes

momo609 force-pushed the ascendupdate2 branch 6 times, most recently from 6c86097 to 0b8002e Compare August 27, 2026 12:27

CalvinXKY reviewed Aug 28, 2026

View reviewed changes

momo609 force-pushed the ascendupdate2 branch 8 times, most recently from 793206b to 0d8e686 Compare September 1, 2026 06:30

wangxiaoxin-sherie added 7 commits September 1, 2026 15:50

Apply-MegatronAdaptor-NPU-migration-to-clean-branch

4057878

Signed-off-by: wangxiaoxin-sherie <wangxiaoxin7@huawei.com>

fix(npu): apply common Megatron patch before NPU patch

dbbe100

Signed-off-by: wangxiaoxin-sherie <wangxiaoxin7@huawei.com>

fix(npu): complete and pin training dependencies

9e2079d

Signed-off-by: wangxiaoxin-sherie <wangxiaoxin7@huawei.com>

feat(npu): support proxied GitHub builds

0d5e894

Signed-off-by: wangxiaoxin-sherie <wangxiaoxin7@huawei.com>

fix(npu): support argparse tuple choices in mindspeed

2b7386e

Signed-off-by: wangxiaoxin-sherie <wangxiaoxin7@huawei.com>

fix(npu): align MindSpeed patch context

3608e57

Signed-off-by: wangxiaoxin-sherie <wangxiaoxin7@huawei.com>

fix(npu): select HCCL for torch dist conversion

2c250e4

Signed-off-by: wangxiaoxin-sherie <wangxiaoxin7@huawei.com>

wangxiaoxin-sherie added 5 commits September 1, 2026 15:50

fix(npu): disable gloo groups during conversion

177b983

Signed-off-by: wangxiaoxin-sherie <wangxiaoxin7@huawei.com>

Revert "fix(npu): disable gloo groups during conversion"

a18b20f

This reverts commit 6d7e969. Signed-off-by: wangxiaoxin-sherie <wangxiaoxin7@huawei.com>

Revert "fix(npu): select HCCL for torch dist conversion"

9843fce

This reverts commit 5e53e29. Signed-off-by: wangxiaoxin-sherie <wangxiaoxin7@huawei.com>

Revert "fix(npu): align MindSpeed patch context"

535f9d4

This reverts commit f5ce534. Signed-off-by: wangxiaoxin-sherie <wangxiaoxin7@huawei.com>

Revert "fix(npu): support argparse tuple choices in mindspeed"

5611e37

This reverts commit 70d763c. Signed-off-by: wangxiaoxin-sherie <wangxiaoxin7@huawei.com>

momo609 force-pushed the ascendupdate2 branch from 0d8e686 to 5611e37 Compare September 1, 2026 07:50

Meihan-chen approved these changes Sep 1, 2026

View reviewed changes

Meihan-chen merged commit 35eac6b into vllm-project:ascend Sep 1, 2026
4 of 5 checks passed

		RUN git clone https://gitcode.com/Ascend/MegatronAdaptor.git /root/MegatronAdaptor && \
		git clone https://gitcode.com/Ascend/TransformerEngineNPU.git /root/TransformerEngineNPU

Conversation

momo609 commented Aug 14, 2026

Summary

Files Changed

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

gemini-code-assist Bot Aug 14, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

read-the-docs-community Bot commented Aug 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Documentation build overview

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

floatlibai Aug 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

read-the-docs-community Bot commented Aug 14, 2026 •

edited

Loading

floatlibai Aug 20, 2026 •

edited

Loading