Repository Structure

This introductory guide presents the structure of the NeMo AutoModel repository, provides a brief overview of its parts, introduces concepts such as components and recipes, and explains how everything fits together.

What is NeMo AutoModel?

NeMo AutoModel is a PyTorch library for fine-tuning and pretraining large-scale models. In particular, it provides:

• Optimized implementations for training efficiency, including fused kernels and memory-saving techniques.
• Day-0 support for LLMs and VLMs available on the Hugging Face Hub.
• Seamless integration with Hugging Face datasets, tokenizers, and related tools.
• Distributed training strategies using FSDP2 and MegatronFSDP across multi-GPU and multi-node environments.
• End-to-end workflows with recipes for data preparation, training, and evaluation.

Repository Structure

The AutoModel source code is available under the nemo_automodel directory. It is organized into three directories:

• components/ - Self-contained modules
• recipes/ - End-to-end training workflows
• cli/ - CLI entry-point and job launcher dispatch.

Components Directory

The components/ directory contains isolated modules used in training loops. Each component is designed to be dependency-light and reusable without cross-module imports.

Directory Structure

The following directory listing shows all components along with explanations of their contents.

$ tree -L 1 nemo_automodel/components/
├── _peft/          - Implementations of PEFT methods, such as LoRA.
├── attention/      - Efficient attention modules and related utilities (e.g., flash attention, rotary embeddings).
├── checkpoint/     - Checkpoint save and load-related logic.
├── config/         - Utils to load YAML files and CLI-parsing helpers.
├── datasets/       - LLM and VLM datasets and utils (collate functions, preprocessing).
├── distributed/    - Distributed processing primitives (DDP, FSDP2, MegatronFSDP) and pipeline parallelism (AutoPipeline).
├── eval/           - Evaluation helpers, including tool-call parsing and scoring.
├── flow_matching/  - Flow-matching pipelines and model adapters for diffusion training.
├── launcher/       - Interactive, SkyPilot (including Kubernetes), and NeMo-Run job launchers.
├── loggers/        - Metric/event logging for Weights & Biases and other tools.
├── loss/           - Loss functions (such as cross-entropy and linear cross-entropy, etc.).
├── models/         - Optimized model implementations for LLMs and VLMs.
├── moe/            - Mixture of Experts modules and routing utilities for scalable model architectures.
├── optim/          - Optimizers and LR schedulers, including fused or second-order variants.
├── quantization/   - Quantization layers and helpers for 4-bit/8-bit or other reduced-precision training and inference.
├── speculative/    - Speculative decoding training, serving, and benchmarking components.
├── training/       - Training and fine-tuning utils.
└── utils/          - Small, dependency-free helpers (seed, profiler, timing, fs).

Key Features

• Each component can be used independently in other projects.
• Each component has its own dependencies, without cross-module imports.
• Unit tests are colocated with the component they cover.

Recipes Directory

Recipes define end-to-end workflows (data and model loading → training with custom loop → saving the output checkpoint) for a variety of tasks, such as training, fine-tuning, and knowledge distillation.

Available Recipes

The following directory listing shows the top-level recipe domains and their non-private Python modules. Package initializers and underscore-prefixed support modules are omitted.

$ tree -L 2 nemo_automodel/recipes/
├── diffusion/
│   └── train.py                 - Diffusion training.
├── dllm/
│   ├── strategy.py              - Shared diffusion language-model strategy helpers.
│   └── train_ft.py              - Diffusion language-model fine-tuning.
├── llm/
│   ├── benchmark.py             - LLM benchmarking.
│   ├── kd.py                    - LLM knowledge distillation.
│   ├── peagle_recipe.py         - Shared PEAGLE recipe behavior.
│   ├── precompute_dspark_dist.py - DSpark distance precomputation.
│   ├── train_dflash.py          - DFlash speculative training.
│   ├── train_domino.py          - Domino speculative training.
│   ├── train_dspark.py          - DSpark speculative training.
│   ├── train_eagle1.py          - EAGLE-1 speculative training.
│   ├── train_eagle2.py          - EAGLE-2 speculative training.
│   ├── train_eagle3.py          - EAGLE-3 speculative training.
│   ├── train_ft.py              - LLM pretraining and fine-tuning.
│   ├── train_jetspec.py         - JetSpec speculative training.
│   └── train_seq_cls.py         - Sequence-classification fine-tuning.
├── multimodal/
│   ├── finetune.py              - Multimodal fine-tuning.
│   └── pretrain.py              - Multimodal pretraining.
├── retrieval/
│   ├── mine_hard_negatives.py   - Hard-negative mining.
│   ├── train_bi_encoder.py      - Bi-encoder training.
│   └── train_cross_encoder.py   - Cross-encoder training.
└── vlm/
    ├── finetune.py              - VLM fine-tuning.
    └── kd.py                    - VLM knowledge distillation.

Run a Recipe

Recipes are launched via the automodel CLI:

$
automodel --nproc-per-node 2 examples/llm_finetune/llama3_2/llama3_2_1b_squad.yaml

The above command will fine-tune the Llama3.2-1B model on the SQuAD dataset with two GPUs using the llama3_2_1b_squad.yaml config. For a single-GPU run, omit --nproc-per-node:

$
automodel examples/llm_finetune/llama3_2/llama3_2_1b_squad.yaml

Each recipe imports the components it needs from the nemo_automodel/components/ catalog. The recipe/components structure enables you to:

• Decouple individual components and replace them with custom implementations when needed.
• Avoid rigid, class-based trainer structures by using linear scripts that expose training logic for maximum flexibility and control.

Configure a Recipe

An example YAML configuration is shown below. The complete config is available here:

1
recipe: TrainFinetuneRecipeForNextTokenPrediction
2
3
step_scheduler:
4
  global_batch_size: 64
5
  local_batch_size: 8
6
  ckpt_every_steps: 1000
7
  val_every_steps: 10  # will run every x number of gradient steps
8
  num_epochs: 1
9
10
model:
11
  _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained
12
  pretrained_model_name_or_path: meta-llama/Llama-3.2-1B
13
14
dataset:
15
  _target_: nemo_automodel.components.datasets.llm.squad.make_squad_dataset
16
  dataset_name: rajpurkar/squad
17
  split: train

More recipe examples are available under the examples/ directory.

CLI Directory (nemo_automodel/cli/)

The automodel (or am) CLI application simplifies job execution across different environments, from single-GPU interactive sessions to batch multi-node runs. It dispatches interactive (local), SkyPilot, and NeMo-Run launchers; Slurm jobs are submitted with sbatch scripts directly. The installed CLI implementation lives in nemo_automodel/cli/, and both console scripts resolve to nemo_automodel.cli.app:main.

Next Steps

Learn how to train models with NeMo AutoModel on:

• Your local workstation: See docs/launcher/local-workstation.mdx.
• A cluster: See docs/launcher/slurm.mdx.