企业官网建设流程全解析

• 人工智能
• 大模型
• 模型优化
• 模型量化
• 模型压缩

导读

本文聚焦 NVIDIA Model-Optimizer 仓库中examples/speculative_decoding示例自带的 SLURM 数据准备方案，讲解如何将大规模合成对话数据的生成任务切分成多个 shard（分片），通过salloc分配多节点、由 vLLM / SGLang 推理服务并行调用基础模型逐片生成、最终合并回完整训练集的全流程。读完本文，你将掌握sharding_utils.py的分片/合并用法、launch.sh的 10 个命令行参数语义、worker 在节点上的真实执行链路（服务启动、温度扫描、断点续跑），并能直接把它复用到你自己的投机解码（EAGLE / DFlash）训练数据生产中。

为什么需要 SLURM 数据准备

投机解码（Speculative Decoding）通过在 GPU 内存中部署一个轻量草稿模型（draft model）与主模型协作，用一次前向验证多个候选 token，从而加速自回归生成。要让草稿模型的输出分布尽可能贴近主模型，训练数据最好由主模型自身生成——这就是"合成数据生成"（Data Synthesis）环节的价值，详见 examples/speculative_decoding/README.md 的 "Data Synthesis" 小节。

但在亿级样本规模下，单机串行生成是不可行的：每个样本都要经过一次完整的推理调用，耗时以小时计。为此，仓库在examples/speculative_decoding/distributed_generate/目录下提供了完整的 SLURM 并行化支持，其基本思路是：

1. 把输入 jsonl 按固定行数切成多个 shard；
2. 用salloc/ SLURM 作业分配 N 个计算节点；
3. 每个节点上启动 vLLM（或 SGLang）推理服务，串行处理分配给自己的若干个 shard；
4. 全部处理完成后把 shard 合并回单个 jsonl，作为后续训练数据。

整个过程假设$SLURM_JOB_ID已存在（即你已在交互式分配或作业脚本环境中），并假定已选定节点 n1、n2、n3、n4。

第一步：用 salloc 分配计算节点

文档给出的分配命令是：

salloc -N4 -A <account> -p <partition> -J <account>-synthetic:data-gen -t 120

各参数含义：

执行成功后，salloc会为当前 shell 绑定一个$SLURM_JOB_ID，后续srun --jobid=$JOB_ID ...才能把任务调度到这批节点上。仓库中还有一份可直接提交的作业脚本示例 examples/speculative_decoding/collect_hidden_states/slurm_dump.sh，它展示了#SBATCH -A <account_name>、#SBATCH --nodes=1 --ntasks-per-node=4 --gpus-per-node=4、-t 04:00:00等典型 SBATCH 写法，可作参考。

第二步：用 sharding_utils.py 创建数据分片

合成数据生成是"每个样本一次推理"，天然适合并行切分。仓库用 sharding_utils.py 完成分片与合并两个方向的操作。

分片（shard）

python3 distributed_generate/sharding_utils.py --input_path /data/train.jsonl --output_dir /data/train/ --max_lines_per_shard 10000

从源码看，该脚本的参数与行为如下：

分片文件的命名规则是train-{shard_idx:05d}-{shard_idx:05d}.jsonl，例如train-00000-00000.jsonl、train-00001-00001.jsonl。也就是说，一个 40 万行的输入文件按max_lines_per_shard=10000切分后，会得到train-00000到train-00039共 40 个 shard——这正是后文"前 40 个 shard"的由来。

合并（combine）

python3 distributed_generate/sharding_utils.py --input_dir /data/output/ --output_path /data/output.jsonl --combine

合并模式（--combine）的逻辑（见combine_jsonl_files函数）值得注意：

• 只收集目录下以.jsonl结尾的文件，并按文件名排序后依次读取；
• 跳过空白行；
• 跳过带有"finished": true标记的记录——这是server_generate.py在全部对话生成完毕后写入的完成哨兵（见下文）；
• 删除每条记录的conversation_id字段后，以 JSON 每行（jsonl）格式写入--output_path。

因此合并输出是"干净的"训练 jsonl，可直接喂给 main.py 或launch_train.sh进入 EAGLE 训练流程。

第三步：用 launch.sh 在 SLURM 上并行生成

命令形态

bash distributed_generate/launch.sh $SLURM_JOB_ID vllm TinyLlama/TinyLlama-1.1B-Chat-v1.0 /data/train/ /data/output /scripts/ 0 10 n1,n2,n3,n4 "\"You are a helpful assistant.\""

对照 launch.sh 第 20 行打印的 Usage，10 个位置参数语义为：

注意$6（scripts 路径）必须指向仓库中 examples/speculative_decoding 的绝对路径，因为 worker 容器内会调用/scripts/distributed_generate/worker.sh与/scripts/scripts/server_generate.py。

执行语义

launch.sh 的核心循环（第 54-66 行）对每个节点执行一次：

