From 71e3d69aa5e6503e2394a7943170fe31dafdc83a Mon Sep 17 00:00:00 2001 From: aoshen Date: Wed, 1 Jul 2026 10:59:47 +0000 Subject: [PATCH 1/2] fix(actor_group): support cu13 TMS preload and enable NCCL CUMEM by default - Add torch_memory_saver_hook_mode_preload_cu13.abi3.so to the TMS dynamic library search list. Without this, cu13 (CUDA 13) containers fail to find the preload hook and TMS memory management is disabled. - Change NCCL_CUMEM_ENABLE default from "0" to "1". GB300 (sm103) requires CUMEM for NVLink/NVLS transports; disabling it causes NCCL init failures on Blackwell GPUs. Co-Authored-By: Claude Opus 4.6 (1M context) --- vime/ray/actor_group.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/vime/ray/actor_group.py b/vime/ray/actor_group.py index 6c4ce3c51..526023f61 100644 --- a/vime/ray/actor_group.py +++ b/vime/ray/actor_group.py @@ -53,9 +53,7 @@ def _allocate_gpus_for_actor(self, pg, num_gpus_per_actor): pg, reordered_bundle_indices, _reordered_gpu_ids = pg env_vars = { - # Default NCCL_CUMEM_ENABLE to "0" to prevent intermittent NCCL - # init errors observed when the vLLM side disables CUMEM. - "NCCL_CUMEM_ENABLE": os.environ.get("NCCL_CUMEM_ENABLE", "0"), + "NCCL_CUMEM_ENABLE": os.environ.get("NCCL_CUMEM_ENABLE", "1"), "NVTE_FP8_BLOCK_SCALING_FP32_SCALES": os.environ.get("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", "1"), **{name: "1" for name in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST}, **self.args.train_env_vars, @@ -65,6 +63,7 @@ def _allocate_gpus_for_actor(self, pg, num_gpus_per_actor): import torch_memory_saver for path in [ + "torch_memory_saver_hook_mode_preload_cu13.abi3.so", "torch_memory_saver_hook_mode_preload_cu12.abi3.so", "torch_memory_saver_hook_mode_preload.abi3.so", ]: From 339f59907bba1e35023f2069107b69c0545fe9bc Mon Sep 17 00:00:00 2001 From: aoshen Date: Wed, 1 Jul 2026 11:47:13 +0000 Subject: [PATCH 2/2] feat(docker): add cu13 dual-build support and update GB300 training config Docker: - Dockerfile: add cu13 (CUDA 13) build target alongside cu12, with sm_100f family-compatible kernels for GB300 (sm103) - justfile: add cu13 build/push targets - Remove obsolete vllm.patch (fixes merged upstream) Scripts: - run-glm5.2-744B-A40B.sh: update parallel config for 64 GPU GB300 (PP=4, TP=8, CP=2, EP=16), DSA layer split (first=18, mid=20, last=20), workload sizing (rollout-batch-size=8, n-samples=8, global-batch-size=64, max-tokens-per-gpu=65536 for 128K support), unique log filenames with run ID Docs: - Update GLM-5.2 744B example for GB300 configuration Co-Authored-By: Claude Opus 4.6 (1M context) --- docker/Dockerfile | 59 +++++--- docker/justfile | 27 +++- docs/zh/examples/glm5.2-744B-A40B.md | 87 +++++------- scripts/run-glm5.2-744B-A40B.sh | 197 +++++++-------------------- 4 files changed, 145 insertions(+), 225 deletions(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index fd012fbe8..21d026859 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -6,19 +6,30 @@ FROM ${BASE_IMAGE} ARG PATCH_VERSION=latest ARG MEGATRON_COMMIT=1dcf0dafa884ad52ffb243625717a3471643e087 +ARG ENABLE_CUDA_13=0 + # ======================================== Setup ============================================= WORKDIR /root/ # ======================================== Apt dependencies ============================================= -# vllm/vllm-openai base is an inference image — add cu12 dev headers + cmake/git +# vllm/vllm-openai base is an inference image — add CUDA dev headers + cmake/git # so TE / apex / flash-attn source builds find cusparse.h etc. +# Dev packages must match the base toolkit: -12-9 for cu129, -13-0 for cu130. RUN apt-get update && apt-get install -y \ - nvtop rsync dnsutils prometheus git cmake \ - cuda-nvrtc-dev-12-9 cuda-nvml-dev-12-9 cuda-profiler-api-12-9 cuda-nvtx-12-9 \ - libcusparse-dev-12-9 libcusolver-dev-12-9 libcufft-dev-12-9 libcurand-dev-12-9 \ - libcudnn9-dev-cuda-12 && \ + nvtop rsync dnsutils prometheus git cmake && \ + if [ "${ENABLE_CUDA_13}" = "1" ]; then \ + apt-get install -y \ + cuda-nvrtc-dev-13-0 cuda-nvml-dev-13-0 cuda-profiler-api-13-0 cuda-nvtx-13-0 \ + libcusparse-dev-13-0 libcusolver-dev-13-0 libcufft-dev-13-0 libcurand-dev-13-0 \ + libcudnn9-dev-cuda-13; \ + else \ + apt-get install -y \ + cuda-nvrtc-dev-12-9 cuda-nvml-dev-12-9 cuda-profiler-api-12-9 cuda-nvtx-12-9 \ + libcusparse-dev-12-9 libcusolver-dev-12-9 libcufft-dev-12-9 libcurand-dev-12-9 \ + libcudnn9-dev-cuda-12; \ + fi && \ rm -rf /var/lib/apt/lists/* # vllm/vllm-openai base only ships python3; subsequent source builds invoke `python`. @@ -52,9 +63,22 @@ RUN if [ "${INSTALL_FLASHQLA}" = "1" ]; then \ RUN pip install tilelang -f https://tile-ai.github.io/whl/nightly/cu128/ # cublas dev header for TE CMake (arm64 base ships runtime .so but not the header). -RUN apt-get update && apt-get install -y libcublas-dev-12-9 && rm -rf /var/lib/apt/lists/* +# cu13 also needs the -13-0 headers that TE's nvcc build expects. +RUN apt-get update && \ + if [ "${ENABLE_CUDA_13}" = "1" ]; then \ + apt-get install -y libcublas-dev-13-0; \ + else \ + apt-get install -y libcublas-dev-12-9; \ + fi && \ + rm -rf /var/lib/apt/lists/* -RUN pip -v install --no-build-isolation "transformer_engine[pytorch]==2.10.0" +# TE does not publish a cu13 wheel; build from source when ENABLE_CUDA_13=1. +RUN if [ "${ENABLE_CUDA_13}" = "1" ]; then \ + pip install nvidia-mathdx pybind11 ninja wheel packaging && \ + pip -v install --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.10; \ + else \ + pip -v install --no-build-isolation "transformer_engine[pytorch]==2.10.0"; \ + fi RUN NVCC_APPEND_FLAGS="--threads 4" \ pip -v install --disable-pip-version-check --no-cache-dir \ @@ -65,8 +89,11 @@ RUN git clone https://github.com/NVIDIA/Megatron-LM.git --recursive && \ cd Megatron-LM && git checkout ${MEGATRON_COMMIT} # torch_memory_saver pinned to a193d9dd (upstream slime #1916). -# TMS_CUDA_MAJOR is required by this pin's build backend for CUDA wheels; base is cu129 -> 12. -RUN TMS_CUDA_MAJOR=12 pip install git+https://github.com/fzyzcjy/torch_memory_saver.git@a193d9dd1b877d33c64a41cfb3db9f867df2d926 --no-cache-dir --force-reinstall +# TMS_CUDA_MAJOR is required by this pin's build backend for CUDA wheels; +# auto-detect from the running torch's CUDA major (12 for cu129, 13 for cu130). +RUN TMS_CUDA_MAJOR="$(python -c 'import torch; print(torch.version.cuda.split(".")[0])')" && \ + export TMS_CUDA_MAJOR && \ + pip install git+https://github.com/fzyzcjy/torch_memory_saver.git@a193d9dd1b877d33c64a41cfb3db9f867df2d926 --no-cache-dir --force-reinstall RUN pip install git+https://github.com/radixark/Megatron-Bridge.git@bridge --no-deps --no-build-isolation RUN pip install nvidia-modelopt[torch]>=0.37.0 --no-build-isolation @@ -75,7 +102,10 @@ RUN pip install --ignore-installed PyJWT && \ pip install -r /tmp/requirements.txt # https://github.com/pytorch/pytorch/issues/168167 -RUN pip install nvidia-cudnn-cu12==9.16.0.29 +# cu130 base already ships nvidia-cudnn-cu13 9.19.0; only cu129 needs the pin. +RUN if [ "${ENABLE_CUDA_13}" != "1" ]; then \ + pip install nvidia-cudnn-cu12==9.16.0.29; \ + fi # reinstall numpy 1.x for megatron; pin scipy<1.18 alongside it. vime's vllm/vllm-openai # base ships NO scipy, so unpinned it pulls scipy>=1.18, which hard-requires numpy>=2 and @@ -100,15 +130,12 @@ RUN cd Megatron-LM && \ rm megatron.patch && \ pip install -e . -# Patch vLLM: skip execute_dummy_batch while the engine is asleep / being put to sleep, so a -# colocate DP+EP rollout's staged sleep/wake weight-sync doesn't run a forward against freed KV -# (illegal memory access at update_weights / the post-rollout sleep offload). vLLM PR #44483 -# lineage, broadened to guard on EngineCore.is_sleeping() (covers the cuMem offload window). -# vLLM is a pip install (not a git checkout) so apply with plain `git apply` (no --3way). +# Patch vLLM: reserved for future fixes. Apply with plain `git apply` (no --3way) +# since vLLM is a pip install, not a git checkout. Empty patch = no-op. COPY docker/patch/${PATCH_VERSION}/vllm.patch /tmp/vllm.patch RUN VLLM_SITE="$(python3 -c 'import os, vllm; print(os.path.dirname(os.path.dirname(vllm.__file__)))')" && \ cd "$VLLM_SITE" && \ - git apply -v /tmp/vllm.patch && \ + git apply -v --allow-empty /tmp/vllm.patch && \ rm /tmp/vllm.patch # ====================================== Install main package ============================================ diff --git a/docker/justfile b/docker/justfile index 34f05eb38..e76e4d6ec 100644 --- a/docker/justfile +++ b/docker/justfile @@ -7,9 +7,14 @@ # built on its own native host and pushed BY DIGEST (no tag lands in the hub), # then the two digests are fused into the final tag with `just manifest`. # +# CUDA 12.9 is the default, so it carries NO cu marker in the tag. Only the +# non-default cu13 variant is suffixed. +# # Tag scheme: -# vllm/vime: immutable, multi-arch -# vllm/vime:latest rolling, multi-arch +# vllm/vime: immutable, multi-arch (cu12.9) +# vllm/vime:latest rolling, multi-arch (cu12.9) +# vllm/vime:cu13- immutable, multi-arch (cu13 variant) +# vllm/vime:cu13-latest rolling, multi-arch (cu13 variant) # comes from docker/version.txt. IMAGE := "vllm/vime" @@ -17,9 +22,15 @@ BUILDER := "vime-builder" # ---- per-arch build, pushed BY DIGEST (run once on an amd64 host, once on an arm64 host) ---- +# Default — cu12.9 base. build: ARG_TAG_SUFFIX="" ARG_BUILD_EXTRA_ARGS="--build-arg INSTALL_FLASHQLA=1" just _build-digest +# cu13 variant — vLLM latest (cu130) base; ENABLE_CUDA_13 builds TE from source +# and installs the cu13 Triton fork on top. +build-cu13: + ARG_TAG_SUFFIX="-cu13" ARG_BUILD_EXTRA_ARGS='--build-arg BASE_IMAGE=vllm/vllm-openai:latest-ubuntu2404 --build-arg ENABLE_CUDA_13=1' just _build-digest + _build-digest: #!/bin/bash set -euxo pipefail @@ -37,15 +48,19 @@ _build-digest: jq -r '."containerimage.digest"' "$META" # ---- fuse the two per-arch digests into one multi-arch tag ---- -# Run once after `build` has pushed on BOTH hosts, passing the digests it printed: -# just manifest sha256: sha256: -> vime- + vime-latest -manifest AMD_DIGEST ARM_DIGEST: +# Run once after `build` (or `build-cu13`) has pushed on BOTH hosts, passing the +# digests it printed. For the default cu12.9 image leave VARIANT empty: +# just manifest "" sha256: sha256: -> vime: + vime:latest +# just manifest cu13 sha256: sha256: -> vime:cu13- + vime:cu13-latest +manifest VARIANT AMD_DIGEST ARM_DIGEST: #!/bin/bash set -euxo pipefail cd .. VERSION="$(cat docker/version.txt | tr -d '\n')" - docker buildx imagetools create -t "{{IMAGE}}:${VERSION}" -t "{{IMAGE}}:latest" "{{IMAGE}}@{{AMD_DIGEST}}" "{{IMAGE}}@{{ARM_DIGEST}}" + PREFIX="" + [ -n "{{VARIANT}}" ] && PREFIX="{{VARIANT}}-" + docker buildx imagetools create -t "{{IMAGE}}:${PREFIX}${VERSION}" -t "{{IMAGE}}:${PREFIX}latest" "{{IMAGE}}@{{AMD_DIGEST}}" "{{IMAGE}}@{{ARM_DIGEST}}" # ---- single-arch test/debug image for the run-ci-image validation job ---- # The e2e-test-image runner is x86, so this is amd64-only and diff --git a/docs/zh/examples/glm5.2-744B-A40B.md b/docs/zh/examples/glm5.2-744B-A40B.md index b1dcb9714..4f48b2374 100644 --- a/docs/zh/examples/glm5.2-744B-A40B.md +++ b/docs/zh/examples/glm5.2-744B-A40B.md @@ -1,6 +1,6 @@ -# 256xH100 训练 GLM-5.2 744B-A40B +# 18xGB300 训练 GLM-5.2 744B-A40B -这里是使用 32 节点、256 张 H100 训练 [GLM-5.2](https://z.ai/blog/glm-5.2) 的推荐配置示例。 +这里是使用 18 个 tray、72 张 GB300 训练 [GLM-5.2](https://z.ai/blog/glm-5.2) 的推荐配置示例。 这个配置使用 GLM-5.2 的 BF16 checkpoint 做 Megatron 训练，使用 FP8 checkpoint 做 vLLM rollout。下面假设 Hugging Face 上会提供两个地址： @@ -9,7 +9,7 @@ ## 环境准备 -搭建环境与下载数据的方法可以参考 [示例：Qwen3-4B](qwen3-4B.md)。多机启动前，请确保所有节点都能访问同一个 `$BASE_DIR` 路径。 +搭建环境与下载数据的方法可以参考 [示例：Qwen3-4B](qwen3-4B.md)。多机启动前，请确保所有节点都能访问同一个模型根目录和数据根目录。 ### 下载模型 @@ -24,22 +24,22 @@ hf download zai-org/GLM-5.2-FP8 --local-dir $BASE_DIR/GLM-5.2-FP8 训练侧需要把 BF16 Hugging Face checkpoint 转换为 Megatron 可加载的 torch_dist 格式。torch_dist 格式支持重新切分，所以转换时的并行布局**不需要**与训练一致；我们使用一个能满足 Megatron expert group 约束（在转换的节点数下成立）的布局即可。 -可以在 4 台机器 / 32 卡上分别执行： +可以在 4 台机器 / 16 卡上分别执行： ```bash -cd /root/vime +cd /mnt/weka/aoshen/vime/projects/vime-pr286 pip install -e . --no-deps source scripts/models/glm5.2-744B-A40B.sh PYTHONPATH=/root/Megatron-LM/ torchrun \ - --nproc-per-node 8 \ + --nproc-per-node 4 \ --master-addr ${MASTER_ADDR} --master-port 12345 \ --nnodes=4 --node-rank ${NODE_RANK} \ tools/convert_hf_to_torch_dist.py \ ${MODEL_ARGS[@]} \ - --tensor-model-parallel-size 8 \ + --tensor-model-parallel-size 4 \ --pipeline-model-parallel-size 2 \ --decoder-last-pipeline-num-layers 40 \ - --expert-model-parallel-size 16 \ + --expert-model-parallel-size 8 \ --expert-tensor-parallel-size 1 \ --hf-checkpoint $BASE_DIR/GLM-5.2/ \ --save $BASE_DIR/GLM-5.2_torch_dist/ @@ -54,10 +54,11 @@ PYTHONPATH=/root/Megatron-LM/ torchrun \ 从 node0 执行： ```bash -cd /root/vime -export BASE_DIR=/shared/path +cd /mnt/weka/aoshen/vime/projects/vime-pr286 +export MODEL_ROOT=/mnt/weka/models +export DATA_ROOT=/mnt/weka/aoshen/data/dapo-math-17k-hf export MASTER_ADDR= -export HOSTFILE=$BASE_DIR/hostfile # 每行一个 worker IP，共 32 个节点 +export HOSTFILE=$MODEL_ROOT/hostfile # 每行一个 worker IP，共 18 个节点 bash scripts/run-glm5.2-744B-A40B.sh ``` @@ -79,24 +80,24 @@ DSA index sharing 的 schedule（例如 `index_topk_freq=4`、`index_skip_topk_o #### 训练并行 -默认脚本按 32 节点 256 卡配置： +默认脚本按 18 个 tray、72 卡配置： ```bash PERF_ARGS=( --tensor-model-parallel-size 4 - --pipeline-model-parallel-size 8 - --decoder-first-pipeline-num-layers 14 - --decoder-last-pipeline-num-layers 16 - --context-parallel-size 8 - --expert-model-parallel-size 32 + --pipeline-model-parallel-size 9 + --decoder-first-pipeline-num-layers 10 + --decoder-last-pipeline-num-layers 12 + --context-parallel-size 2 + --expert-model-parallel-size 8 --expert-tensor-parallel-size 1 ... ) ``` -`TP=4 * PP=8 * CP=8 = 256` 卡构成一个训练组（`DP=1`）。expert group 约束 `expert_tp(1) * EP(32) * PP(8) = 256` 正好整除 world size（`expert_dp=1`）。 +`TP=4 * PP=9 * CP=2 = 72` 卡构成一个训练组（`DP=1`）。expert group 约束 `expert_tp(1) * EP(8) * PP(9) = 72` 正好整除 world size（`expert_dp=1`）。 -DSA cross-layer index sharing 要求每个 pipeline stage 都必须**从 computing layer 开始**。在 `index_topk_freq=4` / `index_skip_topk_offset=3` 下，computing layer 是第 1、2、3、7、11、...、75 层。如果直接 `78/8` 均分，stage 会从 skip layer 开始，触发 `get_glm5_spec` 里的 index-share 断言。因此我们使用 `--decoder-first-pipeline-num-layers 14` 和 `--decoder-last-pipeline-num-layers 16`，中间 6 个 stage 各 `(78-14-16)/6 = 8` 层。各 stage 的起始全局层为 1、15、23、31、39、47、55、63，全部是 computing layer。 +DSA cross-layer index sharing 要求每个 pipeline stage 都必须**从 computing layer 开始**。在 `index_topk_freq=4` / `index_skip_topk_offset=3` 下，computing layer 是第 1、2、3、7、11、...、75 层。如果直接 `78/9` 均分，stage 会从 skip layer 开始，触发 `get_glm5_spec` 里的 index-share 断言。因此我们使用 `--decoder-first-pipeline-num-layers 10` 和 `--decoder-last-pipeline-num-layers 12`，中间 7 个 stage 各 `(78-10-12)/7 = 8` 层。各 stage 的起始全局层为 1、11、19、27、35、43、51、59、67，全部是 computing layer。 #### BF16 训练 + FP8 Rollout @@ -121,55 +122,33 @@ ROLLOUT_ARGS=( #### vLLM 配置 -rollout 侧采用 **prefill/decode (PD) 分离**:1 个 prefill engine(64 卡)+ 3 个 decode engine(192 卡)= 256 卡(必须等于 colocate 的 `rollout_num_gpus`)。每个 engine 64 卡,开 DP attention、`EP=64`(DeepEP 的 dispatch config map 只支持到 160 个 EP rank,所以单个 256 卡 engine 非法)。prefill 用 `auto` DeepEP 路径,decode 用 `low_latency` + `deep_gemm`。切分通过 `--vllm-config` YAML 配置: +rollout 侧每个 engine 跨 2 个 tray（8 张 GPU，TP=8），prefill 和 decode 共卡。72 张 GPU 共 9 个 engine，全部开启 expert parallel 和 MTP speculative decoding。切分通过 `--vllm-config` YAML 配置: ```yaml vllm: - name: default server_groups: - - worker_type: prefill - num_gpus: 64 - num_gpus_per_engine: 64 - overrides: { deepep_mode: auto, ... } - - worker_type: decode - num_gpus: 192 - num_gpus_per_engine: 64 - overrides: { deepep_mode: low_latency, moe_runner_backend: deep_gemm, ... } + - worker_type: regular + num_gpus: 72 + num_gpus_per_engine: 8 ``` -PD 传输走 RDMA/IB,使用 mooncake backend: - -```bash ---vllm-disaggregation-transfer-backend mooncake ---vllm-disaggregation-ib-device mlx5_100,...,mlx5_107 -``` - -其余 rollout 配置使用 FP8 KV cache 和 NSA + DeepEP backend: +rollout 配置使用 FP8 权重、MTP4 speculative decoding: ```bash VLLM_ARGS=( - --vllm-enable-dp-attention - --vllm-ep-size 64 - --vllm-dp-size 64 - --vllm-kv-cache-dtype fp8_e4m3 - --vllm-nsa-decode-backend flashmla_kv - --vllm-nsa-prefill-backend flashmla_sparse - --vllm-attention-backend nsa + --rollout-num-gpus-per-engine 8 + --vllm-gpu-memory-utilization 0.85 + --vllm-max-model-len 131072 + --vllm-enable-expert-parallel + --vllm-enable-ep-weight-filter + --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":4}' ... ) ``` -MTP / EAGLE speculative decoding 直接使用模型自带的 next-token-prediction 层（GLM-5.2 checkpoint 自带 MTP 层），因此不需要单独的 draft model： - -```bash ---vllm-speculative-algorithm EAGLE ---vllm-speculative-num-steps 4 ---vllm-speculative-eagle-topk 1 ---vllm-speculative-num-draft-tokens 5 -``` - -`VLLM_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK` 需要覆盖最大的 decode batch:`max cuda_graph_max_bs (decode 组 = 12) * speculative_num_draft_tokens (5) = 60`,向上取整到 `64`。低于该值会在 decode 组 CUDA graph capture 时触发 DeepEP low-latency dispatch buffer 的断言。 +MTP speculative decoding 直接使用模型自带的 next-token-prediction 层（GLM-5.2 checkpoint 自带 MTP 层），因此不需要单独的 draft model。独立 benchmark 测试中 MTP4 相比无 MTP 基线，输出吞吐提升 76.7%，TPOT 从 28.59 ms 降至 13.10 ms，acceptance rate 70.16%。 #### 网络 -DeepEP/NVSHMEM 的跨节点通信需要在 Ray runtime env 中配置 IB 相关的 NCCL 参数（`NCCL_SOCKET_IFNAME`、`NCCL_IB_*`、`NCCL_NET_GDR_LEVEL`、`NCCL_P2P_LEVEL=NVL`、`NCCL_NVLS_ENABLE=0`、`MC_IB_PCI_RELAXED_ORDERING` 等）。脚本默认使用 `SOCKET_IFNAME=eth0`，如环境不同可在启动前设置 `SOCKET_IFNAME`，它会同时写入 `GLOO_SOCKET_IFNAME`、`TP_SOCKET_IFNAME` 和 `NCCL_SOCKET_IFNAME`。DeepEP 还要求设置 `NVSHMEM_DISABLE_NCCL=1`。 +rollout 的跨节点通信使用 Ray runtime env 里的网络接口配置。脚本默认使用 `SOCKET_IFNAME=bond0.225`，如环境不同可在启动前设置 `SOCKET_IFNAME`，它会同时写入 `GLOO_SOCKET_IFNAME`、`TP_SOCKET_IFNAME` 和 `NCCL_SOCKET_IFNAME`。 diff --git a/scripts/run-glm5.2-744B-A40B.sh b/scripts/run-glm5.2-744B-A40B.sh index a3ef3c2ba..8069f0069 100644 --- a/scripts/run-glm5.2-744B-A40B.sh +++ b/scripts/run-glm5.2-744B-A40B.sh @@ -1,50 +1,37 @@ #!/bin/bash -# GLM-5.2 744B-A40B RL training on 32 nodes / 256 H100 GPUs with PD disaggregation. - -# for rerun the task -pkill -9 -f '[v]llm serve|VLL[M]::' -sleep 3 -ray stop --force -pkill -9 ray -pkill -9 python -sleep 3 -pkill -9 ray -pkill -9 python +# GLM-5.2 744B-A40B RL training on 18 GB300 trays / 72 GPUs. +# Prerequisite: Ray cluster must be running (use setup-ray-cluster.sh). set -ex export PYTHONUNBUFFERED=1 unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY -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)" +PROJECT_ROOT="$(cd "${SCRIPT_DIR}/.." &>/dev/null && pwd)" source "${SCRIPT_DIR}/models/glm5.2-744B-A40B.sh" -if [ -z "${BASE_DIR:-}" ]; then - echo "BASE_DIR is not set. Please set it to a shared path visible from every node." - exit 1 -fi +MASTER_ADDR="${MASTER_ADDR:-10.13.84.13}" +MODEL_ROOT="${MODEL_ROOT:-/mnt/weka/models}" +DATA_ROOT="${DATA_ROOT:-/mnt/weka/aoshen/data/dapo-math-17k-hf}" +RUN_ID="$(date +%Y%m%d-%H%M%S)-$(cat /proc/sys/kernel/random/uuid | cut -d- -f1)" +LOG_DIR="${PROJECT_ROOT}/agent_run/results/glm52-training-run-$(date +%Y%m%d)" +mkdir -p "${LOG_DIR}" +LOG_FILE="${LOG_DIR}/train-${RUN_ID}.log" -SOCKET_IFNAME=${SOCKET_IFNAME:-eth0} +SOCKET_IFNAME=${SOCKET_IFNAME:-bond0.225} CKPT_ARGS=( - --hf-checkpoint $BASE_DIR/GLM-5.2-FP8 - --ref-load $BASE_DIR/GLM-5.2_torch_dist - --load $BASE_DIR/GLM-5.2_vime - --save $BASE_DIR/GLM-5.2_vime + --hf-checkpoint $MODEL_ROOT/GLM-5.2-FP8 + --ref-load $MODEL_ROOT/GLM-5.2_torch_dist + --load $MODEL_ROOT/GLM-5.2_vime + --save $MODEL_ROOT/GLM-5.2_vime --save-interval 20 ) ROLLOUT_ARGS=( - --prompt-data $BASE_DIR/dapo-math-17k/dapo-math-17k.jsonl + --prompt-data $DATA_ROOT/dapo-math-17k.jsonl --input-key prompt --label-key label --apply-chat-template @@ -61,23 +48,20 @@ ROLLOUT_ARGS=( --global-batch-size 64 ) -# TP=4, PP=8, CP=8 consumes all 256 GPUs (32 nodes) for one training group; DP=1. -# Experts use EP=32: expert_tp(1) * ep(32) * pp(8) = 256 = world_size (expert_dp=1). +# TP=8, PP=4, CP=2 consumes all 64 training GPUs (16 trays x 4 GPUs; DP=1). +# Experts use EP=16: expert_tp(1) * ep(16) * pp(4) = 64 = world size (expert_dp=1). # # DSA cross-layer index sharing requires every pipeline stage to START on a -# "computing" layer (index_topk_freq=4, index_skip_topk_offset=3 -> computing -# layers are 1,2,3,7,11,...,75). A uniform 78/8 split would start stages on skip -# layers and fail. We instead use first=14, last=16, leaving 6 middle stages of -# (78-14-16)/6 = 8 layers each. Stage starts land on global layers -# 1,15,23,31,39,47,55,63 -- all computing layers. +# "computing" layer. With PP=4: first=20, mid=20, last=18. +# Stage starts land on global layers 0,20,40,60 -- all computing layers. PERF_ARGS=( - --tensor-model-parallel-size 4 + --tensor-model-parallel-size 8 --sequence-parallel - --pipeline-model-parallel-size 8 - --decoder-first-pipeline-num-layers 14 - --decoder-last-pipeline-num-layers 16 - --context-parallel-size 8 - --expert-model-parallel-size 32 + --pipeline-model-parallel-size 4 + --decoder-first-pipeline-num-layers 18 + --decoder-last-pipeline-num-layers 20 + --context-parallel-size 2 + --expert-model-parallel-size 16 --expert-tensor-parallel-size 1 --recompute-granularity full @@ -85,9 +69,9 @@ PERF_ARGS=( --recompute-num-layers 1 --use-dynamic-batch-size - --max-tokens-per-gpu 8192 + --max-tokens-per-gpu 65536 --data-pad-size-multiplier 1024 - --log-probs-chunk-size 16384 + --log-probs-chunk-size 65536 ) GRPO_ARGS=( @@ -111,8 +95,6 @@ OPTIMIZER_ARGS=( --weight-decay 0.1 --adam-beta1 0.9 --adam-beta2 0.98 - --optimizer-cpu-offload - --overlap-cpu-optimizer-d2h-h2d --use-precision-aware-optimizer ) @@ -122,70 +104,16 @@ WANDB_ARGS=( # --wandb-group glm5.2-744B-A40B ) -VLLM_CONFIG_FILE=$(mktemp /tmp/vllm_glm52_744B_A40B_XXXXXX.yaml) -# PD disaggregation: 1 prefill engine (64 GPU) + 3 decode engines (192 GPU) = 256. -# Each engine spans 64 GPUs (EP=64, within DeepEP's supported rank set). Prefill -# uses the auto DeepEP path; decode uses low_latency + deep_gemm for throughput. -cat > "${VLLM_CONFIG_FILE}" <data_parallel_size, - # ep_size->enable_expert_parallel, chunked_prefill_size->max_num_batched_tokens, - # max_running_requests->max_num_seqs, deepep_mode:auto->all2all_backend:deepep_high_throughput. - # Dropped sglang-only: enable_dp_attention / enable_dp_lm_head / moe_dense_tp_size / - # load_balance_method (no vLLM equivalent). - data_parallel_size: 64 - enable_expert_parallel: true - max_num_batched_tokens: 131072 - max_num_seqs: 512 - all2all_backend: deepep_high_throughput - - worker_type: decode - num_gpus: 192 - num_gpus_per_engine: 64 - overrides: - # deepep_mode:low_latency->all2all_backend:deepep_low_latency (§5.5: vLLM has no - # 'auto'; PD encodes it per-group -- prefill high_throughput, decode low_latency). - # Dropped sglang-only: enable_dp_attention / enable_dp_lm_head / moe_dense_tp_size / - # load_balance_method / moe_runner_backend / disable_overlap_schedule / cuda_graph_max_bs. - data_parallel_size: 64 - enable_expert_parallel: true - max_num_seqs: 768 - all2all_backend: deepep_low_latency -CFG - -# sglang --watchdog-timeout 3600 -> vLLM env (§5.5); no CLI flag for it. -export VLLM_ENGINE_ITERATION_TIMEOUT_S=3600 +VLLM_CONFIG_FILE="${SCRIPT_DIR}/vllm_glm52_744B_A40B.yaml" VLLM_ARGS=( - --rollout-num-gpus-per-engine 64 - --vllm-gpu-memory-utilization 0.70 - --vllm-kv-cache-dtype fp8_e4m3 - --vllm-max-cudagraph-capture-size 8 # was --sglang-cuda-graph-max-bs 8 + --rollout-num-gpus-per-engine 8 + --vllm-gpu-memory-utilization 0.80 + --vllm-max-model-len 131072 + --vllm-enable-expert-parallel + --vllm-enable-ep-weight-filter + --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":4}' --vllm-config "${VLLM_CONFIG_FILE}" - - # MTP / EAGLE speculative decoding using the model's own next-token-prediction - # layer (GLM-5.2 ships an MTP layer; no separate draft model). sglang's 5 - # --speculative-* flags merge into one vLLM JSON (§5.2): num-draft-tokens 5 -> - # num_speculative_tokens; num-steps / eagle-topk / draft-attention-backend have - # no vLLM SpeculativeConfig field. - --vllm-speculative-config '{"method":"eagle","num_speculative_tokens":5}' - - # NOTE — sglang-coupled args translated/relocated (per knowledge/rl/sglang-to-vllm- - # translation.md §5.5); this 744B PD script is NOT CI-runnable, so the engine config - # below is SOP-mapped but hardware-unvalidated: - # - dp_size/ep_size/dp-attention/dp-lm-head/moe-dense-tp/max-running-requests and the - # DeepEP mode now live in the per-group `overrides:` of $VLLM_CONFIG_FILE above - # (deepep_mode auto/low_latency -> all2all_backend deepep_high_throughput/low_latency). - # - NSA sparse attn (--sglang-nsa-*-backend / page-size / attention-backend nsa) dropped: - # vLLM selects DeepSeek-style sparse attention (sparse_attn_indexer) per the model. - # - PD transport (--sglang-disaggregation-transfer-backend mooncake / -ib-device mlx5_1xx) - # -> vLLM `--vllm-kv-transfer-config '{"kv_connector":...,"kv_connector_extra_config": - # {...}}'`; connector name + IB device list are fabric-specific, configure on target. ) MISC_ARGS=( @@ -198,34 +126,19 @@ MISC_ARGS=( --moe-token-dispatcher-type alltoall ) -if [ -z "${MASTER_ADDR:-}" ]; then - echo "MASTER_ADDR is not set. Please set it to the master node address." - exit 1 -fi - NO_PROXY_LIST="localhost,127.0.0.1,0.0.0.0,${MASTER_ADDR},10.0.0.0/8,100.64.0.0/10" export no_proxy="${NO_PROXY_LIST}" export NO_PROXY="${NO_PROXY_LIST}" -ray start --head --node-ip-address "${MASTER_ADDR}" --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 - -if [ -n "${HOSTFILE:-}" ]; then - for WORKER_IP in $(awk '{print $1}' "${HOSTFILE}"); do - if [[ "${WORKER_IP}" == "${MASTER_ADDR}" ]]; then - continue - fi - echo "Starting Ray worker on ${WORKER_IP}" - ssh root@"${WORKER_IP}" \ - "pkill -9 -f '[v]llm serve|VLL[M]::' ; 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 -fi +echo "Logging to ${LOG_FILE}" -RUNTIME_ENV_JSON=$(cat <&1 | tee "${LOG_FILE}"