改進 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 modelns 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 datans prepare_data aime24 aime25# launch evaluationns 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 donens 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_answerpass@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_answerpass@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.pyfrom nemo_skills.pipeline.cli import generate, wrap_argumentscluster = "local"num_gpus = 8postprocess_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 modelns 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 formatns 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 formatns 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 evaluationns 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 donens 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_answerpass@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_answerpass@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，將其用于您自己的項目。
?