srun --output="srun_worker_${node}.log" --jobid="$JOB_ID" -N 1 --ntasks=1 --ntasks-per-node=1 -w "$node" \ --mpi pmix --overlap --container-image="$CONTAINER_IMAGE" \ --container-mounts="$MODEL_PATH":/model/,"$DATA_PATH":/input_data/,"$OUTPUT_PATH":/output_data/,"$SCRIPTS_PATH":/scripts/ \ bash /scripts/distributed_generate/worker.sh "$counter" "$BACKEND" "$JOBS_PER_NODE" "$SYSTEM_PROMPT" &

关键点：

• 每个节点通过-w "$node"精确钉到一个计算节点，-N 1 --ntasks=1保证每个节点只有一个 srun 任务；
• 使用Pyxis 容器：--container-image与--container-mounts。默认镜像按后端区分——vllm 用vllm/vllm-openai:v0.24.0，sglang 用lmsysorg/sglang:v0.5.3-cu129；可通过环境变量CONTAINER_IMAGE覆盖（例如本地.sqsh镜像）；
• 四个目录以只读/可写方式统一挂载：模型、输入分片、输出、脚本；
• worker 以&后台并行启动，每个节点的counter累加JOBS_PER_NODE；
• 脚本最后会wait所有 worker，任一失败则整体以非零码退出，并在日志srun_launch.log中记录每条 srun 命令。

以本例（4 节点、每节点 10 个 shard）为例：节点 n1 处理 shard 0-9，n2 处理 10-19，n3 处理 20-29，n4 处理 30-39，即"前 40 个 shard"。每个 shard 含 10000 行，每节点实际完成的样本数为jobs_per_node × max_lines_per_shard。

第四步：增量续跑——处理后续 shard

一次salloc分配 120 分钟可能不够处理全部数据，但不需要重新申请：worker 会记录已生成的conversation_id并跳过（见下文断点机制），而 launch.sh 的start_shard参数支持精确续跑。处理下一个 40 个 shard（即 shard 40-79）：

bash distributed_generate/launch.sh $SLURM_JOB_ID vllm TinyLlama/TinyLlama-1.1B-Chat-v1.0 /data/train/ /data/output /scripts/ 40 10 n1,n2,n3,n4

这里start_shard=40，于是节点 n1 处理 40-49、n2 处理 50-59、n3 处理 60-69、n4 处理 70-79。只要节点仍在分配时段内，就可以反复这样续跑，直到所有 shard 处理完毕。launch.sh 第 68 行会打印Will process shards $START_SHARD through $((counter - 1))供你核对本次范围。

深入 worker.sh：节点上的实际执行链路

worker.sh 是每个节点上的真正执行者，理解它能帮你排查超时、OOM、端口冲突等问题。

环境变量与启动参数

worker.sh 接收 4 个位置参数：_CURRENT_COUNTER（本节点起始 shard）、BACKEND、JOBS_PER_NODE、SYSTEM_PROMPT，其余行为由环境变量控制：

脚本做了严格的输入校验：backend 必须为 vllm/sglang、GPU_COUNT必须 > 0、TP 规模不得超过 GPU 数、NUM_TEMPERATURES不得超过 GPU 数（TP=1 场景）等。

服务启动与健康检查

当SGLANG_TP_SIZE=1时，worker 会为每张 GPU 启动一个独立服务（vLLM 或 SGLang），端口为BASE_PORT + gpu：

CUDA_VISIBLE_DEVICES=$gpu vllm serve /model/ \ --tensor-parallel-size 1 --served-model-name "$MODEL_NAME" \ --port "$port" --host 0.0.0.0 --trust-remote-code & # sglang 时等价于：python3 -m sglang.launch_server --model-path /model --tp 1 --port $port ...

随后进入健康检查循环，反复 curlhttp://localhost:${port}/health，全部返回 200 才继续（默认 600 秒超时）。全部服务就绪后，只有 MPI rank 0（PMIX_RANK/OMPI_COMM_WORLD_RANK）执行生成逻辑，避免重复。

温度扫描与 shard 处理

worker 把"每个 GPU 上的服务"与"每档温度"一一对应：温度0.0, 0.1, ..., 0.(NUM_TEMPERATURES-1)分别打到不同端口，同一 shard 可生成多个温度变体（对应多条对话），随后依次处理start_shard到end_shard范围内的每个 shard：

• 输入 shard 路径按printf "/input_data/train-%05d-%05d.jsonl"拼出（与 sharding_utils 命名严格对应）；
• 缺失或空 shard 自动跳过（[ ! -s "$shard" ] && continue），因此中途失败不会阻塞；
• 输出文件名为output-%05d-%05d-temp-0.X.jsonl；
• 每个 shard 调用一次：

