diff --git a/docs/en/examples/glm4.7-30B-A3B.md b/docs/en/examples/glm4.7-30B-A3B.md index 1ad877e7e2..aa6f6c70ae 100644 --- a/docs/en/examples/glm4.7-30B-A3B.md +++ b/docs/en/examples/glm4.7-30B-A3B.md @@ -5,6 +5,14 @@ The environment setup, data, and checkpoint conversion are the same as for the Qwen3-4B model. You can refer to [Example: Qwen3-4B Model](qwen3-4B.md), replacing mentions of Qwen3-4B with GLM-4.7-Flash. +### Prerequisites + +GLM-4.7-Flash requires **transformers ≥ 5.0** for the `Glm4MoeLiteForCausalLM` architecture. Install or upgrade: + +```bash +pip install "transformers>=5.0" +``` + ### Download Model ```bash @@ -119,9 +127,9 @@ SPEC_ARGS=( - `--enable-mtp-training`: Enables gradient computation for MTP layers. Without this flag, the MTP layer is loaded but frozen. - `--mtp-loss-scaling-factor 0.2`: Weight of the MTP loss relative to the main policy loss. Default is 0.2. -> ⚠️ **Note**: MTP training for GLM-4.7-Flash is not yet supported because the deepseek_v3 checkpoint bridge does not include MTP weight conversion (`# TODO: mtp` in upstream mbridge). You can still use MTP for speculative decoding during inference — SGLang handles MTP layers internally. +> **Note**: MTP training requires the MTP checkpoint bridge to properly convert weights between HuggingFace and Megatron formats. The `GLM4MoELiteBridge` (in `slime_plugins/mbridge/glm4moe_lite.py`) extends the DeepSeek V3 bridge with dynamic MTP layer indexing to support GLM-4.7-Flash's 47-layer architecture. > -> For models with full MTP training support (e.g., MiMo), see `scripts/run-mimo-7B-rl-eagle.sh` as a reference. +> For other models with MTP training support (e.g., MiMo), see `scripts/run-mimo-7B-rl-eagle.sh` as a reference. ### Multi-Node Support diff --git a/docs/zh/examples/glm4.7-30B-A3B.md b/docs/zh/examples/glm4.7-30B-A3B.md index 0a7a141c27..cf4c1b0602 100644 --- a/docs/zh/examples/glm4.7-30B-A3B.md +++ b/docs/zh/examples/glm4.7-30B-A3B.md @@ -4,6 +4,14 @@ 搭建环境、数据与 ckpt 转换均与 Qwen3-4B 模型相同，可以参考 [示例：Qwen3-4B](qwen3-4B.md)，将文中 Qwen3-4B 的部分转换为 GLM-4.7-Flash 即可。 +### 前置条件 + +GLM-4.7-Flash 需要 **transformers ≥ 5.0** 以支持 `Glm4MoeLiteForCausalLM` 架构。请安装或升级： + +```bash +pip install "transformers>=5.0" +``` + ### 下载模型 ```bash @@ -118,9 +126,9 @@ SPEC_ARGS=( - `--enable-mtp-training`：启用 MTP 层的梯度计算。不设置此标志时，MTP 层会被加载但冻结。 - `--mtp-loss-scaling-factor 0.2`：MTP loss 相对于主策略 loss 的权重，默认为 0.2。 -> ⚠️ **注意**：GLM-4.7-Flash 的 MTP 训练目前尚不支持，因为 deepseek_v3 的 checkpoint bridge 尚未实现 MTP 权重转换（上游 mbridge 中标注为 `# TODO: mtp`）。但推理时的投机解码仍然可用 — SGLang 会内部处理 MTP 层。 +> **注意**：MTP 训练需要 MTP checkpoint bridge 正确转换 HuggingFace 和 Megatron 格式之间的权重。`GLM4MoELiteBridge`（位于 `slime_plugins/mbridge/glm4moe_lite.py`）扩展了 DeepSeek V3 bridge，实现了动态 MTP 层索引以支持 GLM-4.7-Flash 的 47 层架构。 > -> 对于完整支持 MTP 训练的模型（如 MiMo），可参考 `scripts/run-mimo-7B-rl-eagle.sh`。 +> 对于其他支持 MTP 训练的模型（如 MiMo），可参考 `scripts/run-mimo-7B-rl-eagle.sh`。 ### 多机支持 diff --git a/scripts/run-glm4.7-30B-A3B-8gpus.sh b/scripts/run-glm4.7-30B-A3B-8gpus.sh deleted file mode 100755 index 62c56529be..0000000000 --- a/scripts/run-glm4.7-30B-A3B-8gpus.sh +++ /dev/null @@ -1,175 +0,0 @@ -#!/bin/bash - -# for rerun the task -pkill -9 sglang -sleep 3 -ray stop --force -pkill -9 ray -pkill -9 python -sleep 3 -pkill -9 ray -pkill -9 python - -set -ex - -# will prevent ray from buffering stdout/stderr -export PYTHONBUFFERED=16 - -NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) -if [ "$NVLINK_COUNT" -gt 0 ]; then - HAS_NVLINK=1 -else - HAS_NVLINK=0 -fi -echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/models/glm4.7-30B-A3B.sh" - -CKPT_ARGS=( - --hf-checkpoint /root/GLM-4.7-Flash - --ref-load /root/GLM-4.7-Flash_torch_dist - --load /root/GLM-4.7-Flash_slime/ - --save /root/GLM-4.7-Flash_slime/ - --save-interval 20 -) - -ROLLOUT_ARGS=( - --prompt-data /root/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 /root/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 1 - --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 8192 -) - -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 -) - -WANDB_ARGS=( - #--use-wandb - # --wandb-project slime-dev - # --wandb-group glm4.7-flash -) - -SGLANG_ARGS=( - --rollout-num-gpus-per-engine 8 - --sglang-mem-fraction-static 0.7 - --sglang-enable-dp-attention - --sglang-dp-size 8 - --sglang-enable-dp-lm-head - --sglang-moe-dense-tp-size 1 - - # MTP speculative decoding (EAGLE) — speeds up inference - # Uncomment the following lines to enable: - # --sglang-speculative-algorithm EAGLE - # --sglang-speculative-num-steps 2 - # --sglang-speculative-eagle-topk 1 - # --sglang-speculative-num-draft-tokens 3 - - --sglang-cuda-graph-max-bs 16 - --sglang-max-running-requests 64 -) - -MISC_ARGS=( - # default dropout in megatron is 0.1 - --attention-dropout 0.0 - --hidden-dropout 0.0 - # should be good for model performance - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - # need to comment this when using model with MLA - --attention-backend flash - - --moe-token-dispatcher-type alltoall -) - -# MTP training — uncomment to enable MTP layer training -# Requires --mtp-num-layers in MODEL_ARGS (add to glm4.7-30B-A3B.sh or override below) -# MODEL_ARGS+=(--mtp-num-layers 1) -SPEC_ARGS=( - # --enable-mtp-training - # --mtp-loss-scaling-factor 0.2 -) - -# 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} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 - -# Build the runtime environment JSON with proper variable substitution -RUNTIME_ENV_JSON="{ - \"env_vars\": { - \"PYTHONPATH\": \"/root/Megatron-LM/\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\" - } -}" - -ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ - -- python3 train.py \ - --actor-num-nodes 1 \ - --actor-num-gpus-per-node 8 \ - --colocate \ - ${MODEL_ARGS[@]} \ - ${CKPT_ARGS[@]} \ - ${ROLLOUT_ARGS[@]} \ - ${OPTIMIZER_ARGS[@]} \ - ${GRPO_ARGS[@]} \ - ${WANDB_ARGS[@]} \ - ${PERF_ARGS[@]} \ - ${EVAL_ARGS[@]} \ - ${SGLANG_ARGS[@]} \ - ${MISC_ARGS[@]} \ - ${SPEC_ARGS[@]} diff --git a/scripts/run-glm4.7-30B-A3B.sh b/scripts/run-glm4.7-30B-A3B.sh index 11f4f20648..b071dc5fe4 100644 --- a/scripts/run-glm4.7-30B-A3B.sh +++ b/scripts/run-glm4.7-30B-A3B.sh @@ -42,7 +42,6 @@ ROLLOUT_ARGS=( --num-rollout 3000 --rollout-batch-size 128 - #--over-sampling-batch-size 256 --n-samples-per-prompt 8 --rollout-max-response-len 32768 --rollout-temperature 1.0 @@ -56,14 +55,14 @@ EVAL_ARGS=( --eval-prompt-data aime24 $BASE_DIR/rl_data/aime-2024.jsonl --n-samples-per-eval-prompt 2 --eval-max-response-len 16384 - --eval-temperature 0.6 + --eval-temperature 1.0 --eval-top-p 0.95 ) PERF_ARGS=( --tensor-model-parallel-size 4 --sequence-parallel - --pipeline-model-parallel-size 2 + --pipeline-model-parallel-size 1 --context-parallel-size 2 --expert-model-parallel-size 8 --expert-tensor-parallel-size 1 @@ -77,6 +76,12 @@ PERF_ARGS=( --max-tokens-per-gpu 32768 ) +MTP_ARGS=( + --mtp-num-layers 1 + --enable-mtp-training + --mtp-loss-scaling-factor 0.2 +) + GRPO_ARGS=( --advantage-estimator grpo --use-kl-loss @@ -139,56 +144,24 @@ MISC_ARGS=( ) # launch the master node of ray in container -export MASTER_ADDR=${MLP_WORKER_0_HOST} -export no_proxy="127.0.0.1,${MASTER_ADDR}" -ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats - -for WORKER_IP in $(awk '{print $1}' /root/mpi_rack_hostfile); do - if [[ "$WORKER_IP" == "$MLP_WORKER_0_HOST" ]]; then - continue - fi - echo "Starting Ray worker on ${WORKER_IP}" - ssh root@"${WORKER_IP}" \ - "pkill -9 sglang ; ray stop --force ; pkill -9 python ; ray start --address=${MASTER_ADDR}:6379 --num-gpus 8 --node-ip-address ${WORKER_IP} --disable-usage-stats" & -done -wait +export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} +ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 + +# Build the runtime environment JSON with proper variable substitution +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/Megatron-LM/\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\" + } +}" ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json='{ - "env_vars": { - "no_proxy": "localhost,127.0.0.1,0.0.0.0,${MASTER_ADDR}", - "GLOO_SOCKET_IFNAME": "${MLP_SOCKET_IFNAME}", - "TP_SOCKET_IFNAME": "${MLP_SOCKET_IFNAME}", - "MASTER_ADDR": "${MLP_WORKER_0_HOST}", - "PYTHONPATH": "/root/Megatron-LM/", - "NCCL_CUMEM_ENABLE": "0", - "CUDA_DEVICE_MAX_CONNECTIONS": "1", - "NVTE_BWD_LAYERNORM_SM_MARGIN": "20", - "NCCL_IB_TC": "160", - "NCCL_PXN_DISABLE": "0", - "NCCL_IB_GID_INDEX": "3", - "NCCL_NET_GDR_LEVEL": "4", - "NCCL_IB_RETRY_CNT": "7", - "NCCL_IB_TIMEOUT": "32", - "NCCL_IB_QPS_PER_CONNECTION": "8", - "NCCL_P2P_LEVEL": "NVL", - "TORCH_NCCL_AVOID_RECORD_STREAMS": "1", - "NCCL_NVLS_ENABLE": "0", - "NCCL_MIN_CTAS": "4", - "OMPI_MCA_pml": "ob1", - "OMPI_MCA_btl": "^openib", - "OMPI_MCA_routed": "direct", - "OMPI_MCA_routed_radix": "1024", - "OMPI_MCA_plm_rsh_no_tree_spawn": "1", - "OMPI_MCA_oob_tcp_if_include": "${MLP_SOCKET_IFNAME}", - "OMPI_MCA_btl_tcp_if_include": "${MLP_SOCKET_IFNAME}" - } - }' \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ -- python3 train.py \ - --actor-num-nodes 2 \ + --actor-num-nodes 1 \ --actor-num-gpus-per-node 8 \ --colocate \ - --save-debug-rollout-data "data.pt" \ ${MODEL_ARGS[@]} \ ${CKPT_ARGS[@]} \ ${ROLLOUT_ARGS[@]} \ @@ -198,4 +171,5 @@ ray job submit --address="http://127.0.0.1:8265" \ ${PERF_ARGS[@]} \ ${EVAL_ARGS[@]} \ ${SGLANG_ARGS[@]} \ - ${MISC_ARGS[@]} + ${MISC_ARGS[@]} \ + ${MTP_ARGS[@]} diff --git a/slime/backends/megatron_utils/model_provider.py b/slime/backends/megatron_utils/model_provider.py index 31db8b0da8..bd5f711288 100644 --- a/slime/backends/megatron_utils/model_provider.py +++ b/slime/backends/megatron_utils/model_provider.py @@ -96,6 +96,20 @@ def wrapped_model_provider( if getattr(args, "decoder_last_pipeline_num_layers", None) is not None: provider.num_layers_in_last_pipeline_stage = args.decoder_last_pipeline_num_layers provider.finalize() + + if role == "critic": + _original_provide = provider.provide + + def _critic_provide(pre_process=True, post_process=True, vp_stage=None): + model = _original_provide(pre_process=pre_process, post_process=post_process, vp_stage=vp_stage) + if post_process: + model.output_layer = LinearForLastLayer( + input_size=model.config.hidden_size, output_size=1, config=model.config + ) + return model + + return _critic_provide + return provider.provide def model_provider(pre_process: bool = True, post_process: bool = True, vp_stage: int | None = None) -> GPTModel: diff --git a/slime_plugins/mbridge/glm4moe_lite.py b/slime_plugins/mbridge/glm4moe_lite.py index 8306d281ae..ceaa8d86e0 100644 --- a/slime_plugins/mbridge/glm4moe_lite.py +++ b/slime_plugins/mbridge/glm4moe_lite.py @@ -1,7 +1,76 @@ +import torch from mbridge.core import register_model +from mbridge.core.safetensor_io import SafeTensorIO from mbridge.models import DeepseekV3Bridge @register_model("glm4_moe_lite") class GLM4MoELiteBridge(DeepseekV3Bridge): - pass + """ + Bridge for GLM-4.7-Flash (glm4_moe_lite) models. + + Extends DeepseekV3Bridge with: + - Dynamic MTP layer indexing (parent hardcodes layer 61 for DeepSeek V3) + - Standard bf16 safetensor loading (parent uses FP8 dequant for DeepSeek V3) + """ + + def __init__(self, hf_config, **kwargs): + # Patch rope_theta: GLM-4.7-Flash stores it in rope_parameters dict, + # but DeepseekV3Bridge._build_config() expects hf_config.rope_theta directly. + if not hasattr(hf_config, "rope_theta"): + rope_params = getattr(hf_config, "rope_parameters", None) or {} + hf_config.rope_theta = rope_params.get("rope_theta", 1000000) + super().__init__(hf_config, **kwargs) + # Override the shared state dict mapping with dynamic layer index. + # DeepseekV3Bridge hardcodes layer 61; GLM-4.7-Flash uses num_hidden_layers (47). + n = hf_config.num_hidden_layers + if getattr(hf_config, "num_nextn_predict_layers", 0) and n: + self._SHARED_STATE_DICT_MAPPING = { + "embedding.word_embeddings.weight": [ + "model.embed_tokens.weight", + f"model.layers.{n}.embed_tokens.weight", + ], + "output_layer.weight": [ + "lm_head.weight", + f"model.layers.{n}.shared_head.head.weight", + ], + } + + def _get_safetensor_io(self, weights_path: str): + """Use standard SafeTensorIO — GLM-4.7-Flash ships bf16 safetensors, not FP8.""" + return SafeTensorIO(self._get_actual_hf_path(weights_path)) + + def _weight_to_hf_format( + self, mcore_weights_name: str, mcore_weights: torch.Tensor + ) -> tuple[list[str], list[torch.Tensor]]: + """Handle shared embedding/output weights for MTP with dynamic layer count.""" + if ( + self.config.mtp_num_layers is not None + and self.config.mtp_num_layers >= 1 + and mcore_weights_name in self._SHARED_STATE_DICT_MAPPING + ): + hf_names = self._SHARED_STATE_DICT_MAPPING[mcore_weights_name] + return hf_names, [mcore_weights] * len(hf_names) + # Skip DeepseekV3Bridge's _weight_to_hf_format (hardcoded 61) and go to Bridge base + return super(DeepseekV3Bridge, self)._weight_to_hf_format(mcore_weights_name, mcore_weights) + + def _convert_mtp_param(self, name: str) -> list[str]: + """Convert MTP parameter names with dynamic layer count (not hardcoded 61).""" + assert self.config.mtp_num_layers == 1, "only support one mtp layer for now" + n = self.config.num_layers + direct_name_mapping = { + "mtp.layers.0.enorm.weight": f"model.layers.{n}.enorm.weight", + "mtp.layers.0.hnorm.weight": f"model.layers.{n}.hnorm.weight", + "mtp.layers.0.eh_proj.weight": f"model.layers.{n}.eh_proj.weight", + "mtp.layers.0.final_layernorm.weight": f"model.layers.{n}.shared_head.norm.weight", + } + if name in direct_name_mapping: + return [direct_name_mapping[name]] + assert "mtp.layers.0.transformer_layer" in name, f"mtp not found in {name}" + proxy_name = name.replace("mtp.layers.0.transformer_layer", f"decoder.layers.{n}") + if "self_attention" in proxy_name or "input_layernorm.weight" in proxy_name: + return self._weight_name_mapping_attention(proxy_name) + elif "mlp" in proxy_name: + return self._weight_name_mapping_mlp(proxy_name) + else: + raise NotImplementedError(f"Unsupported MTP parameter name: {name}")