NeMoフレームワークを用いたLLMの学習
松尾研究所の尾崎です．25卒でデータサイエンティストをやっています．
元来，LLMの構築にはデータの準備から学習．評価まで様々な工程がありますが，統一的なパイプラインを作るには複数のメーカーの異なるツールや独自実装との混合を検討する必要があります．
たとえばデータ準備は独自実装，事前学習はMegatronベース，事後学習ではtrlベース，評価はlightevalベース，のようにいくつかのツールを組み合わせていくことが常でした．
今回紹介するNVIDIA NeMo Frameworkは，大規模言語モデル（LLM），マルチモーダルモデル，音声AIの開発を効率化するための包括的なツール群を提供しています．単一のGPUから数千ノードのクラスタまでシームレスにスケールでき，データ準備・事前学習からファインチューニング，強化学習，推論デプロイまでの全工程をカバーする予定だそうです．NVIDIAからも「現在再構築中です」と発表がある通り，まだまだ開発途中といったイメージのツール群ではありますが，毎日のようにコミットが飛んでいてどんどんブラッシュアップされているみたいですね．
本記事では，NeMo Frameworkのアーキテクチャと各コンポーネントについて簡単に解説し，後半では実際にNeMo-RLを使ってフルパラメータSFTを実行する方法を，実際のセットアップ経験をもとに解説します．
NeMo Frameworkの全体像
NeMo Frameworkは，AIエージェント(詰まるところLLM)のライフサイクルを管理することを目的としているようで前述の通り，データ準備からデプロイ，その後のモニタリングまで包括的に対応できるツール群です．
LLM開発のライフサイクルと対応するライブラリ
とりわけLLM開発における主要4フェーズは既にカバーされています．
| フェーズ | ライブラリ | 役割 |
|---|---|---|
| Data Curation | Curator, Data Designer | データセットの前処理・品質管理・キュレーション・合成 |
| Pre-training & Fine-tuning | Megatron-Bridge, AutoModel | 事前学習，SFT，LoRA |
| Post-training Alignment | RL | 事後学習全般（SFT，RLHF，DPO等） |
| Inference & Safety | Export-Deploy, Evaluator, Guardrails | モデルのデプロイ，精度評価，安全性制御 |
統一的な実行環境
さらにこれらのライブラリはNeMo Runによって統一的に管理されるようです．NeMo Runにより，ローカルマシン，SLURMクラスタ，クラウド環境（Kubernetes）など，様々な計算環境でシームレスに実行できます．同じコードベースで，単一GPUから数千ノード規模のクラスタまで，スケーラビリティを損なうことなく拡張可能とのことです(が，ぱっと試した限りは重めのエラーで今回は断念．．．)
各ツールの紹介
- LLM学習用の大規模データセットをキュレーションするGPUアクセラレーテッドライブラリ
- テキスト、画像、動画、音声に対応
- 30以上のヒューリスティックフィルタ、品質・安全性フィルタ、PII除去、セマンティック重複排除を搭載
- 従来手法比で約16倍の高速化（8TBのデータセットを約30分で処理）
- 高品質な合成データを生成するためのライブラリ
- シードデータからの拡張やゼロからの生成に対応
- LLMを活用した多様なデータ生成と品質評価（LLM-as-a-judge）
- Python/SQLバリデーターによる出力品質の検証
- Hugging FaceモデルとMegatron形式間の双方向変換を提供
- Tensor/Pipeline並列化による大規模モデル学習に対応
- 混合精度トレーニング（FP8, BF16, FP4）をサポート
- Hugging Faceモデルを即座にGPUアクセラレーテッド学習できるライブラリ
- チェックポイント変換不要で即座に学習開始可能
- LoRAをサポート
- LLMの事後学習向けライブラリ
- SFT，PPO、DPO、GRPO、SteerLM、SPINなど複数のアルゴリズムに対応
- 人間フィードバックから報酬関数を学習し、モデルを最適化
- 学習済みモデルを推論用に最適化・デプロイするライブラリ
- TensorRT-LLM、vLLM、NVIDIA Tritonに対応
- FP8量子化とIn-flightバッチングにより高速推論を実現
- LLMの精度評価を効率化するマイクロサービス
- MMLU、GSM8K等60以上のベンチマークに対応
- LLM-as-a-Judge、RAG評価、カスタムデータセット評価に対応
- LLMアプリケーションに安全性レールを追加するオープンソースツールキット
- ファクトチェック、ハルシネーション検出、有害コンテンツモデレーション機能を搭載
- 複数のガードレールを並列実行しながら低レイテンシを維持
- 様々な計算環境でNeMoワークロードを統一的に管理・実行するツール
- LocalExecutor、SlurmExecutor、SkypilotExecutorに対応
- 自動チェックポイント保存とノード障害時の自動リカバリを実装
これら以外にもいくつかNeMoと名の付くツールはありますが，メインで開発が進んでいるのは前述のツール群のようです．NeMo-Alignerなどは使われているのを見ることもありますが，こちらはNeMo-RLやAutoModelへの統合が進んでいるようです．
NeMo-RLを使ったフルパラSFTについて
ではここからは実際にNeMo-RLを使って学習をやってみます．今回は9B級のモデルを1ノード8GPU(H100 80GB×8)使って数学の問題の英訳タスク(ニッチすぎる笑)について，フルパラメータSFTを行います．
NeMo-RLの主な特徴
| 特徴 | 説明 |
|---|---|
| 複数のアルゴリズム対応 | SFT、GRPO、DPO、On-policy Distillationなど |
| スケーラブル | 単一GPUからマルチノード・マルチGPUまで対応 |
| Rayベースの分散処理 | 効率的なリソース管理と柔軟なデプロイメント |
| Hugging Face統合 | 豊富な事前学習モデルをそのまま使用可能 |
| 高性能バックエンド | DTensor（PyTorch native）とMegatron Coreの2つの学習バックエンドを選択可能 |
全体の流れ
SFTを実行するまでの全体像を把握しておきましょう．
1. 環境構築
├── リポジトリのクローン
├── サブモジュールの初期化
├── 仮想環境の作成（uv）
└── 環境変数の設定
─────────────────────────────────────────────────────────────
2. データ準備
├── Hugging Faceからデータセットをダウンロード
└── JSONL形式に変換
─────────────────────────────────────────────────────────────
3. 設定ファイル作成
├── モデル設定
├── データ設定
├── 分散学習設定
└── ロギング設定
─────────────────────────────────────────────────────────────
4. 学習実行
├── Slurmジョブスクリプト作成
└── ジョブ投入・監視
1. 環境構築
1.1 リポジトリのクローン
まず，NeMo-RLのリポジトリをクローンします．重要なのは --recursive オプションです．これにより，依存するサブモジュールも同時にクローンされます．
git clone git@github.com:NVIDIA-NeMo/RL.git nemo-rl --recursive
cd nemo-rl
確認方法
サブモジュールが正しく初期化されているか確認するには：
git submodule status
正常な場合，各サブモジュールのコミットハッシュの前にスペースが表示されます．-が表示されている場合は未初期化です．
# 正常な状態
1d42deb98... 3rdparty/Automodel-workspace/Automodel
# 未初期化の状態（-が付いている）
-1d42deb98... 3rdparty/Automodel-workspace/Automodel
1.2 uvのインストール
NeMo-RLは，高速なPythonパッケージマネージャ uv を使用して依存関係を管理しています．
# uvのインストール
curl -LsSf https://astral.sh/uv/install.sh | sh
# インストール確認
uv --version
1.3 仮想環境の作成
プロジェクトルートで以下を実行します．
uv sync
1.4 環境変数の設定
学習を実行する前に，いくつかの環境変数を設定する必要があります．~/.bashrcに以下を追加してください．
# ==========================================
# NeMo-RL用環境変数
# ==========================================
# Hugging Faceのモデル・トークナイザーのキャッシュディレクトリ
export HF_HOME=/path/to/your/hf_cache
# Hugging Faceデータセットのキャッシュディレクトリ（任意）
export HF_DATASETS_CACHE=/path/to/your/datasets_cache
# WandB APIキー（学習曲線の可視化に使用）
export WANDB_API_KEY=your_wandb_api_key
設定を反映します．
source ~/.bashrc
1.5 Hugging Faceへのログイン（必要な場合）
Llama等のゲート付きモデルを使用する場合は，Hugging Faceにログインが必要です．
uv run huggingface-cli login
2. データセットの準備
2.1 NeMo-RLが期待するデータ形式
NeMo-RLのSFTは，messages形式のチャットフォーマットを使用します．これは以下のような構造です．
{
"messages": [
{"role": "system", "content": "あなたは数学の問題を解くアシスタントです。"},
{"role": "user", "content": "2 + 3 を計算してください。"},
{"role": "assistant", "content": "2 + 3 = 5 です。"}
]
}
2.2 Hugging Faceデータセットの変換
どうやらHugging Faceにあるデータセットを直接使うことは推奨されておらず，ローカルのJSONLファイルのみを受け付けるため，事前に変換orダウンロードが必要みたいです．
変換スクリプトの作成
scripts/download_dataset.pyを作成します．たとえば以下のような形が考えられます．私の場合は事前に学習データと検証データは切り分けておいたので，ダウンロードのみを行い，JSONL形式に変換しています．
import json
from pathlib import Path
from datasets import load_dataset
# 設定
DATASET_NAME = "your-organization/your-dataset-name"
OUTPUT_DIR = Path("data/your_dataset")
dataset = load_dataset(DATASET_NAME)
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
# 各splitの出力ファイル名を定義
split_to_filename = {
"train": "train.jsonl",
"validation": "val.jsonl",
"test": "test.jsonl",
}
for split, out_name in split_to_filename.items():
if split not in dataset:
continue
with open(OUTPUT_DIR / out_name, "w") as f:
for item in dataset[split]:
f.write(json.dumps({"messages": item["messages"]}, ensure_ascii=False) + "\n")
print(f"{split}: {len(dataset[split])} samples -> {out_name}")
スクリプトの実行
# ディレクトリを作成
mkdir -p scripts
# スクリプトを実行
uv run python scripts/download_dataset.py
実行後，data/your_dataset/に以下のファイルが作成されます．
data/your_dataset/
├── train.jsonl # 学習データ
├── val.jsonl # 検証データ
└── test.jsonl # テストデータ
3. 設定ファイルの作成
NeMo-RLはYAML形式の設定ファイルで学習パラメータを管理します．この辺は事後学習のOSSツールとして有名なAxolotlなどと同じですね．ここでは，各セクションを解説します．
3.1 設定ファイルの全体構造
examples/configs/sft_custom.yamlを作成します．
# ============================================
# SFT（Supervised Fine-Tuning）設定ファイル
# ============================================
# --------------------------------------------
# 1. 学習アルゴリズムの設定
# --------------------------------------------
sft:
# 学習エポック数
# 1エポック = 全学習データを1回学習
max_num_epochs: 1
# 検証（validation）の設定
val_period: 100 # 何ステップごとに検証を行うか
val_batches: 16 # 検証時に使用するバッチ数
val_global_batch_size: 32
val_micro_batch_size: 1
val_at_start: true # 学習開始前に検証を行うか
# 再現性のためのシード値
seed: 42
# --------------------------------------------
# 2. チェックポイント（モデル保存）の設定
# --------------------------------------------
checkpointing:
enabled: true
# チェックポイントの保存先ディレクトリ
checkpoint_dir: "results/sft_custom"
# ベストモデル選択の基準
metric_name: "val:val_loss" # 検証ロスを使用
higher_is_better: false # ロスは低い方が良い
# 保持するチェックポイント数（古いものから削除）
keep_top_k: 3
# 何ステップごとに保存するか
save_period: 1000
# --------------------------------------------
# 3. モデル（Policy）の設定
# --------------------------------------------
policy:
# Hugging Faceのモデル名またはローカルパス
model_name: "your-organization/your-model"
# トークナイザーの設定
tokenizer:
name: ${policy.model_name} # モデルと同じトークナイザーを使用
chat_template: "default" # モデルのデフォルトテンプレートを使用
chat_template_kwargs: null
# バッチサイズの設定
# global_batch_size = micro_batch_size × GPU数 × gradient_accumulation
train_global_batch_size: 32
train_micro_batch_size: 1 # GPU1枚あたりのバッチサイズ
# シーケンス長（トークン数）
# モデルの最大コンテキスト長以下に設定
max_total_sequence_length: 8192
# 計算精度
# bfloat16: メモリ効率と精度のバランスが良い
precision: "bfloat16"
# オプティマイザのオフロード（通常はfalse）
offload_optimizer_for_logprob: false
# --------------------------------------------
# 3.1 分散学習（DTensor）の設定
# --------------------------------------------
dtensor_cfg:
# DTensorのバージョン
# false: v1（安定版）、true: v2（新機能あり）
_v2: false
enabled: true
# 追加の環境変数
env_vars: {}
# CPUオフロード（メモリ不足時に有効化）
cpu_offload: false
# Sequence Parallel
# 注意: PyTorch 2.8/2.9ではTP>1で問題があるため無効化推奨
sequence_parallel: false
# Activation Checkpointing
# メモリ使用量を削減（計算時間は増加）
# 大きなモデルでは必須
activation_checkpointing: true
# Tensor Parallel Size
# モデルをGPU間で分割する数
# 通常はGPU数と同じに設定
tensor_parallel_size: 8
# Context Parallel Size（長文対応）
context_parallel_size: 1
custom_parallel_plan: null
# LoRA設定（Parameter-Efficient Fine-Tuning）
lora_cfg:
enabled: false # フルパラメータSFTの場合はfalse
# --------------------------------------------
# 3.2 Sequence Packing（効率化）
# --------------------------------------------
sequence_packing:
# 複数の短いサンプルを1つのシーケンスにまとめる
# スループットが大幅に向上
enabled: true
train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}}
algorithm: "modified_first_fit_decreasing"
sequence_length_round: 64
# Dynamic Batching（通常は無効）
dynamic_batching:
enabled: false
# シーケンス長をTPサイズで割り切れるようにする
make_sequence_length_divisible_by: ${policy.dtensor_cfg.tensor_parallel_size}
# 勾配クリッピング
max_grad_norm: 1.0
# --------------------------------------------
# 3.3 オプティマイザの設定
# --------------------------------------------
optimizer:
name: "torch.optim.AdamW"
kwargs:
lr: 2.0e-5 # 学習率
weight_decay: 0.1 # 重み減衰（正則化）
betas: [0.9, 0.98] # Adam のモーメンタム係数
eps: 1e-5
foreach: false # DTensorでは false に設定
fused: false # DTensorでは false に設定
# Megatron Core
# DTensorを使う場合は無効化
megatron_cfg:
enabled: false
env_vars: {}
# --------------------------------------------
# 4. データの設定
# --------------------------------------------
data:
# 最大入力シーケンス長
max_input_seq_length: ${policy.max_total_sequence_length}
# 特殊トークンの追加
add_bos: true # Beginning of Sequence
add_eos: true # End of Sequence
add_generation_prompt: false
# データのシャッフル
shuffle: true
# データローダーのワーカー数
num_workers: 4
# データセットタイプ
# "openai_format": OpenAI形式のJSONLファイル
dataset_name: "openai_format"
# データファイルのパス
train_data_path: "data/your_dataset/train.jsonl"
val_data_path: "data/your_dataset/test.jsonl"
# JSONLファイル内のキー名
chat_key: "messages"
system_key: null
system_prompt: null
tool_key: null
use_preserving_dataset: false
# --------------------------------------------
# 5. ロギングの設定
# --------------------------------------------
logger:
log_dir: "logs"
# WandB（Weights & Biases）
wandb_enabled: true
wandb:
entity: "your-wandb-entity" # WandBの組織名
project: "your-project-name" # プロジェクト名
name: "sft-experiment-1" # 実験名
# --------------------------------------------
# 6. クラスタの設定
# --------------------------------------------
cluster:
gpus_per_node: 8 # ノードあたりのGPU数
num_nodes: 1 # 使用するノード数
3.2 重要な設定項目について
Tensor Parallel Size について
tensor_parallel_sizeは，モデルをGPU間で分割する数を指定します．vllmとか馴染みがある引数名ですね．
例: 9Bモデル，8GPU，tensor_parallel_size=8 の場合
GPU 0: モデルの1/8を担当
GPU 1: モデルの1/8を担当
...
GPU 7: モデルの1/8を担当
Activation Checkpointing について
学習中の中間結果（Activation）を保存せず，必要時に再計算する技術です．
大きなモデル（7B以上）では，メモリ不足を防ぐためにtrueを推奨します．
Sequence Packing について
複数の短いサンプルを1つのシーケンスにまとめる技術です．
Packing無効の場合:
[サンプルA][PAD][PAD][PAD] -> GPUの計算が無駄
[サンプルB][PAD][PAD][PAD]
Packing有効の場合:
[サンプルA][サンプルB][サンプルC] -> 効率的
スループットが大幅に向上するため，有効化を推奨します．が，謎のメモリエラーなどを誘発することもあるので，場合によっては有効にしなくてもOKです．
4. Slurmジョブスクリプトの作成
4.1 基本的なジョブスクリプト
今回はslurm環境で学習を行います．ログインノードから1ノードの計算ノードに学習スクリプトをジョブとして投げます．run_sft_slurm.shを作成します．
#!/bin/bash
# ============================================
# Slurm ジョブ設定
# ============================================
#SBATCH --job-name=sft-custom # ジョブ名
#SBATCH --partition=your_partition # パーティション名（環境に合わせて変更）
#SBATCH --nodes=1 # ノード数
#SBATCH --gpus=8 # GPU数
#SBATCH --output=logs/sft-%j.out # 標準出力ログ（%jはジョブID）
#SBATCH --error=logs/sft-%j.err # エラーログ
# エラー時にスクリプトを停止
set -e
# プロジェクトディレクトリに移動
cd /path/to/nemo-rl
# 仮想環境をアクティベート
source .venv/bin/activate
# この環境変数がないと、NeMo-RLが各コンポーネント用に新しい仮想環境を作成しようとして時間がかかる
export NEMO_RL_PY_EXECUTABLES_SYSTEM=1
# ============================================
# 学習の実行
# ============================================
echo "Starting SFT training..."
python examples/run_sft.py \
--config examples/configs/sft_custom.yaml
見ての通り，先ほど作成したconfigを引数として与えるだけなので，簡単です．
またexamplesの中に各学習用のコードが入っています．基本はそのpresetを使うので問題ないです．
4.2 コマンドラインでの設定オーバーライド
設定ファイルを変更せずに，コマンドラインでパラメータをオーバーライド変更することもできます．たとえば以下のように．
python examples/run_sft.py \
--config examples/configs/sft_custom.yaml \
policy.train_global_batch_size=16 \
sft.max_num_epochs=3 \
logger.wandb.name="experiment-v2"
5. 学習の実行
5.1 事前準備
# ログディレクトリを作成
mkdir -p logs
# データが存在することを確認
ls -la data/your_dataset/
5.2 ジョブの投入
sbatch run_sft_slurm.sh
投入後，ジョブIDが表示されます．
Submitted batch job 12345
5.3 ジョブの監視
# ジョブの状態確認
squeue -u $USER
5.4 WandBでの可視化
WandBを有効にしている場合，リアルタイムに学習曲線を確認できます．
たとえば今回だと以下のような学習曲線を観測しました．
今回はタスクがモデルにとって簡単だったみたいで最初からかなりlossが低かったですね．じわじわ下がってはいますね．ちなみにこれは1epochです．
1epochで大体2.5kstepほどの学習でしたが，8GPUで大体1時間半くらいで終わりました．ちなみに今回は一旦FP16でtensor parallel sizeを8にして学習し，それ以外のData parallel，Context parallelは1にしており，Sequence parallelも無効化しています．なのでかなり保守的なメモリ設定ではあります．ちらっと確認した限りは各GPUのメモリは15~17GBくらいでしたね．
あとは評価についてですが，trainとは別途切り分けたデータでvalidation lossを見てみましたが，以下のように減少し続けてはいますね．
5.5 生成結果
ベースモデルと学習後モデルの出力を比較してみました．
良い例
入力: 次の方程式を解きなさい：x² - 5x + 6 = 0
| モデル | 出力 |
|---|---|
| ベースモデル | もちろんです。以下に、与えられた日本語のテキストを英語に翻訳します。 Solve the following equation: x² - 5x + 6 = 0 この翻訳が正確であることを願っています。何か他にお手伝いできることがあれば教えてください。 |
| 学習後モデル | Solve the following equation: x^2 - 5x + 6 = 0 |
ベースモデルは余計な前置きや後書きを付けてしまうのに対し，学習後モデルはシンプルに翻訳のみを出力しています．
悪い例
入力2: 関数 f(x) = 3x² + 2x - 1 を微分し、x = 2 における接線の傾きを求めなさい。
| モデル | 出力 |
|---|---|
| ベースモデル | （翻訳に加えて実際に微分計算を行い、答えまで出力） |
| 学習後モデル | Find the slope of the tangent line to the function I'm not sure how to do this... |
学習後モデルは翻訳自体は正しく行えていますが，末尾に「I'm not sure how to do this...」と余計な文が生成されてしまっています．
考察
ベースモデルの「余計な説明を付ける」という癖は改善されましたが，学習後モデルにも重複表現や不要な文の生成といった課題が残っています．実は今回学習データは合成データを多く含むデータになっていまして，そのときの合成品質などが影響してしまっている可能性はあります．学習ステップ数などのハイパラ調整など，さらなる改善の余地がありそうです．
6. トラブルシューティング
実際の環境構築で個人的に遭遇した問題とその解決策をまとめました．
6.1 Sequence Parallelのエラー
エラーメッセージ:
PyTorch 2.8/2.9でtensor_parallel_size > 1の場合に発生
解決策:
設定ファイルで無効化:
# シーケンス長をGPU間で共有する設定をOFFにします
dtensor_cfg:
sequence_parallel: false
6.2 Out of Memory (OOM)
エラーメッセージ:
CUDA out of memory
解決策（優先度順）:
- バッチサイズを減らす
policy:
train_global_batch_size: 16 # 32から減らす
- Activation Checkpointingを有効化
dtensor_cfg:
activation_checkpointing: true
- シーケンス長を減らす
policy:
max_total_sequence_length: 4096 # 8192から減らす
- CPUオフロードを有効化（最後の手段）
dtensor_cfg:
cpu_offload: true
6.3 DTensor v2のエラー
症状:
_v2: trueで不安定な動作．シーケンス長やバッチ処理ロジックの計算問題に起因しているようですが，特にMoEなどで発生している報告を散見しました．ただ将来的にはv2を使うことが推奨されているようです．
解決策:
v1に戻す:
dtensor_cfg:
_v2: false
7. 学習完了後の操作
7.1 チェックポイントの確認
学習が完了すると，checkpoint_dirにチェックポイントが保存されます．
ls -la results/sft_custom/
results/sft_custom/
├── step_1000/
│ ├── config.yaml
│ └── policy/
│ └── weights/
├── step_2000/
│ └── ...
└── best/
└── ...
7.2 Hugging Face形式への変換
評価や推論のために，チェックポイントをHugging Face形式に変換できます．converterが既にNeMo-RL内に用意されています．ただし，これでもPyTorch形式になるだけ(Dtensor→tensor)なので，safetensor形式にしたい場合はtransformersライブラリなどでモデルロードすると勝手にしてくれます．
uv run python examples/converters/convert_dcp_to_hf.py \
--config results/sft_custom/step_2000/config.yaml \
--dcp-ckpt-path results/sft_custom/step_2000/policy/weights/ \
--hf-ckpt-path results/sft_custom/hf_model
7.3 モデルの評価
変換したモデルを評価:
uv run python examples/run_eval.py \
generation.model_name=results/sft_custom/hf_model
まとめ
本記事では，NeMo Frameworkの全体像から各ツールの役割，そしてNeMo-RLを使った実践的なSFTの実行方法までを解説しました．まだまだNeMo-RLは活発に開発が進んでいるプロジェクトです．最新の情報は公式ドキュメントやGitHubを参照してください．
参考リンク
- NVIDIA NeMo Framework Documentation
- NeMo GitHub Organization
- NeMo-RL GitHub
- NeMo-RL Documentation
- NeMo Curator
- NeMo Data Designer
- NeMo Megatron-Bridge
- NeMo AutoModel
- NeMo Guardrails
- NeMo Evaluator
- NVIDIA Transformer Engine
- trl (Transformer Reinforcement Learning)
- lighteval
- Axolotl
- uv Documentation
- Weights & Biases
- Hugging Face Hub
Discussion