diff --git a/.github/workflows/ci-npu-mindspeed.yml b/.github/workflows/ci-npu-mindspeed.yml
new file mode 100644
index 000000000..0790634c9
--- /dev/null
+++ b/.github/workflows/ci-npu-mindspeed.yml
@@ -0,0 +1,195 @@
+name: MegatronAdaptor NPU Tests
+
+on:
+  workflow_dispatch:
+    inputs:
+      megatron_adaptor_repo:
+        description: "MegatronAdaptor git repository"
+        default: "https://gitcode.com/Ascend/MegatronAdaptor.git"
+      megatron_adaptor_ref:
+        description: "MegatronAdaptor branch, tag, or ref to install"
+        default: "core_r0.17.0"
+      transformer_engine_npu_repo:
+        description: "TransformerEngineNPU git repository"
+        default: "https://gitcode.com/Ascend/TransformerEngineNPU.git"
+      transformer_engine_npu_ref:
+        description: "TransformerEngineNPU branch, tag, or ref to install"
+        default: "main"
+      megatron_core_repo:
+        description: "Megatron-Core git repository"
+        default: "https://github.com/NVIDIA/Megatron-LM.git"
+      megatron_core_ref:
+        description: "Megatron-Core branch, tag, or ref to install"
+        default: "core_r0.17.0"
+  push:
+    branches: [main, npu_ci_all]
+    paths: &npu_paths
+      - ".github/workflows/ci-npu-mindspeed.yml"
+      - "mcore_adapter/**"
+      - "roll/**"
+      - "tests/third_party/megatron/**"
+      - "requirements_common.txt"
+      - "requirements_vision.txt"
+      - "setup.py"
+      - "pyproject.toml"
+  pull_request:
+    branches: [main, npu_ci_all]
+    paths: *npu_paths
+
+permissions:
+  contents: read
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  megatron-adaptor-npu-test:
+    name: MegatronAdaptor 0.17 Core NPU Tests
+    if: github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository
+    runs-on: linux-aarch64-a3-8
+    timeout-minutes: 90
+    container:
+      image: quay.io/ascend/vllm-ascend:v0.18.0-a3
+    env:
+      PIP_CACHE_DIR: ${{ github.workspace }}/.pip-cache
+      PIP_INDEX_URL: https://repo.huaweicloud.com/repository/pypi/simple
+      PIP_TRUSTED_HOST: repo.huaweicloud.com
+      HF_ENDPOINT: https://hf-mirror.com
+      PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True"
+      TASK_QUEUE_ENABLE: "2"
+      MEGATRON_ADAPTOR_REPO: ${{ github.event.inputs.megatron_adaptor_repo || 'https://gitcode.com/Ascend/MegatronAdaptor.git' }}
+      MEGATRON_ADAPTOR_REF: ${{ github.event.inputs.megatron_adaptor_ref || 'core_r0.17.0' }}
+      TRANSFORMER_ENGINE_NPU_REPO: ${{ github.event.inputs.transformer_engine_npu_repo || 'https://gitcode.com/Ascend/TransformerEngineNPU.git' }}
+      TRANSFORMER_ENGINE_NPU_REF: ${{ github.event.inputs.transformer_engine_npu_ref || 'main' }}
+      MEGATRON_CORE_REPO: ${{ github.event.inputs.megatron_core_repo || 'https://github.com/NVIDIA/Megatron-LM.git' }}
+      MEGATRON_CORE_REF: ${{ github.event.inputs.megatron_core_ref || 'core_r0.17.0' }}
+      MEGATRON_ADAPTOR_CACHE_KEY: "core-r0.17.0-te-npu-main"
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Cache NPU pip packages
+        uses: actions/cache@v4
+        with:
+          path: .pip-cache
+          key: ${{ runner.os }}-npu-megatron-adaptor-${{ env.MEGATRON_ADAPTOR_CACHE_KEY }}-${{ hashFiles('requirements_common.txt', 'requirements_vision.txt', 'mcore_adapter/pyproject.toml', 'mcore_adapter/requirements.txt', 'setup.py', 'pyproject.toml', '.github/workflows/ci-npu-mindspeed.yml') }}
+          restore-keys: |
+            ${{ runner.os }}-npu-megatron-adaptor-${{ env.MEGATRON_ADAPTOR_CACHE_KEY }}-
+            ${{ runner.os }}-npu-megatron-adaptor-
+            ${{ runner.os }}-npu-pip-
+
+      - name: Configure Ascend runtime
+        shell: bash
+        run: |
+          for env_file in \
+            /usr/local/Ascend/ascend-toolkit/set_env.sh \
+            /usr/local/Ascend/nnal/atb/set_env.sh; do
+            [ -f "${env_file}" ] && source "${env_file}"
+          done
+
+          ASCEND_HOME_PATH="${ASCEND_HOME_PATH:-/usr/local/Ascend/ascend-toolkit/latest}"
+          ASCEND_TOOLKIT_HOME="${ASCEND_TOOLKIT_HOME:-${ASCEND_HOME_PATH}}"
+          ASCEND_OPP_PATH="${ASCEND_OPP_PATH:-${ASCEND_HOME_PATH}/opp}"
+          ASCEND_AICPU_PATH="${ASCEND_AICPU_PATH:-${ASCEND_HOME_PATH}}"
+          LD_LIBRARY_PATH="${ASCEND_HOME_PATH}/lib64:${ASCEND_HOME_PATH}/runtime/lib64:${ASCEND_HOME_PATH}/runtime/lib64/stub:${ASCEND_HOME_PATH}/tools/hccl/lib64:${ASCEND_HOME_PATH}/hccl/lib64:${LD_LIBRARY_PATH:-}"
+
+          for path in \
+            "${ASCEND_OPP_PATH}/built-in/op_impl/ai_core/tbe" \
+            "${ASCEND_HOME_PATH}/python/site-packages"; do
+            [ -d "${path}" ] && PYTHONPATH="${path}:${PYTHONPATH:-}"
+          done
+
+          {
+            echo "ASCEND_HOME_PATH=${ASCEND_HOME_PATH}"
+            echo "ASCEND_TOOLKIT_HOME=${ASCEND_TOOLKIT_HOME}"
+            echo "ASCEND_OPP_PATH=${ASCEND_OPP_PATH}"
+            echo "ASCEND_AICPU_PATH=${ASCEND_AICPU_PATH}"
+            echo "LD_LIBRARY_PATH=${LD_LIBRARY_PATH}"
+            echo "PYTHONPATH=${PYTHONPATH:-}"
+          } >> "${GITHUB_ENV}"
+          {
+            echo "${ASCEND_HOME_PATH}/bin"
+            echo "${ASCEND_HOME_PATH}/compiler/ccec_compiler/bin"
+          } >> "${GITHUB_PATH}"
+
+      - name: Check NPU environment
+        run: |
+          python3 - <<'PY'
+          import importlib.util
+
+          import torch
+          import torch_npu
+
+          if importlib.util.find_spec("tbe") is None:
+              raise RuntimeError("CANN tbe Python module is not visible in PYTHONPATH")
+          if not torch.npu.is_available():
+              raise RuntimeError("torch.npu.is_available() is False")
+          print(f"npu_device_count={torch.npu.device_count()}")
+          PY
+
+      - name: Install Megatron-Core 0.17
+        shell: bash
+        run: |
+          python3 -m pip install --upgrade pip wheel
+          # Megatron-Core 0.17 metadata requires Python 3.12, while the
+          # Ascend-supported stack and runner image use Python 3.10.
+          python3 -m pip install "setuptools<80" pybind11 "packaging>=24.2"
+          export MEGATRON_CORE_SRC="/tmp/Megatron-LM"
+          rm -rf "${MEGATRON_CORE_SRC}"
+          git clone --depth 1 --branch "${MEGATRON_CORE_REF}" \
+            "${MEGATRON_CORE_REPO}" "${MEGATRON_CORE_SRC}"
+          python3 -m pip install --ignore-requires-python --no-build-isolation --no-deps \
+            -e "${MEGATRON_CORE_SRC}"
+
+      - name: Install ROLL requirements
+        shell: bash
+        run: |
+          # Megatron-Core requires setuptools<80; this also keeps
+          # pkg_resources available for torchair.
+          python3 -m pip install --retries 10 --timeout 120 -r requirements_common.txt
+          python3 -m pip install --retries 10 --timeout 120 deepspeed==0.16.4 tensorboard
+          python3 -m pip install "setuptools<80"
+          python3 -c "import pkg_resources"
+
+      - name: Install Ascend Megatron dependencies
+        shell: bash
+        run: |
+          export MEGATRON_ADAPTOR_SRC="/tmp/MegatronAdaptor"
+          export TRANSFORMER_ENGINE_NPU_SRC="/tmp/TransformerEngineNPU"
+          rm -rf "${MEGATRON_ADAPTOR_SRC}" "${TRANSFORMER_ENGINE_NPU_SRC}"
+          git clone --depth 1 --branch "${TRANSFORMER_ENGINE_NPU_REF}" \
+            "${TRANSFORMER_ENGINE_NPU_REPO}" "${TRANSFORMER_ENGINE_NPU_SRC}"
+          git clone --depth 1 --branch "${MEGATRON_ADAPTOR_REF}" \
+            "${MEGATRON_ADAPTOR_REPO}" "${MEGATRON_ADAPTOR_SRC}"
+          python3 -m pip install --no-build-isolation -e "${TRANSFORMER_ENGINE_NPU_SRC}"
+          python3 -m pip install --no-build-isolation -e "${MEGATRON_ADAPTOR_SRC}"
+
+      - name: Install ROLL
+        run: |
+          python3 -m pip install -e .
+
+      - name: Prepare Megatron test model
+        shell: bash
+        run: |
+          local_model="/data/cpfs_0/common/models/Qwen2.5-0.5B-Instruct"
+          if [ -d "${local_model}" ]; then
+            echo "ROLL_MEGATRON_TEST_MODEL=${local_model}" >> "${GITHUB_ENV}"
+            exit 0
+          fi
+
+          python3 - <<'PY'
+          import os
+          from huggingface_hub import snapshot_download
+
+          model_path = snapshot_download("Qwen/Qwen2.5-0.5B-Instruct")
+          with open(os.environ["GITHUB_ENV"], "a", encoding="utf-8") as env_file:
+              env_file.write(f"ROLL_MEGATRON_TEST_MODEL={model_path}\n")
+          PY
+
+      - name: Run MegatronAdaptor offload tests
+        shell: bash
+        run: |
+          torchrun --standalone --nnodes=1 --nproc-per-node=2 \
+            -m pytest -q --tb=short tests/third_party/megatron/test_offload_states.py
diff --git a/docs_roll/docs/User Guides/Hardware Support/ascend_docker_usage.md b/docs_roll/docs/User Guides/Hardware Support/ascend_docker_usage.md
index c37c92756..628048450 100644
--- a/docs_roll/docs/User Guides/Hardware Support/ascend_docker_usage.md	
+++ b/docs_roll/docs/User Guides/Hardware Support/ascend_docker_usage.md	
@@ -1,6 +1,6 @@
 # Running ROLL on Ascend NPU with Docker
 
