From 37a86412604a0c46acc5e12d48cc9d1d82ba6cfe Mon Sep 17 00:00:00 2001 From: FlamingoPg <1106310035@qq.com> Date: Mon, 29 Jun 2026 18:44:27 +0800 Subject: [PATCH] feat: add GLM-5.2 DSpark draft training configs and deploy plan MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add support artifacts for training a DSpark speculative-decoding draft against GLM-5.2 (744B/40B MoE, glm_moe_dsa) on the TorchSpec online path. - torchspec/config/glm52_dspark_draft_config.json: DSpark draft config sized to GLM-5.2 (hidden 6144, vocab 154880, 78-layer target, aux layers [1,20,38,56,75]); backbone dims aligned with the existing GLM-5.2 DFlash draft config. - configs/sglang_glm52_dspark_8card_colocate.yaml: single-node 8xB200, FP8, colocate (inference TP=8 + FSDP training share the 8 GPUs). - configs/sglang_glm52_dspark_2node.yaml: 16xB200 / 2-node fallback. - torchspec/data/template.py: register the `glm` chat template. - docs/deploy_glm52_dspark_8card.md: end-to-end deploy plan with Go/No-Go. NOT YET VERIFIED — before merge/training, confirm on a GLM-5.2 checkpoint: mask_token_id (placeholder 154820), the glm chat_template, and the embed/lm_head/norm state-dict keys; then run the Step A smoke + dry-run on B200. GLM-5.2 inherits EAGLE3 aux-hidden capture from DeepseekV2ForCausalLM, so no GLM-specific sglang patch is required. Co-Authored-By: Claude Opus 4.8 (1M context) --- configs/sglang_glm52_dspark_2node.yaml | 107 +++++++++++ .../sglang_glm52_dspark_8card_colocate.yaml | 108 ++++++++++++ docs/deploy_glm52_dspark_8card.md | 166 ++++++++++++++++++ .../config/glm52_dspark_draft_config.json | 23 +++ torchspec/data/template.py | 16 ++ 5 files changed, 420 insertions(+) create mode 100644 configs/sglang_glm52_dspark_2node.yaml create mode 100644 configs/sglang_glm52_dspark_8card_colocate.yaml create mode 100644 docs/deploy_glm52_dspark_8card.md create mode 100644 torchspec/config/glm52_dspark_draft_config.json diff --git a/configs/sglang_glm52_dspark_2node.yaml b/configs/sglang_glm52_dspark_2node.yaml new file mode 100644 index 0000000..e003905 --- /dev/null +++ b/configs/sglang_glm52_dspark_2node.yaml @@ -0,0 +1,107 @@ +# DSpark training config for GLM-5.2 (744B/40B MoE, glm_moe_dsa) — 2 node (B200) +# +# DSpark = DFlash block-diffusion drafter + EAGLE-style Markov & confidence +# heads. Shares the DFlash backbone, so the draft dims here MUST match +# torchspec/config/dflash_draft_config_glm52.json (the GLM-5.2 DFlash sibling). +# The L1 / confidence terms need the target's final hidden state, so DSpark sets +# inference.store_last_hidden_states: true (DFlash leaves it false). +# +# GPU allocation (16x B200, 2 nodes): +# - Node 0: 8 GPUs for INFERENCE (SGLang engine, tp_size=8, GLM-5.2 FP8) +# - Node 1: 8 GPUs for TRAINING (FSDP FULL_SHARD, DSpark draft) +# Mirrors configs/sglang_kimi_k25_2node.yaml (Kimi-K2.5 is the closest ~1T MoE). +# +# SERVING NOTE (B200 vs the L20 DFlash work in sglang_glm52_dflash.yaml): +# glm_moe_dsa's DSA kernels are gated to sm90/sm100. B200 IS sm100, so the +# gating is SATISFIED — stock SGLang DSA kernels run and the ~32x L20 ada_dsa +# community port is NOT needed here (that was only for sm89/L20/4090). +# STILL REQUIRED: layer TorchSpec's hidden-capture patch onto GLM-5.2's SGLang +# serving model (deepseek_v2.py-style, since glm_moe_dsa is DeepSeek-like) so +# the 5 aux layers + last hidden stream to Mooncake. This is the gating Step 0. +# +# BEFORE REAL TRAINING — verify these target-specific values (cluster-side): +# 1. mask_token_id in the draft config is the placeholder 154820 (= pad). Replace +# with the real GLM [MASK] id (154820-154879 range): tokenizer.convert_tokens_to_ids('[MASK]'). +# 2. dataset.chat_template 'glm': confirm headers / end_of_turn / thinking vs +# the GLM-5.2 tokenizer_config.json chat_template. +# 3. embedding_key/lm_head_key/norm_key defaults (model.embed_tokens.weight / +# lm_head.weight / model.norm.weight) match the GLM-5.2 state_dict keys. +# +# Usage (2-node ray, see examples/kimi-k25-2node-h200/): +# python -m torchspec.train_entry --config configs/sglang_glm52_dspark_2node.yaml + +model: + target_model_path: zai-org/GLM-5.2-FP8 # FP8 weights so TP=8 fits one node + trust_remote_code: true # glm_moe_dsa is a custom model_type + draft_model_config: torchspec/config/glm52_dspark_draft_config.json + +dataset: + train_data_path: ../examples/data/sample_conversations.jsonl + eval_data_path: null + eval_interval: 100 + chat_template: glm + prompt_key: conversations + min_loss_tokens: 32 + +training: + attention_backend: flex_attention + micro_batch_size: 1 + draft_accumulation_steps: 2 + learning_rate: 6e-4 + min_lr: 6e-5 + weight_decay: 0.0 + max_concurrent_batches: 1 + max_grad_norm: 1.0 + max_seq_length: 4096 + num_epochs: 3 + seed: 42 + training_num_gpus_per_node: 8 + training_num_nodes: 1 + ttt_length: 7 + fsdp_strategy: FULL_SHARD + fsdp_reduce_dtype: bfloat16 + prefetch_depth: 8 + save_interval: 1000 + save_per_epoch: true + max_checkpoints: 2 + warmup_ratio: 0.04 + + # DSpark-specific parameters + dflash_block_size: 7 + dspark_num_anchors: 512 + dspark_num_target_layers: 5 + dspark_loss_decay_gamma: 4.0 + dspark_ce_loss_alpha: 0.1 + dspark_l1_loss_alpha: 0.9 + dspark_confidence_head_alpha: 1.0 + +inference: + inference_engine_type: sgl + store_last_hidden_states: true # DSpark L1/confidence loss needs last hidden + inference_num_gpus: 8 + inference_num_gpus_per_engine: 8 # tp_size=8 → one engine spans a full node + inference_num_gpus_per_node: 8 + max_sample_pool_size: 64 + inference_buffer_threshold: 32 + inference_batch_size: 8 + sglang: + tp_size: 8 + mem_fraction_static: 0.85 + dtype: auto # FP8 follows the checkpoint weights + +mooncake: + master_server_address: null + metadata_server: null + protocol: tcp + global_segment_size: 16GB + local_buffer_size: 4GB + enable_hard_pin: true + +output_dir: ./outputs/glm52-dspark +cache_dir: ./cache/glm52-dspark +model_download_dir: null + +debug: + save_debug_train_data: null + debug_train_only: false + debug_inference_only: false diff --git a/configs/sglang_glm52_dspark_8card_colocate.yaml b/configs/sglang_glm52_dspark_8card_colocate.yaml new file mode 100644 index 0000000..294864f --- /dev/null +++ b/configs/sglang_glm52_dspark_8card_colocate.yaml @@ -0,0 +1,108 @@ +# DSpark training for GLM-5.2 (744B/40B MoE) — 8x B200, SINGLE NODE, FP8, COLOCATE +# +# Fits FP8 GLM-5.2 (needs TP=8 = a full node for inference) + DSpark training on +# ONE 8-GPU node by COLOCATING both on the same 8 GPUs (training.colocate=true). +# placement_group.py:322-331 -> expected GPUs = max(infer, train) = max(8,8) = 8. +# Inference (SGLang TP=8) and training (FSDP) share each B200's 192GB and +# time-slice compute. This is why B200's large VRAM makes it viable; it would NOT +# fit on H100/H200 (80/141GB). colocate is documented as a "Dev" mode +# (docs/ray.md:40) — experimental, expect to tune memory + watch for NCCL/compute +# contention between the SGLang TP group and the FSDP group. +# +# *** CRITICAL KNOB: inference.sglang.mem_fraction_static *** +# FP8 weights take ~93GB/GPU. mem_fraction_static caps SGLang's (weights+KV) use. +# Set it LOW (~0.5) so the co-located FSDP training has room on the same GPUs. +# Default 0.85 WILL OOM the trainer. Budget (per 192GB B200): +# inference ~0.5*192 = 96GB (93 weights + ~3 KV) +# training ~remaining 96GB (FSDP draft shards + embed/lm_head + activations) +# If the trainer OOMs, lower mem_fraction_static further (0.45) and/or reduce +# training.max_seq_length; if KV is too tight for inference, raise it slightly. +# +# Throughput note: colocation time-slices compute, so this is SLOWER than the +# 2-node (8 infer + 8 train) layout in sglang_glm52_dspark_2node.yaml. Use this +# only when you are constrained to a single 8-GPU node. +# +# BEFORE REAL TRAINING — verify (see sglang_glm52_dspark_2node.yaml): +# mask_token_id (placeholder 154820), glm chat_template, embed/lm_head/norm keys, +# and Step 0 capture smoke (GLM inherits DeepseekV2 EAGLE3 capture). +# +# Run (single node; pin GPUs / force a fresh local Ray if the box is shared): +# RAY_ADDRESS=local CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ +# python -m torchspec.train_entry --config configs/sglang_glm52_dspark_8card_colocate.yaml + +model: + target_model_path: zai-org/GLM-5.2-FP8 + trust_remote_code: true + draft_model_config: torchspec/config/glm52_dspark_draft_config.json + +dataset: + train_data_path: ../examples/data/sample_conversations.jsonl + eval_data_path: null + eval_interval: 100 + chat_template: glm + prompt_key: conversations + min_loss_tokens: 32 + +training: + colocate: true # <-- share the 8 GPUs between inference & training + attention_backend: flex_attention + micro_batch_size: 1 + draft_accumulation_steps: 2 + learning_rate: 6e-4 + min_lr: 6e-5 + weight_decay: 0.0 + max_concurrent_batches: 1 + max_grad_norm: 1.0 + max_seq_length: 4096 + num_epochs: 3 + seed: 42 + training_num_gpus_per_node: 8 # same 8 GPUs as inference (colocate) + training_num_nodes: 1 + ttt_length: 7 + fsdp_strategy: FULL_SHARD + fsdp_reduce_dtype: bfloat16 + prefetch_depth: 8 + save_interval: 1000 + save_per_epoch: true + max_checkpoints: 2 + warmup_ratio: 0.04 + + # DSpark-specific parameters + dflash_block_size: 7 + dspark_num_anchors: 512 + dspark_num_target_layers: 5 + dspark_loss_decay_gamma: 4.0 + dspark_ce_loss_alpha: 0.1 + dspark_l1_loss_alpha: 0.9 + dspark_confidence_head_alpha: 1.0 + +inference: + inference_engine_type: sgl + store_last_hidden_states: true # DSpark L1/confidence loss needs last hidden + inference_num_gpus: 8 + inference_num_gpus_per_engine: 8 # tp_size=8 spans the whole node + inference_num_gpus_per_node: 8 + max_sample_pool_size: 64 + inference_buffer_threshold: 32 + inference_batch_size: 8 + sglang: + tp_size: 8 + mem_fraction_static: 0.5 # <-- LOW: leave ~half of each B200 for training + dtype: auto # FP8 follows the checkpoint weights + +mooncake: + master_server_address: null + metadata_server: null + protocol: tcp + global_segment_size: 16GB + local_buffer_size: 4GB + enable_hard_pin: true + +output_dir: ./outputs/glm52-dspark-8card +cache_dir: ./cache/glm52-dspark-8card +model_download_dir: null + +debug: + save_debug_train_data: null + debug_train_only: false + debug_inference_only: false diff --git a/docs/deploy_glm52_dspark_8card.md b/docs/deploy_glm52_dspark_8card.md new file mode 100644 index 0000000..337ea9d --- /dev/null +++ b/docs/deploy_glm52_dspark_8card.md @@ -0,0 +1,166 @@ +# 部署方案:GLM-5.2 DSpark draft 训练(8×B200 单节点 / FP8 / colocate) + +> 目标:在**一台 8×B200(sm100)** 机器上,用 **FP8** 的 GLM-5.2 作为 target,训练一个 **DSpark** speculative-decoding draft,加速 GLM-5.2 推理。 +> 关键手法:`colocate` —— 推理(SGLang TP=8)与训练(FSDP)**共用同一组 8 卡**,而不是各占一组。 + +## 0. 拓扑与原理 + +| 项 | 值 | 依据 | +|---|---|---| +| 节点 | 1 × (8×B200, 192GB/卡) | — | +| target 推理 | SGLang, FP8, **TP=8**(占满 8 卡) | 744GB/8 = 93GB/卡 | +| draft 训练 | FSDP FULL_SHARD, **同 8 卡** | DSpark draft ~4B,分片后每卡 ~6GB | +| 共卡机制 | `training.colocate=true` → GPU 数 = `max(8,8)=8` | `placement_group.py:322-331` | +| 显存预算 | 推理 ~96GB + 训练 ~96GB < 192GB | `mem_fraction_static=0.5` | +| 为什么必须 B200 | H100/H200(80/141GB)装不下「FP8 权重 + 同卡训练」 | — | + +config 文件:**`configs/sglang_glm52_dspark_8card_colocate.yaml`**(已就位)。 + +--- + +## 1. 环境搭建 + +```bash +cd # 本项目分支 proj/glm52-torchspec-training + +# 一键建环境 + 构建 patched sglang(v0.5.14 @ commit 49e384ce)+ 装 torchspec +./tools/build_conda.sh # 新建 torchspec env,装 sglang 后端 +# 或装进当前环境: ./tools/build_conda.sh current sglang +# 可选 Flash-Attention: pip install -e ".[fa]" + +# 激活 +micromamba activate torchspec # 或 conda activate torchspec +``` + +patched sglang 校验(可选,确认 patch 应用成功): +```bash +python tools/test_sglang_engine_patch.py # 仓库自带的 patch 自检 +``` + +> patched sglang 的 base commit 锁在 `docker/sglang/v0.5.14/SGLANG_COMMIT`;手动应用 patch 用 `tools/apply_sglang_patch.sh `。也可直接用 `docker/sglang/v0.5.14/Dockerfile` 起容器。 + +--- + +## 2. 准备:权重 + 核验门槛值 + 数据 + +### 2.1 GLM-5.2 FP8 权重 +```bash +# HF 下载(或用你已有的镜像/路径) +huggingface-cli download zai-org/GLM-5.2-FP8 --local-dir /models/GLM-5.2-FP8 +``` +把 config 里的 `model.target_model_path` 指到这个路径。 + +### 2.2 核验三个 target-相关值(否则静默出错) +```bash +python - <<'PY' +from transformers import AutoTokenizer +import torch, glob, safetensors.torch as st +tok = AutoTokenizer.from_pretrained("/models/GLM-5.2-FP8", trust_remote_code=True) +print("mask_token_id [MASK] =", tok.convert_tokens_to_ids("[MASK]")) # 填进 draft config +print("chat_template head =", (tok.chat_template or "")[:200]) # 对照 glm template +# 抽一个分片看 state_dict key 命名是否 = model.embed_tokens.weight / lm_head.weight / model.norm.weight +f = sorted(glob.glob("/models/GLM-5.2-FP8/*.safetensors"))[0] +ks = list(st.load_file(f).keys()) +print("keys sample:", [k for k in ks if any(s in k for s in ("embed_tokens","lm_head","model.norm"))][:5]) +PY +``` +- `mask_token_id`:把 `torchspec/config/glm52_dspark_draft_config.json` 里的占位 **154820** 换成真值。 +- `chat_template`:核对 `torchspec/data/template.py` 的 `glm` template(header / end token / 是否 thinking)。 +- `embedding_key/lm_head_key/norm_key`:若与默认(`model.embed_tokens.weight` / `lm_head.weight` / `model.norm.weight`)不同,在 config 的 `model:` 段覆盖。 + +### 2.3 训练数据 +- 格式参考 `examples/data/sample_conversations.jsonl`(`prompt_key: conversations`)。 +- 自有数据用 `tools/generate_data.py` 生成 / 转换。 +- 把 `dataset.train_data_path` 指过去。约束:`min_loss_tokens ≥ 2×dflash_block_size`(7→≥14,config 里已是 32 ✓)。 + +--- + +## 3. Step A — 冒烟(门槛,务必先过) + +确认 sglang 能在 B200 上 serve `GlmMoeDsaForCausalLM`(FP8/TP=8)且 capture 通路初始化(GLM 继承 `DeepseekV2ForCausalLM` 的 EAGLE3 capture,理论上零改动)。 + +```bash +# 用 colocate config 但只起推理:临时把 debug.debug_inference_only 设为 true 跑一小步 +RAY_ADDRESS=local CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ + python -m torchspec.train_entry --config configs/sglang_glm52_dspark_8card_colocate.yaml \ + --opts debug.debug_inference_only=true +``` +**看日志确认**(`sgl_engine.py` 启动时会打): +- `SglEngine ... initialized ... (tp_size=8, aux_layers=[...], hidden_size=6144)` → capture 层解析成功。 +- **没有 DSA kernel 崩溃**(sm100 应满足 gating;若崩 → 你的 sglang 太旧或 B200 驱动问题)。 +> `--opts key=value` 覆盖单个字段(见 train_entry 的 `_validate_*` 用法约定);若你的版本不支持 `--opts`,直接复制一份 config 改 `debug_inference_only: true`。 + +--- + +## 4. Step B — 小样本 dry-run + +确认「推理→Mooncake→draft 前向→loss 下降」全链路通 + 显存不 OOM。临时改一份 config: +- `dataset.train_data_path` 指向几十条样本 +- `training.max_seq_length: 1024`、`training.num_epochs: 1` +- 其余保持 colocate / FP8 / 8 卡 + +```bash +RAY_ADDRESS=local CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ + ./examples/qwen3-8b-single-node/run.sh configs/sglang_glm52_dspark_8card_colocate.yaml +``` +盯:每卡显存(`nvidia-smi`)、loss 是否下降、有无 NCCL 报错。**OOM 就降 `mem_fraction_static`(0.5→0.45)或 `max_seq_length`。** + +--- + +## 5. Step C — 正式训练 + +```bash +RAY_ADDRESS=local CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ + ./examples/qwen3-8b-single-node/run.sh configs/sglang_glm52_dspark_8card_colocate.yaml +``` +(`run.sh` 会 `ray start --head --num-gpus 8` → `train_entry` → 退出时 `ray stop`。) + +监控:`docs/performance_metrics.md` 的指标;checkpoint 落 `output_dir`(config 里 `./outputs/glm52-dspark-8card`),按 `save_interval` / `save_per_epoch`。 + +--- + +## 6. Step D — 评估(接受率 vs 自带 MTP) + +```bash +# 评估入口:torchspec/controller/eval.py;评估数据 examples/data/eval_conversations.jsonl +python -m torchspec.controller.eval --help # 先看实际参数 +# 也可用 tools/benchmark_eagle3.py 量 draft 接受率 / 加速比 +python tools/benchmark_eagle3.py --help +``` +目标:DSpark draft 的接受率 / 端到端加速 **优于 GLM-5.2 自带的 MTP(5 token)**,否则没必要替换。 + +--- + +## 7. 产物导出(给推理用) + +```bash +python tools/convert_to_hf.py --help # 把训练出的 draft checkpoint 转成 HF/SGLang 可加载格式 +``` +导出的 draft 即可挂到 SGLang/vLLM/TRT-LLM 做 GLM-5.2 的 speculative decoding。 + +--- + +## 8. 故障排查 + +| 现象 | 处理 | +|---|---| +| 训练 OOM | `inference.sglang.mem_fraction_static` 0.5→0.45→0.4;或降 `training.max_seq_length` | +| 推理 KV 不够 / prefill 失败 | 略升 `mem_fraction_static`,或降 `inference.inference_batch_size` | +| DSA kernel 崩溃 | 确认 B200=sm100 且 sglang ≥ 锁定 commit;**别在 sm89/L20 上跑**(需社区 ada_dsa port) | +| NCCL 端口/超时(colocate 两套通信组冲突) | 设不同 `mooncake` 端口;检查 `ray stop --force` 清干净;`tools/kill_all_torchspec.sh` | +| capture 到的 hidden 维度不对 | 核对 draft config `target_layer_ids`(注意 sglang 内部 +1 偏移)与 `target_hidden_size=6144` | +| 共享机器 Ray 串台 | 用 `RAY_ADDRESS=local` 强制本地实例 | +| `min_loss_tokens` 报错 | 设 ≥ 2×`dflash_block_size` | + +--- + +## 附:本方案涉及的文件 + +| 文件 | 状态 | +|---|---| +| `configs/sglang_glm52_dspark_8card_colocate.yaml` | 新增(主配置) | +| `torchspec/config/glm52_dspark_draft_config.json` | 新增(draft 维度,mask_token_id 待核验) | +| `torchspec/data/template.py` (`glm`) | 改(chat template,待核验) | +| `configs/sglang_glm52_dspark_2node.yaml` | 新增(备选:2 节点更快) | + +> **关键提醒**:`colocate` 是 docs 标注的 "Dev" 模式(`docs/ray.md:40`),实验性。首次务必先过 Step A/B 冒烟,再上正式训练。 diff --git a/torchspec/config/glm52_dspark_draft_config.json b/torchspec/config/glm52_dspark_draft_config.json new file mode 100644 index 0000000..db4a77c --- /dev/null +++ b/torchspec/config/glm52_dspark_draft_config.json @@ -0,0 +1,23 @@ +{ + "architectures": ["DSparkDraftModel"], + "model_type": "dspark", + "hidden_size": 6144, + "intermediate_size": 12288, + "num_hidden_layers": 5, + "num_attention_heads": 32, + "num_key_value_heads": 8, + "vocab_size": 154880, + "rms_norm_eps": 1e-05, + "max_position_embeddings": 65536, + "rope_theta": 8000000.0, + "num_target_layers": 5, + "target_hidden_size": 6144, + "target_num_hidden_layers": 78, + "target_layer_ids": [1, 20, 38, 56, 75], + "mask_token_id": 154820, + "markov_rank": 256, + "markov_head_type": "vanilla", + "enable_confidence_head": true, + "confidence_head_with_markov": true, + "tie_word_embeddings": false +} diff --git a/torchspec/data/template.py b/torchspec/data/template.py index bb788d4..8833230 100644 --- a/torchspec/data/template.py +++ b/torchspec/data/template.py @@ -266,3 +266,19 @@ def get_all_template_names(self) -> List[str]: image_placeholder="", ), ) + +# GLM-4.5 / GLM-5.x family. The `[gMASK]` conversation prefix is added +# automatically by the GLM tokenizer, so it is not encoded in the headers here. +# NOTE: verify the headers / end_of_turn_token / thinking behavior against the +# target checkpoint's tokenizer_config.json `chat_template` before real training. +# For the thinking variant, mirror `qwen3-thinking` (parser_type="thinking", +# enable_thinking=True, assistant_header="<|assistant|>\n\n"). +TEMPLATE_REGISTRY.register( + name="glm", + template=ChatTemplate( + assistant_header="<|assistant|>\n", + user_header="<|user|>\n", + system_prompt="You are a helpful assistant.", + end_of_turn_token="<|endoftext|>", + ), +)