python3 /scripts/scripts/server_generate.py \ --data_path "$shard" --output_path "$output" \ --num_threads "$TEXT_NUM_THREADS" --max_tokens "$MAX_TOKENS" \ --temperature "$temperature" --url "http://localhost:$port/v1" \ --log_empty_conversations [--system_prompt "$SYSTEM_PROMPT"]

TP>1 时（SGLANG_TP_SIZE>1），所有 GPU 组成一个张量并行服务、只有一个端口，温度扫描退化为对该端口的串行调用。

生成脚本 server_generate.py 的断点机制

worker 最终调用的 server_generate.py 是 Medusa 数据生成脚本的适配版，基于openai客户端走 OpenAI 兼容 API。它保证断点续跑安全：

• 启动时扫描已有输出文件，收集已完成的conversation_id并跳过；
• 输出文件已含"finished": true哨兵时直接退出（"All conversations already generated"）；
• 每条完成记录写入conversation_id、conversations（可选truncated: true）；
• 全部处理完后追加一行{"finished": true}——这正是 sharding_utils 合并时用于过滤已完成记录的依据。

因此，即使某个节点在中途挂掉，重新以相同start_shard启动同一批 shard 也不会重复生成，只会补齐缺失的conversation_id。参数方面，--num_threads默认 256（即并发请求数/批大小）、--temperature默认 0.0、--max_tokens默认 2048、--url默认http://localhost:8000/v1、--api_key默认token-abc123，与 README 中vllm serve ... --api-key token-abc123 --port 8000的示例一致。

多模态扩展：launch_multimodal.sh 与 worker_multimodal.sh

除纯文本外，仓库还提供面向视频 VLM 数据的并行生成入口 launch_multimodal.sh。与文本版的主要差异：

• 仅支持sglang后端（使用 SGLang 原生视频客户端）；
• 参数位支持两种顺序：可把media_path放在节点列表前（... <jobs_per_node> <media_path> [num_frames] <nodes> [system_prompt]），也可把节点列表放在前面（... <jobs_per_node> <nodes> <media_path> [num_frames] [system_prompt]），脚本通过判断第 9 参是否含/或.前缀来自动识别；
• 额外的挂载点MEDIA_PATH:/media_data/，并为 OpenAI API 模式起一个本地http.server（默认端口 18080）把媒体目录暴露为可访问 URL；
• 默认SGLANG_TP_SIZE=8、MAX_TOKENS=6144、NUM_FRAMES默认 32，支持API_MODE（openai）、VISION_TOKEN_FORMAT（默认 qwen_vl）等环境变量；
• worker 调用 server_generate_vlm_sglang.py（589 行），负责把视频/图片路径解析、抽帧后构造带视觉 token 的请求。

多模态 shard 的输入输出命名与文本版一致，同样可以复用sharding_utils.py分片、合并。

关联场景：SLURM 上的隐藏状态 dump

数据准备的另一条并行路径是离线训练所需的 base model 隐藏状态导出，见 collect_hidden_states/slurm_dump.sh。这份 SBATCH 脚本示范了 SLURM 数组作业（Array Job）用法：SLURM_ARRAY_TASK_ID作为 DP rank、SLURM_ARRAY_TASK_COUNT作为 DP world size，通过trtllm-llmapi-launch python3 .../compute_hidden_states_trtllm.py --tp 4 --dp-rank $SLURM_ARRAY_TASK_ID --dp-world-size $SLURM_ARRAY_TASK_COUNT实现节点内 TP、跨数组 DP 的并行。它与本文的合成数据生成是互补的两种 SLURM 场景，共同服务于投机解码训练的数据准备。

最佳实践与注意事项小结

1. 先分片再并行：max_lines_per_shard建议根据单 shard 处理时长与作业窗口（-t）估算，避免一个 shard 拖垮整批节点或频繁续跑。
2. 善用start_shard续跑：续跑时从上次结束的 shard 编号开始，worker 与 server_generate.py 的双重断点机制可保证不重复、不遗漏。
3. 容器镜像可控：离线集群可用CONTAINER_IMAGE指向本地.sqsh镜像，避免从 registry 拉取；Pyxis 要求挂载源路径在srun前已存在，因此 launch.sh 会先mkdir -p "$OUTPUT_PATH"。
4. 温度多样性与并发：NUM_TEMPERATURES与 GPU 数绑定（TP=1 时），多温度可提升合成数据多样性；客户端线程数TEXT_NUM_THREADS决定单 shard 内的并发请求量。
5. 完成后务必合并：训练前用--combine把output-*.jsonl合并为单个 jsonl，脚本会自动剔除空对话与finished哨兵。

整套方案全部位于 examples/speculative_decoding/distributed_generate 目录，配合 examples/speculative_decoding/README.md 中的 Data Synthesis 小节（单机版server_generate.py用法）即可从单机平滑扩展到多节点集群，为 EAGLE 系列草稿模型训练生产高质量合成数据。

• 人工智能
• 大模型
• 模型优化
• 模型量化
• 模型压缩