NeMo Automodel

NeMo Automodel est une bibliothèque d’entraînement open source de NVIDIA, native de PyTorch DTensor. Elle prend en charge le pré-entraînement et l’ajustement fin à grande et petite échelle pour les LLMs et les VLMs, permettant des expérimentations rapides dans des environnements de recherche et de production, avec des stratégies de parallélisme incluant FSDP2, le parallélisme de tenseurs, de pipeline, d’experts et de contexte. Pour un débit élevé, elle intègre des noyaux issus de DeepEP et TransformerEngine.

Définissez votre session d’entraînement dans un fichier de configuration YAML (voir le fichier de configuration complet).

# 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

Lancez l'entraînement avec torchrun à l'aide de la commande ci-dessous.

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

Intégration de Transformers

• Tout LLM ou VLM pris en charge par Transformers peut également être instancié via NeMo Automodel. Consultez la couverture complète des modèles.
• Construit sur la base des modèles Hugging Face avec AutoModel.from_pretrained(), avec des substitutions dynamiques de couches haute performance et la prise en charge de parallélismes plus affinés tels que l'Expert Parallelism (EP).
• Détecte le champ architecture dans AutoConfig.from_pretrained() pour charger automatiquement des implémentations personnalisées telles que Nemotron Nano V3.
• Respecte de près l'API de Transformers pour une compatibilité directe.

Ressources

• NeMo Automodel
• API NeMo Transformers
• Benchmarks des modèles denses et Mixture-of-Expert (MoE) benchmarks
• Consultez le guide de pré-entraînement de NeMo pour apprendre à utiliser NeMo pour le pré-entraînement