NeMo Automodel

NeMo Automodel は、NVIDIA によるオープンソースの PyTorch DTensor ネイティブな学習ライブラリです。LLM および VLM 的大規模・小規模な事前学習とファインチューニングをサポートし、FSDP2、テンソル並列、パイプライン並列、エキスパート並列、コンテキスト並列などの並列化戦略を備えています。これにより、研究環境および本番環境での高速な実験が可能です。高スループットを実現するため、DeepEP および TransformerEngine のカーネルを統合しています。

YAML設定ファイルでトレーニング実行を定義します（完全な設定ファイルを参照）。

# Instantiate a Nemotron V3 Nano model
model:
  _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained
  pretrained_model_name_or_path: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16

# Run SFT on HellaSwag
dataset:
  _target_: nemo_automodel.components.datasets.llm.hellaswag.HellaSwag
  path_or_dataset: rowan/hellaswag
  split: train

# Train PEFT adapters
peft:
  _target_: nemo_automodel.components._peft.lora.PeftConfig
  exclude_modules: ["*.out_proj"]  # mamba layers use custom kernels that take in the out_proj.weight directly, thus lora doesn't work here.
  dim: 8
  alpha: 32
  use_triton: True

# Use EP + FSDP2 for training
distributed:
  strategy: fsdp2
  dp_size: none
  tp_size: 1
  cp_size: 1
  ep_size: 4

# ... other parameters

以下のコマンドを使用して、torchrunでトレーニングを開始します。

torchrun --nproc-per-node=4 examples/llm_finetune/finetune.py -c /path/to/yaml

Transformers の統合

Transformers でサポートされている LLM や VLM は、すべて NeMo Automodel 経由でもインスタンス化できます。対応モデルの全リストを参照してください。
Hugging Face のモデルを基盤とし、AutoModel.from_pretrained() を使用して構築されています。動的な高性能レイヤーの差し替えや、Expert Parallelism (EP) といったより洗練された並列化方式に対応しています。
AutoConfig.from_pretrained()内のarchitectureフィールドを検出し、Nemotron Nano V3などのカスタム実装を自動的に読み込みます。
TransformersのAPIに忠実に従っており、そのまま置き換え可能な互換性を備えています。

リソース

NeMo Automodel
NeMo Transformers API
NeMo AutomodelのdenseモデルおよびMixture-of-Expert (MoE)のベンチマーク
NeMoの事前学習ガイドを参照して、NeMoを用いた事前学習の使用方法を学びましょう。

Update on GitHub