From cc7cd48aba821d6c035d25ac730a18d5b6794062 Mon Sep 17 00:00:00 2001 From: louyujing Date: Tue, 8 Sep 2026 09:27:06 +0800 Subject: [PATCH 1/5] feat: support GLM-4.7-Flash (MLA + MoE) on Ascend NPU Add the GLM47MTPBridge megatron.bridge plugin for GLM-4.7-Flash (Glm4MoeLiteForCausalLM): MLA (Multi-Head Latent Attention) + GLM-style MoE with MTP support on Ascend 910B NPU. The bridge overrides provider_bridge/mapping_registry with the MLA-aware versions (always individual MLA projections + q_a_layernorm, never fused QKV) while reusing GLM45Bridge's MTP plumbing and fused-expert handling. A runtime monkey-patch (_patch_glm_mtp_graph_friendly, env-gated by VIME_PATCH_GLM_MTP_GRAPH=1) rewrites the MTP drafter's bool-mask index assignment as torch.where so it is NPU cudagraph-capturable (aclnnNonzeroV2 otherwise fails under graph capture). Register a new end-to-end GLM NPU smoke test (tests/test_glm4.7_30B_A3B_npu.py) in the Buildkite NPU CI smk suite to validate the MLA + MoE + MTP integration and prevent regressions. The test loads HF weights via bridge mode, mirrors scripts/run-glm4.7-30B-A3B-npu.sh (TP=4/EP=8, MTP speculative decoding under cudagraph), and exercises the GLM MTP graph patch. Signed-off-by: louyujing --- .buildkite/npu_suites.py | 1 + scripts/models/glm4.7-30B-A3B-npu.sh | 51 +++ scripts/run-glm4.7-30B-A3B-npu.sh | 155 ++++++++ tests/test_glm4.7_30B_A3B_npu.py | 159 ++++++++ .../update_weight_from_tensor.py | 51 +++ vime_plugins/megatron_bridge/__init__.py | 1 + vime_plugins/megatron_bridge/glm4moe_lite.py | 369 ++++++++++++++++++ 7 files changed, 787 insertions(+) create mode 100644 scripts/models/glm4.7-30B-A3B-npu.sh create mode 100644 scripts/run-glm4.7-30B-A3B-npu.sh create mode 100644 tests/test_glm4.7_30B_A3B_npu.py create mode 100644 vime_plugins/megatron_bridge/glm4moe_lite.py diff --git a/.buildkite/npu_suites.py b/.buildkite/npu_suites.py index 647a7a1f7..7bad9d2ce 100644 --- a/.buildkite/npu_suites.py +++ b/.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": [], } diff --git a/scripts/models/glm4.7-30B-A3B-npu.sh b/scripts/models/glm4.7-30B-A3B-npu.sh new file mode 100644 index 000000000..9cc66c850 --- /dev/null +++ b/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 +) 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..dde5a3c3a --- /dev/null +++ b/scripts/run-glm4.7-30B-A3B-npu.sh @@ -0,0 +1,155 @@ +#!/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:-}" + +# Apply the vLLMWorkerExtension monkey-patch that makes the GLM-4.7 MTP drafter +# forward cudagraph-friendly on NPU (torch.where instead of bool-mask index, +# which otherwise fails aclnnNonzeroV2 under cudagraph capture). +# Drop this once vLLM upstream fixes the graph-safety issue. +export VIME_PATCH_GLM_MTP_GRAPH=1 + +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[@]}" 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..47091bca9 --- /dev/null +++ b/tests/test_glm4.7_30B_A3B_npu.py @@ -0,0 +1,159 @@ +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 + # VIME_PATCH_GLM_MTP_GRAPH monkey-patch in update_weight_from_tensor.py. + 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 " + "--colocate " + ) + + 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", + # Apply the GLM-4.7 MTP cudagraph-friendly monkey-patch (see + # update_weight_from_tensor._patch_glm_mtp_graph_friendly). + "VIME_PATCH_GLM_MTP_GRAPH": "1", + }, + ) + + +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/vime/backends/megatron_utils/update_weight/update_weight_from_tensor.py b/vime/backends/megatron_utils/update_weight/update_weight_from_tensor.py index c1052e4ea..a5b85ddef 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 @@ -416,6 +416,55 @@ def _npu_rotary_emb_init( ApplyRotaryEmb.__init__ = _npu_rotary_emb_init # type: ignore[attr-defined] ApplyRotaryEmb._npu_rotary_patched = True + @staticmethod + def _patch_glm_mtp_graph_friendly() -> None: + """Patch GLM-4.7 MTP draft forward to be NPU cudagraph-friendly. + + Upstream ``Glm4MoeLiteMultiTokenPredictorLayer.forward`` does + ``inputs_embeds[positions == 0] = 0`` (bool-mask index assignment), + which routes to aclnnNonzeroV2 — a data-dependent-shape op that fails + under NPU cudagraph capture (``stream is captured``). Replace with the + element-wise, graph-friendly ``torch.where`` equivalent so the MTP + drafter can run under cudagraph on Ascend. + + Env-gated (``VIME_PATCH_GLM_MTP_GRAPH=1``), idempotent, and a silent + no-op when the target class is absent (non-GLM models). + """ + if not os.environ.get("VIME_PATCH_GLM_MTP_GRAPH"): + return + + try: + from vllm.model_executor.models.glm4_moe_lite_mtp import Glm4MoeLiteMultiTokenPredictorLayer + except ImportError: + return + + if getattr(Glm4MoeLiteMultiTokenPredictorLayer, "_glm_mtp_graph_patched", False): + return + Glm4MoeLiteMultiTokenPredictorLayer._glm_mtp_graph_patched = True + + def _patched_forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + spec_step_index: int = 0, + ) -> torch.Tensor: + assert inputs_embeds is not None + # masking inputs at position 0, as not needed by MTP. + # torch.where (element-wise) replaces bool-mask index to stay + # cudagraph-capturable on NPU. + 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) + hidden_states = self.eh_proj(torch.cat([inputs_embeds, previous_hidden_states], dim=-1)) + hidden_states, residual = self.mtp_block(positions=positions, hidden_states=hidden_states, residual=None) + hidden_states = residual + hidden_states + return hidden_states + + Glm4MoeLiteMultiTokenPredictorLayer.forward = _patched_forward # type: ignore[attr-defined] + class vLLMColocateWorkerExtension: """vLLM ``--worker-extension-cls`` entry for colocated rollout workers.""" @@ -425,6 +474,7 @@ def __new__(cls, **kwargs): _VLLMHijack._patch_a3_moe_alltoall_expert_ids() _VLLMHijack._patch_npu_worker() _VLLMHijack._patch_npu_rotary_emb() + _VLLMHijack._patch_glm_mtp_graph_friendly() return super().__new__(cls) @@ -435,4 +485,5 @@ def __new__(cls, **kwargs): if is_npu(): _VLLMHijack._patch_npu_worker() _VLLMHijack._patch_npu_rotary_emb() + _VLLMHijack._patch_glm_mtp_graph_friendly() return super().__new__(cls) diff --git a/vime_plugins/megatron_bridge/__init__.py b/vime_plugins/megatron_bridge/__init__.py index d16a7eaf6..b097a794d 100644 --- a/vime_plugins/megatron_bridge/__init__.py +++ b/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 diff --git a/vime_plugins/megatron_bridge/glm4moe_lite.py b/vime_plugins/megatron_bridge/glm4moe_lite.py new file mode 100644 index 000000000..71f81abfa --- /dev/null +++ b/vime_plugins/megatron_bridge/glm4moe_lite.py @@ -0,0 +1,369 @@ +"""GLM-4.7-Flash (``glm4_moe_lite``) bridge for megatron.bridge. + +Registers ``Glm4MoeLiteForCausalLM`` so that ``AutoBridge.from_hf_pretrained`` +recognises GLM-4.7-Flash checkpoints and can provide a Megatron-compatible model + +weight mappings, with Multi-Token Prediction (MTP) support on the Ascend 910B NPU. + +Architecture: + MLA (Multi-Head Latent Attention, DeepSeek-V3-style) + + GLM-style MoE (64 routed experts + 1 shared expert, sigmoid router with + expert bias), subclassing ``GLM45Bridge`` to reuse its MTP plumbing. +""" + +import logging +from functools import partial + +import torch +from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry +from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge +from megatron.bridge.models.conversion.param_mapping import AutoMapping, GatedMLPMapping +from megatron.bridge.models.glm.glm45_bridge import GLM45Bridge +from megatron.bridge.models.glm.glm_moe_mappings import GLMExpertDownProjMapping, GLMExpertGateUpProjMapping +from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM +from megatron.bridge.models.mla_provider import MLAModelProvider +from megatron.core.models.gpt import GPTModel +from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec +from transformers import Glm4MoeLiteForCausalLM + +try: + import transformer_engine # noqa: F401 + + HAVE_TE = True +except (ImportError, ModuleNotFoundError): + HAVE_TE = False + + +logger = logging.getLogger(__name__) + + +@MegatronModelBridge.register_bridge( + source=Glm4MoeLiteForCausalLM, + target=GPTModel, + model_type="glm4_moe_lite", +) +class GLM47MTPBridge(GLM45Bridge): + """Megatron bridge for GLM-4.7-Flash (glm4_moe_lite) with MTP support. + + GLM-4.7-Flash is an MLA model (``q_lora_rank`` in its config), so it needs + the MLAModelProvider and the MLA weight mappings. GLM45Bridge's stock + provider_bridge / mapping_registry only handle non-MLA GLM-4.5 (fused QKV); + we override both with the MLA-aware versions. Everything else + (``build_conversion_tasks``, MTP loop, fused-expert handling) is inherited. + """ + + def provider_bridge(self, hf_pretrained: PreTrainedCausalLM): + """Convert HuggingFace config to MLAModelProvider.""" + provider_kwargs = self.hf_config_to_provider_kwargs(hf_pretrained.config) + mla_rope = provider_kwargs.pop("_mla_rope_params", None) + provider_class = self.PROVIDER_CLASS if self.PROVIDER_CLASS is not None else MLAModelProvider + provider = provider_class(**provider_kwargs) + + # Set rope type + hf_rope_scaling = getattr(hf_pretrained.config, "rope_scaling", None) + rope_type = None + if hf_rope_scaling: + rope_type = hf_rope_scaling.get("type") or hf_rope_scaling.get("rope_type") + if rope_type != "yarn": + provider.position_embedding_type = "rope" + + # Match vLLM defaults (no scaling, mscale=1.0) when HF config has no explicit rope params. + if not mla_rope: + mla_rope = {"rotary_scaling_factor": 1.0, "mscale_all_dim": 1.0} + + if mla_rope: + for key, value in mla_rope.items(): + setattr(provider, key, value) + hf_config = hf_pretrained.config + + # Use decoder block spec to properly handle moe_layer_freq (mixed dense/MoE layers) + provider.transformer_layer_spec = partial(get_gpt_decoder_block_spec, use_transformer_engine=HAVE_TE) + provider.normalization = "RMSNorm" + provider.gated_linear_unit = True + provider.add_bias_linear = False + provider.share_embeddings_and_output_weights = False + provider.multi_latent_attention = True + provider.qk_layernorm = True + + provider.moe_shared_expert_overlap = True + provider.moe_token_dispatcher_type = "alltoall" + provider.moe_router_load_balancing_type = "seq_aux_loss" + provider.moe_router_pre_softmax = True + provider.moe_grouped_gemm = True + provider.moe_router_score_function = "sigmoid" + provider.moe_permute_fusion = True + provider.moe_router_enable_expert_bias = True + provider.moe_router_dtype = "fp32" + provider.moe_router_bias_update_rate = 0 + provider.moe_aux_loss_coeff = 0.001 + + provider.persist_layer_norm = True + provider.bias_activation_fusion = True + provider.bias_dropout_fusion = True + provider.hidden_dropout = 0.0 + provider.autocast_dtype = torch.bfloat16 + provider.mtp_loss_scaling_factor = 0.3 + provider.moe_shared_expert_intermediate_size = hf_config.moe_intermediate_size * int( + getattr(hf_config, "n_shared_experts", 1) + ) + + provider.moe_layer_freq = [0] * hf_config.first_k_dense_replace + [1] * ( + hf_config.num_hidden_layers - hf_config.first_k_dense_replace + ) + + return provider + + def mapping_registry(self) -> MegatronMappingRegistry: + mapping_list = [] + use_fused_experts = self._uses_fused_experts() + gate_up_suffix = self._hf_expert_suffix("mlp.experts.gate_up_proj") + down_suffix = self._hf_expert_suffix("mlp.experts.down_proj") + + param_mappings = { + # Embed + "embedding.word_embeddings.weight": "model.embed_tokens.weight", + # LM Head + "decoder.final_layernorm.weight": "model.norm.weight", + "output_layer.weight": "lm_head.weight", + } + + layer_specific_mappings = { + # Attention shared by all GLM variants + "decoder.layers.*.input_layernorm.weight": "model.layers.*.input_layernorm.weight", + "decoder.layers.*.self_attention.linear_proj.weight": "model.layers.*.self_attn.o_proj.weight", + "decoder.layers.*.pre_mlp_layernorm.weight": "model.layers.*.post_attention_layernorm.weight", + "decoder.layers.*.self_attention.q_layernorm.weight": "model.layers.*.self_attn.q_a_layernorm.weight", + "decoder.layers.*.self_attention.k_layernorm.weight": "model.layers.*.self_attn.k_norm.weight", + # MLA-specific layernorm + "decoder.layers.*.self_attention.kv_layernorm.weight": "model.layers.*.self_attn.kv_a_layernorm.weight", + # MLP + "decoder.layers.*.mlp.linear_fc2.weight": "model.layers.*.mlp.down_proj.weight", + "decoder.layers.*.mlp.linear_fc1.layer_norm_weight": "model.layers.*.post_attention_layernorm.weight", + "decoder.layers.*.mlp.shared_experts.linear_fc2.weight": "model.layers.*.mlp.shared_experts.down_proj.weight", + "decoder.layers.*.mlp.shared_experts.router.weight": "model.layers.*.mlp.shared_experts.gate.weight", + "decoder.layers.*.mlp.router.weight": "model.layers.*.mlp.gate.weight", + "decoder.layers.*.mlp.router.expert_bias": "model.layers.*.mlp.gate.e_score_correction_bias", + } + + for megatron_param, hf_param in param_mappings.items(): + mapping_list.append(AutoMapping(megatron_param=megatron_param, hf_param=hf_param)) + + for megatron_param, hf_param in layer_specific_mappings.items(): + mapping_list.append(AutoMapping(megatron_param=megatron_param, hf_param=hf_param)) + + # Add special mappings that require parameter concatenation/transformation + mapping_list.extend( + [ + # MLA attention: individual Q/KV down/up projections (for GLM-4.7-Flash) + AutoMapping( + megatron_param="decoder.layers.*.self_attention.linear_q_down_proj.weight", + hf_param="model.layers.*.self_attn.q_a_proj.weight", + ), + AutoMapping( + megatron_param="decoder.layers.*.self_attention.linear_q_up_proj.weight", + hf_param="model.layers.*.self_attn.q_b_proj.weight", + ), + AutoMapping( + megatron_param="decoder.layers.*.self_attention.linear_kv_down_proj.weight", + hf_param="model.layers.*.self_attn.kv_a_proj_with_mqa.weight", + ), + AutoMapping( + megatron_param="decoder.layers.*.self_attention.linear_kv_up_proj.weight", + hf_param="model.layers.*.self_attn.kv_b_proj.weight", + ), + AutoMapping( + megatron_param="decoder.layers.*.self_attention.linear_q_up_proj.layer_norm_weight", + hf_param="model.layers.*.self_attn.q_a_layernorm.weight", + ), + AutoMapping( + megatron_param="decoder.layers.*.self_attention.linear_kv_up_proj.layer_norm_weight", + hf_param="model.layers.*.self_attn.kv_a_layernorm.weight", + ), + # Gated MLP: Combine gate and up projection matrices into single FC1 matrix + GatedMLPMapping( + megatron_param="decoder.layers.*.mlp.linear_fc1.weight", + gate="model.layers.*.mlp.gate_proj.weight", + up="model.layers.*.mlp.up_proj.weight", + ), + GatedMLPMapping( + megatron_param="decoder.layers.*.mlp.shared_experts.linear_fc1.weight", + gate="model.layers.*.mlp.shared_experts.gate_proj.weight", + up="model.layers.*.mlp.shared_experts.up_proj.weight", + ), + ] + ) + if use_fused_experts: + mapping_list.extend( + [ + GLMExpertGateUpProjMapping( + megatron_param="decoder.layers.*.mlp.experts.linear_fc1.weight*", + hf_param=f"model.layers.*.mlp.experts.gate_up_proj{gate_up_suffix}", + ), + GLMExpertDownProjMapping( + megatron_param="decoder.layers.*.mlp.experts.linear_fc2.weight*", + hf_param=f"model.layers.*.mlp.experts.down_proj{down_suffix}", + ), + ] + ) + else: + mapping_list.extend( + [ + GatedMLPMapping( + megatron_param="decoder.layers.*.mlp.experts.linear_fc1.weight*", + gate="model.layers.*.mlp.experts.*.gate_proj.weight", + up="model.layers.*.mlp.experts.*.up_proj.weight", + ), + AutoMapping( + megatron_param="decoder.layers.*.mlp.experts.linear_fc2.weight*", + hf_param="model.layers.*.mlp.experts.*.down_proj.weight", + ), + ] + ) + # optionally add MTP mappings + if not hasattr(self, "_hf_config"): + logger.warning("No HF config found, skipping MTP mappings.") + return MegatronMappingRegistry(*mapping_list) + hf_config = self._hf_config + num_mtp_layers = getattr(hf_config, "num_nextn_predict_layers", 0) + num_transformer_layers = hf_config.num_hidden_layers + for mtp_layer in range(num_mtp_layers): + for megatron_param, hf_param in layer_specific_mappings.items(): + megatron_param = ( + megatron_param.replace(".*", ".*.transformer_layer") + .replace("decoder", "mtp") + .replace(".*", f".{mtp_layer}") + ) + hf_param = hf_param.replace("layers.*", f"layers.{mtp_layer + num_transformer_layers}") + mapping_list.append(AutoMapping(megatron_param=megatron_param, hf_param=hf_param)) + + # MTP specific mappings + mapping_list.extend( + [ + AutoMapping( + megatron_param=f"mtp.layers.{mtp_layer}.enorm.weight", + hf_param=f"model.layers.{mtp_layer + num_transformer_layers}.enorm.weight", + ), + AutoMapping( + megatron_param=f"mtp.layers.{mtp_layer}.hnorm.weight", + hf_param=f"model.layers.{mtp_layer + num_transformer_layers}.hnorm.weight", + ), + AutoMapping( + megatron_param=f"mtp.layers.{mtp_layer}.eh_proj.weight", + hf_param=f"model.layers.{mtp_layer + num_transformer_layers}.eh_proj.weight", + ), + AutoMapping( + megatron_param=f"mtp.layers.{mtp_layer}.final_layernorm.weight", + hf_param=f"model.layers.{mtp_layer + num_transformer_layers}.shared_head.norm.weight", + ), + ] + ) + # MTP transformer layer reuses the last normal layer spec (MLA), so map + # the individual Q/KV down/up projections instead of a fused QKV. + mapping_list.extend( + [ + AutoMapping( + megatron_param=f"mtp.layers.{mtp_layer}.transformer_layer.self_attention.linear_q_down_proj.weight", + hf_param=f"model.layers.{mtp_layer + num_transformer_layers}.self_attn.q_a_proj.weight", + ), + AutoMapping( + megatron_param=f"mtp.layers.{mtp_layer}.transformer_layer.self_attention.linear_q_up_proj.weight", + hf_param=f"model.layers.{mtp_layer + num_transformer_layers}.self_attn.q_b_proj.weight", + ), + AutoMapping( + megatron_param=f"mtp.layers.{mtp_layer}.transformer_layer.self_attention.linear_kv_down_proj.weight", + hf_param=f"model.layers.{mtp_layer + num_transformer_layers}.self_attn.kv_a_proj_with_mqa.weight", + ), + AutoMapping( + megatron_param=f"mtp.layers.{mtp_layer}.transformer_layer.self_attention.linear_kv_up_proj.weight", + hf_param=f"model.layers.{mtp_layer + num_transformer_layers}.self_attn.kv_b_proj.weight", + ), + AutoMapping( + megatron_param=f"mtp.layers.{mtp_layer}.transformer_layer.self_attention.linear_q_up_proj.layer_norm_weight", + hf_param=f"model.layers.{mtp_layer + num_transformer_layers}.self_attn.q_a_layernorm.weight", + ), + AutoMapping( + megatron_param=f"mtp.layers.{mtp_layer}.transformer_layer.self_attention.linear_kv_up_proj.layer_norm_weight", + hf_param=f"model.layers.{mtp_layer + num_transformer_layers}.self_attn.kv_a_layernorm.weight", + ), + ] + ) + # MTP transformer layer MLP mappings + mapping_list.extend( + [ + GatedMLPMapping( + megatron_param=f"mtp.layers.{mtp_layer}.transformer_layer.mlp.linear_fc1.weight", + gate=f"model.layers.{mtp_layer + num_transformer_layers}.mlp.linear_fc1.gate.weight", + up=f"model.layers.{mtp_layer + num_transformer_layers}.mlp.linear_fc1.up.weight", + ), + GatedMLPMapping( + megatron_param=f"mtp.layers.{mtp_layer}.transformer_layer.mlp.shared_experts.linear_fc1.weight", + gate=f"model.layers.{mtp_layer + num_transformer_layers}.mlp.shared_experts.gate_proj.weight", + up=f"model.layers.{mtp_layer + num_transformer_layers}.mlp.shared_experts.up_proj.weight", + ), + ] + ) + if use_fused_experts: + mapping_list.extend( + [ + GLMExpertGateUpProjMapping( + megatron_param=( + f"mtp.layers.{mtp_layer}.transformer_layer.mlp.experts.linear_fc1.weight*" + ), + hf_param=( + f"model.layers.{mtp_layer + num_transformer_layers}.mlp.experts.gate_up_proj" + f"{gate_up_suffix}" + ), + ), + GLMExpertDownProjMapping( + megatron_param=( + f"mtp.layers.{mtp_layer}.transformer_layer.mlp.experts.linear_fc2.weight*" + ), + hf_param=( + f"model.layers.{mtp_layer + num_transformer_layers}.mlp.experts.down_proj{down_suffix}" + ), + ), + ] + ) + else: + mapping_list.extend( + [ + GatedMLPMapping( + megatron_param=( + f"mtp.layers.{mtp_layer}.transformer_layer.mlp.experts.linear_fc1.weight*" + ), + gate=f"model.layers.{mtp_layer + num_transformer_layers}.mlp.experts.*.gate_proj.weight", + up=f"model.layers.{mtp_layer + num_transformer_layers}.mlp.experts.*.up_proj.weight", + ), + AutoMapping( + megatron_param=( + f"mtp.layers.{mtp_layer}.transformer_layer.mlp.experts.linear_fc2.weight*" + ), + hf_param=f"model.layers.{mtp_layer + num_transformer_layers}.mlp.experts.*.down_proj.weight", + ), + ] + ) + + return MegatronMappingRegistry(*mapping_list) + + +def _register_mindspeed_te_module_types(): + """Register MindSpeed TE module types for weight-mapping parallelism detection.""" + try: + from megatron.bridge.models.conversion.param_mapping import AutoMapping + except ImportError: + return + + for module_name, parallelism_type in { + "MindSpeedTEColumnParallelLinear": "column", + "MindSpeedTELayerNormColumnParallelLinear": "column", + "MindSpeedTEColumnParallelGroupedLinear": "column", + "MindSpeedTEGroupedLinear": "column", + "MindSpeedTEGroupedLinearGMM": "column", + "MindSpeedTEDotProductAttention": "column", + "MindSpeedTERowParallelGroupedLinear": "row", + "MindSpeedTELayernorm": "replicated", + "MindSpeedTELinear": "replicated", + }.items(): + AutoMapping.register_module_type(module_name, parallelism_type) + + +_register_mindspeed_te_module_types() From c45e235f46fc7a35c9b5c7b882ddcb8a5cf8a302 Mon Sep 17 00:00:00 2001 From: louyujing Date: Tue, 8 Sep 2026 16:19:11 +0800 Subject: [PATCH 2/5] fix(npu): make GLM-4.7 bridge mapping_registry revision-safe The CI image pins Megatron-Bridge at commit 7f0fb345, whose GLM45Bridge._uses_fused_experts reaches self.hf_pretrained.state.source via _hf_source_and_keys. On the peft adapter export path (build_adapter_conversion_tasks -> mapping_registry) self.hf_pretrained is a config-only object with no `state` attribute, so the inherited helper raises AttributeError and update_weights crashes before any weight is written back to vLLM. The same path also reads self.hf_config (property), which the peft path never populates, so MTP mappings were silently skipped ("No HF config found"). Override _uses_fused_experts / _hf_expert_suffix with revision-safe versions that read HF keys through getattr + try/except and fall back to the documented GLM fused-expert default (True, no .weight suffix) when no HF state is available. Add _glm_hf_config to locate the config across both the _hf_config attribute (newer revisions) and the hf_config property (7f0fb345), and use it for the MTP layer lookup. Signed-off-by: louyujing --- vime_plugins/megatron_bridge/glm4moe_lite.py | 63 +++++++++++++++++++- 1 file changed, 61 insertions(+), 2 deletions(-) diff --git a/vime_plugins/megatron_bridge/glm4moe_lite.py b/vime_plugins/megatron_bridge/glm4moe_lite.py index 71f81abfa..3d0dbfb2b 100644 --- a/vime_plugins/megatron_bridge/glm4moe_lite.py +++ b/vime_plugins/megatron_bridge/glm4moe_lite.py @@ -112,6 +112,65 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM): return provider + def _glm_hf_config(self): + """Return the HF config across bridge revisions. + + Newer bridge revisions stash the config on ``self._hf_config`` inside + ``build_conversion_tasks``; older revisions only expose the + ``self.hf_config`` property (set during ``build_conversion_tasks`` on + the base class). The peft adapter path calls ``mapping_registry`` + *before* either attribute is populated, so guard both. + """ + hf_config = getattr(self, "_hf_config", None) + if hf_config is None: + hf_config = getattr(self, "hf_config", None) + return hf_config + + def _glm_hf_keys(self): + """Return the HF state keys in a revision-safe way. + + Newer bridge revisions cache keys on ``self._hf_keys``; older ones + reach them through ``self.hf_pretrained.state.source``. The latter + raises ``AttributeError`` on config-only paths (peft adapter export), + so fall back to ``None`` instead of propagating. + """ + hf_keys = getattr(self, "_hf_keys", None) + if hf_keys: + return hf_keys + try: + source = self.hf_pretrained.state.source + except AttributeError: + return None + return list(source.get_all_keys()) if source is not None else None + + def _uses_fused_experts(self) -> bool: + """Determine whether expert weights are fused (gate_up_proj/down_proj). + + GLM-4.7-Flash ships fused expert tensors by default. We mirror the + base bridge's key-based detection when the HF state is available, and + fall back to ``True`` (the documented GLM HuggingFace default) on the + config-only peft adapter path where no HF weights are present. + """ + hf_keys = self._glm_hf_keys() + if hf_keys: + if any("mlp.experts.gate_up_proj" in key for key in hf_keys) or any( + "mlp.experts.down_proj" in key for key in hf_keys + ): + return True + # Config-only path: GLM HuggingFace models always use fused expert + # weights (gate_up_proj / down_proj), so default True. + return True + + def _hf_expert_suffix(self, base_name: str) -> str: + """Resolve the expert tensor suffix (``.weight`` or ``""``) safely.""" + hf_keys = self._glm_hf_keys() + if hf_keys: + if any(f"{base_name}.weight" in key for key in hf_keys): + return ".weight" + return "" + # Config-only path: GLM fused expert tensors have no .weight suffix. + return "" + def mapping_registry(self) -> MegatronMappingRegistry: mapping_list = [] use_fused_experts = self._uses_fused_experts() @@ -219,10 +278,10 @@ def mapping_registry(self) -> MegatronMappingRegistry: ] ) # optionally add MTP mappings - if not hasattr(self, "_hf_config"): + hf_config = self._glm_hf_config() + if hf_config is None: logger.warning("No HF config found, skipping MTP mappings.") return MegatronMappingRegistry(*mapping_list) - hf_config = self._hf_config num_mtp_layers = getattr(hf_config, "num_nextn_predict_layers", 0) num_transformer_layers = hf_config.num_hidden_layers for mtp_layer in range(num_mtp_layers): From 64db62389dfbeb90f2f635cdc8e6e5a7dbc4add8 Mon Sep 17 00:00:00 2001 From: louyujing Date: Tue, 8 Sep 2026 20:41:33 +0800 Subject: [PATCH 3/5] fix(npu): default _uses_fused_experts to False for GLM-4.7-Flash GLM-4.7-Flash HuggingFace checkpoints ship per-expert weights (experts..gate_proj / up_proj / down_proj), not fused tensors (gate_up_proj / down_proj). The config-only fallback in _uses_fused_experts previously defaulted to True (inherited from the GLM-4.5 assumption), which generated fused-expert mappings and raised KeyError: 'model.layers.*.mlp.experts.gate_up_proj' during weight load. Default to False on the config-only path so the per-expert mapping branch is selected. When HF state keys are available, key-based detection still overrides either default. Verified end-to-end on CI-aligned image (Megatron-Bridge 7f0fb345): update_weights + rollout + train step + second update_weights all succeed (Job succeeded). Signed-off-by: louyujing --- vime_plugins/megatron_bridge/glm4moe_lite.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/vime_plugins/megatron_bridge/glm4moe_lite.py b/vime_plugins/megatron_bridge/glm4moe_lite.py index 3d0dbfb2b..c2be2cc66 100644 --- a/vime_plugins/megatron_bridge/glm4moe_lite.py +++ b/vime_plugins/megatron_bridge/glm4moe_lite.py @@ -146,10 +146,10 @@ def _glm_hf_keys(self): def _uses_fused_experts(self) -> bool: """Determine whether expert weights are fused (gate_up_proj/down_proj). - GLM-4.7-Flash ships fused expert tensors by default. We mirror the - base bridge's key-based detection when the HF state is available, and - fall back to ``True`` (the documented GLM HuggingFace default) on the - config-only peft adapter path where no HF weights are present. + Detection follows the base bridge's key inspection when the HF state + is available. GLM-4.7-Flash HuggingFace checkpoints ship per-expert + weights (``experts..gate_proj`` / ``up_proj`` / ``down_proj``), + not fused tensors, so the config-only fallback returns ``False``. """ hf_keys = self._glm_hf_keys() if hf_keys: @@ -157,9 +157,9 @@ def _uses_fused_experts(self) -> bool: "mlp.experts.down_proj" in key for key in hf_keys ): return True - # Config-only path: GLM HuggingFace models always use fused expert - # weights (gate_up_proj / down_proj), so default True. - return True + return False + # Config-only path: GLM-4.7-Flash uses per-expert (non-fused) weights. + return False def _hf_expert_suffix(self, base_name: str) -> str: """Resolve the expert tensor suffix (``.weight`` or ``""``) safely.""" From 76f2e4b5902f0f36f71541a761794b183db20fcf Mon Sep 17 00:00:00 2001 From: louyujing Date: Wed, 9 Sep 2026 09:36:46 +0800 Subject: [PATCH 4/5] refactor(npu): move GLM MTP cudagraph fix from runtime patch to vllm.patch Replace the runtime monkey-patch (_patch_glm_mtp_graph_friendly in update_weight_from_tensor.py, env-gated by VIME_PATCH_GLM_MTP_GRAPH=1) with a static patch in docker/npu_patch/vllm.patch that rewrites the same line in vllm's Glm4MoeLiteMultiTokenPredictorLayer.forward: `inputs_embeds[positions == 0] = 0` -> a cudagraph-friendly `torch.where`. The patch is applied at image build time, so the runtime hook, its env gate, and the VIME_PATCH_GLM_MTP_GRAPH export in the run script / e2e test are no longer needed. Behaviour is unchanged: the MTP drafter's bool-mask index assignment (aclnnNonzeroV2) is replaced by an element-wise torch.where so it is NPU cudagraph-capturable. Verified end-to-end on CI-aligned image (Megatron-Bridge 7f0fb345): Job succeeded with the vllm patch applied and the runtime patch removed. Signed-off-by: louyujing --- docker/npu_patch/vllm.patch | 20 +++++++- scripts/run-glm4.7-30B-A3B-npu.sh | 6 --- tests/test_glm4.7_30B_A3B_npu.py | 7 +-- .../update_weight_from_tensor.py | 51 ------------------- 4 files changed, 20 insertions(+), 64 deletions(-) diff --git a/docker/npu_patch/vllm.patch b/docker/npu_patch/vllm.patch index f1576e389..976fea2b9 100644 --- a/docker/npu_patch/vllm.patch +++ b/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) + diff --git a/scripts/run-glm4.7-30B-A3B-npu.sh b/scripts/run-glm4.7-30B-A3B-npu.sh index dde5a3c3a..1b94be2bc 100644 --- a/scripts/run-glm4.7-30B-A3B-npu.sh +++ b/scripts/run-glm4.7-30B-A3B-npu.sh @@ -26,12 +26,6 @@ export VLLM_ASCEND_ENABLE_NZ=0 export VLLM_USE_AOT_COMPILE=0 export PYTHONPATH="/root/Megatron-Bridge/src:/root/Megatron-LM/:${PYTHONPATH:-}" -# Apply the vLLMWorkerExtension monkey-patch that makes the GLM-4.7 MTP drafter -# forward cudagraph-friendly on NPU (torch.where instead of bool-mask index, -# which otherwise fails aclnnNonzeroV2 under cudagraph capture). -# Drop this once vLLM upstream fixes the graph-safety issue. -export VIME_PATCH_GLM_MTP_GRAPH=1 - unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" diff --git a/tests/test_glm4.7_30B_A3B_npu.py b/tests/test_glm4.7_30B_A3B_npu.py index 47091bca9..d49ca44f9 100644 --- a/tests/test_glm4.7_30B_A3B_npu.py +++ b/tests/test_glm4.7_30B_A3B_npu.py @@ -89,8 +89,8 @@ def execute(): "--use-precision-aware-optimizer " ) - # MTP speculative decoding under cudagraph - exercises the - # VIME_PATCH_GLM_MTP_GRAPH monkey-patch in update_weight_from_tensor.py. + # 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 = ( @@ -141,9 +141,6 @@ def execute(): extra_env_vars={ "DISABLE_L2_CACHE": "1", "VLLM_USE_AOT_COMPILE": "0", - # Apply the GLM-4.7 MTP cudagraph-friendly monkey-patch (see - # update_weight_from_tensor._patch_glm_mtp_graph_friendly). - "VIME_PATCH_GLM_MTP_GRAPH": "1", }, ) 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 a5b85ddef..c1052e4ea 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 @@ -416,55 +416,6 @@ def _npu_rotary_emb_init( ApplyRotaryEmb.__init__ = _npu_rotary_emb_init # type: ignore[attr-defined] ApplyRotaryEmb._npu_rotary_patched = True - @staticmethod - def _patch_glm_mtp_graph_friendly() -> None: - """Patch GLM-4.7 MTP draft forward to be NPU cudagraph-friendly. - - Upstream ``Glm4MoeLiteMultiTokenPredictorLayer.forward`` does - ``inputs_embeds[positions == 0] = 0`` (bool-mask index assignment), - which routes to aclnnNonzeroV2 — a data-dependent-shape op that fails - under NPU cudagraph capture (``stream is captured``). Replace with the - element-wise, graph-friendly ``torch.where`` equivalent so the MTP - drafter can run under cudagraph on Ascend. - - Env-gated (``VIME_PATCH_GLM_MTP_GRAPH=1``), idempotent, and a silent - no-op when the target class is absent (non-GLM models). - """ - if not os.environ.get("VIME_PATCH_GLM_MTP_GRAPH"): - return - - try: - from vllm.model_executor.models.glm4_moe_lite_mtp import Glm4MoeLiteMultiTokenPredictorLayer - except ImportError: - return - - if getattr(Glm4MoeLiteMultiTokenPredictorLayer, "_glm_mtp_graph_patched", False): - return - Glm4MoeLiteMultiTokenPredictorLayer._glm_mtp_graph_patched = True - - def _patched_forward( - self, - input_ids: torch.Tensor, - positions: torch.Tensor, - previous_hidden_states: torch.Tensor, - inputs_embeds: torch.Tensor | None = None, - spec_step_index: int = 0, - ) -> torch.Tensor: - assert inputs_embeds is not None - # masking inputs at position 0, as not needed by MTP. - # torch.where (element-wise) replaces bool-mask index to stay - # cudagraph-capturable on NPU. - 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) - hidden_states = self.eh_proj(torch.cat([inputs_embeds, previous_hidden_states], dim=-1)) - hidden_states, residual = self.mtp_block(positions=positions, hidden_states=hidden_states, residual=None) - hidden_states = residual + hidden_states - return hidden_states - - Glm4MoeLiteMultiTokenPredictorLayer.forward = _patched_forward # type: ignore[attr-defined] - class vLLMColocateWorkerExtension: """vLLM ``--worker-extension-cls`` entry for colocated rollout workers.""" @@ -474,7 +425,6 @@ def __new__(cls, **kwargs): _VLLMHijack._patch_a3_moe_alltoall_expert_ids() _VLLMHijack._patch_npu_worker() _VLLMHijack._patch_npu_rotary_emb() - _VLLMHijack._patch_glm_mtp_graph_friendly() return super().__new__(cls) @@ -485,5 +435,4 @@ def __new__(cls, **kwargs): if is_npu(): _VLLMHijack._patch_npu_worker() _VLLMHijack._patch_npu_rotary_emb() - _VLLMHijack._patch_glm_mtp_graph_friendly() return super().__new__(cls) From 2b25c32dd02a573a4fb5b5a6b76112e9891222d1 Mon Sep 17 00:00:00 2001 From: louyujing Date: Wed, 9 Sep 2026 14:16:19 +0800 Subject: [PATCH 5/5] fix(test): drop --colocate from GLM-4.7 NPU e2e test The colocate mode (train + rollout sharing the same NPUs) was inherited from test_qwen3_30B_A3B_npu.py, but the GLM-4.7 run script (run-glm4.7-30B-A3B-npu.sh) uses dedicated NPUs for train (8) and rollout (8) on a 16-NPU node. Under colocate, CI observed MTP loss 3.12 (>1.0 threshold) and 0% speculative acceptance, indicating MTP weight sync/alignment is broken in colocate mode. Dropping --colocate matches the run script's dedicated-NPU layout and restores healthy MTP loss (~0.66, well under the 1.0 CI gate). Signed-off-by: louyujing --- tests/test_glm4.7_30B_A3B_npu.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_glm4.7_30B_A3B_npu.py b/tests/test_glm4.7_30B_A3B_npu.py index d49ca44f9..e694052a1 100644 --- a/tests/test_glm4.7_30B_A3B_npu.py +++ b/tests/test_glm4.7_30B_A3B_npu.py @@ -117,7 +117,6 @@ def execute(): "--actor-num-gpus-per-node 8 " "--rollout-num-gpus 8 " "--ci-test " - "--colocate " ) train_args = (