feat: support GLM-4.7-Flash (MLA + MoE) on Ascend NPU

.buildkite/npu_suites.py

@@ -32,6 +32,7 @@
         ("test_qwen3_30B_A3B_npu.py", "npu-16", "", {}),
         ("test_qwen3_vl_8B_npu.py", "npu-8", "", {}),
         ("test_qwen3.5_35B_A3B_npu.py", "npu-16", "", {}),
+        ("test_glm4.7_30B_A3B_npu.py", "npu-16", "", {}),
     ],
     "nightly": [],
 }

docker/npu_patch/vllm.patch

@@ -16,11 +16,27 @@ index cb61bca..5c076d5 100644
 --- a/vllm/v1/core/sched/async_scheduler.py
 +++ b/vllm/v1/core/sched/async_scheduler.py
 @@ -51,7 +51,7 @@ class AsyncScheduler(Scheduler):
- 
+
          # Update the number of output placeholders.
          request.num_output_placeholders -= len(new_token_ids)
 -        assert request.num_output_placeholders >= 0
 +        request.num_output_placeholders = max(0, request.num_output_placeholders)
- 
+
          # Cache the new tokens. Preempted requests should be skipped.
          if status_before_update == RequestStatus.RUNNING:
+diff --git a/vllm/model_executor/models/glm4_moe_lite_mtp.py b/vllm/model_executor/models/glm4_moe_lite_mtp.py
+index 596cb48..8c30495 100644
+--- a/vllm/model_executor/models/glm4_moe_lite_mtp.py
++++ b/vllm/model_executor/models/glm4_moe_lite_mtp.py
+@@ -126,7 +126,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
++        # torch.where (element-wise) replaces bool-mask index to stay
++        # cudagraph-capturable on NPU (aclnnNonzeroV2 fails under 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)
+

scripts/models/glm4.7-30B-A3B-npu.sh

@@ -0,0 +1,51 @@
+MOE_ROUTED_EXPERTS=64
+MOE_ACTIVE_ROUTED_EXPERTS=4
+MOE_SHARED_EXPERTS=1
+
+NHIDDEN=2048
+MOE_FFN_HIDDEN=1536
+MOE_SHARED_EXPERT_INTERMEDIATE_SIZE=$((MOE_FFN_HIDDEN * MOE_SHARED_EXPERTS))
+FFN_HIDDEN=10240
+N_DENSE_LAYERS=1
+N_MOE_LAYERS=46
+NHEADS=20
+
+MODEL_ARGS=(
+    --moe-layer-freq "[0]*$N_DENSE_LAYERS+[1]*$N_MOE_LAYERS"
+    --num-experts $MOE_ROUTED_EXPERTS
+    --moe-shared-expert-intermediate-size $MOE_SHARED_EXPERT_INTERMEDIATE_SIZE
+    --moe-router-topk $MOE_ACTIVE_ROUTED_EXPERTS
+    --moe-grouped-gemm
+    --moe-ffn-hidden-size $MOE_FFN_HIDDEN
+    --moe-router-score-function sigmoid
+    --moe-router-pre-softmax
+    --moe-router-enable-expert-bias
+    --moe-router-bias-update-rate 0
+    --moe-router-load-balancing-type seq_aux_loss
+    --moe-router-topk-scaling-factor 1.8
+    --moe-aux-loss-coeff 0
+    --moe-router-dtype fp32
+    --make-vocab-size-divisible-by 64
+    --num-layers $((N_DENSE_LAYERS + N_MOE_LAYERS))
+    --hidden-size $NHIDDEN
+    --ffn-hidden-size $FFN_HIDDEN
+    --num-attention-heads $NHEADS
+    --disable-bias-linear
+    --add-qkv-bias
+    --swiglu
+    --untie-embeddings-and-output-weights
+    --position-embedding-type rope
+    --no-position-embedding
+    --normalization RMSNorm
+    --qk-layernorm
+    --multi-latent-attention
+    --q-lora-rank 768
+    --kv-lora-rank 512
+    --qk-head-dim 192
+    --v-head-dim 256
+    --kv-channels 192
+    --qk-pos-emb-head-dim 64
+    --vocab-size 154880
+    --rotary-base 1000000
+    --no-rope-fusion
+)

scripts/run-glm4.7-30B-A3B-npu.sh

@@ -0,0 +1,149 @@
+#!/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/glm4.7-30B-A3B-npu.sh"
+
+DATA_ROOT="${DATA_ROOT:-/root}"
+
+CKPT_ARGS=(
+   --hf-checkpoint ${DATA_ROOT}/weights/GLM-4.7-Flash/
+   --load ${DATA_ROOT}/weights/GLM-4.7-Flash/
+   --ref-load ${DATA_ROOT}/weights/GLM-4.7-Flash/
+   --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 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
+)
+
+MTP_ARGS=(
+   --mtp-num-layers 1
+   --enable-mtp-training
+   --mtp-loss-scaling-factor 0.2
+)
+
+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)
+   --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":1}'
+)
+
+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
+)
+
+# 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[@]}" \
+   "${MTP_ARGS[@]}"

tests/test_glm4.7_30B_A3B_npu.py

@@ -0,0 +1,155 @@
+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():
+    model_dir = shlex.quote(MODEL_DIR)
+    prompt_data = shlex.quote(f"{DATASET_DIR}/dapo-math-17k.jsonl")
+
+    # NPU skips torch_dist conversion; HF weights load directly via bridge mode.
+    checkpoint_args = (
+        f"--hf-checkpoint {model_dir} "
+        f"--load {model_dir} "
+        f"--ref-load {model_dir} "
+        "--megatron-to-hf-mode bridge "
+        "--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 2048 "
+        "--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 "
+    )
+
+    # MTP speculative decoding under cudagraph - exercises the GLM-4.7 MTP
+    # drafter's graph-friendly forward (patched via docker/npu_patch/vllm.patch).
+    mtp_args = "--mtp-num-layers 1 " "--enable-mtp-training " "--mtp-loss-scaling-factor 0.2 "
+
+    vllm_args = (
+        "--rollout-num-gpus-per-engine 4 "
+        "--vllm-gpu-memory-utilization 0.7 "
+        "--vllm-enable-expert-parallel "
+        "--vllm-cudagraph-capture-sizes 1 2 4 8 "
+        '--vllm-speculative-config \'{"method":"mtp","num_speculative_tokens":1}\' '
+    )
+
+    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 "
+    )
+
+    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()

vime_plugins/megatron_bridge/__init__.py

@@ -1 +1,2 @@
+import vime_plugins.megatron_bridge.glm4moe_lite  # noqa: F401  # register GLM-4.7-Flash bridge
 import vime_plugins.megatron_bridge.glm4v_moe  # noqa: F401  # register GLM-4.6V bridge