NeMo-Skills框架：大模型优化流程的标准化解决方案

1. 大模型优化流程的挑战与NeMo-Skills解决方案

在当前的AI领域，提升大型语言模型（LLM）性能通常需要经历多个复杂阶段：合成数据生成（SDG）、监督微调（SFT）或强化学习（RL）训练，以及模型评估。每个阶段都需要使用不同的工具库，如TensorRT-LLM用于推理加速、NeMo用于模型训练等。这些工具虽然各自强大，但将它们整合到一个连贯的工作流中却面临诸多挑战：

• 环境配置复杂：不同框架依赖的CUDA版本、Python环境可能冲突
• 数据格式转换繁琐：需要在Hugging Face、TensorRT-LLM、NeMo等格式间反复转换
• 资源管理困难：本地开发与集群训练的配置差异大
• 流程碎片化：需要编写大量胶水代码连接各个阶段

NVIDIA NeMo-Skills正是为解决这些问题而设计的统一框架。它提供了三个关键价值：

1. 标准化接口：通过统一API封装不同后端（vLLM/TensorRT-LLM/NeMo）
2. 无缝扩展性：支持从本地8卡GPU到Slurm集群的无缝切换
3. 全流程管理：内置任务依赖跟踪和结果缓存机制

提示：在实际项目中，我们经常遇到模型转换导致精度损失的问题。NeMo-Skills的ns convert命令通过预设最优转换参数，可以最大程度保持模型性能。

2. 环境准备与基础配置

2.1 硬件需求与系统配置

要运行完整的优化流程，你需要以下硬件资源之一：

• 本地工作站：配备8块NVIDIA A100或H100 GPU（至少80GB显存）
• Slurm集群：节点配置与上述相当

系统软件要求：

• Ubuntu 20.04/22.04 LTS
• Docker 20.10+
• NVIDIA Container Toolkit
• CUDA 12.1+

对于Slurm集群，还需安装NVIDIA/pyxis插件以支持GPU容器调度。

2.2 NeMo-Skills安装与初始化

安装过程非常简单，只需执行：

pip install git+https://github.com/NVIDIA/NeMo-Skills.git ns setup

当提示挂载目录时，建议指定一个足够大的空间作为/workspace，后续所有生成的数据和模型都将存储在这里。典型的配置示例：

Mount points: - /mnt/data -> /workspace Environment variables: - HF_HOME=/workspace/huggingface - WANDB_DIR=/workspace/wandb

2.3 多环境支持配置

NeMo-Skills的强大之处在于可以统一管理不同执行环境。通过~/.nemo-skills/config.yaml可以定义多个集群配置：

clusters: local: type: local mounts: - /mnt/data:/workspace slurm_prod: type: slurm ssh: user@cluster-headnode partitions: [a100-80gb, h100] mounts: - /shared/data:/workspace

这样在运行命令时，只需通过--cluster参数指定目标环境，如--cluster=slurm_prod。

3. 基准测试与性能评估

3.1 模型下载与准备

我们以Qwen2.5-14B-Instruct模型为例，首先下载基础模型：

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

这个命令会：

1. 创建名为download-14b的实验记录
2. 将日志保存在指定目录
3. 使用Hugging Face CLI下载模型
4. 自动验证模型完整性

3.2 评估数据集准备

NeMo-Skills内置了多个数学推理基准测试的预处理脚本：

ns prepare_data aime24 aime25

该命令会自动：

• 下载原始AIME24/AIME25数据集
• 转换为标准JSONL格式
• 添加必要的元数据字段
• 存储在/workspace/data/benchmarks目录下

3.3 基准测试执行

使用vLLM后端运行评估：

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

关键参数说明：

• :8表示每个问题生成8个答案
• server_gpus=8指定使用全部8块GPU并行推理
• run_after确保依赖任务先完成

3.4 结果分析与解读

评估完成后，生成汇总报告：

ns summarize_results --cluster=local /workspace/evals/baseline --wandb_name=baseline-evals

典型输出示例：

--------------------------------- 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%

这些指标表明：

• pass@1：单次生成的正确率
• majority@8：8次生成的多数投票正确率
• pass@8：8次生成中至少一次正确的概率

经验分享：在实际应用中，我们发现当pass@8显著高于pass@1时（如本例33% vs 11%），说明模型存在高方差问题，适合通过自洽性(self-consistency)等技术提升性能。

4. 合成数据生成(SDG)实战

4.1 数据准备与预处理

从AoPS论坛提取数学问题：

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'

这个流程会：

1. 下载原始数据处理脚本
2. 运行初步清洗和格式化
3. 提取前1000条作为示例数据集

4.2 问题提取流水线

使用Python API构建SDG流程：

