Nemo Automodel Recipe Development

NeMo AutoModel Recipe Development

<!-- NVSkills signature refresh requested for AM-519. -->

Instructions

For recipe questions, answer with the smallest complete path to action:

1. Name the relevant recipe file or YAML section.
2. List the builder functions or config keys involved.
3. Include a minimal YAML or command example when the question asks how to configure something.
4. End with a local validation command or tiny CPU-compatible test.

For conceptual recipe questions, answer from this skill without inspecting the repository or loading other AutoModel skills unless the user asks you to edit files. Keep the response focused on recipe YAML, builders, CLI routing, tests, and local validation.

Use these compact answer patterns for common questions:

• New finetuning recipe variant: start from the closest file under nemo_automodel/recipes/, update the model, dataset or dataloader, optimizer, loss, LR scheduler, step scheduler, and checkpoint builders, register a recipe alias only if adding a new recipe class, add example YAML under examples/, then add a tiny CPU-compatible unit test and run automodel <config.yaml>.
• _target_ fields: describe _target_ as the fully qualified Python callable, explain that sibling keys become keyword arguments, show optimizer and dataset examples, and mention nested CLI overrides such as --optimizer.lr.
• Validation and checkpointing: name step_scheduler.val_check_interval, step_scheduler.checkpoint_interval, validation_dataset, restore_from.path, and consolidated safetensors; include the minimal YAML snippet from this skill.

For validation and checkpointing, always name:

• step_scheduler.val_check_interval for validation cadence.
• step_scheduler.checkpoint_interval for save cadence.
• validation_dataset as the validation dataloader source.
• restore_from.path for resume.
• Consolidated safetensors as the default checkpoint format for HF ecosystem compatibility.

Routing Boundary

Use this skill for recipe construction and execution-flow questions: YAML structure, _target_ callables, builder functions, validation datasets, checkpoint configuration, CLI route registration, and recipe-specific tests.

Do not use this skill for standalone distributed strategy selection, cluster launcher configuration, or model architecture onboarding unless the user is asking how those choices appear inside an AutoModel recipe YAML.

Recipe Architecture

Execution Flow

CLI (automodel config.yaml)
  -> app.py resolves the config's recipe target
    -> recipe script (e.g. train_ft.py) main(config_path)
      -> Recipe class .setup() builds all components
        -> .run_train_validation_loop() executes training

Recipe Class

Recipes inherit from BaseRecipe and implement two methods:

• setup() -- builds model, optimizer, dataloader, loss, LR scheduler, step scheduler, and checkpoint config via builder functions.
• run_train_validation_loop() -- executes the training and validation loop.

Builder Pattern

All components are constructed through dedicated builder functions:

• build_model() -- instantiates the model from config
• build_optimizer() -- creates optimizer (AdamW, etc.)
• build_dataloader() -- sets up train and validation dataloaders
• build_loss_module() -- creates the loss function
• build_lr_scheduler() -- creates the learning rate scheduler
• build_step_scheduler() -- creates the step scheduler controlling training progression
• CheckpointingConfig -- configures checkpointing (built directly from the YAML checkpoint: block via RecipeConfig.checkpoint)

Infrastructure Application Order

Components are applied in this strict order after building:

1. PEFT (LoRA, etc.)
2. FP8 quantization
3. QAT (quantization-aware training)
4. Checkpoint load / restore
5. Parameter freezing
6. Sharding (FSDP2, Megatron-FSDP, DDP)
7. Device placement
8. torch.compile
9. Context parallelism hooks

YAML Config Anatomy

A complete recipe config follows this structure:

step_scheduler:
  max_steps: 1000
  num_epochs: 1
  grad_accumulation_steps: 4
  val_check_interval: 100
  checkpoint_interval: 500
  log_interval: 10

dist_env:
  master_addr: localhost
  master_port: 29500

rng:
  seed: 42

model:
  _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained
  pretrained_model_name_or_path: meta-llama/Llama-3.2-1B
  dtype: float32
  # additional model kwargs passed to the constructor

compile:
  enabled: false
  backend: inductor

clip_grad_norm:
  max_norm: 1.0

distributed:
  strategy: fsdp2       # fsdp2 | megatron_fsdp | ddp
  dp_size: auto
  tp_size: 1
  cp_size: 1

loss_fn:
  _target_: torch.nn.CrossEntropyLoss

dataset:
  _target_: nemo_automodel.datasets.squad.SquadDataset
  tokenizer_name_or_path: meta-llama/Llama-3.2-1B
  max_seq_length: 2048

validation_dataset:
  _target_: nemo_automodel.datasets.squad.SquadDataset
  split: validation

packed_sequence:
  enabled: false

dataloader:
  batch_size: 4
  num_workers: 4
  pin_memory: true

optimizer:
  _target_: torch.optim.AdamW
  lr: 2.0e-5
  weight_decay: 0.01

lr_scheduler:
  _target_: nemo_automodel.schedulers.CosineAnnealingWarmup
  warmup_steps: 50
  min_lr: 1.0e-6

Full-Parameter Training Precision

For new full-parameter training with torch.optim.Adam/AdamW, explicitly set model.dtype: float32 on NeMoAutoModel loaders for fp32 master weights and Adam moments. Configure compute precision separately (FSDP2: distributed.mp_policy).

PEFT, TE FusedAdam, and diffusion need separate precision choices; see the mixed-precision guide. Validate memory, training behavior, and checkpoint/resume when migrating existing configs.

The _target_ Pattern

The _target_ key specifies a fully qualified Python callable. All remaining keys in that section are passed as keyword arguments:

optimizer:
  _target_: torch.optim.AdamW   # callable
  lr: 2.0e-5                    # kwarg
  weight_decay: 0.01            # kwarg

This is equivalent to: torch.optim.AdamW(lr=2e-5, weight_decay=0.01).

CLI Overrides

Any config value can be overridden from the command line:

automodel config.yaml \
  --optimizer.lr 1e-4 \
  --step_scheduler.max_steps 500 \
  --distributed.tp_size 2

Examples

Validation and checkpointing:

step_scheduler:
  val_check_interval: 100
  checkpoint_interval: 500

validation_dataset:
  _target_: nemo_automodel.datasets.squad.SquadDataset
  split: validation

restore_from:
  path: /checkpoints/step-500

Domain-Specific Notes

LLM

• nemo_automodel/recipes/llm/train_ft.py handles both finetuning and pretraining. The distinction is in the config (dataset, learning rate, etc.).
• nemo_automodel/recipes/llm/kd.py implements knowledge distillation with a teacher and student model.
• nemo_automodel/recipes/llm/benchmark.py runs throughput and latency benchmarks.