-Last updated: 06/23/2026.
+Last updated: 07/16/2026.
 
 This guide explains how to get, build, and run ROLL images on **Huawei Ascend NPU**. Prefer the pre-built image when possible; use `Dockerfile.A2` or `Dockerfile.A3` when you need to customize dependencies. Ascend 950 currently follows the manual installation profile in [ROLL x Ascend](ascend_usage.md).
 
@@ -244,9 +244,9 @@ python -c "import vllm_ascend; print(f'vllm_ascend available')"
 
 ### Important Configuration Notes
 
-Since Megatron-LM is not supported on Ascend NPU, you need to use **FSDP2** as the training backend. Make sure your configuration files use the following settings:
+The bundled RLVR example uses **FSDP2** and runs with the base Ascend image dependencies. Compatible Megatron configurations on A2/A3 require the optional packages from [Install Megatron on Ascend](ascend_usage.md#install-megatron-on-ascend); the A2/A3 Dockerfiles do not install them by default.
 
-1. Set `strategy_args` to use FSDP2
+For the bundled FSDP2 RLVR example, set `strategy_args` to use FSDP2.
 
 ### Example: RLVR Pipeline
 
diff --git a/docs_roll/docs/User Guides/Hardware Support/ascend_npu_env_config.md b/docs_roll/docs/User Guides/Hardware Support/ascend_npu_env_config.md
index e15f3e203..428cf9b2a 100644
--- a/docs_roll/docs/User Guides/Hardware Support/ascend_npu_env_config.md	
+++ b/docs_roll/docs/User Guides/Hardware Support/ascend_npu_env_config.md	
@@ -153,7 +153,7 @@ export CPU_AFFINITY_CONF=1,npu0:0-1,npu1:2-3,npu2:4-5,npu3:6-7
 | -------- | ----------------- | ----------- |
 | `VLLM_USE_V1` | `1` | Enable vLLM V1 architecture. Required for vLLM-Ascend |
 | `VLLM_ATTENTION_BACKEND` | `XFORMERS` | vLLM attention computation backend |
-| `VLLM_ASCEND_ENABLE_FLASHCOMM` | `1` | Enable Ascend FlashComm high-speed communication optimization |
+| `VLLM_ASCEND_ENABLE_FLASHCOMM` | `0` | Enable Ascend FlashComm high-speed communication optimization |
 | `VLLM_ASCEND_ENABLE_PREFETCH_MLP` | `1` | Enable MLP layer weight prefetching. This replaces the older dense optimize toggle in current vLLM-Ascend releases. |
 | `VLLM_ASCEND_ENABLE_TOPK_OPTIMIZE` | `1` | Enable TopK operator fusion optimization for generation decoding |
 | `VLLM_ASCEND_MODEL_EXECUTE_TIME_OBSERVE` | `1` | Print prefill/decode phase timing details (for debugging) |
@@ -165,7 +165,7 @@ Example:
 ```bash
 export VLLM_USE_V1=1
 export VLLM_ATTENTION_BACKEND=XFORMERS
-export VLLM_ASCEND_ENABLE_FLASHCOMM=1
+export VLLM_ASCEND_ENABLE_FLASHCOMM=0
 export VLLM_ASCEND_ENABLE_PREFETCH_MLP=1
 ```
 
@@ -260,7 +260,7 @@ export OMP_NUM_THREADS=1
 
 # vLLM-Ascend inference
 export VLLM_USE_V1=1
-export VLLM_ASCEND_ENABLE_FLASHCOMM=1
+export VLLM_ASCEND_ENABLE_FLASHCOMM=0
 export VLLM_ASCEND_ENABLE_PREFETCH_MLP=1
 
 # Operator compilation cache
diff --git a/docs_roll/docs/User Guides/Hardware Support/ascend_npu_examples.md b/docs_roll/docs/User Guides/Hardware Support/ascend_npu_examples.md
index 60b4be3d6..9c04b5402 100644
--- a/docs_roll/docs/User Guides/Hardware Support/ascend_npu_examples.md	
+++ b/docs_roll/docs/User Guides/Hardware Support/ascend_npu_examples.md	
@@ -1,6 +1,6 @@
 # Ascend NPU End-to-End Configuration Examples
 
-Last updated: 04/27/2026.
+Last updated: 07/16/2026.
 
 This document provides end-to-end configuration examples for running ROLL on Huawei Ascend NPU, including environment setup, resource allocation, and launch commands for both single-node and multi-node scenarios.
 
@@ -12,19 +12,21 @@ Before running these examples, ensure you have:
 2. Verified the environment inside the container (see [Verify the Environment](ascend_docker_usage.md#verify-the-environment)).
 3. Downloaded the model weights to a directory accessible from inside the container.
 
-The repository currently includes a runnable Ascend RLVR example in `examples/ascend_examples`, including `qwen3_30b_rlvr_fsdp2.yaml` and `run_rlvr_pipeline.sh`.
+The repository includes runnable examples in `examples/ascend_examples`: an FSDP2 RLVR example (`qwen3_30b_rlvr_fsdp2.yaml` and `run_rlvr_pipeline.sh`) and a Megatron DPO example (`qwen3_4B_dpo_megatron.yaml` and `run_dpo_pipeline.sh`).
 
 ## Key Differences from GPU
 
-When adapting GPU configurations for NPU, the following changes are **required**:
+When adapting GPU configurations for NPU, select either the FSDP2 path used by the examples below or a compatible Megatron configuration with its optional NPU dependencies:
 
 | Item | GPU | NPU |
 | ---- | --- | --- |
-| Training backend | Megatron or FSDP2 | FSDP2 only (Megatron not supported on NPU) |
+| Training backend | Megatron or FSDP2 | FSDP2 in the examples below; Megatron for compatible A2/A3 configurations after optional dependencies are installed |
 | Attention implementation | `flash_attn` or `fa2` | `fa2` via `transformers` (not `flash_attn` package) |
 | Communication backend | NCCL | HCCL |
 | Device visibility | `CUDA_VISIBLE_DEVICES` | `ASCEND_RT_VISIBLE_DEVICES` |
 
+For the Megatron path, first complete [Install Megatron on Ascend](ascend_usage.md#install-megatron-on-ascend). The following Agentic examples intentionally remain on FSDP2.
+
 ## Example 1: Single-Node Agentic Pipeline (Qwen2.5-0.5B)
 
 This example runs the FrozenLake agentic pipeline on a single 8-NPU node using FSDP2.
@@ -82,7 +84,7 @@ export OMP_NUM_THREADS=1
 # vLLM-Ascend inference
 export VLLM_USE_V1=1
 export VLLM_ASCEND_ENABLE_NZ=0
-export VLLM_ASCEND_ENABLE_FLASHCOMM=1
+export VLLM_ASCEND_ENABLE_FLASHCOMM=0
 export VLLM_ASCEND_ENABLE_PREFETCH_MLP=1
 
 # Operator compilation cache
@@ -175,7 +177,7 @@ actor_train:
       param_dtype: bf16
       reduce_dtype: bf16
       reshard_after_forward: true
-      offload_policy: false    # NPU: Must use FSDP2, NOT megatron_train
+      offload_policy: false    # NPU FSDP2 example
   device_mapping: list(range(0,4))    # NPU: Training on NPUs 0-3
   infer_batch_size: 2
 
@@ -550,7 +552,7 @@ export OMP_NUM_THREADS=1
 # === vLLM-Ascend inference ===
 export VLLM_USE_V1=1
 export VLLM_ASCEND_ENABLE_NZ=0
-export VLLM_ASCEND_ENABLE_FLASHCOMM=1
+export VLLM_ASCEND_ENABLE_FLASHCOMM=0
 export VLLM_ASCEND_ENABLE_PREFETCH_MLP=1
 
 # === Operator compilation cache ===
diff --git a/docs_roll/docs/User Guides/Hardware Support/ascend_npu_faq.md b/docs_roll/docs/User Guides/Hardware Support/ascend_npu_faq.md
index 82122525d..ff543fe62 100644
--- a/docs_roll/docs/User Guides/Hardware Support/ascend_npu_faq.md	
+++ b/docs_roll/docs/User Guides/Hardware Support/ascend_npu_faq.md	
@@ -1,6 +1,6 @@
 # Ascend NPU FAQ
 
-Last updated: 04/27/2026.
+Last updated: 07/16/2026.
 
 This document compiles common issues encountered when running ROLL on Huawei Ascend NPU and their solutions.
 
@@ -93,11 +93,23 @@ pip install triton-ascend==3.2.1 --extra-index-url https://mirrors.huaweicloud.c
 
 ## Training Configuration
 
-### Megatron Strategy Not Supported
+### Megatron Strategy Initialization Error
 
-**Symptom:** Errors when using `strategy: megatron` in configuration on NPU.
+**Symptom:** Import or initialization errors occur when using `megatron_train` or `megatron_infer` on NPU.
 
-**Solution:** Megatron-LM is not supported on Ascend NPU. Use FSDP2 as the training backend:
+**Solution:** Megatron is available for validated A2/A3 configurations, but its optional dependencies are not part of the base Ascend installation. Complete [Install Megatron on Ascend](ascend_usage.md#install-megatron-on-ascend), then verify the three packages in an environment where the Ascend toolkit has been initialized:
+
+```bash
+python - <<'PY'
+import megatron_adaptor
+import megatron.core
+import transformer_engine.pytorch
+
+print("MegatronAdaptor NPU dependencies are available.")
+PY
+```
+
+If you do not need the Megatron path, select FSDP2 instead and keep the matching FSDP2 strategy configuration:
 
 ```yaml
 strategy_args:
diff --git a/docs_roll/docs/User Guides/Hardware Support/ascend_npu_rlvr.md b/docs_roll/docs/User Guides/Hardware Support/ascend_npu_rlvr.md
index 63b520965..9f1637563 100644
--- a/docs_roll/docs/User Guides/Hardware Support/ascend_npu_rlvr.md	
+++ b/docs_roll/docs/User Guides/Hardware Support/ascend_npu_rlvr.md	
@@ -1,6 +1,6 @@
 # Running RLVR Pipeline on Ascend NPU
 
-Last updated: 04/28/2026.
+Last updated: 07/16/2026.
 
 This guide provides a complete end-to-end walkthrough for running the RLVR (Reinforcement Learning with Verifiable Rewards) pipeline on Huawei Ascend NPU, covering environment setup, data preparation, model download, configuration, training launch, monitoring & evaluation, and checkpoint resumption.
 
@@ -220,21 +220,23 @@ reward_pretrain: Qwen/Qwen2.5-7B
 
 ### Key Differences from GPU
 
-When adapting the GPU RLVR configuration for NPU, the following changes are **required**:
+The bundled NPU RLVR example uses FSDP2. When adapting another configuration, choose the training and reference strategies together and apply the corresponding NPU dependencies:
 
 | Item | GPU | NPU |
 | ---- | --- | --- |
-| Training backend | Megatron or FSDP2 | FSDP2 only (Megatron not supported on NPU) |
+| Training backend | Megatron or FSDP2 | FSDP2 in the bundled RLVR example; Megatron for compatible A2/A3 configurations after optional dependencies are installed |
 | Inference backend | vLLM | vLLM-Ascend |
-| Reference model strategy | `megatron_infer` | `fsdp2_infer` |
+| Reference model strategy | `megatron_infer` | `fsdp2_infer` for the FSDP2 path; `megatron_infer` for the Megatron path |
 | Attention implementation | `flash_attn` or `fa2` | `fa2` via `transformers` (not `flash_attn` package) |
 | Communication backend | NCCL | HCCL |
 | Device visibility | `CUDA_VISIBLE_DEVICES` | `ASCEND_RT_VISIBLE_DEVICES` |
-| Sharding config | FSDP2 or Megatron optimizer sharding | FSDP2 with `offload_policy: true` recommended for 7B+ models |
+| Sharding config | FSDP2 or Megatron optimizer sharding | FSDP2 with `offload_policy: true` recommended for 7B+ models, or the matching Megatron parallel configuration |
+
+Before selecting `megatron_train` or `megatron_infer`, complete [Install Megatron on Ascend](ascend_usage.md#install-megatron-on-ascend). The configuration below remains an FSDP2 example and does not require the optional Megatron dependencies.
 
 ### Complete NPU Configuration Example
 
-Below is a complete NPU-adapted configuration (adapted from `examples/ascend_examples/qwen3_30b_rlvr_fsdp2.yaml`), with key differences marked with `# NPU` comments:
+Below is a complete NPU FSDP2 configuration (adapted from `examples/ascend_examples/qwen3_30b_rlvr_fsdp2.yaml`), with key differences marked with `# NPU` comments:
 
 ```yaml
 hydra:
@@ -322,7 +324,7 @@ actor_train:
     interleave_probs: "1.0"
     preprocessing_num_workers: 16
   strategy_args:
-    strategy_name: fsdp2_train      # NPU: Must use FSDP2, NOT megatron_train
+    strategy_name: fsdp2_train      # NPU FSDP2 example
     strategy_config:
       fsdp_size: 16                 # NPU: FSDP2 sharding size
       param_dtype: bf16
@@ -377,7 +379,7 @@ reference:
   data_args:
     template: qwen2_5
   strategy_args:
-    strategy_name: fsdp2_infer          # NPU: Use fsdp2_infer, NOT megatron_infer
+    strategy_name: fsdp2_infer          # NPU FSDP2 reference example
     strategy_config:
       fsdp_size: 16
       param_dtype: bf16
@@ -403,7 +405,7 @@ rewards:
 
 ### Key Configuration Changes Explained
 
-#### 1. Training Strategy: FSDP2 instead of Megatron
+#### 1. Training Strategy: FSDP2 Example
 
 ```yaml
 # GPU (original)
@@ -428,7 +430,7 @@ actor_train:
 
 For 7B models on 4 NPUs, set `offload_policy: true` to enable CPU offloading and avoid OOM. For smaller models (e.g., 0.5B), `offload_policy: false` may be sufficient.
 
-#### 2. Reference Model: fsdp2_infer instead of megatron_infer
+#### 2. Reference Model: FSDP2 Example
 
 ```yaml
 # GPU
@@ -762,17 +764,17 @@ The following RLVR reward workers are supported on NPU:
 When using `LLMJudgeRewardWorker`, the judge model requires its own NPU devices for inference. Ensure you allocate separate NPUs in `device_mapping` for the judge model, and do not share them with `actor_infer` or `actor_train`.
 :::
 
-## GPU-to-NPU Configuration Migration Checklist
+## GPU-to-NPU FSDP2 Migration Checklist
 
-Use this checklist when migrating an existing GPU RLVR configuration to NPU:
+Use this checklist when migrating an existing GPU RLVR configuration to the FSDP2 path shown in this guide. Megatron configurations should retain their Megatron strategy settings and use the dependencies documented in [Install Megatron on Ascend](ascend_usage.md#install-megatron-on-ascend).
 
-- [ ] Change `actor_train.strategy_args.strategy_name` from `megatron_train` to `fsdp2_train`
+- [ ] If starting from a Megatron configuration, change `actor_train.strategy_args.strategy_name` from `megatron_train` to `fsdp2_train`
 - [ ] Change `actor_train.strategy_args.strategy_config` to FSDP2 config (with `offload_policy: true` for 7B+ models)
-- [ ] Change `reference.strategy_args.strategy_name` from `megatron_infer` to `fsdp2_infer`
+- [ ] If starting from a Megatron configuration, change `reference.strategy_args.strategy_name` from `megatron_infer` to `fsdp2_infer`
 - [ ] Set `reference.strategy_args.strategy_config` to FSDP2 config matching `actor_train`
 - [ ] Add `attn_implementation: fa2` to `actor_train.model_args` and `reference.model_args`
 - [ ] Remove any `flash_attn` references
-- [ ] Remove any Megatron-specific config (e.g., `tensor_model_parallel_size`, `pipeline_model_parallel_size`)
+- [ ] For the FSDP2 path, remove Megatron-specific config (e.g., `tensor_model_parallel_size`, `pipeline_model_parallel_size`)
 - [ ] Verify `llm_judge` reward worker has separate NPU allocation (if used)
 
 ## Troubleshooting
diff --git a/docs_roll/docs/User Guides/Hardware Support/ascend_usage.md b/docs_roll/docs/User Guides/Hardware Support/ascend_usage.md
index ab43a1d4a..516ab5fef 100644
--- a/docs_roll/docs/User Guides/Hardware Support/ascend_usage.md	
+++ b/docs_roll/docs/User Guides/Hardware Support/ascend_usage.md	
@@ -1,6 +1,6 @@
 # ROLL x Ascend
 
-Last updated: 06/23/2026.
+Last updated: 07/16/2026.
 
 We have added support for Huawei Ascend devices in ROLL.
 
@@ -115,6 +115,58 @@ pip install -v -e .
 cd ..
 ```
 
+### Install Megatron on Ascend
+
+This step is optional. Install these dependencies when a ROLL configuration uses the `megatron_train` or `megatron_infer` strategy. The versions below match the combination validated by the current NPU CI:
+
+| Software | Version |
+| -------- | ------- |
+| Megatron-Core | `core_r0.17.0` |
+| TransformerEngineNPU | `main` |
+| MegatronAdaptor | `core_r0.17.0` |
+
+Install the packages in the following order. Megatron-Core 0.17 declares Python 3.12 in its package metadata, while the ROLL Ascend environment uses Python 3.11, so its editable installation must ignore the metadata Python requirement:
+
+```
+python -m pip install --upgrade pip wheel
+python -m pip install "setuptools<80" pybind11 "packaging>=24.2"
+
+# Megatron-Core 0.17
+export MEGATRON_CORE_SRC=/tmp/Megatron-LM
+rm -rf "${MEGATRON_CORE_SRC}"
+git clone --depth 1 --branch core_r0.17.0 \
+  https://github.com/NVIDIA/Megatron-LM.git "${MEGATRON_CORE_SRC}"
+python -m pip install --ignore-requires-python --no-build-isolation --no-deps \
+  -e "${MEGATRON_CORE_SRC}"
+
+# Ascend Transformer Engine and Megatron patches
+export TRANSFORMER_ENGINE_NPU_SRC=/tmp/TransformerEngineNPU
+export MEGATRON_ADAPTOR_SRC=/tmp/MegatronAdaptor
+rm -rf "${TRANSFORMER_ENGINE_NPU_SRC}" "${MEGATRON_ADAPTOR_SRC}"
+git clone --depth 1 --branch main \
+  https://gitcode.com/Ascend/TransformerEngineNPU.git \
+  "${TRANSFORMER_ENGINE_NPU_SRC}"
+git clone --depth 1 --branch core_r0.17.0 \
+  https://gitcode.com/Ascend/MegatronAdaptor.git \
+  "${MEGATRON_ADAPTOR_SRC}"
+python -m pip install --no-build-isolation -e "${TRANSFORMER_ENGINE_NPU_SRC}"
+python -m pip install --no-build-isolation -e "${MEGATRON_ADAPTOR_SRC}"
+```
+
+Verify the installation in an environment where the Ascend toolkit has been initialized:
+
+```
+python - <<'PY'
+import megatron_adaptor
+import megatron.core
+import transformer_engine.pytorch
+
+print("MegatronAdaptor NPU dependencies are available.")
+PY
+```
+
+Do not install the NVIDIA `transformer-engine[pytorch]` package for this NPU setup. TransformerEngineNPU provides the Ascend-compatible implementation.
+
 ### Install ROLL
 
 ```
@@ -131,11 +183,11 @@ cd ..
 | --------------------------- | ------------- |
 | transformers                | >= v4.57.6    |
 | flash_attn                  | not supported |
-| transformer-engine[pytorch] | not supported |
+| transformer-engine[pytorch] | not supported on Ascend; use TransformerEngineNPU for Megatron |
 
 1. `transformers` v4.57.6 supports enabling `--flash_attention_2`.
 2. `flash_attn` acceleration is not supported currently.
-3. `transformer-engine[pytorch]` is currently not supported.
+3. The NVIDIA `transformer-engine[pytorch]` package is not supported on Ascend. Megatron configurations use TransformerEngineNPU instead.
 
 ```
 pip install transformers==4.57.6
@@ -143,8 +195,7 @@ pip install transformers==4.57.6
 
 ## Quick Start: Single-Node Deployment
 
-Before full usage, we recommend testing the single-node pipeline to verify your environment and installation.
-Since Megatron-LM is not supported on NPU, first change `strategy_args` in the relevant files to use the `fsdp2` option.
+Before full usage, we recommend testing the single-node pipeline to verify your environment and installation. FSDP2 examples can use the base Ascend environment. Before running a configuration with `megatron_train` or `megatron_infer`, complete [Install Megatron on Ascend](#install-megatron-on-ascend).
 
 
 1. Run the single-node pipeline via shell:
@@ -153,7 +204,13 @@ Since Megatron-LM is not supported on NPU, first change `strategy_args` in the r
 bash examples/agentic_demo/run_agentic_pipeline_frozen_lake_single_node_demo.sh  
 ```
 
-2. Run the agentic pipeline using a config file:
+2. Run the Megatron DPO example:
+
+```
+bash examples/ascend_examples/run_dpo_pipeline.sh
+```
+
+3. Run the agentic pipeline using a config file:
 
 ```
 # Make sure you are in the root directory of the ROLL project
@@ -173,6 +230,7 @@ python examples/start_agentic_pipeline.py \
 | Agentic         | examples/qwen2.5-0.5B-agentic/run_agentic_pipeline_sokoban.sh | FSDP2            | vLLM              | Atlas 900 A2/A3 PODc |
 | Agentic-Rollout | examples/qwen2.5-0.5B-agentic/run_agentic_rollout_sokoban.sh | FSDP2            | vLLM              | Atlas 900 A2/A3 PODc |
 | RLVR            | examples/ascend_examples/run_rlvr_pipeline.sh                | FSDP2            | vLLM              | Atlas 900 A2/A3/Ascend 950 training series |
+| DPO             | examples/ascend_examples/run_dpo_pipeline.sh                 | Megatron         | Megatron          | Atlas 900 A2/A3 training series |
 
 ## Disclaimer
 
diff --git a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Hardware Support/ascend_docker_usage.md b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Hardware Support/ascend_docker_usage.md
index f41c20be6..2dae6f3e0 100644
--- a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Hardware Support/ascend_docker_usage.md	
+++ b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Hardware Support/ascend_docker_usage.md	
@@ -1,6 +1,6 @@
 # 使用 Docker 在昇腾 NPU 上运行 ROLL
 
-最后更新：2026/06/23。
+最后更新：2026/07/16。
 
 本指南介绍如何在**华为昇腾 NPU** 上获取、构建并运行 ROLL 镜像。推荐优先使用预构建镜像；如需自定义依赖，再使用 `Dockerfile.A2` 或 `Dockerfile.A3` 构建。Atlas 950 当前使用 [ROLL x Ascend](ascend_usage.md) 中的手动安装配置。
 
@@ -244,9 +244,9 @@ python -c "import vllm_ascend; print(f'vllm_ascend available')"
 
 ### 重要配置说明
 
-由于昇腾 NPU 上不支持 Megatron-LM 训练，需要使用 **FSDP2** 作为训练后端。请确保配置文件中使用以下设置：
+仓库内置的 RLVR 示例采用 **FSDP2**，使用昇腾基础镜像中的依赖即可运行。A2/A3 上兼容的 Megatron 配置需要按照[在昇腾上安装 Megatron](ascend_usage.md#在昇腾上安装-megatron)补充可选依赖；A2/A3 Dockerfile 默认不会安装这些依赖。
 
-1. 将 `strategy_args` 设置为使用 FSDP2
+运行内置 FSDP2 RLVR 示例时，请将 `strategy_args` 设置为 FSDP2。
 
 
 ### 示例：RLVR 流水线
diff --git a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Hardware Support/ascend_npu_env_config.md b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Hardware Support/ascend_npu_env_config.md
index 8121dfc67..dd764c19d 100644
--- a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Hardware Support/ascend_npu_env_config.md	
+++ b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Hardware Support/ascend_npu_env_config.md	
@@ -153,7 +153,7 @@ export CPU_AFFINITY_CONF=1,npu0:0-1,npu1:2-3,npu2:4-5,npu3:6-7
 | ---- | ------ | ---- |
 | `VLLM_USE_V1` | `1` | 启用 vLLM V1 架构，vLLM-Ascend 必需 |
 | `VLLM_ATTENTION_BACKEND` | `XFORMERS` | vLLM 注意力计算后端 |
-| `VLLM_ASCEND_ENABLE_FLASHCOMM` | `1` | 启用昇腾 FlashComm 高速通信优化 |
+| `VLLM_ASCEND_ENABLE_FLASHCOMM` | `0` | 启用昇腾 FlashComm 高速通信优化 |
 | `VLLM_ASCEND_ENABLE_PREFETCH_MLP` | `1` | 启用 MLP 层权重预取。它替代了较早版本中的 dense optimize 开关。 |
 | `VLLM_ASCEND_ENABLE_TOPK_OPTIMIZE` | `1` | 启用 TopK 算子融合优化，提升生成解码性能 |
 | `VLLM_ASCEND_MODEL_EXECUTE_TIME_OBSERVE` | `1` | 打印 prefill/decode 阶段耗时详情（调试用） |
@@ -165,7 +165,7 @@ export CPU_AFFINITY_CONF=1,npu0:0-1,npu1:2-3,npu2:4-5,npu3:6-7
 ```bash
 export VLLM_USE_V1=1
 export VLLM_ATTENTION_BACKEND=XFORMERS
-export VLLM_ASCEND_ENABLE_FLASHCOMM=1
+export VLLM_ASCEND_ENABLE_FLASHCOMM=0
 export VLLM_ASCEND_ENABLE_PREFETCH_MLP=1
 ```
 
@@ -260,7 +260,7 @@ export OMP_NUM_THREADS=1
 
 # vLLM-Ascend 推理
 export VLLM_USE_V1=1
-export VLLM_ASCEND_ENABLE_FLASHCOMM=1
+export VLLM_ASCEND_ENABLE_FLASHCOMM=0
 export VLLM_ASCEND_ENABLE_PREFETCH_MLP=1
 
 # 算子编译缓存
diff --git a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Hardware Support/ascend_npu_examples.md b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Hardware Support/ascend_npu_examples.md
index a04295c94..eb2ad5f7b 100644
--- a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Hardware Support/ascend_npu_examples.md	
+++ b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Hardware Support/ascend_npu_examples.md	
@@ -1,6 +1,6 @@
 # 昇腾 NPU 端到端配置样例
 
-最后更新：2026/04/27。
+最后更新：2026/07/16。
 
 本文档提供在华为昇腾 NPU 上运行 ROLL 的端到端配置样例，涵盖环境准备、资源切分和启动命令，适用于单机和多机场景。
 
@@ -12,20 +12,22 @@
 2. 已在容器内验证环境（参见 [验证环境](ascend_docker_usage.md#验证环境)）。
 3. 已将模型权重下载到容器可访问的目录。
 
-当前仓库在 `examples/ascend_examples` 中提供可直接运行的昇腾 RLVR 示例，包括 `qwen3_30b_rlvr_fsdp2.yaml` 和 `run_rlvr_pipeline.sh`。
+当前仓库在 `examples/ascend_examples` 中提供可直接运行的示例：FSDP2 RLVR 示例（`qwen3_30b_rlvr_fsdp2.yaml` 和 `run_rlvr_pipeline.sh`）以及 Megatron DPO 示例（`qwen3_4B_dpo_megatron.yaml` 和 `run_dpo_pipeline.sh`）。
 
 
 ## GPU 与 NPU 的关键差异
 
-将 GPU 配置适配到 NPU 时，**必须**进行以下修改：
+将 GPU 配置适配到 NPU 时，可选择下文示例采用的 FSDP2 路径，或为兼容的 Megatron 配置安装可选的 NPU 依赖：
 
 | 项目 | GPU | NPU |
 | ---- | --- | --- |
-| 训练后端 | Megatron 或 FSDP2 | 仅 FSDP2（NPU 不支持 Megatron） |
+| 训练后端 | Megatron 或 FSDP2 | 下文示例采用 FSDP2；安装可选依赖后，兼容的 A2/A3 配置可使用 Megatron |
 | 注意力实现 | `flash_attn` 或 `fa2` | 通过 `transformers` 使用 `fa2`（不能使用 `flash_attn` 包） |
 | 通信后端 | NCCL | HCCL |
 | 设备可见性 | `CUDA_VISIBLE_DEVICES` | `ASCEND_RT_VISIBLE_DEVICES` |
 
+选择 Megatron 路径时，请先完成[在昇腾上安装 Megatron](ascend_usage.md#在昇腾上安装-megatron)。下方 Agentic 示例有意保留为 FSDP2 配置。
+
 ## 样例 1：单机 Agentic 流水线（Qwen2.5-0.5B）
 
 本样例在单个 8 卡 NPU 节点上使用 FSDP2 运行 FrozenLake Agentic 流水线。
@@ -83,7 +85,7 @@ export OMP_NUM_THREADS=1
 # vLLM-Ascend 推理
 export VLLM_USE_V1=1
 export VLLM_ASCEND_ENABLE_NZ=0
-export VLLM_ASCEND_ENABLE_FLASHCOMM=1
+export VLLM_ASCEND_ENABLE_FLASHCOMM=0
 export VLLM_ASCEND_ENABLE_PREFETCH_MLP=1
 
 # 算子编译缓存
@@ -176,7 +178,7 @@ actor_train:
       param_dtype: bf16
       reduce_dtype: bf16
       reshard_after_forward: true
-      offload_policy: false    # NPU: 必须使用 FSDP2，不能用 megatron_train
+      offload_policy: false    # NPU FSDP2 示例
   device_mapping: list(range(0,4))    # NPU: 训练使用 NPU 0-3
   infer_batch_size: 2
 
@@ -551,7 +553,7 @@ export OMP_NUM_THREADS=1
 # === vLLM-Ascend 推理 ===
 export VLLM_USE_V1=1
 export VLLM_ASCEND_ENABLE_NZ=0
-export VLLM_ASCEND_ENABLE_FLASHCOMM=1
+export VLLM_ASCEND_ENABLE_FLASHCOMM=0
 export VLLM_ASCEND_ENABLE_PREFETCH_MLP=1
 
 # === 算子编译缓存 ===
diff --git a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Hardware Support/ascend_npu_faq.md b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Hardware Support/ascend_npu_faq.md
index fdb3e427c..13908056a 100644
--- a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Hardware Support/ascend_npu_faq.md	
+++ b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Hardware Support/ascend_npu_faq.md	
@@ -1,6 +1,6 @@
 # 昇腾 NPU 常见问题
 
-最后更新：2026/04/27。
+最后更新：2026/07/16。
 
 本文档汇总了在华为昇腾 NPU 上运行 ROLL 时可能遇到的常见问题及解决方案。
 
@@ -93,11 +93,23 @@ pip install triton-ascend==3.2.1 --extra-index-url https://mirrors.huaweicloud.c
 
 ## 训练配置
 
-### 不支持 Megatron 策略
+### Megatron 策略初始化错误
 
-**现象：** 在 NPU 上使用 `strategy: megatron` 配置时报错。
+**现象：** 在 NPU 上使用 `megatron_train` 或 `megatron_infer` 时出现导入或初始化错误。
 
-**解决方案：** 昇腾 NPU 上不支持 Megatron-LM 训练，请使用 FSDP2 作为训练后端：
+**解决方案：** 已验证的 A2/A3 配置可以使用 Megatron，但基础昇腾环境不会默认安装其可选依赖。请先完成[在昇腾上安装 Megatron](ascend_usage.md#在昇腾上安装-megatron)，然后在已初始化昇腾 Toolkit 的环境中验证三个依赖包：
+
+```bash
+python - <<'PY'
+import megatron_adaptor
+import megatron.core
+import transformer_engine.pytorch
+
+print("MegatronAdaptor NPU dependencies are available.")
+PY
+```
+
+如果不需要 Megatron 路径，也可以选择 FSDP2，并保留与之匹配的 FSDP2 策略配置：
 
 ```yaml
 strategy_args:
diff --git a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Hardware Support/ascend_npu_rlvr.md b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Hardware Support/ascend_npu_rlvr.md
index 38b59331c..dfaa15aba 100644
--- a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Hardware Support/ascend_npu_rlvr.md	
+++ b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Hardware Support/ascend_npu_rlvr.md	
@@ -1,6 +1,6 @@
 # 在昇腾 NPU 上运行 RLVR 流水线
 
-最后更新：2026/04/28。
+最后更新：2026/07/16。
 
 本文档提供在华为昇腾 NPU 上运行 RLVR（Reinforcement Learning with Verifiable Rewards）流水线的端到端指南，涵盖环境准备、数据准备、模型下载、配置编写、训练启动、监控与评估，以及从 checkpoint 恢复训练。
 
@@ -220,21 +220,23 @@ reward_pretrain: Qwen/Qwen2.5-7B
 
 ### 与 GPU 的关键差异
 
-将 GPU RLVR 配置适配到 NPU 时，**必须**进行以下修改：
+仓库内置的 NPU RLVR 示例采用 FSDP2。适配其他配置时，请配套选择训练和 Reference 策略，并安装对应的 NPU 依赖：
 
 | 项目 | GPU | NPU |
 | ---- | --- | --- |
-| 训练后端 | Megatron 或 FSDP2 | 仅 FSDP2（NPU 不支持 Megatron） |
+| 训练后端 | Megatron 或 FSDP2 | 内置 RLVR 示例采用 FSDP2；安装可选依赖后，兼容的 A2/A3 配置可使用 Megatron |
 | 推理后端 | vLLM | vLLM-Ascend |
-| Reference 模型策略 | `megatron_infer` | `fsdp2_infer` |
+| Reference 模型策略 | `megatron_infer` | FSDP2 路径使用 `fsdp2_infer`；Megatron 路径使用 `megatron_infer` |
 | 注意力实现 | `flash_attn` 或 `fa2` | 通过 `transformers` 使用 `fa2`（不能使用 `flash_attn` 包） |
 | 通信后端 | NCCL | HCCL |
 | 设备可见性 | `CUDA_VISIBLE_DEVICES` | `ASCEND_RT_VISIBLE_DEVICES` |
-| 分片配置 | FSDP2 或 Megatron 优化器分片 | FSDP2，7B+ 模型推荐 `offload_policy: true` |
+| 分片配置 | FSDP2 或 Megatron 优化器分片 | FSDP2 路径中 7B+ 模型推荐 `offload_policy: true`，或使用对应的 Megatron 并行配置 |
+
+选择 `megatron_train` 或 `megatron_infer` 前，请先完成[在昇腾上安装 Megatron](ascend_usage.md#在昇腾上安装-megatron)。下方配置仍是 FSDP2 示例，不需要安装可选的 Megatron 依赖。
 
 ### 完整 NPU 配置样例
 
-下面是一个完整的 NPU 适配配置（改编自 `examples/ascend_examples/qwen3_30b_rlvr_fsdp2.yaml`），关键差异使用 `# NPU` 注释标记：
+下面是一个完整的 NPU FSDP2 配置（改编自 `examples/ascend_examples/qwen3_30b_rlvr_fsdp2.yaml`），关键差异使用 `# NPU` 注释标记：
 
 ```yaml
 hydra:
@@ -322,7 +324,7 @@ actor_train:
     interleave_probs: "1.0"
     preprocessing_num_workers: 16
   strategy_args:
-    strategy_name: fsdp2_train      # NPU：必须使用 FSDP2，不能用 megatron_train
+    strategy_name: fsdp2_train      # NPU FSDP2 示例
     strategy_config:
       fsdp_size: 16                 # NPU：FSDP2 分片大小
       param_dtype: bf16
@@ -377,7 +379,7 @@ reference:
   data_args:
     template: qwen2_5
   strategy_args:
-    strategy_name: fsdp2_infer          # NPU：使用 fsdp2_infer，不能用 megatron_infer
+    strategy_name: fsdp2_infer          # NPU FSDP2 Reference 示例
     strategy_config:
       fsdp_size: 16
       param_dtype: bf16
@@ -403,7 +405,7 @@ rewards:
 
 ### 关键配置变更说明
 
-#### 1. 训练策略：使用 FSDP2 替代 Megatron
+#### 1. 训练策略：FSDP2 示例
 
 ```yaml
 # GPU（原始配置）
@@ -428,7 +430,7 @@ actor_train:
 
 在 4 张 NPU 上运行 7B 模型时，设置 `offload_policy: true` 可以启用 CPU offloading 避免 OOM。对于更小的模型（如 0.5B），`offload_policy: false` 可能已经足够。
 
-#### 2. Reference 模型：使用 fsdp2_infer 替代 megatron_infer
+#### 2. Reference 模型：FSDP2 示例
 
 ```yaml
 # GPU
@@ -762,17 +764,17 @@ NPU 上支持以下 RLVR Reward Worker：
 使用 `LLMJudgeRewardWorker` 时，judge 模型需要独立的 NPU 设备进行推理。请确保在 `device_mapping` 中为 judge 模型分配独立 NPU，不要与 `actor_infer` 或 `actor_train` 共享。
 :::
 
-## GPU 到 NPU 配置迁移 Checklist
+## GPU 到 NPU 的 FSDP2 迁移 Checklist
 
-将已有 GPU RLVR 配置迁移到 NPU 时，可使用以下 checklist：
+将已有 GPU RLVR 配置迁移到本文展示的 FSDP2 路径时，可使用以下 checklist。Megatron 配置应保留 Megatron 策略设置，并使用[在昇腾上安装 Megatron](ascend_usage.md#在昇腾上安装-megatron)中说明的依赖。
 
-- [ ] 将 `actor_train.strategy_args.strategy_name` 从 `megatron_train` 改为 `fsdp2_train`
+- [ ] 如果从 Megatron 配置迁移，将 `actor_train.strategy_args.strategy_name` 从 `megatron_train` 改为 `fsdp2_train`
 - [ ] 将 `actor_train.strategy_args.strategy_config` 改为 FSDP2 配置（7B+ 模型使用 `offload_policy: true`）
-- [ ] 将 `reference.strategy_args.strategy_name` 从 `megatron_infer` 改为 `fsdp2_infer`
+- [ ] 如果从 Megatron 配置迁移，将 `reference.strategy_args.strategy_name` 从 `megatron_infer` 改为 `fsdp2_infer`
 - [ ] 将 `reference.strategy_args.strategy_config` 设置为与 `actor_train` 一致的 FSDP2 配置
 - [ ] 在 `actor_train.model_args` 和 `reference.model_args` 中添加 `attn_implementation: fa2`
 - [ ] 移除所有 `flash_attn` 引用
-- [ ] 移除所有 Megatron 专属配置（如 `tensor_model_parallel_size`、`pipeline_model_parallel_size`）
+- [ ] FSDP2 路径中移除 Megatron 专属配置（如 `tensor_model_parallel_size`、`pipeline_model_parallel_size`）
 - [ ] 如果使用 `llm_judge` reward worker，确认它有独立的 NPU 分配
 
 ## 常见问题
diff --git a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Hardware Support/ascend_usage.md b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Hardware Support/ascend_usage.md
index 7c6f7941b..f945923d5 100644
--- a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Hardware Support/ascend_usage.md	
+++ b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Hardware Support/ascend_usage.md	
@@ -1,6 +1,6 @@
 # ROLL x Ascend
 
-最后更新：2026/06/23。
+最后更新：2026/07/16。
 
 我们在 ROLL 上增加对华为昇腾设备的支持。
 
@@ -116,6 +116,57 @@ pip install -v -e .
 cd ..
 ```
 
+### 在昇腾上安装 Megatron
+
+此步骤为可选项。当 ROLL 配置使用 `megatron_train` 或 `megatron_infer` 策略时，需要安装以下依赖。这里的版本组合与当前 NPU CI 验证环境保持一致：
+
+| 软件 | 版本 |
+| ---- | ---- |
+| Megatron-Core | `core_r0.17.0` |
+| TransformerEngineNPU | `main` |
+| MegatronAdaptor | `core_r0.17.0` |
+
+
+```
+python -m pip install --upgrade pip wheel
+python -m pip install "setuptools<80" pybind11 "packaging>=24.2"
+
+# Megatron-Core 0.17
+export MEGATRON_CORE_SRC=/tmp/Megatron-LM
+rm -rf "${MEGATRON_CORE_SRC}"
+git clone --depth 1 --branch core_r0.17.0 \
+  https://github.com/NVIDIA/Megatron-LM.git "${MEGATRON_CORE_SRC}"
+python -m pip install --ignore-requires-python --no-build-isolation --no-deps \
+  -e "${MEGATRON_CORE_SRC}"
+
+# 昇腾 Transformer Engine 和 Megatron 适配补丁
+export TRANSFORMER_ENGINE_NPU_SRC=/tmp/TransformerEngineNPU
+export MEGATRON_ADAPTOR_SRC=/tmp/MegatronAdaptor
+rm -rf "${TRANSFORMER_ENGINE_NPU_SRC}" "${MEGATRON_ADAPTOR_SRC}"
+git clone --depth 1 --branch main \
+  https://gitcode.com/Ascend/TransformerEngineNPU.git \
+  "${TRANSFORMER_ENGINE_NPU_SRC}"
+git clone --depth 1 --branch core_r0.17.0 \
+  https://gitcode.com/Ascend/MegatronAdaptor.git \
+  "${MEGATRON_ADAPTOR_SRC}"
+python -m pip install --no-build-isolation -e "${TRANSFORMER_ENGINE_NPU_SRC}"
+python -m pip install --no-build-isolation -e "${MEGATRON_ADAPTOR_SRC}"
+```
+
+在已初始化昇腾 Toolkit 环境的终端中验证安装：
+
+```
+python - <<'PY'
+import megatron_adaptor
+import megatron.core
+import transformer_engine.pytorch
+
+print("MegatronAdaptor NPU dependencies are available.")
+PY
+```
+
+此 NPU 环境不要安装 NVIDIA `transformer-engine[pytorch]` 包。TransformerEngineNPU 提供适配昇腾的实现。
+
 ### 安装 ROLL
 
 ```
@@ -132,11 +183,11 @@ cd ..
 | ---- | ---- |
 | transformers | >= v4.57.6 |
 | flash_attn | 不支持 |
-| transformer-engine[pytorch] | 不支持 |
+| transformer-engine[pytorch] | 昇腾上不支持；Megatron 请使用 TransformerEngineNPU |
 
 1. `transformers` v4.57.6 支持启用 `--flash_attention_2`。
 2. 目前不支持 `flash_attn` 加速。
-3. 目前不支持 `transformer-engine[pytorch]`。
+3. 昇腾上不支持 NVIDIA `transformer-engine[pytorch]` 包，Megatron 配置改用 TransformerEngineNPU。
 
 ```
 pip install transformers==4.57.6
@@ -144,8 +195,7 @@ pip install transformers==4.57.6
 
 ## 快速开始：单节点部署指引
 
-正式使用前，建议您通过对单节点流水线的训练尝试以检验环境准备和安装的正确性。
-由于 NPU 上不支持 Megatron-LM 训练，请首先将对应文件中 `strategy_args` 参数修改为 `fsdp2` 选项。
+正式使用前，建议您通过单节点流水线训练检验环境准备和安装是否正确。FSDP2 示例可直接使用基础昇腾环境；运行采用 `megatron_train` 或 `megatron_infer` 的配置前，请先完成[在昇腾上安装 Megatron](#在昇腾上安装-megatron)。
 
 
 1. 使用 shell 执行单节点流水线：
@@ -154,7 +204,13 @@ pip install transformers==4.57.6
 bash examples/agentic_demo/run_agentic_pipeline_frozen_lake_single_node_demo.sh  
 ```
 
-2. 使用配置文件执行 agentic pipeline：
+2. 运行 Megatron DPO 示例：
+
+```
+bash examples/ascend_examples/run_dpo_pipeline.sh
+```
+
+3. 使用配置文件执行 agentic pipeline：
 
 ```
 # 确保当前位于 ROLL 项目目录的根目录下
@@ -174,6 +230,7 @@ python examples/start_agentic_pipeline.py \
 | Agentic | examples/qwen2.5-0.5B-agentic/run_agentic_pipeline_sokoban.sh | FSDP2 | vLLM | Atlas 900 A2/A3/Ascend 950 训练系列 |
 | Agentic-Rollout | examples/qwen2.5-0.5B-agentic/run_agentic_rollout_sokoban.sh | FSDP2 | vLLM | Atlas 900 A2/A3/Ascend 950 训练系列 |
 | RLVR | examples/ascend_examples/run_rlvr_pipeline.sh | FSDP2 | vLLM | Atlas 900 A2/A3/Ascend 950 训练系列 |
+| DPO | examples/ascend_examples/run_dpo_pipeline.sh | Megatron | Megatron | Atlas 900 A2/A3 训练系列 |
 
 ## 声明
 
diff --git a/examples/ascend_examples/qwen3_4B_dpo_megatron.yaml b/examples/ascend_examples/qwen3_4B_dpo_megatron.yaml
new file mode 100644
index 000000000..58f810b17
--- /dev/null
+++ b/examples/ascend_examples/qwen3_4B_dpo_megatron.yaml
@@ -0,0 +1,100 @@
+defaults:
+  - ../config/deepspeed_zero@_here_
+  - ../config/deepspeed_zero2@_here_
+  - ../config/deepspeed_zero3@_here_
+  - ../config/deepspeed_zero3_cpuoffload@_here_
+
+hydra:
+  run:
+    dir: .
+  output_subdir: null
+
+exp_name: "qwen3-4B-dpo-config"
+seed: 42
+logging_dir: ./output/logs
+output_dir: ./output
+system_envs:
+  USE_MODELSCOPE: '1'
+
+checkpoint_config:
+  type: file_system
+  output_dir: ./ckpt
+
+
+track_name: None
+
+
+max_steps: 500
+save_steps: 500
+logging_steps: 1
+eval_steps: 100
+resume_from_checkpoint: false
+
+sequence_length: 512
+train_batch_size: 64
+val_batch_size: 64
+
+# local_rank: -1
+num_nodes: 1
+num_gpus_per_node: 4
+
+pretrain: Qwen/Qwen3-4B
+
+ipo: false
+beta: 0.1
+label_smoothing: 0.0
+
+chosen_key: chosen
+rejected_key: rejected
+
+validation:
+  data_args:
+    template: qwen3
+    file_name: data/comparison_gpt4_data_zh.json
+
+actor_train:
+  model_args:
+    disable_gradient_checkpointing: false
+    dtype: bf16
+    model_type: ~
+  training_args:
+    lr_scheduler_type: constant
+    learning_rate: 1.0e-6
+    weight_decay: 0
+    per_device_train_batch_size: 16
+    gradient_accumulation_steps: 1
+    warmup_steps: 20
+    num_train_epochs: 10
+  data_args:
+    template: qwen3
+    file_name:
+      - data/comparison_gpt4_data_zh.json
+    dataset_dir: data
+    preprocessing_num_workers: 1
+  strategy_args:
+    strategy_name: megatron_train
+    strategy_config:
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      expert_model_parallel_size: 1
+      use_distributed_optimizer: true
+      recompute_granularity: full
+  device_mapping: list(range(0,2))
+  infer_batch_size: 16
+
+
+reference:
+  model_args:
+    disable_gradient_checkpointing: true
+    dtype: bf16
+    model_type: ~
+  data_args:
+    template: qwen3
+  strategy_args:
+    strategy_name: megatron_infer
+    strategy_config:
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      expert_model_parallel_size: 1
+  device_mapping: list(range(2,4))
+  infer_batch_size: 16
\ No newline at end of file
diff --git a/examples/ascend_examples/run_dpo_pipeline.sh b/examples/ascend_examples/run_dpo_pipeline.sh
new file mode 100644
index 000000000..7264e2582
--- /dev/null
+++ b/examples/ascend_examples/run_dpo_pipeline.sh
@@ -0,0 +1,10 @@
+#!/bin/bash
+set +x
+
+export HCCL_NPU_SOCKET_PORT_RANGE="auto"
+export VLLM_ASCEND_ENABLE_NZ=0
+
+CONFIG_PATH=$(basename "$(dirname "$0")")
+python examples/start_dpo_pipeline.py \
+  --config_path "$CONFIG_PATH" \
+  --config_name qwen3_4B_dpo_megatron
diff --git a/mcore_adapter/src/mcore_adapter/initialize.py b/mcore_adapter/src/mcore_adapter/initialize.py
index 7357bfa4c..8817bce45 100644
--- a/mcore_adapter/src/mcore_adapter/initialize.py
+++ b/mcore_adapter/src/mcore_adapter/initialize.py
@@ -13,6 +13,19 @@
 logger = get_logger(__name__)
 
 
+def _load_megatron_adaptor():
+    """Activate Megatron's Ascend NPU patches before initialization."""
+    if not current_platform.is_npu():
+        return
+
+    try:
+        import megatron_adaptor  # noqa: F401
+    except ImportError as exc:
+        raise RuntimeError(
+            "MegatronAdaptor is required to initialize Megatron on Ascend NPU."
+        ) from exc
+
+
 def is_distribute_initialized():
     return mpu.model_parallel_is_initialized()
 
@@ -36,6 +49,7 @@ def _set_random_seed(seed_):
 
 
 def initialize_megatron(args: "TrainingArguments"):
+    _load_megatron_adaptor()
     if not is_distribute_initialized():
         _initialize_distributed(args)
     _set_random_seed(args.seed)
diff --git a/roll/distributed/strategy/megatron_strategy.py b/roll/distributed/strategy/megatron_strategy.py
index a50d29b11..c01a56dab 100644
--- a/roll/distributed/strategy/megatron_strategy.py
+++ b/roll/distributed/strategy/megatron_strategy.py
@@ -172,7 +172,10 @@ def initialize(self, model_provider):
             # R2 mode: init router_replay_action=RouterReplayAction.RECORD
             RouterReplay.set_global_router_replay_action(RouterReplayAction.RECORD)
 
-        logger.info(f"{self.model.get_models()}")
+        if current_platform.is_npu():
+            logger.info("Initialized model chunks: %s", [type(model).__name__ for model in self.models_unwrapped])
+        else:
+            logger.info(f"{self.model.get_models()}")
         dist.barrier()
 
     def _validate_vlm_packing_support(self):
@@ -529,6 +532,13 @@ def inner_forward_step(self, loss_func, data_iterator: Iterator[DataProto], mode
         else:
             input_ids = self._get_feature_on_this_cp_rank(input_ids, "input_ids")
             attention_mask = self._get_feature_on_this_cp_rank(attention_mask, "attention_mask")
+
+            if hasattr(torch, "npu") and torch.npu.is_available() and attention_mask is not None:
+                attention_mask = attention_mask.bool()
+                B, S = attention_mask.shape
+                attention_mask = attention_mask[:, None, None, :]   # [B,1,1,S]
+                attention_mask = attention_mask.expand(B, 1, S, S)        # [B,1,S,S]
+
             if labels is not None:
                 labels = self._get_feature_on_this_cp_rank(labels, "labels")
             loss_mask = self._get_feature_on_this_cp_rank(loss_mask, "loss_mask")
@@ -1250,7 +1260,10 @@ def initialize(self, model_provider):
         # if self.enable_router_replay and self.router_replay_mode == "R3":
             # RouterReplay.set_global_router_replay_action(RouterReplayAction.REPLAY_FORWARD)
 
-        logger.info(f"{self.model.get_models()}")
+        if current_platform.is_npu():
+            logger.info("Initialized model chunks: %s", [type(model).__name__ for model in self.models_unwrapped])
+        else:
+            logger.info(f"{self.model.get_models()}")
         if self.megatron_train_args.compile_warmup and self.worker.rank_info.pp_size > 1:
             compile_warmup_pipeline_stages(self)
 
diff --git a/roll/pipeline/base_worker.py b/roll/pipeline/base_worker.py
index 3e522bea7..e7daf3c67 100644
--- a/roll/pipeline/base_worker.py
+++ b/roll/pipeline/base_worker.py
@@ -457,8 +457,7 @@ async def offload_states_partial(self, target_dp_ranks: List[int]):
 
             # Verify offloaded workers have near-zero GPU memory usage
             if self.rank_info.dp_rank in target_dp_ranks:
-                import torch
-                gpu_memory_gb = torch.cuda.memory_allocated() / 1024**3
+                gpu_memory_gb = current_platform.memory_allocated() / 1024**3
                 if gpu_memory_gb > 1.0:
                     raise RuntimeError(
                         f"GPU memory not properly offloaded for Worker {self.rank} (DP {self.rank_info.dp_rank}): "
diff --git a/roll/platforms/__init__.py b/roll/platforms/__init__.py
index 6869621f4..c9dff3f15 100644
--- a/roll/platforms/__init__.py
+++ b/roll/platforms/__init__.py
@@ -25,26 +25,31 @@ def _init_platform() -> Platform:
     Returns:
         An instance of a subclass of Platform corresponding to the detected hardware.
     """
+    try:
+        import torch_npu  # noqa: F401
+
+        if hasattr(torch, "npu") and torch.npu.is_available():
+            logger.debug("Detected NPU (torch_npu). Initializing NPU platform.")
+            return NpuPlatform()
+    except ImportError:
+        pass
+
     if torch.cuda.is_available():
         device_name = torch.cuda.get_device_name().upper()
         logger.debug(f"Detected CUDA device: {device_name}")
+
         if "NVIDIA" in device_name:
             logger.debug("Initializing CUDA platform (NVIDIA).")
             return CudaPlatform()
         elif "AMD" in device_name:
             logger.debug("Initializing ROCm platform (AMD).")
             return RocmPlatform()
+
         logger.warning("Unrecognized CUDA device. Falling back to UnknownPlatform.")
         return UnknownPlatform()
-    else:
-        try:
-            import torch_npu  # noqa: F401
 
-            logger.debug("Detected torch_npu. Initializing NPU platform.")
-            return NpuPlatform()
-        except ImportError:
-            logger.debug("No supported accelerator detected. Initializing CPU platform.")
-            return CpuPlatform()
+    logger.debug("No supported accelerator detected. Initializing CPU platform.")
+    return CpuPlatform()
 
 
 # Global singleton representing the current platform in use.
diff --git a/roll/utils/functionals.py b/roll/utils/functionals.py
index 66ebe9b34..1dd282cd7 100644
--- a/roll/utils/functionals.py
+++ b/roll/utils/functionals.py
@@ -394,8 +394,9 @@ def reduce_metrics(metrics: dict, reduce_func=np.mean) -> dict:
 
     Notes:
       - The original metric key is preserved (the '@tag' or '_suffix' remains in the key).
-      - Scalar values (int, float, np.number) and torch.Tensor objects are left unchanged.
-      - Values of type list, tuple, or np.ndarray are reduced using the inferred aggregation function.
+      - Scalar values (int, float, np.number) and standalone torch.Tensor objects are left unchanged.
+      - Values of type list, tuple, or np.ndarray are reduced using the inferred aggregation function;
+        on NPU, nested tensors are copied to CPU first so NumPy can aggregate them.
       - If no aggregation tag or suffix is found, the default `reduce_func` is used.
       - Empty sequences are skipped and not modified.
     """
@@ -429,6 +430,13 @@ def _parse_aggregation_func(metric_name: str):
         # No aggregation specifier found → use default
         return reduce_func
 
+    def _to_numpy_compatible(value: object) -> object:
+        if isinstance(value, torch.Tensor):
+            return value.detach().cpu().tolist()
+        if isinstance(value, (list, tuple)):
+            return [_to_numpy_compatible(item) for item in value]
+        return value
+
     for key, val in list(metrics.items()):
         # Skip reduction for scalars and tensors
         if isinstance(val, (int, float, np.number)) or isinstance(val, torch.Tensor):
@@ -439,6 +447,8 @@ def _parse_aggregation_func(metric_name: str):
             if len(val) == 0:
                 continue
             agg_func = _parse_aggregation_func(key)
+            if current_platform.is_npu():
+                val = _to_numpy_compatible(val)
             metrics[key] = float(agg_func(val))
         else:
             # Fallback for other types (e.g., single-element containers)
diff --git a/tests/distributed/strategy/test_vllm_strategy_beam_search.py b/tests/distributed/strategy/test_vllm_strategy_beam_search.py
index d172f31bb..0cce3d7f3 100644
--- a/tests/distributed/strategy/test_vllm_strategy_beam_search.py
+++ b/tests/distributed/strategy/test_vllm_strategy_beam_search.py
@@ -75,6 +75,7 @@ def _install_mock_vllm_modules(monkeypatch):
 
     inputs = ModuleType("vllm.inputs")
     inputs.__path__ = []
+    inputs.TokensPrompt = MockTokensPrompt
     inputs_data = ModuleType("vllm.inputs.data")
     inputs_data.TokensPrompt = MockTokensPrompt
 
diff --git a/tests/third_party/megatron/test_offload_states.py b/tests/third_party/megatron/test_offload_states.py
index cb6416ed9..352c3c00b 100644
--- a/tests/third_party/megatron/test_offload_states.py
+++ b/tests/third_party/megatron/test_offload_states.py
@@ -3,7 +3,6 @@
 
 import pytest
 import torch
-import torch.distributed as dist
 
 from roll.platforms import current_platform
 from megatron.core import DistributedDataParallel
@@ -33,11 +32,18 @@
 )
 from roll.third_party.megatron.optimizer import get_megatron_optimizer
 
+def _default_model_name():
+    local_model = "/data/cpfs_0/common/models/Qwen2.5-0.5B-Instruct"
+    return os.environ.get(
+        "ROLL_MEGATRON_TEST_MODEL",
+        local_model if os.path.isdir(local_model) else "Qwen/Qwen2.5-0.5B-Instruct",
+    )
+
 
 class McaModelCreator:
 
-    def __init__(self, optimizer_type, model_name="/data/cpfs_0/common/models/Qwen2.5-0.5B-Instruct"):
-        self.model_name = model_name
+    def __init__(self, optimizer_type, model_name=None):
+        self.model_name = model_name or _default_model_name()
         if optimizer_type is None:
             self.megatron_train_args = TrainingArguments(
                 output_dir="./output",
@@ -128,7 +134,9 @@ def create_mca_infer_only(self):
 
         self.tokenizer = default_tokenizer_provider(model_args=self.model_args)
         self.model = default_actor_model_provider(
-            tokenizer=self.tokenizer, training_args=self.megatron_train_args, model_args=self.model_args
+            tokenizer=self.tokenizer,
+            training_args=self.megatron_train_args,
+            model_args=self.model_args,
         )
         for module in self.model.get_models():
             module.requires_grad_(False)
@@ -141,8 +149,14 @@ def create_mca_model(self):
 
         self.tokenizer = default_tokenizer_provider(model_args=self.model_args)
         self.model = default_actor_model_provider(
-            tokenizer=self.tokenizer, training_args=self.megatron_train_args, model_args=self.model_args
+            tokenizer=self.tokenizer,
+            training_args=self.megatron_train_args,
+            model_args=self.model_args,
+            is_trainable=True,
         )
+        for module in self.model.get_models():
+            module.train()
+            module.requires_grad_(True)
 
         ddp_config = DistributedDataParallelConfig(
             grad_reduce_in_fp32=self.megatron_train_args.accumulate_allreduce_grads_in_fp32,
@@ -188,26 +202,33 @@ def create_mca_model(self):
         bind_megatron_offload_states_func(optimizer=self.optimizer)
 
         if not isinstance(self.optimizer, ChainedOptimizer):
-            self.scheduler = get_scheduler(
-                "cosine",
-                optimizer=self.optimizer if self.optimizer is None else self.optimizer.optimizer,
-                num_warmup_steps=self.megatron_train_args.get_warmup_steps(self.megatron_train_args.max_steps),
-                num_training_steps=self.megatron_train_args.max_steps,
-                scheduler_specific_kwargs=self.megatron_train_args.lr_scheduler_kwargs,
-            )
-        else:
-            lr_schedulers = []
-            for opt in self.optimizer.chained_optimizers:
-                sch = get_scheduler(
+            base_optimizer = self.optimizer if self.optimizer is None else self.optimizer.optimizer
+            if base_optimizer is not None:
+                self.scheduler = get_scheduler(
                     "cosine",
-                    optimizer=opt if opt is None else opt.optimizer,
+                    optimizer=base_optimizer,
                     num_warmup_steps=self.megatron_train_args.get_warmup_steps(self.megatron_train_args.max_steps),
                     num_training_steps=self.megatron_train_args.max_steps,
                     scheduler_specific_kwargs=self.megatron_train_args.lr_scheduler_kwargs,
                 )
-                lr_schedulers.append(sch)
+        else:
+            lr_schedulers = []
+            for opt in self.optimizer.chained_optimizers:
+                base_optimizer = opt if opt is None else opt.optimizer
+                if base_optimizer is None:
+                    continue
+                lr_schedulers.append(
+                    get_scheduler(
+                        "cosine",
+                        optimizer=base_optimizer,
+                        num_warmup_steps=self.megatron_train_args.get_warmup_steps(self.megatron_train_args.max_steps),
+                        num_training_steps=self.megatron_train_args.max_steps,
+                        scheduler_specific_kwargs=self.megatron_train_args.lr_scheduler_kwargs,
+                    )
+                )
 
-            self.scheduler = ChainedScheduler(lr_schedulers)
+            if lr_schedulers:
+                self.scheduler = ChainedScheduler(lr_schedulers)
 
 
 """
@@ -216,64 +237,24 @@ def create_mca_model(self):
 """
 
 
-def test_megatron_init_memory():
-    MAX_NUM_OF_MEM_EVENTS_PER_SNAPSHOT: int = 100000
-    torch.cuda.memory._record_memory_history(
-        max_entries=MAX_NUM_OF_MEM_EVENTS_PER_SNAPSHOT,
-    )
-
+def test_megatron_dist_optimizer_offload_reload_smoke():
     mca_model = McaModelCreator(optimizer_type="dist_optimizer")
-
-    # buffer_data = []
-    # for buffer in mca_model.optimizer.buffers:
-    #     buffer_data.append(buffer.param_data.data.storage().data_ptr())
-
-    mca_model.optimizer.offload_states(include=[MegatronOffloadStateType.other_params], pin_memory=True)
-
-    t0 = torch.randint(0, 100, (1024, 1024, 1024), device="cuda")
-    del t0
-
-    mca_model.optimizer.reload_states(include=[MegatronOffloadStateType.model_params])
-    if dist.get_rank() == 0:
-        t0 = torch.randint(0, 100, (1024, 1024, 1024), device="cuda")
-        dump_file = f"./memory_dump/snapshot_megatron_init_offload_{os.environ['RANK']}.pickle"
-        os.makedirs(os.path.dirname(dump_file), exist_ok=True)
-        torch.cuda.memory._dump_snapshot(dump_file)
-        torch.cuda.memory._record_memory_history(enabled=None)
-
-    # tensors_group_by_data_ptr = defaultdict(list)
-    # tensors = objgraph.by_type('Tensor')
-    # print(f"len(tensor)={len(tensors)}")
-    # for tensor in tensors:
-    #     tensors_group_by_data_ptr[tensor.storage().data_ptr()].append(tensor)
-    #
-    # for buffer in buffer_data:
-    #     objgraph.show_backrefs(tensors_group_by_data_ptr[buffer], max_depth=10,
-    #                            extra_ignore=[id(locals())],
-    #                            filename=f'/checkpoint/binary/ScaleAligner/memory_dump/buffer_data_group_tensors.param_data_{datetime.now().strftime("%Y%m%d-%H%M%S")}.png')
-
-
-def test_megatron_init_ddp_memory():
-    MAX_NUM_OF_MEM_EVENTS_PER_SNAPSHOT: int = 100000
-    torch.cuda.memory._record_memory_history(
-        max_entries=MAX_NUM_OF_MEM_EVENTS_PER_SNAPSHOT,
+    run_model_dist_optimizer(
+        mca_model,
+        included_state=[MegatronOffloadStateType.other_params],
+        pin_memory=True,
+        non_blocking=True,
     )
 
-    mca_model = McaModelCreator(optimizer_type=None)
-
-    offload_megatron_no_grad_module(model_chunks=mca_model.model.get_models())
 
-    t0 = torch.randint(0, 100, (1024, 1024, 1024), device="cuda")
-    del t0
-
-    reload_megatron_no_grad_module(model_chunks=mca_model.model.get_models())
-
-    if dist.get_rank() == 0:
-        t0 = torch.randint(0, 100, (1024, 1024, 1024), device="cuda")
-        dump_file = f"./memory_dump/snapshot_megatron_init_ddp_offload_{os.environ['RANK']}.pickle"
-        os.makedirs(os.path.dirname(dump_file), exist_ok=True)
-        torch.cuda.memory._dump_snapshot(dump_file)
-        torch.cuda.memory._record_memory_history(enabled=None)
+def test_megatron_no_grad_module_offload_reload_smoke():
+    mca_model = McaModelCreator(optimizer_type=None)
+    run_model_infer(
+        mca_model,
+        included_state=[MegatronOffloadStateType.model_params],
+        pin_memory=True,
+        non_blocking=True,
+    )
 
 
 def check_devices(tensors: List[torch.Tensor], target_device) -> None:
@@ -284,16 +265,32 @@ def check_devices(tensors: List[torch.Tensor], target_device) -> None:
 
 def check_tensors(expected_tensors: List[torch.Tensor], tensors: List[torch.Tensor]) -> None:
     for tensor_expected, tensor_restored in zip(expected_tensors, tensors):
-        assert torch.equal(tensor_expected, tensor_restored)
+        assert torch.equal(tensor_expected.to(tensor_restored.device), tensor_restored)
+
+
+def current_device():
+    return f"{current_platform.device_type}:{current_platform.current_device()}"
+
+
+def prepare_model_batch(batch):
+    input_ids, token_attention_mask = batch
+    input_ids = input_ids.to(current_device())
+    token_attention_mask = token_attention_mask.to(current_device())
+    position_ids = torch.clip(torch.cumsum(token_attention_mask, dim=-1) - 1, min=0, max=None)
+    seq_length = input_ids.size(1)
+    causal_mask = torch.triu(
+        torch.ones((1, 1, seq_length, seq_length), dtype=torch.bool, device=input_ids.device),
+        diagonal=1,
+    )
+    padding_mask = token_attention_mask[:, None, None, :].eq(0)
+    attention_mask = causal_mask | padding_mask
+    return input_ids, attention_mask, position_ids
 
 
 def run_model_infer(mca_model: McaModelCreator, included_state, pin_memory, non_blocking):
     with torch.no_grad():
         for batch in mca_model.data_loader:
-            input_ids, attention_mask = batch
-            input_ids = input_ids.to("cuda")
-            attention_mask = attention_mask.to("cuda")
-            position_ids = torch.clip(torch.cumsum(attention_mask, dim=-1) - 1, min=0, max=None)
+            input_ids, attention_mask, position_ids = prepare_model_batch(batch)
 
             models = mca_model.model.get_models()
             for model in models:
@@ -329,10 +326,7 @@ def run_model_dist_optimizer(mca_model: McaModelCreator, included_state, pin_mem
     assert isinstance(mca_model.optimizer, DistributedOptimizer)
 
     for batch in mca_model.data_loader:
-        input_ids, attention_mask = batch
-        input_ids = input_ids.to("cuda")
-        attention_mask = attention_mask.to("cuda")
-        position_ids = torch.clip(torch.cumsum(attention_mask, dim=-1) - 1, min=0, max=None)
+        input_ids, attention_mask, position_ids = prepare_model_batch(batch)
 
         models = mca_model.model.get_models()
         for model in models:
@@ -534,10 +528,7 @@ def run_model_fp16_optimizer(mca_model: McaModelCreator, included_state, pin_mem
     assert isinstance(mca_model.optimizer, Float16OptimizerWithFloat16Params)
 
     for batch in mca_model.data_loader:
-        input_ids, attention_mask = batch
-        input_ids = input_ids.to("cuda")
-        attention_mask = attention_mask.to("cuda")
-        position_ids = torch.clip(torch.cumsum(attention_mask, dim=-1) - 1, min=0, max=None)
+        input_ids, attention_mask, position_ids = prepare_model_batch(batch)
 
         models = mca_model.model.get_models()
         for model in models:
@@ -710,10 +701,7 @@ def run_model_fp32_optimizer(mca_model: McaModelCreator, included_state, pin_mem
     assert isinstance(mca_model.optimizer, FP32Optimizer)
 
     for batch in mca_model.data_loader:
-        input_ids, attention_mask = batch
-        input_ids = input_ids.to("cuda")
-        attention_mask = attention_mask.to("cuda")
-        position_ids = torch.clip(torch.cumsum(attention_mask, dim=-1) - 1, min=0, max=None)
+        input_ids, attention_mask, position_ids = prepare_model_batch(batch)
 
         models = mca_model.model.get_models()
         for model in models: