[WIP]NPU optimizations for Qwen3.5 9B & 35B

docker/Dockerfile.npu

@@ -12,6 +12,7 @@ ARG NO_PROXY
 ARG http_proxy
 ARG https_proxy
 ARG no_proxy
+ARG SOC_VERSION_FOR_SGL_BUILD="Ascend910_9392"
 
 ENV http_proxy=${HTTP_PROXY} \
     https_proxy=${HTTPS_PROXY} \
@@ -136,39 +137,34 @@ RUN pip_install() { pip install "$@" || pip install --index-url https://pypi.org
     git add . && git commit -m "install info" && \
     git fetch && \
     git cherry-pick ece02ffc9cc32e94382d4f1b553b2c755f83f722 && \
-    patch -p1 < /root/Relax/docker/npu_patch/sglang-npu.patch && \
-    git add . && git commit -m "sglang-npu.patch"
+    git am /root/Relax/docker/npu_patch/sglang-npu.patch
 
 # Install SGLang NPU kernels.
 RUN git clone https://github.com/sgl-project/sgl-kernel-npu /root/sgl-kernel-npu && \
     cd /root/sgl-kernel-npu && git checkout 2026.7.2 && \
-    patch -p1 < /root/Relax/docker/npu_patch/sgl-kernel-npu.patch && \
-    git add . && git commit -m "sgl-kernel-npu.patch" && \
-    bash build.sh && \
+    git submodule update --init third_party/catlass third_party/pto-isa && \
+    git am /root/Relax/docker/npu_patch/sgl-kernel-npu.patch && \
+    bash build.sh -a kernels ${SOC_VERSION_FOR_SGL_BUILD} && bash build.sh -a memory-saver && bash build.sh -a deepep && \
     pip install output/*.whl && \
     cd /root && rm -rf /root/sgl-kernel-npu
 
 # Install AscendC FLA
 RUN git clone https://github.com/flashserve/flash-linear-attention-npu.git /root/flash-linear-attention-npu && \
-    cd /root/flash-linear-attention-npu && git checkout v26.1.0 && \
+    cd /root/flash-linear-attention-npu && git checkout v26.6.0 && \
     apt update && apt-get install -y --no-install-recommends gawk && \
-    # 编译命令，注意--soc=${soc_version}需要指定为当前机器的芯片类型{ascend910b/ascend910_93/ascend950}
-    bash build.sh --soc=ascend910_93 --pkg --ops=causal_conv1d,chunk_bwd_dv_local,chunk_bwd_dqkwg,chunk_gated_delta_rule_bwd_dhu,prepare_wy_repr_bwd_da,prepare_wy_repr_bwd_full,chunk_fwd_o,chunk_gated_delta_rule_fwd_h,recurrent_gated_delta_rule,recompute_wu_fwd && \
-    # 安装run包
-    ./build_out/cann-*.run && \
     source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
     source /usr/local/Ascend/nnal/atb/set_env.sh && \
     source /usr/local/Ascend/cann-9.0.0/share/info/ascendnpu-ir/bin/set_env.sh && \
-    # 一键编译安装脚本，先调用torchnpugen自动接入算子，再运行setup编whl包，最后安装whl包
-    cd torch_custom/fla_npu && bash build.sh
+    FLA_NPU_SOC=ascend910_93 python -m pip wheel --no-build-isolation --no-deps . -w dist && \
+    python -m pip install --force-reinstall --no-deps dist/flash_linear_attention_npu-*.whl
 
 FROM train as relax
 
 WORKDIR /root
 
 COPY requirements.txt /tmp/requirements.txt
 RUN pip install -r /tmp/requirements.txt --no-cache-dir && \
-    pip install --index-url https://pypi.org/simple ray==2.55.1 protobuf==6.33.6
+    pip install --index-url https://pypi.org/simple ray==2.55.1 protobuf==6.33.6 fastokens==0.3.1
 
 # Clean caches.
 RUN pip cache purge && \

relax/utils/reloadable_process_group.py

@@ -338,6 +338,8 @@ def destroy_process_groups(post_destroy_delay: float = 2.0):
     """Destroy all reloadable process groups."""
     if _should_skip_reload_and_destroy():
         return
+    if device_utils.is_npu_available:
+        post_destroy_delay = 0.0
     ReloadableProcessGroup.destroy_process_groups(post_destroy_delay=post_destroy_delay)
 
 

scripts/entrypoint/local-npu-multinode.sh

@@ -0,0 +1,145 @@
+if [ -n "${RELAX_ENTRYPOINT_MODE:-}" ]; then
+    return 0 2>/dev/null || exit 0
+fi
+
+_LOCAL_SH_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)"
+
+# ── delegate to ray-job.sh when inside an existing Ray cluster ─────────────
+# When RAY_ADDRESS is set AND `ray status` succeeds, we're already part of an
+# externally-managed Ray cluster. Skip local Ray startup / process cleanup and
+# fall through to ray-job.sh (source mode) for env setup.
+if [ -n "${RAY_ADDRESS:-}" ] && timeout 5 ray status >/dev/null 2>&1; then
+    echo "=== Detected existing Ray cluster (RAY_ADDRESS=$RAY_ADDRESS); delegating to ray-job.sh ==="
+    source "${_LOCAL_SH_DIR}/ray-job-npu.sh"
+    return 0 2>/dev/null || exit 0
+fi
+
+set -eo pipefail
+
+# ── process cleanup ─────────────────────────────────────────────────────────
+echo "=== Cleaning up stale processes ==="
+pkill -9 sglang 2>/dev/null || true
+sleep 3
+ray stop --force 2>/dev/null || true
+pkill -9 ray 2>/dev/null || true
+pkill -9 python 2>/dev/null || true
+sleep 3
+pkill -9 ray 2>/dev/null || true
+pkill -9 python 2>/dev/null || true
+
+set -x
+
+# ── environment setup ───────────────────────────────────────────────────────
+export PYTHONUNBUFFERED=1
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export MEGATRON=${MEGATRON:-/root/Megatron-LM/}
+export MEGATRON_BRIDGE_SRC=${MEGATRON_BRIDGE_SRC:-/root/Megatron-Bridge/src/}
+export MINDSPEED=${MINDSPEED:-/root/MindSpeed/}
+export RELAX=${RELAX:-${_LOCAL_SH_DIR}/../../}
+export PYTHONPATH=${RELAX}:${MEGATRON_BRIDGE_SRC}:${MINDSPEED}:$MEGATRON:$RELAX:${PYTHONPATH:-}
+export MODEL_CONFIG_DIR="${_LOCAL_SH_DIR}/../models"
+
+# ── Ray cluster startup (multi node) ──────────────────────────────────────
+export MASTER_ADDR_IP=$(ping -c 1 $MASTER_ADDR | head -n 1 | awk -F'[()]' '{print $2}')
+# ── multi-node parameters ──────────────────────────────────────────────────
+NUM_NPUS="${NUM_NPUS:-16}"
+NNODES="${WORLD_SIZE:-2}"
+
+if [ "$MASTER_ADDR" = "$POD_NAME" ]; then
+    # ── HEAD NODE ───────────────────────────────────────────────────────────
+    echo "=== Head node: starting Ray cluster ==="
+    ray start --head \
+        --node-ip-address "${HOST_IP}" \
+        --resources="{\"NPU\": ${NUM_NPUS}}" \
+        --disable-usage-stats \
+        --dashboard-host=0.0.0.0 \
+        --dashboard-port=8265
+
+    sleep 5
+
+    # Wait for all worker nodes to join
+    while true; do
+        ray_status_output=$(ray status)
+        npu_count=$(echo "$ray_status_output" | grep -oP '(?<=/)\d+\.\d+(?=\s*NPU)' | head -n 1)
+        echo "Current NPU count: $npu_count"
+        npu_count_int=$(echo "$npu_count" | awk '{print int($1)}')
+        device_count=$((npu_count_int / ${NUM_NPUS}))
+
+        if [ "$device_count" -eq "$NNODES" ]; then
+            echo "Ray cluster is ready with $device_count devices (from $npu_count NPU resources)."
+            ray status
+            break
+        else
+            echo "Waiting for Ray to allocate $NNODES devices. Current device count: $device_count"
+            sleep 5
+        fi
+    done
+
+    # ── set entrypoint mode ────────────────────────────────────────────────────
+    export RELAX_ENTRYPOINT_MODE="npu-multinode"
+
+    # Runtime env for multi-node 
+    export RUNTIME_ENV_JSON="{
+    \"env_vars\": {
+        \"PYTHONUNBUFFERED\": \"1\",
+        \"PYTHONPATH\": \"${PYTHONPATH}\",
+        \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\",
+        \"RAY_OVERRIDE_JOB_RUNTIME_ENV\": \"1\",
+        \"MASTER_ADDR\": \"${HOST_IP}\",
+        \"RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES\": \"1\"
+    }
+    }"
+
+    echo "=== Head node environment ready ==="
+else
+    # ── WORKER NODE ─────────────────────────────────────────────────────────
+    # NOTE: `set -e` is active, so each retry loop below must keep the
+    # potentially-failing command in a condition position (if/until/||),
+    # otherwise the first failure kills the script and there is no retry.
+    GCS_PORT="${GCS_PORT:-6379}"
+    echo "=== Worker node: waiting for head GCS at ${MASTER_ADDR_IP}:${GCS_PORT} ==="
+    for i in $(seq 1 120); do
+        if timeout 2 bash -c "</dev/tcp/${MASTER_ADDR_IP}/${GCS_PORT}" 2>/dev/null; then
+            echo "Head GCS reachable after ${i} attempt(s)"
+            break
+        fi
+        if [ "$i" -eq 120 ]; then
+            echo "ERROR: head GCS at ${MASTER_ADDR_IP}:${GCS_PORT} unreachable after 10min" >&2
+            exit 1
+        fi
+        sleep 5
+    done
+
+    echo "=== Worker node: joining Ray cluster at ${MASTER_ADDR_IP}:${GCS_PORT} ==="
+    joined=0
+    for i in $(seq 1 30); do
+        ray stop --force >/dev/null 2>&1 || true
+        if ray start \
+            --address="${MASTER_ADDR_IP}:${GCS_PORT}" \
+            --resources="{\"NPU\": ${NUM_NPUS}}" \
+            --node-ip-address "${HOST_IP}" \
+            --disable-usage-stats \
+            --dashboard-host=0.0.0.0 \
+            --dashboard-port=8265; then
+            echo "Joined Ray cluster on attempt ${i}"
+            joined=1
+            break
+        fi
+        echo "ray start failed on attempt ${i}, retrying in 5s..."
+        sleep 5
+    done
+    if [ "$joined" -ne 1 ]; then
+        echo "ERROR: worker failed to join Ray cluster after 30 attempts" >&2
+        exit 1
+    fi
+
+    if ! ray status >/dev/null 2>&1; then
+        echo "ERROR: ray status failed after join" >&2
+        exit 1
+    fi
+    echo "Successfully connected to the Ray cluster!"
+
+    # Worker nodes block indefinitely (training runs on head node)
+    echo "=== Worker node ready, waiting for training to complete ==="
+    sleep inf
+fi