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

• Любая LLM или VLM, поддерживаемая в Transformers, также может быть инициализирована через NeMo Automodel. См. полный список поддерживаемых моделей.
• Построен на базе моделей Hugging Face с использованием AutoModel.from_pretrained(), поддерживает динамическую замену высокопроизводительных слоёв и более продвинутые схемы параллелизма, такие как Expert Parallelism (EP).
• Определяет поле architecture в AutoConfig.from_pretrained() для автоматической загрузки пользовательских реализаций, таких как Nemotron Nano V3.
• Плотно следует API Transformers для обеспечения совместимости «из коробки».

Ресурсы

• NeMo Automodel
• API NeMo Transformers
• Бенчмарки benchmarks плотных моделей и моделей Mixture-of-Expert (MoE) в NeMo Automodel
• См. руководство по предварительному обучению в NeMo, чтобы узнать, как использовать NeMo для предварительного обучения