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 #34269381 | 📁 Comparing 2d6a5fb 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.

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

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>

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 2 times, most recently from a283316 to e91f80c Compare August 26, 2026 08:48

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

da83eef

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

momo609 force-pushed the ascendupdate2 branch from 29cfc9d to 5e639b0 Compare August 27, 2026 03:36

wangxiaoxin-sherie added 2 commits August 27, 2026 17:30

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

ba2b0f8

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

fix(npu): complete and pin training dependencies

f56d647

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

momo609 force-pushed the ascendupdate2 branch 3 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 4 times, most recently from 7db37fb to 971e1af Compare August 28, 2026 02:16

feat(npu): support proxied GitHub builds

2d6a5fb

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

momo609 force-pushed the ascendupdate2 branch from a4bfe86 to 2d6a5fb Compare August 28, 2026 03:38

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

Uh oh!

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!

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