diff --git a/.buildkite/README-NPU-CI.md b/.buildkite/README-NPU-CI.md
new file mode 100644
index 000000000..6cd2e7a21
--- /dev/null
+++ b/.buildkite/README-NPU-CI.md
@@ -0,0 +1,111 @@
+# vime NPU CI on Buildkite
+
+The NPU pipeline lives in [`pipeline-npu.yaml`](./pipeline-npu.yaml). It validates
+vime on Ascend NPU hardware.
+
+The NPU suites are behind a **block step** (`:rocket: Run NPU test suites?`):
+click it in the Buildkite UI like GPU CI test suites.
+
+## Pipeline steps
+
+Three steps run in order:
+
+**`pre-commit-npu`** runs the pre-commit gate on all files, always.
+
+**`npu-gate`** is a block step that pauses the pipeline for manual interaction.
+It appears only on PR triggers (`build.source != "schedule"`).  You select which
+NPU suites to run via a multi-select field:
+
+- **`image-build`** — triggers a fresh image build in the `image-build-npu` step;
+  when selected, `smk` is automatically included and all test steps use the newly
+  built image instead of the pre-built default.
+- **`smk`** — runs the smoke test suite.
+- **`nightly`** — runs the nightly test suite.
+
+**`image-build-npu`** the `image-build-npu` step will build and
+push a fresh NPU test image tagged with the current commit. When `image-build` 
+is selected, all test steps generated by `upload-npu-suites` use the newly built
+image instead of the pre-built default image.
+
+This allows testing code changes that require an updated NPU environment
+(e.g., modifications to `docker/Dockerfile.npu`) before they are merged.
+
+**`upload-npu-suites`** reads the `NPU_SUITES` environment variable (or
+`buildkite-agent meta-data get npu-suites` for PR triggers) and generates
+individual test jobs via [`npu_suites.py`](./npu_suites.py).
+
+## Triggers
+
+The pipeline supports two trigger modes:
+
+**PR trigger.** The `npu-gate` block step appears for manual suite selection.
+Suites not selected in the block step are skipped. The `image-build-npu` step
+builds a new image only when `image-build` is chosen in the block.
+
+**Schedule trigger.** There is no block step — the `npu-gate` step is omitted
+entirely. A new image is **always** built, and the suites are determined by the
+`NPU_SUITES` environment variable (set in the scheduled build's pipeline
+configuration), which lists suite names from the `SUITES` dict.
+
+## Adding a test
+
+Suites and test mappings are defined in [`npu_suites.py`](./npu_suites.py). Two
+suites are predefined — `smk` (runs with the `run-ci-npu-smk` label) and `nightly`
+(runs on schedule or with the `run-ci-npu-nightly` label).
+
+Each entry is a 4-tuple:
+
+```python
+("test-qwen3-4B-npu.py", "npu-8", "", {})
+# (test_name, resource_class, extra_args, env_overrides)
+```
+
+- **`test_name`** — the test script under `tests/`, e.g. `test-qwen3-4B-npu.py`.
+- **`resource_class`** — NPU count for the pod (`npu-2`, `npu-4`, `npu-8`,
+  `npu-16`). 
+- **`extra_args`** — extra CLI arguments passed to the test script.
+- **`env_overrides`** — extra environment variables for the test step, e.g.
+  `{"USE_DEEPEP": "1"}`.
+
+A test script should:
+
+1. **Download models and datasets to `HF_HOME`** before training so subsequent
+   steps reuse the cache. The CI sets `HF_HOME=/root/.cache/huggingface`.
+
+2. **Run training** via `train.py` or `train_async.py` with the appropriate
+   NPU-specific arguments.
+
+To register a new suite, add a key to the `SUITES` dict and update
+`selected_suites()` if it should run conditionally. Also add a corresponding
+option in the block step's `fields[].options` list in
+[`pipeline-npu.yaml`](./pipeline-npu.yaml) so it can be selected in the
+Buildkite UI.
+
+To add a test to an existing suite, simply append a new entry tuple to its
+list in `SUITES`.
+
+## Adding or removing a patch
+
+Patches are applied explicitly by [`docker/Dockerfile.npu`](../docker/Dockerfile.npu)
+and reconciled at runtime in the order declared by
+[`docker/npu_patch/series.conf`](../docker/npu_patch/series.conf).
+Each entry declares `target_worktree|image_patch|source_patch`: the image keeps
+the flat `image_patch` under `/opt/npu_patch`, while runtime CI reads
+`source_patch` relative to the current VIME checkout.
+
+When adding a patch file, add the matching Dockerfile COPY/apply operations and
+the `series.conf` entry in the same order. When deleting a patch, remove the
+Dockerfile operation and `series.conf` entry before deleting the patch file;
+the previous image retains the OLD bytes required for runtime revert. Rename
+and reorder changes must likewise update both declarations. Ordinary patch
+changes must not add patch-specific branches to the CI script.
+
+Patch-only changes reuse the default image and are reconciled when the test
+container starts. Changes that add, remove, or upgrade external repositories,
+modify installed dependencies, or otherwise change the Docker image must be
+validated with `image-build`.
+
+The initial rollout of this mechanism requires publishing a new default image
+that contains `/opt/npu_patch/series.conf`, then updating
+`DEFAULT_CI_IMAGE` in `npu_suites.py`. The legacy default image contains patch
+bytes under `/tmp/npu_patch` but not the OLD series required by this reconciler.
diff --git a/.buildkite/npu_suites.py b/.buildkite/npu_suites.py
new file mode 100644
index 000000000..97517f9dd
--- /dev/null
+++ b/.buildkite/npu_suites.py
@@ -0,0 +1,136 @@
+#!/usr/bin/env python3
+"""Emit Buildkite steps for the NPU suites.
+
+Piped into `buildkite-agent pipeline upload` by the npu step in pipeline.yml.
+The suites and their configurations are defined here.
+
+NPU jobs run on the ascend-a3 queue with resource_class determining NPU count
+(e.g., "npu-8" means 8 NPUs).
+
+The selection is read from the NPU_SUITES environment variable. For local
+testing, set NPU_SUITES=smk,nightly instead of having a buildkite-agent on PATH.
+
+stdlib only — runs with the agent host's python3.
+"""
+
+import json
+import os
+import subprocess
+import sys
+
+NPU_QUEUE = "ascend-a3"
+DEFAULT_CI_IMAGE = "quay.io/ascend/vime:vime-latest"
+IMAGE_REGISTRY = "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry"
+IMAGE_NAME = "vime-ci-npu"
+VIME_IMAGE_TAG = os.environ.get("BUILDKITE_COMMIT", "latest")
+BUILDKITE_SOURCE = os.environ.get("BUILDKITE_SOURCE", "")
+
+# (test_name, resource_class, extra_args, env_overrides)
+SUITES = {
+    "smk": [
+        ("test_qwen3_4B_npu.py", "npu-8", "", {}),
+        ("test_qwen3_30B_A3B_npu.py", "npu-16", "", {}),
+        ("test_qwen3_vl_8B_npu.py", "npu-8", "", {}),
+        ("test_glm4.7_30B_A3B_npu.py", "npu-16", "", {}),
+    ],
+    "nightly": [],
+}
+
+
+def _read_suite_values() -> list[str]:
+    raw = os.environ.get("NPU_SUITES")
+    if raw is None:
+        try:
+            raw = subprocess.run(
+                ["buildkite-agent", "meta-data", "get", "npu-suites"],
+                check=True,
+                capture_output=True,
+                text=True,
+            ).stdout
+        except subprocess.CalledProcessError:
+            raw = ""
+    return [v.strip() for v in raw.replace(",", "\n").splitlines()]
+
+
+def _ci_image() -> str:
+    values = _read_suite_values()
+    if ("image-build" in values) or (BUILDKITE_SOURCE == "schedule"):
+        return f"{IMAGE_REGISTRY}/{IMAGE_NAME}:{VIME_IMAGE_TAG}"
+    return DEFAULT_CI_IMAGE
+
+
+def selected_suites() -> list:
+    values = _read_suite_values()
+    unknown = [v for v in values if v and v not in SUITES and v != "image-build"]
+    if unknown:
+        raise SystemExit(f"unknown suite(s) {unknown}; expected {sorted(SUITES)}")
+    if "image-build" in values:
+        # image-build auto-includes smk tests
+        values.append("smk")
+    return [s for s in SUITES if s in values]
+
+
+def npu_step(suite: str, test_name: str, resource_class: str, extra_args: str, env: dict) -> dict:
+    step_env = {
+        "VIME_TEST_ENABLE_INFINITE_RUN": "false",
+        "BUILDKITE_PULL_REQUEST": os.environ.get("BUILDKITE_PULL_REQUEST", "false"),
+        "BUILDKITE_COMMIT": os.environ.get("BUILDKITE_COMMIT", ""),
+        "HF_TOKEN": "${HF_TOKEN}",
+        "HF_ENDPOINT": "https://hf-mirror.com",
+        "ASCEND_RT_VISIBLE_DEVICES": "0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15",
+        "IMAGE_REGISTRY": IMAGE_REGISTRY,
+        "IMAGE_NAME": IMAGE_NAME,
+        "VIME_IMAGE_TAG": VIME_IMAGE_TAG,
+        "HF_HOME": "/root/.cache/huggingface",
+        **env,
+    }
+
+    commands = "\n".join(
+        [
+            'echo "INFO: update NPU environment"',
+            'if [ -n "${BUILDKITE_COMMIT}" ]; then',
+            "  source /workspace/build/buildkite/.buildkite/scripts/update-npu-environment.sh",
+            "fi",
+            "export HF_HUB_OFFLINE=0",
+            f"python tests/{test_name}{' ' + extra_args if extra_args else ''}",
+        ]
+    )
+    command = f"bash -c '{commands}'"
+
+    label = f":fire: {suite}: {test_name}{' ' + extra_args if extra_args else ''}"
+    step = {
+        "label": label,
+        "depends_on": "image-build-npu",
+        "command": command,
+        "agents": {
+            "queue": NPU_QUEUE,
+            "resource_class": resource_class,
+        },
+        "timeout_in_minutes": 180,
+        "image": _ci_image(),
+        "plugins": [
+            {
+                "kubernetes": {
+                    "podSpecPatch": {
+                        "imagePullSecrets": [{"name": "swr-secret"}],
+                    },
+                }
+            }
+        ],
+        "env": step_env,
+    }
+    return step
+
+
+def main() -> None:
+    steps = [npu_step(suite, *entry) for suite in selected_suites() for entry in SUITES[suite]]
+    json_str = json.dumps({"steps": steps}, indent=2)
+
+    print("--- Generated Pipeline JSON:", file=sys.stderr)
+    print(json_str, file=sys.stderr)
+
+    print(json_str)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/.buildkite/pipeline-npu-image.yaml b/.buildkite/pipeline-npu-image.yaml
new file mode 100644
index 000000000..22b9ca82e
--- /dev/null
+++ b/.buildkite/pipeline-npu-image.yaml
@@ -0,0 +1,88 @@
+steps:
+  - label: ":docker: Build and Push NPU Test Image - A3"
+    key: image-build-npu
+    depends_on: 
+      - pre-commit-npu
+      - npu-gate
+    timeout_in_minutes: 240
+    skip: __SKIP_IMAGE_BUILD__
+    agents:
+      queue: "ascend-a3"
+      resource_class: "npu-2"
+    plugins:
+      - kubernetes:
+          metadata:
+            annotations:
+              vault.hashicorp.com/agent-init-first: "true"
+              vault.hashicorp.com/agent-inject: "true"
+              vault.hashicorp.com/agent-inject-perms-ca.pem: "0400"
+              vault.hashicorp.com/agent-inject-perms-cert.pem: "0400"
+              vault.hashicorp.com/agent-inject-perms-config.json: "0400"
+              vault.hashicorp.com/agent-inject-perms-key.pem: "0400"
+              vault.hashicorp.com/agent-inject-secret-ca.pem: internal/data/ascend/buildkitd
+              vault.hashicorp.com/agent-inject-secret-cert.pem: internal/data/ascend/buildkitd
+              vault.hashicorp.com/agent-inject-secret-config.json: internal/data/ascend/buildkitd
+              vault.hashicorp.com/agent-inject-secret-key.pem: internal/data/ascend/buildkitd
+              vault.hashicorp.com/agent-inject-template-ca.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.RootCA }}\n{{- end }}"
+              vault.hashicorp.com/agent-inject-template-cert.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.ClientCaCert }}\n{{- end }}"
+              vault.hashicorp.com/agent-inject-template-config.json: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.dockerConfig }}\n{{- end }}"
+              vault.hashicorp.com/agent-inject-template-key.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.ClientCaKey }}\n{{- end }}"
+              vault.hashicorp.com/agent-pre-populate-only: "true"
+              vault.hashicorp.com/agent-run-as-group: "1000"
+              vault.hashicorp.com/agent-run-as-user: "1000"
+              vault.hashicorp.com/agent-service-account-token-volume-name: token-vol
+              vault.hashicorp.com/role: ascend-gha-runners
+              vault.hashicorp.com/secret-volume-path: /home/user/.docker/
+              vault.hashicorp.com/tls-skip-verify: "true"
+          podSpecPatch:
+            volumes:
+            - name: token-vol
+              projected:
+                defaultMode: 420
+                sources:
+                - serviceAccountToken:
+                    audience: api
+                    expirationSeconds: 600
+                    path: token
+    env:
+      VIME_IMAGE_TAG: "${BUILDKITE_COMMIT}"
+      IMAGE_NAME: "vime-ci-npu"
+      IMAGE_REGISTRY: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry"
+      BUILDKITD_ADDR: "tcp://buildkitd-service.buildkitd:1234"
+    command: |
+      set -ex
+      echo "--- Building and pushing NPU Test Image"
+      echo "Image: $${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VIME_IMAGE_TAG}"
+      echo "buildkitd address: $${BUILDKITD_ADDR}"
+
+      if ! command -v buildctl &> /dev/null; then
+        echo "Installing buildctl..."
+        mkdir -p /tmp/buildkit
+        BUILDKIT_VERSION="v0.29.0"
+        wget -q "https://gh-proxy.test.osinfra.cn/https://github.com/moby/buildkit/releases/download/$${BUILDKIT_VERSION}/buildkit-$${BUILDKIT_VERSION}.linux-arm64.tar.gz" -O /tmp/buildkit.tar.gz
+        tar -xzf /tmp/buildkit.tar.gz -C /tmp/buildkit
+        cp /tmp/buildkit/bin/buildctl /usr/local/bin/
+      fi
+
+      sed -i '/^RUN git config --global http.sslVerify false/i RUN git config --global url."https://gh-proxy.test.osinfra.cn/https://github.com/".insteadOf "https://github.com/"' docker/Dockerfile.npu
+      sed -i '/^# syntax=docker\/dockerfile:1\.7$$/d' docker/Dockerfile.npu
+
+      export DOCKER_CONFIG=/home/user/.docker
+      buildctl \
+        --addr="$${BUILDKITD_ADDR}" \
+        --tlscacert=/home/user/.docker/ca.pem \
+        --tlscert=/home/user/.docker/cert.pem \
+        --tlskey=/home/user/.docker/key.pem \
+        build \
+        --frontend dockerfile.v0 \
+        --local context=. \
+        --local dockerfile=./docker \
+        --opt filename=Dockerfile.npu \
+        --opt build-arg:APTMIRROR=http://cache-service.nginx-pypi-cache.svc.cluster.local:8081 \
+        --opt build-arg:PIP_INDEX_URL=http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
+        --secret id=dockerconfig,src=/home/user/.docker/config.json \
+        --output type=image,name=$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VIME_IMAGE_TAG},push=true \
+        --progress=plain
+
+      echo "--- Image pushed successfully"
+      echo "$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VIME_IMAGE_TAG}"
\ No newline at end of file
diff --git a/.buildkite/pipeline-npu.yaml b/.buildkite/pipeline-npu.yaml
new file mode 100644
index 000000000..244e860c6
--- /dev/null
+++ b/.buildkite/pipeline-npu.yaml
@@ -0,0 +1,73 @@
+# Usage:  buildkite-agent pipeline upload .buildkite/pipeline-npu.yaml
+#
+# This pipeline is for running NPU tests on PRs.
+# Scheduled builds always build a fresh image. PR builds reuse DEFAULT_CI_IMAGE
+# unless image-build is selected; patch-only changes are reconciled at runtime.
+
+steps:
+  - group: ":pipeline: Run NPU test"
+    key: npu-tests
+    steps:
+      # Before run any test, check pre-commit
+      - label: ":lint-roller: pre-commit-npu"
+        key: pre-commit-npu
+        agents:
+          queue: small_cpu_queue_premerge
+        timeout_in_minutes: 15
+        retry:
+          automatic:
+            - exit_status: -1  # agent lost (fresh instance failed to boot)
+              limit: 2
+        command: |
+          docker run --rm \
+            -e GIT_CONFIG_PARAMETERS="'safe.directory=/workspace'" \
+            -v "$$PWD:/workspace" -w /workspace \
+            python:3.11 bash -c '
+              set -euo pipefail
+              pip install -q pre-commit
+              pre-commit run --all-files --show-diff-on-failure --color=always
+            '
+
+      - block: ":rocket: Run NPU test suites?"
+        key: npu-gate
+        depends_on: pre-commit-npu
+        if: build.source != "schedule"
+        blocked_state: passed
+        prompt: "Select the npu suites to run."
+        fields:
+          - select: "NPU suites"
+            key: npu-suites
+            multiple: true
+            required: true
+            options:
+              - label: "run-npu-ci-image-build (test new image, auto include smk)"
+                value: image-build
+              - label: "run-npu-ci-smk"
+                value: smk
+              - label: "run-npu-ci-nightly"
+                value: nightly
+
+      - label: ":pipeline: upload NPU suites"
+        key: upload-npu-suites
+        depends_on: 
+          - npu-gate
+        agents:
+          queue: "ascend-a3"
+          resource_class: "npu-2"
+        timeout_in_minutes: 180
+        command:  |
+          NPU_SUITES=$$(buildkite-agent meta-data get "npu-suites" --default "")
+          if [[ "$$BUILDKITE_SOURCE" == "schedule" ]]; then
+            echo "Scheduled build — building the image."
+            echo "BUILDKITE_SOURCE: $$BUILDKITE_SOURCE"
+            SKIP_IMAGE_BUILD=false
+          elif [[ "$$NPU_SUITES" == *"image-build"* ]]; then
+            echo "image-build selected — building the image."
+            echo "NPU_SUITES: $$NPU_SUITES"
+            SKIP_IMAGE_BUILD=false
+          else
+            echo "Skipping image build because image-build is not present in the npu-suites."
+            SKIP_IMAGE_BUILD=true
+          fi
+          sed -e "s/__SKIP_IMAGE_BUILD__/$${SKIP_IMAGE_BUILD}/g" .buildkite/pipeline-npu-image.yaml | buildkite-agent pipeline upload
+          python .buildkite/npu_suites.py | buildkite-agent pipeline upload
diff --git a/.buildkite/scripts/update-npu-environment.sh b/.buildkite/scripts/update-npu-environment.sh
new file mode 100644
index 000000000..761f2406d
--- /dev/null
+++ b/.buildkite/scripts/update-npu-environment.sh
@@ -0,0 +1,238 @@
+#!/bin/bash
+# Purpose: Updates an NPU test container to match the requested VIME commit.
+#          - Reads the image's persisted OLD patch series and exact patch bytes
+#          - Updates VIME, then reconciles OLD -> NEW in declared series order
+#          - Installs the current VIME checkout and normalizes visible devices
+# Usage: Called by Buildkite pipeline during NPU test runs
+set -e -o pipefail
+
+VIME_DIR="${VIME_DIR:-/root/vime}"
+VIME_NPU_PATCH_STATE_DIR="${VIME_NPU_PATCH_STATE_DIR:-/opt/npu_patch}"
+VIME_NPU_PATCH_SOURCE_ROOT="${VIME_NPU_PATCH_SOURCE_ROOT:-${VIME_DIR}}"
+PATCH_SERIES_RELATIVE_PATH="docker/npu_patch/series.conf"
+
+sha256_stdin() {
+    if command -v sha256sum >/dev/null 2>&1; then
+        sha256sum | awk '{print $1}'
+    else
+        shasum -a 256 | awk '{print $1}'
+    fi
+}
+
+sha256_file() {
+    local path="$1"
+    if command -v sha256sum >/dev/null 2>&1; then
+        sha256sum "$path" | awk '{print $1}'
+    else
+        shasum -a 256 "$path" | awk '{print $1}'
+    fi
+}
+
+SERIES_ENTRIES=()
+
+load_series() {
+    local series_file="$1"
+    local line target image_patch source_patch extra
+
+    SERIES_ENTRIES=()
+    if [ ! -f "$series_file" ]; then
+        echo "ERROR: Patch series not found: $series_file" >&2
+        return 1
+    fi
+
+    while IFS= read -r line || [ -n "$line" ]; do
+        if [[ "$line" =~ ^[[:space:]]*$ || "$line" =~ ^[[:space:]]*# ]]; then
+            continue
+        fi
+
+        IFS='|' read -r target image_patch source_patch extra <<< "$line"
+        if [ -z "$target" ] || [ -z "$image_patch" ] || [ -z "$source_patch" ] || [ -n "$extra" ]; then
+            echo "ERROR: Invalid patch series entry: $line" >&2
+            return 1
+        fi
+        if [[ "$image_patch" = */* ]]; then
+            echo "ERROR: Image patch must be a flat file name: $image_patch" >&2
+            return 1
+        fi
+        if [[ "$source_patch" = /* || "/$source_patch/" = *"/../"* ]]; then
+            echo "ERROR: Source patch must stay under the VIME root: $source_patch" >&2
+            return 1
+        fi
+        SERIES_ENTRIES+=("${target}|${image_patch}|${source_patch}")
+    done < "$series_file"
+}
+
+validate_series_entry() {
+    local target="$1"
+    local patch_path="$2"
+
+    if ! git -C "$target" rev-parse --is-inside-work-tree >/dev/null 2>&1; then
+        echo "ERROR: Patch target is not a Git worktree: $target" >&2
+        return 1
+    fi
+    if [ ! -f "$patch_path" ]; then
+        echo "ERROR: Patch file not found: $patch_path" >&2
+        return 1
+    fi
+}
+
+series_digest() {
+    local series_file="$1"
+    local patch_root="$2"
+    local path_kind="$3"
+    local entry target image_patch source_patch patch_ref patch_path patch_sha
+
+    load_series "$series_file"
+    for entry in "${SERIES_ENTRIES[@]}"; do
+        IFS='|' read -r target image_patch source_patch <<< "$entry"
+        case "$path_kind" in
+            image) patch_ref="$image_patch" ;;
+            source) patch_ref="$source_patch" ;;
+            *) echo "ERROR: Unknown patch path kind: $path_kind" >&2; return 1 ;;
+        esac
+        patch_path="${patch_root}/${patch_ref}"
+        if [ ! -f "$patch_path" ]; then
+            echo "ERROR: Patch file not found: $patch_path" >&2
+            return 1
+        fi
+    done
+
+    {
+        printf 'vime-npu-patch-series-v1\0'
+        for entry in "${SERIES_ENTRIES[@]}"; do
+            IFS='|' read -r target image_patch source_patch <<< "$entry"
+            if [ "$path_kind" = "image" ]; then
+                patch_ref="$image_patch"
+            else
+                patch_ref="$source_patch"
+            fi
+            patch_path="${patch_root}/${patch_ref}"
+            patch_sha=$(sha256_file "$patch_path")
+            printf '%s\0%s\0%s\0%s\0' "$target" "$image_patch" "$source_patch" "$patch_sha"
+        done
+    } | sha256_stdin
+}
+
+apply_series() {
+    local series_file="$1"
+    local source_root="$2"
+    local entry target image_patch source_patch patch_path
+
+    load_series "$series_file"
+    for entry in "${SERIES_ENTRIES[@]}"; do
+        IFS='|' read -r target image_patch source_patch <<< "$entry"
+        patch_path="${source_root}/${source_patch}"
+        validate_series_entry "$target" "$patch_path"
+        echo "INFO: Applying $source_patch to $target"
+        git -C "$target" apply --whitespace=nowarn "$patch_path"
+    done
+}
+
+revert_series() {
+    local series_file="$1"
+    local image_root="$2"
+    local i entry target image_patch source_patch patch_path
+
+    load_series "$series_file"
+    for ((i=${#SERIES_ENTRIES[@]}-1; i>=0; i--)); do
+        entry="${SERIES_ENTRIES[$i]}"
+        IFS='|' read -r target image_patch source_patch <<< "$entry"
+        patch_path="${image_root}/${image_patch}"
+        validate_series_entry "$target" "$patch_path"
+        echo "INFO: Reverting $image_patch from $target"
+        git -C "$target" apply --reverse --whitespace=nowarn "$patch_path"
+    done
+}
+
+reconcile_series() {
+    local old_series="$1"
+    local old_root="$2"
+    local new_series="$3"
+    local new_root="$4"
+    local old_digest new_digest
+
+    old_digest=$(series_digest "$old_series" "$old_root" image)
+    new_digest=$(series_digest "$new_series" "$new_root" source)
+    if [ "$old_digest" = "$new_digest" ]; then
+        echo "INFO: Patch series is unchanged"
+        return
+    fi
+
+    revert_series "$old_series" "$old_root"
+    apply_series "$new_series" "$new_root"
+}
+
+update_vime_code() {
+    echo "INFO: Updating VIME code..."
+
+    if [ -n "${BUILDKITE_COMMIT:-}" ]; then
+        echo "INFO: Fetching and checking out commit ${BUILDKITE_COMMIT}"
+        git -C "$VIME_DIR" fetch origin "${BUILDKITE_COMMIT}"
+        git -C "$VIME_DIR" checkout "${BUILDKITE_COMMIT}"
+    else
+        echo "INFO: BUILDKITE_COMMIT not set, skipping code update"
+    fi
+}
+
+install_vime_code() {
+    pip install -e "$VIME_DIR" --no-deps --break-system-packages || pip install -e "$VIME_DIR" --no-deps
+}
+
+sort_ascend_visible_devices() {
+    export ASCEND_VISIBLE_DEVICES="${ASCEND_VISIBLE_DEVICES:-${ASCEND_RT_VISIBLE_DEVICES:-}}"
+    echo "Value: ${ASCEND_VISIBLE_DEVICES}"
+    if [ -n "${ASCEND_VISIBLE_DEVICES}" ]; then
+        SORTED_DEVICES=$(echo "${ASCEND_VISIBLE_DEVICES}" | tr ',' '\n' | sort -n | tr '\n' ',')
+        SORTED_DEVICES=${SORTED_DEVICES%,}
+        export ASCEND_VISIBLE_DEVICES=$SORTED_DEVICES
+        echo "Sorted ASCEND_VISIBLE_DEVICES: $ASCEND_VISIBLE_DEVICES"
+    fi
+}
+
+main() {
+    local old_series="${VIME_NPU_PATCH_STATE_DIR}/series.conf"
+    local new_series="${VIME_NPU_PATCH_SOURCE_ROOT}/${PATCH_SERIES_RELATIVE_PATH}"
+
+    echo "=== Step 1: Sort ASCEND_VISIBLE_DEVICES ==="
+    sort_ascend_visible_devices
+
+    echo "=== Step 2: Update VIME code ==="
+    update_vime_code
+
+    if [ ! -f "$old_series" ]; then
+        echo "ERROR: The selected image does not contain NPU patch state: $old_series" >&2
+        echo "ERROR: Build a patch-state-enabled NPU image before running this commit." >&2
+        return 1
+    fi
+
+    echo "=== Step 3: Reconcile image patches with current VIME patches ==="
+    reconcile_series "$old_series" "$VIME_NPU_PATCH_STATE_DIR" "$new_series" "$VIME_NPU_PATCH_SOURCE_ROOT"
+
+    echo "=== Step 4: Install current VIME code ==="
+    install_vime_code
+
+    echo "INFO: NPU environment update completed successfully"
+}
+
+if [[ "${BASH_SOURCE[0]}" == "$0" ]]; then
+    case "${1:-}" in
+        series-digest)
+            if [ "$#" -ne 4 ]; then
+                echo "Usage: $0 series-digest SERIES_FILE PATCH_ROOT image|source" >&2
+                exit 2
+            fi
+            series_digest "$2" "$3" "$4"
+            exit
+            ;;
+        reconcile)
+            if [ "$#" -ne 5 ]; then
+                echo "Usage: $0 reconcile OLD_SERIES OLD_ROOT NEW_SERIES NEW_ROOT" >&2
+                exit 2
+            fi
+            reconcile_series "$2" "$3" "$4" "$5"
+            exit
+            ;;
+    esac
+fi
+
+main
diff --git a/.gitattributes b/.gitattributes
new file mode 100644
index 000000000..9812ceb1f
--- /dev/null
+++ b/.gitattributes
@@ -0,0 +1 @@
+*.patch text eol=lf
diff --git a/docker/Dockerfile.npu b/docker/Dockerfile.npu
new file mode 100644
index 000000000..f9e42c9aa
--- /dev/null
+++ b/docker/Dockerfile.npu
@@ -0,0 +1,148 @@
+# syntax=docker/dockerfile:1.7
+
+ARG BASE_IMAGE=quay.io/atlas-ci/vllm-ascend
+ARG BASE_IMAGE_TAG=v0.28.0-fd81546-a3
+FROM ${BASE_IMAGE}:${BASE_IMAGE_TAG}
+
+SHELL ["/bin/bash", "-o", "pipefail", "-c"]
+WORKDIR /root
+
+ARG MEGATRON_COMMIT=1dcf0dafa884ad52ffb243625717a3471643e087
+ARG MEGATRON_BRIDGE_COMMIT=3fd3768045422d0aa5c97e90a4e6c659aea9acb9
+ARG MINDSPEED_COMMIT=fc63de5c48426dd019c3b3f39e65f5bdf56e4086
+ARG MEGATRON_ADAPTOR_COMMIT=15582addff3f3d4680e350826fa70d012b475509
+ARG TRANSFORMER_ENGINE_NPU_COMMIT=d743c83d060d5edc48867ecb9e93ec80d81860e4
+ARG MBRIDGE_COMMIT=89eb10887887bc74853f89a4de258c0702932a1c
+ARG SOC_VERSION="ascend910_9391"
+ARG PIP_INDEX_URL="https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple"
+ARG APTMIRROR=""
+
+ENV SOC_VERSION=$SOC_VERSION \
+    DEBIAN_FRONTEND=noninteractive \
+    PIP_NO_CACHE_DIR=1 \
+    ASCEND_TOOLKIT_HOME=/usr/local/Ascend/ascend-toolkit/latest \
+    ASCEND_OPP_PATH=/usr/local/Ascend/ascend-toolkit/latest/opp \
+    ASCEND_AICPU_PATH=/usr/local/Ascend/ascend-toolkit/latest \
+    ASCEND_HOME_PATH=/usr/local/Ascend/ascend-toolkit/latest \
+    VLLM_ASCEND_ENABLE_NZ=0 \
+    RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 \
+    HCCL_HOST_SOCKET_PORT_RANGE=60000-60050 \
+    HCCL_NPU_SOCKET_PORT_RANGE=61000-61050 \
+    PYTORCH_NPU_ALLOC_CONF=expandable_segments:True \
+    HYDRA_FULL_ERROR=1 \
+    PYTHONPATH=/root/Megatron-Bridge/src:/root/Megatron-LM:/root/MegatronAdaptor:/root/TransformerEngineNPU:/root/vime
+
+# PATCH MAINTENANCE: keep patch COPY/apply operations and
+# docker/npu_patch/series.conf synchronized.
+COPY docker/npu_patch /opt/npu_patch
+COPY docker/patch/latest/megatron.patch /opt/npu_patch/megatron-common.patch
+
+RUN git config --global http.sslVerify false
+
+# System and pip dependencies.
+RUN if [ -n "$APTMIRROR" ];then sed -i "s@^\(deb.*\)https\?://[a-z0-9.-]*\.ubuntu\.com@\1$APTMIRROR@g" /etc/apt/sources.list; \
+    else sed -i '/ports\.ubuntu\.com/ {h;s|ports\.ubuntu\.com|mirrors.tuna.tsinghua.edu.cn|g;G}' /etc/apt/sources.list; fi && \
+    apt-get update && \
+    apt-get install -y --no-install-recommends \
+      build-essential cmake curl git libnuma-dev ninja-build patch rsync wget && \
+    rm -rf /var/lib/apt/lists/* && \
+    pip config set global.index-url ${PIP_INDEX_URL} && \
+    if [ "${PIP_INDEX_URL#http://}" != "${PIP_INDEX_URL}" ]; then pip config set global.trusted-host "$(echo ${PIP_INDEX_URL} | awk -F/ '{print $3}')"; fi && \
+    pip config set global.extra-index-url \
+      "https://download.pytorch.org/whl/cpu/ https://mirrors.huaweicloud.com/ascend/repos/pypi https://mirrors.aliyun.com/pypi/simple/" && \
+    git config --global http.version HTTP/1.1
+
+# vllm and vllm-ascend are installed editable (-e) in the base image; apply the
+# NPU colocate patches directly to their source trees.
+RUN git -C /vllm-workspace/vllm apply --check --whitespace=nowarn \
+      /opt/npu_patch/vllm.patch && \
+    git -C /vllm-workspace/vllm apply --whitespace=nowarn \
+      /opt/npu_patch/vllm.patch
+
+RUN git -C /vllm-workspace/vllm-ascend apply --check --whitespace=nowarn \
+      /opt/npu_patch/vllm-ascend.patch && \
+    git -C /vllm-workspace/vllm-ascend apply --whitespace=nowarn \
+      /opt/npu_patch/vllm-ascend.patch
+
+# Protect the serving stack while installing Vime dependencies.
+RUN python3 -c 'import importlib.metadata as m; names = ["numpy", "ray", "torch", "torch-npu", "torchvision", "transformers", "triton-ascend", "vllm", "vllm-ascend"]; open("/tmp/vime-npu-constraints.txt", "w").write("\n".join(f"{name}=={m.version(name)}" for name in names) + "\n")'
+
+# Training source dependencies.
+RUN git clone https://github.com/NVIDIA/Megatron-LM.git /root/Megatron-LM && \
+    git -C /root/Megatron-LM checkout "${MEGATRON_COMMIT}"
+
+RUN git clone --branch bridge https://github.com/radixark/Megatron-Bridge.git \
+      /root/Megatron-Bridge && \
+    git -C /root/Megatron-Bridge checkout "${MEGATRON_BRIDGE_COMMIT}"
+
+RUN git clone https://gitcode.com/Ascend/MindSpeed.git /root/MindSpeed && \
+    git -C /root/MindSpeed checkout "${MINDSPEED_COMMIT}"
+
+RUN git clone https://gitcode.com/Ascend/MegatronAdaptor.git /root/MegatronAdaptor && \
+    git -C /root/MegatronAdaptor checkout "${MEGATRON_ADAPTOR_COMMIT}" && \
+    git clone https://gitcode.com/Ascend/TransformerEngineNPU.git /root/TransformerEngineNPU && \
+    git -C /root/TransformerEngineNPU checkout "${TRANSFORMER_ENGINE_NPU_COMMIT}"
+
+RUN git clone https://github.com/ISEEKYAN/mbridge.git /root/mbridge && \
+    git -C /root/mbridge checkout "${MBRIDGE_COMMIT}"
+
+# Apply NPU training-stack patches from the build-context snapshot.
+# The NPU Megatron patch is based on the common Vime Megatron patch, so the
+# common patch must be applied first.
+RUN git -C /root/Megatron-LM apply --check --whitespace=nowarn \
+      /opt/npu_patch/megatron-common.patch && \
+    git -C /root/Megatron-LM apply --whitespace=nowarn \
+      /opt/npu_patch/megatron-common.patch && \
+    git -C /root/Megatron-LM apply --check --whitespace=nowarn \
+      /opt/npu_patch/megatron.patch && \
+    git -C /root/Megatron-LM apply --whitespace=nowarn \
+      /opt/npu_patch/megatron.patch && \
+    git -C /root/Megatron-Bridge apply --check --whitespace=nowarn \
+      /opt/npu_patch/megatron-bridge.patch && \
+    git -C /root/Megatron-Bridge apply --whitespace=nowarn \
+      /opt/npu_patch/megatron-bridge.patch && \
+    git -C /root/MindSpeed apply --whitespace=nowarn \
+      /opt/npu_patch/mindspeed.patch
+
+# Megatron-Bridge is used directly from PYTHONPATH. Installing its package
+# metadata would pull CUDA-only dependencies into the Ascend environment.
+RUN pip install --constraint /tmp/vime-npu-constraints.txt --no-build-isolation \
+      "nvidia-modelopt==0.46.0" "nvdlfw-inspect==0.2.2" && \
+    pip install --no-deps --no-build-isolation -e /root/mbridge && \
+    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
+
+# Defaults to the ascend branch for local builds. Release workflows should
+# pass an immutable commit SHA for reproducible images.
+ARG VIME_COMMIT=ascend
+RUN test -n "${VIME_COMMIT}" && \
+    git init /root/vime && \
+    git -C /root/vime remote add origin https://github.com/vllm-project/vime.git && \
+    git -C /root/vime fetch --depth 1 origin "${VIME_COMMIT}" && \
+    git -C /root/vime checkout --detach FETCH_HEAD
+
+# ring_flash_attn is a CUDA extension and is not used on Ascend.
+RUN pip install \
+      --constraint /tmp/vime-npu-constraints.txt \
+      --requirement /root/vime/requirements.txt && \
+    pip install --no-deps --no-build-isolation -e /root/vime
+
+# Vime currently imports torch_memory_saver from its training actor.
+RUN git clone --depth 1 --branch 2026.6.0 \
+      https://github.com/sgl-project/sgl-kernel-npu.git /root/sgl-kernel-npu && \
+    cd /root/sgl-kernel-npu/contrib/torch_memory_saver/python && \
+    python3 setup.py bdist_wheel && \
+    python3 -m pip install --no-deps \
+      dist/torch_memory_saver-*.whl && \
+    cd /root && \
+    rm -rf /root/sgl-kernel-npu
+
+# Minimal import check.
+RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+    python3 -c 'import megatron, mindspeed, megatron_adaptor, transformer_engine, torch_memory_saver, vime, vllm, vllm_ascend;'
+
+WORKDIR /root/vime
+ENTRYPOINT []
+CMD ["/bin/bash"]
diff --git a/docker/npu_patch/README.md b/docker/npu_patch/README.md
new file mode 100644
index 000000000..e81b1daaf
--- /dev/null
+++ b/docker/npu_patch/README.md
@@ -0,0 +1,155 @@
+# Vime NPU Patch Installation Guide
+
+This guide provides instructions for installing Vime with NPU support, including all required dependencies and patches.
+
+> S7 closeout (2026-09-09): retain Ascend #385 (training stack) and #396
+> (torch_dist/ref-load), with native HF loading and main's shared orchestration.
+> Revert #409 (`f5b84916`) and its follow-up Qwen3.5 NPU adaptations; defer that
+> model to the next stage in a fresh, matched environment. Main's Qwen3.5 model
+> code is retained. Existing Qwen3-4B, Qwen3-30B-A3B, Qwen3-VL-8B and the 30B
+> torch_dist/ref-load run passed before the Qwen3.5 environment changes; this
+> does not certify a fresh image or a post-revert E2E run. No installed packages
+> or vendor source trees are rolled back as part of this source-only closeout.
+> Post-revert checks: 183 grouped CPU tests passed. Common → NPU Megatron
+> patches and the reverted Bridge patch pass apply checks on their pinned
+> clean source revisions. Serving patches and `docker/patch/latest` are unchanged.
+
+## Component Version Mapping
+
+| Component       | Version/Commit                           | Source                                                                                                              |
+| --------------- | ---------------------------------------- | ------------------------------------------------------------------------------------------------------------------- |
+| vime            | main                                     | [GitHub](https://github.com/vllm-project/vime/tree/main)                                                            |
+| vLLM | e6bfe03ad73a3330cb427885aa90d97a12e1c704 + NPU patch | S6 serving baseline, retained for S7 |
+| vLLM-Ascend | fd815467c221ee600137f6bdd53fe354d5e7c999 + NPU patch | S6 serving baseline, retained for S7 |
+| Megatron-Bridge | 3fd3768045422d0aa5c97e90a4e6c659aea9acb9 | [GitHub](https://github.com/radixark/Megatron-Bridge)                                                               |
+| Megatron-LM     | 1dcf0dafa884ad52ffb243625717a3471643e087 | [GitHub](https://github.com/NVIDIA/Megatron-LM)                                                                     |
+| MegatronAdaptor | 15582addff3f3d4680e350826fa70d012b475509 | [GitCode](https://gitcode.com/Ascend/MegatronAdaptor)                                                               |
+| TransformerEngineNPU | d743c83d060d5edc48867ecb9e93ec80d81860e4 | [GitCode](https://gitcode.com/Ascend/TransformerEngineNPU)                                                          |
+| MindSpeed       | fc63de5c48426dd019c3b3f39e65f5bdf56e4086 | [GitCode](https://gitcode.com/Ascend/MindSpeed)                                                                     |
+| HDK             | 25.3.RC1                                 | [Ascend](https://www.hiascend.com/hardware/firmware-drivers/commercial?product=7\&model=33)                         |
+| CANN            | 9.0.0                                    | [Ascend](https://www.hiascend.com/developer/download/community/result?module=cann\&cann=9.0.0\&product=7\&model=33) |
+
+## Preparing the Running Environment
+
+Run the steps below in a Python 3.12 environment with CANN 9.0.0. A
+`quay.io/ascend/vllm-ascend:nightly-main-a3` container can be used as the base.
+
+```bash
+export WORKSPACE=/root
+cd "${WORKSPACE}"
+```
+
+Vime's Ascend NPU adaptation lives on the **`ascend`** branch, so clone that
+branch (not `main`):
+
+```bash
+git clone --branch ascend https://github.com/vllm-project/vime.git "${WORKSPACE}/vime"
+export PATCH_DIR="${WORKSPACE}/vime/docker/npu_patch"
+```
+
+#### 1. Megatron-Bridge (legacy build dependency, not the native loader)
+
+The source PR used this via `PYTHONPATH` (no editable install) and required
+`nvidia-modelopt`. This is not a prerequisite for Vime's native HF loader;
+whether to retain it in the S7 image remains under review.
+
+```bash
+export MEGATRON_BRIDGE_COMMIT=3fd3768045422d0aa5c97e90a4e6c659aea9acb9
+export MBRIDGE_COMMIT=89eb10887887bc74853f89a4de258c0702932a1c
+pip install "git+https://github.com/ISEEKYAN/mbridge.git@${MBRIDGE_COMMIT}" --no-deps
+git clone --branch bridge https://github.com/radixark/Megatron-Bridge.git "${WORKSPACE}/Megatron-Bridge"
+git -C "${WORKSPACE}/Megatron-Bridge" checkout "${MEGATRON_BRIDGE_COMMIT}"
+
+git -C "${WORKSPACE}/Megatron-Bridge" apply --whitespace=nowarn "${PATCH_DIR}/megatron-bridge.patch"
+
+pip install --no-build-isolation "nvidia-modelopt[torch]>=0.37.0"
+```
+
+#### 2. Megatron-LM
+
+```bash
+export MEGATRON_COMMIT=1dcf0dafa884ad52ffb243625717a3471643e087
+git clone https://github.com/NVIDIA/Megatron-LM.git "${WORKSPACE}/Megatron-LM"
+git -C "${WORKSPACE}/Megatron-LM" checkout "${MEGATRON_COMMIT}"
+
+git -C "${WORKSPACE}/Megatron-LM" apply --whitespace=nowarn "${WORKSPACE}/vime/docker/patch/latest/megatron.patch"
+git -C "${WORKSPACE}/Megatron-LM" apply --whitespace=nowarn "${PATCH_DIR}/megatron.patch"
+
+pip install --no-deps --no-build-isolation -e "${WORKSPACE}/Megatron-LM"
+```
+
+#### 3. MegatronAdaptor and TransformerEngineNPU
+
+The NPU training stack now uses the two source repositories directly. The mainline Megatron patch is applied first; `docker/npu_patch/megatron.patch` contains only the NPU-specific changes rebased onto that mainline patch:
+
+pip install --no-deps --no-build-isolation -e ${WORKSPACE}/MegatronAdaptor
+pip install --no-deps --no-build-isolation -e ${WORKSPACE}/TransformerEngineNPU
+
+Do not install the CUDA TransformerEngine package in the same environment.
+
+#### 4. MindSpeed
+
+```bash
+export MINDSPEED_COMMIT=fc63de5c48426dd019c3b3f39e65f5bdf56e4086
+git clone https://gitcode.com/Ascend/MindSpeed.git "${WORKSPACE}/MindSpeed"
+git -C "${WORKSPACE}/MindSpeed" checkout "${MINDSPEED_COMMIT}"
+
+git -C "${WORKSPACE}/MindSpeed" apply --whitespace=nowarn "${PATCH_DIR}/mindspeed.patch"
+
+pip install --no-deps --no-build-isolation -e "${WORKSPACE}/MindSpeed"
+```
+
+#### 5. Vime
+
+```bash
+pip install -r "${WORKSPACE}/vime/requirements.txt"
+pip install "vllm-router>=0.1.14"
+pip install --no-deps --no-build-isolation -e "${WORKSPACE}/vime"
+```
+
+The NPU training region and optimizer state use Ascend `torch_memory_saver`.
+Retain the working build in an existing environment; the source build recipe is:
+
+```bash
+git clone --branch 2026.6.0 https://github.com/sgl-project/sgl-kernel-npu.git "${WORKSPACE}/sgl-kernel-npu"
+cd "${WORKSPACE}/sgl-kernel-npu"
+bash build.sh -a kernels
+bash build.sh -a memory-saver
+pip install --no-deps output/torch_memory_saver-0.0.8-cp312-cp312-linux_aarch64.whl
+```
+
+#### 5. Install vLLM and vLLM Ascend
+
+```bash
+export VLLM_COMMIT=e6bfe03ad73a3330cb427885aa90d97a12e1c704
+export VLLM_ASCEND_COMMIT=fd815467c221ee600137f6bdd53fe354d5e7c999
+
+git clone https://github.com/vllm-project/vllm.git "${WORKSPACE}/vllm"
+git -C "${WORKSPACE}/vllm" checkout "${VLLM_COMMIT}"
+VLLM_TARGET_DEVICE=empty pip install -v -e "${WORKSPACE}/vllm"
+
+git clone https://github.com/vllm-project/vllm-ascend.git "${WORKSPACE}/vllm-ascend"
+git -C "${WORKSPACE}/vllm-ascend" checkout "${VLLM_ASCEND_COMMIT}"
+git -C "${WORKSPACE}/vllm-ascend" submodule update --init --recursive
+pip install -v -e "${WORKSPACE}/vllm-ascend"
+```
+
+Apply `vllm.patch` and `vllm-ascend.patch` to those exact revisions before
+validation. Do not replace the existing source trees during conflict resolution.
+
+For image patch reconciliation, persist the common Megatron patch as
+`/opt/npu_patch/megatron-common.patch`. `series.conf` applies common → NPU and
+reverts in reverse order. This reconciles patch bytes, not repository versions;
+an old Megatron checkout cannot be upgraded by the patch reconciler alone.
+
+## Additional Dependencies
+
+The source PR specified the following versions. They are not an instruction to
+upgrade the existing S6 environment; in particular, validate the new NPU kernel
+requirements before changing torch-npu:
+
+```shell
+pip install torch-npu==2.10.0
+pip install torchvision==0.25.0
+pip install numpy==1.26.4
+```
diff --git a/docker/npu_patch/megatron-bridge.patch b/docker/npu_patch/megatron-bridge.patch
new file mode 100644
index 000000000..adb77c908
--- /dev/null
+++ b/docker/npu_patch/megatron-bridge.patch
@@ -0,0 +1,375 @@
+diff --git a/src/megatron/bridge/models/conversion/param_mapping.py b/src/megatron/bridge/models/conversion/param_mapping.py
+index a0421273..7f9203ab 100644
+--- a/src/megatron/bridge/models/conversion/param_mapping.py
++++ b/src/megatron/bridge/models/conversion/param_mapping.py
+@@ -1097,15 +1097,20 @@ class AutoMapping(MegatronParamMapping[torch.Tensor]):
+             "LinearCrossEntropyModule",
+             "TEColumnParallelLinear",
++            "MindSpeedTEColumnParallelLinear",
+             "TELayerNormColumnParallelLinear",
++            "MindSpeedTELayerNormColumnParallelLinear",
+             "TEColumnParallelGroupedLinear",
++            "MindSpeedTEColumnParallelGroupedLinear",
+             "VocabParallelEmbedding",
+             "DotProductAttention",  # for attention sink only
+             "TEDotProductAttention",  # for attention sink only
++            "MindSpeedTEDotProductAttention",
+         },
+         "row": {
+             "RowParallelLinear",
+             "TERowParallelLinear",
+             "TERowParallelGroupedLinear",
++            "MindSpeedTERowParallelGroupedLinear",
+         },
+         "replicated": {
+             # Normalization layers
+@@ -1176,7 +1180,7 @@ class AutoMapping(MegatronParamMapping[torch.Tensor]):
+         # Handle fused modules like TELayerNormColumnParallelLinear
+         # These modules have both column-parallel weights (weight, bias)
+         # and replicated layer norm weights (layer_norm_weight, layer_norm_bias)
+-        if module_type == "TELayerNormColumnParallelLinear":
++        if module_type == "TELayerNormColumnParallelLinear" or module_type == "MindSpeedTELayerNormColumnParallelLinear":
+             # Check the actual parameter name to determine the correct parallelism type
+             if self.megatron_param and (
+                 self.megatron_param.endswith("layer_norm_weight") or self.megatron_param.endswith("layer_norm_bias")
+@@ -1207,7 +1211,7 @@ class AutoMapping(MegatronParamMapping[torch.Tensor]):
+             return "replicated"
+ 
+         # Check parallel_mode for TELinear
+-        if module_type == "TELinear":
++        if module_type == "TELinear" or module_type == "MindSpeedTELinear":
+             if module.parallel_mode == "column":
+                 return "column"
+             elif module.parallel_mode == "row":
+diff --git a/src/megatron/bridge/models/qwen/__init__.py b/src/megatron/bridge/models/qwen/__init__.py
+index b3656b6d..382845cc 100644
+--- a/src/megatron/bridge/models/qwen/__init__.py
++++ b/src/megatron/bridge/models/qwen/__init__.py
+@@ -15,7 +15,7 @@
+ from megatron.bridge.models.qwen.qwen2_bridge import Qwen2Bridge  # noqa: F401
+ from megatron.bridge.models.qwen.qwen3_bridge import Qwen3Bridge  # noqa: F401
+ from megatron.bridge.models.qwen.qwen3_moe_bridge import Qwen3MoEBridge  # noqa: F401
+-from megatron.bridge.models.qwen.qwen3_next_bridge import Qwen3NextBridge
++# from megatron.bridge.models.qwen.qwen3_next_bridge import Qwen3NextBridge
+ from megatron.bridge.models.qwen.qwen_provider import (
+     Qwen2ModelProvider,
+     Qwen2ModelProvider1P5B,
+@@ -32,8 +32,8 @@ from megatron.bridge.models.qwen.qwen_provider import (
+     Qwen3MoEModelProvider,
+     Qwen3MoEModelProvider30B_A3B,
+     Qwen3MoEModelProvider235B_A22B,
+-    Qwen3NextModelProvider,
+-    Qwen3NextModelProvider80B_A3B,
++    # Qwen3NextModelProvider,
++    # Qwen3NextModelProvider80B_A3B,
+     Qwen25ModelProvider1P5B,
+     Qwen25ModelProvider3B,
+     Qwen25ModelProvider7B,
+@@ -67,6 +67,6 @@ __all__ = [
+     "Qwen3MoEModelProvider",
+     "Qwen3MoEModelProvider30B_A3B",
+     "Qwen3MoEModelProvider235B_A22B",
+-    "Qwen3NextModelProvider",
+-    "Qwen3NextModelProvider80B_A3B",
++    # "Qwen3NextModelProvider",
++    # "Qwen3NextModelProvider80B_A3B",
+ ]
+diff --git a/src/megatron/bridge/models/qwen/qwen_provider.py b/src/megatron/bridge/models/qwen/qwen_provider.py
+index 775b9765..6200103c 100644
+--- a/src/megatron/bridge/models/qwen/qwen_provider.py
++++ b/src/megatron/bridge/models/qwen/qwen_provider.py
+@@ -18,9 +18,9 @@ from typing import TYPE_CHECKING, Callable, Optional
+ 
+ import torch
+ import torch.nn.functional as F
+-from megatron.core.models.gpt.experimental_attention_variant_module_specs import (
+-    get_transformer_block_with_experimental_attention_variant_spec,
+-)
++# from megatron.core.models.gpt.experimental_attention_variant_module_specs import (
++#     get_transformer_block_with_experimental_attention_variant_spec,
++# )
+ from megatron.core.transformer.spec_utils import ModuleSpec
+ 
+ from megatron.bridge.models.gpt_provider import GPTModelProvider
+@@ -430,53 +430,53 @@ class Qwen3MoEModelProvider235B_A22B(Qwen3MoEModelProvider):
+ # =============================================================================
+ 
+ 
+-@dataclass
+-class Qwen3NextModelProvider(Qwen3MoEModelProvider):
+-    """Base provider for Qwen 3 Next Models."""
+-
+-    transformer_layer_spec: ModuleSpec | Callable[["GPTModelProvider"], ModuleSpec] = (
+-        get_transformer_block_with_experimental_attention_variant_spec
+-    )
+-
+-    layernorm_zero_centered_gamma: bool = True  # Zero-centered RMSNorm
+-    kv_channels: int | None = 256
+-    num_query_groups: int = 2
+-    seq_length: int = 262144  # 256k tokens
+-    rotary_base: float = 10000000.0
+-    rotary_percent: float = 0.25  # 25% of the hidden size is used for RoPE
+-    attention_output_gate: bool = True  # Gated Attention
+-
+-    # MoE specific parameters
+-    num_moe_experts: int = 512
+-    moe_router_topk: int = 10  # 10 routed experts per token
+-    moe_shared_expert_gate: bool = True  # Qwen3-Next uses a gate for the shared expert
+-    moe_router_dtype: str = "fp32"
+-    moe_router_load_balancing_type: str = "global_aux_loss"  # Qwen3-Next uses global aux loss for load balancing
+-
+-    # Linear Attention specific parameters
+-    experimental_attention_variant: str = "gated_delta_net"  # Gated Delta Net used in 75% of the model layers
+-    linear_attention_freq: int | list[int] = 4  # 1 gated standard attention layer per 4 layers
+-    linear_conv_kernel_dim: int = 4
+-    linear_key_head_dim: int = 128
+-    linear_value_head_dim: int = 128
+-    linear_num_key_heads: int = 16
+-    linear_num_value_heads: int = 32
+-
+-    # Checkpointing
+-    hetereogenous_dist_checkpoint: bool = True
+-
+-
+-@dataclass
+-class Qwen3NextModelProvider80B_A3B(Qwen3NextModelProvider):
+-    """
+-    Provider for Qwen 3 Next 80B-A3B: https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Instruct and https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Thinking
+-    """
+-
+-    num_layers: int = 48
+-    hidden_size: int = 2048
+-    num_attention_heads: int = 16
+-    num_query_groups: int = 2
+-    ffn_hidden_size: int = 5120
+-    moe_ffn_hidden_size: int = 512
+-    moe_shared_expert_intermediate_size: int = 512
+-    mtp_num_layers: Optional[int] = None
++# @dataclass
++# class Qwen3NextModelProvider(Qwen3MoEModelProvider):
++#     """Base provider for Qwen 3 Next Models."""
++
++#     transformer_layer_spec: ModuleSpec | Callable[["GPTModelProvider"], ModuleSpec] = (
++#         get_transformer_block_with_experimental_attention_variant_spec
++#     )
++
++#     layernorm_zero_centered_gamma: bool = True  # Zero-centered RMSNorm
++#     kv_channels: int | None = 256
++#     num_query_groups: int = 2
++#     seq_length: int = 262144  # 256k tokens
++#     rotary_base: float = 10000000.0
++#     rotary_percent: float = 0.25  # 25% of the hidden size is used for RoPE
++#     attention_output_gate: bool = True  # Gated Attention
++
++#     # MoE specific parameters
++#     num_moe_experts: int = 512
++#     moe_router_topk: int = 10  # 10 routed experts per token
++#     moe_shared_expert_gate: bool = True  # Qwen3-Next uses a gate for the shared expert
++#     moe_router_dtype: str = "fp32"
++#     moe_router_load_balancing_type: str = "global_aux_loss"  # Qwen3-Next uses global aux loss for load balancing
++
++#     # Linear Attention specific parameters
++#     experimental_attention_variant: str = "gated_delta_net"  # Gated Delta Net used in 75% of the model layers
++#     linear_attention_freq: int | list[int] = 4  # 1 gated standard attention layer per 4 layers
++#     linear_conv_kernel_dim: int = 4
++#     linear_key_head_dim: int = 128
++#     linear_value_head_dim: int = 128
++#     linear_num_key_heads: int = 16
++#     linear_num_value_heads: int = 32
++
++#     # Checkpointing
++#     hetereogenous_dist_checkpoint: bool = True
++
++
++# @dataclass
++# class Qwen3NextModelProvider80B_A3B(Qwen3NextModelProvider):
++#     """
++#     Provider for Qwen 3 Next 80B-A3B: https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Instruct and https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Thinking
++#     """
++
++#     num_layers: int = 48
++#     hidden_size: int = 2048
++#     num_attention_heads: int = 16
++#     num_query_groups: int = 2
++#     ffn_hidden_size: int = 5120
++#     moe_ffn_hidden_size: int = 512
++#     moe_shared_expert_intermediate_size: int = 512
++#     mtp_num_layers: Optional[int] = None
+diff --git a/src/megatron/bridge/models/qwen_vl/qwen35_vl_provider.py b/src/megatron/bridge/models/qwen_vl/qwen35_vl_provider.py
+index 6c74827c..75973903 100644
+--- a/src/megatron/bridge/models/qwen_vl/qwen35_vl_provider.py
++++ b/src/megatron/bridge/models/qwen_vl/qwen35_vl_provider.py
+@@ -36,9 +36,9 @@ from typing import Any, Callable, List, Optional
+ 
+ import transformers
+ from megatron.core.models.gpt import GPTModel as MCoreGPTModel
+-from megatron.core.models.gpt.experimental_attention_variant_module_specs import (
+-    get_transformer_block_with_experimental_attention_variant_spec,
+-)
++# from megatron.core.models.gpt.experimental_attention_variant_module_specs import (
++#     get_transformer_block_with_experimental_attention_variant_spec,
++# )
+ from megatron.core.transformer.spec_utils import ModuleSpec
+ from megatron.core.transformer.transformer_block import TransformerBlockSubmodules
+ from packaging.version import Version as PkgVersion
+@@ -105,9 +105,10 @@ class Qwen35VLModelProvider(GPTModelProvider):
+     # =========================================================================
+     # Hybrid Architecture (Qwen3-Next style)
+     # =========================================================================
+-    transformer_layer_spec: ModuleSpec | Callable[["GPTModelProvider"], ModuleSpec] = (
+-        get_transformer_block_with_experimental_attention_variant_spec
+-    )
++    # transformer_layer_spec: ModuleSpec | Callable[["GPTModelProvider"], ModuleSpec] = (
++    #     get_transformer_block_with_experimental_attention_variant_spec
++    # )
++    transformer_layer_spec: ModuleSpec = None
+     layernorm_zero_centered_gamma: bool = True
+     attention_output_gate: bool = True
+     experimental_attention_variant: str = "gated_delta_net"
+@@ -261,9 +262,10 @@ class Qwen35VLMoEModelProvider(GPTModelProvider):
+     # =========================================================================
+     # Hybrid Architecture (Qwen3-Next style)
+     # =========================================================================
+-    transformer_layer_spec: ModuleSpec | Callable[["GPTModelProvider"], ModuleSpec] = (
+-        get_transformer_block_with_experimental_attention_variant_spec
+-    )
++    # transformer_layer_spec: ModuleSpec | Callable[["GPTModelProvider"], ModuleSpec] = (
++    #     get_transformer_block_with_experimental_attention_variant_spec
++    # )
++    transformer_layer_spec: ModuleSpec = None
+     layernorm_zero_centered_gamma: bool = True
+     attention_output_gate: bool = True
+     experimental_attention_variant: str = "gated_delta_net"
+diff --git a/src/megatron/bridge/models/transformer_config.py b/src/megatron/bridge/models/transformer_config.py
+index 618700ed..3133f4ea 100644
+--- a/src/megatron/bridge/models/transformer_config.py
++++ b/src/megatron/bridge/models/transformer_config.py
+@@ -47,6 +47,10 @@ def _safe_asdict(obj, skip_keys: set[str]) -> dict:
+         return obj.__class__((_safe_asdict(k, skip_keys), _safe_asdict(v, skip_keys)) for k, v in obj.items())
+     return obj
+ 
++from dataclasses import dataclass, field
++
++class MyConfig:
++    pass
+ 
+ @dataclass
+ class TransformerConfig(MCoreTransformerConfig):
+@@ -70,6 +74,13 @@ class TransformerConfig(MCoreTransformerConfig):
+     """
+ 
+     _NO_COPY_KEYS = {"_pg_collection"}
++    # vllm_eplb_config: MyConfig = field(default_factory=MyConfig)
++    # vllm_ir_op_priority: MyConfig = field(default_factory=MyConfig)
++    # vllm_kernel_config: MyConfig = field(default_factory=MyConfig)
++    # vllm_profiler_config: MyConfig = field(default_factory=MyConfig)
++    # vllm_structured_outputs_config: MyConfig = field(default_factory=MyConfig)
++    # vllm_compilation_config: MyConfig = field(default_factory=MyConfig)
++    # vllm_attention_config: MyConfig = field(default_factory=MyConfig)
+ 
+     def __post_init__(self) -> None:
+         """Skip MCore post_init during initial construction.
+diff --git a/src/megatron/bridge/peft/utils.py b/src/megatron/bridge/peft/utils.py
+index 1ca5b18b..3d8ec12a 100644
+--- a/src/megatron/bridge/peft/utils.py
++++ b/src/megatron/bridge/peft/utils.py
+@@ -62,7 +62,7 @@ HAVE_TE = all(
+     )
+ )
+ 
+-MixedFusedLayerNorm, HAVE_APEX = safe_import_from("apex.normalization.fused_layer_norm", "MixedFusedLayerNorm")
++# MixedFusedLayerNorm, HAVE_APEX = safe_import_from("apex.normalization.fused_layer_norm", "MixedFusedLayerNorm")
+ 
+ TECL = (TEColumnParallelLinear, TELayerNormColumnParallelLinear, TEColumnParallelGroupedLinear)
+ TERL = (TERowParallelLinear, TERowParallelGroupedLinear)
+diff --git a/src/megatron/bridge/training/mlm_compat/model.py b/src/megatron/bridge/training/mlm_compat/model.py
+index 60cc091c..1dd5c808 100644
+--- a/src/megatron/bridge/training/mlm_compat/model.py
++++ b/src/megatron/bridge/training/mlm_compat/model.py
+@@ -21,9 +21,9 @@ from megatron.core import tensor_parallel
+ from megatron.core.enums import ModelType
+ from megatron.core.fp8_utils import correct_amax_history_if_needed
+ from megatron.core.models.gpt import GPTModel
+-from megatron.core.models.gpt.experimental_attention_variant_module_specs import (
+-    get_transformer_block_with_experimental_attention_variant_spec,
+-)
++# from megatron.core.models.gpt.experimental_attention_variant_module_specs import (
++#     get_transformer_block_with_experimental_attention_variant_spec,
++# )
+ from megatron.core.models.gpt.gpt_layer_specs import (
+     get_gpt_decoder_block_spec,
+     get_gpt_layer_local_spec,
+diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
+index 7775c11c..c7b094cf 100644
+--- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
++++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
+@@ -26,6 +26,7 @@ from megatron.core.packed_seq_params import PackedSeqParams
+ from megatron.core.process_groups_config import ProcessGroupCollection
+ from megatron.core.transformer import MegatronModule
+ from megatron.core.transformer.spec_utils import ModuleSpec
++from megatron.core.utils import nvtx_range_pop, nvtx_range_push
+ from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig as Qwen3VLConfigHF
+ 
+ from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.attention import Qwen3VLSelfAttention
+@@ -294,7 +295,7 @@ class Qwen3VLModel(MegatronModule):
+         # position ids is computed within the model
+         position_ids = None
+ 
+-        torch.cuda.nvtx.range_push("Qwen3VLModel.forward.pre_process")
++        nvtx_range_push(msg="Qwen3VLModel.forward.pre_process")
+ 
+         cp_rank = self.pg_collection.cp.rank()
+         cp_size = self.pg_collection.cp.size()
+@@ -387,7 +388,7 @@ class Qwen3VLModel(MegatronModule):
+                 combined_embeddings = split_data_cp_rank(combined_embeddings, cp_size, 0, cp_rank)
+             if packed_seq_params is not None:
+                 if attention_mask is None:
+-                    attention_mask = torch.ones_like(input_ids, dtype=torch.int32, device=input_ids.device)
++                    attention_mask = torch.ones_like(input_ids, dtype=torch.bool, device=input_ids.device)
+                 input_ids_thd, _ = preprocess_packed_seqs(
+                     input_ids, attention_mask, pre_process=True, pg_collection=self.pg_collection
+                 )
+@@ -443,7 +444,7 @@ class Qwen3VLModel(MegatronModule):
+             # convert lm_input_ids to THD format so it matches position_ids.
+             if packed_seq_params is not None:
+                 if attention_mask is None:
+-                    attention_mask = torch.ones_like(input_ids, dtype=torch.int32, device=input_ids.device)
++                    attention_mask = torch.ones_like(input_ids, dtype=torch.bool, device=input_ids.device)
+                 lm_input_ids, _ = preprocess_packed_seqs(
+                     input_ids, attention_mask, pre_process=True, pg_collection=self.pg_collection
+                 )
+@@ -499,8 +500,8 @@ class Qwen3VLModel(MegatronModule):
+                 attention_mask = None
+                 self.language_model.rotary_pos_emb.is_thd_format = True
+ 
+-        torch.cuda.nvtx.range_pop()
+-        torch.cuda.nvtx.range_push("Qwen3VLModel.forward.language_model")
++        nvtx_range_pop(msg="Qwen3VLModel.forward.pre_process")
++        nvtx_range_push(msg="Qwen3VLModel.forward.language_model")
+ 
+         output = self.language_model(
+             input_ids=lm_input_ids,
+@@ -516,6 +517,6 @@ class Qwen3VLModel(MegatronModule):
+             **(extra_block_kwargs or {}),
+             **kwargs,
+         )
+-        torch.cuda.nvtx.range_pop()
++        nvtx_range_pop(msg="Qwen3VLModel.forward.language_model")
+ 
+         return output
+diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py
+index b714d0f7..42fcab74 100644
+--- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py
++++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py
+@@ -668,6 +668,11 @@ def preprocess_packed_seqs(
+     """
+     batch_size = input_ids.shape[0]
+ 
++    # Ensure boolean dtype for correct advanced indexing (bool → mask select,
++    # int → fancy index which silently corrupts data when values are 0/1).
++    if attention_mask.dtype != torch.bool:
++        attention_mask = attention_mask.bool()
++
+     seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
+     if pg_collection is not None:
+         tp_size = pg_collection.tp.size()
diff --git a/docker/npu_patch/megatron.patch b/docker/npu_patch/megatron.patch
new file mode 100644
index 000000000..e676ac3f6
--- /dev/null
+++ b/docker/npu_patch/megatron.patch
@@ -0,0 +1,661 @@
+diff --git a/megatron/core/activations.py b/megatron/core/activations.py
+index 8b422d73a..58fba4667 100644
+--- a/megatron/core/activations.py
++++ b/megatron/core/activations.py
+@@ -5,19 +5,19 @@ import torch.nn.functional as F
+ from megatron.core.jit import jit_fuser
+ 
+ 
+-@jit_fuser
++
+ def squared_relu(x: torch.Tensor) -> torch.Tensor:
+     """Squared ReLU activation"""
+     return torch.pow(F.relu(x), 2)
+ 
+ 
+-@jit_fuser
++
+ def quick_gelu(x: torch.Tensor) -> torch.Tensor:
+     """Quick GELU activation"""
+     return x * torch.sigmoid(1.702 * x)
+ 
+ 
+-@jit_fuser
++
+ def fast_gelu(x: torch.Tensor) -> torch.Tensor:
+     """Fast GELU activation"""
+     return 0.5 * x * (1.0 + torch.tanh(x * 0.7978845608 * (1.0 + 0.044715 * x * x)))
+diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py
+index a9982e176..807c4a88d 100644
+--- a/megatron/core/distributed/param_and_grad_buffer.py
++++ b/megatron/core/distributed/param_and_grad_buffer.py
+@@ -787,45 +787,19 @@ class _ParamAndGradBuffer:
+             # Individual param/grad contexts below handle TMS regions separately.
+             mem_alloc_context = nullcontext
+ 
+-        def _make_no_backup_context(tag, disable, flag_name="disable_grad_buffers_cpu_backup"):
+-            if disable:
+-                try:
+-                    from torch_memory_saver import torch_memory_saver
+-                except ImportError as e:
+-                    raise ImportError(
+-                        f"{flag_name}=True requires torch_memory_saver. "
+-                        "Install with: pip install torch-memory-saver"
+-                    ) from e
+-                return partial(
+-                    torch_memory_saver.region,
+-                    tag=tag,
+-                    enable_cpu_backup=False,
+-                )
+-            return nullcontext
+-        grad_mem_alloc_context = _make_no_backup_context(
+-            "grad_buffer", disable_grad_buffers_cpu_backup
+-        )
+-        param_mem_alloc_context = _make_no_backup_context(
+-            "param_buffer", disable_param_buffers_cpu_backup, "disable_param_buffers_cpu_backup"
+-        )
+-
++        # NPU already wraps model construction in the outer VIME training
++        # torch_memory_saver region. Do not open nested mem-pool regions here.
+         with mem_alloc_context():
+             # For MXFP8 param: Create a shared buffer for param AG and grad RS for memory efficiency
+             # The buffer is mapped to weight gradients whose dtype is either bf16 or FP32.
+             # It can be temporarily reused by param AG.
+             if self.ddp_config.use_distributed_optimizer and any(is_mxfp8tensor(p) for p in params):
+-                shared_mem_alloc_context = (
+-                    param_mem_alloc_context
+-                    if disable_param_buffers_cpu_backup
+-                    else grad_mem_alloc_context
++                self.shared_buffer = torch.zeros(
++                    self.numel,
++                    dtype=self.grad_dtype,
++                    device=torch.cuda.current_device(),
++                    requires_grad=False,
+                 )
+-                with shared_mem_alloc_context():
+-                    self.shared_buffer = torch.zeros(
+-                        self.numel,
+-                        dtype=self.grad_dtype,
+-                        device=torch.cuda.current_device(),
+-                        requires_grad=False,
+-                    )
+                 # For FP32 weight grads, only half of the buffer is used to store params in bf16.
+                 if self.grad_dtype == torch.float32:
+                     self.param_data = self.shared_buffer[: math.ceil(self.numel / 2)].view(
+@@ -837,20 +811,18 @@ class _ParamAndGradBuffer:
+             else:
+                 # Only re-map param tensors if using distributed optimizer.
+                 if self.ddp_config.use_distributed_optimizer:
+-                    with param_mem_alloc_context():
+-                        self.param_data = torch.zeros(
+-                            self.numel,
+-                            dtype=self.param_dtype,
+-                            device=torch.cuda.current_device(),
+-                            requires_grad=False,
+-                        )
+-                with grad_mem_alloc_context():
+-                    self.grad_data = torch.zeros(
++                    self.param_data = torch.zeros(
+                         self.numel,
+-                        dtype=self.grad_dtype,
++                        dtype=self.param_dtype,
+                         device=torch.cuda.current_device(),
+                         requires_grad=False,
+                     )
++                self.grad_data = torch.zeros(
++                    self.numel,
++                    dtype=self.grad_dtype,
++                    device=torch.cuda.current_device(),
++                    requires_grad=False,
++                )
+ 
+         self.grad_data_size = 0
+         self.param_data_size = 0
+diff --git a/megatron/core/fusions/fused_bias_dropout.py b/megatron/core/fusions/fused_bias_dropout.py
+index 336452562..614ee1a48 100644
+--- a/megatron/core/fusions/fused_bias_dropout.py
++++ b/megatron/core/fusions/fused_bias_dropout.py
+@@ -64,14 +64,14 @@ def bias_dropout_add_unfused(training):
+     return _bias_dropout_add
+ 
+ 
+-@jit_fuser
++
+ def bias_dropout_add_fused_train(
+     x_with_bias: Tuple[torch.Tensor, Optional[torch.Tensor]], residual: torch.Tensor, prob: float
+ ) -> torch.Tensor:
+     return _bias_dropout_add_func(x_with_bias, residual, prob, True)
+ 
+ 
+-@jit_fuser
++
+ def bias_dropout_add_fused_inference(
+     x_with_bias: Tuple[torch.Tensor, Optional[torch.Tensor]], residual: torch.Tensor, prob: float
+ ) -> torch.Tensor:
+diff --git a/megatron/core/fusions/fused_bias_geglu.py b/megatron/core/fusions/fused_bias_geglu.py
+index 7a7fbe7f9..f9a438953 100644
+--- a/megatron/core/fusions/fused_bias_geglu.py
++++ b/megatron/core/fusions/fused_bias_geglu.py
+@@ -13,7 +13,7 @@ from megatron.core.jit import jit_fuser
+ # x * 0.5 * (1.0 + torch.erf(x * 0.70710678))
+ 
+ 
+-@jit_fuser
++
+ def geglu(y):
+     """Performs GEGLU (GELU-Gated Linear Unit) activation.
+ 
+@@ -27,7 +27,7 @@ def geglu(y):
+     return (y_1 * 0.5 * (1.0 + torch.tanh(0.79788456 * y_1 * (1 + 0.044715 * y_1 * y_1)))) * y_2
+ 
+ 
+-@jit_fuser
++
+ def bias_geglu(bias, y):
+     """Performs GEGLU activation with bias addition.
+ 
+@@ -45,7 +45,7 @@ def bias_geglu(bias, y):
+ # gradient of tanh approximation of gelu
+ # gradient of actual gelu is:
+ # 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x)
+-@jit_fuser
++
+ def geglu_back(g, y):
+     """Computes the gradient for the GEGLU activation.
+ 
+@@ -65,7 +65,7 @@ def geglu_back(g, y):
+     return torch.cat(((g * y_2) * ff, g * (y_1 * 0.5 * (1.0 + tanh_out))), -1)
+ 
+ 
+-@jit_fuser
++
+ def bias_geglu_back(g, y, bias):
+     """Computes the gradient for the biased GEGLU activation.
+ 
+@@ -181,13 +181,13 @@ def bias_geglu_impl(input, bias):
+ # ------------------------- QUICK GEGLU FUSION --------------------------
+ 
+ 
+-@jit_fuser
++
+ def quick_gelu(y: torch.Tensor) -> torch.Tensor:
+     """Sigmoid approximation of gelu"""
+     return y * torch.sigmoid(1.702 * y)
+ 
+ 
+-@jit_fuser
++
+ def quick_geglu(y: torch.Tensor, linear_offset: float = 0.0) -> torch.Tensor:
+     """Performs Quick-GELU-based GEGLU activation : quick_gelu(y1) * (y2 + offset).
+ 
+@@ -202,7 +202,7 @@ def quick_geglu(y: torch.Tensor, linear_offset: float = 0.0) -> torch.Tensor:
+     return quick_gelu(y_1) * (y_2 + linear_offset)
+ 
+ 
+-@jit_fuser
++
+ def weighted_quick_geglu(
+     y: torch.Tensor, weights: torch.Tensor, linear_offset: float = 0.0
+ ) -> torch.Tensor:
+@@ -217,7 +217,7 @@ def weighted_quick_geglu(
+ 
+ 
+ # gradient of sigmoid approximation of gelu
+-@jit_fuser
++
+ def quick_geglu_back(g, y, linear_offset: float = 0.0) -> torch.Tensor:
+     """Backward helper for Quick-GEGLU.
+ 
+@@ -236,7 +236,7 @@ def quick_geglu_back(g, y, linear_offset: float = 0.0) -> torch.Tensor:
+     return torch.cat((dy_1, dy_2), -1)
+ 
+ 
+-@jit_fuser
++
+ def weighted_quick_geglu_back(g, y, weights, linear_offset: float = 0.0):
+     """Backward helper for weighted Quick-GEGLU.
+     Returns gradient w.r.t input `y` and `weights`.
+@@ -255,7 +255,7 @@ def weighted_quick_geglu_back(g, y, weights, linear_offset: float = 0.0):
+ # ---------------- Weighted Bias Quick-GEGLU helpers -----------------
+ 
+ 
+-@jit_fuser
++
+ def weighted_bias_quick_geglu(
+     y: torch.Tensor, bias: torch.Tensor, weights: torch.Tensor, linear_offset: float = 0.0
+ ) -> torch.Tensor:
+@@ -275,7 +275,7 @@ def weighted_bias_quick_geglu(
+     return res.to(dtype)
+ 
+ 
+-@jit_fuser
++
+ def weighted_bias_quick_geglu_back(g, y, bias, weights, linear_offset: float = 0.0):
+     """Backward helper for weighted Quick-GEGLU with bias.
+ 
+diff --git a/megatron/core/fusions/fused_bias_gelu.py b/megatron/core/fusions/fused_bias_gelu.py
+index 8cc90f617..fda8f2f5f 100644
+--- a/megatron/core/fusions/fused_bias_gelu.py
++++ b/megatron/core/fusions/fused_bias_gelu.py
+@@ -13,7 +13,7 @@ from megatron.core.jit import jit_fuser
+ # x * 0.5 * (1.0 + torch.erf(x * 0.70710678))
+ 
+ 
+-@jit_fuser
++
+ def bias_gelu(bias, y):
+     x = bias + y
+     return x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))
+@@ -22,7 +22,7 @@ def bias_gelu(bias, y):
+ # gradient of tanh approximation of gelu
+ # gradient of actual gelu is:
+ # 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x)
+-@jit_fuser
++
+ def bias_gelu_back(g, bias, y):
+     x = bias + y
+     tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))
+diff --git a/megatron/core/fusions/fused_bias_swiglu.py b/megatron/core/fusions/fused_bias_swiglu.py
+index 632470876..105936786 100644
+--- a/megatron/core/fusions/fused_bias_swiglu.py
++++ b/megatron/core/fusions/fused_bias_swiglu.py
+@@ -12,7 +12,7 @@ from megatron.core.utils import nvtx_decorator
+ ###### BIAS SWIGLU FUSION/ NO AUTOGRAD ################
+ 
+ 
+-@jit_fuser
++
+ def swiglu(y):
+     """Performs SwiGLU (Swish-Gated Linear Unit) activation function.
+ 
+@@ -26,7 +26,7 @@ def swiglu(y):
+     return F.silu(y_1) * y_2
+ 
+ 
+-@jit_fuser
++
+ def bias_swiglu(y, bias):
+     """Performs SwiGLU activation with bias addition.
+ 
+@@ -41,7 +41,7 @@ def bias_swiglu(y, bias):
+     return swiglu(y)
+ 
+ 
+-@jit_fuser
++
+ def weighted_swiglu(y, weights):
+     dtype = y.dtype
+     res = swiglu(y) * weights
+@@ -51,7 +51,7 @@ def weighted_swiglu(y, weights):
+ # gradient of tanh approximation of gelu
+ # gradient of actual gelu is:
+ # 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x)
+-@jit_fuser
++
+ def swiglu_back(g, y):
+     """Computes the gradient for the SwiGLU activation function.
+ 
+@@ -69,7 +69,7 @@ def swiglu_back(g, y):
+     )
+ 
+ 
+-@jit_fuser
++
+ def bias_swiglu_back(g, y, bias):
+     """Computes the gradient for the biased SwiGLU activation function.
+ 
+@@ -86,7 +86,7 @@ def bias_swiglu_back(g, y, bias):
+     return swiglu_back(g, y)
+ 
+ 
+-@jit_fuser
++
+ def weighted_swiglu_back(g, y, weights):
+     input_dtype = y.dtype
+     w_dtype = weights.dtype
+diff --git a/megatron/core/fusions/fused_cross_entropy.py b/megatron/core/fusions/fused_cross_entropy.py
+index 23e4b6031..4d447e0e0 100644
+--- a/megatron/core/fusions/fused_cross_entropy.py
++++ b/megatron/core/fusions/fused_cross_entropy.py
+@@ -9,7 +9,7 @@ from megatron.core.tensor_parallel.cross_entropy import VocabParallelCrossEntrop
+ from megatron.core.tensor_parallel.utils import VocabUtility
+ 
+ 
+-@jit_fuser
++
+ def calculate_logits_max(vocab_parallel_logits: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+     """
+     Calculates the maximum logits of the predicted tokens.
+@@ -22,7 +22,7 @@ def calculate_logits_max(vocab_parallel_logits: torch.Tensor) -> Tuple[torch.Ten
+     return vocab_parallel_logits, logits_max
+ 
+ 
+-@jit_fuser
++
+ def calculate_predicted_logits(
+     vocab_parallel_logits: torch.Tensor,
+     target: torch.Tensor,
+@@ -44,7 +44,7 @@ def calculate_predicted_logits(
+     return target_mask, masked_target_1d, predicted_logits_sum_exp_logits, exp_logits
+ 
+ 
+-@jit_fuser
++
+ def calculate_cross_entropy_loss(
+     exp_logits: torch.Tensor, predicted_logits_sum_exp_logits: torch.Tensor
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+@@ -61,7 +61,7 @@ def calculate_cross_entropy_loss(
+     return exp_logits, loss
+ 
+ 
+-@jit_fuser
++
+ def calculate_gradients(
+     softmax: torch.Tensor,
+     grad_output: torch.Tensor,
+diff --git a/megatron/core/fusions/fused_pad_routing_map.py b/megatron/core/fusions/fused_pad_routing_map.py
+index c382178b6..563279edd 100644
+--- a/megatron/core/fusions/fused_pad_routing_map.py
++++ b/megatron/core/fusions/fused_pad_routing_map.py
+@@ -70,7 +70,7 @@ def _pad_routing_map_kernel(
+     tl.store(output_row_ptr + token_indices, output_row, mask=token_mask)
+ 
+ 
+-@jit_fuser
++
+ def fused_pad_routing_map(routing_map: torch.Tensor, pad_multiple: int) -> torch.Tensor:
+     """Fused version of pad_routing_map.
+     Args:
+diff --git a/megatron/core/fusions/fused_weighted_squared_relu.py b/megatron/core/fusions/fused_weighted_squared_relu.py
+index 02dabc14c..137022386 100644
+--- a/megatron/core/fusions/fused_weighted_squared_relu.py
++++ b/megatron/core/fusions/fused_weighted_squared_relu.py
+@@ -10,7 +10,7 @@ from megatron.core.utils import nvtx_decorator
+ ######################  WEIGHTED SQUARED ReLU FUSION  ######################
+ 
+ 
+-@jit_fuser
++
+ def weighted_squared_relu(x: torch.Tensor, weights: torch.Tensor) -> torch.Tensor:
+     """Element-wise weight applied after Squared-ReLU.
+ 
+@@ -28,7 +28,7 @@ def weighted_squared_relu(x: torch.Tensor, weights: torch.Tensor) -> torch.Tenso
+     return res.to(out_dtype)
+ 
+ 
+-@jit_fuser
++
+ def _squared_relu_back(g: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
+     """Gradient of Squared-ReLU.
+ 
+@@ -37,7 +37,7 @@ def _squared_relu_back(g: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
+     return g * 2 * F.relu(x)
+ 
+ 
+-@jit_fuser
++
+ def weighted_squared_relu_back(g: torch.Tensor, x: torch.Tensor, weights: torch.Tensor):
+     """Backward for weighted Squared-ReLU.
+ 
+diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py
+index 2b241c86d..cca4503d5 100644
+--- a/megatron/core/inference/contexts/dynamic_context.py
++++ b/megatron/core/inference/contexts/dynamic_context.py
+@@ -51,12 +51,8 @@ try:
+ except:
+     HAVE_PACKAGING = False
+ 
+-try:
+-    import flashinfer  # pylint: disable=unused-import
+ 
+-    HAVE_FLASHINFER = True
+-except ImportError:
+-    HAVE_FLASHINFER = False
++HAVE_FLASHINFER = False
+ 
+ try:
+     from torch_memory_saver import torch_memory_saver
+diff --git a/megatron/core/models/gpt/gpt_layer_specs.py b/megatron/core/models/gpt/gpt_layer_specs.py
+index a9d3583e5..f2adaceb2 100755
+--- a/megatron/core/models/gpt/gpt_layer_specs.py
++++ b/megatron/core/models/gpt/gpt_layer_specs.py
+@@ -209,7 +209,7 @@ def get_gpt_layer_with_transformer_engine_spec(
+             'The fp8 argument in "get_gpt_layer_with_transformer_engine_spec" has been deprecated'
+             " and will be removed soon. Please update your code accordingly."
+         )
+-
++    from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider
+     if use_kitchen:
+         assert HAVE_KITCHEN
+         backend: BackendSpecProvider = KitchenSpecProvider(
+diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py
+index dd8590a79..3736b0e31 100644
+--- a/megatron/core/optimizer/distrib_optimizer.py
++++ b/megatron/core/optimizer/distrib_optimizer.py
+@@ -659,10 +659,12 @@ class DistributedOptimizer(MixedPrecisionOptimizer):
+                     break
+         elif USING_TE_OPTIMIZER or USING_APEX_OPTIMIZER:
+             # Extract 'step', for TE FusedAdam support.
++            # Use int() to convert tensors to Python scalars so set() deduplicates
++            # by value (PyTorch tensor hashing is identity-based).
+             steps = list(
+                 set(
+                     [
+-                        g["step"]
++                        int(g["step"])
+                         for g in inner_state_dict["param_groups"]
+                         if len(g["params"]) > 0 and "step" in g
+                     ]
+@@ -824,7 +826,7 @@ class DistributedOptimizer(MixedPrecisionOptimizer):
+ 
+         # Extract 'step', for non-Apex/TE support.
+         if not HAVE_APEX_OR_TE:
+-            steps = list(set([g["step"] for g in state_dict["optimizer"]["param_groups"]]))
++            steps = list(set([int(g["step"]) for g in state_dict["optimizer"]["param_groups"]]))
+             assert len(steps) == 1
+             step = torch.tensor(steps[0], dtype=torch.float)
+ 
+diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py
+index 601a72a43..ec114cf38 100644
+--- a/megatron/core/ssm/gated_delta_net.py
++++ b/megatron/core/ssm/gated_delta_net.py
+@@ -465,7 +465,7 @@ class GatedDeltaNet(MegatronModule):
+ 
+         return out, out_bias
+ 
+-    @jit_fuser
++    
+     def _apply_gated_norm(self, x, gate):
+         # Output Norm
+         x_dtype = x.dtype
+diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py
+index bc5e4e2ee..d223e774e 100644
+--- a/megatron/core/transformer/attention.py
++++ b/megatron/core/transformer/attention.py
+@@ -1219,7 +1219,7 @@ class Attention(MegatronModule, ABC):
+ 
+         return output, bias
+ 
+-    @jit_fuser
++    
+     def _apply_output_gate(self, x, gate):
+         x_dtype = x.dtype
+         gate = gate.contiguous()
+diff --git a/megatron/core/transformer/module.py b/megatron/core/transformer/module.py
+index c30c107e7..f4340f4d1 100644
+--- a/megatron/core/transformer/module.py
++++ b/megatron/core/transformer/module.py
+@@ -17,9 +17,9 @@ from megatron.core.transformer.utils import (
+     sharded_state_dict_default,
+ )
+ 
+-_FLOAT_TYPES = (torch.FloatTensor, torch.cuda.FloatTensor)
+-_HALF_TYPES = (torch.HalfTensor, torch.cuda.HalfTensor)
+-_BF16_TYPES = (torch.BFloat16Tensor, torch.cuda.BFloat16Tensor)
++_FLOAT_TYPES = (torch.FloatTensor, torch.cuda.FloatTensor, torch.npu.FloatTensor)
++_HALF_TYPES = (torch.HalfTensor, torch.cuda.HalfTensor, torch.npu.HalfTensor)
++_BF16_TYPES = (torch.BFloat16Tensor, torch.cuda.BFloat16Tensor, torch.npu.BFloat16Tensor)
+ 
+ 
+ def param_is_not_shared(param):  # pylint: disable=missing-function-docstring
+diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py
+index d8e753422..aff929985 100644
+--- a/megatron/core/transformer/moe/experts.py
++++ b/megatron/core/transformer/moe/experts.py
+@@ -92,7 +92,7 @@ class GroupedMLP(MegatronModule):
+             if self.config.activation_func not in (F.silu, F.gelu):
+                 raise ValueError("Activation function must be silu or gelu when using GroupedMLP.")
+ 
+-            @jit_fuser
++            
+             def glu(x):
+                 x = torch.chunk(x, 2, dim=-1)
+                 return self.config.activation_func(x[0]) * x[1]
+@@ -109,7 +109,7 @@ class GroupedMLP(MegatronModule):
+                 "moe_act recompute for fp8 or fp4 cannot work with the legacy GroupedMLP."
+             )
+ 
+-        @jit_fuser
++        
+         def activation_func_with_probs(x, probs):
+             dtype = x.dtype
+             res = self.activation_func(x) * probs
+diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py
+index 327dbc8a3..b007d6662 100644
+--- a/megatron/core/transformer/moe/token_dispatcher.py
++++ b/megatron/core/transformer/moe/token_dispatcher.py
+@@ -1402,7 +1402,7 @@ class MoEFlexTokenDispatcher(MoETokenDispatcher):
+         ).contiguous()
+         return routing_map, probs
+ 
+-    @jit_fuser
++    
+     def dispatch_preprocess(
+         self, hidden_states: torch.Tensor, routing_map: torch.Tensor, probs: torch.Tensor
+     ):
+diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
+index 63f81465d..c47baeaea 100755
+--- a/megatron/core/transformer/multi_token_prediction.py
++++ b/megatron/core/transformer/multi_token_prediction.py
+@@ -7,6 +7,7 @@ from typing import Callable, List, Optional, Union
+ 
+ import torch
+ from torch import Tensor
++import warnings
+ 
+ from megatron.core import InferenceParams, parallel_state, tensor_parallel
+ from megatron.core.dist_checkpointing.mapping import ShardedStateDict
+diff --git a/megatron/core/transformer/torch_norm.py b/megatron/core/transformer/torch_norm.py
+index d0ceca7af..f16796680 100644
+--- a/megatron/core/transformer/torch_norm.py
++++ b/megatron/core/transformer/torch_norm.py
+@@ -69,7 +69,7 @@ class L2Norm(torch.nn.Module):
+         self.hidden_size = hidden_size
+         self.eps = eps
+ 
+-    @jit_fuser
++    
+     def _norm(self, x):
+         """
+         Performs the actual L2 normalization.
+diff --git a/megatron/core/transformer/utils.py b/megatron/core/transformer/utils.py
+index 880c53099..dbc95736e 100644
+--- a/megatron/core/transformer/utils.py
++++ b/megatron/core/transformer/utils.py
+@@ -51,7 +51,7 @@ def attention_mask_func(attention_scores, attention_mask):
+     return attention_scores
+ 
+ 
+-@jit_fuser
++
+ def gelu_impl(x):
+     """OpenAI's gelu implementation."""
+     return 0.5 * x * (1.0 + torch.tanh(0.7978845608028654 * x * (1.0 + 0.044715 * x * x)))
+@@ -65,7 +65,7 @@ def openai_gelu(x):
+ # This is actually Python equivalent of torch.nn.functional.gelu(), also with
+ # type hints for ONNX exporter
+ # pylint: disable=missing-function-docstring
+-@jit_fuser
++
+ def erf_gelu(x):
+     return (
+         x * 0.5 * (torch.erf(x / 1.41421).to(dtype=x.dtype) + torch.ones_like(x).to(dtype=x.dtype))
+diff --git a/megatron/legacy/model/fused_bias_gelu.py b/megatron/legacy/model/fused_bias_gelu.py
+index e00e63148..ffe4b7ec6 100644
+--- a/megatron/legacy/model/fused_bias_gelu.py
++++ b/megatron/legacy/model/fused_bias_gelu.py
+@@ -12,7 +12,7 @@ from megatron.core.jit import jit_fuser
+ # actual gelu is:
+ # x * 0.5 * (1.0 + torch.erf(x * 0.70710678))
+ 
+-@jit_fuser
++
+ def bias_gelu(bias, y):
+     x = bias + y
+     return  x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))
+@@ -20,7 +20,7 @@ def bias_gelu(bias, y):
+ # gradient of tanh approximation of gelu
+ # gradient of actual gelu is:
+ # 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x)
+-@jit_fuser
++
+ def bias_gelu_back(g, bias, y):
+     x = bias + y
+     tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))
+diff --git a/megatron/legacy/model/transformer.py b/megatron/legacy/model/transformer.py
+index ca3414eec..5dd9e4798 100644
+--- a/megatron/legacy/model/transformer.py
++++ b/megatron/legacy/model/transformer.py
+@@ -856,7 +856,7 @@ def get_bias_dropout_add(training):
+     return _bias_dropout_add
+ 
+ 
+-@jit_fuser
++
+ def bias_dropout_add_fused_train(x: torch.Tensor,
+                                  bias: Optional[torch.Tensor],
+                                  residual: torch.Tensor,
+@@ -864,7 +864,7 @@ def bias_dropout_add_fused_train(x: torch.Tensor,
+     return bias_dropout_add(x, bias, residual, prob, True)
+ 
+ 
+-@jit_fuser
++
+ def bias_dropout_add_fused_inference(x: torch.Tensor,
+                                      bias: Optional[torch.Tensor],
+                                      residual: torch.Tensor,
+diff --git a/megatron/legacy/model/utils.py b/megatron/legacy/model/utils.py
+index 5762000d5..534858df7 100644
+--- a/megatron/legacy/model/utils.py
++++ b/megatron/legacy/model/utils.py
+@@ -43,7 +43,7 @@ def get_linear_layer(rows, columns, init_method):
+     return layer
+ 
+ 
+-@jit_fuser
++
+ def gelu_impl(x):
+     """OpenAI's gelu implementation."""
+     return 0.5 * x * (1.0 + torch.tanh(0.7978845608028654 * x *
+@@ -54,7 +54,7 @@ def openai_gelu(x):
+ 
+ 
+ #This is actually Python equivalent of torch.nn.functional.gelu(), also with type hints for ONNX exporter
+-@jit_fuser
++
+ def erf_gelu(x):
+     return x * 0.5 * (torch.erf(x / 1.41421).to(dtype=x.dtype)+torch.ones_like(x).to(dtype=x.dtype))
+ 
+diff --git a/megatron/training/utils.py b/megatron/training/utils.py
+index 7709f6513..65b8f7fca 100644
+--- a/megatron/training/utils.py
++++ b/megatron/training/utils.py
+@@ -446,6 +446,8 @@ def is_first_or_last_pipeline_stage(vp_stage):
+ 
+ def get_device_arch_version():
+     """Returns GPU arch version (8: Ampere, 9: Hopper, 10: Blackwell, ...)"""
++    if hasattr(torch, 'npu') and torch.npu.is_available():
++        return 10  # NPU: treat as Blackwell to avoid CUDA restrictions
+     return torch.cuda.get_device_properties(torch.device("cuda:0")).major
+ 
+ 
diff --git a/docker/npu_patch/mindspeed.patch b/docker/npu_patch/mindspeed.patch
new file mode 100644
index 000000000..8a017c8e2
--- /dev/null
+++ b/docker/npu_patch/mindspeed.patch
@@ -0,0 +1,299 @@
+diff --git a/mindspeed/core/fusions/fused_moe_permute.py b/mindspeed/core/fusions/fused_moe_permute.py
+index bb007b44..98708a5b 100644
+--- a/mindspeed/core/fusions/fused_moe_permute.py
++++ b/mindspeed/core/fusions/fused_moe_permute.py
+@@ -100,9 +100,9 @@ def sort_chunks_by_idxs_wrapper(fn):
+ def moe_alltoall_token_dispatcher_init_wrapper(fn):
+     @wraps(fn)
+     def wrapper(
+-            self, num_local_experts, local_expert_indices, config, model_comm_pgs=None
++            self, num_local_experts, local_expert_indices, config, pg_collection=None
+     ) -> None:
+-        fn(self, num_local_experts, local_expert_indices, config, model_comm_pgs)
++        fn(self, num_local_experts, local_expert_indices, config, pg_collection)
+         # Since fused_sort_chunks_by_index is not currently supported, set self.permute_idx_device to None
+         self.permute_idx_device = None
+         input_chunk_idxs = torch.arange(
+diff --git a/mindspeed/core/megatron_basic/arguments_basic.py b/mindspeed/core/megatron_basic/arguments_basic.py
+index 8ea25b9f..054313cb 100644
+--- a/mindspeed/core/megatron_basic/arguments_basic.py
++++ b/mindspeed/core/megatron_basic/arguments_basic.py
+@@ -91,5 +91,11 @@ def transformer_config_init_wrapper(fn):
+         known_config = {}
+         unknown_config = {}
+         ignore_config = ['rope_type']
+-        full_args = vars(get_full_args()).copy()
++        # vLLM serving arguments share the process-level Namespace in Vime,
++        # but they are not Megatron TransformerConfig fields.
++        full_args = {
++            key: value
++            for key, value in vars(get_full_args()).items()
++            if not key.startswith("vllm_")
++        }
+         full_args.update(dict(kwargs))
+@@ -113,3 +119,41 @@ def transformer_config_init_wrapper(fn):
+         fn(self, *args, **known_config)
+ 
+     return wrapper
++
++
++def transformer_config_getattr(self, name):
++    """Resolve MindSpeed extension fields for configs created before patching."""
++    full_args = vars(get_full_args())
++    if name in full_args:
++        return full_args[name]
++    raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
++
++
++def transformer_config_init_subclass(cls, **kwargs):
++    mutable_types = (list, dict, set, bytearray)
++    unknown_config = {}
++    # Keep serving-only vLLM values out of dynamically created Megatron
++    # dataclass fields. They remain available on the original Vime args.
++    full_args = {
++        key: value
++        for key, value in vars(get_full_args()).items()
++        if not key.startswith("vllm_")
++    }
++    full_args.update(kwargs)
++
++    config_key = inspect.signature(cls).parameters
++    for key, value in full_args.items():
++        if key not in config_key:
++            unknown_config[key] = value
++
++    for key, value in unknown_config.items():
++        if not hasattr(cls, key):
++            cls.__annotations__[key] = type(value)
++            value = field(default_factory=value) if callable(value) and not isinstance(value, type) else value
++            if callable(value) and not isinstance(value, type):
++                value = field(default_factory=value)
++            elif type(value) in mutable_types:
++                value = field(default_factory=lambda: value)
++            else:
++                value = value
++            setattr(cls, key, value)
+diff --git a/mindspeed/features_manager/megatron_basic/megatron_basic.py b/mindspeed/features_manager/megatron_basic/megatron_basic.py
+index 355913e1..41678a03 100644
+--- a/mindspeed/features_manager/megatron_basic/megatron_basic.py
++++ b/mindspeed/features_manager/megatron_basic/megatron_basic.py
+@@ -43,8 +43,13 @@ class MegatronBasicFeature(MindSpeedFeature):
+ 
+     def register_mcore_basic_patches(self, pm, args):
+         # configuration patches
+-        from mindspeed.core.megatron_basic.arguments_basic import transformer_config_init_wrapper, transformer_config_post_init_wrapper
++        from mindspeed.core.megatron_basic.arguments_basic import (transformer_config_init_wrapper,
++                                                                    transformer_config_getattr,
++                                                                    transformer_config_post_init_wrapper,
++                                                                    transformer_config_init_subclass)
+         pm.register_patch("megatron.core.transformer.transformer_config.TransformerConfig.__init__", transformer_config_init_wrapper)
++        pm.register_patch("megatron.core.transformer.transformer_config.TransformerConfig.__init_subclass__", classmethod(transformer_config_init_subclass))
++        pm.register_patch("megatron.core.transformer.transformer_config.TransformerConfig.__getattr__", transformer_config_getattr, create_dummy=True)
+         pm.register_patch("megatron.core.transformer.transformer_config.MLATransformerConfig.__init__", transformer_config_init_wrapper)
+         pm.register_patch("megatron.core.transformer.transformer_config.TransformerConfig.__post_init__", transformer_config_post_init_wrapper)
+ 
+diff --git a/mindspeed/patch_utils.py b/mindspeed/patch_utils.py
+index a489d58c..d9328555 100644
+--- a/mindspeed/patch_utils.py
++++ b/mindspeed/patch_utils.py
+@@ -2,7 +2,7 @@ import importlib
+ import sys
+ import types
+ from typing import List, Dict, Union
+-
++import inspect
+ _MEGATRON_TRAINING_AVAILABLE = None
+ 
+ 
+@@ -93,8 +93,10 @@ class Patch:
+ 
+     def remove_patch(self):
+         for key, value in sys.modules.copy().items():
+-            if 'mindspeed' in key:
++            if 'mindspeed' in key or 'torch.classes' == key:
+                 continue
++            if inspect.isclass(self.orig_module) and hasattr(value, self.orig_module_name.split('.')[-1]):
++ 	            value = getattr(value, self.orig_module_name.split('.')[-1])
+             if self.orig_func_name is not None and hasattr(value, self.orig_func_name) \
+                     and id(getattr(value, self.orig_func_name)) == id(self.final_patch_func):
+                 setattr(value, self.orig_func_name, self.orig_func)
+diff --git a/mindspeed/core/fusions/fused_rope.py b/mindspeed/core/fusions/fused_rope.py
+index a6f02e07..70f7cb08 100644
+--- a/mindspeed/core/fusions/fused_rope.py
++++ b/mindspeed/core/fusions/fused_rope.py
+@@ -126,5 +126,6 @@ def apply_rotary_pos_emb(
+         freqs,
+         rotary_interleaved=config.rotary_interleaved,
+         multi_latent_attention=config.multi_latent_attention,
+-        mscale=mscale
++        mscale=mscale,
++        cp_group=cp_group
+     )
+diff --git a/mindspeed/megatron_adaptor.py b/mindspeed/megatron_adaptor.py
+index f14b231d..4615590e 100644
+--- a/mindspeed/megatron_adaptor.py
++++ b/mindspeed/megatron_adaptor.py
+@@ -57,6 +57,7 @@ def delete_lock_file():
+ def repatch(args):
+     MindSpeedFeaturesManager.remove_patches()
+     full_args = get_full_args()
++    args = vars(args)
+     for k, v in args.items():
+         setattr(full_args, k, v)
+     MindSpeedFeaturesManager.apply_features_pre_patches(full_args)
+diff --git a/mindspeed/te/pytorch/attention/dot_product_attention/dot_product_attention.py b/mindspeed/te/pytorch/attention/dot_product_attention/dot_product_attention.py
+index ac4eabe5..f82a581b 100644
+--- a/mindspeed/te/pytorch/attention/dot_product_attention/dot_product_attention.py
++++ b/mindspeed/te/pytorch/attention/dot_product_attention/dot_product_attention.py
+@@ -330,6 +330,8 @@ class DotProductAttention(torch.nn.Module):
+         inference_params: Any = None,
+         pad_between_seqs: Optional[bool] = None,
+         fp8_output: Optional[bool] = False,
++        local_cp_size=None,
++        cp_group=None,
+     ) -> torch.Tensor:
+         """
+         Dot Product Attention Layer.
+@@ -653,6 +655,30 @@ class MindSpeedTEDotProductAttention(DotProductAttention):
+         packed_seq_params: PackedSeqParams = None,
+         ):
+         """Forward."""
++        packed_seq_kwargs = (
++            {key: getattr(packed_seq_params, key) for key in self.kept_packed_seq_params}
++            if packed_seq_params is not None
++            else {}
++        )
++
++        # Honor the per-call attn_mask_type. Megatron passes AttnMaskType.no_mask for the
++        # (bidirectional) vision tower and AttnMaskType.causal for the causal LM decoder.
++        # The previous code unconditionally built a triu(2048) causal mask and forwarded
++        # config.attention_mask_type, forcing CAUSAL attention even for no_mask -> the
++        # Qwen3-VL vision tower's full attention was computed causally (wrong image embeds).
++        if attn_mask_type == AttnMaskType.no_mask:
++            # Full (bidirectional) attention: no causal mask; per-image segmentation is
++            # handled by cu_seqlens in packed_seq_params (sparse_mode=0 via 'no_mask').
++            core_attn_out = super().forward(
++                query,
++                key,
++                value,
++                None,
++                attn_mask_type='no_mask',
++                **packed_seq_kwargs,
++            )
++            return core_attn_out
++
+         if (
+             attention_mask is None and
+             self.attn_mask_type == AttnMaskType.causal
+@@ -660,12 +686,6 @@ class MindSpeedTEDotProductAttention(DotProductAttention):
+             self.config.sparse_mode = 2
+             attention_mask = get_attention_mask(self.config)
+ 
+-        packed_seq_kwargs = (
+-            {key: getattr(packed_seq_params, key) for key in self.kept_packed_seq_params}
+-            if packed_seq_params is not None
+-            else {}
+-        )
+-
+         core_attn_out = super().forward(
+             query,
+             key,
+diff --git a/mindspeed/core/fusions/fused_rope.py b/mindspeed/core/fusions/fused_rope.py
+index 70f7cb08..15189c2d 100644
+--- a/mindspeed/core/fusions/fused_rope.py
++++ b/mindspeed/core/fusions/fused_rope.py
+@@ -65,7 +65,7 @@ def apply_rotary_pos_emb_bshd(
+     cos_ = (torch.cos(freqs) * _mscale).to(t.dtype)
+     sin_ = (torch.sin(freqs) * _mscale).to(t.dtype)
+ 
+-    if getattr(args, "use_fused_rotary_pos_emb"):
++    if getattr(args, "use_fused_rotary_pos_emb", False):
+         mode = 1 if rotary_interleaved else 0
+         t = npu_rotary_position_embedding(t.contiguous(), cos_, sin_, mode).to(t.dtype)
+     else:
+@@ -82,7 +82,7 @@ def transformer_config_post_init_wrapper(fn):
+         self.apply_rope_fusion = False
+         fn(self)
+         self.apply_rope_fusion = ori_apply_rope_fusion
+-        if ((getattr(self, "multi_head_latent_attention") or getattr(self, "multi_latent_attention"))
++        if ((getattr(self, "multi_head_latent_attention", False) or getattr(self, "multi_latent_attention", False))
+                 and self.rope_type == "yarn"):
+             self.apply_rope_fusion = False
+         del ori_apply_rope_fusion
+diff --git a/mindspeed/core/models/common/embeddings/rotary_pos_embedding.py b/mindspeed/core/models/common/embeddings/rotary_pos_embedding.py
+index c0319906..c0741691 100644
+--- a/mindspeed/core/models/common/embeddings/rotary_pos_embedding.py
++++ b/mindspeed/core/models/common/embeddings/rotary_pos_embedding.py
+@@ -79,7 +79,7 @@ def apply_rotary_pos_emb_bshd(t: Tensor, freqs: Tensor, rotary_interleaved: bool
+     cos_ = (torch.cos(freqs) * _mscale).to(t.dtype)
+     sin_ = (torch.sin(freqs) * _mscale).to(t.dtype)
+ 
+-    if getattr(args, "use_fused_rotary_pos_emb"):
++    if getattr(args, "use_fused_rotary_pos_emb", False):
+         mode = 1 if rotary_interleaved else 0
+         t = npu_rotary_position_embedding(t.contiguous(), cos_, sin_, mode).to(t.dtype)
+     else:
+diff --git a/mindspeed/features_manager/features_manager.py b/mindspeed/features_manager/features_manager.py
+index ba476adc..f5924b98 100644
+--- a/mindspeed/features_manager/features_manager.py
++++ b/mindspeed/features_manager/features_manager.py
+@@ -43,7 +43,10 @@ class MindSpeedFeaturesManager:
+             post_validate_features_args(args=args)  # args.x = old_x
+         """
+         for feature in cls.FEATURES_LIST:
+-            feature.pre_validate_args(args)
++            try:
++                feature.pre_validate_args(args)
++            except AttributeError:
++                pass
+ 
+     @classmethod
+     def post_validate_features_args(cls, args):
+@@ -54,13 +57,19 @@ class MindSpeedFeaturesManager:
+             post_validate_features_args(args=args)  # args.x = old_x
+         """
+         for feature in cls.FEATURES_LIST:
+-            feature.post_validate_args(args)
++            try:
++                feature.post_validate_args(args)
++            except AttributeError:
++                pass
+ 
+     @classmethod
+     def validate_features_args(cls, args):
+         """Validate arguments of all features."""
+         for feature in cls.FEATURES_LIST:
+-            feature.validate_args(args)
++            try:
++                feature.validate_args(args)
++            except AttributeError:
++                pass
+ 
+     @classmethod
+     def remove_patches(cls):
+diff --git a/mindspeed/features_manager/moe/fb_overlap.py b/mindspeed/features_manager/moe/fb_overlap.py
+index 6c76d891..e099c51d 100644
+--- a/mindspeed/features_manager/moe/fb_overlap.py
++++ b/mindspeed/features_manager/moe/fb_overlap.py
+@@ -16,6 +16,8 @@ class MoEFwdBwdOverlapFeature(MindSpeedFeature):
+         group.add_argument('--moe-unperm2-mem-optim-swap', action='store_true')
+ 
+     def validate_args(self, args):
++        if not hasattr(args, 'moe_fb_overlap'):
++            return
+         self.incompatible_check(args, 'moe_alltoall_overlap_comm')
+         self.incompatible_check(args, 'overlap_grad_reduce')
+         self.incompatible_check(args, 'moe_hierarchical_alltoallv')
+diff --git a/mindspeed/features_manager/moe/moe_zero_memory.py b/mindspeed/features_manager/moe/moe_zero_memory.py
+index 4d95e64e..94f26109 100644
+--- a/mindspeed/features_manager/moe/moe_zero_memory.py
++++ b/mindspeed/features_manager/moe/moe_zero_memory.py
+@@ -23,6 +23,8 @@ class MoEZeroMemoryFeature(MindSpeedFeature):
+                                 'in each pp stage.')
+ 
+     def pre_validate_args(self, args):
++        if not hasattr(args, "moe_zero_memory_num_layers") or not hasattr(args, "moe_zero_memory"):
++            return
+         #Zero Memory check.
+         if args.moe_zero_memory_num_layers is not None:
+             num_layers_per_pipeline_stage = args.num_layers // args.pipeline_model_parallel_size
diff --git a/docker/npu_patch/series.conf b/docker/npu_patch/series.conf
new file mode 100644
index 000000000..10d543dee
--- /dev/null
+++ b/docker/npu_patch/series.conf
@@ -0,0 +1,17 @@
+# Ordered NPU patch series used by the runtime CI reconciler.
+# Format: target_worktree|image_patch|source_patch.
+# Apply top-to-bottom from source_patch; revert bottom-to-top from image_patch.
+#
+# PATCH MAINTENANCE:
+# - Add: add the required COPY/apply operations to docker/Dockerfile.npu, then
+#   add this entry in the same application order.
+# - Delete: remove the explicit Dockerfile operation and this entry, then delete
+#   the patch file. The previous image keeps the OLD bytes needed for revert.
+# - Rename/reorder: update both Dockerfile.npu and this file in the same change.
+# Ordinary patch changes must not add patch-specific logic to the CI script.
+/vllm-workspace/vllm|vllm.patch|docker/npu_patch/vllm.patch
+/vllm-workspace/vllm-ascend|vllm-ascend.patch|docker/npu_patch/vllm-ascend.patch
+/root/Megatron-LM|megatron-common.patch|docker/patch/latest/megatron.patch
+/root/Megatron-LM|megatron.patch|docker/npu_patch/megatron.patch
+/root/Megatron-Bridge|megatron-bridge.patch|docker/npu_patch/megatron-bridge.patch
+/root/MindSpeed|mindspeed.patch|docker/npu_patch/mindspeed.patch
diff --git a/docker/npu_patch/vllm-ascend.patch b/docker/npu_patch/vllm-ascend.patch
new file mode 100644
index 000000000..768852a59
--- /dev/null
+++ b/docker/npu_patch/vllm-ascend.patch
@@ -0,0 +1,472 @@
+diff --git a/vllm_ascend/distributed/weight_transfer/__init__.py b/vllm_ascend/distributed/weight_transfer/__init__.py
+index d6434f05f..fd7223836 100644
+--- a/vllm_ascend/distributed/weight_transfer/__init__.py
++++ b/vllm_ascend/distributed/weight_transfer/__init__.py
+@@ -33,6 +33,11 @@ def register_engine():
+         "vllm_ascend.distributed.weight_transfer.npu_ipc_engine",
+         "NPUIPCWeightTransferEngine",
+     )
++    WeightTransferTrainerFactory.register_engine(
++        "hccl",
++        "vllm_ascend.distributed.weight_transfer.hccl_engine",
++        "HCCLTrainerWeightTransferEngine",
++    )
+     WeightTransferTrainerFactory.register_engine(
+         "npu_ipc",
+         "vllm_ascend.distributed.weight_transfer.npu_ipc_engine",
+diff --git a/vllm_ascend/distributed/weight_transfer/hccl_engine.py b/vllm_ascend/distributed/weight_transfer/hccl_engine.py
+index 023ffd1ae..7b184e279 100644
+--- a/vllm_ascend/distributed/weight_transfer/hccl_engine.py
++++ b/vllm_ascend/distributed/weight_transfer/hccl_engine.py
+@@ -3,8 +3,9 @@
+ """HCCL-based weight transfer engine."""
+ 
+ from collections.abc import Callable, Iterator
+-from dataclasses import dataclass
+-from typing import TYPE_CHECKING, Any
++from concurrent.futures import ThreadPoolExecutor
++from dataclasses import asdict, dataclass
++from typing import TYPE_CHECKING, Any, ClassVar
+ 
+ import torch
+ 
+@@ -14,6 +15,11 @@ if TYPE_CHECKING:
+ from vllm.config import VllmConfig
+ from vllm.config.weight_transfer import WeightTransferConfig
+ from vllm.distributed.weight_transfer.base import (
++    ParamMeta,
++    TrainerInitInfo,
++    TrainerWeightTransferEngine,
++    VLLMWeightSyncClient,
++    WeightSource,
+     WeightTransferEngine,
+     WeightTransferInitInfo,
+     WeightTransferUpdateInfo,
+@@ -26,6 +32,19 @@ from vllm_ascend.distributed.weight_transfer.packed_tensor import (
+ )
+ 
+ 
++@dataclass
++class HCCLTrainerInitInfo(TrainerInitInfo):
++    """Stateful trainer configuration; rank 0 owns the HCCL endpoint."""
++
++    backend: ClassVar[str] = "hccl"
++    master_address: str
++    master_port: int
++    world_size: int
++    packed: bool = True
++    packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES
++    packed_num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS
++
++
+ @dataclass
+ class HCCLWeightTransferInitInfo(WeightTransferInitInfo):
+     """Initialization info for HCCL weight transfer backend."""
+@@ -338,3 +357,114 @@ class HCCLWeightTransferEngine(WeightTransferEngine[HCCLWeightTransferInitInfo,
+         pg = StatelessProcessGroup.create(host=master_address, port=master_port, rank=rank, world_size=world_size)
+         pyhccl = PyHcclCommunicator(pg, device=device)
+         return pyhccl
++
++
++class HCCLTrainerWeightTransferEngine(TrainerWeightTransferEngine[HCCLTrainerInitInfo]):
++    """Stateful control plane over the existing HCCL broadcast transport.
++
++    Every trainer rank replays the source collectives. Only rank 0 opens the
++    transfer communicator and drives worker RPCs. The frozen receiver expects
++    packed geometry on each update, so derive it from the trainer init info.
++    """
++
++    init_info_cls = HCCLTrainerInitInfo
++
++    def __init__(self, *, init_info: HCCLTrainerInitInfo, client: VLLMWeightSyncClient, source: WeightSource) -> None:
++        super().__init__(client=client, source=source, is_sender=init_info.is_sender)
++        self.init_info = init_info
++        self.model_update_group: PyHcclCommunicator | None = None
++
++    @classmethod
++    def trainer_init(
++        cls,
++        init_info: HCCLTrainerInitInfo,
++        *,
++        client: VLLMWeightSyncClient,
++        source: WeightSource,
++    ) -> "HCCLTrainerWeightTransferEngine":
++        engine = cls(init_info=init_info, client=client, source=source)
++        if not engine.is_sender:
++            return engine
++
++        worker_info = HCCLWeightTransferInitInfo(
++            master_address=init_info.master_address,
++            master_port=init_info.master_port,
++            rank_offset=1,
++            world_size=init_info.world_size,
++        )
++        # Both endpoints must rendezvous concurrently.
++        executor = ThreadPoolExecutor(max_workers=1)
++        try:
++            future = executor.submit(client.init_weight_transfer_engine, asdict(worker_info))
++            if future.done():
++                future.result()
++            engine.model_update_group = HCCLWeightTransferEngine.trainer_init(worker_info)
++            future.result()
++        finally:
++            executor.shutdown(wait=False)
++        return engine
++
++    def send_weights(self) -> None:
++        # Metadata export can itself contain Megatron collectives.
++        meta = self.source.metadata()
++        if not self.is_sender:
++            for _ in self.source:
++                pass
++            torch.npu.current_stream().synchronize()
++            return
++
++        assert self.model_update_group is not None, "HCCL trainer has been shut down."
++        info = self.init_info
++        update_info = HCCLWeightTransferUpdateInfo(
++            names=[m.name for m in meta],
++            dtype_names=[str(m.dtype).split(".")[-1] for m in meta],
++            shapes=[list(m.shape) for m in meta],
++            packed=info.packed,
++            packed_buffer_size_bytes=info.packed_buffer_size_bytes,
++            packed_num_buffers=info.packed_num_buffers,
++        )
++        self.client.start_weight_update()
++        executor = ThreadPoolExecutor(max_workers=1)
++        try:
++            future = executor.submit(self.client.update_weights, asdict(update_info))
++            if future.done():
++                future.result()
++            HCCLWeightTransferEngine.trainer_send_weights(
++                self._checked_iter(self.source, meta),
++                HCCLTrainerSendWeightsArgs(
++                    group=self.model_update_group,
++                    packed=info.packed,
++                    packed_buffer_size_bytes=info.packed_buffer_size_bytes,
++                    packed_num_buffers=info.packed_num_buffers,
++                ),
++            )
++            future.result()
++        finally:
++            # A failed broadcast can leave the worker RPC waiting in HCCL.
++            # Surface the error; joining that thread here would hide it forever.
++            executor.shutdown(wait=False)
++        self.client.finish_weight_update()
++        torch.npu.current_stream().synchronize()
++
++    @staticmethod
++    def _checked_iter(source: WeightSource, meta: list[ParamMeta]) -> Iterator[tuple[str, torch.Tensor]]:
++        # Receive buffer sizes and packed chunk boundaries come from metadata.
++        sent = 0
++        for name, tensor in source:
++            if sent >= len(meta):
++                raise ValueError(f"WeightSource yielded more parameters than metadata(): {name!r}.")
++            expected = meta[sent]
++            if name != expected.name or tensor.dtype != expected.dtype or tuple(tensor.shape) != expected.shape:
++                raise ValueError(
++                    f"WeightSource metadata() disagrees with iteration at index {sent}: "
++                    f"expected {expected.name!r} {expected.dtype} {expected.shape}, "
++                    f"got {name!r} {tensor.dtype} {tuple(tensor.shape)}."
++                )
++            sent += 1
++            # Unpacked HCCL reads contiguous storage from data_ptr().
++            yield name, tensor if tensor.is_contiguous() else tensor.contiguous()
++        if sent != len(meta):
++            raise ValueError(f"WeightSource yielded {sent} parameters but metadata() declared {len(meta)}.")
++
++    def shutdown(self) -> None:
++        self.model_update_group = None
+diff --git a/vllm_ascend/distributed/weight_transfer/npu_ipc_engine.py b/vllm_ascend/distributed/weight_transfer/npu_ipc_engine.py
+index 7a63adf7b..d3e3a7fc4 100644
+--- a/vllm_ascend/distributed/weight_transfer/npu_ipc_engine.py
++++ b/vllm_ascend/distributed/weight_transfer/npu_ipc_engine.py
+@@ -156,12 +156,14 @@ class NPUIPCWeightTransferEngine(  # type: ignore[no-redef]
+         self.packed = init_info.packed
+ 
+     def start_weight_update(self) -> None:
+-        """No-op for NPU IPC engine (no layerwise reloading)."""
+-        pass
++        from vllm.model_executor.model_loader.reload import initialize_layerwise_reload
++
++        initialize_layerwise_reload(self.model)
+ 
+     def finish_weight_update(self) -> None:
+-        """No-op for NPU IPC engine (no layerwise reloading)."""
+-        pass
++        from vllm.model_executor.model_loader.reload import finalize_layerwise_reload
++
++        finalize_layerwise_reload(self.model, self.model_config)
+ 
+     def receive_weights(self, update_info: NPUIPCWeightTransferUpdateInfo) -> None:
+         """Receive weights from the trainer via NPU IPC handles.
+@@ -170,6 +172,10 @@ class NPUIPCWeightTransferEngine(  # type: ignore[no-redef]
+             update_info: NPU IPC update info containing parameter names,
+                 dtypes, shapes, and IPC handles.
+         """
++        # A rank-local slot may have no weights in this chunk (and no handle).
++        if not update_info.names:
++            return
++
+         # Use the worker's assigned device rather than the ambient current
+         # device: the receive path is no longer wrapped in
+         # ``with torch.device(self.device)`` by the caller, so the current
+@@ -219,7 +225,7 @@ class NPUIPCWeightTransferEngine(  # type: ignore[no-redef]
+                 weight = rebuild_npu_tensor(*list_args)
+                 weights.append((name, weight))
+ 
+-            self.model.load_weights(weights)
++        self.model.load_weights(weights)
+ 
+     def shutdown(self) -> None:
+         pass
+@@ -288,6 +294,8 @@ class NPUIPCTrainerWeightTransferEngine(IPCTrainerWeightTransferEngine):
+             self.client.finish_weight_update()
+         self._post_send_sync()
+         del weight_refs
++        torch.npu.ipc_collect()
++        torch.npu.empty_cache()
+ 
+     def _send(self, source: "WeightSource") -> list[torch.Tensor] | None:
+         if self.packed:
+diff --git a/vllm_ascend/distributed/weight_transfer/packed_tensor.py b/vllm_ascend/distributed/weight_transfer/packed_tensor.py
+index a35d9af8d..d55c988c3 100644
+--- a/vllm_ascend/distributed/weight_transfer/packed_tensor.py
++++ b/vllm_ascend/distributed/weight_transfer/packed_tensor.py
+@@ -39,6 +39,7 @@ def packed_broadcast_producer(
+     target_packed_tensor_size = buffer_size_bytes
+ 
+     streams = [torch.npu.Stream() for _ in range(num_buffers)]
++    source_stream = torch.npu.current_stream()
+     buffer_idx = 0
+ 
+     packing_tensor_list: list[list[torch.Tensor]] = [[] for _ in range(num_buffers)]
+@@ -50,6 +51,8 @@ def packed_broadcast_producer(
+         # Synchronize the current stream (waits for previous
+         # iteration's work on this buffer to finish)
+         streams[buffer_idx].synchronize()
++        # Wait for tensors produced on the caller's stream before packing.
++        streams[buffer_idx].wait_stream(source_stream)
+         # Start tasks for the new buffer in a new stream
+         with torch.npu.stream(streams[buffer_idx]):
+             # Initialize the packing tensor list and sizes
+@@ -206,10 +209,10 @@ def packed_npu_ipc_producer(
+ ) -> Iterator[dict[str, Any]]:
+     """Pack tensors into a reusable NPU IPC buffer and yield chunks.
+ 
+-    Allocates a single NPU buffer of ``buffer_size_bytes`` and registers
+-    it for IPC once via ``reduce_tensor``.  Each chunk's packed data is
+-    copied into this buffer before yielding, so only one IPC-shared
+-    allocation exists for the lifetime of the transfer.
++    Allocates a single NPU buffer of ``buffer_size_bytes``. Each chunk
++    publishes a fresh IPC reference via ``reduce_tensor`` so its consumer
++    releases that reference exactly once. The underlying buffer is reused
++    for the lifetime of the transfer.
+ 
+     Args:
+         iterator: Iterator of (name, tensor) pairs.
+@@ -218,9 +221,6 @@ def packed_npu_ipc_producer(
+         buffer_size_bytes: Exact capacity of the reusable IPC buffer.
+     """
+     ipc_buffer = torch.empty(buffer_size_bytes, dtype=torch.uint8, device="npu")
+-    # Store only the rebuild args (drop the func); the consumer rebuilds with
+-    # the well-known ``rebuild_npu_tensor``, mirroring upstream's CUDA IPC engine.
+-    _, ipc_args = reduce_tensor(ipc_buffer)
+ 
+     names: list[str] = []
+     shapes: list[list[int]] = []
+@@ -240,6 +240,7 @@ def packed_npu_ipc_producer(
+ 
+         if total_bytes and total_bytes + flat.numel() > buffer_size_bytes:
+             torch.npu.current_stream().synchronize()
++            _, ipc_args = reduce_tensor(ipc_buffer)
+             yield {
+                 "names": names,
+                 "shapes": shapes,
+@@ -259,6 +260,7 @@ def packed_npu_ipc_producer(
+ 
+     if total_bytes:
+         torch.npu.current_stream().synchronize()
++        _, ipc_args = reduce_tensor(ipc_buffer)
+         yield {
+             "names": names,
+             "shapes": shapes,
+diff --git a/vllm_ascend/worker/v2/model_runner.py b/vllm_ascend/worker/v2/model_runner.py
+index c388619f2..38a7d2fb7 100644
+--- a/vllm_ascend/worker/v2/model_runner.py
++++ b/vllm_ascend/worker/v2/model_runner.py
+@@ -17,7 +17,7 @@
+ # This file is a part of the vllm-ascend project.
+ #
+ 
+-from contextlib import contextmanager
++from contextlib import AbstractContextManager, contextmanager
+ 
+ import numpy as np
+ import torch
+@@ -239,9 +239,13 @@ class NPUModelRunner(GPUModelRunner):
+             self.pp_handler.broadcast_draft_tokens()
+         return output
+ 
+-    def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
++    def initialize_kv_cache(
++        self,
++        kv_cache_config: KVCacheConfig,
++        kv_cache_allocation_context: AbstractContextManager | None = None,
++    ) -> None:
+         with graph_manager_wrapper(self):
+-            super().initialize_kv_cache(kv_cache_config)
++            super().initialize_kv_cache(kv_cache_config, kv_cache_allocation_context=kv_cache_allocation_context)
+             if self.pcp_manager is not None:
+                 assert isinstance(self.pcp_manager, AscendPCPManager)
+                 self.pcp_manager.vllm_config = self.vllm_config
+diff --git a/vllm_ascend/worker/worker.py b/vllm_ascend/worker/worker.py
+index 6d99ac76a..b7df4df22 100644
+--- a/vllm_ascend/worker/worker.py
++++ b/vllm_ascend/worker/worker.py
+@@ -321,7 +321,48 @@ class NPUWorker(WorkerBase):
+
+     def start_weight_update(self) -> None:
+         """Begin a new weight update; prepares the model for layerwise reload."""
++        with set_current_vllm_config(self.vllm_config):
++            self._start_weight_update()
++
++    def get_draft_model(self) -> nn.Module | None:
++        return self.model_runner.get_draft_model()
++
++    def supports_draft_weight_updates(self) -> bool:
++        engine = self.weight_transfer_engine
++        speculative_config = self.speculative_config
++        get_draft_model = getattr(self.model_runner, "get_draft_model", None)
++        return (
++            engine is not None
++            and engine.supports_draft_weight_update
++            and callable(get_draft_model)
++            and get_draft_model() is not None
++            and speculative_config is not None
++            and speculative_config.draft_model_config is not None
++        )
++
++    def _set_draft_weight_update_target(self) -> None:
++        assert self.weight_transfer_engine is not None
++        draft_model = self.get_draft_model()
++        if draft_model is None:
++            raise RuntimeError("Draft model weight update requested, but no draft model is configured.")
++        speculative_config = self.speculative_config
++        if speculative_config is None or speculative_config.draft_model_config is None:
++            raise RuntimeError("Draft model weight update requested, but no draft model config is configured.")
++        self.weight_transfer_engine.set_weight_update_target(draft_model, speculative_config.draft_model_config)
++
++    def start_draft_weight_update(self) -> None:
++        """Retarget the native engine at the draft model for this session."""
++        with set_current_vllm_config(self.vllm_config):
++            self._start_weight_update(is_draft=True)
++
++    def _start_weight_update(self, is_draft: bool = False) -> None:
+         self._check_weight_transfer_engine()
++        assert self.weight_transfer_engine is not None
++
++        if is_draft and not self.weight_transfer_engine.supports_draft_weight_update:
++            raise RuntimeError(
++                f"{type(self.weight_transfer_engine).__name__} does not support draft model weight updates."
++            )
+
+         if self._weight_update_active:
+             raise RuntimeError(
+@@ -330,11 +371,16 @@ class NPUWorker(WorkerBase):
+
+         self._check_nz_disabled()
+
+-        assert self.weight_transfer_engine is not None
+-        self.weight_transfer_engine.start_weight_update()
++        try:
++            if is_draft:
++                self._set_draft_weight_update_target()
++            self.weight_transfer_engine.start_weight_update()
++        except BaseException:
++            self.weight_transfer_engine.reset_weight_update_target()
++            raise
+         self._weight_update_active = True
+
+-    def update_weights(self, update_info: dict) -> None:
++    def update_weights(self, update_info: dict | list[dict]) -> None:
+         """Receive a chunk of weights from the trainer and load them in place."""
+         self._check_weight_transfer_engine()
+         assert self.weight_transfer_engine is not None
+@@ -343,11 +389,19 @@ class NPUWorker(WorkerBase):
+         if not self._weight_update_active:
+             raise RuntimeError("start_weight_update must be called before update_weights.")
+
+-        try:
+-            self.weight_transfer_engine.update_weights(update_info)
+-        except BaseException:
+-            self._weight_update_active = False
+-            raise
++        with set_current_vllm_config(self.vllm_config):
++            try:
++                if isinstance(update_info, list):
++                    parallel_config = self.vllm_config.parallel_config
++                    worker_rank = parallel_config.data_parallel_rank * parallel_config.world_size + self.rank
++                    local_update_info = update_info[worker_rank]
++                else:
++                    local_update_info = update_info
++                self.weight_transfer_engine.update_weights(local_update_info)
++            except BaseException:
++                self._weight_update_active = False
++                self.weight_transfer_engine.reset_weight_update_target()
++                raise
+
+     def finish_weight_update(self) -> None:
+         """Finish the current weight update; runs layerwise postprocessing."""
+@@ -357,8 +411,12 @@ class NPUWorker(WorkerBase):
+             raise RuntimeError("start_weight_update must be called before finish_weight_update.")
+
+         assert self.weight_transfer_engine is not None
+-        self.weight_transfer_engine.finish_weight_update()
+-        self._weight_update_active = False
++        with set_current_vllm_config(self.vllm_config):
++            try:
++                self.weight_transfer_engine.finish_weight_update()
++            finally:
++                self._weight_update_active = False
++                self.weight_transfer_engine.reset_weight_update_target()
+
+     def shutdown(self) -> None:
+         if ensure_kv_transfer_shutdown is not None:
+@@ -448,7 +506,9 @@ class NPUWorker(WorkerBase):
+         # take current memory snapshot
+         self.init_snapshot = MemorySnapshot(device=device)
+         self.requested_memory = self.init_snapshot.total_memory * self.cache_config.gpu_memory_utilization
+-        if self.init_snapshot.free_memory < self.requested_memory:
++        weight_transfer_config = self.vllm_config.weight_transfer_config
++        uses_ipc_weight_transfer = weight_transfer_config is not None and weight_transfer_config.backend == "npu_ipc"
++        if not uses_ipc_weight_transfer and self.init_snapshot.free_memory < self.requested_memory:
+             GiB = lambda b: round(b / GiB_bytes, 2)
+             raise ValueError(
+                 f"Free memory on device "
+@@ -597,7 +657,9 @@ class NPUWorker(WorkerBase):
+         self.non_torch_memory = profile_result.non_torch_increase
+
+         free_gpu_memory = profile_result.after_profile.free_memory
+-        assert self.init_snapshot.free_memory > free_gpu_memory, (
++        weight_transfer_config = self.vllm_config.weight_transfer_config
++        uses_ipc_weight_transfer = weight_transfer_config is not None and weight_transfer_config.backend == "npu_ipc"
++        assert uses_ipc_weight_transfer or self.init_snapshot.free_memory > free_gpu_memory, (
+             "Error in memory profiling. "
+             f"Initial free memory {GiB(self.init_snapshot.free_memory)} GiB, "
+             f"current free memory {GiB(free_gpu_memory)} GiB. "
+@@ -1118,8 +1180,12 @@ class NPUWorker(WorkerBase):
+             from contextlib import nullcontext
+
+             context = nullcontext()  # type: ignore
+-        with context:
+-            self.model_runner.initialize_kv_cache(kv_cache_config)
++        if self.use_v2_model_runner:
++            # MRV2 bookkeeping must survive sleep; pool only the KV data.
++            self.model_runner.initialize_kv_cache(kv_cache_config, kv_cache_allocation_context=context)
++        else:
++            with context:
++                self.model_runner.initialize_kv_cache(kv_cache_config)
+
+         # MRV2's scheduler emits new_block_ids_to_zero whenever this flag is
+         # set, so its worker-side consumer must use the same condition. Keep the
diff --git a/docker/npu_patch/vllm.patch b/docker/npu_patch/vllm.patch
new file mode 100644
index 000000000..142941525
--- /dev/null
+++ b/docker/npu_patch/vllm.patch
@@ -0,0 +1,158 @@
+diff --git a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py
+index f304bf677b..71a17e9363 100644
+--- a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py
++++ b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py
+@@ -240,6 +240,8 @@ class GenerateStreamResponse(BaseModel):
+     )
+     choices: list[GenerateResponseStreamChoice]
+     usage: UsageInfo | None = Field(default=None)
++    weight_version: str | None = None
++    request_spec_decode_stats: dict[str, Any] | None = Field(default=None)
+
+
+ class GenerateResponse(BaseModel):
+@@ -255,6 +257,8 @@ class GenerateResponse(BaseModel):
+     created: int | None = None
+     choices: list[GenerateResponseChoice]
+     usage: UsageInfo | None = Field(default=None)
++    weight_version: str | None = None
++    request_spec_decode_stats: dict[str, Any] | None = Field(default=None)
+     prompt_logprobs: list[dict[int, Logprob] | None] | None = None
+
+     kv_transfer_params: dict[str, Any] | None = Field(
+diff --git a/vllm/entrypoints/scale_out/token_in_token_out/serving.py b/vllm/entrypoints/scale_out/token_in_token_out/serving.py
+index bbbd85137c..809f1a66ea 100644
+--- a/vllm/entrypoints/scale_out/token_in_token_out/serving.py
++++ b/vllm/entrypoints/scale_out/token_in_token_out/serving.py
+@@ -14,6 +14,7 @@ from vllm.engine.protocol import EngineClient
+ from vllm.entrypoints.chat_utils import AsyncMultiModalItemTracker
+ from vllm.entrypoints.generate.base.serving import (
+     GenerateBaseServing,
++    build_spec_decoding_metrics,
+     clamp_prompt_logprobs,
+ )
+ from vllm.entrypoints.openai.chat_completion.protocol import (
+@@ -261,6 +262,7 @@ class ServingTokens(GenerateBaseServing):
+         )
+
+         assert result_generator is not None
++        weight_version = await self.engine_client.get_weight_version()
+
+         if request.stream:
+             return self.serve_tokens_stream_generator(
+@@ -269,10 +271,16 @@ class ServingTokens(GenerateBaseServing):
+                 request_id,
+                 model_name,
+                 request_metadata,
++                weight_version,
+             )
+
+         return await self.serve_tokens_full_generator(
+-            request, result_generator, request_id, model_name, request_metadata
++            request,
++            result_generator,
++            request_id,
++            model_name,
++            request_metadata,
++            weight_version,
+         )
+
+     async def serve_tokens_full_generator(
+@@ -282,6 +290,7 @@ class ServingTokens(GenerateBaseServing):
+         request_id: str,
+         model_name: str,
+         request_metadata: RequestResponseMetadata,
++        weight_version: str | None,
+     ) -> ErrorResponse | GenerateResponse:
+         created_time = int(time.time())
+         final_res: RequestOutput | None = None
+@@ -355,6 +364,11 @@ class ServingTokens(GenerateBaseServing):
+                 cached_tokens=final_res.num_cached_tokens
+             )
+
++        spec_decode_metrics = build_spec_decoding_metrics(final_res)
++        request_spec_decode_stats = (
++            spec_decode_metrics.model_dump() if spec_decode_metrics else None
++        )
++
+         request_metadata.final_usage_info = usage
+
+         response = GenerateResponse(
+@@ -363,6 +377,8 @@ class ServingTokens(GenerateBaseServing):
+             model=model_name,
+             choices=choices,
+             usage=usage,
++            weight_version=weight_version,
++            request_spec_decode_stats=request_spec_decode_stats,
+             prompt_logprobs=clamp_prompt_logprobs(final_res.prompt_logprobs),
+             kv_transfer_params=final_res.kv_transfer_params,
+             ec_transfer_params=final_res.ec_transfer_params,
+@@ -396,11 +412,13 @@ class ServingTokens(GenerateBaseServing):
+         request_id: str,
+         model_name: str,
+         request_metadata: RequestResponseMetadata,
++        weight_version: str | None,
+     ) -> AsyncGenerator[str, None]:
+         num_prompt_tokens = 0
+         num_generated_tokens: list[int] = []
+         first_iteration = True
+         num_cached_tokens = None
++        request_spec_decode_stats: dict[str, object] | None = None
+         sampling_params: SamplingParams = request.sampling_params
+
+         include_usage, include_continuous_usage = should_include_usage(
+@@ -409,6 +427,9 @@ class ServingTokens(GenerateBaseServing):
+
+         try:
+             async for res in result_generator:
++                spec_decode_metrics = build_spec_decoding_metrics(res)
++                if spec_decode_metrics is not None:
++                    request_spec_decode_stats = spec_decode_metrics.model_dump()
+                 if first_iteration:
+                     if res.prompt_token_ids is not None:
+                         num_prompt_tokens = len(res.prompt_token_ids)
+@@ -448,6 +469,8 @@ class ServingTokens(GenerateBaseServing):
+
+                     chunk = GenerateStreamResponse(
+                         request_id=request_id,
++                        weight_version=weight_version,
++                        request_spec_decode_stats=request_spec_decode_stats,
+                         choices=[
+                             GenerateResponseStreamChoice(
+                                 index=i,
+@@ -482,6 +505,8 @@ class ServingTokens(GenerateBaseServing):
+             if include_usage:
+                 final_chunk = GenerateStreamResponse(
+                     request_id=request_id,
++                    weight_version=weight_version,
++                    request_spec_decode_stats=request_spec_decode_stats,
+                     choices=[],
+                     usage=final_usage_info,
+                 )
+diff --git a/vllm/model_executor/model_loader/reload/meta.py b/vllm/model_executor/model_loader/reload/meta.py
+index a8f8023cf2..d62f39a2b8 100644
+--- a/vllm/model_executor/model_loader/reload/meta.py
++++ b/vllm/model_executor/model_loader/reload/meta.py
+@@ -30,5 +30,6 @@ SKIP_LOAD_TENSORS: set[str] = {
+     "expert_global_to_physical",
+     "expert_physical_to_global",
+     "expert_local_to_global",
++    "expert_ids_per_ep_rank",
+     "e_score_correction_bias",
+ }
+diff --git a/vllm/model_executor/models/glm4_moe_lite_mtp.py b/vllm/model_executor/models/glm4_moe_lite_mtp.py
+index 1d30a0eaf6..37239f7f62 100644
+--- a/vllm/model_executor/models/glm4_moe_lite_mtp.py
++++ b/vllm/model_executor/models/glm4_moe_lite_mtp.py
+@@ -129,6 +129,10 @@ class Glm4MoeLiteMultiTokenPredictorLayer(nn.Module):
+     ) -> torch.Tensor:
+         assert inputs_embeds is not None
+         # masking inputs at position 0, as not needed by MTP
+-        inputs_embeds[positions == 0] = 0
++        # Avoid dynamic bool indexing during NPU graph capture.
++        mask = (positions == 0).unsqueeze(-1)
++        inputs_embeds = torch.where(
++            mask, torch.zeros_like(inputs_embeds), inputs_embeds
++        )
+         inputs_embeds = self.enorm(inputs_embeds)
+         previous_hidden_states = self.hnorm(previous_hidden_states)
diff --git a/docs/en/get_started/NPU.md b/docs/en/get_started/NPU.md
new file mode 100644
index 000000000..b6084fd3e
--- /dev/null
+++ b/docs/en/get_started/NPU.md
@@ -0,0 +1,205 @@
+# NPU
+
+⚠️ If you encounter problems running vime on Ascend NPU, feel free to open an
+issue on [vllm-project/vime](https://github.com/vllm-project/vime/issues).
+
+## Introduction
+
+If you are running vime on Ascend NPU, please refer to the following materials.
+This tutorial explains how to set up the runtime environment and provides an
+end-to-end example for running GRPO training. It uses the **Megatron** training
+backend together with the **vLLM Ascend** rollout backend, synchronizing actor 
+weights to vLLM through the native HCCL weight-sync path.
+
+The current NPU support targets Ascend **Atlas A2 / A3** (aarch64) hosts with the
+Ascend driver and **CANN 9.0.0** (Toolkit, Kernels, and NNAL/ATB) installed.
+Only `python==3.12` is supported.
+
+## Docker
+
+The recommended path for validation is the published vime NPU image.
+
+```bash
+export IMAGE=quay.io/ascend/vime:vime-latest
+# A2:  export IMAGE=quay.io/ascend/vime:vime-a2-latest
+
+docker pull "${IMAGE}"
+```
+
+For source builds and dependency debugging, the patch list and pinned commits are
+documented in [`docker/npu_patch/README.md`](https://github.com/vllm-project/vime/blob/npu/docker/npu_patch/README.md).
+
+## Quick Start
+
+### Environment Setup
+
+Start the container, mounting the Ascend devices and driver files. Device names
+and driver mount paths vary by host; reuse the mounts from a known working vLLM
+Ascend container if the layout differs.
+
+```bash
+docker run -d --name vime-npu -it --net=host --shm-size=1024g \
+    --privileged=true \
+    --cap-add=SYS_PTRACE \
+    --device=/dev/davinci_manager \
+    --device=/dev/hisi_hdc \
+    --device=/dev/devmm_svm \
+    -v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
+    -v /usr/local/dcmi:/usr/local/dcmi \
+    -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
+    -v /usr/local/sbin:/usr/local/sbin \
+    -v /home:/home \
+    -v /mnt:/mnt \
+    -v /tmp:/tmp \
+    -v /data:/data \
+    -v /path/to:/path/to \
+    -v /usr/share/zoneinfo/Asia/Shanghai:/etc/localtime \
+    "${IMAGE}"
+
+docker exec -it vime-npu bash
+```
+
+Inside the container, initialize the CANN environment before training:
+
+```bash
+source /usr/local/Ascend/ascend-toolkit/set_env.sh
+source /usr/local/Ascend/nnal/atb/set_env.sh
+```
+
+### Prepare Model and Data
+
+Set `MODEL_ROOT` to a host-visible directory that will hold both the checkpoint
+and the dataset, then download the Qwen3-4B checkpoint and the DAPO Math 17K
+dataset:
+
+```bash
+export MODEL_ROOT=/root
+mkdir -p ${MODEL_ROOT}/models ${MODEL_ROOT}/datasets
+
+# hf checkpoint
+hf download Qwen/Qwen3-4B \
+  --local-dir ${MODEL_ROOT}/models/Qwen3-4B
+
+# train data
+hf download --repo-type dataset zhuzilin/dapo-math-17k \
+  --local-dir ${MODEL_ROOT}/datasets/dapo-math-17k
+```
+
+### Example: Qwen3-4B
+
+We provide an example to run GRPO training with
+[Qwen3-4B](https://huggingface.co/Qwen/Qwen3-4B) on 8 NPUs (4 for the actor,
+4 for rollout), please refer to:
+[scripts/run-qwen3-4B-npu.sh](../../../scripts/run-qwen3-4B-npu.sh).
+Just run:
+
+```bash
+cd /root/vime
+
+# Source these explicitly if not already initialized by the image.
+source /usr/local/Ascend/ascend-toolkit/set_env.sh
+source /usr/local/Ascend/nnal/atb/set_env.sh
+
+DATA_ROOT="${MODEL_ROOT:-/root}" bash scripts/run-qwen3-4B-npu.sh \
+  2>&1 | tee /root/vime/train_qwen3_4b_vllm.log
+```
+
+The full log is written to `/root/vime/train_qwen3_4b_vllm.log`.
+
+⚠️ Note: The main difference between the NPU training script and the NVIDIA one
+is the Ascend-specific environment variables — `ASCEND_RT_VISIBLE_DEVICES`
+selects the NPUs, and `RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1` lets
+Ray schedule them correctly. The reference target is an Atlas A3 host with 16
+visible NPUs; on an 8-NPU host, set
+`ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7`.
+
+We show the training script below:
+
+```bash
+export SLIME_SCRIPT_TRAIN_BACKEND=megatron
+export PYTHONPATH="/root/Megatron-Bridge/src:/root/Megatron-LM/:$PYTHONPATH"
+export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1
+export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050
+export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050
+export HYDRA_FULL_ERROR=1
+export MASTER_PORT=$(shuf -i 20000-65000 -n 1)  # or any free port
+export DISABLE_L2_CACHE=1
+export VLLM_ASCEND_ENABLE_NZ=0
+
+SCRIPT_DIR="/root/vime/scripts/"
+source "${SCRIPT_DIR}/models/qwen3-4B.sh"
+LOG_FILE="/root/vime/train_qwen3_4b_vllm.log"
+MODEL_ROOT="${MODEL_ROOT:-/root}"
+
+python /root/vime/train.py \
+  --train-backend megatron \
+  --actor-num-nodes 1 \
+  --actor-num-gpus-per-node 4 \
+  --rollout-num-gpus 4 \
+  --rollout-num-gpus-per-engine 4 \
+  ${MODEL_ARGS[@]} \
+  \
+  --hf-checkpoint ${MODEL_ROOT}/models/Qwen3-4B/ \
+  \
+  --prompt-data ${MODEL_ROOT}/datasets/dapo-math-17k/dapo-math-17k.jsonl \
+  --input-key prompt \
+  --label-key label \
+  --apply-chat-template \
+  --rollout-shuffle \
+  --rm-type math \
+  \
+  --rollout-backend vllm \
+  --vllm-weight-sync-mode native \
+  --vllm-gpu-memory-utilization 0.6 \
+  --vllm-enable-sleep-mode \
+  --vllm-max-model-len 4096 \
+  \
+  --num-rollout 200 \
+  --rollout-batch-size 32 \
+  --n-samples-per-prompt 8 \
+  --rollout-max-response-len 2048 \
+  --rollout-temperature 1.0 \
+  --global-batch-size 256 \
+  --balance-data \
+  \
+  --advantage-estimator grpo \
+  --kl-loss-coef 0.0 \
+  --kl-loss-type low_var_kl \
+  --kl-coef 0.00 \
+  --entropy-coef 0.0 \
+  --eps-clip 0.2 \
+  --eps-clip-high 0.28 \
+  \
+  --optimizer adam \
+  --lr 1e-6 \
+  --lr-decay-style constant \
+  --weight-decay 0.1 \
+  --adam-beta1 0.9 \
+  --adam-beta2 0.98 \
+  \
+  --tensor-model-parallel-size 4 \
+  --pipeline-model-parallel-size 1 \
+  --context-parallel-size 1 \
+  --expert-model-parallel-size 1 \
+  --expert-tensor-parallel-size 1 \
+  --recompute-granularity full \
+  --recompute-method uniform \
+  --recompute-num-layers 1 \
+  --use-dynamic-batch-size \
+  --max-tokens-per-gpu 8192 \
+  --load ${MODEL_ROOT}/models/Qwen3-4B \
+  --megatron-to-hf-mode bridge \
+  \
+  --attention-dropout 0.0 \
+  --hidden-dropout 0.0 \
+  --accumulate-allreduce-grads-in-fp32 \
+  --attention-softmax-in-fp32 \
+  --attention-backend flash \
+  --micro-batch-size 1 \
+  --use-flash-attn \
+  \
+  --train-memory-margin-bytes 2147483648 \
+  2>&1 | tee -a "$LOG_FILE"
+```
diff --git a/examples/fully_async/run-qwen3-4B-fully_async-npu.sh b/examples/fully_async/run-qwen3-4B-fully_async-npu.sh
new file mode 100755
index 000000000..9c05eff36
--- /dev/null
+++ b/examples/fully_async/run-qwen3-4B-fully_async-npu.sh
@@ -0,0 +1,140 @@
+#!/bin/bash
+# Tiny end-to-end fully-async GRPO example using Qwen3-4B-Instruct on the
+# dapo-math-17k dataset. Designed to run on a single 4-GPU node in a few minutes
+#
+# Prerequisites:
+#   /root/models/Qwen3-4B/             (HF checkpoint)
+#   /root/datasets/dapo-math-17k/dapo-math-17k.jsonl
+
+# clean any leftover ray/vllm
+pkill -9 -f '[v]llm serve|VLL[M]::'
+pkill -9 -f VLLM
+sleep 3
+ray stop --force
+pkill -9 ray
+pkill -9 python
+sleep 3
+pkill -9 ray
+pkill -9 python
+pkill -9 redis
+
+set -ex
+
+export PYTHONUNBUFFERED=1
+export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15
+export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1
+export RAY_USE_UVLOOP=0 # upstream: Ray's uvloop integration has caused intermittent async actor issues.
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050
+export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050
+export HYDRA_FULL_ERROR=1
+export DISABLE_L2_CACHE=1
+export VLLM_ASCEND_ENABLE_NZ=0
+export VLLM_USE_AOT_COMPILE=0
+export PYTHONPATH="/root/Megatron-Bridge/src:/root/Megatron-LM/:${PYTHONPATH:-}"
+
+unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY
+
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
+source "${SCRIPT_DIR}/../../scripts/models/qwen3-4B.sh"
+
+MODEL_DIR=${MODEL_DIR:-/root/models/Qwen3-4B}
+DATA_PATH=${DATA_PATH:-/root/datasets/dapo-math-17k/dapo-math-17k.jsonl}
+
+CKPT_ARGS=(
+   --hf-checkpoint "${MODEL_DIR}"
+   --load "${MODEL_DIR}"
+   --ref-load "${MODEL_DIR}"
+   --megatron-to-hf-mode bridge
+   --save /tmp/vime_fully_async_demo/
+   --save-interval 9999
+)
+
+ROLLOUT_ARGS=(
+   # ↓↓↓ This is the only knob you need to flip to go fully-async ↓↓↓
+   --rollout-function-path vime.rollout.fully_async_rollout.generate_rollout_fully_async
+
+   --prompt-data "${DATA_PATH}"
+   --input-key prompt
+   --label-key label
+   --apply-chat-template
+   --rollout-shuffle
+   --rm-type math
+   --num-rollout 200
+   --rollout-batch-size 32
+   --n-samples-per-prompt 8
+   --rollout-max-response-len 2048
+   --rollout-temperature 1
+   --global-batch-size 256
+   --balance-data
+)
+
+PERF_ARGS=(
+   --tensor-model-parallel-size 2
+   --sequence-parallel
+   --pipeline-model-parallel-size 1
+   --context-parallel-size 1
+   --expert-model-parallel-size 1
+   --expert-tensor-parallel-size 1
+   --recompute-granularity full
+   --recompute-method uniform
+   --recompute-num-layers 1
+   --use-dynamic-batch-size
+   --max-tokens-per-gpu 8192
+)
+
+GRPO_ARGS=(
+   --advantage-estimator grpo
+   --kl-loss-coef 0.00
+   --kl-loss-type low_var_kl
+   --kl-coef 0.00
+   --entropy-coef 0.00
+   --eps-clip 0.2
+   --eps-clip-high 0.28
+)
+
+OPTIMIZER_ARGS=(
+   --optimizer adam
+   --lr 1e-6
+   --lr-decay-style constant
+   --weight-decay 0.1
+   --adam-beta1 0.9
+   --adam-beta2 0.98
+   --optimizer-cpu-offload
+   --overlap-cpu-optimizer-d2h-h2d
+   --use-precision-aware-optimizer
+)
+
+VLLM_ARGS=(
+   --rollout-num-gpus-per-engine 2
+   --vllm-gpu-memory-utilization 0.6
+)
+
+MISC_ARGS=(
+   --attention-dropout 0.0
+   --hidden-dropout 0.0
+   --accumulate-allreduce-grads-in-fp32
+   --attention-softmax-in-fp32
+   --attention-backend flash
+   --use-flash-attn
+)
+
+ray start --head --node-ip-address 127.0.0.1 --disable-usage-stats
+
+# fully-async splits actor / rollout onto disjoint GPUs (no colocation).
+ACTOR_GPUS=2
+ROLLOUT_GPUS=2
+
+ray job submit --address="http://127.0.0.1:8265" \
+   -- python3 train_async.py \
+   --actor-num-nodes 1 \
+   --actor-num-gpus-per-node "${ACTOR_GPUS}" \
+   --rollout-num-gpus "${ROLLOUT_GPUS}" \
+   ${MODEL_ARGS[@]} \
+   ${CKPT_ARGS[@]} \
+   ${ROLLOUT_ARGS[@]} \
+   ${OPTIMIZER_ARGS[@]} \
+   ${GRPO_ARGS[@]} \
+   ${PERF_ARGS[@]} \
+   ${VLLM_ARGS[@]} \
+   ${MISC_ARGS[@]}
diff --git a/examples/geo3k_vlm/run_geo3k_vlm_npu.sh b/examples/geo3k_vlm/run_geo3k_vlm_npu.sh
new file mode 100644
index 000000000..911035b22
--- /dev/null
+++ b/examples/geo3k_vlm/run_geo3k_vlm_npu.sh
@@ -0,0 +1,216 @@
+#!/bin/bash
+
+# Single-turn Qwen3-VL GRPO on geo3k for Ascend NPU.
+set -eo pipefail
+
+MODEL_NAME=${VIME_SCRIPT_MODEL_NAME:-Qwen3-VL-8B-Instruct}
+DATASET_NAME=${VIME_SCRIPT_DATASET_NAME:-chenhegu/geo3k_imgurl}
+NUM_GPUS=${VIME_SCRIPT_NUM_GPUS:-8}
+NUM_ROLLOUT=${VIME_SCRIPT_NUM_ROLLOUT:-3000}
+DATA_ROOT="/root/datasets/$(basename "$DATASET_NAME")"
+MODEL_ROOT="/root/models/${MODEL_NAME}"
+
+if ! [[ "$NUM_GPUS" =~ ^[1-9][0-9]*$ && "$NUM_ROLLOUT" =~ ^[1-9][0-9]*$ ]]; then
+   echo "Error: VIME_SCRIPT_NUM_GPUS and VIME_SCRIPT_NUM_ROLLOUT must be positive integers"
+   exit 1
+fi
+
+VALID_MODELS="
+  Qwen3-VL-2B-Instruct
+  Qwen3-VL-4B-Instruct
+  Qwen3-VL-8B-Instruct
+  Qwen3-VL-2B-Thinking
+  Qwen3-VL-4B-Thinking
+  Qwen3-VL-8B-Thinking
+"
+if ! echo "$VALID_MODELS" | grep -qw "$MODEL_NAME"; then
+   echo "Error: unsupported MODEL_NAME=${MODEL_NAME}"
+   exit 1
+fi
+
+if [ -z "${VIME_SCRIPT_EXTERNAL_RAY:-}" ] || [ "$VIME_SCRIPT_EXTERNAL_RAY" = "0" ]; then
+   USE_EXTERNAL_RAY=0
+else
+   USE_EXTERNAL_RAY=1
+fi
+
+export PYTHONUNBUFFERED=1
+export PYTORCH_ALLOC_CONF=expandable_segments:True
+export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
+export PYTHONPATH="/root/Megatron-Bridge/src:/root/Megatron-LM/:$PYTHONPATH"
+export ASCEND_RT_VISIBLE_DEVICES=${ASCEND_RT_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15}
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1
+export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050
+export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050
+export HYDRA_FULL_ERROR=1
+export MASTER_PORT=${MASTER_PORT:-$(shuf -i 20000-65000 -n 1)}
+export VLLM_ASCEND_ENABLE_NZ=0
+export VLLM_USE_AOT_COMPILE=0
+export ASCEND_TOOLKIT_HOME=/usr/local/Ascend/ascend-toolkit/latest/
+export ASCEND_OPP_PATH=/usr/local/Ascend/ascend-toolkit/latest/opp/
+export ASCEND_AICPU_PATH=/usr/local/Ascend/ascend-toolkit/latest/
+export ASCEND_HOME_PATH=/usr/local/Ascend/ascend-toolkit/latest/
+export set_env_path=/usr/local/Ascend/nnal/atb/set_env.sh
+
+IFS=',' read -r -a VISIBLE_NPUS <<<"${ASCEND_RT_VISIBLE_DEVICES}"
+REQUIRED_NPUS=$((NUM_GPUS * 2))
+if [ "${#VISIBLE_NPUS[@]}" -lt "$REQUIRED_NPUS" ]; then
+   echo "Error: actor and rollout require ${REQUIRED_NPUS} visible NPUs"
+   exit 1
+fi
+
+unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY ALL_PROXY all_proxy
+
+mkdir -p /root/models /root/datasets
+if [ ! -d "$MODEL_ROOT" ]; then
+   hf download "Qwen/${MODEL_NAME}" --local-dir "$MODEL_ROOT"
+fi
+if [ ! -d "$DATA_ROOT" ]; then
+   hf download --repo-type dataset "$DATASET_NAME" --local-dir "$DATA_ROOT"
+fi
+
+MODEL_LOAD_ARGS=(
+   --hf-checkpoint "$MODEL_ROOT"
+   --rotary-base 5000000
+)
+
+ROLLOUT_ARGS=(
+   --prompt-data "${DATA_ROOT}/train.parquet"
+   --input-key problem
+   --label-key answer
+   --apply-chat-template
+   --rollout-shuffle
+   --rm-type math
+   --num-rollout "$NUM_ROLLOUT"
+   --rollout-batch-size 64
+   --n-samples-per-prompt 8
+   --rollout-max-response-len 4096
+   --rollout-temperature 1.0
+   --global-batch-size 512
+)
+
+GRPO_ARGS=(
+   --advantage-estimator grpo
+   --kl-loss-coef 0.00
+   --kl-loss-type low_var_kl
+   --kl-coef 0.00
+   --entropy-coef 0.00
+   --eps-clip 0.2
+   --eps-clip-high 0.28
+)
+
+OPTIMIZER_ARGS=(
+   --optimizer adam
+   --lr 1e-6
+   --lr-decay-style constant
+   --weight-decay 0.1
+   --adam-beta1 0.9
+   --adam-beta2 0.98
+)
+
+VLLM_ARGS=(
+   --rollout-num-gpus-per-engine 1
+   --vllm-gpu-memory-utilization "${VLLM_GPU_MEMORY_UTILIZATION:-0.8}"
+   --vllm-max-model-len 16384
+   --vllm-generation-config auto
+   --vllm-logprobs-mode processed_logprobs
+)
+
+if [ -n "${WANDB_API_KEY:-}" ]; then
+   WANDB_ARGS=(
+      --use-wandb
+      --wandb-project vime-geo3k-vlm
+      --wandb-group "${MODEL_NAME,,}-megatron-vllm-${NUM_GPUS}npu"
+      --wandb-key "$WANDB_API_KEY"
+      --disable-wandb-random-suffix
+   )
+else
+   WANDB_ARGS=()
+fi
+
+BACKEND_ARGS=(
+   --train-backend megatron
+   --load "$MODEL_ROOT"
+   --tensor-model-parallel-size 4
+   --sequence-parallel
+   --pipeline-model-parallel-size 1
+   --context-parallel-size 1
+   --expert-model-parallel-size 1
+   --expert-tensor-parallel-size 1
+   --recompute-granularity full
+   --recompute-method uniform
+   --recompute-num-layers 1
+   --use-dynamic-batch-size
+   --max-tokens-per-gpu 4096
+   --attention-dropout 0.0
+   --hidden-dropout 0.0
+   --accumulate-allreduce-grads-in-fp32
+   --attention-softmax-in-fp32
+   --attention-backend flash
+   --megatron-to-hf-mode bridge
+)
+
+VIME_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../.." &>/dev/null && pwd)"
+MODEL_ARGS_FILE=$(echo "$MODEL_NAME" | sed 's/-Instruct//g; s/-Thinking//g; s/Qwen3-VL-/qwen3-/g; s/-2B/-1.7B/g')
+MODEL_ARGS_ROTARY_BASE=5000000 source "${VIME_DIR}/scripts/models/${MODEL_ARGS_FILE}.sh"
+
+pkill -9 -f '[v]llm serve|VLL[M]::' || true
+if [ "$USE_EXTERNAL_RAY" = "0" ]; then
+   ray stop --force || true
+   pkill -9 ray || true
+fi
+pkill -9 vime || true
+pkill -9 redis || true
+
+export MASTER_ADDR=${MASTER_ADDR:-127.0.0.1}
+export no_proxy="127.0.0.1,${MASTER_ADDR}"
+if [ "$USE_EXTERNAL_RAY" = "0" ]; then
+   ray start --head --node-ip-address "$MASTER_ADDR" --disable-usage-stats \
+      --dashboard-host=0.0.0.0 --dashboard-port=8265
+fi
+
+RUNTIME_ENV_KEYS=(
+   CUDA_DEVICE_MAX_CONNECTIONS
+   RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES
+   PYTHONPATH
+   PYTORCH_ALLOC_CONF
+   PYTORCH_NPU_ALLOC_CONF
+   VLLM_ASCEND_ENABLE_NZ
+   ASCEND_TOOLKIT_HOME
+   ASCEND_OPP_PATH
+   ASCEND_AICPU_PATH
+   ASCEND_HOME_PATH
+   set_env_path
+   HYDRA_FULL_ERROR
+   HCCL_HOST_SOCKET_PORT_RANGE
+   HCCL_NPU_SOCKET_PORT_RANGE
+   no_proxy
+   MASTER_ADDR
+)
+RUNTIME_ENV_JSON=$(python - "${RUNTIME_ENV_KEYS[@]}" <<'PY'
+import json
+import os
+import sys
+
+print(json.dumps({"env_vars": {key: os.environ[key] for key in sys.argv[1:]}}))
+PY
+)
+
+ray job submit --address=http://127.0.0.1:8265 \
+   --runtime-env-json="$RUNTIME_ENV_JSON" \
+   -- "$(command -v python)" train.py \
+   --actor-num-nodes 1 \
+   --actor-num-gpus-per-node "$NUM_GPUS" \
+   --rollout-num-gpus "$NUM_GPUS" \
+   --multimodal-keys '{"image": "images"}' \
+   "${MODEL_ARGS[@]}" \
+   "${MODEL_LOAD_ARGS[@]}" \
+   "${ROLLOUT_ARGS[@]}" \
+   "${GRPO_ARGS[@]}" \
+   "${OPTIMIZER_ARGS[@]}" \
+   "${VLLM_ARGS[@]}" \
+   "${WANDB_ARGS[@]}" \
+   "${BACKEND_ARGS[@]}" \
+   --no-gradient-accumulation-fusion \
+   --use-flash-attn
diff --git a/examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn_grpo_npu.py b/examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn_grpo_npu.py
new file mode 100644
index 000000000..35124a53a
--- /dev/null
+++ b/examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn_grpo_npu.py
@@ -0,0 +1,154 @@
+import os
+
+from vime.utils.external_utils.command_utils import execute_train
+
+MODEL_NAME = os.environ.get("VIME_SCRIPT_MODEL_NAME", "Qwen3-VL-8B-Instruct")
+SUPPORTED_MODELS = {
+    "Qwen3-VL-2B-Instruct",
+    "Qwen3-VL-4B-Instruct",
+    "Qwen3-VL-8B-Instruct",
+    "Qwen3-VL-2B-Thinking",
+    "Qwen3-VL-4B-Thinking",
+    "Qwen3-VL-8B-Thinking",
+}
+if MODEL_NAME not in SUPPORTED_MODELS:
+    raise ValueError(f"Unsupported VIME_SCRIPT_MODEL_NAME={MODEL_NAME}")
+
+NUM_ROLLOUT = int(os.environ.get("VIME_SCRIPT_NUM_ROLLOUT", "3000"))
+if NUM_ROLLOUT <= 0:
+    raise ValueError("VIME_SCRIPT_NUM_ROLLOUT must be positive")
+
+MODEL_ROOT = os.environ.get("VIME_SCRIPT_MODEL_ROOT", "/root/models")
+DATA_ROOT = os.environ.get("VIME_SCRIPT_DATA_ROOT", "/root/datasets/geo3k_imgurl_processed")
+TRAIN_DATA_PATH = os.path.join(DATA_ROOT, "train.parquet")
+
+
+def get_megatron_model_type(model_name: str) -> str:
+    model_type = model_name.replace("-Instruct", "").replace("-Thinking", "")
+    model_type = model_type.replace("Qwen3-VL-", "qwen3-")
+    return model_type.replace("-2B", "-1.7B")
+
+
+def execute():
+    model_path = os.path.join(MODEL_ROOT, MODEL_NAME)
+    if not os.path.isdir(model_path):
+        raise FileNotFoundError(f"Model not found: {model_path}")
+    if not os.path.isfile(TRAIN_DATA_PATH):
+        raise FileNotFoundError(f"Dataset not found: {TRAIN_DATA_PATH}")
+
+    wandb_api_key = os.environ.get("WANDB_API_KEY")
+    wandb_args = (
+        (
+            "--use-wandb "
+            "--wandb-project vime-dev "
+            "--wandb-group geo3k_vlm_multi_turn "
+            f"--wandb-key '{wandb_api_key}' "
+        )
+        if wandb_api_key
+        else ""
+    )
+
+    rollout_args = (
+        f"--prompt-data {TRAIN_DATA_PATH} "
+        "--input-key problem "
+        "--label-key answer "
+        '--multimodal-keys \'{"image": "images"}\' '
+        "--rm-type math "
+        "--custom-generate-function-path examples.geo3k_vlm_multi_turn.rollout.generate "
+        "--custom-config-path examples/geo3k_vlm_multi_turn/geo3k_vlm_multi_turn_config.yaml "
+        "--rollout-shuffle "
+        f"--num-rollout {NUM_ROLLOUT} "
+        "--rollout-batch-size 32 "
+        "--n-samples-per-prompt 8 "
+        "--rollout-max-response-len 4096 "
+        "--rollout-temperature 1 "
+        "--global-batch-size 256 "
+    )
+
+    grpo_args = (
+        "--advantage-estimator grpo "
+        "--kl-loss-coef 0.00 "
+        "--kl-loss-type low_var_kl "
+        "--kl-coef 0.00 "
+        "--entropy-coef 0.00 "
+        "--eps-clip 0.2 "
+        "--eps-clip-high 0.28 "
+        "--use-kl-loss "
+    )
+
+    optimizer_args = (
+        "--optimizer adam "
+        "--lr 1e-6 "
+        "--lr-decay-style constant "
+        "--weight-decay 0.1 "
+        "--adam-beta1 0.9 "
+        "--adam-beta2 0.98 "
+        "--optimizer-cpu-offload "
+        "--overlap-cpu-optimizer-d2h-h2d "
+        "--use-precision-aware-optimizer "
+    )
+
+    vllm_args = (
+        "--rollout-num-gpus-per-engine 1 "
+        "--vllm-gpu-memory-utilization 0.6 "
+        "--vllm-max-model-len 16384 "
+        "--vllm-generation-config auto "
+        "--vllm-logprobs-mode processed_logprobs "
+    )
+
+    megatron_args = (
+        "--train-backend megatron "
+        f"--load {model_path} "
+        f"--ref-load {model_path} "
+        "--tensor-model-parallel-size 4 "
+        "--sequence-parallel "
+        "--pipeline-model-parallel-size 1 "
+        "--context-parallel-size 1 "
+        "--expert-model-parallel-size 1 "
+        "--expert-tensor-parallel-size 1 "
+        "--recompute-granularity full "
+        "--recompute-method uniform "
+        "--recompute-num-layers 1 "
+        "--use-dynamic-batch-size "
+        "--max-tokens-per-gpu 16384 "
+        "--balance-data "
+        "--attention-dropout 0.0 "
+        "--hidden-dropout 0.0 "
+        "--accumulate-allreduce-grads-in-fp32 "
+        "--attention-softmax-in-fp32 "
+        "--attention-backend flash "
+        "--megatron-to-hf-mode bridge "
+    )
+
+    misc_args = (
+        "--actor-num-nodes 1 "
+        "--actor-num-gpus-per-node 8 "
+        "--rollout-num-gpus 8 "
+        "--no-gradient-accumulation-fusion "
+        "--use-flash-attn "
+    )
+
+    megatron_model_type = get_megatron_model_type(MODEL_NAME)
+    os.environ["MODEL_ARGS_ROTARY_BASE"] = "5000000"
+
+    train_args = (
+        f"--hf-checkpoint {model_path} "
+        f"{rollout_args} "
+        f"{optimizer_args} "
+        f"{grpo_args} "
+        f"{vllm_args} "
+        f"{megatron_args} "
+        f"{misc_args} "
+        f"{wandb_args} "
+    )
+
+    execute_train(
+        train_args=train_args,
+        num_gpus_per_node=16,  # actor 8 + rollout 8 (non-colocate)
+        megatron_model_type=megatron_model_type,
+        extra_env_vars={"WANDB_API_KEY": wandb_api_key} if wandb_api_key else {},
+    )
+
+
+if __name__ == "__main__":
+    execute()
diff --git a/examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn_ppo_npu.py b/examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn_ppo_npu.py
new file mode 100644
index 000000000..43c04ba4e
--- /dev/null
+++ b/examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn_ppo_npu.py
@@ -0,0 +1,160 @@
+import os
+import tempfile
+
+from vime.utils.external_utils.command_utils import execute_train
+
+MODEL_NAME = os.environ.get("VIME_SCRIPT_MODEL_NAME", "Qwen3-VL-2B-Instruct")
+assert MODEL_NAME in {
+    "Qwen3-VL-2B-Instruct",
+    "Qwen3-VL-4B-Instruct",
+    "Qwen3-VL-8B-Instruct",
+    "Qwen3-VL-2B-Thinking",
+    "Qwen3-VL-4B-Thinking",
+    "Qwen3-VL-8B-Thinking",
+}
+
+EXTERNAL_RAY = int(os.environ.get("VIME_SCRIPT_EXTERNAL_RAY", "0"))
+TRAIN_BACKEND = os.environ.get("VIME_SCRIPT_TRAIN_BACKEND", "fsdp").lower()
+assert TRAIN_BACKEND in {"fsdp", "megatron"}
+
+DATASET_NAME = "VeraIsHere/geo3k_imgurl_processed"
+DATA_ROOT = "/path/to/datasets/geo3k_imgurl_processed"
+TRAIN_DATA_PATH = os.path.join(DATA_ROOT, "train.parquet")
+
+
+def get_megatron_model_type(model_name: str) -> str:
+    model_type = model_name.replace("-Instruct", "").replace("-Thinking", "")
+    model_type = model_type.replace("Qwen3-VL-", "qwen3-")
+    return model_type.replace("-2B", "-1.7B")
+
+
+def execute():
+    megatron_config = tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False)
+    megatron_config.write(
+        """
+megatron:
+  - name: default
+    role: critic
+    overrides:
+      lr: 1e-5
+"""
+    )
+    megatron_config.close()
+
+    ckpt_args = f"--hf-checkpoint /path/to/model/checkpoints/{MODEL_NAME} "
+
+    wandb_args = (
+        (
+            "--use-wandb "
+            "--wandb-project vime-dev "
+            "--wandb-group geo3k_vlm_multi_turn "
+            f"--wandb-key '{wandb_api_key}' "
+        )
+        if (wandb_api_key := os.environ.get("WANDB_API_KEY"))
+        else ""
+    )
+
+    rollout_args = (
+        f"--prompt-data {TRAIN_DATA_PATH} "
+        "--input-key problem "
+        "--label-key answer "
+        '--multimodal-keys \'{"image": "images"}\' '
+        "--rm-type math "
+        "--apply-chat-template "
+        "--custom-generate-function-path examples.geo3k_vlm_multi_turn.rollout.generate "
+        "--custom-config-path examples/geo3k_vlm_multi_turn/geo3k_vlm_multi_turn_config.yaml "
+        "--rollout-shuffle "
+        "--num-rollout 3000 "
+        "--rollout-batch-size 32 "
+        "--n-samples-per-prompt 8 "
+        "--rollout-max-response-len 4096 "
+        "--rollout-temperature 1 "
+        "--global-batch-size 256 "
+    )
+
+    ppo_args = (
+        "--advantage-estimator ppo "
+        "--kl-loss-coef 0.00 "
+        "--kl-loss-type k1 "
+        "--kl-coef 0.00 "
+        "--entropy-coef 0.00 "
+        "--eps-clip 4e-4 "
+        "--num-critic-only-steps 1 "
+        "--normalize-advantages "
+    )
+
+    optimizer_args = (
+        "--optimizer adam "
+        "--lr 1e-6 "
+        "--lr-decay-style constant "
+        "--weight-decay 0.1 "
+        "--adam-beta1 0.9 "
+        "--adam-beta2 0.98 "
+        "--optimizer-cpu-offload "
+        "--overlap-cpu-optimizer-d2h-h2d "
+        "--use-precision-aware-optimizer "
+    )
+
+    vllm_args = "--rollout-num-gpus-per-engine 1 " "--vllm-gpu-memory-utilization 0.6 "
+
+    megatron_args = (
+        "--train-backend megatron "
+        f"--load /path/to/model/checkpoints/{MODEL_NAME} "
+        f"--ref-load /path/to/model/checkpoints/{MODEL_NAME} "
+        "--tensor-model-parallel-size 4 "
+        "--sequence-parallel "
+        "--pipeline-model-parallel-size 1 "
+        "--context-parallel-size 1 "
+        "--expert-model-parallel-size 1 "
+        "--expert-tensor-parallel-size 1 "
+        "--recompute-granularity full "
+        "--recompute-method uniform "
+        "--recompute-num-layers 1 "
+        "--use-dynamic-batch-size "
+        "--max-tokens-per-gpu 16384 "
+        "--balance-data "
+        "--attention-dropout 0.0 "
+        "--hidden-dropout 0.0 "
+        "--accumulate-allreduce-grads-in-fp32 "
+        "--attention-softmax-in-fp32 "
+        "--attention-backend flash "
+        "--megatron-to-hf-mode bridge "
+    )
+
+    misc_args = (
+        "--actor-num-nodes 1 "
+        "--actor-num-gpus-per-node 8 "
+        "--rollout-num-gpus 8 "
+        "--no-gradient-accumulation-fusion "
+        "--use-flash-attn "
+    )
+
+    if TRAIN_BACKEND == "megatron":
+        backend_args = megatron_args
+        megatron_model_type = get_megatron_model_type(MODEL_NAME)
+        os.environ["MODEL_ARGS_ROTARY_BASE"] = "5000000"
+    else:
+        exit()
+
+    train_args = (
+        f"--megatron-config-path {megatron_config.name} "
+        f"{ckpt_args} "
+        f"{rollout_args} "
+        f"{optimizer_args} "
+        f"{ppo_args} "
+        f"{vllm_args} "
+        f"{backend_args} "
+        f"{misc_args} "
+        f"{wandb_args} "
+    )
+
+    execute_train(
+        train_args=train_args,
+        num_gpus_per_node=16,  # actor 8 + rollout 8 (non-colocate)
+        megatron_model_type=megatron_model_type,
+        extra_env_vars=({"WANDB_API_KEY": os.environ["WANDB_API_KEY"]} if os.environ.get("WANDB_API_KEY") else {}),
+    )
+
+
+if __name__ == "__main__":
+    execute()
diff --git a/examples/multi_agent/agent_system.py b/examples/multi_agent/agent_system.py
index ebd521900..c04640e6f 100644
--- a/examples/multi_agent/agent_system.py
+++ b/examples/multi_agent/agent_system.py
@@ -12,12 +12,18 @@
 from .prompts import SOLVER_PROMPT_TEMPLATE, generate_rewriter_template, generate_select_template
 
 
-async def generate_response(args, prompt, key):
+async def generate_response(args, prompt, key, worker_id: int | None = None):
     try:
         sampling_params = args.sampling_params
         tokenizer = args.tokenizer
         max_context_length = args.rollout_max_context_len
         sample = deepcopy(args.sample)
+        sample.metadata = dict(sample.metadata or {})
+        sample.metadata["multi_agent_role"] = key
+        sample.metadata["multi_agent_parent_group_index"] = sample.group_index
+        sample.metadata["multi_agent_parent_index"] = sample.index
+        if worker_id is not None:
+            sample.metadata["multi_agent_worker_id"] = worker_id
 
         url = f"http://{args.vllm_router_ip}:{args.vllm_router_port}/inference/v1/generate"
 
@@ -90,11 +96,11 @@ class Agent:
     def __init__(self):
         pass
 
-    async def run(self, args, prompt, max_retries: int = 1, key: str = None) -> str:
+    async def run(self, args, prompt, max_retries: int = 1, key: str = None, worker_id: int | None = None) -> str:
         """Runs the agent by sending a prompt to the LLM."""
         for _i in range(max_retries):
             try:
-                response = await generate_response(args, prompt, key=key)
+                response = await generate_response(args, prompt, key=key, worker_id=worker_id)
                 return response
             except Exception as e:
                 print(f"Error querying LLM: {e}")
@@ -109,10 +115,10 @@ class SolverAgent(Agent):
     def __init__(self):
         super().__init__()
 
-    async def generate_initial_solution(self, args, problem_statement) -> str:
+    async def generate_initial_solution(self, args, problem_statement, worker_id: int) -> str:
         """Generates the first solution attempt."""
         prompt = SOLVER_PROMPT_TEMPLATE.format(problem_statement=problem_statement)
-        return await self.run(args, prompt, max_retries=3, key="solver")
+        return await self.run(args, prompt, max_retries=3, key="solver", worker_id=worker_id)
 
 
 class RewriterAgent(Agent):
@@ -121,7 +127,7 @@ class RewriterAgent(Agent):
     def __init__(self):
         super().__init__()
 
-    async def rewrite(self, args, problem_statement, previous_solutions: list[str]) -> str:
+    async def rewrite(self, args, problem_statement, previous_solutions: list[str], worker_id: int) -> str:
         """Generates the rewrited solution."""
 
         # Build the prompt template dynamically.
@@ -133,7 +139,7 @@ async def rewrite(self, args, problem_statement, previous_solutions: list[str])
             format_params[f"solution{i+1}"] = solution
 
         prompt = template.format(**format_params)
-        return await self.run(args, prompt, max_retries=1, key="rewriter")
+        return await self.run(args, prompt, max_retries=1, key="rewriter", worker_id=worker_id)
 
 
 class SelectorAgent(Agent):
@@ -142,7 +148,7 @@ class SelectorAgent(Agent):
     def __init__(self):
         super().__init__()
 
-    async def select(self, args, problem_statement, candidate_solutions: list[str]) -> str:
+    async def select(self, args, problem_statement, candidate_solutions: list[str], worker_id: int = 0) -> str:
         """Generates the rewrited solution."""
 
         # Build the prompt template dynamically.
@@ -154,7 +160,7 @@ async def select(self, args, problem_statement, candidate_solutions: list[str])
             format_params[f"solution{i+1}"] = solution
 
         prompt = template.format(**format_params)
-        return await self.run(args, prompt, max_retries=10, key="selector")
+        return await self.run(args, prompt, max_retries=10, key="selector", worker_id=worker_id)
 
     def extract_selected_solution_idx(self, response: str, candidate_solutions: list[str]) -> int:
         """Extracts the selected solution ID from the response."""
@@ -173,7 +179,7 @@ def extract_selected_solution_idx(self, response: str, candidate_solutions: list
 
 async def rewrite_worker(args, previous_solutions, problem_statement, worker_id):
     rewriter = RewriterAgent()
-    new_solution = await rewriter.rewrite(args, problem_statement, previous_solutions)
+    new_solution = await rewriter.rewrite(args, problem_statement, previous_solutions, worker_id)
     return new_solution
 
 
@@ -184,7 +190,7 @@ async def solver_worker(args, problem_statement, worker_id):
 
     try:
         solver = SolverAgent()
-        current_solution = await solver.generate_initial_solution(args, problem_statement)
+        current_solution = await solver.generate_initial_solution(args, problem_statement, worker_id)
         return current_solution
 
     except Exception as e:
@@ -245,7 +251,7 @@ def reward_adjustment(samples, reward_weight):
 
     # Selection
     selector = SelectorAgent()
-    response = await selector.select(args, problem_statement, rewrited_solutions)
+    response = await selector.select(args, problem_statement, rewrited_solutions, worker_id=0)
     if len(args.results_dict["selector"]) == 0:
         reward_adjustment(args.results_dict["solver"], args.incorrect_reward_weight)
         reward_adjustment(args.results_dict["rewriter"], args.incorrect_reward_weight)
@@ -254,13 +260,18 @@ def reward_adjustment(samples, reward_weight):
     assert (
         len(args.results_dict["selector"]) == 1
     ), f"selector should only return one solution, but got {len(args.results_dict['selector'])}"
+    selector_sample = args.results_dict["selector"][0]
     if response is None:
-        args.results_dict["selector"][0].reward = 0
+        selector_sample.reward = 0
+        selector_sample.metadata["selector_parse_success"] = False
     else:
         selected_solution_idx = selector.extract_selected_solution_idx(response, rewrited_solutions)
         if selected_solution_idx is None:
-            args.results_dict["selector"][0].reward = 0
+            selector_sample.reward = 0
+            selector_sample.metadata["selector_parse_success"] = False
         else:
+            selector_sample.metadata["selector_parse_success"] = True
+            selector_sample.metadata["selector_choice"] = selected_solution_idx + 1
             selected_solution = rewrited_solutions[selected_solution_idx]
             for sample in args.results_dict["rewriter"]:
                 if sample.response_content is not None and selected_solution in sample.response_content:
diff --git a/examples/multi_agent/rollout_with_multi_agents.py b/examples/multi_agent/rollout_with_multi_agents.py
index d8a843c37..b3e4725cb 100644
--- a/examples/multi_agent/rollout_with_multi_agents.py
+++ b/examples/multi_agent/rollout_with_multi_agents.py
@@ -28,6 +28,21 @@ async def generate_with_multi_agents(args, sample: Sample, sampling_params, eval
     custom_multi_agent_func = load_function(args.custom_multi_agent_function_path)
     samples = await custom_multi_agent_func(args, sample)
 
+    # VIME compact rollouts return multiple training samples from one source
+    # sample. Newer VIME requires all siblings to share a rollout_id so loss
+    # reduction counts the source rollout once instead of over-counting agents.
+    compact_rollout_id = (
+        sample.rollout_id
+        if sample.rollout_id is not None
+        else (
+            sample.index
+            if sample.index is not None
+            else sample.group_index if sample.group_index is not None else id(sample)
+        )
+    )
+    for sibling in samples:
+        sibling.rollout_id = compact_rollout_id
+
     random.shuffle(samples)
 
     return samples
diff --git a/examples/multi_agent/run-qwen3-4B-multi-agent-npu.sh b/examples/multi_agent/run-qwen3-4B-multi-agent-npu.sh
new file mode 100755
index 000000000..9b8424be3
--- /dev/null
+++ b/examples/multi_agent/run-qwen3-4B-multi-agent-npu.sh
@@ -0,0 +1,150 @@
+#!/bin/bash
+
+# for rerun the task
+pkill -9 -f '[v]llm serve|VLL[M]::'
+pkill -9 -f VLLM
+sleep 3
+ray stop --force
+pkill -9 ray
+pkill -9 python
+sleep 3
+pkill -9 ray
+pkill -9 python
+pkill -9 redis
+
+set -ex
+
+export PYTHONUNBUFFERED=1
+
+export SLIME_SCRIPT_TRAIN_BACKEND=megatron
+export PYTHONPATH="/workspace/wky/Megatron-Bridge/src:/workspace/wky/Megatron-LM/:${PYTHONPATH:-}"
+export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1
+export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050
+export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050
+export HYDRA_FULL_ERROR=1
+export DISABLE_L2_CACHE=1
+export VLLM_ASCEND_ENABLE_NZ=0
+export VLLM_USE_AOT_COMPILE=0
+
+unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY
+
+VIME_ROOT="${VIME_ROOT:-/workspace/wky/vime-ascend}"
+SCRIPT_DIR="${VIME_ROOT}/scripts"
+WEIGHT_DIR="${WEIGHT_DIR:-/home/data/weights/Qwen3-4B}"
+DATA_FILE="${DATA_FILE:-/home/w00893744/dataset/dapo-math-17k.jsonl}"
+RUN_TS="${RUN_TS:-$(date +%Y%m%d_%H%M%S)}"
+LOG_FILE="${LOG_FILE:-/home/w00893744/train_qwen3_4b_multi_agent_vllm_${RUN_TS}.log}"
+MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}"
+
+cd "${VIME_ROOT}"
+source "${SCRIPT_DIR}/models/qwen3-4B.sh"
+
+CKPT_ARGS=(
+   --hf-checkpoint "${WEIGHT_DIR}"
+   --load "${WEIGHT_DIR}"
+   --megatron-to-hf-mode bridge
+)
+
+ROLLOUT_ARGS=(
+   --custom-generate-function-path examples.multi_agent.rollout_with_multi_agents.generate_with_multi_agents
+   --prompt-data "${DATA_FILE}"
+   --input-key prompt
+   --label-key label
+   --apply-chat-template
+   --rollout-shuffle
+   --rm-type math
+
+   --rollout-backend vllm
+   --vllm-weight-sync-mode native
+   --vllm-gpu-memory-utilization 0.6
+   --vllm-enable-sleep-mode
+   --vllm-max-model-len 4096
+   --vllm-enforce-eager
+
+   --num-rollout 200
+   --rollout-batch-size 32
+   --n-samples-per-prompt 8
+   --rollout-max-context-len 4096
+   --rollout-max-response-len 2048
+   --rollout-temperature 1.0
+
+   --global-batch-size 256
+   --balance-data
+)
+
+EVAL_ARGS=(
+)
+
+PERF_ARGS=(
+   --tensor-model-parallel-size 4
+   --pipeline-model-parallel-size 1
+   --context-parallel-size 1
+   --expert-model-parallel-size 1
+   --expert-tensor-parallel-size 1
+
+   --recompute-granularity full
+   --recompute-method uniform
+   --recompute-num-layers 1
+
+   --use-dynamic-batch-size
+   --max-tokens-per-gpu 8192
+   --micro-batch-size 1
+)
+
+GRPO_ARGS=(
+   --advantage-estimator grpo
+   --kl-loss-coef 0.0
+   --kl-loss-type low_var_kl
+   --kl-coef 0.00
+   --entropy-coef 0.0
+   --eps-clip 0.2
+   --eps-clip-high 0.28
+)
+
+OPTIMIZER_ARGS=(
+   --optimizer adam
+   --lr 1e-6
+   --lr-decay-style constant
+   --weight-decay 0.1
+   --adam-beta1 0.9
+   --adam-beta2 0.98
+)
+
+WANDB_ARGS=(
+)
+
+VLLM_ARGS=(
+   --rollout-num-gpus-per-engine 4
+)
+
+MISC_ARGS=(
+   --attention-dropout 0.0
+   --hidden-dropout 0.0
+   --accumulate-allreduce-grads-in-fp32
+   --attention-softmax-in-fp32
+   --attention-backend flash
+   --use-flash-attn
+   --train-memory-margin-bytes 2147483648
+)
+
+export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
+ray start --head --node-ip-address ${MASTER_ADDR} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265
+
+ray job submit --address="http://127.0.0.1:8265" \
+   -- python3 train.py \
+   --train-backend megatron \
+   --actor-num-nodes 1 \
+   --actor-num-gpus-per-node 4 \
+   --rollout-num-gpus 4 \
+   ${MODEL_ARGS[@]} \
+   ${CKPT_ARGS[@]} \
+   ${ROLLOUT_ARGS[@]} \
+   ${OPTIMIZER_ARGS[@]} \
+   ${GRPO_ARGS[@]} \
+   ${WANDB_ARGS[@]} \
+   ${PERF_ARGS[@]} \
+   ${EVAL_ARGS[@]} \
+   ${VLLM_ARGS[@]} \
+   ${MISC_ARGS[@]}
diff --git a/examples/retool/README.md b/examples/retool/README.md
new file mode 100644
index 000000000..cab318ded
--- /dev/null
+++ b/examples/retool/README.md
@@ -0,0 +1,88 @@
+# Retool: from SFT to RL
+This example _**(Retool)**_ demonstrates how to use the retool functionality for tool-enabled language model generation.
+
+## Overview
+The retool example provides:
+
+- Safe Python code execution in a sandbox environment
+- Tool registry for managing available tools
+- Integration with language model generation
+- Reward calculation for tool usage
+
+## Files
+- `generate_with_retool.py`: Main generation function with tool support
+- `tool_sandbox.py`: Tool execution and safety management
+- `sft_data_processing.py`: Process SFT dataset
+
+## Reward Design
+The RL reward function (`generate_with_retool.reward_func`) uses a **tool-aware reward shaping** strategy on top of the math accuracy reward:
+
+| Answer | Tool Used | Reward | Rationale |
+|--------|-----------|--------|-----------|
+| ✅ Correct | No | 1.0 | Pure reasoning — the ideal case |
+| ✅ Correct | Yes | 1.0 + min(0.2, turns × 0.05) | Bonus for effective tool use (capped at 0.2) |
+| ❌ Wrong | No | 0.0 | Neutral — model didn't attempt tools |
+| ❌ Wrong | Yes | min(0.1, turns × 0.02) | Small positive to encourage exploration |
+
+This design encourages the model to **explore tool calling** during early RL training without letting it reward-hack by preferring tool calls over correct answers. As training progresses and accuracy improves, the accuracy reward dominates and the tool bonus becomes a tiebreaker.
+
+## Usage
+### 1. Setup
+```
+git clone -b ascend https://github.com/vllm-project/vime.git
+cd vime
+docker build -f docker/Dockerfile.npu -t vime-ascend:latest .
+```
+```
+# Update the vime image
+export IMAGE=vime-ascend:latest
+
+docker run -d --name vime-npu -it --net=host --shm-size=1024g \
+    --privileged=true \
+    --cap-add=SYS_PTRACE \
+    --device=/dev/davinci_manager \
+    --device=/dev/hisi_hdc \
+    --device=/dev/devmm_svm \
+    -v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
+    -v /usr/local/dcmi:/usr/local/dcmi \
+    -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
+    -v /usr/local/sbin:/usr/local/sbin \
+    -v /home:/home \
+    -v /mnt:/mnt \
+    -v /tmp:/tmp \
+    -v /data:/data \
+    -v /usr/share/zoneinfo/Asia/Shanghai:/etc/localtime \
+    $IMAGE
+
+docker exec -it vime-npu bash
+```
+### 2. Download
+```
+# For SFT part, you can use later model to RL directly and skip SFT.
+hf download --repo-type dataset JoeYing/ReTool-SFT  --local-dir /path/to/ReTool-SFT
+hf download Qwen/Qwen3-4B-Instruct-2507 --local-dir /path/to/Qwen3-4B-Instruct-2507
+
+# For RL part
+hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /path/to/dapo-math-17k
+hf download --repo-type dataset zhuzilin/aime-2024  --local-dir /path/to/aime-2024
+# download our SFT model if you want to skip SFT
+hf download font-info/qwen3-4b-sft-SGLang-RL --local-dir /path/to/qwen3-4b-sft
+```
+### 3. Preprocessing
+```
+cd /root/vime
+# Replace the save path with /path/to/ReTool-SFT.parquet
+python examples/retool/sft_data_processing.py
+```
+### 4. SFT
+```
+cd /root/vime
+# Replace the model and data loading/saving paths
+python examples/retool/retool_qwen3_4b_sft.sh
+```
+### 5. RL
+```
+cd /root/vime
+# Replace the model and data loading/saving paths
+python examples/retool/retool_qwen3_4b_rl.sh
+```
diff --git a/examples/retool/README_zh.md b/examples/retool/README_zh.md
new file mode 100644
index 000000000..4d933e42d
--- /dev/null
+++ b/examples/retool/README_zh.md
@@ -0,0 +1,93 @@
+# Retool：从 SFT 到 RL
+本示例 **(Retool)** 演示了如何使用 retool 功能进行带工具调用的语言模型生成。
+
+## 概述
+Retool 示例提供了：
+
+- 沙盒环境中的安全 Python 代码执行
+- 工具注册表，用于管理可用工具
+- 与语言模型生成的集成
+- 工具使用的奖励计算
+
+## 文件说明
+- `generate_with_retool.py`：主生成函数，支持工具调用
+- `tool_sandbox.py`：工具执行与安全管理
+- `sft_data_processing.py`：处理 SFT 数据集
+
+## 奖励设计
+RL 奖励函数（`generate_with_retool.reward_func`）在数学正确性奖励的基础上采用了**工具感知的奖励塑形**策略：
+
+| 答案 | 是否使用工具 | 奖励值 | 说明 |
+|------|-------------|--------|------|
+| ✅ 正确 | 否 | 1.0 | 纯推理——理想情况 |
+| ✅ 正确 | 是 | 1.0 + min(0.2, 轮次 × 0.05) | 有效工具使用的额外奖励（上限 0.2） |
+| ❌ 错误 | 否 | 0.0 | 中性——模型未尝试使用工具 |
+| ❌ 错误 | 是 | min(0.1, 轮次 × 0.02) | 小额正奖励以鼓励探索 |
+
+该设计鼓励模型在 RL 训练早期**探索工具调用**，同时避免模型通过偏好工具调用而非正确答案来进行奖励投机。随着训练推进和准确率提升，正确性奖励将占主导，工具奖励变为辅助项。
+
+## 使用方法
+### 1. 环境搭建
+```bash
+git clone -b ascend https://github.com/vllm-project/vime.git
+cd vime
+docker build -f docker/Dockerfile.npu -t vime-ascend:latest .
+```
+
+```bash
+# 更新 vime 镜像
+export IMAGE=vime-ascend:latest
+
+docker run -d --name vime-npu -it --net=host --shm-size=1024g \
+    --privileged=true \
+    --cap-add=SYS_PTRACE \
+    --device=/dev/davinci_manager \
+    --device=/dev/hisi_hdc \
+    --device=/dev/devmm_svm \
+    -v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
+    -v /usr/local/dcmi:/usr/local/dcmi \
+    -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
+    -v /usr/local/sbin:/usr/local/sbin \
+    -v /home:/home \
+    -v /mnt:/mnt \
+    -v /tmp:/tmp \
+    -v /data:/data \
+    -v /usr/share/zoneinfo/Asia/Shanghai:/etc/localtime \
+    $IMAGE
+
+docker exec -it vime-npu bash
+```
+
+### 2. 下载
+```bash
+# SFT 部分，你也可以直接使用后续模型进行 RL 而跳过 SFT。
+hf download --repo-type dataset JoeYing/ReTool-SFT  --local-dir /path/to/ReTool-SFT
+hf download Qwen/Qwen3-4B-Instruct-2507 --local-dir /path/to/Qwen3-4B-Instruct-2507
+
+# RL 部分
+hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /path/to/dapo-math-17k
+hf download --repo-type dataset zhuzilin/aime-2024  --local-dir /path/to/aime-2024
+# 如果你想跳过 SFT，可下载我们的 SFT 模型
+hf download font-info/qwen3-4b-sft-SGLang-RL --local-dir /path/to/qwen3-4b-sft
+```
+
+### 3. 数据预处理
+```bash
+cd /root/vime
+# 将保存路径替换为 /path/to/ReTool-SFT.parquet
+python examples/retool/sft_data_processing.py
+```
+
+### 4. SFT
+```bash
+cd /root/vime
+# 替换模型和数据加载/保存路径
+python examples/retool/retool_qwen3_4b_sft.sh
+```
+
+### 5. RL
+```bash
+cd /root/vime
+# 替换模型和数据加载/保存路径
+python examples/retool/retool_qwen3_4b_rl.sh
+```
diff --git a/examples/retool/generate_with_retool.py b/examples/retool/generate_with_retool.py
new file mode 100644
index 000000000..e5b748324
--- /dev/null
+++ b/examples/retool/generate_with_retool.py
@@ -0,0 +1,571 @@
+# Adapted from https://github.com/volcengine/verl/blob/cb809d66e46dfd3342d008628891a14a054fa424/recipe/retool/retool.py
+# Adapted for vLLM /inference/v1/generate endpoint (disagg API)
+import logging
+import random
+import re
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+try:
+    from jinja2 import Template
+except ImportError as e:
+    raise ImportError("Jinja2 is required. Please install it with: pip install jinja2") from e
+
+from vime.rollout.vllm_rollout import GenerateState, _build_inference_sampling_params
+from vime.utils.http_utils import post
+from vime.utils.types import Sample
+
+# Import reward models
+try:
+    from vime.rollout.rm_hub.math_dapo_utils import compute_score as math_dapo_compute_score
+except ImportError as e:
+    raise ImportError("MathDapo is not installed") from e
+
+# Import tool sandbox functionality
+from tool_sandbox import SEMAPHORE, TOOL_CONFIGS, tool_registry
+
+# ── Sample-level verbose logging ──────────────────────────────────────────
+# Roughly 1/20 samples are logged so the output stays readable.
+_LOG_SAMPLE_PROB = 0.05
+_LOG_WIDTH = 300
+
+# ── Log probability collection ───────────────────────────────────────────
+# When True, collect log probabilities for TIS (Trajectory Importance Sampling).
+# When True, we CANNOT postprocess the decoded text because that would break
+# token/logp alignment. Instead, we rely on stop strings to ensure the engine
+# stops at tool/answer boundaries.
+RETURN_LOGPROB = True
+
+
+def _trunc(s: str, n: int = 300) -> str:
+    """Truncate *s* to at most *n* characters for display."""
+    if len(s) <= n:
+        return s
+    return s[:n] + f"…[+{len(s) - n}]"
+
+
+# Jinja2 template for tool-enabled conversations
+TOOL_TEMPLATE = """<|im_start|>system
+{%- if messages[0]['role'] == 'system' %}
+{{- messages[0]['content'] }}
+{%- else %}
+You are a helpful assistant.
+{%- endif %}
+{%- if tools %}
+# Tools
+
+You can write Python code to solve problems. Put your code between <code> and </code> tags. The code will be executed and you will see the output.
+When you have the final answer, write it as: Answer: \boxed{your_answer}
+{%- endif %}
+<|im_end|>
+{%- for message in messages %}
+{%- if message['role'] == 'user' %}
+<|im_start|>user
+{{- message['content'] }}<|im_end|>
+{%- elif message['role'] == 'assistant' %}
+<|im_start|>assistant
+{{- message['content'] }}<|im_end|>
+{%- endif %}
+{%- endfor %}
+<|im_start|>assistant
+"""
+
+
+def format_conversation_with_tools(
+    prompt: str, tools: list[dict[str, Any]] = None, system_prompt: str = None, messages: list[dict[str, Any]] = None
+) -> str:
+    """Format conversation using Jinja2 template with tool support"""
+    template = Template(TOOL_TEMPLATE)
+
+    # Prepare messages
+    messages_to_render = []
+
+    # Always add system message - use provided one or default
+    if system_prompt:
+        system_content = system_prompt
+    else:
+        system_content = "You are a helpful assistant that solves " "mathematical problems step by step."
+
+    messages_to_render.append({"role": "system", "content": system_content})
+
+    # Add user message if provided
+    if prompt:
+        messages_to_render.append({"role": "user", "content": prompt})
+
+    # Add assistant responses from previous turns if provided
+    if messages:
+        messages_to_render.extend(messages)
+
+    # Render template
+    formatted_text = template.render(messages=messages_to_render, tools=tools or [])
+
+    return formatted_text
+
+
+def postprocess_predictions(prediction: str):
+    """Extract action and content from prediction string"""
+    # Check for Answer: \boxed{...} format (only format we need for math_dapo)
+    # Use a more robust regex that handles nested braces
+    answer_pattern = r"Answer:\s*\\boxed\{((?:[^{}]|\{[^{}]*\})*)\}"
+    answer_match = re.search(answer_pattern, prediction, re.DOTALL)
+    if answer_match:
+        content = answer_match.group(1).strip()
+        return "answer", content
+
+    # Then check for <tool_call> tags (new format from Jinja2 template)
+    tool_call_pattern = r"<tool_call>\s*(\{.*?\})\s*</tool_call>"
+    tool_call_match = re.search(tool_call_pattern, prediction, re.DOTALL)
+    if tool_call_match:
+        try:
+            import json
+
+            # Clean up the JSON string by removing newlines and extra
+            # whitespace
+            json_str = tool_call_match.group(1)
+            # Replace newlines in string values with \n
+            json_str = json_str.replace("\n", "\\n")
+            tool_call_data = json.loads(json_str)
+            tool_name = tool_call_data.get("name")
+            arguments = tool_call_data.get("arguments", {})
+
+            if tool_name == "code_interpreter":
+                code = arguments.get("code", "")
+                if code.strip():
+                    return "code", code
+        except (json.JSONDecodeError, KeyError, AttributeError):
+            pass
+
+    # Then check for <code> tags
+    code_pattern = r"<code>(.*?)</code>"
+    code_match = re.search(code_pattern, prediction, re.DOTALL)
+    if code_match:
+        content = code_match.group(1).strip()
+        return "code", content
+
+    # Finally check for ```python code blocks (lowest priority)
+    python_code_pattern = r"```python\s*(.*?)\s*```"
+    python_code_match = re.search(python_code_pattern, prediction, re.DOTALL)
+    if python_code_match:
+        content = python_code_match.group(1).strip()
+        return "code", content
+
+    return None, ""
+
+
+def postprocess_responses(resp: str) -> str:
+    """Post-process response to ensure tag completeness.
+
+    IMPORTANT: Only used when RETURN_LOGPROB is False. When log prob collection
+    is enabled, we cannot postprocess the decoded text because that would break
+    token/logp alignment. Instead, we rely on stop strings to ensure the engine
+    stops at tool/answer boundaries.
+    """
+    # Handle <tool_call> tags (new format from Jinja2 template)
+    if "<tool_call>" in resp:
+        # Find the last occurrence of <tool_call>...</tool_call>
+        tool_call_pattern = r"<tool_call>\s*\{.*?\}\s*</tool_call>"
+        matches = list(re.finditer(tool_call_pattern, resp, re.DOTALL))
+        if matches:
+            last_match = matches[-1]
+            return resp[: last_match.end()]
+
+    # Handle <code> tags
+    if "</code>" in resp:
+        return resp.split("</code>")[0] + "</code>"
+
+    # Handle ```python code blocks
+    if "```python" in resp:
+        # Find the last occurrence of ```python...```
+        python_pattern = r"```python\s*.*?```"
+        matches = list(re.finditer(python_pattern, resp, re.DOTALL))
+        if matches:
+            last_match = matches[-1]
+            return resp[: last_match.end()]
+
+    # Handle Answer: \boxed{...} format (only format we need for math_dapo)
+    if "Answer:" in resp and "\\boxed{" in resp:
+        # Find the last occurrence of Answer: \boxed{...} with nested braces support
+        answer_pattern = r"Answer:\s*\\boxed\{((?:[^{}]|\{[^{}]*\})*)\}"
+        matches = list(re.finditer(answer_pattern, resp, re.DOTALL))
+        if matches:
+            last_match = matches[-1]
+            return resp[: last_match.end()]
+
+    return resp
+
+
+async def execute_predictions(prediction: str) -> str:
+    """Execute predictions and return results"""
+    action, content = postprocess_predictions(prediction)
+
+    if action == "code":
+        # Content is already the Python code (extracted by
+        # postprocess_predictions)
+        code = content.strip()
+        if code:
+            async with SEMAPHORE:
+                result = await tool_registry.execute_tool("code_interpreter", {"code": code})
+            next_obs = f"\n\n<interpreter>\n{result}\n</interpreter>\n\n"
+            done = False
+        else:
+            next_obs = "\n\n<interpreter>\nError: No Python code found\n</interpreter>\n\n"
+            done = False
+    elif action == "answer":
+        next_obs = ""
+        done = True
+    else:
+        next_obs = (
+            "\nMy previous action is invalid. "
+            "If I want to execute code, I should put the code between "
+            "<code> and </code>. "
+            "If I want to give the final answer, I should write "
+            "'Answer: ' followed by a boxed expression. Let me try again.\n"
+        )
+        done = False
+
+    return next_obs, done
+
+
+# Stop tags: make the inference engine STOP at the tool boundary.
+# "Answer:" was previously included, but it caused the engine to stop
+# immediately after "Answer:" before the model could generate "\boxed{...}",
+# resulting in invalid actions and reward=-1 every time.  The model now
+# completes its answer naturally; execute_predictions detects the
+# Answer: \boxed{...} pattern and sets done=True, ending the turn cleanly.
+_STOP_TAGS = ["</code>"]
+
+
+async def generate(args, sample: Sample, sampling_params) -> Sample:
+    """Custom generation function supporting tool calls (vLLM version)"""
+    assert not args.partial_rollout, "Partial rollout is not supported for " "this function at the moment."
+
+    state = GenerateState(args)
+    router_ip = args.vllm_router_ip
+    router_port = args.vllm_router_port
+    url = f"http://{router_ip}:{router_port}/inference/v1/generate"
+
+    # Set up the initial prompt with system prompt and tools (outside the loop)
+    tool_specs = tool_registry.get_tool_specs()
+
+    # Extract the raw user content from sample.prompt to avoid duplicate
+    # <|im_start|>user tags — the Jinja2 template adds them itself.
+    _raw_prompt = sample.prompt
+    if isinstance(_raw_prompt, list):
+        # List of message dicts, e.g. [{"role": "user", "content": "..."}]
+        _raw_prompt = "".join(m.get("content", "") for m in _raw_prompt)
+    # If the prompt string already contains chat-format tags, strip them so
+    # the template can re-add them consistently.
+    if isinstance(_raw_prompt, str) and "<|im_start|>" in _raw_prompt:
+        _raw_prompt = re.sub(r"<\|im_start\|>(?:system|user|assistant)\n?", "", _raw_prompt)
+        _raw_prompt = _raw_prompt.replace("<|im_end|>", "")
+        _raw_prompt = _raw_prompt.strip()
+
+    prompt = format_conversation_with_tools(prompt=_raw_prompt, tools=tool_specs)
+
+    prompt_tokens_ids = state.tokenizer(prompt, add_special_tokens=False)["input_ids"]
+    response = ""
+    response_token_ids = []
+    loss_masks = []
+    rollout_log_probs = [] if RETURN_LOGPROB else None
+    tool_call_count = 0  # Track actual tool call rounds
+    consecutive_memory_errors = 0  # Track consecutive "Memory usage too high" errors
+    obs_truncated = False  # Flag: obs caused total length to exceed max_context_length
+    # Calculate max context length once at the beginning
+    max_context_length = len(prompt_tokens_ids) + args.rollout_max_response_len
+    logger.debug("max_context_length is set to %d", max_context_length)
+
+    # Randomly select a small fraction of samples for detailed turn-by-turn logging.
+    verbose = random.random() < _LOG_SAMPLE_PROB
+    if verbose:
+        _sep = "═" * _LOG_WIDTH
+        _prompt_display = (
+            "".join(m.get("content", "") for m in sample.prompt) if isinstance(sample.prompt, list) else sample.prompt
+        )
+        logger.debug(
+            "\n%s\n[ReTool LOG] prompt (%d tokens): %s\n%s",
+            _sep,
+            len(prompt_tokens_ids),
+            _trunc(_prompt_display, 200),
+            _sep,
+        )
+
+    # Add stop tags to sampling_params so the engine stops at tool/answer boundaries.
+    # This is critical for keeping token/logp alignment when RETURN_LOGPROB is enabled.
+    _existing_stop = sampling_params.get("stop") or []
+    if isinstance(_existing_stop, str):
+        _existing_stop = [_existing_stop]
+    sampling_params = {**sampling_params, "stop": list(dict.fromkeys([*_existing_stop, *_STOP_TAGS]))}
+
+    # Build vLLM-style sampling params (maps max_new_tokens -> max_tokens, adds logprobs, etc.)
+    inference_sampling_params = _build_inference_sampling_params(sampling_params)
+
+    for turn in range(TOOL_CONFIGS["max_turns"]):
+        # vLLM /inference/v1/generate requires token_ids instead of text.
+        # For multi-turn, re-tokenize the full context each time.
+        full_text = prompt + response
+        current_token_ids = state.tokenizer(full_text, add_special_tokens=False)["input_ids"]
+
+        # Check if total length exceeds max context length
+        total_length = len(current_token_ids)
+        if total_length >= max_context_length:
+            sample.status = Sample.Status.TRUNCATED
+            break
+
+        # Dynamically calculate remaining token budget for this turn
+        remaining_tokens = max_context_length - total_length
+
+        # Update max_tokens for this turn to respect the remaining budget
+        current_inference_params = dict(inference_sampling_params)
+        current_inference_params["max_tokens"] = min(
+            inference_sampling_params["max_tokens"],
+            remaining_tokens,
+        )
+
+        # Check if we have budget for more tokens
+        if current_inference_params["max_tokens"] <= 0:
+            sample.status = Sample.Status.TRUNCATED
+            break
+
+        payload = {
+            "token_ids": current_token_ids,
+            "sampling_params": current_inference_params,
+        }
+        if hasattr(args, "hf_checkpoint"):
+            payload["model"] = args.hf_checkpoint
+
+        # Log payload to wandb for debugging
+        try:
+            import wandb
+
+            if wandb.run is not None:
+                # Count available tools (from tool_specs)
+                available_tools = len(tool_specs)
+                # Count tools used in the current response
+                tools_used = response.count("<interpreter>")
+
+                wandb.log(
+                    {
+                        "debug/payload_length": len(prompt + response),
+                        "debug/available_tools": available_tools,
+                        "debug/tools_used": tools_used,
+                        "debug/turn": turn,
+                    }
+                )
+        except ImportError:
+            pass  # wandb not available
+
+        output = await post(url, payload)
+
+        # Parse vLLM GenerateResponse: {"choices": [{"token_ids": [...], "logprobs": ..., "finish_reason": "stop"}]}
+        choice = output["choices"][0]
+        finish_reason = choice.get("finish_reason", "stop")
+
+        # Handle abort
+        if finish_reason == "abort":
+            sample.status = Sample.Status.ABORTED
+            return sample
+
+        # Extract token IDs from vLLM response
+        cur_response_token_ids = choice.get("token_ids") or []
+
+        # Decode text from token_ids
+        skip_sp = current_inference_params.get("skip_special_tokens")
+        skip_decode = True if skip_sp is None else bool(skip_sp)
+        cur_response = (
+            state.tokenizer.decode(cur_response_token_ids, skip_special_tokens=skip_decode)
+            if cur_response_token_ids
+            else ""
+        )
+
+        # Extract log probs if enabled
+        if RETURN_LOGPROB:
+            cur_response_log_probs: list[float] = []
+            lp = choice.get("logprobs")
+            if isinstance(lp, dict):
+                content_items = lp.get("content") or []
+                cur_response_log_probs = [
+                    float(item.get("logprob", 0.0)) if isinstance(item, dict) else 0.0 for item in content_items
+                ]
+            if not cur_response_log_probs:
+                cur_response_log_probs = [0.0] * len(cur_response_token_ids)
+        else:
+            # When not collecting log probs, we can safely postprocess the response
+            cur_response = postprocess_responses(cur_response)
+            # Re-tokenize after postprocessing
+            cur_response_token_ids = state.tokenizer(cur_response, add_special_tokens=False)["input_ids"]
+            cur_response_log_probs = None
+
+        response += cur_response
+        response_token_ids += cur_response_token_ids
+        loss_masks += [1] * len(cur_response_token_ids)
+
+        # Add log probs if enabled
+        if RETURN_LOGPROB:
+            rollout_log_probs += cur_response_log_probs
+
+        # verbose: show what the model generated this turn
+        if verbose:
+            n_tok = len(cur_response_token_ids)
+            logger.debug(
+                "\n%s\n[Turn %d] model output (%d tok, finish=%s):\n  %s",
+                "─" * _LOG_WIDTH,
+                turn + 1,
+                n_tok,
+                finish_reason,
+                _trunc(cur_response).replace("\n", "\n  "),
+            )
+
+        # Check length limit
+        if finish_reason == "length":
+            if verbose:
+                logger.debug("[Turn %d] → length limit reached, stopping.", turn + 1)
+            break
+
+        next_obs, done = await execute_predictions(cur_response)
+
+        # Track consecutive memory errors to break the vicious retry cycle.
+        # When sandbox keeps returning "Memory usage too high", the model
+        # generates another code attempt which also fails, consuming more
+        # tokens and memory. Force done=True after N consecutive failures.
+        if "Memory usage too high" in next_obs:
+            consecutive_memory_errors += 1
+            if consecutive_memory_errors >= TOOL_CONFIGS.get("max_consecutive_memory_errors", 3):
+                done = True
+        else:
+            consecutive_memory_errors = 0
+
+        # verbose: show action and observation
+        if verbose:
+            if done:
+                logger.debug("[Turn %d] → answer detected (DONE)", turn + 1)
+            elif "<interpreter>" in next_obs:
+                obs_display = "  " + _trunc(next_obs, 300).replace("\n", "\n  ")
+                logger.debug("[Turn %d] → code executed, observation:\n%s", turn + 1, obs_display)
+            else:
+                logger.debug("[Turn %d] → invalid action (no recognized code or answer)", turn + 1)
+
+        if done:
+            break
+
+        # Count tool calls (when we get interpreter output, it means a tool
+        # was called)
+        if "<interpreter>" in next_obs:
+            tool_call_count += 1
+
+        assert next_obs != "", "Next observation should not be empty."
+        obs_tokens_ids = state.tokenizer(next_obs, add_special_tokens=False)["input_ids"]
+        response += next_obs
+        response_token_ids += obs_tokens_ids
+        loss_masks += [0] * len(obs_tokens_ids)
+
+        # Add dummy log probs for observation tokens if enabled (they won't be used due to loss_mask=0)
+        if RETURN_LOGPROB:
+            rollout_log_probs += [0.0] * len(obs_tokens_ids)
+
+            # Verify alignment when collecting log probs
+            assert len(response_token_ids) == len(
+                rollout_log_probs
+            ), f"Token/logp length mismatch at turn {turn}: {len(response_token_ids)} tokens vs {len(rollout_log_probs)} logps"
+
+        # Truncate if obs pushed total response tokens beyond max_context_length
+        max_response_tokens = max_context_length - len(prompt_tokens_ids)
+        if len(response_token_ids) > max_response_tokens:
+            response_token_ids = response_token_ids[:max_response_tokens]
+            loss_masks = loss_masks[:max_response_tokens]
+            if RETURN_LOGPROB:
+                rollout_log_probs = rollout_log_probs[:max_response_tokens]
+            obs_truncated = True
+            break
+
+        if tool_call_count >= TOOL_CONFIGS["max_tool_calls"]:
+            break
+
+    if verbose:
+        logger.debug(
+            "\n%s\n[ReTool LOG] finished | tool_calls=%d | " "response_tokens=%d | finish=%s\n%s",
+            "═" * _LOG_WIDTH,
+            tool_call_count,
+            len(response_token_ids),
+            finish_reason,
+            "═" * _LOG_WIDTH,
+        )
+
+    # Set sample attributes
+    sample.tokens = prompt_tokens_ids + response_token_ids
+    sample.response_length = len(response_token_ids)
+    sample.response = response
+    sample.loss_mask = loss_masks
+    sample.prompt = prompt
+
+    # Store log probs if enabled
+    if RETURN_LOGPROB:
+        sample.rollout_log_probs = rollout_log_probs if rollout_log_probs else None
+
+    # Store payload information for wandb logging
+    sample.payload_text = prompt + response
+    sample.payload_has_system = "<|im_start|>system" in prompt + response
+    sample.payload_has_tools = "# Tools" in prompt + response
+
+    # Store tool call count for reward calculation
+    sample.tool_call_count = tool_call_count
+
+    # Set status
+    # vLLM finish_reason is a string: "stop", "length", or "abort"
+    match finish_reason:
+        case "length":
+            sample.status = Sample.Status.TRUNCATED
+        case "abort":
+            sample.status = Sample.Status.ABORTED
+        case "stop":
+            sample.status = Sample.Status.COMPLETED
+
+    if obs_truncated:
+        sample.status = Sample.Status.TRUNCATED
+
+    return sample
+
+
+async def reward_func(args, sample, **kwargs):
+    """Tool call reward function using math_dapo as primary reward model"""
+    if not isinstance(sample, Sample):
+        raise TypeError("Sample must be an instance of Sample class.")
+
+    # Truncated samples: neutral reward — not right, not wrong, just incomplete.
+    if sample.status == Sample.Status.TRUNCATED:
+        return {"score": 0.0, "acc": False, "pred": ""}
+
+    # Build complete solution string.
+    # sample.prompt may be a list of message dicts; flatten to plain text.
+    if isinstance(sample.prompt, list):
+        prompt_str = "".join(m.get("content", "") for m in sample.prompt)
+    else:
+        prompt_str = sample.prompt
+    solution_str = prompt_str + sample.response
+
+    # Get ground truth answer - label is a string, not a dict
+    ground_truth = sample.label if sample.label is not None else ""
+
+    # Get tool call count as num_turns
+    num_turns = getattr(sample, "tool_call_count", 0)
+
+    # use \\boxed{...} answer
+    result = math_dapo_compute_score(solution_str, ground_truth, strict_box_verify=True)
+
+    # Reward shaping:
+    #   Correct + no tools    → 1.0  (pure reasoning, best)
+    #   Correct + tools       → 1.0 + bonus (tools helped, still good)
+    #   Wrong + no tools      → 0.0  (neutral — model didn't even try tools)
+    #   Wrong + tools         → small positive (encourages exploration,
+    #                           but capped low to avoid reward hacking:
+    #                           the model should not prefer tool-calling
+    #                           over getting the right answer)
+    if result["score"] > 0:
+        result["score"] = 1.0 + min(0.2, num_turns * 0.05)
+    else:
+        result["score"] = min(0.1, num_turns * 0.02)
+
+    if result["pred"] is None:
+        result["pred"] = ""
+
+    return result
diff --git a/examples/retool/requirements.txt b/examples/retool/requirements.txt
new file mode 100644
index 000000000..dd93aaab7
--- /dev/null
+++ b/examples/retool/requirements.txt
@@ -0,0 +1,3 @@
+jinja2>=3.0.0
+psutil>=5.8.0
+pytest>=7.0.0
diff --git a/examples/retool/retool_qwen3_4b_rl.sh b/examples/retool/retool_qwen3_4b_rl.sh
new file mode 100644
index 000000000..42d4eb2da
--- /dev/null
+++ b/examples/retool/retool_qwen3_4b_rl.sh
@@ -0,0 +1,206 @@
+#!/bin/bash
+set -ex
+ulimit -u 65535
+
+# cleanup
+pkill -9 -f "vllm serve" 2>/dev/null || true
+sleep 2
+npu-smi info 2>/dev/null | grep rayWorker | awk '{print $4}' | xargs -r kill -9 2>/dev/null || true
+sleep 3
+
+# Ray isolation: independent temp-dir, ports, and cleanup
+export RAY_TMPDIR=/tmp/ray_vime_npu_retool
+export RAY_PORT=6379
+export RAY_DASHBOARD_PORT=8265
+export RAY_AGENT_PORT=52378
+unset RAY_ADDRESS RAY_REDIS_ADDRESS
+
+ray stop --force 2>/dev/null || true
+rm -rf "${RAY_TMPDIR}"
+sleep 2
+
+project_name="vime"
+exp_name="qwen3-4b-retool-rl"
+RAY_DATA_HOME=${RAY_DATA_HOME:-"/root/logs"}
+start_time=$(date +"%Y%m%d_%H%M%S")
+LOG_DIR=${LOG_DIR:-"${RAY_DATA_HOME}/${project_name}/${exp_name}"}
+mkdir -p "${LOG_DIR}"
+LOG_FILE="${LOG_DIR}/${start_time}.log"
+
+echo "Experiment Log will be saved to: ${LOG_FILE}"
+VIME_DIR="/root/vime"
+
+# NPU environment
+source /usr/local/Ascend/driver/bin/setenv.bash
+source /usr/local/Ascend/ascend-toolkit/set_env.sh
+source /usr/local/Ascend/nnal/atb/set_env.sh
+export PYTHONPATH="${VIME_DIR}:${VIME_DIR}/examples/retool:/root/Megatron-LM:/root/vllm:/root/vllm-ascend:/root/Megatron-Bridge:/root/mbridge:/root/MegatronAdaptor:/root/TransformerEngineNPU:/usr/local/Ascend/ascend-toolkit/latest/python/site-packages:/usr/local/Ascend/ascend-toolkit/latest/tools/ms_fmk_transplt/torch_npu_bridge:${PYTHONPATH}"
+export PYTHONUNBUFFERED=1
+export PYTORCH_NPU_ALLOC_CONF=expandable_segments:False
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15
+export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050
+export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050
+export HCCL_CONNECT_TIMEOUT=7200
+export HCCL_DETERMINISTIC=true
+export VLLM_ASCEND_ENABLE_NZ=0
+export ASCEND_COREDUMP_SIGNAL=None
+export ATB_MATMUL_SHUFFLE_K_ENABLE=0
+export ATB_LLM_LCOC_ENABLE=0
+export TASK_QUEUE_ENABLE=1
+export RAY_DISABLE_SIGINT_OVERRIDE=1
+export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1
+export LD_LIBRARY_PATH=/usr/local/Ascend/driver/lib64:/usr/local/Ascend/ascend-toolkit/latest/lib64:/usr/local/Ascend/nnal/atb/latest/atb/cxx_abi_1/lib:/usr/local/Ascend/cann/lib64:${LD_LIBRARY_PATH}
+export VLLM_DISABLE_COMPILE_CACHE=1
+export TRANSFORMERS_VERBOSITY=error
+export RUST_LOG=vllm_router_rs=warn
+
+NUM_NPUS=16
+source "${VIME_DIR}/scripts/models/qwen3-4B-Instruct-2507.sh"
+
+CKPT_ARGS=(
+   --hf-checkpoint /path/to/Qwen3-4B_sft_vime_hf
+   --ref-load /path/to/Qwen3-4B_sft_vime_hf
+   --load /path/to/Qwen3-4B_vime_npu/
+   --save /path/to/Qwen3-4B_vime_npu/
+   --save-interval 20
+   --no-load-optim
+   --megatron-to-hf-mode bridge
+)
+
+ROLLOUT_ARGS=(
+   --prompt-data /path/to/dapo-math-17k/dapo-math-17k.jsonl
+   --input-key prompt
+   --label-key label
+   --apply-chat-template
+   --rollout-shuffle
+   --reward-key score
+   --num-rollout 200
+   --rollout-batch-size 32
+   --n-samples-per-prompt 8
+   --rollout-max-response-len 8192
+   --rollout-temperature 1
+   --global-batch-size 256
+   --balance-data
+)
+
+EVAL_ARGS=(
+   --eval-interval 50
+   --eval-prompt-data aime /path/to/aime-2024/aime-2024.jsonl
+   --n-samples-per-eval-prompt 16
+   --eval-max-response-len 16384
+   --eval-top-p 1
+)
+
+PERF_ARGS=(
+   --tensor-model-parallel-size 4
+   --pipeline-model-parallel-size 1
+   --context-parallel-size 1
+   --expert-model-parallel-size 1
+   --expert-tensor-parallel-size 1
+
+   --recompute-granularity full
+   --recompute-method uniform
+   --recompute-num-layers 1
+
+   --micro-batch-size 1
+   --use-dynamic-batch-size
+   --max-tokens-per-gpu 9216
+)
+
+GRPO_ARGS=(
+   --advantage-estimator grpo
+   --use-kl-loss
+   --kl-loss-coef 0.00
+   --kl-loss-type low_var_kl
+   --entropy-coef 0.00
+   --eps-clip 0.2
+   --eps-clip-high 0.28
+)
+
+OPTIMIZER_ARGS=(
+   --optimizer adam
+   --lr 1e-6
+   --lr-decay-style constant
+   --weight-decay 0.1
+   --adam-beta1 0.9
+   --adam-beta2 0.98
+   --optimizer-cpu-offload
+   --overlap-cpu-optimizer-d2h-h2d
+   --use-precision-aware-optimizer
+)
+
+VLLM_ARGS=(
+   --rollout-num-gpus-per-engine 4
+   --vllm-gpu-memory-utilization 0.7
+   --vllm-enable-sleep-mode
+   --vllm-weight-sync-mode native
+   --vllm-max-model-len 16384
+)
+
+MISC_ARGS=(
+   --attention-dropout 0.0
+   --hidden-dropout 0.0
+   --accumulate-allreduce-grads-in-fp32
+   --attention-softmax-in-fp32
+   --attention-backend flash
+   --use-flash-attn
+)
+
+CUSTOM_ARGS=(
+   --custom-generate-function-path generate_with_retool.generate
+   --custom-rm-path generate_with_retool.reward_func
+)
+
+# launch the master node of ray in container
+unset https_proxy http_proxy proxy
+ray start --head \
+	--temp-dir="${RAY_TMPDIR}" \
+	--port="${RAY_PORT}" \
+	--dashboard-port="${RAY_DASHBOARD_PORT}" \
+	--dashboard-agent-listen-port="${RAY_AGENT_PORT}" \
+	--node-ip-address 127.0.0.1 \
+	--num-gpus 0 \
+	--resources "{\"NPU\": $NUM_NPUS}" \
+	--disable-usage-stats \
+	--dashboard-host=0.0.0.0
+
+# Build the runtime environment JSON with proper variable substitution
+RUNTIME_ENV_JSON=$(cat << 'EOF'
+{
+  "env_vars": {
+    "PYTHONPATH": "${VIME_DIR}:${VIME_DIR}/examples/retool:/root/Megatron-LM:/root/vllm:/root/vllm-ascend:/root/Megatron-Bridge:/root/mbridge:/root/MegatronAdaptor:/root/TransformerEngineNPU:/usr/local/Ascend/ascend-toolkit/latest/python/site-packages:/usr/local/Ascend/ascend-toolkit/latest/tools/ms_fmk_transplt/torch_npu_bridge",
+    "CUDA_DEVICE_MAX_CONNECTIONS": "1",
+    "HCCL_HOST_SOCKET_PORT_RANGE": "60000-60050",
+    "HCCL_NPU_SOCKET_PORT_RANGE": "61000-61050",
+    "HCCL_CONNECT_TIMEOUT": "7200",
+    "PYTORCH_NPU_ALLOC_CONF": "expandable_segments:False",
+    "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES": "1",
+    "VLLM_DISABLE_COMPILE_CACHE": "1",
+    "TRANSFORMERS_VERBOSITY": "error",
+    "RUST_LOG": "vllm_router_rs=warn",
+    "LD_LIBRARY_PATH": "/usr/local/Ascend/driver/lib64:/usr/local/Ascend/driver/lib64/driver:/usr/local/Ascend/driver/lib64/common:/usr/local/Ascend/ascend-toolkit/latest/lib64:/usr/local/Ascend/ascend-toolkit/latest/compiler/lib64/plugin/opskernel:/usr/local/Ascend/ascend-toolkit/latest/compiler/lib64/plugin/nnengine:/usr/local/Ascend/ascend-toolkit/latest/opp/built-in/op_impl/ai_core/tbe/op_tiling/lib/:/usr/local/Ascend/nnal/atb/latest/atb/cxx_abi_1/lib:/usr/local/Ascend/cann/lib64:/usr/local/Ascend/cann/aarch64-linux/devlib"
+  }
+}
+EOF
+)
+
+ray job submit --address="http://127.0.0.1:${RAY_DASHBOARD_PORT}" \
+   --runtime-env-json="${RUNTIME_ENV_JSON}" \
+   --working-dir="${VIME_DIR}" \
+   -- python3 -u train.py \
+   --train-backend megatron \
+   --actor-num-nodes 1 \
+   --actor-num-gpus-per-node 8 \
+   --rollout-num-gpus 8 \
+   ${MODEL_ARGS[@]} \
+   ${CKPT_ARGS[@]} \
+   ${ROLLOUT_ARGS[@]} \
+   ${OPTIMIZER_ARGS[@]} \
+   ${GRPO_ARGS[@]} \
+   ${PERF_ARGS[@]} \
+   ${EVAL_ARGS[@]} \
+   ${VLLM_ARGS[@]} \
+   ${MISC_ARGS[@]} \
+   ${CUSTOM_ARGS[@]} \
+   2>&1 | tee "${LOG_FILE}"
diff --git a/examples/retool/retool_qwen3_4b_sft.sh b/examples/retool/retool_qwen3_4b_sft.sh
new file mode 100644
index 000000000..cf067ab80
--- /dev/null
+++ b/examples/retool/retool_qwen3_4b_sft.sh
@@ -0,0 +1,170 @@
+#!/bin/bash
+set -ex
+ulimit -u 65535
+
+# cleanup
+pkill -9 -f "vllm serve" 2>/dev/null || true
+sleep 2
+npu-smi info 2>/dev/null | grep rayWorker | awk '{print $4}' | xargs -r kill -9 2>/dev/null || true
+sleep 3
+
+# Ray isolation: independent temp-dir, ports, and cleanup
+export RAY_TMPDIR=/tmp/ray_vime_npu_retool
+export RAY_PORT=6379
+export RAY_DASHBOARD_PORT=8265
+export RAY_AGENT_PORT=52378
+unset RAY_ADDRESS RAY_REDIS_ADDRESS
+
+ray stop --force 2>/dev/null || true
+rm -rf "${RAY_TMPDIR}"
+sleep 2
+
+project_name="vime"
+exp_name="qwen3-4b-retool-sft"
+RAY_DATA_HOME=${RAY_DATA_HOME:-"/root/logs"}
+start_time=$(date +"%Y%m%d_%H%M%S")
+LOG_DIR=${LOG_DIR:-"${RAY_DATA_HOME}/${project_name}/${exp_name}"}
+mkdir -p "${LOG_DIR}"
+LOG_FILE="${LOG_DIR}/${start_time}.log"
+
+echo "Experiment Log will be saved to: ${LOG_FILE}"
+VIME_DIR="/root/vime"
+
+# NPU environment
+source /usr/local/Ascend/driver/bin/setenv.bash
+source /usr/local/Ascend/ascend-toolkit/set_env.sh
+source /usr/local/Ascend/nnal/atb/set_env.sh
+export PYTHONPATH="${VIME_DIR}:${VIME_DIR}/examples/retool:/root/Megatron-LM:/root/vllm:/root/vllm-ascend:/root/Megatron-Bridge:/root/mbridge:/root/MegatronAdaptor:/root/TransformerEngineNPU:/usr/local/Ascend/ascend-toolkit/latest/python/site-packages:/usr/local/Ascend/ascend-toolkit/latest/tools/ms_fmk_transplt/torch_npu_bridge:${PYTHONPATH}"
+export PYTHONUNBUFFERED=1
+export PYTORCH_NPU_ALLOC_CONF=expandable_segments:False
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15
+export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050
+export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050
+export HCCL_CONNECT_TIMEOUT=7200
+export HCCL_DETERMINISTIC=true
+export VLLM_ASCEND_ENABLE_NZ=0
+export ASCEND_COREDUMP_SIGNAL=None
+export ATB_MATMUL_SHUFFLE_K_ENABLE=0
+export ATB_LLM_LCOC_ENABLE=0
+export TASK_QUEUE_ENABLE=1
+export RAY_DISABLE_SIGINT_OVERRIDE=1
+export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1
+export LD_LIBRARY_PATH=/usr/local/Ascend/driver/lib64:/usr/local/Ascend/ascend-toolkit/latest/lib64:/usr/local/Ascend/nnal/atb/latest/atb/cxx_abi_1/lib:/usr/local/Ascend/cann/lib64:${LD_LIBRARY_PATH}
+export VLLM_DISABLE_COMPILE_CACHE=1
+export RUST_LOG=vllm_router_rs=warn
+export TRANSFORMERS_VERBOSITY=error
+
+NUM_NPUS=16
+source "${VIME_DIR}/scripts/models/qwen3-4B-Instruct-2507.sh"
+
+CKPT_ARGS=(
+   --hf-checkpoint /path/to/Qwen3-4B-Instruct-2507
+   --ref-load /path/to/Qwen3-4B-Instruct-2507
+   --save /path/to/Qwen3-4B_sft_vime/
+   --save-interval 1000
+   --save-hf /path/to/Qwen3-4B_sft_vime_hf/
+   --no-load-optim
+   --megatron-to-hf-mode bridge
+)
+
+SFT_ARGS=(
+   --rollout-function-path vime.rollout.sft_rollout.generate_rollout
+   --prompt-data /path/to/ReTool-SFT/ReTool-SFT.parquet
+   --input-key messages
+   --rollout-shuffle
+   --num-epoch 3
+   --rollout-batch-size 128
+   --global-batch-size 128
+   --loss-type sft_loss
+   --calculate-per-token-loss
+   --disable-compute-advantages-and-returns
+   --debug-train-only
+)
+
+PERF_ARGS=(
+   --tensor-model-parallel-size 4
+   --pipeline-model-parallel-size 1
+   --context-parallel-size 1
+   --expert-model-parallel-size 1
+   --expert-tensor-parallel-size 1
+
+   --recompute-granularity full
+   --recompute-method uniform
+   --recompute-num-layers 1
+
+   --micro-batch-size 1
+   --use-dynamic-batch-size
+   --max-tokens-per-gpu 9216
+)
+
+OPTIMIZER_ARGS=(
+   --optimizer adam
+   --lr 1e-5
+   --lr-decay-style cosine
+   --min-lr 1e-6
+   --lr-warmup-fraction 0.1
+   --weight-decay 0.1
+   --adam-beta1 0.9
+   --adam-beta2 0.95
+   --optimizer-cpu-offload
+   --overlap-cpu-optimizer-d2h-h2d
+   --use-precision-aware-optimizer
+)
+
+MISC_ARGS=(
+   --attention-dropout 0.0
+   --hidden-dropout 0.0
+   --accumulate-allreduce-grads-in-fp32
+   --attention-softmax-in-fp32
+   --attention-backend flash
+   --use-flash-attn
+)
+
+# launch the master node of ray in container
+unset https_proxy http_proxy proxy
+ray start --head \
+	--temp-dir="${RAY_TMPDIR}" \
+	--port="${RAY_PORT}" \
+	--dashboard-port="${RAY_DASHBOARD_PORT}" \
+	--dashboard-agent-listen-port="${RAY_AGENT_PORT}" \
+	--node-ip-address 127.0.0.1 \
+	--num-gpus 0 \
+	--resources "{\"NPU\": $NUM_NPUS}" \
+	--disable-usage-stats \
+	--dashboard-host=0.0.0.0
+
+# Build the runtime environment JSON with proper variable substitution
+RUNTIME_ENV_JSON=$(cat << 'EOF'
+{
+  "env_vars": {
+    "PYTHONPATH": "${VIME_DIR}:${VIME_DIR}/examples/retool:/root/Megatron-LM:/root/vllm:/root/vllm-ascend:/root/Megatron-Bridge:/root/mbridge:/root/MegatronAdaptor:/root/TransformerEngineNPU:/usr/local/Ascend/ascend-toolkit/latest/python/site-packages:/usr/local/Ascend/ascend-toolkit/latest/tools/ms_fmk_transplt/torch_npu_bridge",
+    "CUDA_DEVICE_MAX_CONNECTIONS": "1",
+    "HCCL_HOST_SOCKET_PORT_RANGE": "60000-60050",
+    "HCCL_NPU_SOCKET_PORT_RANGE": "61000-61050",
+    "HCCL_CONNECT_TIMEOUT": "7200",
+    "PYTORCH_NPU_ALLOC_CONF": "expandable_segments:False",
+    "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES": "1",
+    "VLLM_DISABLE_COMPILE_CACHE": "1",
+    "TRANSFORMERS_VERBOSITY": "error",
+    "RUST_LOG": "vllm_router_rs=warn",
+    "LD_LIBRARY_PATH": "/usr/local/Ascend/driver/lib64:/usr/local/Ascend/driver/lib64/driver:/usr/local/Ascend/driver/lib64/common:/usr/local/Ascend/ascend-toolkit/latest/lib64:/usr/local/Ascend/ascend-toolkit/latest/compiler/lib64/plugin/opskernel:/usr/local/Ascend/ascend-toolkit/latest/compiler/lib64/plugin/nnengine:/usr/local/Ascend/ascend-toolkit/latest/opp/built-in/op_impl/ai_core/tbe/op_tiling/lib/:/usr/local/Ascend/nnal/atb/latest/atb/cxx_abi_1/lib:/usr/local/Ascend/cann/lib64:/usr/local/Ascend/cann/aarch64-linux/devlib"
+  }
+}
+EOF
+)
+
+ray job submit --address="http://127.0.0.1:${RAY_DASHBOARD_PORT}" \
+   --runtime-env-json="${RUNTIME_ENV_JSON}" \
+   --working-dir="${VIME_DIR}" \
+   -- python3 -u train.py \
+   --train-backend megatron \
+   --actor-num-nodes 1 \
+   --actor-num-gpus-per-node 16 \
+   ${MODEL_ARGS[@]} \
+   ${CKPT_ARGS[@]} \
+   ${SFT_ARGS[@]} \
+   ${OPTIMIZER_ARGS[@]} \
+   ${PERF_ARGS[@]} \
+   ${MISC_ARGS[@]} \
+   2>&1 | tee "${LOG_FILE}"
diff --git a/examples/retool/rl_data_preprocess.py b/examples/retool/rl_data_preprocess.py
new file mode 100644
index 000000000..da4f75631
--- /dev/null
+++ b/examples/retool/rl_data_preprocess.py
@@ -0,0 +1,21 @@
+from datasets import load_dataset
+
+# Load the original dataset
+ds = load_dataset("BytedTsinghua-SIA/DAPO-Math-17k", split="train")
+
+
+# Map to extract the ground_truth from the reward_model dict and create a new 'label' field
+def transform(example):
+    return {
+        "prompt": example["prompt"][0]["content"] if example["prompt"] else None,
+        "label": example["reward_model"]["ground_truth"],
+    }
+
+
+ds2 = ds.map(transform, remove_columns=ds.column_names)
+
+# Optionally, verify the first few entries
+print(ds2[0])
+
+# save to jsonl
+ds2.to_json("/path/to/dapo-math-17k/dapo-math-17k.jsonl", orient="records", lines=True)
diff --git a/examples/retool/sft_data_processing.py b/examples/retool/sft_data_processing.py
new file mode 100644
index 000000000..38bbccd29
--- /dev/null
+++ b/examples/retool/sft_data_processing.py
@@ -0,0 +1,31 @@
+from datasets import load_dataset
+
+ds = load_dataset("/path/to/ReTool-SFT")["train"]
+
+
+def convert(sample):
+    conversations = sample["messages"]
+
+    def convert_role(role):
+        if role == "user":
+            return "user"
+        elif role == "assistant":
+            return "assistant"
+        elif role == "system":
+            return "system"
+        else:
+            raise ValueError(f"Unknown role: {role}")
+
+    messages = [
+        {
+            "role": convert_role(turn["role"]),
+            "content": turn["content"],
+        }
+        for turn in conversations
+    ]
+
+    return {"messages": messages}
+
+
+ds = ds.map(convert)
+ds.to_parquet("/path/to/ReTool-SFT/ReTool-SFT.parquet")
diff --git a/examples/retool/tool_sandbox.py b/examples/retool/tool_sandbox.py
new file mode 100644
index 000000000..5ccf48006
--- /dev/null
+++ b/examples/retool/tool_sandbox.py
@@ -0,0 +1,367 @@
+"""
+Tool sandbox module for safe code execution and tool management.
+
+This module provides:
+- PythonSandbox: Safe Python code execution environment
+- ToolRegistry: Tool registration and execution management
+- Memory management utilities
+"""
+
+import asyncio
+import gc
+import os
+import re
+import subprocess
+import tempfile
+from contextlib import contextmanager
+from typing import Any
+
+import psutil
+
+# Configuration for tool execution
+TOOL_CONFIGS = {
+    "max_turns": 4,
+    "max_tool_calls": 4,
+    "max_consecutive_memory_errors": 3,  # Break retry cycle after N consecutive memory errors
+    "tool_concurrency": 8,  # Conservative: avoid RSS pressure on RolloutManager
+    # Python interpreter settings
+    "python_timeout": 120,  # 2 minutes for complex calculations
+    "python_memory_limit": "4GB",  # 4GB per Python process
+    "python_cpu_limit": 1,
+    # Memory management settings
+    "max_memory_usage": 12288,  # 12GB total (75% of 16GB)
+    "cleanup_threshold": 6144,  # 6GB
+    "aggressive_cleanup_threshold": 3072,  # 3GB
+    "force_cleanup_threshold": 9216,  # 9GB
+}
+
+# Global semaphore for controlling concurrent tool executions
+SEMAPHORE = asyncio.Semaphore(TOOL_CONFIGS["tool_concurrency"])
+
+
+def get_memory_usage() -> float:
+    """Get current memory usage in MB"""
+    process = psutil.Process()
+    return process.memory_info().rss / 1024 / 1024
+
+
+def cleanup_memory():
+    """Force garbage collection to free memory"""
+    gc.collect()
+
+
+def aggressive_cleanup_memory():
+    """More aggressive memory cleanup"""
+    # Force multiple garbage collection cycles
+    for _ in range(3):
+        gc.collect()
+
+    # Force glibc to return freed memory to the OS.
+    # Without this, Python's allocator holds onto freed pages due to
+    # heap fragmentation, causing RSS to stay high even after gc.collect().
+    # Try multiple libc paths for compatibility across different Linux distros
+    # (including Ascend NPU servers).
+    for libc_path in ("libc.so.6", "libc.musl.so.1", "libc.so"):
+        try:
+            import ctypes
+
+            ctypes.CDLL(libc_path).malloc_trim(0)
+            break
+        except OSError:
+            continue
+
+
+def check_and_cleanup_memory():
+    """Check memory usage and perform appropriate cleanup"""
+    current_memory = get_memory_usage()
+
+    if current_memory > TOOL_CONFIGS["force_cleanup_threshold"]:
+        # Force aggressive cleanup
+        aggressive_cleanup_memory()
+        return f"Warning: High memory usage ({current_memory:.1f}MB), performed aggressive cleanup"
+    elif current_memory > TOOL_CONFIGS["cleanup_threshold"]:
+        # Normal cleanup
+        cleanup_memory()
+        return f"Info: Memory usage ({current_memory:.1f}MB), performed cleanup"
+    elif current_memory > TOOL_CONFIGS["aggressive_cleanup_threshold"]:
+        # Light cleanup
+        gc.collect()
+        return f"Info: Memory usage ({current_memory:.1f}MB), performed light cleanup"
+
+    return None
+
+
+class PythonSandbox:
+    """Python code sandbox, provides safe code execution environment"""
+
+    def __init__(self, timeout: int = 10, memory_limit: str = "100MB"):
+        self.timeout = timeout
+        self.memory_limit = memory_limit
+        self.allowed_modules = {
+            "math",
+            "random",
+            "datetime",
+            "collections",
+            "itertools",
+            "functools",
+            "operator",
+            "statistics",
+            "decimal",
+            "fractions",
+        }
+
+    def _check_code_safety(self, code: str) -> tuple[bool, str]:
+        """Check code safety by scanning for dangerous patterns"""
+        # Check for dangerous operations
+        dangerous_patterns = [
+            r"import\s+os",
+            r"import\s+sys",
+            r"import\s+subprocess",
+            r"import\s+shutil",
+            r"import\s+glob",
+            r"import\s+pathlib",
+            r"__import__",
+            r"eval\s*\(",
+            r"exec\s*\(",
+            r"open\s*\(",
+            r"file\s*\(",
+            r"input\s*\(",
+            r"raw_input\s*\(",
+            r"compile\s*\(",
+            r"execfile\s*\(",
+            r"getattr\s*\(",
+            r"setattr\s*\(",
+            r"delattr\s*\(",
+            r"hasattr\s*\(",
+            r"globals\s*\(",
+            r"locals\s*\(",
+            r"vars\s*\(",
+            r"dir\s*\(",
+            r"type\s*\(",
+            r"isinstance\s*\(",
+            r"issubclass\s*\(",
+            r"super\s*\(",
+            r"property\s*\(",
+            r"staticmethod\s*\(",
+            r"classmethod\s*\(",
+            r"__\w+__",  # double underscore methods
+        ]
+
+        for pattern in dangerous_patterns:
+            if re.search(pattern, code, re.IGNORECASE):
+                return False, f"Code contains dangerous pattern: {pattern}"
+
+        # Check imported modules
+        import_pattern = r"import\s+(\w+)"
+        from_pattern = r"from\s+(\w+)"
+
+        imports = re.findall(import_pattern, code)
+        froms = re.findall(from_pattern, code)
+
+        all_imports = set(imports + froms)
+        for imp in all_imports:
+            if imp not in self.allowed_modules:
+                return False, f"Import of '{imp}' is not allowed"
+
+        return True, "Code is safe"
+
+    @contextmanager
+    def _create_safe_environment(self):
+        """Create safe execution environment with temporary directory"""
+        # Create temporary directory
+        temp_dir = tempfile.mkdtemp(prefix="python_sandbox_")
+
+        try:
+            # Create safe Python script
+            script_path = os.path.join(temp_dir, "code.py")
+
+            # Set environment variables
+            env = os.environ.copy()
+            env["PYTHONPATH"] = temp_dir
+            env["PYTHONUNBUFFERED"] = "1"
+
+            yield script_path, env, temp_dir
+
+        finally:
+            # Clean up temporary directory
+            try:
+                import shutil
+
+                shutil.rmtree(temp_dir)
+            except Exception:
+                pass
+
+    async def execute_code(self, code: str) -> str:
+        """Execute Python code in sandbox with safety checks"""
+        # Check memory usage before execution
+        current_memory = get_memory_usage()
+        if current_memory > TOOL_CONFIGS["max_memory_usage"]:
+            aggressive_cleanup_memory()
+            # Re-check after cleanup — gc.collect + malloc_trim may free enough
+            current_memory = get_memory_usage()
+            if current_memory > TOOL_CONFIGS["max_memory_usage"]:
+                return "Error: Memory usage too high, please try again"
+
+        # Check code safety
+        is_safe, message = self._check_code_safety(code)
+        if not is_safe:
+            return f"Error: {message}"
+
+        # Add necessary wrapper code with memory limits
+        # Properly indent the user code within the try block
+        # Handle indentation properly by adding 4 spaces to each line
+        indented_code = "\n".join("    " + line for line in code.split("\n"))
+
+        wrapped_code = f"""import sys
+import traceback
+from io import StringIO
+import resource
+
+# Set memory limit (4GB)
+try:
+    resource.setrlimit(resource.RLIMIT_AS, (4 * 1024 * 1024 * 1024, -1))
+except Exception:
+    pass
+
+# Redirect stdout and stderr
+old_stdout = sys.stdout
+old_stderr = sys.stderr
+stdout_capture = StringIO()
+stderr_capture = StringIO()
+sys.stdout = stdout_capture
+sys.stderr = stderr_capture
+
+try:
+    # User code
+{indented_code}
+
+    # Get output
+    stdout_output = stdout_capture.getvalue()
+    stderr_output = stderr_capture.getvalue()
+
+    # Restore standard output
+    sys.stdout = old_stdout
+    sys.stderr = old_stderr
+
+    # Return result
+    result = ""
+    if stdout_output:
+        result += f"Output:\\n{{stdout_output}}"
+    if stderr_output:
+        result += f"\\nErrors:\\n{{stderr_output}}"
+
+    print(result)
+
+except Exception as e:
+    # Restore standard output
+    sys.stdout = old_stdout
+    sys.stderr = old_stderr
+
+    # Return error information
+    error_msg = f"Error: {{str(e)}}\\nTraceback:\\n{{traceback.format_exc()}}"
+    print(error_msg)"""
+
+        with self._create_safe_environment() as (script_path, env, temp_dir):
+            # Write code to file
+            with open(script_path, "w") as f:
+                f.write(wrapped_code)
+
+            try:
+                # Use subprocess to run code
+                process = subprocess.Popen(
+                    ["python3", script_path],
+                    stdout=subprocess.PIPE,
+                    stderr=subprocess.PIPE,
+                    env=env,
+                    cwd=temp_dir,
+                    text=True,
+                )
+
+                # Set timeout
+                try:
+                    stdout, stderr = process.communicate(timeout=self.timeout)
+
+                    if process.returncode == 0:
+                        result = stdout.strip()
+                    else:
+                        result = f"Error: Process exited with code {process.returncode}\n{stderr}"
+
+                except subprocess.TimeoutExpired:
+                    process.kill()
+                    result = f"Error: Code execution timed out after {self.timeout} seconds"
+
+                # Explicitly release subprocess resources to reduce RSS
+                del process
+
+            except Exception as e:
+                result = f"Error: Failed to execute code: {str(e)}"
+
+            # Check memory usage after execution and cleanup if needed
+            cleanup_message = check_and_cleanup_memory()
+            if cleanup_message:
+                print(f"Memory cleanup: {cleanup_message}")
+
+            return result
+
+
+class ToolRegistry:
+    """Tool registry, manages available tools and their execution"""
+
+    def __init__(self):
+        self.tools = {}
+        self.python_sandbox = PythonSandbox(
+            timeout=TOOL_CONFIGS["python_timeout"], memory_limit=TOOL_CONFIGS["python_memory_limit"]
+        )
+        self._register_default_tools()
+
+    def _register_default_tools(self):
+        """Register default tools in the registry"""
+        # Python code interpreter
+        self.register_tool(
+            "code_interpreter",
+            {
+                "type": "function",
+                "function": {
+                    "name": "code_interpreter",
+                    "description": "A tool for executing Python code in a safe sandbox environment.",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {"code": {"type": "string", "description": "The Python code to execute"}},
+                        "required": ["code"],
+                    },
+                },
+            },
+        )
+
+    def register_tool(self, name: str, tool_spec: dict[str, Any]):
+        """Register a new tool in the registry"""
+        self.tools[name] = tool_spec
+
+    def get_tool_specs(self) -> list[dict[str, Any]]:
+        """Get all tool specifications as a list"""
+        return list(self.tools.values())
+
+    async def execute_tool(self, tool_name: str, arguments: dict[str, Any]) -> str:
+        """Execute a tool call with the given arguments"""
+        if tool_name not in self.tools:
+            return f"Error: Tool '{tool_name}' not found"
+
+        async with SEMAPHORE:
+            if tool_name == "code_interpreter":
+                return await self._execute_python(arguments)
+            else:
+                return f"Error: Tool '{tool_name}' not implemented"
+
+    async def _execute_python(self, arguments: dict[str, Any]) -> str:
+        """Execute Python code using the sandbox"""
+        code = arguments.get("code", "")
+        if not code.strip():
+            return "Error: No code provided"
+
+        # Execute code in sandbox
+        result = await self.python_sandbox.execute_code(code)
+        return result
+
+
+# Global tool registry instance
+tool_registry = ToolRegistry()
diff --git a/examples/search-r1/README.md b/examples/search-r1/README.md
new file mode 100644
index 000000000..05228154f
--- /dev/null
+++ b/examples/search-r1/README.md
@@ -0,0 +1,236 @@
+# Search-R1 lite
+
+This example **(Search-R1 lite)** demonstrates how to use vime for tool-enabled language model generation with search/retrieval capabilities, based on a minimal reproduction of [Search-R1](https://github.com/PeterGriffinJin/Search-R1).
+
+## Overview
+
+The Search-R1 example provides:
+
+- **Multi-turn conversation** with tool-calling (search/answer actions)
+- **Dual search backend support**: local dense retriever (FAISS + E5) or Google Search (serper.dev)
+- **GRPO-based RL training** with exact-match (EM) reward for QA tasks
+- **TIS (Trajectory Importance Sampling)** support for handling train/inference mismatch
+- **Format-aware reward** that evaluates both answer correctness and output structure
+
+## Files
+
+| File                                        | Description                                                                        |
+|---------------------------------------------|------------------------------------------------------------------------------------|
+| `generate_with_search.py`                   | Main generation function with multi-turn search + answer tool-calling, and reward function |
+| `google_search_server.py`                   | Google Search backend via serper.dev API                                           |
+| `local_search_server.py`                    | Local search backend that wraps the retrieval server                               |
+| `qa_em_format.py`                           | QA exact-match scoring with format validation and retrieval correctness check      |
+| `run_qwen3_4b_npu.sh`                       | Training launch script (NPU 8-card, Qwen3-4B-Instruct-2507, GRPO)                                        |
+| `local_dense_retriever/retrieval_server.py` | Dense retriever server (FAISS + E5 model)                                          |
+| `local_dense_retriever/download.py`         | Download wiki-18 index and corpus from HuggingFace                                 |
+
+---
+
+## Usage
+
+### 1. Setup
+
+```bash
+git clone -b ascend https://github.com/vllm-project/vime.git
+cd vime
+docker build -f docker/Dockerfile.npu -t vime-ascend:latest .
+```
+
+```bash
+# Update the vime image
+export IMAGE=vime-ascend:latest
+
+docker run -d --name vime-npu -it --net=host --shm-size=1024g \
+    --privileged=true \
+    --cap-add=SYS_PTRACE \
+    --device=/dev/davinci_manager \
+    --device=/dev/hisi_hdc \
+    --device=/dev/devmm_svm \
+    -v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
+    -v /usr/local/dcmi:/usr/local/dcmi \
+    -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
+    -v /usr/local/sbin:/usr/local/sbin \
+    -v /home:/home \
+    -v /mnt:/mnt \
+    -v /tmp:/tmp \
+    -v /data:/data \
+    -v /usr/share/zoneinfo/Asia/Shanghai:/etc/localtime \
+    $IMAGE
+
+docker exec -it vime-npu bash
+```
+
+```bash
+# For retriever
+pip install faiss-cpu==1.13.2
+```
+
+### 2. Download
+
+#### 2.1 Data
+
+**Option A: Online Auto Download**
+
+```bash
+# Training data (NQ + HotpotQA)
+cd /root
+git clone https://github.com/PeterGriffinJin/Search-R1.git
+cd Search-R1
+pip install -e . --no-deps
+pip install tensordict
+pip install chardet
+
+# Set your working directory
+WORK_DIR=/root/Search-R1
+LOCAL_DIR=/path/to/nq_hotpotqa_train
+
+# Process multiple dataset search format train file
+DATA=nq,hotpotqa
+python $WORK_DIR/scripts/data_process/qa_search_train_merge.py \
+    --local_dir $LOCAL_DIR \
+    --data_sources $DATA
+
+# (Optional) Process multiple dataset search format test file
+DATA=nq,triviaqa,popqa,hotpotqa,2wikimultihopqa,musique,bamboogle
+python $WORK_DIR/scripts/data_process/qa_search_test_merge.py \
+    --local_dir $LOCAL_DIR \
+    --data_sources $DATA
+```
+
+**Option B: Offline Manual Download**
+
+- Download full dataset assets from Hugging Face repo: [PeterJinGo/nq_hotpotqa_train](https://huggingface.co/datasets/PeterJinGo/nq_hotpotqa_train)
+- Upload all downloaded files to your target directory `LOCAL_DIR=/path/to/nq_hotpotqa_train`
+
+#### 2.2 Model
+
+**Option A: Online Auto Download**
+
+```bash
+hf download Qwen/Qwen3-4B-Instruct- --local-dir /path/to/Qwen3-4B-Instruct-2507
+```
+
+**Option B: Offline Manual Download**
+
+- Download full model weights from Hugging Face repo: [Qwen/Qwen3-4B-Instruct-2507](https://huggingface.co/Qwen/Qwen3-4B-Instruct-2507)
+- Upload all downloaded model files to your target directory `MODEL_DIR=/path/to/Qwen3-4B-Instruct-2507`
+
+#### 2.3 Local Retrieval Server (Optional)
+
+> Only needed if using the local search backend instead of Google Search.
+
+**(1) Data**
+
+**Option A: Online Auto Download**
+
+```bash
+# Download index and corpus (~60-70 GB download)
+SAVE_PATH=/path/to/Index
+python /root/vime/examples/search-r1/local_dense_retriever/download.py --save_path $SAVE_PATH
+```
+
+**Option B: Offline Manual Download**
+
+- Download full index assets from Hugging Face repo: [PeterJinGo/wiki-18-e5-index](https://huggingface.co/datasets/PeterJinGo/wiki-18-e5-index) and [PeterJinGo/wiki-18-corpus](https://huggingface.co/datasets/PeterJinGo/wiki-18-corpus)
+- Upload all downloaded files to your target directory `SAVE_PATH=/path/to/Index`
+
+> No matter which download option you use (Option A or B), you must run the following commands after the download completes to merge the index shards and unpack the corpus dataset.
+
+```bash
+SAVE_PATH=/path/to/Index
+cat $SAVE_PATH/part_* > $SAVE_PATH/e5_Flat.index
+gzip -d $SAVE_PATH/wiki-18.jsonl.gz
+```
+
+**(2) Model**
+
+**Option A: Online Auto Download**
+
+```bash
+hf download intfloat/e5-base-v2 --local-dir /path/to/e5-base-v2
+```
+
+**Option B: Offline Manual Download**
+
+- Download full model weights from Hugging Face repo: [intfloat/e5-base-v2](https://huggingface.co/intfloat/e5-base-v2)
+- Upload all downloaded model files to your target directory `MODEL_DIR=/path/to/e5-base-v2`
+
+### 3. Configure Search Backend
+
+The `generate_with_search.py` file supports both **local** search and **Google** search backends. Configure via the `SEARCH_R1_CONFIGS` dictionary:
+
+```python
+SEARCH_R1_CONFIGS = {
+    # ============== General Configuration ==============
+    "max_turns": 2,
+    "topk": 3,
+    "search_concurrency": 8,
+
+    # ============== Search Backend Selection ==============
+    "search_backend": "local",  # Options: "local" or "google"
+
+    # ============== Local Search Configuration ==============
+    # (Only used when search_backend="local")
+    "local": {
+        "search_url": "http://127.0.0.1:8000/retrieve",  # URL of your local retrieval server
+        "proxy": None,
+    },
+
+    # ============== Google Search Configuration ==============
+    # (Only used when search_backend="google")
+    "google": {
+        "api_key": "your_api_key_here",  # Replace with your actual serper.dev API key
+        "snippet_only": True,
+        "proxy": None,
+    },
+
+    # ============== Log Probability Collection ==============
+    "return_logprob": True,  # Set to True to collect log probabilities (required for TIS)
+
+    # ============== Reward Model Configuration ==============
+    "format_score": 0.2,
+}
+```
+
+#### Using Local Search
+
+- Set `"search_backend": "local"`
+- Configure `"local"` section with your local retrieval server URL
+- Start your local search server before running the training script
+
+```bash
+# Set paths
+SAVE_PATH=/path/to/Index
+INDEX_FILE=$SAVE_PATH/e5_Flat.index
+CORPUS_FILE=$SAVE_PATH/wiki-18.jsonl
+RETRIEVER_NAME=e5
+RETRIEVER_PATH=/path/to/e5-base-v2
+
+# Start the retrieval server
+python /root/vime/examples/search-r1/local_dense_retriever/retrieval_server.py \
+    --index_path $INDEX_FILE\
+    --corpus_path $CORPUS_FILE\
+    --topk 3 \
+    --retriever_name $RETRIEVER_NAME \
+    --retriever_model $RETRIEVER_PATH
+```
+
+> **Note:**
+> - First startup will download the model and load the index, which may take a few minutes
+> - Normal startup time (excluding downloads): 1-2 minutes
+> - The local search engine's Python process will not terminate when the shell closes
+> - To restart the server: `lsof -i :8000` to find the PID, then kill it and restart
+
+#### Using Google Search
+
+- Set `"search_backend": "google"`
+- Configure `"google"` section with your serper.dev API key
+- Get your API key from [serper.dev](https://serper.dev/)
+
+### 4. Run Training
+
+```bash
+cd /root/vime
+# Replace the model and data loading/saving paths
+bash examples/search-r1/run_qwen3_4b_npu.sh
+```
diff --git a/examples/search-r1/README_zh.md b/examples/search-r1/README_zh.md
new file mode 100644
index 000000000..7d196ea0e
--- /dev/null
+++ b/examples/search-r1/README_zh.md
@@ -0,0 +1,234 @@
+# Search-R1 lite
+
+本示例 **(Search-R1 lite)** 演示了如何使用 vime 进行带工具调用的语言模型生成，具备搜索/检索能力，是基于 [Search-R1](https://github.com/PeterGriffinJin/Search-R1) 的最小复现。
+
+## 概述
+
+Search-R1 示例提供了：
+
+- **多轮对话**：支持工具调用（搜索/回答）
+- **双搜索后端支持**：本地稠密检索器（FAISS + E5）或 Google 搜索（serper.dev）
+- **基于 GRPO 的强化学习训练**：使用精确匹配（EM）计算问答任务奖励
+- **TIS（轨迹重要性采样）**：用于处理训推不一致问题
+- **格式感知奖励**：同时评估答案正确性和输出结构
+
+## 文件说明
+
+| 文件                                          | 描述 |
+|---------------------------------------------|------|
+| `generate_with_search.py`                   | 主生成函数，支持多轮搜索 + 回答工具调用，以及奖励函数 |
+| `google_search_server.py`                   | 通过 serper.dev API 实现 Google 搜索后端 |
+| `local_search_server.py`                    | 本地搜索后端，封装检索服务 |
+| `qa_em_format.py`                           | 问答精确匹配评分，包含格式验证和检索正确性检查 |
+| `run_qwen3_4b_npu.sh`                       | 训练启动脚本（NPU 8卡，Qwen3-4B-Instruct-2507，GRPO） |
+| `local_dense_retriever/retrieval_server.py` | 稠密检索服务器（FAISS + E5 模型） |
+| `local_dense_retriever/download.py`         | 从 HuggingFace 下载 wiki-18 索引和语料库 |
+
+---
+
+## 使用方法
+
+### 1. 环境搭建
+
+```bash
+git clone -b ascend https://github.com/vllm-project/vime.git
+cd vime
+docker build -f docker/Dockerfile.npu -t vime-ascend:latest .
+```
+
+```bash
+export IMAGE=vime-ascend:latest
+
+docker run -d --name vime-npu -it --net=host --shm-size=1024g \
+    --privileged=true \
+    --cap-add=SYS_PTRACE \
+    --device=/dev/davinci_manager \
+    --device=/dev/hisi_hdc \
+    --device=/dev/devmm_svm \
+    -v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
+    -v /usr/local/dcmi:/usr/local/dcmi \
+    -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
+    -v /usr/local/sbin:/usr/local/sbin \
+    -v /home:/home \
+    -v /mnt:/mnt \
+    -v /tmp:/tmp \
+    -v /data:/data \
+    -v /usr/share/zoneinfo/Asia/Shanghai:/etc/localtime \
+    $IMAGE
+
+docker exec -it vime-npu bash
+```
+
+```bash
+# 检索依赖
+pip install faiss-cpu==1.13.2
+```
+
+### 2. 下载
+
+#### 2.1 数据
+
+**方式 A：在线自动下载**
+
+```bash
+cd /root
+git clone https://github.com/PeterGriffinJin/Search-R1.git
+cd Search-R1
+pip install -e . --no-deps
+pip install tensordict
+pip install chardet
+
+# 设置工作目录
+WORK_DIR=/root/Search-R1
+LOCAL_DIR=/path/to/nq_hotpotqa_train
+
+# 处理多数据集搜索格式训练文件
+DATA=nq,hotpotqa
+python $WORK_DIR/scripts/data_process/qa_search_train_merge.py \
+    --local_dir $LOCAL_DIR \
+    --data_sources $DATA
+
+# （可选）处理多数据集搜索格式测试文件
+DATA=nq,triviaqa,popqa,hotpotqa,2wikimultihopqa,musique,bamboogle
+python $WORK_DIR/scripts/data_process/qa_search_test_merge.py \
+    --local_dir $LOCAL_DIR \
+    --data_sources $DATA
+```
+
+**方式 B：离线手动下载**
+
+- 从 Hugging Face 仓库下载完整数据集资源：[PeterJinGo/nq_hotpotqa_train](https://huggingface.co/datasets/PeterJinGo/nq_hotpotqa_train)
+- 将所有下载的文件上传到目标目录 `LOCAL_DIR=/path/to/nq_hotpotqa_train`
+
+#### 2.2 模型
+
+**方式 A：在线自动下载**
+
+```bash
+hf download Qwen/Qwen3-4B-Instruct-2507 --local-dir /path/to/Qwen3-4B-Instruct-2507
+```
+
+**方式 B：离线手动下载**
+
+- 从 Hugging Face 仓库下载完整模型权重：[Qwen/Qwen3-4B-Instruct-2507](https://huggingface.co/Qwen/Qwen3-4B-Instruct-2507)
+- 将所有下载的模型文件上传到目标目录 `MODEL_DIR=/path/to/Qwen3-4B-Instruct-2507`
+
+#### 2.3 本地检索服务器（可选）
+
+> 仅在使用本地搜索后端而非 Google 搜索时需要。
+
+**(1) 数据**
+
+**方式 A：在线自动下载**
+
+```bash
+# 下载索引和语料库（约 60-70 GB 下载量）
+SAVE_PATH=/path/to/Index
+python /root/vime/examples/search-r1/local_dense_retriever/download.py --save_path $SAVE_PATH
+```
+
+**方式 B：离线手动下载**
+
+- 从 Hugging Face 仓库下载完整索引资源：[PeterJinGo/wiki-18-e5-index](https://huggingface.co/datasets/PeterJinGo/wiki-18-e5-index) 和 [PeterJinGo/wiki-18-corpus](https://huggingface.co/datasets/PeterJinGo/wiki-18-corpus)
+- 将所有下载的文件上传到目标目录 `SAVE_PATH=/path/to/Index`
+
+> **注意：**无论采用上述哪种下载方式（方案 A 或方案 B），在数据下载完成后，均须执行以下命令以完成索引分片的合并以及语料库的解压。
+
+```bash
+SAVE_PATH=/path/to/Index
+cat $SAVE_PATH/part_* > $SAVE_PATH/e5_Flat.index
+gzip -d $SAVE_PATH/wiki-18.jsonl.gz
+```
+
+**(2) 模型**
+
+**方式 A：在线自动下载**
+
+```bash
+hf download intfloat/e5-base-v2 --local-dir /path/to/e5-base-v2
+```
+
+**方式 B：离线手动下载**
+
+- 从 Hugging Face 仓库下载完整模型权重：[intfloat/e5-base-v2](https://huggingface.co/intfloat/e5-base-v2)
+- 将所有下载的模型文件上传到目标目录 `MODEL_DIR=/path/to/e5-base-v2`
+
+### 3. 配置搜索后端
+
+`generate_with_search.py` 文件同时支持**本地**搜索和 **Google** 搜索后端。通过 `SEARCH_R1_CONFIGS` 字典进行配置：
+
+```python
+SEARCH_R1_CONFIGS = {
+    # ============== 通用配置 ==============
+    "max_turns": 2,
+    "topk": 3,
+    "search_concurrency": 8,
+
+    # ============== 搜索后端选择 ==============
+    "search_backend": "local",  # 选项："local" 或 "google"
+
+    # ============== 本地搜索配置 ==============
+    # （仅在 search_backend="local" 时使用）
+    "local": {
+        "search_url": "http://127.0.0.1:8000/retrieve",  # 本地检索服务器 URL
+        "proxy": None,
+    },
+
+    # ============== Google 搜索配置 ==============
+    # （仅在 search_backend="google" 时使用）
+    "google": {
+        "api_key": "your_api_key_here",  # 替换为你的 serper.dev API 密钥
+        "snippet_only": True,
+        "proxy": None,
+    },
+
+    # ============== 对数概率收集 ==============
+    "return_logprob": True,  # 设为 True 以收集对数概率（TIS 所需）
+
+    # ============== 奖励模型配置 ==============
+    "format_score": 0.2,
+}
+```
+
+#### 使用本地搜索
+
+- 设置 `"search_backend": "local"`
+- 在 `"local"` 部分配置本地检索服务器 URL
+- 在运行训练脚本之前启动本地搜索服务器
+
+```bash
+# 设置路径
+SAVE_PATH=/path/to/Index
+INDEX_FILE=$SAVE_PATH/e5_Flat.index
+CORPUS_FILE=$SAVE_PATH/wiki-18.jsonl
+RETRIEVER_NAME=e5
+RETRIEVER_PATH=/path/to/e5-base-v2
+
+# 启动检索服务器
+python /root/vime/examples/search-r1/local_dense_retriever/retrieval_server.py \
+    --index_path $INDEX_FILE\
+    --corpus_path $CORPUS_FILE\
+    --topk 3 \
+    --retriever_name $RETRIEVER_NAME \
+    --retriever_model $RETRIEVER_PATH
+```
+
+> **注意：**
+> - 首次启动将下载模型并加载索引，可能需要几分钟
+> - 正常启动时间（不含下载）：1-2 分钟
+> - 本地搜索引擎的 Python 进程不会在 Shell 关闭时终止
+> - 重启服务器：使用 `lsof -i :8000` 查找进程 PID，然后终止并重启
+
+#### 使用 Google 搜索
+
+- 设置 `"search_backend": "google"`
+- 在 `"google"` 部分配置你的 serper.dev API 密钥
+- 从 [serper.dev](https://serper.dev/) 获取你的 API 密钥
+
+### 4. 运行训练
+
+```bash
+cd /root/vime
+# 替换模型和数据加载/保存路径
+bash examples/search-r1/run_qwen3_4b_npu.sh
+```
diff --git a/examples/search-r1/generate_with_search.py b/examples/search-r1/generate_with_search.py
new file mode 100644
index 000000000..e2ac4aac2
--- /dev/null
+++ b/examples/search-r1/generate_with_search.py
@@ -0,0 +1,307 @@
+# Adapted from https://github.com/PeterGriffinJin/Search-R1/blob/ceee7b89655ed52f205b9beb98e1190c3eedcfb0/search_r1/llm_agent/generation.py
+# This is a unified version supporting both local search and Google search, with optional log probability collection
+# Adapted for vLLM /inference/v1/generate endpoint (disagg API)
+
+import asyncio
+import re
+
+from qa_em_format import compute_score_em
+
+from vime.rollout.vllm_rollout import GenerateState, _build_inference_sampling_params
+from vime.utils.http_utils import post
+from vime.utils.types import Sample
+
+# Configuration for Search-R1
+SEARCH_R1_CONFIGS = {
+    # ============== General Configuration ==============
+    "max_turns": 2,
+    "topk": 3,
+    "search_concurrency": 8,  # TODO: temporarily lowered for CPU-based Faiss retrieval; restore to 256 once NPU-accelerated Faiss is available
+    # ============== Search Backend Selection ==============
+    "search_backend": "local",  # Options: "local" or "google"
+    # ============== Local Search Configuration ==============
+    # (Only used when search_backend="local")
+    "local": {
+        "search_url": "http://127.0.0.1:8000/retrieve",  # URL of your local retrieval server
+        "proxy": None,  # Set to your proxy if needed
+    },
+    # ============== Google Search Configuration ==============
+    # (Only used when search_backend="google")
+    "google": {
+        "api_key": "your_api_key_here",  # Replace with your actual API key
+        "snippet_only": True,  # Set to True to only return snippets
+        "proxy": None,  # Set to your proxy if needed
+    },
+    # ============== Log Probability Collection ==============
+    "return_logprob": True,  # Set to True to collect log probabilities for TIS metrics
+    # ============== Reward Model Configuration ==============
+    "format_score": 0.2,
+}
+
+
+SEMAPHORE = asyncio.Semaphore(SEARCH_R1_CONFIGS["search_concurrency"])
+
+
+def _passages2string(retrieval_result):
+    """
+    Convert retrieval results to a formatted string.
+    This function works with both google_search and local_search results.
+    """
+    format_reference = ""
+    for idx, doc_item in enumerate(retrieval_result):
+        content = doc_item["document"]["contents"]
+        title = content.split("\n")[0]
+        text = "\n".join(content.split("\n")[1:])
+        format_reference += f"Doc {idx+1}(Title: {title}) {text}\n"
+
+    return format_reference
+
+
+async def search(query: str) -> str:
+    """
+    Perform search using either local search engine or Google search.
+    The search backend is determined by SEARCH_R1_CONFIGS["search_backend"].
+    """
+    backend = SEARCH_R1_CONFIGS["search_backend"]
+
+    if backend == "local":
+        from local_search_server import local_search
+
+        local_config = SEARCH_R1_CONFIGS["local"]
+        result = await local_search(
+            local_config["search_url"],
+            query,
+            SEARCH_R1_CONFIGS["topk"],
+            proxy=local_config["proxy"],
+        )
+    elif backend == "google":
+        from google_search_server import google_search
+
+        google_config = SEARCH_R1_CONFIGS["google"]
+        result = await google_search(
+            google_config["api_key"],
+            query,
+            SEARCH_R1_CONFIGS["topk"],
+            snippet_only=google_config["snippet_only"],
+            proxy=google_config["proxy"],
+        )
+    else:
+        raise ValueError(f"Unknown search backend: {backend}. " f"Must be either 'local' or 'google'.")
+
+    return _passages2string(result)
+
+
+# IMPORTANT: When we need to collect log probabilities (logp), we CANNOT do any postprocessing
+# on the strings decoded from token_ids. This is because:
+# 1. We don't know how to truncate the corresponding tokens/logp arrays to match the modified string
+# 2. Re-tokenizing the postprocessed string may produce different tokens than what the engine generated,
+#    leading to misalignment between tokens and their log probabilities
+# Therefore, postprocess_responses is only used when return_logprob=False.
+def postprocess_responses(resp: str) -> str:
+    """
+    Post-process response to ensure tag completeness.
+    Only used when SEARCH_R1_CONFIGS["return_logprob"] is False.
+    """
+    return (
+        resp.split("</search>")[0] + "</search>"
+        if "</search>" in resp
+        else resp.split("</answer>")[0] + "</answer>" if "</answer>" in resp else resp
+    )
+
+
+def postprocess_predictions(prediction: str):
+    pattern = r"<(search|answer)>(.*?)</\1>"
+    match = re.search(pattern, prediction, re.DOTALL)
+    if match:
+        content = match.group(2).strip()  # Return only the content inside the tags
+        action = match.group(1)
+    else:
+        content = ""
+        action = None
+
+    return action, content
+
+
+async def execute_predictions(prediction: str) -> tuple[str, bool]:
+    action, content = postprocess_predictions(prediction)
+
+    if action == "search":
+        search_query = content
+        async with SEMAPHORE:
+            search_results = await search(search_query)
+        next_obs = f"\n\n<information>{search_results.strip()}</information>\n\n"
+        done = False
+    elif action == "answer":
+        next_obs = ""
+        done = True
+    else:
+        next_obs = "\nMy previous action is invalid. \
+If I want to search, I should put the query between <search> and </search>. \
+If I want to give the final answer, I should put the answer between <answer> and </answer>. Let me try again.\n"
+        done = False
+
+    return next_obs, done
+
+
+async def generate(args, sample: Sample, sampling_params) -> Sample:
+    assert not args.partial_rollout, "Partial rollout is not supported for this function at the moment."
+
+    state = GenerateState(args)
+
+    router_ip = args.vllm_router_ip
+    router_port = args.vllm_router_port
+    url = f"http://{router_ip}:{router_port}/inference/v1/generate"
+
+    # Handle partial rollout samples: continue generation from existing response
+    prompt_text = sample.prompt
+    prompt_tokens_ids = state.tokenizer(prompt_text, add_special_tokens=False)["input_ids"]
+    sample.tokens = list(prompt_tokens_ids)
+    sample.loss_mask = []
+    response = ""
+    response_token_ids = []
+    loss_mask = []
+    rollout_log_probs = [] if SEARCH_R1_CONFIGS["return_logprob"] else None
+    sample.rollout_top_p_token_ids = None
+    sample.rollout_top_p_token_offsets = None
+
+    # BUGFIX: make the inference engine STOP at the tool/answer boundary.
+    # Without a stop, the engine keeps emitting tokens after </search> / </answer>
+    # (junk, even fabricated new "Question:"s). The example only trimmed that junk
+    # via postprocess_responses when return_logprob=False; with return_logprob=True
+    # (TIS) trimming is disabled to keep token/logp aligned, so the junk stayed in
+    # the trajectory and got trained on (loss_mask=1) AND broke is_valid_sequence
+    # (trailing content after </answer> -> format invalid -> lower reward).
+    # Stopping at the tag avoids all of that and keeps token/logp aligned natively.
+    _stop_tags = ["</search>", "</answer>"]
+    _existing_stop = sampling_params.get("stop") or []
+    if isinstance(_existing_stop, str):
+        _existing_stop = [_existing_stop]
+    sampling_params = {**sampling_params, "stop": list(dict.fromkeys([*_existing_stop, *_stop_tags]))}
+
+    # Build vLLM-style sampling params (maps max_new_tokens -> max_tokens, adds logprobs, etc.)
+    inference_sampling_params = _build_inference_sampling_params(sampling_params)
+
+    for _turn_idx in range(SEARCH_R1_CONFIGS["max_turns"]):
+        # vLLM /inference/v1/generate requires token_ids instead of text.
+        # For multi-turn, re-tokenize the full context each time.
+        full_text = prompt_text + response
+        full_token_ids = state.tokenizer(full_text, add_special_tokens=False)["input_ids"]
+
+        payload = {
+            "token_ids": full_token_ids,
+            "sampling_params": inference_sampling_params,
+        }
+        if hasattr(args, "hf_checkpoint"):
+            payload["model"] = args.hf_checkpoint
+
+        output = await post(url, payload)
+
+        # Parse vLLM GenerateResponse: {"choices": [{"token_ids": [...], "logprobs": ..., "finish_reason": "stop"}]}
+        choice = output["choices"][0]
+        finish_reason = choice.get("finish_reason", "stop")
+
+        # abort
+        if finish_reason == "abort":
+            sample.status = Sample.Status.ABORTED
+            return sample
+
+        # Extract token IDs from vLLM response
+        cur_response_token_ids = choice.get("token_ids") or []
+
+        # Decode text from token_ids
+        skip_sp = inference_sampling_params.get("skip_special_tokens")
+        skip_decode = True if skip_sp is None else bool(skip_sp)
+        cur_response = (
+            state.tokenizer.decode(cur_response_token_ids, skip_special_tokens=skip_decode)
+            if cur_response_token_ids
+            else ""
+        )
+
+        # Extract log probs if enabled
+        if SEARCH_R1_CONFIGS["return_logprob"]:
+            cur_response_log_probs: list[float] = []
+            lp = choice.get("logprobs")
+            if isinstance(lp, dict):
+                content_items = lp.get("content") or []
+                cur_response_log_probs = [
+                    float(item.get("logprob", 0.0)) if isinstance(item, dict) else 0.0 for item in content_items
+                ]
+            if not cur_response_log_probs:
+                cur_response_log_probs = [0.0] * len(cur_response_token_ids)
+        else:
+            # When not collecting log probs, we can safely postprocess the response
+            cur_response = postprocess_responses(cur_response)
+            # Re-tokenize after postprocessing
+            cur_response_token_ids = state.tokenizer(cur_response, add_special_tokens=False)["input_ids"]
+            cur_response_log_probs = None
+
+        response += cur_response
+        response_token_ids += cur_response_token_ids
+        loss_mask += [1] * len(cur_response_token_ids)
+
+        # Add log probs if enabled
+        if SEARCH_R1_CONFIGS["return_logprob"]:
+            rollout_log_probs += cur_response_log_probs
+
+        if finish_reason == "length":
+            break
+
+        next_obs, done = await execute_predictions(cur_response)
+        if done:
+            break
+
+        assert next_obs != "", "Next observation should not be empty."
+        obs_tokens_ids = state.tokenizer(next_obs, add_special_tokens=False)["input_ids"]
+        response += next_obs
+        response_token_ids += obs_tokens_ids
+        loss_mask += [0] * len(obs_tokens_ids)
+
+        # Add dummy log probs for observation tokens if enabled (they won't be used due to loss_mask=0)
+        if SEARCH_R1_CONFIGS["return_logprob"]:
+            rollout_log_probs += [0.0] * len(obs_tokens_ids)
+
+            # Verify alignment when collecting log probs
+            assert len(response_token_ids) == len(
+                rollout_log_probs
+            ), f"Token/logp length mismatch: {len(response_token_ids)} tokens vs {len(rollout_log_probs)} logps"
+
+    # Store statistics for wandb logging
+    sample.tokens = prompt_tokens_ids + response_token_ids
+    sample.response_length = len(response_token_ids)
+    sample.response = response
+    sample.loss_mask = loss_mask
+    sample.prompt = prompt_text
+
+    # Store log probs if enabled
+    if SEARCH_R1_CONFIGS["return_logprob"]:
+        sample.rollout_log_probs = rollout_log_probs if rollout_log_probs else None
+
+    # vLLM finish_reason is a string: "stop", "length", or "abort"
+    match finish_reason:
+        case "length":
+            sample.status = Sample.Status.TRUNCATED
+        case "abort":
+            sample.status = Sample.Status.ABORTED
+        case "stop":
+            sample.status = Sample.Status.COMPLETED
+
+    return sample
+
+
+async def reward_func(args, sample, **kwargs):
+    """The reward function for retrieval-based question answering.
+
+    Args:
+        args: the arguments
+        sample: the sample to evaluate
+    """
+    if not isinstance(sample, Sample):
+        raise TypeError("Sample must be an instance of Sample class.")
+
+    score = compute_score_em(
+        solution_str=sample.prompt + sample.response,
+        ground_truth=sample.label["ground_truth"],
+        format_score=SEARCH_R1_CONFIGS["format_score"],
+    )
+
+    return score
diff --git a/examples/search-r1/google_search_server.py b/examples/search-r1/google_search_server.py
new file mode 100644
index 000000000..394315ddb
--- /dev/null
+++ b/examples/search-r1/google_search_server.py
@@ -0,0 +1,149 @@
+import asyncio
+import os
+import random
+import re
+
+import aiohttp
+import chardet
+
+
+# --- Utilities ---
+def parse_snippet(snippet: str) -> list[str]:
+    segments = snippet.split("...")
+    return [s.strip() for s in segments if len(s.strip().split()) > 5]
+
+
+def sanitize_search_query(query: str) -> str:
+    # Remove or replace special characters that might cause issues.
+    # This is a basic example; you might need to add more characters or patterns.
+    sanitized_query = re.sub(r"[^\w\s]", " ", query)  # Replace non-alphanumeric and non-whitespace with spaces.
+    sanitized_query = re.sub(
+        r"[\t\r\f\v\n]", " ", sanitized_query
+    )  # replace tab, return, formfeed, vertical tab with spaces.
+    sanitized_query = re.sub(
+        r"\s+", " ", sanitized_query
+    ).strip()  # remove duplicate spaces, and trailing/leading spaces.
+
+    return sanitized_query
+
+
+def filter_links(search_results: list[dict]) -> list[str]:
+    links = []
+    for result in search_results:
+        for item in result.get("items", []):
+            if "mime" in item:
+                continue
+            ext = os.path.splitext(item["link"])[1]
+            if ext in ["", ".html", ".htm", ".shtml"]:
+                links.append(item["link"])
+    return links
+
+
+async def fetch(session: aiohttp.ClientSession, url: str, semaphore: asyncio.Semaphore) -> str:
+    if url == "":
+        return ""
+    user_agents = [
+        "Mozilla/5.0 (Linux; Android 6.0.1; Nexus 5X Build/MMB29P)...",
+        "Mozilla/5.0 AppleWebKit/537.36...",
+        "Mozilla/5.0 (compatible; Googlebot/2.1; +https://www.google.com/bot.html)",
+    ]
+    headers = {"User-Agent": random.choice(user_agents)}
+
+    async with semaphore:
+        try:
+            async with session.get(url, headers=headers) as response:
+                raw = await response.read()
+                detected = chardet.detect(raw)
+                encoding = detected["encoding"] or "utf-8"
+                return raw.decode(encoding, errors="ignore")
+        except (aiohttp.ClientError, asyncio.TimeoutError):
+            return ""
+
+
+async def fetch_all(urls: list[str], limit: int = 8) -> list[str]:
+    semaphore = asyncio.Semaphore(limit)
+    timeout = aiohttp.ClientTimeout(total=5)
+    connector = aiohttp.TCPConnector(limit_per_host=limit, force_close=True)
+
+    async with aiohttp.ClientSession(timeout=timeout, connector=connector) as session:
+        tasks = [fetch(session, url, semaphore) for url in urls]
+        return await asyncio.gather(*tasks)
+
+
+def collect_context(snippet: str, doc: str) -> str:
+    snippets = parse_snippet(snippet)
+    ctx_paras = []
+
+    for s in snippets:
+        pos = doc.replace("\n", " ").find(s)
+        if pos == -1:
+            continue
+        sta = pos
+        while sta > 0 and doc[sta] != "\n":
+            sta -= 1
+        end = pos + len(s)
+        while end < len(doc) and doc[end] != "\n":
+            end += 1
+        para = doc[sta:end].strip()
+        if para not in ctx_paras:
+            ctx_paras.append(para)
+
+    return "\n".join(ctx_paras)
+
+
+async def google_search(api_key, query, top_k=5, timeout: int = 60, proxy=None, snippet_only=False) -> list[dict]:
+    timeout_obj = aiohttp.ClientTimeout(total=timeout)
+    session_kwargs = {}
+    if proxy:
+        session_kwargs["proxy"] = proxy
+    async with aiohttp.ClientSession(**session_kwargs) as session:
+        async with session.post(
+            "https://google.serper.dev/search",
+            json={
+                "q": query,
+                "num": top_k,
+                "gl": "us",
+                "hl": "en",
+            },
+            headers={
+                "Content-Type": "application/json",
+                "X-API-KEY": api_key,
+            },
+            timeout=timeout_obj,
+        ) as resp:
+            resp.raise_for_status()
+            response = await resp.json()
+            items = response.get("organic", [])
+
+    contexts = []
+    if snippet_only:
+        for item in items:
+            title = item.get("title", "")
+            context = " ".join(parse_snippet(item.get("snippet", "")))
+            if title != "" or context != "":
+                title = "No title." if not title else title
+                context = "No snippet available." if not context else context
+                contexts.append(
+                    {
+                        "document": {"contents": f'"{title}"\n{context}'},
+                    }
+                )
+    else:
+        links = [item.get("link", "") for item in items]
+        web_contents = await fetch_all(links)
+        contexts = []
+        for i, item in enumerate(items):
+            title = item.get("title", "")
+            snippet = item.get("snippet", "")
+
+            context = collect_context(snippet, web_contents[i])
+            if title != "" or context != "":
+                title = "No title." if not title else title
+                context = "No snippet available." if not context else context
+                contexts.append(
+                    {
+                        "document": {"contents": f'"{title}"\n{context}'},
+                    }
+                )
+
+    return contexts
diff --git a/examples/search-r1/local_dense_retriever/download.py b/examples/search-r1/local_dense_retriever/download.py
new file mode 100644
index 000000000..6fe554936
--- /dev/null
+++ b/examples/search-r1/local_dense_retriever/download.py
@@ -0,0 +1,44 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+# Copyright 2023-2024 SGLang Team
+# Copyright 2025 Search-R1 Contributors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# Adapted from https://github.com/PeterGriffinJin/Search-R1/blob/main/scripts/download.py
+
+
+import argparse
+
+from huggingface_hub import hf_hub_download
+
+parser = argparse.ArgumentParser(description="Download files from a Hugging Face dataset repository.")
+parser.add_argument("--repo_id", type=str, default="PeterJinGo/wiki-18-e5-index", help="Hugging Face repository ID")
+parser.add_argument("--save_path", type=str, required=True, help="Local directory to save files")
+
+args = parser.parse_args()
+
+repo_id = "PeterJinGo/wiki-18-e5-index"
+for file in ["part_aa", "part_ab"]:
+    hf_hub_download(
+        repo_id=repo_id,
+        filename=file,  # e.g., "e5_Flat.index"
+        repo_type="dataset",
+        local_dir=args.save_path,
+    )
+
+repo_id = "PeterJinGo/wiki-18-corpus"
+hf_hub_download(
+    repo_id=repo_id,
+    filename="wiki-18.jsonl.gz",
+    repo_type="dataset",
+    local_dir=args.save_path,
+)
diff --git a/examples/search-r1/local_dense_retriever/retrieval_server.py b/examples/search-r1/local_dense_retriever/retrieval_server.py
new file mode 100644
index 000000000..5604bba29
--- /dev/null
+++ b/examples/search-r1/local_dense_retriever/retrieval_server.py
@@ -0,0 +1,435 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+# Copyright 2023-2024 SGLang Team
+# Copyright 2025 Search-R1 Contributors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# Adapted from https://github.com/PeterGriffinJin/Search-R1/blob/main/search_r1/search/retrieval_server.py
+
+import argparse
+import json
+import warnings
+
+import datasets
+import faiss
+import numpy as np
+import torch
+import uvicorn
+from fastapi import FastAPI
+from pydantic import BaseModel
+from tqdm import tqdm
+from transformers import AutoModel, AutoTokenizer
+
+
+def load_corpus(corpus_path: str):
+    corpus = datasets.load_dataset("json", data_files=corpus_path, split="train", num_proc=4)
+    return corpus
+
+
+def load_docs(corpus, doc_idxs):
+    results = [corpus[int(idx)] for idx in doc_idxs]
+    return results
+
+
+def load_model(model_path: str, use_fp16: bool = False, device: torch.device | None = None):
+    """Load transformer model and tokenizer."""
+    if device is None:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    model = AutoModel.from_pretrained(model_path, trust_remote_code=True)
+    model.eval()
+    model.to(device)
+
+    if use_fp16 and device.type == "cuda":
+        model = model.half()
+    elif use_fp16 and device.type != "cuda":
+        warnings.warn("FP16 requested but CUDA is not available; running in FP32 on CPU.", stacklevel=2)
+
+    tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True, trust_remote_code=True)
+    return model, tokenizer
+
+
+def pooling(pooler_output, last_hidden_state, attention_mask=None, pooling_method="mean"):
+    if pooling_method == "mean":
+        last_hidden = last_hidden_state.masked_fill(~attention_mask[..., None].bool(), 0.0)
+        return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
+    elif pooling_method == "cls":
+        return last_hidden_state[:, 0]
+    elif pooling_method == "pooler":
+        return pooler_output
+    else:
+        raise NotImplementedError("Pooling method not implemented!")
+
+
+class Encoder:
+    def __init__(self, model_name, model_path, pooling_method, max_length, use_fp16):
+        self.model_name = model_name
+        self.model_path = model_path
+        self.pooling_method = pooling_method
+        self.max_length = max_length
+        self.use_fp16 = use_fp16
+        # Set the device
+        if torch.cuda.is_available():
+            self.device = torch.device("cuda")
+        else:
+            self.device = torch.device("cpu")
+        self.model, self.tokenizer = load_model(model_path=model_path, use_fp16=use_fp16, device=self.device)
+        self.model.eval()
+
+    @torch.no_grad()
+    def encode(self, query_list: list[str], is_query=True) -> np.ndarray:
+        # processing query for different encoders
+        if isinstance(query_list, str):
+            query_list = [query_list]
+
+        if "e5" in self.model_name.lower():
+            if is_query:
+                query_list = [f"query: {query}" for query in query_list]
+            else:
+                query_list = [f"passage: {query}" for query in query_list]
+
+        if "bge" in self.model_name.lower():
+            if is_query:
+                query_list = [
+                    f"Represent this sentence for searching relevant passages: {query}" for query in query_list
+                ]
+
+        inputs = self.tokenizer(
+            query_list, max_length=self.max_length, padding=True, truncation=True, return_tensors="pt"
+        )
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
+
+        if "T5" in type(self.model).__name__:
+            # T5-based retrieval model
+            decoder_input_ids = torch.zeros((inputs["input_ids"].shape[0], 1), dtype=torch.long).to(
+                inputs["input_ids"].device
+            )
+            output = self.model(**inputs, decoder_input_ids=decoder_input_ids, return_dict=True)
+            query_emb = output.last_hidden_state[:, 0, :]
+        else:
+            output = self.model(**inputs, return_dict=True)
+            query_emb = pooling(
+                output.pooler_output, output.last_hidden_state, inputs["attention_mask"], self.pooling_method
+            )
+            if "dpr" not in self.model_name.lower():
+                query_emb = torch.nn.functional.normalize(query_emb, dim=-1)
+
+        query_emb = query_emb.detach().cpu().numpy()
+        query_emb = query_emb.astype(np.float32, order="C")
+
+        del inputs, output
+        if self.device.type == "cuda":
+            torch.cuda.empty_cache()
+
+        return query_emb
+
+
+class BaseRetriever:
+    def __init__(self, config):
+        self.config = config
+        self.retrieval_method = config.retrieval_method
+        self.topk = config.retrieval_topk
+
+        self.index_path = config.index_path
+        self.corpus_path = config.corpus_path
+
+    def _search(self, query: str, num: int, return_score: bool):
+        raise NotImplementedError
+
+    def _batch_search(self, query_list: list[str], num: int, return_score: bool):
+        raise NotImplementedError
+
+    def search(self, query: str, num: int = None, return_score: bool = False):
+        return self._search(query, num, return_score)
+
+    def batch_search(self, query_list: list[str], num: int = None, return_score: bool = False):
+        return self._batch_search(query_list, num, return_score)
+
+
+class BM25Retriever(BaseRetriever):
+    def __init__(self, config):
+        super().__init__(config)
+        from pyserini.search.lucene import LuceneSearcher
+
+        self.searcher = LuceneSearcher(self.index_path)
+        self.contain_doc = self._check_contain_doc()
+        if not self.contain_doc:
+            self.corpus = load_corpus(self.corpus_path)
+        self.max_process_num = 8
+
+    def _check_contain_doc(self):
+        return self.searcher.doc(0).raw() is not None
+
+    def _search(self, query: str, num: int = None, return_score: bool = False):
+        if num is None:
+            num = self.topk
+        hits = self.searcher.search(query, num)
+        if len(hits) < 1:
+            if return_score:
+                return [], []
+            else:
+                return []
+        scores = [hit.score for hit in hits]
+        if len(hits) < num:
+            warnings.warn("Not enough documents retrieved!", stacklevel=2)
+        else:
+            hits = hits[:num]
+
+        if self.contain_doc:
+            all_contents = [json.loads(self.searcher.doc(hit.docid).raw())["contents"] for hit in hits]
+            results = [
+                {
+                    "title": content.split("\n")[0].strip('"'),
+                    "text": "\n".join(content.split("\n")[1:]),
+                    "contents": content,
+                }
+                for content in all_contents
+            ]
+        else:
+            results = load_docs(self.corpus, [hit.docid for hit in hits])
+
+        if return_score:
+            return results, scores
+        else:
+            return results
+
+    def _batch_search(self, query_list: list[str], num: int = None, return_score: bool = False):
+        results = []
+        scores = []
+        for query in query_list:
+            item_result, item_score = self._search(query, num, True)
+            results.append(item_result)
+            scores.append(item_score)
+        if return_score:
+            return results, scores
+        else:
+            return results
+
+
+class DenseRetriever(BaseRetriever):
+    def __init__(self, config):
+        super().__init__(config)
+        self.index = faiss.read_index(self.index_path)
+        if config.faiss_gpu:
+            co = faiss.GpuMultipleClonerOptions()
+            co.useFloat16 = True
+            co.shard = True
+            self.index = faiss.index_cpu_to_all_gpus(self.index, co=co)
+
+        self.corpus = load_corpus(self.corpus_path)
+        self.encoder = Encoder(
+            model_name=self.retrieval_method,
+            model_path=config.retrieval_model_path,
+            pooling_method=config.retrieval_pooling_method,
+            max_length=config.retrieval_query_max_length,
+            use_fp16=config.retrieval_use_fp16,
+        )
+        self.topk = config.retrieval_topk
+        self.batch_size = config.retrieval_batch_size
+
+    def _search(self, query: str, num: int = None, return_score: bool = False):
+        if num is None:
+            num = self.topk
+        query_emb = self.encoder.encode(query)
+        scores, idxs = self.index.search(query_emb, k=num)
+        idxs = idxs[0]
+        scores = scores[0]
+        results = load_docs(self.corpus, idxs)
+        if return_score:
+            return results, scores.tolist()
+        else:
+            return results
+
+    def _batch_search(self, query_list: list[str], num: int = None, return_score: bool = False):
+        if isinstance(query_list, str):
+            query_list = [query_list]
+        if num is None:
+            num = self.topk
+
+        results = []
+        scores = []
+        for start_idx in tqdm(range(0, len(query_list), self.batch_size), desc="Retrieval process: "):
+            query_batch = query_list[start_idx : start_idx + self.batch_size]
+            batch_emb = self.encoder.encode(query_batch)
+            batch_scores, batch_idxs = self.index.search(batch_emb, k=num)
+            batch_scores = batch_scores.tolist()
+            batch_idxs = batch_idxs.tolist()
+
+            # load_docs is not vectorized, but is a python list approach
+            flat_idxs = sum(batch_idxs, [])
+            batch_results = load_docs(self.corpus, flat_idxs)
+            # chunk them back
+            batch_results = [batch_results[i * num : (i + 1) * num] for i in range(len(batch_idxs))]
+
+            results.extend(batch_results)
+            scores.extend(batch_scores)
+
+            del batch_emb, batch_scores, batch_idxs, query_batch, flat_idxs, batch_results
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+
+        if return_score:
+            return results, scores
+        else:
+            return results
+
+
+def get_retriever(config):
+    if config.retrieval_method == "bm25":
+        return BM25Retriever(config)
+    else:
+        return DenseRetriever(config)
+
+
+#####################################
+# FastAPI server below
+#####################################
+
+
+class Config:
+    """
+    Minimal config class (simulating your argparse)
+    Replace this with your real arguments or load them dynamically.
+    """
+
+    def __init__(
+        self,
+        retrieval_method: str = "bm25",
+        retrieval_topk: int = 10,
+        index_path: str = "./index/bm25",
+        corpus_path: str = "./data/corpus.jsonl",
+        dataset_path: str = "./data",
+        data_split: str = "train",
+        faiss_gpu: bool = True,
+        retrieval_model_path: str = "./model",
+        retrieval_pooling_method: str = "mean",
+        retrieval_query_max_length: int = 256,
+        retrieval_use_fp16: bool = False,
+        retrieval_batch_size: int = 128,
+    ):
+        self.retrieval_method = retrieval_method
+        self.retrieval_topk = retrieval_topk
+        self.index_path = index_path
+        self.corpus_path = corpus_path
+        self.dataset_path = dataset_path
+        self.data_split = data_split
+        self.faiss_gpu = faiss_gpu
+        self.retrieval_model_path = retrieval_model_path
+        self.retrieval_pooling_method = retrieval_pooling_method
+        self.retrieval_query_max_length = retrieval_query_max_length
+        self.retrieval_use_fp16 = retrieval_use_fp16
+        self.retrieval_batch_size = retrieval_batch_size
+
+
+class QueryRequest(BaseModel):
+    queries: list[str]
+    topk: int | None = None
+    return_scores: bool = False
+
+
+app = FastAPI()
+
+
+@app.post("/retrieve")
+def retrieve_endpoint(request: QueryRequest):
+    """
+    Endpoint that accepts queries and performs retrieval.
+
+    Input format:
+    {
+      "queries": ["What is Python?", "Tell me about neural networks."],
+      "topk": 3,
+      "return_scores": true
+    }
+
+    Output format (when return_scores=True，similarity scores are returned):
+    {
+        "result": [
+            [   # Results for each query
+                {
+                    {"document": doc, "score": score}
+                },
+                # ... more documents
+            ],
+            # ... results for other queries
+        ]
+    }
+    """
+    if not request.topk:
+        request.topk = config.retrieval_topk  # fallback to default
+
+    # Perform batch retrieval
+    tmp = retriever.batch_search(query_list=request.queries, num=request.topk, return_score=request.return_scores)
+
+    scores = []
+    try:
+        results, scores = tmp
+    except ValueError:
+        results = tmp
+
+    # Format response
+    resp = []
+    for i, single_result in enumerate(results):
+        if scores:
+            # If scores are returned, combine them with results
+            combined = []
+            for doc, score in zip(single_result, scores[i], strict=True):
+                combined.append({"document": doc, "score": score})
+            resp.append(combined)
+        else:
+            resp.append(single_result)
+    return {"result": resp}
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Launch the local faiss retriever.")
+    parser.add_argument(
+        "--index_path",
+        type=str,
+        default="./e5_Flat.index",
+        help="Corpus indexing file.",
+    )
+    parser.add_argument(
+        "--corpus_path",
+        type=str,
+        default="./wiki-18.jsonl",
+        help="Local corpus file.",
+    )
+    parser.add_argument("--topk", type=int, default=3, help="Number of retrieved passages for one query.")
+    parser.add_argument("--retriever_name", type=str, default="e5", help="Name of the retriever model.")
+    parser.add_argument(
+        "--retriever_model", type=str, default="intfloat/e5-base-v2", help="Path of the retriever model."
+    )
+    parser.add_argument("--faiss_gpu", action="store_true", help="Use GPU for computation")
+
+    args = parser.parse_args()
+
+    # 1) Build a config (could also parse from arguments).
+    #    In real usage, you'd parse your CLI arguments or environment variables.
+    config = Config(
+        retrieval_method=args.retriever_name,  # or "dense"
+        index_path=args.index_path,
+        corpus_path=args.corpus_path,
+        retrieval_topk=args.topk,
+        faiss_gpu=args.faiss_gpu,
+        retrieval_model_path=args.retriever_model,
+        retrieval_pooling_method="mean",
+        retrieval_query_max_length=256,
+        retrieval_use_fp16=True,
+        retrieval_batch_size=512,
+    )
+
+    # 2) Instantiate a global retriever so it is loaded once and reused.
+    retriever = get_retriever(config)
+
+    # 3) Launch the server. By default, it listens on http://127.0.0.1:8000
+    uvicorn.run(app, host="0.0.0.0", port=8000)
diff --git a/examples/search-r1/local_search_server.py b/examples/search-r1/local_search_server.py
new file mode 100644
index 000000000..4771556a0
--- /dev/null
+++ b/examples/search-r1/local_search_server.py
@@ -0,0 +1,97 @@
+"""
+Local Search Server for Search-R1
+
+This module provides a local search engine interface that mimics the google_search_server.py API.
+It sends requests to a local retrieval server (e.g., running retrieval_server.py from Search-R1)
+and formats the results to match the expected output format.
+
+Usage:
+    In your generate_with_search.py, replace:
+        from google_search_server import google_search
+    with:
+        from local_search_server import local_search as google_search
+
+    And update SEARCH_R1_CONFIGS:
+        SEARCH_R1_CONFIGS = {
+            "search_url": "http://127.0.0.1:8000/retrieve",  # URL of local retrieval server
+            "topk": 3,
+            ...
+        }
+"""
+
+import aiohttp
+
+
+async def local_search(
+    search_url: str,
+    query: str,
+    top_k: int = 5,
+    timeout: int = 1200,
+    proxy: str | None = None,
+) -> list[dict]:
+    """
+    Call local search engine server and format results to match google_search_server.py output.
+
+    This function provides the same interface as google_search() from google_search_server.py,
+    making it a drop-in replacement. The only difference is that instead of using an API key,
+    it uses a search_url parameter.
+
+    Args:
+        search_url: URL of the local retrieval server (e.g., "http://127.0.0.1:8000/retrieve")
+        query: Search query string
+        top_k: Number of results to retrieve
+        timeout: Request timeout in seconds (default: 60)
+        proxy: Proxy URL if needed (not used for local retrieval, kept for API compatibility)
+        snippet_only: If True, only return snippet (kept for API compatibility with google_search)
+
+    Returns:
+        List of dictionaries with format: [{"document": {"contents": '"<title>"\n<text>'}}]
+        This matches the output format of google_search() from google_search_server.py
+    """
+    # Prepare request payload for local retrieval server
+    payload = {
+        "queries": [query],
+        "topk": top_k,
+        "return_scores": False,  # We don't need scores for compatibility with google_search_server
+    }
+
+    # Send async request to local retrieval server
+    timeout_obj = aiohttp.ClientTimeout(total=timeout)
+    session_kwargs = {}
+    # Note: proxy parameter is kept for API compatibility but typically not needed for local server
+    if proxy:
+        session_kwargs["proxy"] = proxy
+
+    try:
+        async with aiohttp.ClientSession(**session_kwargs) as session:
+            async with session.post(search_url, json=payload, timeout=timeout_obj) as resp:
+                resp.raise_for_status()
+                result = await resp.json()
+    except Exception as e:
+        print(f"Error calling local search engine at {search_url}: {repr(e)}")
+        return []
+
+    # Parse retrieval results
+    # Format from retrieval_server.py: {"result": [[{"document": {"id": "...", "contents": "..."}}]]}
+    retrieval_results = result.get("result", [[]])[0]
+    # Format to match google_search_server.py output
+    # Google format: [{"document": {"contents": '"<title>"\n<context>'}}]
+    contexts = []
+
+    for item in retrieval_results:
+        # Extract contents from retrieval result
+        # retrieval_server returns: {"document": {"id": "...", "contents": '"Title"\nText...'}}
+        if isinstance(item, dict):
+            # Access the document dict first, then get contents
+            content = item.get("contents", "")
+
+            if content:
+                # The contents are already in the correct format: '"Title"\nText content...'
+                # Just pass through as-is to match google_search format
+                contexts.append({"document": {"contents": content}})
+            else:
+                # Empty content case - provide default values
+                contexts.append({"document": {"contents": '"No title."\nNo snippet available.'}})
+
+    # If no results found, return empty list (consistent with google_search_server.py)
+    return contexts
diff --git a/examples/search-r1/qa_em_format.py b/examples/search-r1/qa_em_format.py
new file mode 100644
index 000000000..afd9049b9
--- /dev/null
+++ b/examples/search-r1/qa_em_format.py
@@ -0,0 +1,208 @@
+# Adapt from https://github.com/PeterGriffinJin/Search-R1/blob/ceee7b89655ed52f205b9beb98e1190c3eedcfb0/verl/utils/reward_score/qa_em_format.py
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import random
+import re
+import string
+
+
+def normalize_answer(s):
+    def remove_articles(text):
+        return re.sub(r"\b(a|an|the)\b", " ", text)
+
+    def white_space_fix(text):
+        return " ".join(text.split())
+
+    def remove_punc(text):
+        exclude = set(string.punctuation)
+        return "".join(ch for ch in text if ch not in exclude)
+
+    def lower(text):
+        return text.lower()
+
+    return white_space_fix(remove_articles(remove_punc(lower(s))))
+
+
+def em_check(prediction, golden_answers):
+    if isinstance(golden_answers, str):
+        golden_answers = [golden_answers]
+    normalized_prediction = normalize_answer(prediction)
+    score = 0
+    for golden_answer in golden_answers:
+        golden_answer = normalize_answer(golden_answer)
+        if golden_answer == normalized_prediction:
+            score = 1
+            break
+    return score
+
+
+def is_valid_sequence(text):
+    # Find the position of "<|im_start|>assistant" with potential whitespace
+    assistant_pattern = r"<\|im_start\|>assistant\s*"
+    assistant_match = re.search(assistant_pattern, text)
+
+    if not assistant_match:
+        return False, "Missing assistant marker"
+
+    # Extract the content after the assistant marker
+    start_pos = assistant_match.end()
+    content = text[start_pos:]
+
+    # Check for balanced tags
+    tags_to_check = ["think", "search", "information", "answer"]
+    for tag in tags_to_check:
+        opening_count = len(re.findall(f"<{tag}>", content))
+        closing_count = len(re.findall(f"</{tag}>", content))
+        if opening_count != closing_count:
+            return False, f"Mismatch in {tag} tags: {opening_count} opening vs {closing_count} closing tags"
+
+    # Now check for proper sequence pattern and no extraneous content
+
+    # 1. First split the content by any tags we recognize
+    split_pattern = r"(</?(?:think|search|information|answer)>)"
+    parts = re.split(split_pattern, content)
+
+    # 2. Keep track of the current position in the expected sequence
+    state = "start"  # start -> think -> search -> information -> think -> ... -> answer -> end
+
+    # 3. Check each part
+    for _i, part in enumerate(parts):
+        # Skip empty parts
+        if not part.strip():
+            continue
+
+        # Check if this is a tag
+        if re.match(r"</?(?:think|search|information|answer)>", part):
+            # This is a tag, check if it's valid in the current state
+            if part == "<think>" and state in ["start", "information"]:
+                state = "in_think"
+            elif part == "</think>" and state == "in_think":
+                state = "after_think"
+            elif part == "<search>" and state == "after_think":
+                state = "in_search"
+            elif part == "</search>" and state == "in_search":
+                state = "after_search"
+            elif part == "<information>" and state == "after_search":
+                state = "in_information"
+            elif part == "</information>" and state == "in_information":
+                state = "information"
+            elif part == "<answer>" and state == "after_think":
+                state = "in_answer"
+            elif part == "</answer>" and state == "in_answer":
+                state = "end"
+            else:
+                return False, f"Unexpected tag {part} in state {state}"
+        else:
+            # This is content, check if it's valid in the current state
+            if state in ["in_think", "in_search", "in_information", "in_answer"]:
+                # Content is allowed inside tags
+                pass
+            elif state in ["start", "after_think", "after_search", "information"]:
+                # Only whitespace is allowed between tags
+                if part.strip():
+                    return False, f"Unexpected content '{part.strip()}' between tags (state: {state})"
+            else:
+                return False, f"Unexpected content in state {state}"
+
+    # Check final state
+    if state != "end":
+        return False, f"Incomplete sequence, ended in state {state}"
+
+    return True, "Valid sequence format"
+
+
+def extract_solution(solution_str):
+    """Extract the equation from the solution string."""
+
+    answer_pattern = r"<answer>(.*?)</answer>"
+    match = re.finditer(answer_pattern, solution_str, re.DOTALL)
+    matches = list(match)
+
+    # If there are 0 or exactly 1 matches, return None
+    if len(matches) <= 1:
+        return None
+
+    # If there are 2 or more matches, return the last one
+    return matches[-1].group(1).strip()
+
+
+def extract_information_blocks(text: str) -> list[str]:
+    pattern = r"<information>(.*?)</information>"
+    matches = re.findall(pattern, text, re.DOTALL)
+    return [match.strip() for match in matches]
+
+
+def is_retrieval_correct(text: str, golden_answers: list[str]) -> bool:
+    seqs = extract_information_blocks(text)
+    for seq in seqs:
+        for golden_answer in golden_answers:
+            if normalize_answer(golden_answer) in normalize_answer(seq):
+                return True
+    return False
+
+
+def compute_score_em(
+    solution_str,
+    ground_truth,
+    method="strict",
+    structure_format_score=0,
+    final_format_score=0,
+    retrieval_score=0,
+    format_score=0,
+    score=1.0,
+):
+    """The scoring function for exact match (EM).
+
+    Args:
+        solution_str: the solution text
+        ground_truth: the ground truth
+        method: the method to extract the solution, choices are 'strict' and 'flexible'
+        format_score: the score for the format
+        score: the score for the correct answer
+    """
+    is_valid_format, _ = is_valid_sequence(solution_str)
+    retrieval_correct = False
+    if is_valid_format:
+        retrieval_correct = is_retrieval_correct(solution_str, ground_truth["target"])
+    answer = extract_solution(solution_str=solution_str)
+    do_print = random.randint(1, 64) == 1
+
+    if do_print:
+        print("--------------------------------")
+        print(f"Golden answers: {ground_truth['target']}")
+        print(f"Extracted answer: {answer}")
+        print(f"Solution string: {solution_str}")
+
+    if answer is None:
+        if is_valid_format:
+            if retrieval_correct:
+                return structure_format_score + retrieval_score  # 0.3
+            else:
+                return structure_format_score  # 0.2
+        else:
+            return 0
+    else:
+        if em_check(answer, ground_truth["target"]):
+            if is_valid_format:
+                return score  # 1
+            else:
+                return score - structure_format_score  # 0.8
+        elif is_valid_format:
+            if retrieval_correct:
+                return structure_format_score + retrieval_score  # 0.3
+            else:
+                return structure_format_score  # 0.2
+        else:
+            return final_format_score  # 0.1
diff --git a/examples/search-r1/run_qwen3_4b_npu.sh b/examples/search-r1/run_qwen3_4b_npu.sh
new file mode 100644
index 000000000..8f2aa6dea
--- /dev/null
+++ b/examples/search-r1/run_qwen3_4b_npu.sh
@@ -0,0 +1,213 @@
+#!/bin/bash
+set -ex
+ulimit -u 65535
+
+# cleanup
+pkill -9 -f "vllm serve" 2>/dev/null || true
+sleep 2
+npu-smi info 2>/dev/null | grep rayWorker | awk '{print $4}' | xargs -r kill -9 2>/dev/null || true
+sleep 3
+
+# Ray isolation: independent temp-dir, ports, and cleanup
+export RAY_TMPDIR=/tmp/ray_vime_npu_search_r1
+export RAY_PORT=6379
+export RAY_DASHBOARD_PORT=8265
+export RAY_AGENT_PORT=52378
+unset RAY_ADDRESS RAY_REDIS_ADDRESS
+
+ray stop --force 2>/dev/null || true
+rm -rf "${RAY_TMPDIR}"
+sleep 2
+
+project_name="vime"
+exp_name="qwen3-4b-search-r1"
+RAY_DATA_HOME=${RAY_DATA_HOME:-"/root/logs"}
+start_time=$(date +"%Y%m%d_%H%M%S")
+LOG_DIR=${LOG_DIR:-"${RAY_DATA_HOME}/${project_name}/${exp_name}"}
+mkdir -p "${LOG_DIR}"
+LOG_FILE="${LOG_DIR}/${start_time}.log"
+
+echo "Experiment Log will be saved to: ${LOG_FILE}"
+VIME_DIR="/root/vime"
+
+# NPU environment
+source /usr/local/Ascend/driver/bin/setenv.bash
+source /usr/local/Ascend/ascend-toolkit/set_env.sh
+source /usr/local/Ascend/nnal/atb/set_env.sh
+export PYTHONPATH="${VIME_DIR}/examples/search-r1:/root/Megatron-LM:/root/vllm:/root/vllm-ascend:${VIME_DIR}:/root/Megatron-Bridge:/root/mbridge:/root/MegatronAdaptor:/root/TransformerEngineNPU:/usr/local/Ascend/ascend-toolkit/latest/python/site-packages:/usr/local/Ascend/ascend-toolkit/latest/tools/ms_fmk_transplt/torch_npu_bridge:${PYTHONPATH}"
+export PYTHONUNBUFFERED=1
+export PYTORCH_NPU_ALLOC_CONF=expandable_segments:False
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15
+export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050
+export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050
+export HCCL_CONNECT_TIMEOUT=7200
+export HCCL_DETERMINISTIC=true
+export VLLM_ASCEND_ENABLE_NZ=0
+export ASCEND_COREDUMP_SIGNAL=None
+export ATB_MATMUL_SHUFFLE_K_ENABLE=0
+export ATB_LLM_LCOC_ENABLE=0
+export TASK_QUEUE_ENABLE=1
+export RAY_DISABLE_SIGINT_OVERRIDE=1
+export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1
+export LD_LIBRARY_PATH=/usr/local/Ascend/driver/lib64:/usr/local/Ascend/ascend-toolkit/latest/lib64:/usr/local/Ascend/nnal/atb/latest/atb/cxx_abi_1/lib:/usr/local/Ascend/cann/lib64:${LD_LIBRARY_PATH}
+export VLLM_USE_AOT_COMPILE=0
+export VLLM_DISABLE_COMPILE_CACHE=1
+export TRANSFORMERS_VERBOSITY=error
+export RUST_LOG=vllm_router_rs=warn
+
+NUM_NPUS=16
+source "${VIME_DIR}/scripts/models/qwen3-4B-Instruct-2507.sh"
+
+CKPT_ARGS=(
+   --hf-checkpoint /path/to/Qwen3-4B-Instruct-2507
+   --ref-load  /path/to/Qwen3-4B-Instruct-2507
+   --load /path/to/Qwen3-4B-Instruct-2507_vime_npu/
+   --save /path/to/Qwen3-4B-Instruct-2507_vime_npu/
+   --save-interval 100
+   --no-load-optim
+   --megatron-to-hf-mode bridge
+)
+
+ROLLOUT_ARGS=(
+   --prompt-data /path/to/nq_hotpotqa_train/train.parquet
+   --input-key prompt
+   --label-key reward_model
+   --apply-chat-template
+   --rollout-shuffle
+   --num-rollout 200
+   --rollout-batch-size 32
+   --n-samples-per-prompt 8
+   --rollout-max-response-len 512
+   --rollout-temperature 1
+   --global-batch-size 256
+   --balance-data
+)
+
+EVAL_ARGS=(
+   --eval-interval 50
+   --eval-prompt-data nq_test /path/to/nq_hotpotqa_train/test.parquet@[0:500]
+   --n-samples-per-eval-prompt 1
+   --eval-input-key prompt
+   --eval-label-key reward_model
+   --eval-top-p 1
+)
+
+PERF_ARGS=(
+   --tensor-model-parallel-size 4
+   --pipeline-model-parallel-size 1
+   --context-parallel-size 1
+   --expert-model-parallel-size 1
+   --expert-tensor-parallel-size 1
+
+   --recompute-granularity full
+   --recompute-method uniform
+   --recompute-num-layers 1
+
+   --micro-batch-size 1
+   --use-dynamic-batch-size
+   --max-tokens-per-gpu 9216
+)
+
+GRPO_ARGS=(
+   --advantage-estimator grpo
+   --use-kl-loss
+   --kl-loss-coef 0.001
+   --kl-loss-type low_var_kl
+   --entropy-coef 0.00
+   --eps-clip 0.2
+   --eps-clip-high 0.28
+
+   # TIS (Trajectory Importance Sampling)
+   # --use-tis
+)
+
+OPTIMIZER_ARGS=(
+   --optimizer adam
+   --lr 1e-6
+   --lr-decay-style constant
+   --weight-decay 0.1
+   --adam-beta1 0.9
+   --adam-beta2 0.98
+   --optimizer-cpu-offload
+   --overlap-cpu-optimizer-d2h-h2d
+   --use-precision-aware-optimizer
+)
+
+VLLM_ARGS=(
+   --rollout-num-gpus-per-engine 4
+   --vllm-gpu-memory-utilization 0.7
+   --vllm-enable-sleep-mode
+   --vllm-weight-sync-mode native
+)
+
+MISC_ARGS=(
+   --attention-dropout 0.0
+   --hidden-dropout 0.0
+   --accumulate-allreduce-grads-in-fp32
+   --attention-softmax-in-fp32
+   --attention-backend flash
+   --use-flash-attn
+)
+
+CUSTOM_ARGS=(
+   --custom-generate-function-path generate_with_search.generate
+   --custom-rm-path generate_with_search.reward_func
+
+   # TIS (Trajectory Importance Sampling)
+   # --custom-config-path examples/train_infer_mismatch_helper/mis.yaml
+   # --custom-tis-function-path examples.train_infer_mismatch_helper.mis.compute_mis_weights_with_cp
+)
+
+# launch the master node of ray in container
+unset https_proxy http_proxy proxy
+ray start --head \
+    --temp-dir="${RAY_TMPDIR}" \
+    --port="${RAY_PORT}" \
+    --dashboard-port="${RAY_DASHBOARD_PORT}" \
+    --dashboard-agent-listen-port="${RAY_AGENT_PORT}" \
+    --node-ip-address 127.0.0.1 \
+    --num-gpus 0 \
+    --resources "{\"NPU\": $NUM_NPUS}" \
+    --disable-usage-stats \
+    --dashboard-host=0.0.0.0
+
+# Build the runtime environment JSON with proper variable substitution
+RUNTIME_ENV_JSON=$(cat << 'EOF'
+{
+  "env_vars": {
+    "PYTHONPATH": "${VIME_DIR}/examples/search-r1:/root/Megatron-LM:/root/vllm:/root/vllm-ascend:${VIME_DIR}:/root/Megatron-Bridge:/root/mbridge:/root/MegatronAdaptor:/root/TransformerEngineNPU:/usr/local/Ascend/ascend-toolkit/latest/python/site-packages:/usr/local/Ascend/ascend-toolkit/latest/tools/ms_fmk_transplt/torch_npu_bridge",
+    "CUDA_DEVICE_MAX_CONNECTIONS": "1",
+    "HCCL_HOST_SOCKET_PORT_RANGE": "60000-60050",
+    "HCCL_NPU_SOCKET_PORT_RANGE": "61000-61050",
+    "HCCL_CONNECT_TIMEOUT": "7200",
+    "PYTORCH_NPU_ALLOC_CONF": "expandable_segments:False",
+    "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES": "1",
+    "VLLM_DISABLE_COMPILE_CACHE": "1",
+    "TRANSFORMERS_VERBOSITY": "error",
+    "RUST_LOG": "vllm_router_rs=warn",
+    "LD_LIBRARY_PATH": "/usr/local/Ascend/driver/lib64:/usr/local/Ascend/driver/lib64/driver:/usr/local/Ascend/driver/lib64/common:/usr/local/Ascend/ascend-toolkit/latest/lib64:/usr/local/Ascend/ascend-toolkit/latest/compiler/lib64/plugin/opskernel:/usr/local/Ascend/ascend-toolkit/latest/compiler/lib64/plugin/nnengine:/usr/local/Ascend/ascend-toolkit/latest/opp/built-in/op_impl/ai_core/tbe/op_tiling/lib/:/usr/local/Ascend/nnal/atb/latest/atb/cxx_abi_1/lib:/usr/local/Ascend/cann/lib64:/usr/local/Ascend/cann/aarch64-linux/devlib"
+  }
+}
+EOF
+)
+
+ray job submit --address="http://127.0.0.1:${RAY_DASHBOARD_PORT}" \
+   --runtime-env-json="${RUNTIME_ENV_JSON}" \
+   --working-dir="${VIME_DIR}" \
+   -- python3 -u train.py \
+   --train-backend megatron \
+   --actor-num-nodes 1 \
+   --actor-num-gpus-per-node 8 \
+   --rollout-num-gpus 8 \
+   ${MODEL_ARGS[@]} \
+   ${CKPT_ARGS[@]} \
+   ${ROLLOUT_ARGS[@]} \
+   ${OPTIMIZER_ARGS[@]} \
+   ${GRPO_ARGS[@]} \
+   ${PERF_ARGS[@]} \
+   ${EVAL_ARGS[@]} \
+   ${VLLM_ARGS[@]} \
+   ${MISC_ARGS[@]} \
+   ${CUSTOM_ARGS[@]} \
+   2>&1 | tee "${LOG_FILE}"
diff --git a/examples/tau-bench/run_qwen3_4B_npu.sh b/examples/tau-bench/run_qwen3_4B_npu.sh
new file mode 100644
index 000000000..c5bad23fc
--- /dev/null
+++ b/examples/tau-bench/run_qwen3_4B_npu.sh
@@ -0,0 +1,148 @@
+#!/bin/bash
+
+if grep -q $'\r' "$0" 2>/dev/null; then
+  exec bash <(sed 's/\r$//' "$0") "$@"
+fi
+
+# for rerun the task
+pkill -9 vllm 2>/dev/null || true
+pkill -9 VLLM 2>/dev/null || true
+sleep 3
+ray stop --force 2>/dev/null || true
+pkill -9 ray 2>/dev/null || true
+pkill -9 -f 'python3 train.py' 2>/dev/null || true
+sleep 3
+
+set -ex
+
+export PYTHONUNBUFFERED=1
+export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15
+export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050
+export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050
+export HYDRA_FULL_ERROR=1
+export VLLM_ASCEND_ENABLE_NZ=0
+export VLLM_USE_AOT_COMPILE=0
+export VIME_VLLM_SERVER_HEALTH_TIMEOUT_SEC=900
+
+unset PYTORCH_CUDA_ALLOC_CONF PYTORCH_ALLOC_CONF
+unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY
+
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
+source "${SCRIPT_DIR}/../../scripts/models/qwen3-4B-Instruct-2507.sh"
+
+export PYTHONPATH="${SCRIPT_DIR}:/root/Megatron-Bridge/src:/root/Megatron-LM:${PYTHONPATH:-}"
+
+DATA_ROOT="${DATA_ROOT:-/root}"
+TAU_BENCH_ROOT="${TAU_BENCH_ROOT:-/root/tau-bench}"
+
+CKPT_ARGS=(
+   --hf-checkpoint ${DATA_ROOT}/weights/Qwen3-4B-Instruct-2507/
+   --load ${DATA_ROOT}/weights/Qwen3-4B-Instruct-2507/
+   --ref-load ${DATA_ROOT}/weights/Qwen3-4B-Instruct-2507/
+   --megatron-to-hf-mode bridge
+)
+
+ROLLOUT_ARGS=(
+   --prompt-data ${TAU_BENCH_ROOT}/retail_train_tasks.jsonl
+   --input-key index
+   --rollout-shuffle
+   --num-rollout 500
+   --rollout-batch-size 32
+   --n-samples-per-prompt 8
+   --rollout-max-response-len 4096
+   --rollout-max-context-len 16384
+   --rollout-temperature 0.7
+   --global-batch-size 256
+   --dynamic-sampling-filter-path vime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std
+   --balance-data
+)
+
+EVAL_ARGS=(
+   --eval-interval 5
+   --eval-prompt-data retail-dev ${TAU_BENCH_ROOT}/retail_dev_tasks.jsonl
+   --n-samples-per-eval-prompt 1
+   --eval-max-response-len 4096
+   --eval-top-k 1
+)
+
+PERF_ARGS=(
+   --tensor-model-parallel-size 2
+   --sequence-parallel
+   --pipeline-model-parallel-size 1
+   --context-parallel-size 1
+   --expert-model-parallel-size 1
+   --expert-tensor-parallel-size 1
+   --recompute-granularity full
+   --recompute-method uniform
+   --recompute-num-layers 1
+   --use-dynamic-batch-size
+   --max-tokens-per-gpu 9216
+)
+
+GRPO_ARGS=(
+   --advantage-estimator grpo
+   --use-kl-loss
+   --kl-loss-coef 0.001
+   --kl-loss-type low_var_kl
+   --entropy-coef 0.01
+   --eps-clip 0.2
+   --eps-clip-high 0.28
+)
+
+OPTIMIZER_ARGS=(
+   --optimizer adam
+   --lr 5e-6
+   --lr-decay-style constant
+   --weight-decay 0.1
+   --adam-beta1 0.9
+   --adam-beta2 0.98
+)
+
+VLLM_ARGS=(
+   --rollout-num-gpus-per-engine 1
+   --vllm-gpu-memory-utilization 0.7
+   --vllm-max-model-len 16384
+)
+
+MISC_ARGS=(
+   --attention-dropout 0.0
+   --hidden-dropout 0.0
+   --accumulate-allreduce-grads-in-fp32
+   --attention-softmax-in-fp32
+   --attention-backend flash
+   --use-flash-attn
+   --no-gradient-accumulation-fusion
+)
+
+CUSTOM_ARGS=(
+   --custom-generate-function-path generate_with_tau.generate
+   --custom-rm-path generate_with_tau.batched_tau_bench_rm
+)
+
+export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
+
+ray start --head \
+  --node-ip-address "${MASTER_ADDR}" \
+  --disable-usage-stats \
+  --dashboard-host=0.0.0.0 \
+  --dashboard-port=8265
+
+ray job submit --address="http://127.0.0.1:8265" \
+   -- python3 train.py \
+   --train-backend megatron \
+   --actor-num-nodes 1 \
+   --actor-num-gpus-per-node 4 \
+   --rollout-num-gpus 4 \
+   "${MODEL_ARGS[@]}" \
+   "${CKPT_ARGS[@]}" \
+   "${ROLLOUT_ARGS[@]}" \
+   "${EVAL_ARGS[@]}" \
+   "${OPTIMIZER_ARGS[@]}" \
+   "${GRPO_ARGS[@]}" \
+   "${PERF_ARGS[@]}" \
+   "${VLLM_ARGS[@]}" \
+   "${CUSTOM_ARGS[@]}" \
+   "${MISC_ARGS[@]}"
+
diff --git a/scripts/models/qwen3-30B-A3B-npu.sh b/scripts/models/qwen3-30B-A3B-npu.sh
new file mode 100644
index 000000000..c4512cd4e
--- /dev/null
+++ b/scripts/models/qwen3-30B-A3B-npu.sh
@@ -0,0 +1,48 @@
+NLAYERS=48
+FIRST_K_DENSE_REPLACE=0
+
+arr=()
+for ((i=0; i<NLAYERS; i++)); do
+  if (( i < FIRST_K_DENSE_REPLACE )); then
+    arr+=(0)
+  else
+    arr+=(1)
+  fi
+done
+
+printf -v MOE_LAYER_FREQ "[%s]" "$(IFS=', '; echo "${arr[*]}")"
+
+
+MODEL_ARGS=(
+   --disable-bias-linear
+   --qk-layernorm
+   --group-query-attention
+   --num-attention-heads 32
+   --num-query-groups 4
+   --kv-channels 128
+   --num-layers 48
+   --hidden-size 2048
+   --ffn-hidden-size 6144
+
+   --normalization RMSNorm
+   --position-embedding-type rope
+   --norm-epsilon 1e-6
+   --rotary-percent 1.0
+   --swiglu
+   --untie-embeddings-and-output-weights
+   --vocab-size 151936
+
+   --rotary-base "${MODEL_ARGS_ROTARY_BASE:-1000000}"
+
+   # moe
+   --moe-ffn-hidden-size 768
+   --moe-router-score-function softmax
+   --moe-token-dispatcher-type alltoall
+   --moe-router-topk 8
+   --moe-layer-freq "$MOE_LAYER_FREQ"
+   --num-experts 128
+   --moe-grouped-gemm
+   --moe-token-drop-policy probs
+   --moe-router-dtype fp32
+   --moe-aux-loss-coeff 0
+)
diff --git a/scripts/run-glm4.7-30B-A3B-npu.sh b/scripts/run-glm4.7-30B-A3B-npu.sh
new file mode 100644
index 000000000..9432812d0
--- /dev/null
+++ b/scripts/run-glm4.7-30B-A3B-npu.sh
@@ -0,0 +1,146 @@
+#!/bin/bash
+
+# for rerun the task
+pkill -9 -f '[v]llm serve|VLL[M]::'
+pkill -9 -f VLLM
+sleep 3
+ray stop --force
+pkill -9 ray
+pkill -9 python
+sleep 3
+pkill -9 ray
+pkill -9 python
+pkill -9 redis
+
+set -ex
+
+export PYTHONUNBUFFERED=1
+export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15
+export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050
+export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050
+export HYDRA_FULL_ERROR=1
+export DISABLE_L2_CACHE=1
+export VLLM_ASCEND_ENABLE_NZ=0
+export VLLM_USE_AOT_COMPILE=0
+export PYTHONPATH="/root/Megatron-LM:/root/MegatronAdaptor:/root/TransformerEngineNPU:${PYTHONPATH:-}"
+
+unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY
+
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
+source "${SCRIPT_DIR}/models/glm4.7-30B-A3B.sh"
+
+DATA_ROOT="${DATA_ROOT:-/root}"
+MODEL_DIR="${DATA_ROOT}/weights/GLM-4.7-Flash"
+DATASET_DIR="${DATA_ROOT}/datasets/dapo-math-17k"
+
+CKPT_ARGS=(
+   --hf-checkpoint "${MODEL_DIR}"
+   --load "${MODEL_DIR}"
+   --ref-load "${MODEL_DIR}"
+)
+
+ROLLOUT_ARGS=(
+   --prompt-data "${DATASET_DIR}/dapo-math-17k.jsonl"
+   --input-key prompt
+   --label-key label
+   --apply-chat-template
+   --rollout-shuffle
+   --rm-type deepscaler
+   --num-rollout 3000
+   --rollout-batch-size 32
+   --n-samples-per-prompt 8
+   --rollout-max-response-len 8192
+   --rollout-temperature 1
+   --global-batch-size 256
+   --balance-data
+)
+
+EVAL_ARGS=(
+   --eval-interval 20
+   --eval-prompt-data aime ${DATA_ROOT}/datasets/aime-2024/aime-2024.jsonl
+   --n-samples-per-eval-prompt 16
+   --eval-max-response-len 16384
+   --eval-top-p 1
+)
+
+PERF_ARGS=(
+   --tensor-model-parallel-size 4
+   --sequence-parallel
+   --pipeline-model-parallel-size 1
+   --context-parallel-size 1
+   --expert-model-parallel-size 8
+   --expert-tensor-parallel-size 1
+
+   --recompute-granularity full
+   --recompute-method uniform
+   --recompute-num-layers 1
+
+   --use-dynamic-batch-size
+   --max-tokens-per-gpu 20480
+   --seq-length 24576
+)
+
+GRPO_ARGS=(
+   --advantage-estimator grpo
+   --use-kl-loss
+   --kl-loss-coef 0.00
+   --kl-loss-type low_var_kl
+   --entropy-coef 0.00
+   --eps-clip 0.2
+   --eps-clip-high 0.28
+)
+
+OPTIMIZER_ARGS=(
+   --optimizer adam
+   --lr 1e-6
+   --lr-decay-style constant
+   --weight-decay 0.1
+   --adam-beta1 0.9
+   --adam-beta2 0.98
+   --optimizer-cpu-offload
+   --overlap-cpu-optimizer-d2h-h2d
+   --use-precision-aware-optimizer
+)
+
+
+VLLM_ARGS=(
+   --vllm-additional-config '{"weight_nz_mode":0}'
+   --rollout-num-gpus-per-engine 4
+   --vllm-gpu-memory-utilization 0.7
+   --vllm-enable-expert-parallel
+   --vllm-cudagraph-capture-sizes 1 2 4 8 $(seq 16 8 256)
+)
+
+MISC_ARGS=(
+   # Match GLM's unscaled RoPE without changing main's shared model script.
+   --rope-type rope
+   --attention-dropout 0.0
+   --hidden-dropout 0.0
+   --accumulate-allreduce-grads-in-fp32
+   --attention-softmax-in-fp32
+   --attention-backend flash
+
+   --use-flash-attn
+   --no-gradient-accumulation-fusion
+)
+
+# launch the master node of ray in container
+export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
+ray start --head --node-ip-address ${MASTER_ADDR} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265
+
+ray job submit --address="http://127.0.0.1:8265" \
+   -- python3 train.py \
+   --actor-num-nodes 1 \
+   --actor-num-gpus-per-node 8 \
+   --rollout-num-gpus 8 \
+   "${MODEL_ARGS[@]}" \
+   "${CKPT_ARGS[@]}" \
+   "${ROLLOUT_ARGS[@]}" \
+   "${OPTIMIZER_ARGS[@]}" \
+   "${GRPO_ARGS[@]}" \
+   "${PERF_ARGS[@]}" \
+   "${EVAL_ARGS[@]}" \
+   "${VLLM_ARGS[@]}" \
+   "${MISC_ARGS[@]}"
diff --git a/scripts/run-qwen3-30B-A3B-npu.sh b/scripts/run-qwen3-30B-A3B-npu.sh
new file mode 100644
index 000000000..915b4c8bd
--- /dev/null
+++ b/scripts/run-qwen3-30B-A3B-npu.sh
@@ -0,0 +1,137 @@
+#!/bin/bash
+
+# for rerun the task
+pkill -9 -f '[v]llm serve|VLL[M]::'
+sleep 3
+ray stop --force
+pkill -9 ray
+pkill -9 python
+sleep 3
+pkill -9 ray
+pkill -9 python
+pkill -9 redis
+
+set -ex
+
+export PYTHONUNBUFFERED=1
+export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15
+export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050
+export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050
+export HYDRA_FULL_ERROR=1
+export DISABLE_L2_CACHE=1
+export VLLM_ASCEND_ENABLE_NZ=0
+export VLLM_USE_AOT_COMPILE=0
+export PYTHONPATH="/root/Megatron-Bridge/src:/root/Megatron-LM/:${PYTHONPATH:-}"
+
+unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY
+
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
+source "${SCRIPT_DIR}/models/qwen3-30B-A3B.sh"
+
+DATA_ROOT="${DATA_ROOT:-/root}"
+
+CKPT_ARGS=(
+   --hf-checkpoint ${DATA_ROOT}/weights/Qwen3-30B-A3B/
+   --load ${DATA_ROOT}/weights/Qwen3-30B-A3B/
+   --ref-load ${DATA_ROOT}/weights/Qwen3-30B-A3B/
+   --megatron-to-hf-mode bridge
+)
+
+ROLLOUT_ARGS=(
+   --prompt-data ${DATA_ROOT}/datasets/dapo-math-17k/dapo-math-17k.jsonl
+   --input-key prompt
+   --label-key label
+   --apply-chat-template
+   --rollout-shuffle
+   --rm-type deepscaler
+   --num-rollout 3000
+   --rollout-batch-size 32
+   --n-samples-per-prompt 8
+   --rollout-max-response-len $((1024 * 8))
+   --rollout-temperature 1
+   --global-batch-size 256
+   --balance-data
+)
+
+EVAL_ARGS=(
+   --eval-interval 20
+   --eval-prompt-data aime ${DATA_ROOT}/datasets/aime-2024/aime-2024.jsonl
+   --n-samples-per-eval-prompt 16
+   --eval-max-response-len 16384
+   --eval-top-p 1
+)
+
+PERF_ARGS=(
+   --tensor-model-parallel-size 4
+   --sequence-parallel
+   --pipeline-model-parallel-size 1
+   --context-parallel-size 1
+   --expert-model-parallel-size 8
+   --expert-tensor-parallel-size 1
+
+   --recompute-granularity full
+   --recompute-method uniform
+   --recompute-num-layers 1
+
+   --use-dynamic-batch-size
+   --max-tokens-per-gpu 20480
+)
+
+GRPO_ARGS=(
+   --advantage-estimator grpo
+   --use-kl-loss
+   --kl-loss-coef 0.00
+   --kl-loss-type low_var_kl
+   --entropy-coef 0.00
+   --eps-clip 0.2
+   --eps-clip-high 0.28
+)
+
+OPTIMIZER_ARGS=(
+   --optimizer adam
+   --lr 1e-6
+   --lr-decay-style constant
+   --weight-decay 0.1
+   --adam-beta1 0.9
+   --adam-beta2 0.98
+
+   --optimizer-cpu-offload
+   --overlap-cpu-optimizer-d2h-h2d
+   --use-precision-aware-optimizer
+)
+
+VLLM_ARGS=(
+   --rollout-num-gpus-per-engine 4
+   --vllm-gpu-memory-utilization 0.7
+   --vllm-cudagraph-capture-sizes 1 2 4 8 $(seq 16 8 256)
+)
+
+MISC_ARGS=(
+   --attention-dropout 0.0
+   --hidden-dropout 0.0
+   --accumulate-allreduce-grads-in-fp32
+   --attention-softmax-in-fp32
+   --attention-backend flash
+
+   --use-flash-attn
+   --no-gradient-accumulation-fusion
+)
+
+ray start --head --node-ip-address 127.0.0.1 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265
+
+ray job submit --address="http://127.0.0.1:8265" \
+   -- python3 train.py \
+   --actor-num-nodes 1 \
+   --actor-num-gpus-per-node 8 \
+   --rollout-num-gpus 8 \
+   ${MODEL_ARGS[@]} \
+   ${CKPT_ARGS[@]} \
+   ${ROLLOUT_ARGS[@]} \
+   ${OPTIMIZER_ARGS[@]} \
+   ${GRPO_ARGS[@]} \
+   ${PERF_ARGS[@]} \
+   ${EVAL_ARGS[@]} \
+   ${VLLM_ARGS[@]} \
+   ${MISC_ARGS[@]}
diff --git a/scripts/run-qwen3-4B-npu.sh b/scripts/run-qwen3-4B-npu.sh
new file mode 100644
index 000000000..73b672866
--- /dev/null
+++ b/scripts/run-qwen3-4B-npu.sh
@@ -0,0 +1,124 @@
+#!/bin/bash
+
+# for rerun the task
+pkill -9 -f '[v]llm serve|VLL[M]::'
+pkill -9 -f VLLM
+sleep 3
+ray stop --force
+pkill -9 ray
+pkill -9 python
+sleep 3
+pkill -9 ray
+pkill -9 python
+pkill -9 redis
+
+set -ex
+
+export PYTHONUNBUFFERED=1
+export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15
+export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050
+export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050
+export HYDRA_FULL_ERROR=1
+export DISABLE_L2_CACHE=1
+export VLLM_ASCEND_ENABLE_NZ=0
+export VLLM_USE_AOT_COMPILE=0
+export PYTHONPATH="/root/Megatron-Bridge/src:/root/Megatron-LM/:${PYTHONPATH:-}"
+
+unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY
+
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
+source "${SCRIPT_DIR}/models/qwen3-4B.sh"
+
+DATA_ROOT="${DATA_ROOT:-/root}"
+
+CKPT_ARGS=(
+   --hf-checkpoint ${DATA_ROOT}/models/Qwen3-4B/
+   --load ${DATA_ROOT}/models/Qwen3-4B/
+   --ref-load ${DATA_ROOT}/models/Qwen3-4B/
+   --megatron-to-hf-mode bridge
+)
+
+ROLLOUT_ARGS=(
+   --prompt-data ${DATA_ROOT}/datasets/dapo-math-17k/dapo-math-17k.jsonl
+   --input-key prompt
+   --label-key label
+   --apply-chat-template
+   --rollout-shuffle
+   --rm-type math
+   --num-rollout 200
+   --rollout-batch-size 32
+   --n-samples-per-prompt 8
+   --rollout-max-response-len 2048
+   --rollout-temperature 1
+   --global-batch-size 256
+   --balance-data
+)
+
+PERF_ARGS=(
+   --tensor-model-parallel-size 4
+   --pipeline-model-parallel-size 1
+   --context-parallel-size 1
+   --expert-model-parallel-size 1
+   --expert-tensor-parallel-size 1
+   --recompute-granularity full
+   --recompute-method uniform
+   --recompute-num-layers 1
+   --use-dynamic-batch-size
+   --max-tokens-per-gpu 8192
+   --megatron-to-hf-mode bridge
+)
+
+GRPO_ARGS=(
+   --advantage-estimator grpo
+   --kl-loss-coef 0.0
+   --kl-loss-type low_var_kl
+   --kl-coef 0.00
+   --entropy-coef 0.0
+   --eps-clip 0.2
+   --eps-clip-high 0.28
+)
+
+OPTIMIZER_ARGS=(
+   --optimizer adam
+   --lr 1e-6
+   --lr-decay-style constant
+   --weight-decay 0.1
+   --adam-beta1 0.9
+   --adam-beta2 0.98
+   --optimizer-cpu-offload
+   --overlap-cpu-optimizer-d2h-h2d
+   --use-precision-aware-optimizer
+)
+
+VLLM_ARGS=(
+   --rollout-num-gpus-per-engine 4
+   --vllm-gpu-memory-utilization 0.6
+)
+
+MISC_ARGS=(
+   --attention-dropout 0.0
+   --hidden-dropout 0.0
+   --accumulate-allreduce-grads-in-fp32
+   --attention-softmax-in-fp32
+   --attention-backend flash
+   --micro-batch-size 1
+   --use-flash-attn
+)
+
+ray start --head --node-ip-address 127.0.0.1 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265
+
+ray job submit --address="http://127.0.0.1:8265" \
+-- python3 train.py \
+--actor-num-nodes 1 \
+--actor-num-gpus-per-node 4 \
+--rollout-num-gpus 4 \
+${MODEL_ARGS[@]} \
+${CKPT_ARGS[@]} \
+${ROLLOUT_ARGS[@]} \
+${OPTIMIZER_ARGS[@]} \
+${GRPO_ARGS[@]} \
+${PERF_ARGS[@]} \
+${VLLM_ARGS[@]} \
+${MISC_ARGS[@]}
diff --git a/tests/_unit_stubs.py b/tests/_unit_stubs.py
index 2f0bb62ba..827c615d2 100644
--- a/tests/_unit_stubs.py
+++ b/tests/_unit_stubs.py
@@ -35,9 +35,8 @@
 
 def real_module_available(name: str) -> bool:
     """True when the real package is importable and should not be shadowed."""
-    if name in sys.modules:
-        return True
     try:
+        # Earlier test modules may have installed a stub with no import spec.
         return importlib.util.find_spec(name) is not None
     except (ImportError, ValueError):
         return False
diff --git a/tests/test_empty_colocated_weight_bucket.py b/tests/test_empty_colocated_weight_bucket.py
index 14c8589dd..4162d3003 100644
--- a/tests/test_empty_colocated_weight_bucket.py
+++ b/tests/test_empty_colocated_weight_bucket.py
@@ -44,6 +44,9 @@ def _install_fake_deps(monkeypatch):
     accelerator_mod.current_device = lambda: "cuda:0"
     accelerator_mod.ipc_collect = lambda: None
 
+    platform_mod = types.ModuleType("vime.platforms")
+    platform_mod.current_platform = lambda: types.SimpleNamespace(is_npu=False)
+
     dist_mod = types.ModuleType("torch.distributed")
 
     def gather_object(obj, object_gather_list, dst, group):
@@ -110,6 +113,7 @@ def gather_object(obj, object_gather_list, dst, group):
     update_from_distributed_mod.update_weights_from_distributed = lambda *args, **kwargs: []
 
     monkeypatch.setitem(sys.modules, "vime", vime_pkg)
+    monkeypatch.setitem(sys.modules, "vime.platforms", platform_mod)
     monkeypatch.setitem(sys.modules, "vime.backends", vime_backends_pkg)
     monkeypatch.setitem(sys.modules, "vime.backends.megatron_utils", megatron_utils_pkg)
     monkeypatch.setitem(sys.modules, "vime.backends.megatron_utils.update_weight", update_weight_pkg)
diff --git a/tests/test_glm4.7_30B_A3B_npu.py b/tests/test_glm4.7_30B_A3B_npu.py
new file mode 100644
index 000000000..b333e849e
--- /dev/null
+++ b/tests/test_glm4.7_30B_A3B_npu.py
@@ -0,0 +1,156 @@
+import os
+import shlex
+
+import vime.utils.external_utils.command_utils as U
+
+
+TEST_ROOT = os.environ.get("HF_HOME") or "/root"
+MODEL_DIR = f"{TEST_ROOT}/models/GLM-4.7-Flash"
+DATASET_DIR = f"{TEST_ROOT}/datasets/dapo-math-17k"
+
+
+def prepare():
+    models_dir = shlex.quote(f"{TEST_ROOT}/models")
+    datasets_dir = shlex.quote(f"{TEST_ROOT}/datasets")
+    model_dir = shlex.quote(MODEL_DIR)
+    dataset_dir = shlex.quote(DATASET_DIR)
+
+    U.exec_command(f"mkdir -p {models_dir} {datasets_dir}")
+    U.exec_command(f"hf download zai-org/GLM-4.7-Flash --local-dir {model_dir}")
+    U.exec_command("hf download --repo-type dataset zhuzilin/dapo-math-17k " f"--local-dir {dataset_dir}")
+
+
+def execute():
+    # Default to G1; G2 and eager diagnostics are explicit, independent opt-ins.
+    enable_mtp = os.environ.get("VIME_TEST_GLM_MTP", "0") == "1"
+    enforce_eager = os.environ.get("VIME_TEST_GLM_EAGER", "0") == "1"
+    model_dir = shlex.quote(MODEL_DIR)
+    prompt_data = shlex.quote(f"{DATASET_DIR}/dapo-math-17k.jsonl")
+
+    # G1 loads HF weights through the native loader, without torch_dist conversion.
+    checkpoint_args = f"--hf-checkpoint {model_dir} --load {model_dir} --ref-load {model_dir} --no-load-optim "
+
+    # Smoke-scaled rollout (num-rollout/batch/n-samples trimmed like test_qwen3_30B_A3B_npu).
+    rollout_args = (
+        f"--prompt-data {prompt_data} "
+        "--input-key prompt "
+        "--label-key label "
+        "--apply-chat-template "
+        "--rollout-shuffle "
+        "--rm-type deepscaler "
+        "--num-rollout 2 "
+        "--rollout-batch-size 4 "
+        "--n-samples-per-prompt 4 "
+        "--rollout-max-response-len 128 "
+        "--rollout-temperature 1 "
+        "--global-batch-size 16 "
+        "--balance-data "
+    )
+
+    # TP=4/EP=8 mirrors scripts/run-glm4.7-30B-A3B-npu.sh.
+    parallel_args = (
+        "--tensor-model-parallel-size 4 "
+        "--sequence-parallel "
+        "--pipeline-model-parallel-size 1 "
+        "--context-parallel-size 1 "
+        "--expert-model-parallel-size 8 "
+        "--expert-tensor-parallel-size 1 "
+        "--moe-token-dispatcher-type alltoall "
+        "--recompute-granularity full "
+        "--recompute-method uniform "
+        "--recompute-num-layers 1 "
+        "--use-dynamic-batch-size "
+        "--max-tokens-per-gpu 20480 "
+        "--micro-batch-size 1 "
+    )
+
+    grpo_args = (
+        "--advantage-estimator grpo "
+        "--use-kl-loss "
+        "--kl-loss-coef 0.00 "
+        "--kl-loss-type low_var_kl "
+        "--entropy-coef 0.00 "
+        "--eps-clip 0.2 "
+        "--eps-clip-high 0.28 "
+    )
+
+    optimizer_args = (
+        "--optimizer adam "
+        "--lr 1e-6 "
+        "--lr-decay-style constant "
+        "--weight-decay 0.1 "
+        "--adam-beta1 0.9 "
+        "--adam-beta2 0.98 "
+        "--optimizer-cpu-offload "
+        "--overlap-cpu-optimizer-d2h-h2d "
+        "--use-precision-aware-optimizer "
+    )
+
+    vllm_args = (
+        "--vllm-additional-config '{\"weight_nz_mode\":0}' "
+        "--rollout-num-gpus-per-engine 4 "
+        "--vllm-gpu-memory-utilization 0.7 "
+        "--vllm-enable-expert-parallel "
+        "--vllm-cudagraph-capture-sizes 1 2 4 8 "
+    )
+    mtp_args = ""
+    if enable_mtp:
+        mtp_args = "--mtp-num-layers 1 --enable-mtp-training --mtp-loss-scaling-factor 0.2 "
+        vllm_args += '--vllm-speculative-config \'{"method":"mtp","num_speculative_tokens":1}\' '
+    if enforce_eager:
+        vllm_args += "--vllm-enforce-eager "
+
+    model_args = (
+        # GLM-4.7-Flash has no HF rope_scaling; MLA otherwise defaults to YaRN.
+        "--rope-type rope "
+        "--attention-dropout 0.0 "
+        "--hidden-dropout 0.0 "
+        "--accumulate-allreduce-grads-in-fp32 "
+        "--attention-softmax-in-fp32 "
+        "--attention-backend flash "
+        "--use-flash-attn "
+        "--no-gradient-accumulation-fusion "
+    )
+
+    runtime_args = (
+        "--train-backend megatron "
+        "--actor-num-nodes 1 "
+        "--actor-num-gpus-per-node 8 "
+        "--rollout-num-gpus 8 "
+        "--ci-test "
+    )
+
+    train_args = (
+        checkpoint_args
+        + rollout_args
+        + parallel_args
+        + grpo_args
+        + optimizer_args
+        + mtp_args
+        + vllm_args
+        + model_args
+        + runtime_args
+    )
+    # Model architecture (num-experts, moe-*, multi-latent-attention, q-lora-rank,
+    # kv-lora-rank, ...) is injected by sourcing scripts/models/glm4.7-30B-A3B.sh
+    # via ${MODEL_ARGS[@]}, so only runtime/training args are passed here.
+    U.execute_train(
+        train_args=train_args,
+        num_gpus_per_node=16,
+        megatron_model_type="glm4.7-30B-A3B",
+        extra_env_vars={
+            "DISABLE_L2_CACHE": "1",
+            "VLLM_USE_AOT_COMPILE": "0",
+        },
+    )
+
+
+def main():
+    prepare()
+    for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"):
+        os.environ.pop(proxy_var, None)
+    execute()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/test_hf_to_megatron.py b/tests/test_hf_to_megatron.py
index 07c0d8a52..412fab08a 100644
--- a/tests/test_hf_to_megatron.py
+++ b/tests/test_hf_to_megatron.py
@@ -22,7 +22,7 @@
 
 from vime.backends.megatron_utils import megatron_to_hf as megatron_to_hf_module
 from vime.backends.megatron_utils.hf_to_megatron import _LOADERS
-from vime.backends.megatron_utils.hf_to_megatron.common import SafetensorReader
+from vime.backends.megatron_utils.hf_to_megatron.common import SafetensorReader, shard_mcore_tensor
 from vime.backends.megatron_utils.hf_to_megatron.deepseek import deepseek_hf_tensor
 from vime.backends.megatron_utils.hf_to_megatron.glm import glm4_hf_tensor, glm4_moe_hf_tensor
 from vime.backends.megatron_utils.hf_to_megatron.qwen import (
@@ -33,6 +33,7 @@
 )
 from vime.backends.megatron_utils.hf_to_megatron.qwen3_next import qwen3_next_hf_tensor
 from vime.backends.megatron_utils.hf_to_megatron.qwen3_omni import qwen3_omni_hf_tensor
+from vime.backends.megatron_utils.hf_to_megatron.qwen3_vl import qwen3_vl_hf_tensor
 from vime.backends.megatron_utils.megatron_to_hf import _convert_to_hf_core, convert_to_hf
 from vime.backends.megatron_utils.megatron_to_hf.deepseekv3 import convert_deepseekv3_to_hf
 from vime.backends.megatron_utils.megatron_to_hf.glm4 import convert_glm4_to_hf
@@ -42,6 +43,7 @@
 from vime.backends.megatron_utils.megatron_to_hf.qwen2 import convert_qwen2_to_hf
 from vime.backends.megatron_utils.megatron_to_hf.qwen3_next import convert_qwen3_next_to_hf
 from vime.backends.megatron_utils.megatron_to_hf.qwen3_omni import convert_qwen3_omni_to_hf
+from vime.backends.megatron_utils.megatron_to_hf.qwen3_vl import convert_qwen3vl_to_hf
 from vime.backends.megatron_utils.megatron_to_hf.qwen3moe import convert_qwen3moe_to_hf
 
 NUM_GPUS = 0
@@ -222,7 +224,7 @@ def test_qwen3_omni_encoder_mapping_is_replicated():
 
 
 @pytest.mark.unit
-@pytest.mark.parametrize("model_name", ["deepseekv32config", "kimik2config"])
+@pytest.mark.parametrize("model_name", ["deepseekv32config", "kimik2config", "glm4moeliteconfig"])
 def test_deepseek_family_parameter_updates_use_the_direct_exporter(model_name):
     parameter = torch.randn(8, 8)
 
@@ -257,6 +259,58 @@ def test_qwen2_moe_parameter_updates_use_the_moe_exporter():
     assert torch.equal(converted[1][1], parameter[6:])
 
 
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    "rest,shape",
+    [
+        ("input_layernorm.weight", (8,)),
+        ("self_attention.linear_q_down_proj.weight", (4, 8)),
+        ("self_attention.linear_q_up_proj.weight", (8, 4)),
+        ("self_attention.linear_q_up_proj.layer_norm_weight", (4,)),
+        ("self_attention.linear_kv_down_proj.weight", (6, 8)),
+        ("self_attention.linear_kv_up_proj.weight", (8, 4)),
+        ("self_attention.linear_kv_up_proj.layer_norm_weight", (4,)),
+        ("self_attention.linear_proj.weight", (8, 8)),
+        ("pre_mlp_layernorm.weight", (8,)),
+        ("mlp.linear_fc1.weight", (12, 8)),
+        ("mlp.linear_fc2.weight", (8, 6)),
+        ("mlp.shared_experts.linear_fc1.weight", (12, 8)),
+        ("mlp.shared_experts.linear_fc2.weight", (8, 6)),
+        ("mlp.experts.linear_fc1.weight3", (12, 8)),
+        ("mlp.experts.linear_fc2.weight3", (8, 6)),
+        ("mlp.router.weight", (4, 8)),
+        ("mlp.router.expert_bias", (4,)),
+    ],
+)
+@pytest.mark.parametrize("mtp", [False, True])
+def test_glm_lite_native_mla_and_moe_round_trip(rest, shape, mtp):
+    prefix = "mtp.layers.0.transformer_layer" if mtp else "decoder.layers.1"
+    name = f"module.module.{prefix}.{rest}"
+    parameter = torch.randn(shape)
+    exported = _convert_to_hf_core(_EXPORT_ARGS, "glm4moeliteconfig", name, parameter)
+    loaded = _LOADERS["glm4_moe_lite"](name, Reader(**dict(exported)), _config("glm4_moe_lite"))
+    assert torch.equal(loaded, parameter)
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    "rest,hf_rest,shape",
+    [
+        ("eh_proj.weight", "eh_proj.weight", (8, 16)),
+        ("enorm.weight", "enorm.weight", (8,)),
+        ("hnorm.weight", "hnorm.weight", (8,)),
+        ("final_layernorm.weight", "shared_head.norm.weight", (8,)),
+    ],
+)
+def test_glm_lite_native_mtp_projection_and_norm_round_trip(rest, hf_rest, shape):
+    name = f"module.module.mtp.layers.0.{rest}"
+    parameter = torch.randn(shape)
+    exported = _convert_to_hf_core(_EXPORT_ARGS, "glm4moeliteconfig", name, parameter)
+    assert [key for key, _ in exported] == [f"model.layers.{_EXPORT_ARGS.num_layers}.{hf_rest}"]
+    loaded = _LOADERS["glm4_moe_lite"](name, Reader(**dict(exported)), _config("glm4_moe_lite"))
+    assert torch.equal(loaded, parameter)
+
+
 @pytest.mark.unit
 def test_qwen_and_llama_share_the_basic_qkv_mapping():
     q = torch.arange(16).view(4, 4)
@@ -398,9 +452,108 @@ def test_loader_scope_stays_explicit():
         "qwen3_moe",
         "qwen3_next",
         "qwen3_omni_moe",
+        "qwen3_vl",
     }
 
 
+@pytest.mark.unit
+@pytest.mark.parametrize(("model_name", "prefix"), [("qwen3", ""), ("qwen3_vl", "language_model.")])
+@pytest.mark.parametrize(
+    ("name", "shape", "is_vocab"),
+    [
+        ("embedding.word_embeddings.weight", (16, 8), True),
+        ("output_layer.weight", (16, 8), True),
+        ("decoder.final_layernorm.weight", (16,), False),
+    ],
+)
+def test_native_export_removes_only_vocab_padding(model_name, prefix, name, shape, is_vocab):
+    args = types.SimpleNamespace(vocab_size=8)
+    weight = torch.randn(shape)
+    [(_, exported)] = convert_to_hf(args, model_name, "module.module." + prefix + name, weight)
+    expected = weight[: args.vocab_size] if is_vocab else weight
+    assert torch.equal(exported, expected)
+    if not is_vocab:
+        assert exported is weight
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    ("name", "shape"),
+    [
+        ("embedding.word_embeddings.weight", (16, 8)),
+        ("output_layer.weight", (16, 8)),
+        ("decoder.final_layernorm.weight", (8,)),
+        ("decoder.layers.0.self_attention.linear_qkv.weight", (16, 8)),
+        ("decoder.layers.0.self_attention.linear_qkv.bias", (16,)),
+        ("decoder.layers.0.self_attention.linear_proj.weight", (8, 8)),
+        ("decoder.layers.0.self_attention.linear_qkv.layer_norm_weight", (8,)),
+        ("decoder.layers.0.self_attention.q_layernorm.weight", (2,)),
+        ("decoder.layers.0.self_attention.k_layernorm.weight", (2,)),
+        ("decoder.layers.0.mlp.linear_fc1.weight", (24, 8)),
+        ("decoder.layers.0.mlp.linear_fc2.weight", (8, 12)),
+        ("decoder.layers.0.mlp.linear_fc1.layer_norm_weight", (8,)),
+    ],
+)
+def test_qwen3_vl_language_round_trip(name, shape):
+    name = "module.module.language_model." + name
+    weight = torch.randn(shape)
+    config = types.SimpleNamespace(
+        tie_word_embeddings=False,
+        text_config=types.SimpleNamespace(
+            hidden_size=8, num_attention_heads=4, num_key_value_heads=2, head_dim=2, tie_word_embeddings=False
+        ),
+    )
+    reader = Reader(**dict(convert_qwen3vl_to_hf(_EXPORT_ARGS, name, weight)))
+    assert torch.equal(qwen3_vl_hf_tensor(name, reader, config), weight)
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    ("top_tied", "text_tied", "has_head"), [(True, False, True), (False, True, True), (False, False, False)]
+)
+def test_qwen3_vl_tied_output_uses_language_embedding(top_tied, text_tied, has_head):
+    embedding = torch.randn(8, 4)
+    tensors = {"model.language_model.embed_tokens.weight": embedding}
+    if has_head:
+        tensors["lm_head.weight"] = torch.zeros_like(embedding)
+    config = types.SimpleNamespace(
+        tie_word_embeddings=top_tied, text_config=types.SimpleNamespace(tie_word_embeddings=text_tied)
+    )
+    assert qwen3_vl_hf_tensor("output_layer.weight", Reader(**tensors), config) is embedding
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("platform", ["cuda", "npu"])
+@pytest.mark.parametrize("expert", [False, True])
+@pytest.mark.parametrize("parallel_size", [1, 2, 4])
+def test_native_fc1_loading_uses_platform_partition_metadata(monkeypatch, platform, expert, parallel_size):
+    mpu = pytest.importorskip("megatron.core").mpu
+    monkeypatch.setenv("VIME_PLATFORM", platform)
+    monkeypatch.setattr(mpu, "get_tensor_model_parallel_world_size", lambda: 8 if expert else parallel_size)
+    monkeypatch.setattr(mpu, "get_tensor_model_parallel_rank", lambda: 0 if expert else parallel_size - 1)
+    monkeypatch.setattr(mpu, "get_expert_tensor_parallel_world_size", lambda: parallel_size if expert else 8)
+    monkeypatch.setattr(mpu, "get_expert_tensor_parallel_rank", lambda: parallel_size - 1 if expert else 0)
+    name = "decoder.layers.0.mlp." + ("experts.linear_fc1.weight0" if expert else "linear_fc1.weight")
+    weight = torch.arange(16 * 8).reshape(16, 8)
+    # MindSpeed grouped column-parallel weights carry dim=1, but store [out, in].
+    parameter = types.SimpleNamespace(
+        tensor_model_parallel=True, partition_dim=1 if platform == "npu" and expert else 0, partition_stride=1
+    )
+    shard = shard_mcore_tensor(name, weight, parameter)
+    gate, up = weight.chunk(2)
+    assert torch.equal(shard, torch.cat((gate.chunk(parallel_size)[-1], up.chunk(parallel_size)[-1])))
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("parallel_mode", [None, "duplicated"])
+def test_native_loading_does_not_shard_replicated_parameters(monkeypatch, parallel_mode):
+    pytest.importorskip("megatron.core")
+    monkeypatch.setenv("VIME_PLATFORM", "npu")
+    parameter = types.SimpleNamespace(tensor_model_parallel=parallel_mode is not None, parallel_mode=parallel_mode)
+    weight = torch.randn(4, 8)
+    assert shard_mcore_tensor("model.visual.blocks.0.mlp.linear_fc1.weight", weight, parameter) is weight
+
+
 @pytest.mark.unit
 def test_reader_dequantizes_block_scaled_fp8(tmp_path):
     weight = torch.linspace(-2, 2, 128 * 128).view(128, 128).to(torch.float8_e4m3fn)
diff --git a/tests/test_qwen3_30B_A3B_npu.py b/tests/test_qwen3_30B_A3B_npu.py
new file mode 100644
index 000000000..e3310096a
--- /dev/null
+++ b/tests/test_qwen3_30B_A3B_npu.py
@@ -0,0 +1,171 @@
+import os
+import shlex
+import sys
+import tempfile
+from pathlib import Path
+
+import vime.utils.external_utils.command_utils as U
+
+
+TEST_ROOT = os.environ.get("HF_HOME") or "/root"
+MODEL_DIR = f"{TEST_ROOT}/models/Qwen3-30B-A3B"
+DATASET_DIR = f"{TEST_ROOT}/datasets/dapo-math-17k"
+
+
+def prepare(torch_dist_ref_load=True):
+    models_dir = shlex.quote(f"{TEST_ROOT}/models")
+    datasets_dir = shlex.quote(f"{TEST_ROOT}/datasets")
+    model_dir = shlex.quote(MODEL_DIR)
+    dataset_dir = shlex.quote(DATASET_DIR)
+
+    U.exec_command(f"mkdir -p {models_dir} {datasets_dir}")
+    U.exec_command(f"hf download Qwen/Qwen3-30B-A3B --local-dir {model_dir}")
+    U.exec_command("hf download --repo-type dataset zhuzilin/dapo-math-17k " f"--local-dir {dataset_dir}")
+    if not torch_dist_ref_load:
+        return None
+
+    # Retain conversion artifacts for inspection; never delete an existing checkpoint.
+    checkpoint_path = Path(tempfile.mkdtemp(prefix="Qwen3-30B-A3B_torch_dist_", dir=f"{TEST_ROOT}/models"))
+    checkpoint_dir = shlex.quote(str(checkpoint_path))
+    U.exec_command(
+        "source scripts/models/qwen3-30B-A3B.sh && "
+        "TRANSFORMERS_VERBOSITY=error "
+        f"VIME_PLATFORM=npu PYTHONPATH={shlex.quote(str(U.repo_base_dir))}:/root/Megatron-LM:${{PYTHONPATH:-}} "
+        f"{shlex.quote(sys.executable)} -m torch.distributed.run --nproc-per-node 8 "
+        "tools/convert_hf_to_torch_dist.py "
+        "${MODEL_ARGS[@]} "
+        f"--hf-checkpoint {model_dir} --save {checkpoint_dir}"
+    )
+
+    tracker = checkpoint_path / "latest_checkpointed_iteration.txt"
+    assert tracker.read_text().strip() == "release"
+    weight_files = [
+        path
+        for path in checkpoint_path.rglob("*")
+        if path.is_file() and path.name != "latest_checkpointed_iteration.txt"
+    ]
+    assert weight_files, f"No checkpoint weights found under {checkpoint_path}"
+    return str(checkpoint_path)
+
+
+def execute(torch_dist_checkpoint=None):
+    model_dir = shlex.quote(MODEL_DIR)
+    prompt_data = shlex.quote(f"{DATASET_DIR}/dapo-math-17k.jsonl")
+
+    checkpoint_args = f"--hf-checkpoint {model_dir} --load {model_dir} --ref-load {model_dir} --no-load-optim "
+    if torch_dist_checkpoint is not None:
+        checkpoint_args = (
+            f"--hf-checkpoint {model_dir} --ref-load {shlex.quote(torch_dist_checkpoint)} --no-load-optim "
+        )
+
+    rollout_args = (
+        f"--prompt-data {prompt_data} "
+        "--input-key prompt "
+        "--label-key label "
+        "--apply-chat-template "
+        "--rollout-shuffle "
+        "--rm-type deepscaler "
+        "--num-rollout 2 "
+        "--rollout-batch-size 4 "
+        "--n-samples-per-prompt 4 "
+        "--rollout-max-response-len 8192 "
+        "--rollout-temperature 1 "
+        "--global-batch-size 16 "
+        "--balance-data "
+    )
+
+    parallel_args = (
+        "--tensor-model-parallel-size 4 "
+        "--sequence-parallel "
+        "--pipeline-model-parallel-size 1 "
+        "--context-parallel-size 1 "
+        "--expert-model-parallel-size 8 "
+        "--expert-tensor-parallel-size 1 "
+        "--moe-token-dispatcher-type alltoall "
+        "--recompute-granularity full "
+        "--recompute-method uniform "
+        "--recompute-num-layers 1 "
+        "--use-dynamic-batch-size "
+        "--max-tokens-per-gpu 20480 "
+        "--micro-batch-size 1 "
+    )
+
+    grpo_args = (
+        "--advantage-estimator grpo "
+        "--use-kl-loss "
+        "--kl-loss-coef 0.00 "
+        "--kl-loss-type low_var_kl "
+        "--entropy-coef 0.00 "
+        "--eps-clip 0.2 "
+        "--eps-clip-high 0.28 "
+    )
+
+    optimizer_args = (
+        "--optimizer adam "
+        "--lr 1e-6 "
+        "--lr-decay-style constant "
+        "--weight-decay 0.1 "
+        "--adam-beta1 0.9 "
+        "--adam-beta2 0.98 "
+        "--optimizer-cpu-offload "
+        "--overlap-cpu-optimizer-d2h-h2d "
+        "--use-precision-aware-optimizer "
+    )
+
+    vllm_args = (
+        "--vllm-additional-config '{\"weight_nz_mode\":0}' "
+        "--rollout-num-gpus-per-engine 4 "
+        "--vllm-enable-sleep-mode "
+        "--vllm-enable-expert-parallel "
+        "--vllm-gpu-memory-utilization 0.7 "
+    )
+
+    model_args = (
+        "--attention-dropout 0.0 "
+        "--hidden-dropout 0.0 "
+        "--accumulate-allreduce-grads-in-fp32 "
+        "--attention-softmax-in-fp32 "
+        "--attention-backend flash "
+        "--use-flash-attn "
+        "--no-gradient-accumulation-fusion "
+    )
+
+    runtime_args = (
+        "--train-backend megatron "
+        "--actor-num-nodes 1 "
+        "--actor-num-gpus-per-node 8 "
+        "--rollout-num-gpus 8 "
+        "--ci-test "
+        "--colocate "
+    )
+
+    train_args = (
+        checkpoint_args
+        + rollout_args
+        + parallel_args
+        + grpo_args
+        + optimizer_args
+        + vllm_args
+        + model_args
+        + runtime_args
+    )
+    U.execute_train(
+        train_args=train_args,
+        num_gpus_per_node=16,
+        megatron_model_type="qwen3-30B-A3B",
+        extra_env_vars={
+            "DISABLE_L2_CACHE": "1",
+            "VLLM_USE_AOT_COMPILE": "0",
+        },
+    )
+
+
+def main():
+    checkpoint = prepare(torch_dist_ref_load=os.environ.get("VIME_TEST_TORCH_DIST_REF_LOAD", "1") == "1")
+    for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"):
+        os.environ.pop(proxy_var, None)
+    execute(checkpoint)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/test_qwen3_4B_npu.py b/tests/test_qwen3_4B_npu.py
new file mode 100644
index 000000000..ea94cf6d1
--- /dev/null
+++ b/tests/test_qwen3_4B_npu.py
@@ -0,0 +1,132 @@
+import os
+import shlex
+
+import vime.utils.external_utils.command_utils as U
+
+
+TEST_ROOT = os.environ.get("HF_HOME") or "/root"
+MODEL_DIR = f"{TEST_ROOT}/models/Qwen3-4B"
+DATASET_DIR = f"{TEST_ROOT}/datasets/dapo-math-17k"
+
+
+def prepare():
+    models_dir = shlex.quote(f"{TEST_ROOT}/models")
+    datasets_dir = shlex.quote(f"{TEST_ROOT}/datasets")
+    model_dir = shlex.quote(MODEL_DIR)
+    dataset_dir = shlex.quote(DATASET_DIR)
+
+    U.exec_command(f"mkdir -p {models_dir} {datasets_dir}")
+    U.exec_command(f"hf download Qwen/Qwen3-4B --local-dir {model_dir}")
+    U.exec_command("hf download --repo-type dataset zhuzilin/dapo-math-17k " f"--local-dir {dataset_dir}")
+
+
+def execute():
+    model_dir = shlex.quote(MODEL_DIR)
+    prompt_data = shlex.quote(f"{DATASET_DIR}/dapo-math-17k.jsonl")
+
+    checkpoint_args = f"--hf-checkpoint {model_dir} --ref-load {model_dir} --load {model_dir} --no-load-optim "
+
+    rollout_args = (
+        f"--prompt-data {prompt_data} "
+        "--input-key prompt "
+        "--label-key label "
+        "--apply-chat-template "
+        "--rollout-shuffle "
+        "--rm-type math "
+        "--num-rollout 3 "
+        "--rollout-batch-size 4 "
+        "--n-samples-per-prompt 4 "
+        "--rollout-max-response-len 2048 "
+        "--rollout-temperature 1 "
+        "--global-batch-size 16 "
+        "--balance-data "
+    )
+
+    parallel_args = (
+        "--tensor-model-parallel-size 4 "
+        "--pipeline-model-parallel-size 1 "
+        "--context-parallel-size 1 "
+        "--expert-model-parallel-size 1 "
+        "--expert-tensor-parallel-size 1 "
+        "--recompute-granularity full "
+        "--recompute-method uniform "
+        "--recompute-num-layers 1 "
+        "--use-dynamic-batch-size "
+        "--max-tokens-per-gpu 8192 "
+    )
+
+    grpo_args = (
+        "--advantage-estimator grpo "
+        "--kl-loss-coef 0.0 "
+        "--kl-loss-type low_var_kl "
+        "--kl-coef 0.00 "
+        "--entropy-coef 0.0 "
+        "--eps-clip 0.2 "
+        "--eps-clip-high 0.28 "
+    )
+
+    optimizer_args = (
+        "--optimizer adam "
+        "--lr 1e-6 "
+        "--lr-decay-style constant "
+        "--weight-decay 0.1 "
+        "--adam-beta1 0.9 "
+        "--adam-beta2 0.98 "
+        "--optimizer-cpu-offload "
+        "--overlap-cpu-optimizer-d2h-h2d "
+        "--use-precision-aware-optimizer "
+    )
+
+    vllm_args = (
+        "--vllm-additional-config '{\"weight_nz_mode\":0}' "
+        "--rollout-num-gpus-per-engine 4 "
+        "--vllm-enable-sleep-mode "
+        "--vllm-gpu-memory-utilization 0.6 "
+        "--vllm-max-model-len 4096 "
+    )
+
+    model_args = (
+        "--attention-dropout 0.0 "
+        "--hidden-dropout 0.0 "
+        "--accumulate-allreduce-grads-in-fp32 "
+        "--attention-softmax-in-fp32 "
+        "--attention-backend flash "
+        "--micro-batch-size 1 "
+        "--use-flash-attn "
+    )
+
+    runtime_args = (
+        "--train-backend megatron "
+        "--actor-num-nodes 1 "
+        "--actor-num-gpus-per-node 4 "
+        "--rollout-num-gpus 4 "
+        "--ci-test "
+    )
+
+    train_args = (
+        checkpoint_args
+        + rollout_args
+        + parallel_args
+        + grpo_args
+        + optimizer_args
+        + vllm_args
+        + model_args
+        + runtime_args
+    )
+    U.execute_train(
+        train_args=train_args,
+        num_gpus_per_node=8,
+        megatron_model_type="qwen3-4B",
+        extra_env_vars={},
+    )
+
+
+def main():
+    prepare()
+    for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"):
+        os.environ.pop(proxy_var, None)
+    execute()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/test_qwen3_vl_8B_npu.py b/tests/test_qwen3_vl_8B_npu.py
new file mode 100644
index 000000000..c117a6744
--- /dev/null
+++ b/tests/test_qwen3_vl_8B_npu.py
@@ -0,0 +1,139 @@
+import os
+import shlex
+
+import vime.utils.external_utils.command_utils as U
+
+
+# Single-turn Qwen3-VL GRPO on geo3k (mirrors examples/geo3k_vlm/run_geo3k_vlm_npu.sh).
+# Qwen3-VL-8B uses the qwen3-8B language config and the native VL provider.
+MODEL_NAME = "Qwen3-VL-8B-Instruct"
+MODEL_TYPE = "qwen3-8B"
+TEST_ROOT = os.environ.get("HF_HOME") or "/root"
+MODEL_DIR = f"{TEST_ROOT}/models/{MODEL_NAME}"
+DATASET_DIR = f"{TEST_ROOT}/datasets/geo3k_imgurl"
+
+
+def prepare():
+    models_dir = shlex.quote(f"{TEST_ROOT}/models")
+    datasets_dir = shlex.quote(f"{TEST_ROOT}/datasets")
+    model_dir = shlex.quote(MODEL_DIR)
+    dataset_dir = shlex.quote(DATASET_DIR)
+
+    U.exec_command(f"mkdir -p {models_dir} {datasets_dir}")
+    U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir {model_dir}")
+    U.exec_command(f"hf download --repo-type dataset chenhegu/geo3k_imgurl --local-dir {dataset_dir}")
+
+
+def execute():
+    model_dir = shlex.quote(MODEL_DIR)
+    prompt_data = shlex.quote(f"{DATASET_DIR}/train.parquet")
+
+    checkpoint_args = f"--hf-checkpoint {model_dir} --load {model_dir} --no-load-optim "
+
+    rollout_args = (
+        f"--prompt-data {prompt_data} "
+        "--input-key problem "
+        "--label-key answer "
+        '--multimodal-keys \'{"image": "images"}\' '
+        "--apply-chat-template "
+        "--rollout-shuffle "
+        "--rm-type math "
+        "--num-rollout 2 "
+        "--rollout-batch-size 4 "
+        "--n-samples-per-prompt 4 "
+        "--rollout-max-response-len 4096 "
+        "--rollout-temperature 1 "
+        "--global-batch-size 16 "
+    )
+
+    parallel_args = (
+        "--tensor-model-parallel-size 4 "
+        "--sequence-parallel "
+        "--pipeline-model-parallel-size 1 "
+        "--context-parallel-size 1 "
+        "--expert-model-parallel-size 1 "
+        "--expert-tensor-parallel-size 1 "
+        "--recompute-granularity full "
+        "--recompute-method uniform "
+        "--recompute-num-layers 1 "
+        "--use-dynamic-batch-size "
+        "--max-tokens-per-gpu 4096 "
+    )
+
+    grpo_args = (
+        "--advantage-estimator grpo "
+        "--kl-loss-coef 0.00 "
+        "--kl-loss-type low_var_kl "
+        "--kl-coef 0.00 "
+        "--entropy-coef 0.00 "
+        "--eps-clip 0.2 "
+        "--eps-clip-high 0.28 "
+    )
+
+    optimizer_args = (
+        "--optimizer adam "
+        "--lr 1e-6 "
+        "--lr-decay-style constant "
+        "--weight-decay 0.1 "
+        "--adam-beta1 0.9 "
+        "--adam-beta2 0.98 "
+    )
+
+    vllm_args = (
+        "--vllm-additional-config '{\"weight_nz_mode\":0}' "
+        "--rollout-num-gpus-per-engine 1 "
+        "--vllm-gpu-memory-utilization 0.7 "
+        "--vllm-max-model-len 16384 "
+        "--vllm-generation-config auto "
+        "--vllm-logprobs-mode processed_logprobs "
+    )
+
+    model_args = (
+        "--spec vime_plugins.models.qwen3_vl get_qwen3_vl_model_provider "
+        "--attention-dropout 0.0 "
+        "--hidden-dropout 0.0 "
+        "--accumulate-allreduce-grads-in-fp32 "
+        "--attention-softmax-in-fp32 "
+        "--attention-backend flash "
+        "--use-flash-attn "
+        "--no-gradient-accumulation-fusion "
+    )
+
+    runtime_args = (
+        "--train-backend megatron "
+        "--actor-num-nodes 1 "
+        "--actor-num-gpus-per-node 8 "
+        "--rollout-num-gpus 8 "
+        "--colocate "
+        "--ci-test "
+    )
+
+    train_args = (
+        checkpoint_args
+        + rollout_args
+        + parallel_args
+        + grpo_args
+        + optimizer_args
+        + vllm_args
+        + model_args
+        + runtime_args
+    )
+    # qwen3-8B.sh builds MODEL_ARGS with --rotary-base ${MODEL_ARGS_ROTARY_BASE}; Qwen3-VL needs 5e6.
+    os.environ["MODEL_ARGS_ROTARY_BASE"] = "5000000"
+    U.execute_train(
+        train_args=train_args,
+        num_gpus_per_node=8,
+        megatron_model_type=MODEL_TYPE,
+        extra_env_vars={},
+    )
+
+
+def main():
+    prepare()
+    for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"):
+        os.environ.pop(proxy_var, None)
+    execute()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/test_qwen3_vl_native.py b/tests/test_qwen3_vl_native.py
new file mode 100644
index 000000000..135fcd5b5
--- /dev/null
+++ b/tests/test_qwen3_vl_native.py
@@ -0,0 +1,318 @@
+from types import SimpleNamespace
+
+import pytest
+import torch
+from safetensors.torch import save_file
+
+NUM_GPUS = 0
+pytestmark = pytest.mark.unit
+
+
+@pytest.fixture
+def native(monkeypatch):
+    monkeypatch.setenv("VIME_PLATFORM", "cuda")
+    pytest.importorskip("megatron.core")
+    from vime_plugins.models import qwen3_vl
+
+    return qwen3_vl
+
+
+@pytest.fixture
+def hf_config():
+    from transformers import Qwen3VLConfig
+
+    return Qwen3VLConfig(
+        image_token_id=10,
+        video_token_id=20,
+        vision_start_token_id=30,
+        text_config={
+            "hidden_size": 8,
+            "intermediate_size": 16,
+            "num_hidden_layers": 2,
+            "num_attention_heads": 2,
+            "num_key_value_heads": 1,
+            "head_dim": 4,
+            "vocab_size": 32,
+            "rope_scaling": {"rope_type": "default", "mrope_section": [1, 1, 0], "mrope_interleaved": True},
+            "rope_theta": 5000000,
+        },
+        vision_config={
+            "hidden_size": 8,
+            "intermediate_size": 16,
+            "depth": 2,
+            "num_heads": 2,
+            "out_hidden_size": 8,
+            "patch_size": 2,
+            "temporal_patch_size": 2,
+            "spatial_merge_size": 2,
+            "num_position_embeddings": 4,
+            "deepstack_visual_indexes": [0],
+        },
+    )
+
+
+def test_vision_native_load_export_and_backward(native, hf_config, tmp_path, monkeypatch):
+    from vime.backends.megatron_utils.hf_to_megatron.common import load_model_hf_weights
+    from vime.backends.megatron_utils.hf_to_megatron.qwen3_vl import qwen3_vl_hf_tensor
+    from vime.backends.megatron_utils.megatron_to_hf import convert_to_hf
+    from vime.backends.megatron_utils.update_weight import common
+
+    config = SimpleNamespace(use_cpu_initialization=True, params_dtype=torch.float32, recompute_granularity="full")
+    vision = native._load_vision_model(hf_config, config)
+    tensors = {"model.visual." + name: param.detach().clone() for name, param in vision.named_parameters()}
+    save_file(tensors, tmp_path / "model.safetensors")
+    with torch.no_grad():
+        for parameter in vision.parameters():
+            parameter.zero_()
+    named = [("module.module.model.visual." + name, param) for name, param in vision.named_parameters()]
+    monkeypatch.setattr(common, "named_params_and_buffers", lambda args, model: iter(named))
+    load_model_hf_weights(SimpleNamespace(), [vision], tmp_path, hf_config, qwen3_vl_hf_tensor)
+
+    for name, parameter in named:
+        [(hf_name, exported)] = convert_to_hf(None, "qwen3_vl", name, parameter)
+        assert torch.equal(exported, tensors[hf_name])
+        assert parameter.requires_grad
+        assert not parameter.tensor_model_parallel
+        assert parameter.partition_dim == -1
+    output = vision(torch.randn(4, 24), grid_thw=torch.tensor([[1, 2, 2]]))
+    features, deepstack = (
+        (output.pooler_output, output.deepstack_features) if hasattr(output, "pooler_output") else output
+    )
+    (features.square().sum() + deepstack[0].square().sum()).backward()
+    assert vision.patch_embed.proj.weight.grad.abs().sum() > 0
+    assert vision.deepstack_merger_list[0].linear_fc2.weight.grad.abs().sum() > 0
+
+
+def _injection_model(native, *, sequence_parallel=False):
+    class Embedding:
+        def __call__(self, input_ids, position_ids):
+            return input_ids.T[..., None].float().expand(-1, -1, 2).clone()
+
+    class Vision:
+        dtype = torch.float32
+
+        def __call__(self, values, grid_thw):
+            return SimpleNamespace(pooler_output=values, deepstack_features=[values * 2, values * 3])
+
+    return SimpleNamespace(
+        config=SimpleNamespace(sequence_parallel=sequence_parallel),
+        language_model=SimpleNamespace(embedding=Embedding()),
+        model=SimpleNamespace(visual=Vision()),
+        image_token_id=10,
+        video_token_id=20,
+    )
+
+
+def test_vision_and_deepstack_follow_token_order_and_keep_gradients(native):
+    model = _injection_model(native)
+    image = torch.tensor([[100.0, 101.0]], requires_grad=True)
+    video = torch.tensor([[200.0, 201.0]], requires_grad=True)
+    embeddings, mask, deepstack = native.Qwen3VLModel._inject_vision_embeddings(
+        model, torch.tensor([[7, 20, 8, 10]]), image, video, torch.tensor([[1, 2, 2]]), torch.tensor([[1, 2, 2]])
+    )
+    expected = torch.cat((video, image))
+    assert torch.equal(embeddings[:, 0][mask[0]], expected)
+    assert torch.equal(deepstack[0], expected * 2)
+    (embeddings.sum() + sum(t.sum() for t in deepstack)).backward()
+    assert torch.equal(image.grad, torch.full_like(image, 6))
+    assert torch.equal(video.grad, torch.full_like(video, 6))
+
+
+def test_deepstack_sp_routes_gradients_through_tp_copy(native, monkeypatch):
+    copied = []
+    monkeypatch.setattr(
+        native.tensor_parallel, "copy_to_tensor_model_parallel_region", lambda value: copied.append(value) or value
+    )
+    monkeypatch.setattr(native.tensor_parallel, "scatter_to_sequence_parallel_region", lambda value: value.chunk(4)[1])
+    monkeypatch.setattr(native.mpu, "get_tensor_model_parallel_world_size", lambda: 4)
+    monkeypatch.setattr(native.mpu, "get_tensor_model_parallel_rank", lambda: 1)
+    model = _injection_model(native, sequence_parallel=True)
+    image = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+    _, mask, deepstack = native.Qwen3VLModel._inject_vision_embeddings(
+        model, torch.tensor([[10, 7, 10, 7, 7, 7, 7, 7]]), image, None, torch.tensor([[1, 2, 4]]), None
+    )
+    assert len(copied) == 2
+    assert mask.tolist() == [[True, False]]
+    assert torch.equal(deepstack[0], image[1:] * 2)
+
+
+def test_missing_vision_is_not_silently_treated_as_text(native):
+    with pytest.raises(ValueError, match="pixel values"):
+        native.Qwen3VLModel._inject_vision_embeddings(
+            _injection_model(native), torch.tensor([[10]]), None, None, None, None
+        )
+
+
+def _vision_tp_worker(rank, rendezvous):
+    from datetime import timedelta
+    from unittest.mock import patch
+
+    import torch.distributed as dist
+
+    from vime_plugins.models import qwen3_vl as native
+
+    torch.set_num_threads(1)
+    dist.init_process_group("gloo", init_method=rendezvous, rank=rank, world_size=4, timeout=timedelta(seconds=45))
+    try:
+        copy = native.tensor_parallel.copy_to_tensor_model_parallel_region
+        scatter = native.tensor_parallel.scatter_to_sequence_parallel_region
+        with (
+            patch.object(torch.cuda, "current_device", lambda: torch.device("cpu")),
+            patch.object(
+                native.tensor_parallel,
+                "copy_to_tensor_model_parallel_region",
+                lambda x: copy(x, group=dist.group.WORLD),
+            ),
+            patch.object(
+                native.tensor_parallel,
+                "scatter_to_sequence_parallel_region",
+                lambda x: scatter(x, group=dist.group.WORLD),
+            ),
+            patch.object(native.mpu, "get_tensor_model_parallel_world_size", lambda: 4),
+            patch.object(native.mpu, "get_tensor_model_parallel_rank", lambda: rank),
+        ):
+            image = torch.arange(8, dtype=torch.float32).view(4, 2).requires_grad_()
+            embeddings, _, deepstack = native.Qwen3VLModel._inject_vision_embeddings(
+                _injection_model(native, sequence_parallel=True),
+                torch.tensor([[10, 7, 10, 7, 10, 7, 10, 7]]),
+                image,
+                None,
+                torch.tensor([[1, 4, 4]]),
+                None,
+            )
+            (embeddings.sum() + sum(t.sum() for t in deepstack)).backward()
+            # All four replicas receive the same complete gradient, including
+            # image features consumed by other SP ranks (1 + 2 + 3 = 6).
+            torch.testing.assert_close(image.grad, torch.full_like(image, 6))
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.integration
+def test_trainable_vision_tp4_gradients_on_cpu(native, tmp_path):
+    torch.multiprocessing.spawn(_vision_tp_worker, args=((tmp_path / "gloo").as_uri(),), nprocs=4)
+
+
+def test_deepstack_gradients_survive_main_recompute(native, monkeypatch):
+    from torch.utils.checkpoint import checkpoint
+
+    from vime_plugins.models.qwen3_omni_transformer import Qwen3OmniTransformerBlock
+
+    class Layer(torch.nn.Module):
+        def __init__(self, number):
+            super().__init__()
+            self.layer_number = number
+
+        def forward(self, hidden_states, **kwargs):
+            return hidden_states * 2, None
+
+    block = Qwen3OmniTransformerBlock.__new__(Qwen3OmniTransformerBlock)
+    torch.nn.Module.__init__(block)
+    block.config = SimpleNamespace(
+        fp8=False, distribute_saved_activations=False, recompute_method="uniform", recompute_num_layers=1
+    )
+    block.pre_process = True
+    block.layers = torch.nn.ModuleList([Layer(1), Layer(2)])
+    block.num_layers_per_pipeline_rank = 2
+    monkeypatch.setattr(
+        native.tensor_parallel, "checkpoint", lambda fn, distribute, *args: checkpoint(fn, *args, use_reentrant=True)
+    )
+    hidden = torch.ones(4, 1, 2, requires_grad=True)
+    features = [torch.ones(1, 2, requires_grad=True), torch.ones(1, 2, requires_grad=True)]
+    output = block._checkpointed_forward(
+        hidden,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,
+        False,
+        visual_pos_masks=torch.tensor([[False, True, False, False]]),
+        deepstack_visual_embeds=features,
+    )
+    output.sum().backward()
+    torch.testing.assert_close(hidden.grad, torch.full_like(hidden, 4))
+    torch.testing.assert_close(features[0].grad, torch.full_like(features[0], 2))
+    torch.testing.assert_close(features[1].grad, torch.ones_like(features[1]))
+
+
+def test_interleaved_mrope_matches_hf_for_packed_images(native, hf_config):
+    from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLTextRotaryEmbedding
+
+    from vime_plugins.models.qwen3_omni_moe import Qwen3OmniMultimodalRotaryEmbedding
+
+    positions = native.build_packed_mrope_position_ids(
+        torch.tensor([[30, 10, 10, 10, 10, 7, 30, 10, 8]]),
+        [0, 6, 9],
+        torch.tensor([[1, 4, 4], [1, 2, 2]]),
+        None,
+        image_token_id=10,
+        video_token_id=20,
+        vision_start_token_id=30,
+        spatial_merge_size=2,
+    )
+    assert positions[:, 0, 6].tolist() == [0, 0, 0]
+    hf_rope = Qwen3VLTextRotaryEmbedding(hf_config.text_config)
+    # Exercise the real main RoPE arithmetic without allocating a CUDA buffer.
+    rope = Qwen3OmniMultimodalRotaryEmbedding.__new__(Qwen3OmniMultimodalRotaryEmbedding)
+    torch.nn.Module.__init__(rope)
+    rope.inv_freq = hf_rope.inv_freq
+    rope.seq_len_interpolation_factor = None
+    rope.cp_group = None
+    rope.is_thd_format = True
+    freqs = rope(positions, [1, 1, 0])[:, 0, 0].unsqueeze(0)
+    cos, sin = hf_rope(torch.zeros(1, 9, 4), positions)
+    torch.testing.assert_close(freqs.cos(), cos)
+    torch.testing.assert_close(freqs.sin(), sin)
+
+
+@pytest.mark.parametrize(("pp", "cp", "mtp"), [(2, 1, None), (1, 2, None), (1, 1, 1)])
+def test_provider_rejects_unimplemented_topologies(native, pp, cp, mtp):
+    with pytest.raises(ValueError, match="PP=1 and CP=1|MTP"):
+        native.get_qwen3_vl_model_provider(
+            SimpleNamespace(mtp_num_layers=mtp),
+            SimpleNamespace(pipeline_model_parallel_size=pp, context_parallel_size=cp),
+            None,
+        )
+
+
+def test_provider_uses_main_dense_deepstack_gpt(native, hf_config, monkeypatch):
+    calls = {}
+
+    def gpt(**kwargs):
+        calls.update(kwargs)
+        return SimpleNamespace(share_embeddings_and_output_weights=False)
+
+    monkeypatch.setattr(native, "Qwen3OmniMoeGPTModel", gpt)
+    monkeypatch.setattr(native, "_load_vision_model", lambda *args: torch.nn.Linear(8, 8))
+    monkeypatch.setattr(native.AutoConfig, "from_pretrained", lambda *args, **kwargs: hf_config)
+    monkeypatch.setattr(
+        native, "get_gpt_layer_with_transformer_engine_spec", lambda *, qk_layernorm: {"qk_layernorm": qk_layernorm}
+    )
+    args = SimpleNamespace(
+        hf_checkpoint="unused",
+        mtp_num_layers=None,
+        transformer_impl="transformer_engine",
+        normalization="RMSNorm",
+        padded_vocab_size=32,
+        max_position_embeddings=128,
+        fp16_lm_cross_entropy=False,
+        untie_embeddings_and_output_weights=True,
+        rotary_percent=1.0,
+        rotary_base=1000000,
+    )
+    config = SimpleNamespace(pipeline_model_parallel_size=1, context_parallel_size=1, normalization="RMSNorm")
+    model = native.get_qwen3_vl_model_provider(args, config, None)()
+    assert calls["transformer_layer_spec"] == {"qk_layernorm": True}
+    assert calls["config"] is config
+    assert calls["config"].normalization == "RMSNorm"
+    assert calls["position_embedding_type"] == "mrope"
+    assert calls["rotary_base"] == 5000000
+    assert calls["scatter_embedding_sequence_parallel"] is False
+    assert config.mrope_section == [1, 1, 0]
+    assert all(parameter.requires_grad for parameter in model.model.visual.parameters())
+
+
+if __name__ == "__main__":
+    raise SystemExit(pytest.main([__file__]))
diff --git a/tests/utils/test_megatron_role_config.py b/tests/utils/test_megatron_role_config.py
index fe1c847db..f6ebc5f4b 100644
--- a/tests/utils/test_megatron_role_config.py
+++ b/tests/utils/test_megatron_role_config.py
@@ -15,6 +15,7 @@
 import _unit_stubs
 
 _unit_stubs.install_rollout_optional_stubs()
+_unit_stubs.install_vllm_cli_stubs()
 
 
 def _write_yaml(data: dict) -> str:
diff --git a/tests/utils/test_npu_accelerator.py b/tests/utils/test_npu_accelerator.py
new file mode 100644
index 000000000..709d62955
--- /dev/null
+++ b/tests/utils/test_npu_accelerator.py
@@ -0,0 +1,170 @@
+"""CPU contracts for NPU selection and pre-Megatron bootstrap ordering."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vime.platforms import current_platform, get_platform, npu, reset_platform_cache
+from vime.platforms.npu import NPUAccelerator
+from vime.utils import accelerator
+
+
+@pytest.fixture(autouse=True)
+def npu_runtime(monkeypatch):
+    reset_platform_cache()
+    monkeypatch.setenv("VIME_PLATFORM", "npu")
+    monkeypatch.delenv("VIME_ACCELERATOR", raising=False)
+    monkeypatch.delenv("MUSA_VISIBLE_DEVICES", raising=False)
+    monkeypatch.delenv("MUSA_PATCH_PATH", raising=False)
+    monkeypatch.setattr(accelerator, "_REGISTRY", {})
+    monkeypatch.setattr(accelerator, "_ACCELERATOR", None)
+    monkeypatch.setattr(accelerator, "_cuda_available", lambda: False)
+    monkeypatch.setattr(accelerator, "is_musa_available", lambda: False)
+    calls = []
+    fake_npu = SimpleNamespace(
+        is_available=lambda: True,
+        current_device=lambda: 1,
+        device_count=lambda: 2,
+        set_device=lambda index: calls.append(("set_device", index)),
+        synchronize=lambda: calls.append("synchronize"),
+        empty_cache=lambda: calls.append("empty_cache"),
+        ipc_collect=lambda: calls.append("ipc_collect"),
+    )
+    monkeypatch.setattr(torch, "npu", fake_npu, raising=False)
+    yield calls
+    reset_platform_cache()
+
+
+def test_platform_registers_npu_before_main_auto_selection(monkeypatch, npu_runtime):
+    # MindSpeed may also make CUDA's availability probe return true.
+    monkeypatch.setattr(accelerator, "_cuda_available", lambda: True)
+    assert current_platform().is_npu
+    assert isinstance(accelerator.initialize_accelerator(), NPUAccelerator)
+    assert accelerator.device_type() == "npu"
+    assert accelerator.process_group_backend() == "hccl"
+    assert accelerator.process_group_backend("gloo") == "gloo"
+    assert accelerator.is_accelerator_backend("cpu:gloo,npu:hccl")
+    assert not accelerator.is_accelerator_backend("gloo")
+    assert accelerator.distributed_device_id() is None
+    assert accelerator.visible_devices_env_key() == "ASCEND_RT_VISIBLE_DEVICES"
+    monkeypatch.setenv("ASCEND_RT_VISIBLE_DEVICES", "4,7")
+    assert accelerator.resolve_visible_device_id("7") == 1
+    accelerator.set_device(1)
+    accelerator.synchronize()
+    accelerator.ipc_collect()
+    accelerator.empty_cache()
+    assert npu_runtime == [("set_device", 1), "synchronize", "ipc_collect", "empty_cache"]
+
+
+def test_npu_allocator_remains_owned_by_existing_runtime_hooks(monkeypatch):
+    monkeypatch.setenv("VIME_ENABLE_EXPANDABLE_SEGMENTS", "1")
+    assert NPUAccelerator().set_allocator_expandable_segments() is False
+
+
+def test_main_npu_override_selects_matching_platform(monkeypatch):
+    monkeypatch.delenv("VIME_PLATFORM")
+    monkeypatch.setenv("VIME_ACCELERATOR", "npu")
+    assert current_platform().is_npu
+    assert accelerator.get_accelerator().name == "npu"
+
+
+@pytest.mark.parametrize("platform,backend", [("npu", "cuda"), ("cuda", "npu"), ("npu", "musa")])
+def test_conflicting_overrides_fail_before_bootstrap(monkeypatch, platform, backend):
+    monkeypatch.setenv("VIME_PLATFORM", platform)
+    monkeypatch.setenv("VIME_ACCELERATOR", backend)
+    with pytest.raises(ValueError, match="Conflicting VIME_PLATFORM"):
+        current_platform()
+
+
+def test_registered_npu_does_not_override_explicit_cuda_platform(monkeypatch):
+    get_platform("npu")
+    monkeypatch.setenv("VIME_PLATFORM", "cuda")
+    monkeypatch.setattr(accelerator, "_cuda_available", lambda: True)
+    monkeypatch.setattr(accelerator.CUDAAccelerator, "is_available", lambda self: True)
+    assert current_platform().name == "cuda"
+    assert accelerator.initialize_accelerator().name == "cuda"
+
+
+def test_bootstrap_selects_npu_before_adaptor_and_attention(monkeypatch):
+    events = []
+    bootstrap = current_platform().megatron
+    monkeypatch.setattr(npu, "_ensure_torch_npu", lambda: events.append("torch_npu"))
+    monkeypatch.setattr(npu, "_install_safe_empty_cache", lambda: events.append("empty_cache_guard"))
+    original_import = npu.importlib.import_module
+
+    def import_module(name, *args, **kwargs):
+        if name in {"megatron_adaptor", "vime.backends.megatron_utils.npu_attention_patch"}:
+            assert accelerator.get_accelerator().name == "npu"
+            events.append(name)
+            bootstrap.bootstrap()  # Recursive imports must not repeat initialization.
+            return SimpleNamespace()
+        return original_import(name, *args, **kwargs)
+
+    monkeypatch.setattr(npu.importlib, "import_module", import_module)
+    bootstrap.bootstrap()
+    bootstrap.bootstrap()
+    assert events == [
+        "torch_npu",
+        "empty_cache_guard",
+        "megatron_adaptor",
+        "vime.backends.megatron_utils.npu_attention_patch",
+    ]
+
+
+def test_bootstrap_rejects_preselected_cuda_without_replacing_it(monkeypatch):
+    selected = accelerator.CUDAAccelerator()
+    monkeypatch.setattr(accelerator, "_ACCELERATOR", selected)
+    monkeypatch.setattr(npu, "_ensure_torch_npu", lambda: None)
+    bootstrap = current_platform().megatron
+    with pytest.raises(RuntimeError, match="already selected 'cuda'"):
+        bootstrap.bootstrap()
+    assert accelerator._ACCELERATOR is selected
+    assert not bootstrap._bootstrapping
+    assert not bootstrap._bootstrapped
+
+
+def test_repatch_passes_typed_args_and_restores_attention(monkeypatch):
+    events = []
+    full_args = SimpleNamespace(adaptor_default=True)
+    typed_config = {"weight_nz_mode": 0}
+    args = SimpleNamespace(vllm_additional_config=typed_config, tensor_model_parallel_size=4)
+    original_forward = object()
+    vime_forward = object()
+    attention_class = SimpleNamespace(forward=vime_forward)
+
+    def apply_features(config):
+        assert config is full_args
+        assert config.vllm_additional_config is typed_config
+        assert config.tensor_model_parallel_size == 4
+        assert config.adaptor_default
+        attention_class.forward = original_forward
+        events.append("features")
+
+    modules = {
+        "megatron_adaptor.features_manager.features_manager": SimpleNamespace(
+            FeaturesManager=SimpleNamespace(
+                remove_patches=lambda: events.append("remove"),
+                apply_features_pre_patches=lambda config: events.append(("pre", config)),
+                apply_features_patches=apply_features,
+            )
+        ),
+        "megatron_adaptor.utils.args_utils": SimpleNamespace(get_full_args=lambda: full_args),
+        "vime.backends.megatron_utils.npu_attention_patch": SimpleNamespace(
+            DotProductAttention=attention_class,
+            npu_dot_product_attention_forward=vime_forward,
+        ),
+    }
+    bootstrap = current_platform().megatron
+    original_import = npu.importlib.import_module
+
+    def import_module(name, *args, **kwargs):
+        if name in modules:
+            return modules[name]
+        return original_import(name, *args, **kwargs)
+
+    monkeypatch.setattr(npu.importlib, "import_module", import_module)
+    bootstrap.repatch(args)
+    assert events == ["remove", ("pre", full_args), "features"]
+    assert attention_class.forward is vime_forward
+    assert args.vllm_additional_config is typed_config
diff --git a/tests/utils/test_npu_sync_scripts.py b/tests/utils/test_npu_sync_scripts.py
new file mode 100644
index 000000000..10f6ccaac
--- /dev/null
+++ b/tests/utils/test_npu_sync_scripts.py
@@ -0,0 +1,252 @@
+"""CPU contracts for the S7 checkpoint test modes and patch ordering."""
+
+import ast
+import importlib.util
+import shlex
+import textwrap
+from pathlib import Path
+
+import pytest
+
+REPO = Path(__file__).resolve().parents[2]
+
+
+@pytest.fixture
+def glm_loader(monkeypatch, tmp_path):
+    monkeypatch.setenv("HF_HOME", str(tmp_path))
+    monkeypatch.delenv("VIME_TEST_GLM_MTP", raising=False)
+    monkeypatch.delenv("VIME_TEST_GLM_EAGER", raising=False)
+
+    def load():
+        spec = importlib.util.spec_from_file_location("glm_npu_case", REPO / "tests/test_glm4.7_30B_A3B_npu.py")
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+        return module
+
+    return load
+
+
+def test_glm_g1_uses_native_non_colocate_without_mtp(glm_loader, monkeypatch):
+    glm = glm_loader()
+    launches = []
+    monkeypatch.setattr(glm.U, "execute_train", lambda **kwargs: launches.append(kwargs))
+    glm.execute()
+    launch = launches[0]
+    tokens = shlex.split(launch["train_args"])
+    expected = {
+        "--hf-checkpoint": glm.MODEL_DIR,
+        "--ref-load": glm.MODEL_DIR,
+        "--actor-num-gpus-per-node": "8",
+        "--rollout-num-gpus": "8",
+        "--rollout-num-gpus-per-engine": "4",
+        "--tensor-model-parallel-size": "4",
+        "--expert-model-parallel-size": "8",
+        "--rope-type": "rope",
+        "--vllm-additional-config": '{"weight_nz_mode":0}',
+        "--num-rollout": "2",
+    }
+    for flag, value in expected.items():
+        assert tokens[tokens.index(flag) + 1] == value
+    assert "--vllm-enable-expert-parallel" in tokens
+    assert "--ci-test" in tokens
+    assert "--vllm-enforce-eager" not in tokens
+    assert not {
+        "--colocate",
+        "--megatron-to-hf-mode",
+        "--mtp-num-layers",
+        "--enable-mtp-training",
+        "--vllm-speculative-config",
+        "--dspark-enabled",
+    }.intersection(tokens)
+    assert launch["num_gpus_per_node"] == 16
+    assert launch["megatron_model_type"] == "glm4.7-30B-A3B"
+
+
+@pytest.mark.parametrize("mtp,eager", [("0", "0"), ("0", "1"), ("1", "0"), ("1", "1")])
+def test_glm_explicit_mtp_and_eager_modes(glm_loader, monkeypatch, mtp, eager):
+    monkeypatch.setenv("VIME_TEST_GLM_MTP", mtp)
+    monkeypatch.setenv("VIME_TEST_GLM_EAGER", eager)
+    glm = glm_loader()
+    launches = []
+    monkeypatch.setattr(glm.U, "execute_train", lambda **kwargs: launches.append(kwargs))
+    glm.execute()
+    tokens = shlex.split(launches[0]["train_args"])
+    assert ("--vllm-enforce-eager" in tokens) == (eager == "1")
+    for flag in (
+        "--mtp-num-layers",
+        "--enable-mtp-training",
+        "--mtp-loss-scaling-factor",
+        "--vllm-speculative-config",
+    ):
+        assert (flag in tokens) == (mtp == "1")
+    if mtp == "1":
+        assert tokens[tokens.index("--mtp-num-layers") + 1] == "1"
+        assert tokens[tokens.index("--mtp-loss-scaling-factor") + 1] == "0.2"
+        assert tokens[tokens.index("--vllm-speculative-config") + 1] == '{"method":"mtp","num_speculative_tokens":1}'
+    assert "--colocate" not in tokens
+    assert "--dspark-enabled" not in tokens
+    assert tokens[tokens.index("--actor-num-gpus-per-node") + 1] == "8"
+    assert tokens[tokens.index("--rollout-num-gpus") + 1] == "8"
+    assert tokens[tokens.index("--rollout-num-gpus-per-engine") + 1] == "4"
+
+
+def test_glm_prepare_preserves_ci_download_defaults(glm_loader, monkeypatch):
+    glm = glm_loader()
+    commands = []
+    monkeypatch.setattr(glm.U, "exec_command", commands.append)
+    glm.prepare()
+    assert commands == [
+        f"mkdir -p {shlex.quote(f'{glm.TEST_ROOT}/models')} {shlex.quote(f'{glm.TEST_ROOT}/datasets')}",
+        f"hf download zai-org/GLM-4.7-Flash --local-dir {shlex.quote(glm.MODEL_DIR)}",
+        "hf download --repo-type dataset zhuzilin/dapo-math-17k " f"--local-dir {shlex.quote(glm.DATASET_DIR)}",
+    ]
+
+
+def test_glm_shell_keeps_g1_model_and_serving_configuration():
+    script = (REPO / "scripts/run-glm4.7-30B-A3B-npu.sh").read_text()
+    assert 'source "${SCRIPT_DIR}/models/glm4.7-30B-A3B.sh"' in script
+    assert "--rope-type rope" in script
+    assert "--vllm-additional-config '{\"weight_nz_mode\":0}'" in script
+    assert "--vllm-enable-expert-parallel" in script
+    for flag in (
+        "--colocate",
+        "--megatron-to-hf-mode",
+        "--mtp-num-layers",
+        "--enable-mtp-training",
+        "--vllm-speculative-config",
+    ):
+        assert flag not in script
+
+
+@pytest.fixture
+def qwen30(monkeypatch, tmp_path):
+    monkeypatch.setenv("HF_HOME", str(tmp_path))
+    (tmp_path / "models").mkdir()
+    spec = importlib.util.spec_from_file_location("qwen30_npu_case", REPO / "tests/test_qwen3_30B_A3B_npu.py")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def test_explicit_30b_hf_mode_skips_conversion(qwen30, monkeypatch):
+    commands = []
+    launches = []
+    monkeypatch.setattr(qwen30.U, "exec_command", commands.append)
+    monkeypatch.setattr(qwen30.U, "execute_train", lambda **kwargs: launches.append(kwargs))
+    assert qwen30.prepare(torch_dist_ref_load=False) is None
+    qwen30.execute()
+    assert not any("torch.distributed.run" in cmd or "rm -rf" in cmd for cmd in commands)
+    args = launches[0]["train_args"]
+    assert f"--ref-load {shlex.quote(qwen30.MODEL_DIR)} " in args
+    assert "--colocate " in args
+    assert "--tensor-model-parallel-size 4 " in args
+    assert "--expert-model-parallel-size 8 " in args
+    assert "weight_nz_mode" in args
+
+
+def test_default_torch_dist_mode_uses_new_output_and_ref_load(qwen30, monkeypatch, tmp_path):
+    commands = []
+    launches = []
+    existing = tmp_path / "models/Qwen3-30B-A3B_torch_dist"
+    existing.mkdir()
+    sentinel = existing / "keep"
+    sentinel.write_text("existing checkpoint")
+
+    def execute(command):
+        commands.append(command)
+        if "torch.distributed.run" in command:
+            tokens = shlex.split(command)
+            target = Path(tokens[tokens.index("--save") + 1])
+            (target / "latest_checkpointed_iteration.txt").write_text("release")
+            (target / ".metadata").write_bytes(b"test fixture")
+
+    monkeypatch.setattr(qwen30.U, "exec_command", execute)
+    monkeypatch.setattr(qwen30.U, "execute_train", lambda **kwargs: launches.append(kwargs))
+    checkpoint = qwen30.prepare()
+    qwen30.execute(checkpoint)
+    assert Path(checkpoint) != existing
+    assert sentinel.read_text() == "existing checkpoint"
+    assert not any("rm -rf" in command for command in commands)
+    conversion = next(command for command in commands if "torch.distributed.run" in command)
+    assert "VIME_PLATFORM=npu" in conversion
+    assert "--nproc-per-node 8 " in conversion
+    args = launches[0]["train_args"]
+    assert f"--ref-load {shlex.quote(checkpoint)} " in args
+    assert "--load " not in args
+    assert "--colocate " in args
+    assert "weight_nz_mode" in args
+
+
+@pytest.mark.parametrize("override,enabled", [(None, True), ("1", True), ("0", False)])
+def test_30b_main_checkpoint_mode(qwen30, monkeypatch, override, enabled):
+    monkeypatch.delenv("VIME_TEST_TORCH_DIST_REF_LOAD", raising=False)
+    if override is not None:
+        monkeypatch.setenv("VIME_TEST_TORCH_DIST_REF_LOAD", override)
+    for key in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"):
+        monkeypatch.delenv(key, raising=False)
+    checkpoint = object()
+    launches = []
+
+    def prepare(*, torch_dist_ref_load):
+        assert torch_dist_ref_load is enabled
+        return checkpoint if enabled else None
+
+    monkeypatch.setattr(qwen30, "prepare", prepare)
+    monkeypatch.setattr(qwen30, "execute", launches.append)
+    qwen30.main()
+    assert launches == [checkpoint if enabled else None]
+
+
+def test_converter_bootstraps_before_first_megatron_import():
+    tree = ast.parse((REPO / "tools/convert_hf_to_torch_dist.py").read_text())
+    first_megatron = next(
+        node.lineno
+        for node in tree.body
+        if isinstance(node, ast.ImportFrom) and (node.module or "").startswith("megatron.")
+    )
+    bootstrap = next(
+        node.lineno
+        for node in ast.walk(tree)
+        if isinstance(node, ast.Import) and any(alias.name == "vime.backends.megatron_utils" for alias in node.names)
+    )
+    assert bootstrap < first_megatron
+    assert "vime.utils.common" not in ast.unparse(tree)
+
+
+def test_common_megatron_patch_is_snapshotted_before_npu_patch():
+    entries = [
+        line.split("|")
+        for line in (REPO / "docker/npu_patch/series.conf").read_text().splitlines()
+        if line and not line.startswith("#")
+    ]
+    megatron = [entry for entry in entries if entry[0] == "/root/Megatron-LM"]
+    assert megatron == [
+        ["/root/Megatron-LM", "megatron-common.patch", "docker/patch/latest/megatron.patch"],
+        ["/root/Megatron-LM", "megatron.patch", "docker/npu_patch/megatron.patch"],
+    ]
+    dockerfile = (REPO / "docker/Dockerfile.npu").read_text()
+    assert "COPY docker/patch/latest/megatron.patch /opt/npu_patch/megatron-common.patch" in dockerfile
+    assert "/opt/vime_patch/megatron.patch" not in dockerfile
+
+
+def _megatron_patch_additions(path, patch_path="docker/npu_patch/megatron.patch"):
+    patch = (REPO / patch_path).read_text()
+    section = patch.split(f"diff --git a/{path} b/{path}\n", 1)[1].split("diff --git ", 1)[0]
+    return "\n".join(line[1:] for line in section.splitlines() if line.startswith("+") and not line.startswith("+++"))
+
+
+def test_npu_patch_keeps_public_transformer_layer():
+    patch = (REPO / "docker/npu_patch/megatron.patch").read_text()
+    assert "diff --git a/megatron/core/transformer/transformer_layer.py " not in patch
+
+
+def test_post_layernorm_flags_remain_dataclass_generated():
+    additions = _megatron_patch_additions(
+        "megatron/core/transformer/transformer_config.py", "docker/patch/latest/megatron.patch"
+    )
+    fields = ast.parse(textwrap.dedent(additions)).body
+    defaults = {node.target.id: ast.literal_eval(node.value) for node in fields if isinstance(node, ast.AnnAssign)}
+    npu_patch = (REPO / "docker/npu_patch/megatron.patch").read_text()
+    for name in ("post_self_attn_layernorm", "post_mlp_layernorm"):
+        assert defaults[name] is False
+        assert f"--{name.replace('_', '-')}" not in npu_patch
diff --git a/tests/utils/test_platform_contract.py b/tests/utils/test_platform_contract.py
new file mode 100644
index 000000000..c930fed82
--- /dev/null
+++ b/tests/utils/test_platform_contract.py
@@ -0,0 +1,237 @@
+from __future__ import annotations
+
+import sys
+from argparse import Namespace
+from types import ModuleType, SimpleNamespace
+
+import pytest
+
+from vime.platforms import current_platform, reset_platform_cache
+
+
+@pytest.fixture(autouse=True)
+def _reset_platform_selection():
+    reset_platform_cache()
+    yield
+    reset_platform_cache()
+
+
+def test_vime_platform_override_selects_cuda(monkeypatch):
+    monkeypatch.setenv("VIME_PLATFORM", "cuda")
+
+    platform = current_platform()
+
+    assert platform.name == "cuda"
+    assert platform.ray.resource_name == "GPU"
+    assert platform.ray.visible_devices_env == "CUDA_VISIBLE_DEVICES"
+    assert platform.checkpoint.default_megatron_to_hf_mode == "raw"
+
+
+def test_vime_platform_override_selects_npu_without_vendor_import(monkeypatch):
+    monkeypatch.setenv("VIME_PLATFORM", "npu")
+    before = {name for name in ("torch_npu", "vllm_ascend", "mindspeed") if name in sys.modules}
+
+    platform = current_platform()
+
+    after = {name for name in ("torch_npu", "vllm_ascend", "mindspeed") if name in sys.modules}
+    assert platform.name == "npu"
+    assert platform.ray.resource_name == "NPU"
+    assert platform.checkpoint.default_megatron_to_hf_mode == "bridge"
+    assert before == after
+
+
+def test_unknown_explicit_platform_has_clear_error(monkeypatch):
+    monkeypatch.setenv("VIME_PLATFORM", "not-registered")
+
+    with pytest.raises(ValueError, match="Unknown Vime platform"):
+        current_platform()
+
+
+def test_npu_auto_detection_does_not_import_vendor_without_device_nodes(monkeypatch):
+    from vime.platforms import npu
+
+    monkeypatch.setattr(npu.os.path, "exists", lambda path: False)
+    monkeypatch.setattr(npu, "glob", lambda pattern: [])
+
+    def fail_import(name):
+        raise AssertionError(f"unexpected import during negative NPU detection: {name}")
+
+    monkeypatch.setattr(npu.importlib, "import_module", fail_import)
+
+    assert npu.detect_npu() is False
+
+
+def test_npu_safe_empty_cache_wraps_original_once(monkeypatch):
+    from vime.platforms import npu
+
+    calls = []
+
+    def original_empty_cache():
+        calls.append("original")
+        raise RuntimeError("allocator is between offload states")
+
+    fake_torch = SimpleNamespace(
+        npu=SimpleNamespace(empty_cache=original_empty_cache),
+        cuda=SimpleNamespace(empty_cache=lambda: None),
+    )
+    monkeypatch.setattr(npu.importlib, "import_module", lambda name: fake_torch if name == "torch" else None)
+
+    npu._install_safe_empty_cache()
+    wrapped = fake_torch.npu.empty_cache
+    wrapped()
+    npu._install_safe_empty_cache()
+
+    assert calls == ["original"]
+    assert fake_torch.npu.empty_cache is wrapped
+    assert fake_torch.cuda.empty_cache is wrapped
+
+
+@pytest.mark.parametrize(
+    ("name", "bundle", "actor_options"),
+    [
+        ("cuda", {"GPU": 2, "CPU": 3}, {"num_gpus": 0.4}),
+        ("npu", {"NPU": 2, "CPU": 3}, {"resources": {"NPU": 0.4}}),
+    ],
+)
+def test_ray_resource_contract(monkeypatch, name, bundle, actor_options):
+    monkeypatch.setenv("VIME_PLATFORM", name)
+    ray_spec = current_platform().ray
+
+    assert ray_spec.bundle_resources(device_count=2, cpu_count=3) == bundle
+    assert ray_spec.actor_options(0.4) == actor_options
+
+
+def test_npu_runtime_env_is_scoped_to_npu_provider(monkeypatch, tmp_path):
+    toolkit = tmp_path / "toolkit"
+    (toolkit / "python" / "site-packages" / "acl").mkdir(parents=True)
+    monkeypatch.setenv("ASCEND_TOOLKIT_HOME", str(toolkit))
+    monkeypatch.setenv("VIME_PLATFORM", "npu")
+    args = Namespace(offload_train=True, train_backend="megatron", colocate=True)
+
+    train_env = current_platform().ray.train_runtime_env(args, {"BASE": "1"})
+    rollout_env = current_platform().ray.rollout_runtime_env(args, {"BASE": "1"})
+
+    assert train_env["TMS_HOOK_MODE"] == "torch"
+    assert train_env["TMS_REGION_TAG"] == "training"
+    assert train_env["TMS_ENABLE_CPU_BACKUP"] == "1"
+    assert train_env["PYTORCH_NPU_ALLOC_CONF"] == "expandable_segments:False"
+    assert str(toolkit / "python" / "site-packages") in train_env["PYTHONPATH"]
+    assert rollout_env["VLLM_USE_AOT_COMPILE"] == "0"
+    assert rollout_env["PYTORCH_NPU_ALLOC_CONF"] == "expandable_segments:False"
+
+
+def test_npu_vllm_env_replaces_cuda_and_rocm_visibility(monkeypatch):
+    monkeypatch.setenv("VIME_PLATFORM", "npu")
+    platform = current_platform()
+
+    env = platform.vllm.subprocess_env(
+        {
+            "KEEP": "1",
+            "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
+            "CUDA_VISIBLE_DEVICES": "0,1",
+            "HIP_VISIBLE_DEVICES": "0,1",
+        },
+        visible_devices="4,5",
+        colocate=True,
+    )
+
+    assert env["KEEP"] == "1"
+    assert "PYTORCH_CUDA_ALLOC_CONF" not in env
+    assert "CUDA_VISIBLE_DEVICES" not in env
+    assert "HIP_VISIBLE_DEVICES" not in env
+    assert env["ASCEND_RT_VISIBLE_DEVICES"] == "4,5"
+    assert env["PYTORCH_NPU_ALLOC_CONF"] == "expandable_segments:False"
+
+
+@pytest.mark.parametrize("name,backends", [("cuda", ("nccl", "ipc")), ("npu", ("hccl", "npu_ipc"))])
+def test_weight_transfer_provider_selects_matching_backend_and_init_info(monkeypatch, name, backends):
+    from vime.platforms import get_platform
+
+    plugin_calls = []
+    for colocate, backend in zip((False, True), backends, strict=True):
+        if name == "cuda":
+            module_name = f"vllm.distributed.weight_transfer.{backend}_engine"
+            cls_name = "IPCTrainerInitInfo" if colocate else "NCCLTrainerInitInfo"
+        else:
+            module_name = f"vllm_ascend.distributed.weight_transfer.{backend}_engine"
+            cls_name = "NPUIPCTrainerInitInfo" if colocate else "HCCLTrainerInitInfo"
+        module = ModuleType(module_name)
+        info_cls = type(cls_name, (SimpleNamespace,), {"backend": backend})
+        setattr(module, cls_name, info_cls)
+        monkeypatch.setitem(sys.modules, module_name, module)
+
+        plugins = ModuleType("vllm.plugins")
+        plugins.load_general_plugins = lambda: plugin_calls.append("load")
+        monkeypatch.setitem(sys.modules, "vllm.plugins", plugins)
+        monkeypatch.setitem(sys.modules, "torch_npu", ModuleType("torch_npu"))
+
+        ops = get_platform(name).weight_transfer
+        info = ops.trainer_init_info(colocate=colocate, rank=3, packed=True)
+        assert ops.backend("ipc" if colocate else "nccl") == info.backend == backend
+        assert isinstance(info, info_cls)
+        assert (info.rank, info.packed) == (3, True)
+
+    assert plugin_calls == (["load", "load"] if name == "npu" else [])
+    for backend in ("npu_ipc", "hccl", "custom"):
+        assert ops.backend(backend) == backend
+
+
+@pytest.mark.parametrize("platform", ["cuda", "npu"])
+@pytest.mark.parametrize("chained", [False, True])
+def test_optimizer_state_initialization_reuses_megatron_callback(monkeypatch, platform, chained):
+    monkeypatch.setenv("VIME_PLATFORM", platform)
+    calls = []
+
+    def init_state(optimizer, config):
+        calls.append((optimizer, config))
+
+    optimizers = [
+        SimpleNamespace(optimizer=object(), config=object(), init_state_fn=init_state)
+        for _ in range(2 if chained else 1)
+    ]
+    optimizer = SimpleNamespace(chained_optimizers=optimizers) if chained else optimizers[0]
+
+    current_platform().megatron.initialize_optimizer_state(optimizer)
+
+    expected = [(opt.optimizer, opt.config) for opt in optimizers] if platform == "npu" else []
+    assert calls == expected
+
+
+@pytest.mark.parametrize("empty_optimizer", [False, True])
+def test_npu_optimizer_state_initialization_skips_missing_state_or_optimizer(monkeypatch, empty_optimizer):
+    monkeypatch.setenv("VIME_PLATFORM", "npu")
+
+    def unexpected_init(*args):
+        raise AssertionError("an empty optimizer must not initialize state")
+
+    optimizer = SimpleNamespace(
+        optimizer=None if empty_optimizer else object(),
+        config=object(),
+        init_state_fn=unexpected_init if empty_optimizer else None,
+    )
+
+    current_platform().megatron.initialize_optimizer_state(optimizer)
+
+
+@pytest.mark.parametrize("platform", ["cuda", "npu"])
+def test_eval_only_has_no_optimizer_state_to_initialize(monkeypatch, platform):
+    monkeypatch.setenv("VIME_PLATFORM", platform)
+    current_platform().megatron.initialize_optimizer_state(None)
+
+
+def test_memory_utils_keep_main_accelerator_surface(monkeypatch):
+    from vime.utils import memory_utils
+
+    calls = []
+    fake_accelerator = SimpleNamespace(
+        synchronize=lambda: calls.append("synchronize"),
+        empty_cache=lambda: calls.append("empty_cache"),
+    )
+    fake_torch = SimpleNamespace(_C=SimpleNamespace(_host_emptyCache=lambda: calls.append("empty_host_cache")))
+    monkeypatch.setattr(memory_utils, "accelerator", fake_accelerator)
+    monkeypatch.setattr(memory_utils, "torch", fake_torch)
+    monkeypatch.setattr(memory_utils.gc, "collect", lambda: calls.append("gc"))
+
+    memory_utils.clear_memory(clear_host_memory=True)
+
+    assert calls == ["synchronize", "gc", "empty_cache", "empty_host_cache"]
diff --git a/tests/utils/test_ray_platform_integration.py b/tests/utils/test_ray_platform_integration.py
new file mode 100644
index 000000000..e5d19dcf4
--- /dev/null
+++ b/tests/utils/test_ray_platform_integration.py
@@ -0,0 +1,274 @@
+from __future__ import annotations
+
+import sys
+import types
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+
+
+def _fake_platform(ray_ops, *, is_npu=False):
+    return SimpleNamespace(ray=ray_ops, is_npu=is_npu)
+
+
+@pytest.mark.parametrize(
+    ("platform_name", "visible_env", "assigned_id", "expected_resource"),
+    [
+        ("cuda", "CUDA_VISIBLE_DEVICES", "7", "GPU"),
+        ("npu", "ASCEND_RT_VISIBLE_DEVICES", "9", "NPU"),
+    ],
+)
+def test_platform_ray_accelerator_ids_and_local_mapping(
+    monkeypatch,
+    platform_name,
+    visible_env,
+    assigned_id,
+    expected_resource,
+):
+    import ray
+
+    from vime.platforms import get_platform
+
+    platform = get_platform(platform_name)
+    monkeypatch.setenv(visible_env, f"unused,{assigned_id}")
+    if platform_name == "cuda":
+        monkeypatch.setattr(ray, "get_gpu_ids", lambda: [assigned_id])
+    else:
+        context = SimpleNamespace(get_accelerator_ids=lambda: {"NPU": [assigned_id]})
+        monkeypatch.setattr(ray, "get_runtime_context", lambda: context)
+
+    assert platform.ray.resource_name == expected_resource
+    assert platform.ray.accelerator_ids() == [assigned_id]
+    assert platform.ray.local_device_id() == 1
+
+
+def test_placement_group_uses_platform_ray_resource_contract(monkeypatch):
+    from vime.ray import placement_group as placement_group_module
+
+    ray_ops = SimpleNamespace(
+        resource_name="ACCEL",
+        bundle_resources=Mock(side_effect=lambda: {"ACCEL": 1, "CPU": 1}),
+        actor_options=Mock(side_effect=lambda fraction: {"resources": {"ACCEL": fraction}}),
+    )
+    monkeypatch.setattr(placement_group_module, "current_platform", lambda: _fake_platform(ray_ops, is_npu=True))
+
+    created = {}
+
+    class FakePlacementGroup:
+        def ready(self):
+            return "ready"
+
+    def fake_placement_group(bundles, strategy):
+        created["bundles"] = bundles
+        created["strategy"] = strategy
+        return FakePlacementGroup()
+
+    actor_options = []
+
+    class FakeInfoActor:
+        def __init__(self, result):
+            self.get_ip_and_gpu_id = SimpleNamespace(remote=lambda: result)
+
+    class FakeInfoActorClass:
+        results = iter([("10.0.0.1", "3"), ("10.0.0.1", "1")])
+
+        @classmethod
+        def options(cls, **options):
+            actor_options.append(options)
+            result = next(cls.results)
+            return SimpleNamespace(remote=lambda: FakeInfoActor(result))
+
+    monkeypatch.setattr(placement_group_module, "placement_group", fake_placement_group)
+    monkeypatch.setattr(placement_group_module, "InfoActor", FakeInfoActorClass)
+    wait_results = iter([([], ["ready"]), (["ready"], [])])
+    monkeypatch.setattr(placement_group_module.ray, "wait", lambda *_args, **_kwargs: next(wait_results))
+    monkeypatch.setattr(placement_group_module.ray, "cluster_resources", lambda: {"ACCEL": 2})
+    monkeypatch.setattr(placement_group_module.ray, "available_resources", lambda: {"ACCEL": 2})
+    monkeypatch.setattr(placement_group_module.ray, "get", lambda value: value)
+    monkeypatch.setattr(placement_group_module.ray, "kill", lambda _actor: None)
+
+    pg, reordered_indices, reordered_ids = placement_group_module._create_placement_group(2)
+
+    assert isinstance(pg, FakePlacementGroup)
+    assert created == {
+        "bundles": [{"ACCEL": 1, "CPU": 1}, {"ACCEL": 1, "CPU": 1}],
+        "strategy": "PACK",
+    }
+    assert reordered_indices == [1, 0]
+    assert reordered_ids == ["1", "3"]
+    assert [options["resources"] for options in actor_options] == [{"ACCEL": 1}, {"ACCEL": 1}]
+    assert [options["num_gpus"] for options in actor_options] == [0, 0]
+    assert ray_ops.bundle_resources.call_count == 2
+    assert [entry.args for entry in ray_ops.actor_options.call_args_list] == [(1,), (1,)]
+
+
+def test_ray_noset_visible_devices_keeps_ascend_entry():
+    from vime.ray.utils import NOSET_VISIBLE_DEVICES_ENV_VARS_LIST
+
+    assert "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES" in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST
+
+
+def test_train_group_uses_npu_runtime_env_and_actor_resources(monkeypatch):
+    from vime.ray import actor_group as actor_group_module
+
+    runtime_env_inputs = []
+    ray_ops = SimpleNamespace(
+        train_runtime_env=lambda args, env: runtime_env_inputs.append((args, dict(env)))
+        or {**env, "PLATFORM_ENV": "1"},
+        actor_options=Mock(side_effect=lambda fraction: {"resources": {"ACCEL": fraction}}),
+    )
+    monkeypatch.setattr(actor_group_module, "current_platform", lambda: _fake_platform(ray_ops, is_npu=True))
+
+    actor_module = types.ModuleType("vime.backends.megatron_utils.actor")
+    actor_module.MegatronTrainRayActor = object
+    monkeypatch.setitem(sys.modules, "vime.backends.megatron_utils.actor", actor_module)
+
+    remote_declarations = []
+    actor_allocations = []
+
+    class FakeActorHandle:
+        get_master_addr_and_port = SimpleNamespace(remote=lambda: ("127.0.0.1", 20000))
+        init = SimpleNamespace(remote=lambda *_args, **_kwargs: 0)
+
+    class FakeRemoteActor:
+        def options(self, **options):
+            actor_allocations.append(options)
+            return self
+
+        def remote(self, *_args):
+            return FakeActorHandle()
+
+    def fake_remote(**options):
+        remote_declarations.append(options)
+        return lambda _actor_impl: FakeRemoteActor()
+
+    monkeypatch.setattr(actor_group_module.ray, "remote", fake_remote)
+    monkeypatch.setattr(actor_group_module.ray, "get", lambda value: value)
+
+    args = SimpleNamespace(
+        train_env_vars={"USER_ENV": "yes"},
+        offload_train=True,
+        train_backend="megatron",
+        colocate=False,
+        use_routing_replay=False,
+    )
+    group = actor_group_module.RayTrainGroup(
+        args=args,
+        num_nodes=1,
+        num_gpus_per_node=1,
+        pg=(object(), [0], [7]),
+        num_gpus_per_actor=0.4,
+    )
+    assert group.create() == [0]
+
+    assert runtime_env_inputs[0][0] is args
+    assert runtime_env_inputs[0][1]["USER_ENV"] == "yes"
+    assert remote_declarations[0]["runtime_env"]["env_vars"]["PLATFORM_ENV"] == "1"
+    assert remote_declarations[0]["num_gpus"] == 1
+    assert actor_allocations[0]["resources"] == {"ACCEL": 0.4}
+    assert actor_allocations[0]["num_gpus"] == 0
+    ray_ops.actor_options.assert_called_once_with(0.4)
+
+
+@pytest.mark.parametrize("is_npu", [False, True])
+def test_rollout_engine_delegates_runtime_env_and_actor_options(monkeypatch, is_npu):
+    from vime.backends.vllm_utils import engine_group as rollout_module
+
+    runtime_env_inputs = []
+    ray_ops = SimpleNamespace(
+        rollout_runtime_env=lambda args, env: runtime_env_inputs.append((args, dict(env)))
+        or {**env, "PLATFORM_ENV": "rollout"},
+        actor_options=Mock(side_effect=lambda fraction: {"resources": {"ACCEL": fraction}}),
+    )
+    monkeypatch.setattr(rollout_module, "current_platform", lambda: _fake_platform(ray_ops, is_npu=is_npu))
+
+    actor_options = []
+
+    class FakeEngine:
+        init = SimpleNamespace(remote=lambda **_kwargs: "init-ref")
+
+    class FakeRemoteActor:
+        def options(self, **options):
+            actor_options.append(options)
+            return self
+
+        def remote(self, *_args, **_kwargs):
+            return FakeEngine()
+
+    monkeypatch.setattr(rollout_module.ray, "remote", lambda _actor_impl: FakeRemoteActor())
+    monkeypatch.setattr(
+        rollout_module,
+        "_allocate_rollout_engine_addr_and_ports_normal",
+        lambda **_kwargs: ({0: {}}, {0: 15001}),
+    )
+
+    args = SimpleNamespace(
+        debug_train_only=False,
+        num_gpus_per_node=8,
+        rollout_num_gpus=1,
+        rollout_num_gpus_per_engine=1,
+        rollout_external=False,
+        colocate=True,
+    )
+    group = rollout_module.ServerGroup(
+        args=args,
+        pg=(object(), [0], [7]),
+        all_engines=[None],
+        num_gpus_per_engine=1,
+        num_new_engines=1,
+    )
+
+    handles, cursors = group.start_engines()
+
+    assert handles == ["init-ref"]
+    assert cursors == {0: 15001}
+    if is_npu:
+        assert runtime_env_inputs[0][0] is args
+        assert actor_options[0]["runtime_env"]["env_vars"]["PLATFORM_ENV"] == "rollout"
+        assert actor_options[0]["resources"] == {"ACCEL": 0.2}
+        assert actor_options[0]["num_gpus"] == 0
+        ray_ops.actor_options.assert_called_once_with(0.2)
+    else:
+        assert runtime_env_inputs == []
+        assert "resources" not in actor_options[0]
+        assert "PLATFORM_ENV" not in actor_options[0]["runtime_env"]["env_vars"]
+        assert actor_options[0]["num_gpus"] == 0.2
+        ray_ops.actor_options.assert_not_called()
+
+    # The main placement check must still run before an invalid slot is used.
+    group.all_engines = [None]
+    group.gpu_offset = 1
+    with pytest.raises(ValueError, match="Invalid rollout server group GPU placement"):
+        group.start_engines()
+
+
+def test_train_actor_uses_npu_local_device_mapping(monkeypatch):
+    from vime.ray import train_actor as train_actor_module
+
+    local_device_id = Mock(return_value=5)
+    monkeypatch.setattr(
+        train_actor_module,
+        "current_platform",
+        lambda: _fake_platform(SimpleNamespace(local_device_id=local_device_id), is_npu=True),
+    )
+
+    assert train_actor_module.get_local_gpu_id() == 5
+    local_device_id.assert_called_once_with()
+
+
+def test_train_actor_keeps_main_cuda_local_device_mapping(monkeypatch):
+    from vime.ray import train_actor as train_actor_module
+
+    monkeypatch.setattr(
+        train_actor_module,
+        "current_platform",
+        lambda: _fake_platform(SimpleNamespace(), is_npu=False),
+    )
+    monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "7,3")
+    monkeypatch.setattr(
+        train_actor_module.accelerator, "_ACCELERATOR", train_actor_module.accelerator.CUDAAccelerator()
+    )
+    monkeypatch.setattr(train_actor_module.ray, "get_gpu_ids", lambda: ["3"])
+
+    assert train_actor_module.get_local_gpu_id() == 1
diff --git a/tests/utils/test_update_weight_from_distributed.py b/tests/utils/test_update_weight_from_distributed.py
index d456cc226..f63a72b2b 100644
--- a/tests/utils/test_update_weight_from_distributed.py
+++ b/tests/utils/test_update_weight_from_distributed.py
@@ -187,7 +187,10 @@ def get_hf_weight_chunks(self, weights):
 
 @pytest.mark.unit
 def test_nccl_trainer_uses_single_packed_buffer(update_module, monkeypatch):
+    from vime.platforms import get_platform
+
     adapter = sys.modules[update_module.create_nccl_trainer.__module__]
+    monkeypatch.setattr(adapter, "current_platform", lambda: get_platform("cuda"))
     created = []
 
     class NCCLTrainerInitInfo:
diff --git a/tests/utils/test_update_weight_from_tensor.py b/tests/utils/test_update_weight_from_tensor.py
index 815c9dd35..0e8ac0d61 100644
--- a/tests/utils/test_update_weight_from_tensor.py
+++ b/tests/utils/test_update_weight_from_tensor.py
@@ -188,6 +188,9 @@ def trainer_init(init_info, *, client, source):
 
 @pytest.mark.unit
 def test_connect_uses_native_ipc_and_nccl_trainers(update_module, monkeypatch):
+    from vime.platforms import get_platform
+
+    monkeypatch.setattr(update_module, "current_platform", lambda: get_platform("cuda"))
     updater = _updater(update_module)
     old_trainer = RecordingTrainer(object())
     updater._native_trainers = [old_trainer]
diff --git a/tests/utils/test_vllm_arguments.py b/tests/utils/test_vllm_arguments.py
index 743f73831..6eeb10424 100644
--- a/tests/utils/test_vllm_arguments.py
+++ b/tests/utils/test_vllm_arguments.py
@@ -5,7 +5,7 @@
 import argparse
 import sys
 from pathlib import Path
-from types import SimpleNamespace
+from types import ModuleType, SimpleNamespace
 
 _tests_root = Path(__file__).resolve().parents[1]
 if str(_tests_root) not in sys.path:
@@ -22,6 +22,20 @@
 NUM_GPUS = 0
 
 
+@pytest.mark.unit
+@pytest.mark.parametrize("preloaded", [False, True])
+def test_real_module_available_rejects_missing_package_and_stub(monkeypatch, preloaded):
+    name = "_vime_missing_optional_dependency"
+    if preloaded:
+        monkeypatch.setitem(sys.modules, name, ModuleType(name))
+    assert not _unit_stubs.real_module_available(name)
+
+
+@pytest.mark.unit
+def test_real_module_available_accepts_loaded_real_module():
+    assert _unit_stubs.real_module_available("sys")
+
+
 @pytest.fixture(scope="module")
 def args_mod():
     from vime.backends.vllm_utils import arguments as mod  # noqa: PLC0415
@@ -172,11 +186,13 @@ def test_add_vllm_arguments_overrides_router_balance_threshold_defaults(args_mod
 
 
 def _patch_device_config(monkeypatch):
-    """Patch DeviceConfig.__post_init__ to avoid GPU device detection on CPU CI."""
+    """Avoid device detection and third-party plugin loading in parser tests."""
     try:
         from vllm.config.device import DeviceConfig
+        from vllm.engine import arg_utils
 
         monkeypatch.setattr(DeviceConfig, "__post_init__", lambda self: setattr(self, "device_type", "cpu"))
+        monkeypatch.setattr(arg_utils, "load_general_plugins", lambda: None)
     except ImportError:
         pass
 
diff --git a/tests/utils/test_vllm_engine.py b/tests/utils/test_vllm_engine.py
index 5f599b397..bd540b0a5 100644
--- a/tests/utils/test_vllm_engine.py
+++ b/tests/utils/test_vllm_engine.py
@@ -349,6 +349,38 @@ def test_compute_server_args_adds_sleep_mode_for_offload_rollout(vllm_args):
     assert vllm_args.vllm_enable_sleep_mode is True
 
 
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    "platform_name,colocate,backend",
+    [("cuda", False, "nccl"), ("cuda", True, "ipc"), ("npu", False, "hccl"), ("npu", True, "npu_ipc")],
+)
+def test_compute_server_args_uses_native_platform_backend(vllm_args, monkeypatch, platform_name, colocate, backend):
+    from vime.platforms import get_platform
+
+    monkeypatch.setattr(mod, "_VLLM_SERVER_FIELDS", frozenset({"worker_extension_cls"}))
+    vllm_args.vllm_worker_extension_cls = ""
+    vllm_args.colocate = colocate
+    monkeypatch.setattr(mod, "current_platform", lambda: get_platform(platform_name))
+
+    sa, _ = mod._compute_server_args(vllm_args, rank=0, dist_init_addr=None, host="127.0.0.1", port=8000)
+
+    assert not sa.get("worker_extension_cls")
+    assert sa["weight_transfer_config"] == {"backend": backend}
+
+
+@pytest.mark.unit
+def test_compute_server_args_keeps_user_worker_extension(vllm_args, monkeypatch):
+    from vime.platforms import get_platform
+
+    monkeypatch.setattr(mod, "_VLLM_SERVER_FIELDS", frozenset({"worker_extension_cls"}))
+    vllm_args.vllm_worker_extension_cls = "example.UserWorkerExtension"
+    monkeypatch.setattr(mod, "current_platform", lambda: get_platform("npu"))
+
+    sa, _ = mod._compute_server_args(vllm_args, rank=0, dist_init_addr=None, host="127.0.0.1", port=8000)
+
+    assert sa["worker_extension_cls"] == "example.UserWorkerExtension"
+
+
 @pytest.mark.unit
 def test_compute_server_args_no_sleep_mode_from_colocate(vllm_args):
     vllm_args.colocate = True
@@ -454,6 +486,7 @@ def fake_post(endpoint: str, payload: dict):
 
     assert posted[0][0] == "update_weights"
     sent = posted[0][1]["update_info"]
+    # ipc_handles are pickled for native vLLM/vLLM-Ascend parse_update_info.
     # ipc_handles got cloudpickle'd into ipc_handles_pickled
     assert "ipc_handles" not in sent
     assert isinstance(sent["ipc_handles_pickled"], str)
diff --git a/tools/convert_hf_to_torch_dist.py b/tools/convert_hf_to_torch_dist.py
index 798e79ef6..45c7360ac 100644
--- a/tools/convert_hf_to_torch_dist.py
+++ b/tools/convert_hf_to_torch_dist.py
@@ -5,6 +5,11 @@
 
 import torch
 import torch.distributed as dist
+from vime.platforms import current_platform
+
+if current_platform().is_npu:
+    import vime.backends.megatron_utils  # noqa: F401
+
 from megatron.core.enums import ModelType
 from megatron.training.arguments import parse_args, validate_args
 from megatron.training.checkpointing import get_checkpoint_name, get_checkpoint_tracker_filename, save_checkpoint
diff --git a/train.py b/train.py
index 2130030f5..9b17f2ff4 100644
--- a/train.py
+++ b/train.py
@@ -1,5 +1,10 @@
 import ray
 
+from vime.platforms import current_platform
+
+if current_platform().is_npu:
+    import vime.backends.megatron_utils  # noqa: F401
+
 from vime.observability.logging_utils import configure_logger, finish_tracking, init_tracking
 from vime.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_models
 from vime.utils.arguments import parse_args
diff --git a/train_async.py b/train_async.py
index 4cbc55bbc..5d58ecfb6 100644
--- a/train_async.py
+++ b/train_async.py
@@ -1,5 +1,10 @@
 import ray
 
+from vime.platforms import current_platform
+
+if current_platform().is_npu:
+    import vime.backends.megatron_utils  # noqa: F401
+
 from vime.observability.logging_utils import configure_logger, finish_tracking, init_tracking
 from vime.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_models
 from vime.utils.arguments import parse_args
diff --git a/vime/backends/megatron_utils/__init__.py b/vime/backends/megatron_utils/__init__.py
index 471d92043..feb890366 100644
--- a/vime/backends/megatron_utils/__init__.py
+++ b/vime/backends/megatron_utils/__init__.py
@@ -2,6 +2,11 @@
 
 import torch
 
+from vime.platforms import current_platform
+
+# Load NPU prerequisites before the shared Megatron patches.
+current_platform().megatron.bootstrap()
+
 from vime.utils import accelerator
 
 accelerator.initialize_accelerator()
diff --git a/vime/backends/megatron_utils/actor.py b/vime/backends/megatron_utils/actor.py
index 13ded4186..e14682b22 100644
--- a/vime/backends/megatron_utils/actor.py
+++ b/vime/backends/megatron_utils/actor.py
@@ -16,6 +16,7 @@
 from vime.observability.logging_utils import init_tracking
 from vime.observability.profile_utils import TrainProfiler
 from vime.observability.timer import Timer, inverse_timer, timer, with_defer
+from vime.platforms import current_platform
 from vime.ray.train_actor import TrainRayActor
 from vime.utils import accelerator
 from vime.utils.data import process_rollout_data
@@ -77,6 +78,7 @@ def init(
 
         init(args)
 
+        current_platform().megatron.repatch(args)
         if is_megatron_main_rank():
             init_tracking(args, primary=False, role=role)
 
@@ -91,9 +93,12 @@ def init(
 
         dist.barrier(group=get_gloo_group())
 
-        self.model, self.optimizer, self.opt_param_scheduler, loaded_rollout_id = initialize_model_and_optimizer(
-            args, role
-        )
+        with current_platform().megatron.training_context(args.offload_train):
+            self.model, self.optimizer, self.opt_param_scheduler, loaded_rollout_id = initialize_model_and_optimizer(
+                args, role
+            )
+            if args.offload_train:
+                current_platform().megatron.initialize_optimizer_state(self.optimizer)
 
         vpp_size = mpu.get_virtual_pipeline_model_parallel_world_size() or 1
         if vpp_size > 1:
@@ -602,7 +607,11 @@ def update_weights(self) -> None:
             if dist.get_rank() == 0:
                 ray.get(self.rollout_manager.clear_updatable_num_new_engines.remote())
 
-        with torch_memory_saver.disable() if self.args.offload_train else nullcontext():
+        with (
+            torch_memory_saver.disable()
+            if (self.args.offload_train and not current_platform().is_npu)
+            else nullcontext()
+        ):
             if self.args.dspark_enabled and self.args.offload_train:
                 backup = self.weights_backuper.get("actor")
                 for name, param in named_params_and_buffers(self.args, self.model):
diff --git a/vime/backends/megatron_utils/checkpoint.py b/vime/backends/megatron_utils/checkpoint.py
index 122e97232..5ee3c7de1 100644
--- a/vime/backends/megatron_utils/checkpoint.py
+++ b/vime/backends/megatron_utils/checkpoint.py
@@ -8,6 +8,8 @@
 from megatron.training.checkpointing import save_checkpoint
 from megatron.training.global_vars import get_args
 
+from vime.platforms import current_platform
+
 try:
     # Here we patch out the `validate_non_overlapping_shards_metadata` in both functions
     # because it is really slow for large models with many shards.
@@ -19,6 +21,7 @@
     from torch.distributed._shard.sharded_tensor.shard import Shard
     from torch.distributed._shard.sharded_tensor.utils import _parse_and_validate_remote_device
     from torch.distributed._shard.sharding_spec.api import EnumerableShardingSpec
+    from torch.distributed.checkpoint import default_planner
 
     def __post_init__(self):
         pass
@@ -84,6 +87,8 @@ def _init_from_local_shards_and_global_metadata(  # type: ignore[override]
 
     ShardedTensor._init_from_local_shards_and_global_metadata = _init_from_local_shards_and_global_metadata
 
+    current_platform().checkpoint.patch_default_planner(default_planner)
+
 except ImportError:
     pass
 
diff --git a/vime/backends/megatron_utils/hf_to_megatron/__init__.py b/vime/backends/megatron_utils/hf_to_megatron/__init__.py
index 5276163f6..a5398ae97 100644
--- a/vime/backends/megatron_utils/hf_to_megatron/__init__.py
+++ b/vime/backends/megatron_utils/hf_to_megatron/__init__.py
@@ -9,6 +9,7 @@
 from .qwen3_5 import qwen3_5_hf_tensor
 from .qwen3_next import qwen3_next_hf_tensor
 from .qwen3_omni import qwen3_omni_hf_tensor
+from .qwen3_vl import qwen3_vl_hf_tensor
 
 _LOADERS = {
     "deepseek_v3": deepseek_hf_tensor,
@@ -29,6 +30,7 @@
     "qwen3_moe": qwen_moe_hf_tensor,
     "qwen3_next": qwen3_next_hf_tensor,
     "qwen3_omni_moe": qwen3_omni_hf_tensor,
+    "qwen3_vl": qwen3_vl_hf_tensor,
 }
 
 
diff --git a/vime/backends/megatron_utils/hf_to_megatron/common.py b/vime/backends/megatron_utils/hf_to_megatron/common.py
index 06b874b20..8adff4198 100644
--- a/vime/backends/megatron_utils/hf_to_megatron/common.py
+++ b/vime/backends/megatron_utils/hf_to_megatron/common.py
@@ -9,6 +9,8 @@
 import torch.nn.functional as F
 from safetensors import safe_open
 
+from vime.platforms import current_platform
+
 
 class SafetensorReader:
     def __init__(self, path: str | Path):
@@ -129,7 +131,7 @@ def shard_mcore_tensor(name: str, tensor: torch.Tensor, parameter: torch.Tensor)
         tensor,
         parallel_size=parallel_size,
         parallel_rank=parallel_rank,
-        partition_dim=parameter.partition_dim,
+        partition_dim=current_platform().megatron.adjust_tp_partition_dim(name, parameter.partition_dim),
         partition_stride=parameter.partition_stride,
     )
 
diff --git a/vime/backends/megatron_utils/hf_to_megatron/qwen3_vl.py b/vime/backends/megatron_utils/hf_to_megatron/qwen3_vl.py
new file mode 100644
index 000000000..a5b350628
--- /dev/null
+++ b/vime/backends/megatron_utils/hf_to_megatron/qwen3_vl.py
@@ -0,0 +1,26 @@
+from .common import strip_mcore_wrappers
+from .qwen import qwen_hf_tensor
+
+
+class _LanguageReader:
+    def __init__(self, reader):
+        self.reader = reader
+
+    @staticmethod
+    def _key(name):
+        return name.replace("model.", "model.language_model.", 1) if name.startswith("model.") else name
+
+    def __contains__(self, name):
+        return self._key(name) in self.reader
+
+    def get_tensor(self, name):
+        return self.reader.get_tensor(self._key(name))
+
+
+def qwen3_vl_hf_tensor(name, reader, config):
+    name = strip_mcore_wrappers(name)
+    if name.startswith("model.visual."):
+        return reader.get_tensor(name)
+    if name == "output_layer.weight" and getattr(config, "tie_word_embeddings", False):
+        return reader.get_tensor("model.language_model.embed_tokens.weight")
+    return qwen_hf_tensor(name, _LanguageReader(reader), config.text_config)
diff --git a/vime/backends/megatron_utils/megatron_to_hf/processors/padding_remover.py b/vime/backends/megatron_utils/megatron_to_hf/processors/padding_remover.py
index 741584fa4..4e548adc9 100644
--- a/vime/backends/megatron_utils/megatron_to_hf/processors/padding_remover.py
+++ b/vime/backends/megatron_utils/megatron_to_hf/processors/padding_remover.py
@@ -7,6 +7,7 @@ def remove_padding(name: str, param: torch.Tensor, vocab_size: int) -> torch.Ten
     """
     Remove vocab padding: param[:vocab_size] for embedding/output layers, else unchanged.
     """
-    if strip_param_name_prefix(name) in {"embedding.word_embeddings.weight", "output_layer.weight"}:
+    name = strip_param_name_prefix(name).removeprefix("language_model.")
+    if name in {"embedding.word_embeddings.weight", "output_layer.weight"}:
         return param[:vocab_size]
     return param
diff --git a/vime/backends/megatron_utils/npu_attention_patch.py b/vime/backends/megatron_utils/npu_attention_patch.py
new file mode 100644
index 000000000..6cfc23222
--- /dev/null
+++ b/vime/backends/megatron_utils/npu_attention_patch.py
@@ -0,0 +1,81 @@
+import torch_npu
+from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.transformer.enums import AttnMaskType
+from torch import Tensor
+
+try:
+    from einops import rearrange
+except ImportError:
+    rearrange = None
+
+
+def npu_dot_product_attention_forward(
+    self,
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    attention_mask: Tensor,
+    attn_mask_type: AttnMaskType = None,
+    attention_bias: Tensor = None,
+    packed_seq_params: PackedSeqParams | None = None,
+):
+    assert attention_bias is None, "Attention bias is not supported for DotProductAttention."
+
+    if packed_seq_params is None:
+        n_head = query.shape[2]
+    else:
+        n_head = query.shape[1]
+
+    sparse_mode = getattr(self.config, "sparse_mode", 0)
+    if attn_mask_type == AttnMaskType.no_mask:
+        sparse_mode = 0
+
+    scale = self.softmax_scale
+
+    pre_tockens = getattr(self.config, "pre_tockens", 65536)
+    next_tockens = getattr(self.config, "next_tockens", 0)
+
+    if packed_seq_params is not None:
+        if isinstance(packed_seq_params.cu_seqlens_q, list):
+            actual_seq_qlen = packed_seq_params.cu_seqlens_q
+            actual_seq_kvlen = packed_seq_params.cu_seqlens_kv
+        else:
+            actual_seq_qlen = packed_seq_params.cu_seqlens_q.tolist()
+            actual_seq_kvlen = packed_seq_params.cu_seqlens_kv.tolist()
+        shape_order = "TND"
+    else:
+        actual_seq_qlen = None
+        actual_seq_kvlen = None
+        if rearrange is not None:
+            query, key, value = [rearrange(x, "s b h d -> s b (h d)") for x in [query, key, value]]
+        else:
+            query = query.reshape(query.shape[0], query.shape[1], -1)
+            key = key.reshape(key.shape[0], key.shape[1], -1)
+            value = value.reshape(value.shape[0], value.shape[1], -1)
+        shape_order = "SBH"
+
+    output = torch_npu.npu_fusion_attention(
+        query,
+        key,
+        value,
+        n_head,
+        shape_order,
+        pse=None,
+        padding_mask=None,
+        atten_mask=attention_mask,
+        scale=scale,
+        pre_tockens=pre_tockens,
+        next_tockens=next_tockens,
+        keep_prob=1 - self.attention_dropout.p,
+        inner_precise=0,
+        sparse_mode=sparse_mode,
+        actual_seq_qlen=actual_seq_qlen,
+        actual_seq_kvlen=actual_seq_kvlen,
+    )[0]
+
+    return output
+
+
+from megatron.core.transformer.dot_product_attention import DotProductAttention
+
+DotProductAttention.forward = npu_dot_product_attention_forward
diff --git a/vime/backends/megatron_utils/update_weight/common.py b/vime/backends/megatron_utils/update_weight/common.py
index 14a2560f8..548f148c9 100644
--- a/vime/backends/megatron_utils/update_weight/common.py
+++ b/vime/backends/megatron_utils/update_weight/common.py
@@ -10,6 +10,7 @@
 from megatron.core import mpu
 from megatron.core.transformer.transformer_layer import get_transformer_layer_offset
 
+from vime.platforms import current_platform
 from vime.utils.distributed_utils import get_gloo_group
 from vime.utils.types import ParamInfo
 
@@ -51,6 +52,7 @@ def all_gather_param(name: str, param: torch.nn.Parameter) -> torch.Tensor:
     if "linear_fc1.weight" in name or "linear_fc1.bias" in name:
         param_partitions = [p.chunk(2, dim=0) for p in param_partitions]
         param_partitions = [p[0] for p in param_partitions] + [p[1] for p in param_partitions]
+        partition_dim = current_platform().megatron.adjust_tp_partition_dim(name, partition_dim)
     # this is bug in megatron's grouped moe.
     if "linear_fc2.weight" in name:
         if partition_dim == 0:
@@ -118,6 +120,7 @@ def all_gather_params_async(
             if "linear_fc1.weight" in info.name or "linear_fc1.bias" in info.name:
                 param_partitions = [p.chunk(2, dim=0) for p in param_partitions]
                 param_partitions = [p[0] for p in param_partitions] + [p[1] for p in param_partitions]
+                partition_dim = current_platform().megatron.adjust_tp_partition_dim(info.name, partition_dim)
             # this is bug in megatron's grouped moe.
             if "linear_fc2.weight" in info.name:
                 if partition_dim == 0:
@@ -289,7 +292,6 @@ def create_nccl_trainer(
 ):
     import ray
     from vllm.distributed.weight_transfer.factory import WeightTransferTrainerFactory
-    from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo
 
     rendezvous = [None]
     if dist.get_rank() == 0:
@@ -299,7 +301,8 @@ def create_nccl_trainer(
     dist.broadcast_object_list(rendezvous, src=0, group=get_gloo_group())
     master_address, master_port = rendezvous[0]
     return WeightTransferTrainerFactory.trainer_init(
-        NCCLTrainerInitInfo(
+        current_platform().weight_transfer.trainer_init_info(
+            colocate=False,
             master_address=master_address,
             master_port=master_port,
             world_size=sum(engine_gpu_counts) + 1,
diff --git a/vime/backends/megatron_utils/update_weight/update_weight_from_disk_delta.py b/vime/backends/megatron_utils/update_weight/update_weight_from_disk_delta.py
index d3725dffb..c98997aa4 100644
--- a/vime/backends/megatron_utils/update_weight/update_weight_from_disk_delta.py
+++ b/vime/backends/megatron_utils/update_weight/update_weight_from_disk_delta.py
@@ -208,20 +208,16 @@ def _encode_delta(self) -> None:
 
         # Pinned host-buffer pool: a pinned non_blocking GPU->CPU copy is far faster than .cpu().
         max_bytes = max((int(v.nbytes) for v in snapshot.values()), default=0)
-        free_q: queue.Queue = queue.Queue()
-        use_pinned = True
-        try:
-            for _ in range(max(4, min(2 * NUM_WORKERS, (32 << 30) // max(max_bytes, 1)))):
-                free_q.put(torch.empty(max_bytes, dtype=torch.uint8, pin_memory=True))
-        except RuntimeError as e:  # low memlock limit
-            logger.warning("pinned host buffers unavailable (%s); using pageable .cpu()", e)
-            use_pinned = False
+        free_q = _make_pinned_pool(max_bytes)
+        use_pinned = not free_q.empty()
 
         def diff_and_compress(name, buf, nbytes, pinned):
             if pinned:  # copy out and free the pinned buffer before the heavy diff/compress
-                new = np.empty(nbytes, dtype=np.uint8)
-                np.copyto(new, buf.numpy()[:nbytes])
-                free_q.put(buf)
+                try:
+                    new = np.empty(nbytes, dtype=np.uint8)
+                    np.copyto(new, buf.numpy()[:nbytes])
+                finally:
+                    free_q.put(buf)
             else:
                 new = buf
             old = snapshot[name]
@@ -253,15 +249,23 @@ def collect(fut):
             for name, tensor in self._iter_hf_tensors():
                 flat = tensor.detach().contiguous().view(torch.uint8).reshape(-1)
                 nbytes = int(flat.numel())
-                if use_pinned and nbytes <= max_bytes:
-                    buf = free_q.get()  # blocks when all buffers are in flight -> backpressures the gather
-                    buf[:nbytes].copy_(flat, non_blocking=True)
-                    accelerator.current_stream().synchronize()
-                    payload, pinned = buf, True
-                else:
-                    payload, pinned = flat.cpu().numpy(), False
-                self.total_bytes += nbytes
-                inflight.append(pool.submit(diff_and_compress, name, payload, nbytes, pinned))
+                buf = None
+                submitted = False
+                try:
+                    if use_pinned and nbytes <= max_bytes:
+                        buf = free_q.get()  # backpressure until a worker returns a buffer
+                        buf[:nbytes].copy_(flat, non_blocking=True)
+                        accelerator.current_stream().synchronize()
+                        payload, pinned = buf, True
+                    else:
+                        payload, pinned = flat.cpu().numpy().copy(), False
+                    self.total_bytes += nbytes
+                    future = pool.submit(diff_and_compress, name, payload, nbytes, pinned)
+                    submitted = True  # the worker now owns returning the buffer
+                    inflight.append(future)
+                finally:
+                    if buf is not None and not submitted:
+                        free_q.put(buf)
                 if len(inflight) >= 2 * NUM_WORKERS:
                     collect(inflight.popleft())
             while inflight:
@@ -291,6 +295,20 @@ def _record_metrics(self) -> None:
             )
 
 
+def _make_pinned_pool(max_bytes: int) -> queue.Queue:
+    """Limit this pool's requested pinned storage to 8 GiB, excluding other CPU state."""
+    free_q: queue.Queue = queue.Queue()
+    num_buffers = min(2 * NUM_WORKERS, (8 << 30) // max_bytes) if max_bytes > 0 else 0
+    try:
+        for _ in range(num_buffers):
+            free_q.put(torch.empty(max_bytes, dtype=torch.uint8, pin_memory=True))
+    except (RuntimeError, MemoryError) as e:
+        while not free_q.empty():
+            free_q.get_nowait()
+        logger.warning("pinned host buffers unavailable (%s); using pageable .cpu()", e)
+    return free_q
+
+
 def _atomic_write(path: str, data: bytes) -> None:
     tmp = path + ".tmp"
     with open(tmp, "wb") as f:
diff --git a/vime/backends/megatron_utils/update_weight/update_weight_from_tensor.py b/vime/backends/megatron_utils/update_weight/update_weight_from_tensor.py
index c88da13d0..22c7ad8ef 100644
--- a/vime/backends/megatron_utils/update_weight/update_weight_from_tensor.py
+++ b/vime/backends/megatron_utils/update_weight/update_weight_from_tensor.py
@@ -13,6 +13,7 @@
 from ray.actor import ActorHandle
 from tqdm import tqdm
 
+from vime.platforms import current_platform
 from vime.utils import accelerator
 from vime.utils.distributed_utils import get_gloo_group
 from vime.utils.types import ParamInfo
@@ -25,6 +26,16 @@
 from .update_weight_from_distributed import post_process_weights
 
 
+def _current_gpu_uuid() -> str:
+    platform = current_platform()
+    if platform.is_npu:
+        return platform.weight_transfer.current_device_uuid()
+
+    device_index = torch.cuda.current_device()
+    props = torch.cuda.get_device_properties(device_index)
+    return str(props.uuid)
+
+
 def _native_ipc_buffer_size(args: Namespace, param_info_buckets: Sequence[Sequence[ParamInfo]] | None) -> int:
     buffer_size = args.update_weight_buffer_size
     if not param_info_buckets:
@@ -64,7 +75,7 @@ def _build_packed_ipc_update_info(
     )
     assert chunk is not None
     _, ipc_args = reduce_tensor(chunk.packed_tensor)
-    gpu_uuid = str(torch.cuda.get_device_properties(torch.cuda.current_device()).uuid)
+    gpu_uuid = _current_gpu_uuid()
     return (
         {
             "names": chunk.names,
@@ -188,11 +199,11 @@ def connect_rollout_engines(
         if not self._expert_transfer_plan:
             if self.rollout_engines:
                 from vllm.distributed.weight_transfer.factory import WeightTransferTrainerFactory
-                from vllm.distributed.weight_transfer.ipc_engine import IPCTrainerInitInfo
 
                 client = VimeRayWeightSyncClient(self.rollout_engines, lambda: self.weight_version)
                 trainer = WeightTransferTrainerFactory.trainer_init(
-                    IPCTrainerInitInfo(
+                    current_platform().weight_transfer.trainer_init_info(
+                        colocate=True,
                         rank=dist.get_rank(),
                         packed=True,
                         packed_buffer_size_bytes=_native_ipc_buffer_size(
diff --git a/vime/backends/vllm_utils/engine_group.py b/vime/backends/vllm_utils/engine_group.py
index 324208aa3..4142c7859 100644
--- a/vime/backends/vllm_utils/engine_group.py
+++ b/vime/backends/vllm_utils/engine_group.py
@@ -10,6 +10,7 @@
 
 from vime.backends.vllm_utils.vllm_config import ServerGroupConfig
 from vime.backends.vllm_utils.vllm_engine import VLLMEngine, _resolve_parallel_sizes
+from vime.platforms import current_platform
 from vime.ray.utils import NOSET_VISIBLE_DEVICES_ENV_VARS_LIST, add_default_ray_env_vars
 
 GPU_MEMORY_TYPE_KV_CACHE = "kv_cache"
@@ -117,6 +118,7 @@ def start_engines(self, port_cursors: dict[int, int] | None = None) -> tuple[lis
         RolloutRayActor = ray.remote(VLLMEngine)
 
         rollout_engines = []
+        platform = current_platform()
         for i in range(len(self.all_engines)):
             if self.all_engines[i] is not None:
                 continue
@@ -142,13 +144,18 @@ def start_engines(self, port_cursors: dict[int, int] | None = None) -> tuple[lis
             env_vars["PYTORCH_CUDA_ALLOC_CONF"] = ",".join(
                 kv for kv in _alloc.split(",") if kv and not kv.strip().startswith("expandable_segments")
             )
+            if platform.is_npu:
+                env_vars = platform.ray.rollout_runtime_env(self.args, env_vars)
+            resource_options = {"num_gpus": num_gpus}
+            if platform.is_npu:
+                resource_options = {"num_gpus": 0, **platform.ray.actor_options(num_gpus)}
             rollout_engine = RolloutRayActor.options(
                 num_cpus=num_cpus,
-                num_gpus=num_gpus,
                 scheduling_strategy=scheduling_strategy,
                 runtime_env={
                     "env_vars": add_default_ray_env_vars(env_vars),
                 },
+                **resource_options,
             ).remote(
                 self.args,
                 rank=global_rank,
diff --git a/vime/backends/vllm_utils/vllm_engine.py b/vime/backends/vllm_utils/vllm_engine.py
index 82c07e6ff..eaa4efddc 100644
--- a/vime/backends/vllm_utils/vllm_engine.py
+++ b/vime/backends/vllm_utils/vllm_engine.py
@@ -13,6 +13,7 @@
 from vllm.utils.system_utils import kill_process_tree
 
 from vime.backends.vllm_utils.external import get_server_info
+from vime.platforms import current_platform
 from vime.ray.ray_actor import RayActor
 from vime.utils.http_utils import _wrap_ipv6, get_host_info
 
@@ -67,6 +68,11 @@ def _build_subprocess_env(server_args_dict: dict[str, Any]) -> dict[str, str]:
     env["CUDA_VISIBLE_DEVICES"] = server_args_dict["_visible_devices"]
     # ROCm: keep HIP visibility in sync with CUDA (no-op on CUDA).
     env["HIP_VISIBLE_DEVICES"] = server_args_dict["_visible_devices"]
+    env = current_platform().vllm.subprocess_env(
+        env,
+        visible_devices=server_args_dict["_visible_devices"],
+        colocate=getattr(args, "colocate", False),
+    )
     env.setdefault("VLLM_SERVER_DEV_MODE", "1")
     env["VLLM_USE_V2_MODEL_RUNNER"] = "1"
     if getattr(args, "vllm_enable_deterministic_inference", False):
@@ -690,6 +696,9 @@ def _compute_server_args(
         kwargs["weight_transfer_config"] = {"backend": "ipc"}
     else:
         kwargs["weight_transfer_config"] = {"backend": "nccl"}
+    kwargs["weight_transfer_config"]["backend"] = current_platform().weight_transfer.backend(
+        kwargs["weight_transfer_config"]["backend"]
+    )
 
     if worker_type == "encoder":
         # vLLM EPD producers have no language-model KV cache groups. Prefix
diff --git a/vime/platforms/__init__.py b/vime/platforms/__init__.py
new file mode 100644
index 000000000..8cec41787
--- /dev/null
+++ b/vime/platforms/__init__.py
@@ -0,0 +1,83 @@
+"""Accelerator platform discovery and narrow capability providers.
+
+``VIME_PLATFORM`` is the explicit override.  When it is not set, NPU detection
+is lazy and failure-safe; CUDA is the default.
+"""
+
+from __future__ import annotations
+
+import os
+from functools import cache
+
+from .base import (
+    CheckpointCapabilities,
+    Platform,
+    RayResourceSpec,
+    TrainingBootstrap,
+    VLLMLaunchPlatformOps,
+    WeightTransferPlatformOps,
+)
+from .cuda import create_cuda_platform
+from .npu import create_npu_platform, detect_npu
+
+_PLATFORM_FACTORIES = {
+    "cuda": create_cuda_platform,
+    "npu": create_npu_platform,
+}
+
+
+@cache
+def get_platform(name: str) -> Platform:
+    normalized = name.strip().lower()
+    try:
+        factory = _PLATFORM_FACTORIES[normalized]
+    except KeyError as exc:
+        available = ", ".join(_PLATFORM_FACTORIES)
+        raise ValueError(f"Unknown Vime platform {name!r}; registered platforms: {available}") from exc
+    return factory()
+
+
+@cache
+def _resolve_platform(override: str | None) -> Platform:
+    if override:
+        return get_platform(override)
+
+    try:
+        if detect_npu():
+            return get_platform("npu")
+    except Exception:  # noqa: BLE001 - a failed detector must not break imports
+        pass
+    return get_platform("cuda")
+
+
+def current_platform() -> Platform:
+    """Resolve the active platform without probing hardware at module import."""
+    raw_override = os.environ.get("VIME_PLATFORM")
+    override = raw_override.strip().lower() if raw_override and raw_override.strip() else None
+    accelerator_override = os.environ.get("VIME_ACCELERATOR", "").strip().lower()
+    if accelerator_override in {"npu", "cuda", "musa"}:
+        accelerator_platform = "npu" if accelerator_override == "npu" else "cuda"
+        if override in _PLATFORM_FACTORIES and override != accelerator_platform:
+            raise ValueError(f"Conflicting VIME_PLATFORM={override!r} and VIME_ACCELERATOR={accelerator_override!r}")
+        if override is None:
+            override = accelerator_platform
+    return _resolve_platform(override)
+
+
+def reset_platform_cache() -> None:
+    """Clear resolver/factory caches (primarily for tests and plugin registration)."""
+    get_platform.cache_clear()
+    _resolve_platform.cache_clear()
+
+
+__all__ = [
+    "CheckpointCapabilities",
+    "Platform",
+    "RayResourceSpec",
+    "TrainingBootstrap",
+    "VLLMLaunchPlatformOps",
+    "WeightTransferPlatformOps",
+    "current_platform",
+    "get_platform",
+    "reset_platform_cache",
+]
diff --git a/vime/platforms/base.py b/vime/platforms/base.py
new file mode 100644
index 000000000..c0775e5db
--- /dev/null
+++ b/vime/platforms/base.py
@@ -0,0 +1,143 @@
+"""Small contracts for behavior that genuinely differs by platform."""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Mapping
+from contextlib import nullcontext
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass(frozen=True)
+class RayResourceSpec:
+    """How one accelerator is represented and addressed by Ray."""
+
+    resource_name: str
+    visible_devices_env: str
+    uses_ray_gpu_resource: bool
+
+    def bundle_resources(self, device_count: float = 1, cpu_count: float = 1) -> dict[str, float]:
+        return {self.resource_name: device_count, "CPU": cpu_count}
+
+    def actor_options(self, fraction: float) -> dict[str, object]:
+        if self.uses_ray_gpu_resource:
+            return {"num_gpus": fraction}
+        return {"resources": {self.resource_name: fraction}}
+
+    def accelerator_ids(self) -> list[str]:
+        import ray
+
+        if self.uses_ray_gpu_resource:
+            ids = ray.get_gpu_ids()
+        else:
+            ids = ray.get_runtime_context().get_accelerator_ids().get(self.resource_name, [])
+        return [str(device_id) for device_id in ids]
+
+    def local_device_id(self) -> int | str:
+        device_ids = self.accelerator_ids()
+        if not device_ids:
+            raise RuntimeError(f"No {self.resource_name} accelerator IDs are assigned to this Ray actor")
+
+        assigned_id = device_ids[0]
+        visible_devices = os.environ.get(self.visible_devices_env)
+        if visible_devices is None:
+            try:
+                return int(assigned_id)
+            except ValueError:
+                return assigned_id
+
+        visible_ids = [value.strip() for value in visible_devices.split(",") if value.strip()]
+        try:
+            return visible_ids.index(assigned_id)
+        except ValueError as exc:
+            raise RuntimeError(
+                f"Ray assigned {self.resource_name} id {assigned_id}, but it is absent from "
+                f"{self.visible_devices_env}={visible_devices!r}"
+            ) from exc
+
+    def train_runtime_env(
+        self,
+        args: Any,
+        env_vars: Mapping[str, str] | None = None,
+    ) -> dict[str, str]:
+        return dict(env_vars or {})
+
+    def rollout_runtime_env(
+        self,
+        args: Any,
+        env_vars: Mapping[str, str] | None = None,
+    ) -> dict[str, str]:
+        return dict(env_vars or {})
+
+
+class WeightTransferPlatformOps:
+    """Select vendor backends without changing the shared trainer lifecycle."""
+
+    def backend(self, backend: str) -> str:
+        return backend
+
+    def trainer_init_info(self, *, colocate: bool, **kwargs):
+        if colocate:
+            from vllm.distributed.weight_transfer.ipc_engine import IPCTrainerInitInfo
+
+            return IPCTrainerInitInfo(**kwargs)
+        from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo
+
+        return NCCLTrainerInitInfo(**kwargs)
+
+
+class VLLMLaunchPlatformOps:
+    """Platform additions to the common vLLM launch command and environment."""
+
+    def subprocess_env(
+        self,
+        base_env: Mapping[str, str],
+        *,
+        visible_devices: str,
+        colocate: bool,
+    ) -> dict[str, str]:
+        return dict(base_env)
+
+
+class TrainingBootstrap:
+    """Lazy Megatron/vendor initialization hooks."""
+
+    def bootstrap(self) -> None:
+        return None
+
+    def repatch(self, args: Any) -> None:
+        return None
+
+    def adjust_tp_partition_dim(self, name: str, partition_dim: int) -> int:
+        return partition_dim
+
+    def training_context(self, offload_train: bool):
+        return nullcontext()
+
+    def initialize_optimizer_state(self, optimizer: Any) -> None:
+        return None
+
+
+@dataclass(frozen=True)
+class CheckpointCapabilities:
+    default_megatron_to_hf_mode: str = "raw"
+
+    def patch_default_planner(self, default_planner: Any) -> None:
+        return None
+
+
+@dataclass(frozen=True)
+class Platform:
+    """Aggregate only the providers whose semantics differ on Ascend."""
+
+    name: str
+    ray: RayResourceSpec
+    weight_transfer: WeightTransferPlatformOps
+    vllm: VLLMLaunchPlatformOps
+    megatron: TrainingBootstrap
+    checkpoint: CheckpointCapabilities
+
+    @property
+    def is_npu(self) -> bool:
+        return self.name == "npu"
diff --git a/vime/platforms/cuda.py b/vime/platforms/cuda.py
new file mode 100644
index 000000000..b068eebab
--- /dev/null
+++ b/vime/platforms/cuda.py
@@ -0,0 +1,25 @@
+"""Default CUDA platform assembled from the shared CUDA-compatible behavior."""
+
+from __future__ import annotations
+
+from .base import (
+    CheckpointCapabilities,
+    Platform,
+    RayResourceSpec,
+    TrainingBootstrap,
+    VLLMLaunchPlatformOps,
+    WeightTransferPlatformOps,
+)
+
+
+def create_cuda_platform() -> Platform:
+    return Platform(
+        name="cuda",
+        ray=RayResourceSpec(
+            resource_name="GPU", visible_devices_env="CUDA_VISIBLE_DEVICES", uses_ray_gpu_resource=True
+        ),
+        weight_transfer=WeightTransferPlatformOps(),
+        vllm=VLLMLaunchPlatformOps(),
+        megatron=TrainingBootstrap(),
+        checkpoint=CheckpointCapabilities(),
+    )
diff --git a/vime/platforms/npu.py b/vime/platforms/npu.py
new file mode 100644
index 000000000..94d2fea2a
--- /dev/null
+++ b/vime/platforms/npu.py
@@ -0,0 +1,292 @@
+"""Ascend NPU implementation of the Vime platform contracts."""
+
+from __future__ import annotations
+
+import importlib
+import logging
+import os
+from contextlib import nullcontext
+from glob import glob
+from typing import Any
+
+from vime.utils import accelerator
+from vime.utils.accelerator.torch_accelerator import TorchAccelerator
+
+from .base import (
+    CheckpointCapabilities,
+    Platform,
+    RayResourceSpec,
+    TrainingBootstrap,
+    VLLMLaunchPlatformOps,
+    WeightTransferPlatformOps,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class NPUAccelerator(TorchAccelerator):
+    name = "npu"
+    device_type = "npu"
+    communication_backend_name = "hccl"
+
+    def _module(self):
+        return getattr(importlib.import_module("torch"), "npu", None)
+
+    @property
+    def visible_devices_env(self) -> str:
+        return "ASCEND_RT_VISIBLE_DEVICES"
+
+    def distributed_device_id(self, index=None):
+        # Preserve lazy HCCL initialization instead of CUDA's eager device binding.
+        return None
+
+    def set_allocator_expandable_segments(self) -> bool:
+        # NPU allocator policy belongs to the existing TMS/runtime-env hooks.
+        return False
+
+
+def register_npu_accelerator() -> None:
+    accelerator.register_accelerator(
+        "npu",
+        NPUAccelerator,
+        is_available=lambda: os.environ.get("VIME_PLATFORM", "").strip().lower() != "cuda"
+        and NPUAccelerator().is_available(),
+        priority=300,
+        communication_backends=("hccl",),
+    )
+
+
+def detect_npu() -> bool:
+    """Return whether a usable NPU is visible, without leaking probe errors."""
+    # Do not import torch/torch_npu on an unselected CUDA host merely because
+    # torch_npu happens to be installed.  Its import has process-wide monkey
+    # patch side effects.  An explicit VIME_PLATFORM=npu override bypasses
+    # detection, while automatic selection first requires an exposed device.
+    if not (os.path.exists("/dev/davinci_manager") or glob("/dev/davinci[0-9]*")):
+        return False
+    try:
+        torch = importlib.import_module("torch")
+        if getattr(torch, "npu", None) is None:
+            importlib.import_module("torch_npu")
+        npu = getattr(torch, "npu", None)
+        return bool(npu is not None and npu.is_available())
+    except Exception:  # noqa: BLE001 - detection must be safe on non-NPU hosts
+        return False
+
+
+def _ensure_torch_npu() -> None:
+    importlib.import_module("torch_npu")
+
+
+def _install_safe_empty_cache() -> None:
+    """Preserve the Ascend allocator guard required by MindSpeed/TMS callers."""
+    torch = importlib.import_module("torch")
+    original_empty_cache = torch.npu.empty_cache
+    if not getattr(original_empty_cache, "_vime_safe_empty_cache", False):
+
+        def _safe_empty_cache(_original=original_empty_cache) -> None:
+            try:
+                _original()
+            except RuntimeError:
+                pass
+
+        _safe_empty_cache._vime_safe_empty_cache = True
+        torch.npu.empty_cache = _safe_empty_cache
+
+    # Some shared dependencies still call the CUDA spelling after torch_npu's
+    # compatibility patching. Keep that alias local to NPU-bootstrapped jobs.
+    torch.cuda.empty_cache = torch.npu.empty_cache
+
+
+def _cann_python_site_packages() -> str | None:
+    candidates: list[str] = []
+    for env_key in ("ASCEND_TOOLKIT_HOME", "ASCEND_HOME_PATH"):
+        base = os.environ.get(env_key)
+        if not base:
+            continue
+        candidates.extend(
+            [
+                os.path.join(base, "python", "site-packages"),
+                os.path.normpath(os.path.join(base, "..", "python", "site-packages")),
+            ]
+        )
+    candidates.append("/usr/local/Ascend/ascend-toolkit/latest/python/site-packages")
+    for path in candidates:
+        if os.path.isdir(os.path.join(path, "acl")):
+            return path
+    return None
+
+
+def _prepend_pythonpath(env: dict[str, str], *paths: str) -> None:
+    existing = env.get("PYTHONPATH", os.environ.get("PYTHONPATH", ""))
+    existing_parts = {part for part in existing.split(os.pathsep) if part}
+    prefix_parts = [path for path in paths if path and path not in existing_parts]
+    if prefix_parts:
+        env["PYTHONPATH"] = os.pathsep.join([*prefix_parts, existing] if existing else prefix_parts)
+
+
+class NpuRayResourceSpec(RayResourceSpec):
+    def __init__(self) -> None:
+        super().__init__(
+            resource_name="NPU",
+            visible_devices_env="ASCEND_RT_VISIBLE_DEVICES",
+            uses_ray_gpu_resource=False,
+        )
+
+    def train_runtime_env(self, args: Any, env_vars=None) -> dict[str, str]:
+        env = dict(env_vars or {})
+        if not (getattr(args, "offload_train", False) and getattr(args, "train_backend", None) == "megatron"):
+            return env
+
+        env["TMS_HOOK_MODE"] = "torch"
+        env["TMS_REGION_TAG"] = "training"
+        env["TMS_ENABLE_CPU_BACKUP"] = "1"
+        if getattr(args, "colocate", False):
+            env["PYTORCH_NPU_ALLOC_CONF"] = "expandable_segments:False"
+        cann_python_path = _cann_python_site_packages()
+        if cann_python_path is not None:
+            _prepend_pythonpath(env, cann_python_path)
+        return env
+
+    def rollout_runtime_env(self, args: Any, env_vars=None) -> dict[str, str]:
+        env = dict(env_vars or {})
+        cann_python_path = _cann_python_site_packages()
+        if cann_python_path is not None:
+            _prepend_pythonpath(env, cann_python_path)
+        if getattr(args, "colocate", False):
+            env["PYTORCH_NPU_ALLOC_CONF"] = "expandable_segments:False"
+        env["VLLM_USE_AOT_COMPILE"] = "0"
+        return env
+
+
+class NpuWeightTransferPlatformOps(WeightTransferPlatformOps):
+    def current_device_uuid(self) -> str:
+        # Reuse vLLM Ascend's canonical host-IP/physical-chip identifier so
+        # trainer and receiver always use the exact same mapping.
+        from vllm_ascend.distributed.weight_transfer.npu_ipc_engine import npu_generate_uuid
+
+        return npu_generate_uuid()
+
+    def backend(self, backend: str) -> str:
+        return {"ipc": "npu_ipc", "nccl": "hccl"}.get(backend, backend)
+
+    def trainer_init_info(self, *, colocate: bool, **kwargs):
+        _ensure_torch_npu()
+        from vllm.plugins import load_general_plugins
+
+        # Trainers, unlike vLLM workers, may not have loaded general plugins yet.
+        load_general_plugins()
+        if colocate:
+            from vllm_ascend.distributed.weight_transfer.npu_ipc_engine import NPUIPCTrainerInitInfo
+
+            return NPUIPCTrainerInitInfo(**kwargs)
+        from vllm_ascend.distributed.weight_transfer.hccl_engine import HCCLTrainerInitInfo
+
+        return HCCLTrainerInitInfo(**kwargs)
+
+
+class NpuVLLMLaunchPlatformOps(VLLMLaunchPlatformOps):
+    def subprocess_env(self, base_env, *, visible_devices: str, colocate: bool) -> dict[str, str]:
+        env = dict(base_env)
+        env.pop("PYTORCH_CUDA_ALLOC_CONF", None)
+        env.pop("CUDA_VISIBLE_DEVICES", None)
+        env.pop("HIP_VISIBLE_DEVICES", None)
+        env["ASCEND_RT_VISIBLE_DEVICES"] = visible_devices
+        env["VLLM_USE_AOT_COMPILE"] = "0"
+        cann_python_path = _cann_python_site_packages()
+        if cann_python_path is not None:
+            _prepend_pythonpath(env, cann_python_path)
+        if colocate:
+            env["PYTORCH_NPU_ALLOC_CONF"] = "expandable_segments:False"
+        return env
+
+
+class NpuTrainingBootstrap(TrainingBootstrap):
+    def __init__(self) -> None:
+        self._bootstrapping = False
+        self._bootstrapped = False
+
+    def bootstrap(self) -> None:
+        if self._bootstrapped or self._bootstrapping:
+            return
+        self._bootstrapping = True
+        try:
+            _ensure_torch_npu()
+            # Select NPU before MegatronAdaptor can make torch.cuda appear available.
+            register_npu_accelerator()
+            selected = accelerator.get_accelerator()
+            if selected.name != "npu":
+                raise RuntimeError(f"NPU bootstrap cannot use an already selected {selected.name!r} accelerator")
+            _install_safe_empty_cache()
+            # MegatronAdaptor must install its pre-patches before any Megatron module
+            # is imported. Apply the NPU attention override afterwards.
+            importlib.import_module("megatron_adaptor")
+            importlib.import_module("vime.backends.megatron_utils.npu_attention_patch")
+        except Exception:
+            # A failed bootstrap may be retried after the runtime environment is
+            # corrected; never leave a partially initialized success marker.
+            raise
+        else:
+            self._bootstrapped = True
+        finally:
+            self._bootstrapping = False
+
+    def repatch(self, args: Any) -> None:
+        features_manager = importlib.import_module(
+            "megatron_adaptor.features_manager.features_manager"
+        ).FeaturesManager
+        full_args = importlib.import_module("megatron_adaptor.utils.args_utils").get_full_args()
+        for key, value in vars(args).items():
+            setattr(full_args, key, value)
+        features_manager.remove_patches()
+        features_manager.apply_features_pre_patches(full_args)
+        features_manager.apply_features_patches(full_args)
+        # Repatch may replace attention again; importing a cached module alone
+        # does not reinstall Vime's existing override.
+        attention = importlib.import_module("vime.backends.megatron_utils.npu_attention_patch")
+        attention.DotProductAttention.forward = attention.npu_dot_product_attention_forward
+
+    def adjust_tp_partition_dim(self, name: str, partition_dim: int) -> int:
+        if "linear_fc1.weight" in name or "linear_fc1.bias" in name:
+            return 0
+        return partition_dim
+
+    def training_context(self, offload_train: bool):
+        if not offload_train:
+            return nullcontext()
+        from torch_memory_saver import torch_memory_saver
+
+        return torch_memory_saver.region(tag="training", enable_cpu_backup=True)
+
+    def initialize_optimizer_state(self, optimizer: Any) -> None:
+        """Create lazy optimizer state before leaving the training memory pool."""
+        if optimizer is None:
+            return
+        for opt in getattr(optimizer, "chained_optimizers", [optimizer]):
+            if opt.optimizer is not None and opt.init_state_fn is not None:
+                opt.init_state_fn(opt.optimizer, opt.config)
+
+
+class NpuCheckpointCapabilities(CheckpointCapabilities):
+    def patch_default_planner(self, default_planner: Any) -> None:
+        if not hasattr(default_planner, "_validate_global_plan"):
+            return
+
+        def _validate_global_plan(global_plan, metadata):
+            logger.info("[NPU checkpoint] Skipping validate_access_integrity")
+            return True
+
+        default_planner._validate_global_plan = _validate_global_plan
+
+
+def create_npu_platform() -> Platform:
+    # Ray workers also resolve the platform without importing Megatron.
+    register_npu_accelerator()
+    return Platform(
+        name="npu",
+        ray=NpuRayResourceSpec(),
+        weight_transfer=NpuWeightTransferPlatformOps(),
+        vllm=NpuVLLMLaunchPlatformOps(),
+        megatron=NpuTrainingBootstrap(),
+        checkpoint=NpuCheckpointCapabilities(default_megatron_to_hf_mode="bridge"),
+    )
diff --git a/vime/ray/actor_group.py b/vime/ray/actor_group.py
index c3725d14c..2bb8ef5d6 100644
--- a/vime/ray/actor_group.py
+++ b/vime/ray/actor_group.py
@@ -7,6 +7,7 @@
 from ray.util.placement_group import PlacementGroup
 from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
 
+from vime.platforms import current_platform
 from vime.ray.utils import NOSET_VISIBLE_DEVICES_ENV_VARS_LIST, add_default_ray_env_vars
 
 
@@ -70,28 +71,24 @@ def _allocate_gpus_for_actor(self, pg, num_gpus_per_actor):
             **self.args.train_env_vars,
         }
 
-        if self.args.offload_train:
-            import torch_memory_saver
-
-            for path in [
-                "torch_memory_saver_hook_mode_preload_cu13.abi3.so",
-                "torch_memory_saver_hook_mode_preload_cu12.abi3.so",
-                "torch_memory_saver_hook_mode_preload.abi3.so",
-            ]:
-                dynlib_path = os.path.join(
-                    os.path.dirname(os.path.dirname(torch_memory_saver.__file__)),
-                    path,
-                )
-                if os.path.exists(dynlib_path):
-                    break
+        platform = current_platform()
+        if self.args.offload_train and self.args.train_backend == "megatron":
+            if platform.is_npu:
+                env_vars = platform.ray.train_runtime_env(self.args, env_vars)
             else:
-                raise FileNotFoundError(
-                    "Cannot find torch_memory_saver dynamic library. Please make sure torch_memory_saver is properly installed."
-                )
-
-            env_vars["LD_PRELOAD"] = dynlib_path
-            env_vars["TMS_INIT_ENABLE"] = "1"
-            env_vars["TMS_INIT_ENABLE_CPU_BACKUP"] = "1"
+                import torch_memory_saver
+
+                for path in [
+                    "torch_memory_saver_hook_mode_preload_cu13.abi3.so",
+                    "torch_memory_saver_hook_mode_preload_cu12.abi3.so",
+                    "torch_memory_saver_hook_mode_preload.abi3.so",
+                ]:
+                    dynlib_path = os.path.join(
+                        os.path.dirname(os.path.dirname(torch_memory_saver.__file__)),
+                        path,
+                    )
+                    if os.path.exists(dynlib_path):
+                        break
 
         # We cannot do routing replay for critic.
         if self.args.use_routing_replay and self.role == "actor":
@@ -116,13 +113,16 @@ def _allocate_gpus_for_actor(self, pg, num_gpus_per_actor):
         self._actor_handlers = []
         master_addr, master_port = None, None
         for rank in range(world_size):
+            resource_options = {"num_gpus": num_gpus_per_actor}
+            if platform.is_npu:
+                resource_options = {"num_gpus": 0, **platform.ray.actor_options(num_gpus_per_actor)}
             actor = TrainRayActor.options(
                 num_cpus=num_gpus_per_actor,
-                num_gpus=num_gpus_per_actor,
                 scheduling_strategy=PlacementGroupSchedulingStrategy(
                     placement_group=pg,
                     placement_group_bundle_index=reordered_bundle_indices[rank],
                 ),
+                **resource_options,
             ).remote(world_size, rank, master_addr, master_port)
             if rank == 0:
                 master_addr, master_port = ray.get(actor.get_master_addr_and_port.remote())
diff --git a/vime/ray/placement_group.py b/vime/ray/placement_group.py
index 145677b62..a08886e81 100644
--- a/vime/ray/placement_group.py
+++ b/vime/ray/placement_group.py
@@ -6,6 +6,8 @@
 from ray.util.placement_group import placement_group
 from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
 
+from vime.platforms import current_platform
+
 from .actor_group import RayTrainGroup
 from .utils import add_default_ray_env_vars
 
@@ -15,6 +17,17 @@
 @ray.remote(num_gpus=1)
 class InfoActor:
     def get_ip_and_gpu_id(self):
+        platform = current_platform()
+        if platform.is_npu:
+            accelerator_ids = platform.ray.accelerator_ids()
+            if accelerator_ids:
+                return ray.util.get_node_ip_address(), accelerator_ids[0]
+
+            raise RuntimeError(
+                f"No {platform.ray.resource_name} accelerator IDs found. "
+                f"Accelerator IDs: {ray.get_runtime_context().get_accelerator_ids()}"
+            )
+
         return ray.util.get_node_ip_address(), ray.get_gpu_ids()[0]
 
 
@@ -44,7 +57,11 @@ def _create_placement_group(num_gpus):
     if num_gpus == 0:
         return None, [], []
 
+    platform = current_platform()
+    resource_name = platform.ray.resource_name
     bundles = [{"GPU": 1, "CPU": 1} for _ in range(num_gpus)]
+    if platform.is_npu:
+        bundles = [platform.ray.bundle_resources() for _ in range(num_gpus)]
     pg = placement_group(bundles, strategy="PACK")
     num_bundles = len(bundles)
 
@@ -59,22 +76,24 @@ def _create_placement_group(num_gpus):
     log_interval = 30
     while not ray.wait([ready_ref], timeout=log_interval)[0]:
         elapsed += log_interval
-        total = ray.cluster_resources().get("GPU", 0)
-        available = ray.available_resources().get("GPU", 0)
+        total = ray.cluster_resources().get(resource_name, 0)
+        available = ray.available_resources().get(resource_name, 0)
         logger.info(
-            f"Waiting for placement group of {num_gpus} GPUs (elapsed {elapsed}s): "
-            f"{total:g} GPUs registered with Ray, {available:g} available."
+            f"Waiting for placement group of {num_gpus} {resource_name} devices (elapsed {elapsed}s): "
+            f"{total:g} registered with Ray, {available:g} available."
         )
 
     # use info actor to get the GPU id
     info_actors = []
     for i in range(num_bundles):
+        resource_options = {"num_gpus": 0, **platform.ray.actor_options(1)} if platform.is_npu else {}
         info_actors.append(
             InfoActor.options(
                 scheduling_strategy=PlacementGroupSchedulingStrategy(
                     placement_group=pg,
                     placement_group_bundle_index=i,
                 ),
+                **resource_options,
             ).remote()
         )
     gpu_ids = ray.get([actor.get_ip_and_gpu_id.remote() for actor in info_actors])
diff --git a/vime/ray/train_actor.py b/vime/ray/train_actor.py
index 72fea191c..a9da1b065 100644
--- a/vime/ray/train_actor.py
+++ b/vime/ray/train_actor.py
@@ -10,6 +10,7 @@
 
 import vime.utils.eval_config
 from vime.observability.logging_utils import configure_logger
+from vime.platforms import current_platform
 from vime.ray.ray_actor import RayActor
 from vime.utils import accelerator
 from vime.utils.distributed_utils import init_gloo_group
@@ -19,6 +20,10 @@
 
 
 def get_local_gpu_id():
+    platform = current_platform()
+    if platform.is_npu:
+        return platform.ray.local_device_id()
+
     return accelerator.resolve_visible_device_id(ray.get_gpu_ids()[0])
 
 
diff --git a/vime/utils/disk_delta.py b/vime/utils/disk_delta.py
index c3abc65b1..e6631b794 100644
--- a/vime/utils/disk_delta.py
+++ b/vime/utils/disk_delta.py
@@ -59,15 +59,41 @@ def checksum(algorithm: str, buf) -> str:
 
 def _tensor_locations(ckpt_dir: str) -> dict[str, tuple[str, int, int]]:
     """Map each tensor name to (file, byte offset, nbytes) by reading every safetensors header."""
+    paths = sorted(glob.glob(os.path.join(ckpt_dir, "*.safetensors")))
+    if not paths:
+        raise FileNotFoundError(f"No .safetensors files found in checkpoint directory: {ckpt_dir}")
     locations: dict[str, tuple[str, int, int]] = {}
-    for path in glob.glob(os.path.join(ckpt_dir, "*.safetensors")):
-        with open(path, "rb") as f:
-            (header_len,) = struct.unpack("<Q", f.read(8))
-            header = json.loads(f.read(header_len))
+    for path in paths:
+        try:
+            with open(path, "rb") as f:
+                file_size = os.fstat(f.fileno()).st_size
+                prefix = f.read(8)
+                if len(prefix) != 8:
+                    raise ValueError("truncated header length field")
+                (header_len,) = struct.unpack("<Q", prefix)
+                if header_len > file_size - 8:
+                    raise ValueError("declared header length exceeds file size")
+                header_bytes = f.read(header_len)
+                if len(header_bytes) != header_len:
+                    raise ValueError("truncated header")
+                header = json.loads(header_bytes)
+                if not isinstance(header, dict):
+                    raise ValueError("header must be a JSON object")
+        except (ValueError, UnicodeError, struct.error) as e:
+            raise RuntimeError(f"Failed to parse safetensors header from {path}: {e}") from e
+        data_size = file_size - 8 - header_len
         for name, info in header.items():
             if name == "__metadata__":
                 continue
-            begin, end = info["data_offsets"]
+            offsets = info.get("data_offsets") if isinstance(info, dict) else None
+            if (
+                not isinstance(offsets, list)
+                or len(offsets) != 2
+                or any(type(value) is not int for value in offsets)
+                or not 0 <= offsets[0] <= offsets[1] <= data_size
+            ):
+                raise RuntimeError(f"Invalid data_offsets for tensor {name!r} in {path}: {offsets!r}")
+            begin, end = offsets
             locations[name] = (path, 8 + header_len + begin, end - begin)
     return locations
 
@@ -81,6 +107,9 @@ def read(name: str) -> np.ndarray:
         path, offset, nbytes = locations[name]
         with open(path, "rb") as f:
             f.seek(offset)
-            return np.frombuffer(f.read(nbytes), dtype=np.uint8)
+            data = f.read(nbytes)
+            if len(data) != nbytes:
+                raise RuntimeError(f"Truncated tensor {name!r} in {path}: expected {nbytes} bytes, read {len(data)}")
+            return np.frombuffer(data, dtype=np.uint8)
 
     return read
diff --git a/vime/utils/external_utils/command_utils.py b/vime/utils/external_utils/command_utils.py
index bcfacd4b8..504772da5 100644
--- a/vime/utils/external_utils/command_utils.py
+++ b/vime/utils/external_utils/command_utils.py
@@ -10,6 +10,7 @@
 from dataclasses import dataclass
 from pathlib import Path
 
+from vime.utils.external_utils.launch import current_platform, launch_commands
 from vime.utils.external_utils.typer_utils import dataclass_cli
 from vime.utils.misc import exec_command
 
@@ -27,6 +28,11 @@ def convert_checkpoint(
     dir_dst: str = "/root",
     hf_checkpoint: str | None = None,
 ):
+    # Platforms without automatic conversion use native HF loading by default.
+    if not current_platform().torch_dist_convert:
+        print(f"convert_checkpoint skip on {current_platform().name} (native HF load)")
+        return
+
     hf_checkpoint = hf_checkpoint or f"/root/models/{model_name}"
     normalized_extra_args = f" {extra_args.strip()}" if extra_args.strip() else ""
 
@@ -104,6 +110,29 @@ def execute_train(
         extra_env_vars = {}
     if config is None:
         config = ExecuteTrainConfig()
+
+    # Platform seam: non-CUDA delegates to the launcher; the CUDA path below stays
+    # byte-identical to upstream.
+    platform = current_platform()
+    if platform.name != "cuda":
+        cmds = launch_commands(
+            platform,
+            train_args=train_args,
+            num_devices=num_gpus_per_node,
+            megatron_model_type=megatron_model_type,
+            repo_base_dir=str(repo_base_dir),
+            train_script=train_script,
+            extra_env={**extra_env_vars, **_parse_extra_env_vars(config.extra_env_vars)},
+            external_ray=get_bool_env_var("VIME_SCRIPT_EXTERNAL_RAY"),
+            master_addr=os.environ.get("MASTER_ADDR", "127.0.0.1"),
+        )
+        for cmd in cmds[:-1]:
+            exec_command(cmd)
+        if before_ray_job_submit is not None:
+            before_ray_job_submit()
+        exec_command(cmds[-1])
+        return
+
     external_ray = get_bool_env_var("VIME_SCRIPT_EXTERNAL_RAY")
     master_addr = os.environ.get("MASTER_ADDR", "127.0.0.1")
 
diff --git a/vime/utils/external_utils/launch.py b/vime/utils/external_utils/launch.py
new file mode 100644
index 000000000..34368ea60
--- /dev/null
+++ b/vime/utils/external_utils/launch.py
@@ -0,0 +1,143 @@
+"""Device-specific launch policy for the vime test/example harness.
+
+Used only by the test and example launch utilities (`command_utils.execute_train`), not by
+vime core: a `Platform` describes one accelerator for the purpose of *launching a job* — how
+Ray advertises its devices, the device runtime env, whether torch_dist checkpoint conversion
+works, and how to construct the launch command. It adapts the core ``vime.platforms``
+selection to shell commands.
+
+Imports stay stdlib-only (torch imports are lazy) so the module is unit-testable in isolation;
+the actual `exec_command` calls live in `command_utils`.
+"""
+
+import json
+import shlex
+from dataclasses import dataclass, field
+
+
+# ── Platform contract ──────────────────────────────────────────────────────
+
+
+@dataclass(frozen=True)
+class Platform:
+    name: str
+    ray_args: str  # ray-start resource flags, "{n}"-templated with the device count
+    env: dict = field(default_factory=dict)  # device runtime env (into runtime_env + raylet)
+    torch_dist_convert: bool = True  # False -> use native HF loading without automatic conversion
+
+    def ray_start_args(self, num_devices: int) -> str:
+        return self.ray_args.format(n=num_devices)
+
+
+# ── Registry ────────────────────────────────────────────────────────────────
+
+PLATFORMS: dict[str, Platform] = {}
+
+
+def register(platform: Platform) -> None:
+    PLATFORMS[platform.name] = platform
+
+
+# ── Registered platforms (add a new accelerator here) ─────────────────────────
+
+
+register(Platform(name="cuda", ray_args="--num-gpus {n}"))  # default; other fields unused for cuda
+
+register(
+    Platform(
+        name="npu",
+        # vime requests NPU bundles, not GPU (see ray/placement_group.py), so advertise
+        # the custom NPU resource rather than Ray GPU capacity.
+        ray_args="--num-gpus 0 --resources '{{\"NPU\": {n}}}'",
+        torch_dist_convert=False,  # Keep HF tests unchanged; torch_dist has a separate opt-in test.
+        env={
+            "PYTHONPATH": (
+                "/root/Megatron-LM:/root/vime:"
+                "/root/Megatron-Bridge/src:/root/mbridge:"
+                "/root/MegatronAdaptor:/root/TransformerEngineNPU:"
+                "/usr/local/Ascend/ascend-toolkit/latest/python/site-packages:"
+                "/usr/local/Ascend/ascend-toolkit/latest/tools/ms_fmk_transplt/torch_npu_bridge"
+            ),
+            "LD_LIBRARY_PATH": (
+                "/usr/local/Ascend/driver/lib64:/usr/local/Ascend/driver/lib64/driver:"
+                "/usr/local/Ascend/driver/lib64/common:/usr/local/Ascend/ascend-toolkit/latest/lib64:"
+                "/usr/local/Ascend/ascend-toolkit/latest/compiler/lib64/plugin/opskernel:"
+                "/usr/local/Ascend/ascend-toolkit/latest/compiler/lib64/plugin/nnengine:"
+                "/usr/local/Ascend/ascend-toolkit/latest/opp/built-in/op_impl/ai_core/tbe/op_tiling/lib/:"
+                "/usr/local/Ascend/nnal/atb/latest/atb/cxx_abi_1/lib:"
+                "/usr/local/Ascend/cann/lib64:/usr/local/Ascend/cann/aarch64-linux/devlib"
+            ),
+            "CUDA_DEVICE_MAX_CONNECTIONS": "1",
+            "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES": "1",
+            "HCCL_HOST_SOCKET_PORT_RANGE": "60000-60050",
+            "HCCL_NPU_SOCKET_PORT_RANGE": "61000-61050",
+            "HCCL_CONNECT_TIMEOUT": "7200",
+            "PYTORCH_NPU_ALLOC_CONF": "expandable_segments:False",
+            "VLLM_DISABLE_COMPILE_CACHE": "1",
+            "VLLM_ASCEND_ENABLE_NZ": "0",
+            "HYDRA_FULL_ERROR": "1",
+            "TRANSFORMERS_VERBOSITY": "error",  # silence transformers image-processing log spam
+        },
+    )
+)
+
+
+# ── Resolver + launcher (platform-agnostic; unchanged when adding a platform) ──
+
+
+def current_platform() -> Platform:
+    """Adapt the selected core runtime platform to the launch contract."""
+    from vime.platforms import current_platform as current_runtime_platform
+
+    selected = current_runtime_platform().name
+    try:
+        return PLATFORMS[selected]
+    except KeyError as exc:
+        raise ValueError(f"No launcher adapter is registered for platform {selected!r}") from exc
+
+
+def launch_commands(
+    platform: Platform,
+    train_args: str,
+    num_devices: int,
+    megatron_model_type: str | None,
+    repo_base_dir: str,
+    train_script: str = "train.py",
+    extra_env: dict | None = None,
+    external_ray: bool = False,
+    master_addr: str = "127.0.0.1",
+) -> list:
+    """Ordered shell commands to launch a training job on `platform` (pure; caller execs).
+
+    Driven entirely by platform data (ray resources + device env), so any non-CUDA platform
+    works with no change here. Lives outside `execute_train` so its CUDA body stays
+    byte-identical to upstream and only adds a one-line seam.
+    """
+    extra_env = extra_env or {}
+    all_env = {**platform.env, **extra_env}
+    all_env["VIME_PLATFORM"] = platform.name
+    cmds: list = []
+    cmds.append(
+        "pkill -9 -f '[v]llm serve|VLL[M]::'; sleep 3; "
+        + ("" if external_ray else "ray stop --force; pkill -9 ray; ")
+        + "pkill -9 vime; sleep 3; true; "
+    )
+    if not external_ray:
+        # Export the device env BEFORE `ray start` so the raylet and the actors it spawns
+        # (e.g. vLLM engines) inherit it; a job's runtime_env does not reach those actors.
+        exports = "".join(f"export {k}={shlex.quote(str(v))} && " for k, v in all_env.items())
+        cmds.append(
+            f"{exports}export PYTHONUNBUFFERED=1 && ray start --head --node-ip-address {master_addr} "
+            f"{platform.ray_start_args(num_devices)} --disable-usage-stats"
+        )
+
+    runtime_env = {"env_vars": {"no_proxy": f"127.0.0.1,{master_addr}", "MASTER_ADDR": master_addr, **all_env}}
+    src = f'source "{repo_base_dir}/scripts/models/{megatron_model_type}.sh" && ' if megatron_model_type else ""
+    model_args = "${MODEL_ARGS[@]}" if megatron_model_type else ""
+    cmds.append(
+        f"export no_proxy=127.0.0.1 && export PYTHONUNBUFFERED=1 && {src}"
+        f'ray job submit --address="http://127.0.0.1:8265" '
+        f"--runtime-env-json={shlex.quote(json.dumps(runtime_env))} "
+        f"-- python3 {train_script} {model_args} {train_args}"
+    )
+    return cmds
diff --git a/vime_plugins/models/qwen3_vl.py b/vime_plugins/models/qwen3_vl.py
new file mode 100644
index 000000000..de5782c6f
--- /dev/null
+++ b/vime_plugins/models/qwen3_vl.py
@@ -0,0 +1,201 @@
+"""Native Qwen3-VL: Megatron language model and a replicated HF vision tower."""
+
+from __future__ import annotations
+
+import torch
+from megatron.core import mpu, tensor_parallel
+from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec
+from megatron.core.transformer.module import MegatronModule
+from transformers import AutoConfig
+
+from .qwen3_5_vl_utils import build_packed_mrope_position_ids
+from .qwen3_omni_moe import Qwen3OmniMoeGPTModel
+from .qwen3_omni_transformer import split_deepstack_embeddings
+
+
+def _load_vision_model(hf_config, config):
+    from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLVisionModel
+
+    device = (
+        torch.device("cpu") if config.use_cpu_initialization else torch.device("cuda", torch.cuda.current_device())
+    )
+    with device:
+        vision_model = Qwen3VLVisionModel._from_config(hf_config.vision_config, attn_implementation="sdpa")
+    vision_model.to(dtype=config.params_dtype)
+    if config.recompute_granularity == "full":
+        vision_model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
+    for parameter in vision_model.parameters():
+        parameter.tensor_model_parallel = False
+        parameter.partition_dim = -1
+        parameter.partition_stride = 1
+    return vision_model
+
+
+class Qwen3VLModel(MegatronModule):
+    def __init__(self, config, language_layer_spec, hf_config, args, *, pre_process, post_process, vp_stage):
+        super().__init__(config=config)
+        self.pre_process = pre_process
+        self.post_process = post_process
+        self.image_token_id = hf_config.image_token_id
+        self.video_token_id = hf_config.video_token_id
+        self.vision_start_token_id = hf_config.vision_start_token_id
+        self.spatial_merge_size = hf_config.vision_config.spatial_merge_size
+
+        text_config = hf_config.text_config
+        rope = getattr(text_config, "rope_parameters", None) or text_config.rope_scaling
+        config.mrope_section = list(rope["mrope_section"])
+        config.position_embedding_type = "mrope"
+        config.rotary_base = rope.get("rope_theta", getattr(text_config, "rope_theta", args.rotary_base))
+        config.apply_rope_fusion = False
+        # The main Omni GPT class is also usable with a dense Qwen layer spec:
+        # its additions are interleaved MRoPE and checkpoint-aware DeepStack.
+        self.language_model = Qwen3OmniMoeGPTModel(
+            config=config,
+            transformer_layer_spec=language_layer_spec,
+            vocab_size=args.padded_vocab_size,
+            max_sequence_length=args.max_position_embeddings,
+            pre_process=pre_process,
+            post_process=post_process,
+            fp16_lm_cross_entropy=args.fp16_lm_cross_entropy,
+            parallel_output=True,
+            share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights,
+            position_embedding_type="mrope",
+            rotary_percent=args.rotary_percent,
+            rotary_base=config.rotary_base,
+            scatter_embedding_sequence_parallel=False,
+            vp_stage=vp_stage,
+        )
+        self.model = torch.nn.Module()
+        self.model.visual = _load_vision_model(hf_config, config) if pre_process else None
+        self.share_embeddings_and_output_weights = self.language_model.share_embeddings_and_output_weights
+
+    @property
+    def decoder(self):
+        return self.language_model.decoder
+
+    def shared_embedding_or_output_weight(self):
+        return self.language_model.shared_embedding_or_output_weight()
+
+    def set_input_tensor(self, input_tensor):
+        self.language_model.set_input_tensor(input_tensor)
+
+    def _inject_vision_embeddings(self, input_ids, pixel_values, pixel_values_videos, image_grid_thw, video_grid_thw):
+        embeddings = self.language_model.embedding(input_ids=input_ids, position_ids=None).transpose(0, 1).clone()
+        visual_mask = torch.zeros_like(input_ids, dtype=torch.bool)
+        positions, deepstack = [], []
+        for values, grids, token_id in (
+            (pixel_values, image_grid_thw, self.image_token_id),
+            (pixel_values_videos, video_grid_thw, self.video_token_id),
+        ):
+            mask = input_ids == token_id
+            if values is None:
+                if grids is not None or mask.any():
+                    raise ValueError("Qwen3-VL vision tokens/grids require matching pixel values")
+                continue
+            if grids is None:
+                raise ValueError("Qwen3-VL pixel values require matching grid_thw")
+            output = self.model.visual(values.to(dtype=self.model.visual.dtype), grid_thw=grids)
+            if hasattr(output, "pooler_output"):
+                features, layer_features = output.pooler_output, output.deepstack_features
+            else:
+                features, layer_features = output
+            if mask.sum().item() != features.shape[0]:
+                raise ValueError("Qwen3-VL token/features count mismatch")
+            embeddings[mask] = features.to(embeddings)
+            visual_mask |= mask
+            positions.append(mask.flatten().nonzero(as_tuple=False).flatten())
+            deepstack.append(layer_features)
+
+        deepstack_features = None
+        if deepstack:
+            # Images and videos are encoded separately but can interleave in a sample.
+            order = torch.cat(positions).argsort()
+            deepstack_features = [
+                torch.cat(features)[order].to(embeddings) for features in zip(*deepstack, strict=True)
+            ]
+        embeddings = embeddings.transpose(0, 1).contiguous()
+        if self.config.sequence_parallel:
+            embeddings = tensor_parallel.scatter_to_sequence_parallel_region(embeddings).contiguous()
+            if deepstack_features is not None:
+                # Unlike Omni's frozen tower, this tower remains trainable. Each
+                # SP rank uses only part of DeepStack; sum its feature gradients
+                # before backpropagating through the replicated vision tower.
+                deepstack_features = [
+                    tensor_parallel.copy_to_tensor_model_parallel_region(features) for features in deepstack_features
+                ]
+                visual_mask, deepstack_features = split_deepstack_embeddings(
+                    visual_mask,
+                    deepstack_features,
+                    tp_size=mpu.get_tensor_model_parallel_world_size(),
+                    tp_rank=mpu.get_tensor_model_parallel_rank(),
+                    sequence_parallel=True,
+                )
+        return embeddings, visual_mask if deepstack_features is not None else None, deepstack_features
+
+    def forward(
+        self,
+        input_ids,
+        position_ids=None,
+        attention_mask=None,
+        labels=None,
+        packed_seq_params=None,
+        loss_mask=None,
+        pixel_values=None,
+        pixel_values_videos=None,
+        image_grid_thw=None,
+        video_grid_thw=None,
+        **kwargs,
+    ):
+        if packed_seq_params is None or packed_seq_params.qkv_format != "thd":
+            raise ValueError("Qwen3-VL native training requires THD packed sequences")
+        if position_ids is None:
+            cu_seqlens = packed_seq_params.cu_seqlens_q_padded
+            if cu_seqlens is None:
+                cu_seqlens = packed_seq_params.cu_seqlens_q
+            position_ids = build_packed_mrope_position_ids(
+                input_ids,
+                cu_seqlens,
+                image_grid_thw,
+                video_grid_thw,
+                image_token_id=self.image_token_id,
+                video_token_id=self.video_token_id,
+                vision_start_token_id=self.vision_start_token_id,
+                spatial_merge_size=self.spatial_merge_size,
+            )
+        embeddings, visual_mask, deepstack = self._inject_vision_embeddings(
+            input_ids, pixel_values, pixel_values_videos, image_grid_thw, video_grid_thw
+        )
+        self.language_model.rotary_pos_emb.is_thd_format = True
+        return self.language_model(
+            input_ids=input_ids,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            decoder_input=embeddings,
+            labels=labels,
+            packed_seq_params=packed_seq_params,
+            loss_mask=loss_mask,
+            visual_pos_masks=visual_mask,
+            deepstack_visual_embeds=deepstack,
+            **kwargs,
+        )
+
+
+def get_qwen3_vl_model_provider(args, config, vp_stage):
+    """Use main's --spec provider interface without adding a Bridge branch."""
+    if config.pipeline_model_parallel_size != 1 or config.context_parallel_size != 1:
+        raise ValueError("Qwen3-VL native training currently supports PP=1 and CP=1")
+    if args.mtp_num_layers:
+        raise ValueError("Qwen3-VL native MTP is not supported")
+    if args.transformer_impl != "transformer_engine":
+        raise ValueError("Qwen3-VL native training requires the TE/MindSpeed layer spec")
+    hf_config = AutoConfig.from_pretrained(args.hf_checkpoint, trust_remote_code=True)
+    if hf_config.model_type != "qwen3_vl":
+        raise ValueError(f"{args.hf_checkpoint} is not a Qwen3-VL checkpoint")
+    layer_spec = get_gpt_layer_with_transformer_engine_spec(qk_layernorm=True)
+
+    def model_provider(pre_process=True, post_process=True, vp_stage=None):
+        return Qwen3VLModel(
+            config, layer_spec, hf_config, args, pre_process=pre_process, post_process=post_process, vp_stage=vp_stage
+        )
+
+    return model_provider