改进 LLM 的典型方法涉及多个阶段：合成数据生成 (SDG) 、通过监督式微调 (SFT) 或强化学习 (RL) 进行模型训练以及模型评估。每个阶段都需要使用不同的库，这些库通常很难设置，而且难以一起使用。
例如，您可以使用 NVIDIA TensorRT-LLM 或 vLLM 进行 SDG 训练，使用 NVIDIA NeMo 或 verl 进行训练。在这种情况下，您需要调用许多不同的脚本和容器，将 Hugging Face 检查点转换为 TensorRT-LLM，执行大规模 SDG，将数据和模型转换为 NeMo 格式，并在各种基准测试中运行训练和评估。
为了简化这一复杂的工作流程，NVIDIA 开发了 NeMo 技能库。它提供了无缝连接不同框架的高级抽象，使其能够以统一和可互换的方式使用。NeMo-Skill 还可以轻松地从快速本地原型设计过渡到在 Slurm 集群上编排大规模作业。
本文将向您介绍帮助 NVIDIA 团队赢得 AIMO2 Kaggle 比赛的简化版制作流程。该过程从数学推理能力有限的模型开始。然后，通过一系列 NeMo 技能作业来增强这些技能。
如果您继续关注，则需要使用配备 8 个 NVIDIA A100 (或更新版本) GPU 的 NVIDIA DGX 机箱，或具有类似配置节点的 Slurm 集群。演示中使用的所有命令都使用 NeMo 技能进行了测试。
在本地或 Slurm 上设置 NeMo 技能
为了编排复杂的作业，NeMo-Skills 使用 Docker 容器。如果在本地运行，或使用支持 NVIDIA/pyxis 的 Slurm 集群，则需要安装 NVIDIA Container Toolkit。在这两种情况下，建议您在本地工作站上设置 NeMo-Skill，并将其配置为通过 SSH 访问 Slurm 集群。它会负责上传您的代码并调度作业。
在本地运行以下命令以完成设置：
pip install git+https://github.com/NVIDIA/NeMo-Skills.git
ns setup
当系统提示您添加挂载时，请将文件夹定义为/workspace。后续命令中将使用此文件夹。有关更多详细信息，请参阅 NeMo-Skills 配置文档。
在以下各节中，我们将始终使用带有 --cluster=local 参数的命令，如果这些命令在 Slurm 上运行，您需要将其更改为 --cluster=slurm (或您在设置过程中命名的任何配置) 。使用 Slurm 时，所有命令将立即完成，并在集群队列中调度作业。
权重和偏差 (W&B) 将用于方便地记录评估结果和模型输出。您可以通过从命令中删除所有与 W&B 相关的参数来禁用此功能。
建立基准
在努力提升 LLM 技能之前，请先评估原始模型，看看它所处的位置。请注意，本教程适用于 Qwen2.5 14B Instruct，并使用 AIME24 和 AIME25 评估模型的数学推理能力。vLLM 用作推理库。
# download the model
ns run_cmd --expname=download-14b --log_dir=/workspace/Qwen2.5-14B-Instruct --cluster=local \
huggingface-cli download Qwen/Qwen2.5-14B-Instruct --local-dir /workspace/Qwen2.5-14B-Instruct
# prepare benchmark data
ns prepare_data aime24 aime25
# launch evaluation
ns eval \
--cluster=local \
--expname=baseline-eval \
--run_after=download-14b \
--model=/workspace/Qwen2.5-14B-Instruct \
--server_type=vllm \
--server_gpus=8 \
--benchmarks=aime24:8,aime25:8 \
--output_dir=/workspace/evals/baseline
# summarize results, after the evaluation job is done
ns summarize_results --cluster=local /workspace/evals/baseline --wandb_name=baseline-evals
在 aime24/ 25 基准测试中，ns eval 命令将为每个样本运行八代，summarize_results 将报告平均通过率 = 1、通过率 = 8 和大多数 = 8 个指标。
--------------------------------- aime24 --------------------------------
evaluation_mode | num_entries | avg_tokens | symbolic_correct | no_answer
pass@1[8] | 30 | 829 | 11.67% | 0.00%
majority@8 | 30 | 829 | 13.33% | 0.00%
pass@8 | 30 | 829 | 33.33% | 0.00%
--------------------------------- aime25 --------------------------------
evaluation_mode | num_entries | avg_tokens | symbolic_correct | no_answer
pass@1[8] | 30 | 834 | 11.67% | 0.42%
majority@8 | 30 | 834 | 20.00% | 0.00%
pass@8 | 30 | 834 | 26.67% | 0.00%
请注意，由于 LLM 代次的随机性质，您可能无法获得完全相同的数字。有关 ns eval 工作流选项的更多信息，请参阅 NeMo 技能评估文档。
将 SDG 与 TensorRT-LLM 或 vLLM 结合使用
为改进既定基准，您可以生成一些合成数学数据。按照 OpenMathReasoning recipe，使用少量 AoPS 论坛讨论，并使用 Qwen2.5 14B Instruct 从中提取问题。然后，使用 QwQ 32B 生成新的“长推理”解决方案。这些问题解决方案对将用于训练。
这种简化的工作流非常基础，并且忽略了多个重要步骤 (例如，提取真值答案和筛选正确性) 。但是，教会 14B 模型如何使用长推理并显著改善基准结果就足够了。
首先，下载数据以及问题提取提示和后处理脚本：
ns run_cmd --expname=prepare-data --log_dir=/workspace/prepare-data --cluster=local \
'cd /workspace && \
export DOWNLOAD_PREFIX=https://raw.githubusercontent.com/NVIDIA/NeMo-Skills/refs/heads/main/recipes/openmathreasoning && \
wget $DOWNLOAD_PREFIX/scripts/prepare_raw_data.py && \
wget $DOWNLOAD_PREFIX/prompts/extract-problems.yaml && \
wget $DOWNLOAD_PREFIX/scripts/postprocess_problem_extraction.py && \
python prepare_raw_data.py && \
head -n 1000 raw_aops_data.jsonl > data.jsonl'
data.jsonl 中的字段将用于填充 extract-problems.yaml 中的提示词，并将此最终提示词传递给 LLM。如需了解详情，您可以查看数据文件并提示脚本。有关提示格式的更多详细信息，请参阅 NeMo 技能提示文档。
接下来，使用 NeMo-Skills Python API 运行生成工作流：
# run_sdg.py
from nemo_skills.pipeline.cli import generate, wrap_arguments
cluster = "local"
num_gpus = 8
postprocess_cmd = (
f"python /workspace/postprocess_problem_extraction.py "
f" /workspace/sdg/problems/output.jsonl "
f" /workspace/sdg/extracted-problems.jsonl "
)
generate(
ctx=wrap_arguments(
f"++prompt_config=/workspace/extract-problems.yaml "
f"++prompt_template=qwen-instruct "
),
cluster=cluster,
input_file="/workspace/data.jsonl",
output_dir="/workspace/sdg/problems",
postprocess_cmd=postprocess_cmd,
expname="problem-extraction",
run_after=["prepare-data", "download-14b"],
model="/workspace/Qwen2.5-14B-Instruct",
server_type="vllm",
server_gpus=num_gpus,
# remove these parameters to disable wandb logging
log_samples=True,
wandb_group="sdg",
)
您可以检查 sdg/extracted-problems.yaml 以查看输出。应有一个新字段，其中包含提取的问题。使用 QwQ 32B 模型生成这些问题的解决方案。由于此模型会生成包含许多 token 的长推理解决方案，因此请将检查点转换为 TensorRT-LLM 格式，以获得更快的推理速度。
# download the model
ns run_cmd --expname=download-qwq --log_dir=/workspace/QwQ-32B --cluster=local \
huggingface-cli download Qwen/QwQ-32B --local-dir /workspace/QwQ-32B
# convert to trtllm format
ns convert \
--cluster=local \
--expname=convert-qwq-trtllm \
--run_after=download-qwq \
--input_model=/workspace/QwQ-32B \
--output_model=/workspace/qwq32b-trtllm \
--convert_from=hf \
--convert_to=trtllm \
--num_gpus=8 \
--model_type=qwen \
--hf_model_name=Qwen/QwQ-32B \
--max_seq_len 10000
下一步是生成解决方案。将以下代码添加到 sdg.py 脚本的末尾并重新运行。默认情况下，它将跳过问题提取步骤 (如果已完成) ，因为 NeMo-Skill 可以检测生成是否已经完成。
generate(
ctx=wrap_arguments(
f"++prompt_config=generic/math "
f"++inference.temperature=0.6 "
f"++inference.tokens_to_generate=8192 "
f"++prompt_template=qwen-instruct "
),
cluster=cluster,
input_file="/workspace/sdg/extracted-problems.jsonl",
output_dir="/workspace/sdg/solutions",
expname="solution-generation",
run_after=["problem-extraction", "convert-qwq-trtllm"],
model="/workspace/qwq32b-trtllm",
server_type="trtllm",
server_gpus=num_gpus,
# remove these parameters to disable wandb logging
log_samples=True,
wandb_group="sdg",
)
这项工作可能需要几个小时才能完成。如果您能够在 Slurm 集群的多个节点上运行，则可以通过添加 num_chunks=N 在 N 个独立作业中对其进行并行化。您可以在 NeMo 技能生成文档中详细了解此参数和其他参数。
如果您已启用 W&B 日志记录，则可以检查其中的几代产品。图 1 显示了 W&B 控制面板中 ns 生成工作流的输出。您可以打开一个实验，切换到“Files” (文件) 选项卡，然后单击 samples.json 来找到它。
ns generate 管道的输出使用 NeMo 进行模型训练
生成合成数据后，您可以使用它来微调模型。以下各节将展示如何使用 NeMo-Aligner 或 NeMo-RL 来执行此操作。
首先，按照所需格式准备数据：
ns run_cmd --log_dir=/workspace/prepare-sft-data --expname=prepare-sft-data --run_after=solution-generation --cluster=local \
'python -m nemo_skills.training.prepare_data \
++input_files=/workspace/sdg/solutions/output.jsonl \
++output_path=/workspace/sft-data.jsonl \
++prompt_config=generic/math \
++prompt_template=qwen-instruct \
++filters.remove_contaminated=false \
++add_unlabeled=true \
++filters.remove_no_think_tags=true \
++filters.trim_solutions=false'
接下来，将模型转换为 NeMo 格式。您可以跳过此步骤进行 NeMo-RL 训练。
ns convert \
--cluster=local \
--expname=convert-14b-nemo \
--run_after=download-14b \
--input_model=/workspace/Qwen2.5-14B-Instruct \
--output_model=/workspace/qwen2.5-14b-instruct-nemo \
--convert_from=hf \
--convert_to=nemo \
--num_gpus=8 \
--model_type=qwen \
--hf_model_name=Qwen/Qwen2.5-14B-Instruct
对于 NeMo-Aligner 后端，请使用以下训练命令。添加 --disable_wandb 以禁用 W&B 日志记录。
ns train \
--cluster=local \
--expname=training \
--run_after=convert-14b-nemo \
--run_after=prepare-sft-data \
--output_dir=/workspace/training \
--nemo_model=/workspace/qwen2.5-14b-instruct-nemo \
--num_nodes=1 \
--num_gpus=8 \
--training_data=/workspace/sft-data.jsonl \
++model.data.train_ds.max_seq_length=8192 \
++model.data.train_ds.global_batch_size=32 \
++model.tensor_model_parallel_size=4 \
++model.context_parallel_size=2 \
++model.optim.lr=1e-5 \
++trainer.sft.max_epochs=2
对于 NeMo-RL 后端，请使用以下训练命令。添加 --disable_wandb 以禁用 W&B 日志记录。仅运行其中一个训练命令，不能同时运行 (或相应地更改路径和表达式名称) 。
ns nemo_rl sft \
--cluster=local \
--expname=training \
--run_after=download-14b \
--run_after=prepare-sft-data \
--output_dir=/workspace/training \
--hf_model=/workspace/Qwen2.5-14B-Instruct \
--num_nodes=1 \
--num_gpus=8 \
--training_data=/workspace/sft-data.jsonl \
--cache_dir=/workspace/nemo-rl-cache \
--final_hf_path=/workspace/training/qwen2.5-14b-improved-hf \
++sft.max_num_epochs=4 \
++policy.dtensor_cfg.tensor_parallel_size=8 \
++policy.max_total_sequence_length=8192 \
++policy.train_global_batch_size=32 \
++policy.optimizer.kwargs.lr=1e-5 \
++policy.dtensor_cfg.sequence_parallel=true \
++policy.dtensor_cfg.activation_checkpointing=true
如需了解有关 SFT 配置的更多信息，请参阅 NeMo-Skills 训练文档。如果您已启用 W&B 日志记录，则可以检查此处的训练指标。
最终评估
要检查模型改进，请运行另一个评估。将检查点转换回 Hugging Face 格式，以加快评估速度。如果使用 NeMo-RL 后端进行训练，则可以跳过转换步骤。
# converting back to HF format
ns convert \
--cluster=local \
--expname=convert-14b-hf \
--run_after=training \
--input_model=/workspace/training/model-averaged-nemo \
--output_model=/workspace/training/qwen2.5-14b-improved-hf \
--convert_from=nemo \
--convert_to=hf \
--num_gpus=8 \
--model_type=qwen \
--hf_model_name=Qwen/Qwen2.5-14B-Instruct
# launching evaluation
ns eval \
--cluster=local \
--expname=final-eval \
--run_after=convert-14b-hf \
--run_after=training \
--model=/workspace/training/qwen2.5-14b-improved-hf \
--server_type=vllm \
--server_gpus=8 \
--benchmarks=aime24:8,aime25:8 \
--output_dir=/workspace/evals/after-training \
++inference.tokens_to_generate=16384
# summarize results, after the evaluation job is done
ns summarize_results --cluster=local /workspace/evals/after-training --wandb_name=after-training-evals
此评估应显示两个基准都有不错的改进。图 3 显示了 W&B 控制面板中的评估结果。切换到“Runs” (运行) 面板并单击“Columns” (列) 以自定义显示的指标。
--------------------------------- aime24 --------------------------------
evaluation_mode | num_entries | avg_tokens | symbolic_correct | no_answer
pass@1[8] | 30 | 13362 | 27.92% | 55.83%
majority@8 | 30 | 13362 | 40.00% | 16.67%
pass@8 | 30 | 13362 | 50.00% | 16.67%
--------------------------------- aime25 --------------------------------
evaluation_mode | num_entries | avg_tokens | symbolic_correct | no_answer
pass@1[8] | 30 | 13445 | 17.92% | 53.33%
majority@8 | 30 | 13445 | 26.67% | 10.00%
pass@8 | 30 | 13445 | 36.67% | 10.00%
使用 NeMo-Skills 改进任何 LLM
借助 NeMo-Skills，您可以连接提高 LLM 能力所需的各个阶段，从而轻松构建复杂的工作流。这使您能够在不同的训练和推理框架之间无缝切换。本教程中使用的所有命令都可以组合成一个脚本来调度整个作业。只需更改一次生产线，您就可以从本地工作站上的快速原型设计过渡到 Slurm 集群上的大规模实验。
作为练习，请尝试添加 NeMo-Skills 文档中提到的额外过滤步骤。您还可以尝试为每个问题生成多个解决方案，并检查这会如何影响最终评估结果。如您所见，只需一个脚本即可运行从数据生成到模型训练再到评估的所有内容，因此您可以轻松地对管道的任何部分进行更改迭代。
NVIDIA 团队成功使用 NeMo-Skills 开发了多个热门模型和数据集。具体来说，它用于：
- 为 NVIDIA 在 AIMO-2 竞赛中的获奖作品提供支持，并开发 OpenMathReasoning 数据集和 OpenMath-Nemotron 模型。
- 创建 OpenCodeReasoning 和 OpenCodeInstruct 集合。
- 为 Llama-Nemotron 模型系列整理部分后训练数据。
开始使用 NeMo-Skills，将其用于您自己的项目。