# run_sdg.py from nemo_skills.pipeline.cli import generate, wrap_arguments generate( ctx=wrap_arguments( "++prompt_config=/workspace/extract-problems.yaml " "++prompt_template=qwen-instruct" ), cluster="local", input_file="/workspace/data.jsonl", output_dir="/workspace/sdg/problems", postprocess_cmd="python /workspace/postprocess_problem_extraction.py " "/workspace/sdg/problems/output.jsonl " "/workspace/sdg/extracted-problems.jsonl", expname="problem-extraction", run_after=["prepare-data", "download-14b"], model="/workspace/Qwen2.5-14B-Instruct", server_type="vllm", server_gpus=8, log_samples=True, wandb_group="sdg", )

关键组件解析：

• prompt_config：定义问题提取的提示模板
• postprocess_cmd：指定后处理脚本清洗输出
• wandb_group：在W&B中分组显示相关实验

4.3 高性能推理优化

对于生成长推理链的解决方案，我们转换到TensorRT-LLM格式提升效率：

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

转换参数说明：

• max_seq_len 10000：支持超长文本生成
• num_gpus=8：并行加速转换过程
• model_type=qwen：指定模型架构优化方案

4.4 解决方案生成

扩展SDG脚本生成详细解答：

generate( ctx=wrap_arguments( "++prompt_config=generic/math " "++inference.temperature=0.6 " "++inference.tokens_to_generate=8192 " "++prompt_template=qwen-instruct" ), cluster="local", 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=8, log_samples=True, wandb_group="sdg", )

性能技巧：当处理大规模数据时，添加num_chunks=N参数可以将任务分割到N个Slurm节点并行执行，显著缩短总运行时间。

5. 模型训练与优化

5.1 训练数据准备

将生成的解决方案转换为训练格式：

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'

数据预处理关键步骤：

1. 根据prompt_template格式化输入
2. 过滤低质量样本（如不含推理步骤）
3. 添加辅助训练样本平衡分布
4. 保留原始文本不做截断

5.2 模型格式转换

将基础模型转换为NeMo格式：

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

5.3 监督微调(SFT)训练

使用NeMo-Aligner后端进行训练：

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

关键训练参数解析：

• tensor_model_parallel_size=4：4-way张量并行
• context_parallel_size=2：2-way上下文并行
• global_batch_size=32：总批次大小
• max_seq_length=8192：支持长上下文训练

5.4 强化学习(RL)训练方案

替代的NeMo-RL训练方案：

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

RL训练特有配置：

• sequence_parallel=true：优化长序列处理
• activation_checkpointing=true：节省显存
• dtensor_cfg：分布式张量配置

6. 最终评估与结果分析

6.1 模型格式回转换

将训练好的模型转换回Hugging Face格式：

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

6.2 性能对比评估

运行最终评估并比较结果：

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 ns summarize_results --cluster=local /workspace/evals/after-training --wandb_name=after-training-evals

典型改进结果：

--------------------------------- aime24 -------------------------------- | 原始模型 | 优化后模型 pass@1[8] | 11.67% | 27.92% (+16.25) majority@8 | 13.33% | 40.00% (+26.67) pass@8 | 33.33% | 50.00% (+16.67)

6.3 错误分析与改进方向

通过W&B分析错误样本，我们发现主要进步来自：

1. 推理步骤更完整（平均生成长度从829→13362 tokens）
2. 数学符号使用更准确
3. 减少了"不知道"的回答

进一步改进建议：

• 增加解决方案多样性（多采样温度）
• 添加验证步骤过滤错误解答
• 引入外部符号计算验证

7. 生产级部署建议

7.1 全流程自动化脚本

将整个流程整合为单个可执行脚本：

# pipeline.py from nemo_skills.pipeline.cli import run_pipeline run_pipeline( stages=[ ('download', {'model': 'Qwen/Qwen2.5-14B-Instruct'}), ('evaluate', {'benchmarks': ['aime24:8', 'aime25:8']}), ('sdg', {'prompt_config': 'recipes/openmathreasoning/extract-problems.yaml'}), ('train', {'method': 'nemo-aligner', 'epochs': 2}), ('evaluate', {'compare_to': 'baseline'}) ], cluster='slurm', # 一键切换到集群运行 wandb_project='llm-optimization' )

7.2 持续改进策略

建立迭代优化机制：

1. 每周自动生成新数据
2. 增量训练模型
3. A/B测试评估改进
4. 模型滚动更新

7.3 性能监控看板

建议监控的关键指标：

• 生成多样性（unique n-gram比率）
• 推理步骤长度分布
• 符号计算准确率
• 用户反馈评分

这些可以通过Grafana+Prometheus实现自动化监控。