Supervised Fine-Tuning (SFT) and Parameter-Efficient Fine-Tuning (PEFT) with NeMo AutoModel

View as Markdown

Introduction

Pretrained language models are general-purpose: they know a lot about language but nothing about your particular domain, terminology, or task. Fine-tuning bridges that gap. You fine-tune the model on your own examples so it produces answers that are accurate and relevant for your use case, without the cost of training a model from scratch. The result is a model optimized for your data that you can evaluate, publish, and deploy. This guide walks you through that process end-to-end with NeMo AutoModel (from installation through training, evaluation, and deployment) using Meta Llama 3.2 1B and the SQuAD v1.1 dataset as a running example.

NeMo AutoModel supports two fine-tuning modes:

  • Supervised Fine-Tuning (SFT) updates all model parameters. Use SFT when you need maximum accuracy and have sufficient compute.
  • Parameter-Efficient Fine-Tuning (PEFT) using LoRA freezes the base model and trains small low-rank adapters. PEFT reduces trainable parameters to less than 1% of the original model, lowering memory and storage costs.

Workflow Overview

┌──────────────┐ ┌──────────────┐ ┌──────────────┐ ┌──────────────┐ ┌──────────────┐ ┌──────────────┐ ┌──────────────┐
│ 1. Install │--->│ 2. Configure │--->│ 3. Train │--->│ 4. Inference │--->│ 5. Evaluate │--->│ 6. Publish │--->│ 7. Deploy │
│ │ │ │ │ │ │ │ │ │ │ (optional) │ │ (optional) │
│ uv venv + │ │ YAML config │ │ automodel CLI│ │ HF generate │ │ Val loss + │ │ HF Hub │ │ vLLM serving │
│ uv pip │ │ Choose SFT │ │ or torchrun │ │ API │ │ lm-eval- │ │ upload │ │ │
│ or Docker │ │ or PEFT │ │ │ │ │ │ harness │ │ │ │ │
└──────────────┘ └──────────────┘ └──────────────┘ └──────────────┘ └──────────────┘ └──────────────┘ └──────────────┘
StepSectionSFTPEFT
1. InstallInstall NeMo AutoModelSameSame
2. ConfigureConfigure Your Training RecipeYAML without peft: sectionYAML with peft: section
3. TrainFine-Tune the ModelSame command for both modesSame command for both modes
4. InferenceRun InferenceLoad consolidated checkpoint directlyLoad base model + adapter
5. EvaluateEvaluate the Fine-Tuned ModelValidation loss during training; lm-eval-harness post-trainingSame
6. PublishPublish to HF HubUpload model/consolidated/Upload model/ (adapter only)
7. DeployDeploy with vLLMvllm.LLM(model=...)Base model + native vLLM LoRARequest

Install NeMo AutoModel

$uv venv --python 3.12 --seed
$source .venv/bin/activate
$uv pip install "nemo-automodel" --torch-backend=auto

Alternatively, if you run into dependency or driver issues, use the pre-built Docker container:

$docker pull nvcr.io/nvidia/nemo-automodel:26.06.00
$docker run --gpus all -it --rm --shm-size=8g -v $(pwd)/checkpoints:/opt/Automodel/checkpoints nvcr.io/nvidia/nemo-automodel:26.06.00

Docker containers are ephemeral, and files written inside the container are lost when it stops. The -v flag in the docker run command above bind-mounts a local checkpoints/ directory into the container so that saved checkpoints persist across runs. For more details, see Save Checkpoints When Using Docker.

For the full set of installation methods, see the installation guide.

Configure Your Training Recipe

Training is configured through a YAML config file. For the automodel CLI, the file must name a recipe. The fine-tuning runtime also requires model, dataset, dataloader, loss_fn, optimizer, and distributed sections. If you configure a validation_dataset, also configure a validation_dataloader. The step_scheduler section is optional because it has typed defaults, but set its global and local batch sizes explicitly for a predictable effective batch. Add peft only for LoRA. For the complete copy-pastable file, see Full Config YAML.

Under the hood, both SFT and PEFT are executed by a recipe: a self-contained Python class that wires together model loading, dataset preparation, training, checkpointing, and logging. The fine-tuning recipe is TrainFinetuneRecipeForNextTokenPrediction. The config file tells the recipe what to build; the recipe decides how to build it.

NeMo AutoModel configs use a convention borrowed from Hydra: the special _target_ key tells the framework which Python class or function to call, and every other key in the same YAML block is passed as a keyword argument to that call. For example:

1optimizer:
2 _target_: torch.optim.Adam
3 lr: 1.0e-5
4 weight_decay: 0

At the YAML boundary, RecipeConfig normalizes this target and its fields into an OptimizerConfig wrapper. The recipe then supplies the runtime model parameters when it builds the optimizer:

1optimizer_config = recipe_config.optimizer
2optimizers = optimizer_config.build(model, device_mesh=device_mesh)

The _target_ value is a dotted Python import path: the same string you would use in an import statement. The framework resolves it at runtime by importing the module and looking up the attribute. This means you can point _target_ at any class constructor or factory function, and the remaining keys become its arguments.

For raw model, dataset, dataloader, and PEFT sections, inspect the Python signature of _target_. For typed sections such as optimizer, loss, step scheduler, LR scheduler, and checkpoint, inspect the corresponding config dataclass fields because RecipeConfig validates and owns those settings.

From YAML to running code. Here is the path a config takes through the framework:

finetune_config.yaml
┌──────────────┐ load_yaml_config() parses the file into
│ ConfigNode │◄─── a tree of ConfigNode objects, one per
└──────┬───────┘ YAML section.
┌──────────────┐ The recipe wraps the raw tree in RecipeConfig.
│ RecipeConfig │◄─── Known component sections become typed configs;
└──────┬───────┘ raw construction sections remain ConfigNodes.
┌────┴───────────────────────────────┐
▼ ▼
typed component configs raw ConfigNode sections
optimizer, loss_fn, model, dataset,
step_scheduler, lr_scheduler, dataloader, peft
checkpoint
│ │
▼ ▼
config.build(runtime arguments) recipe builders call instantiate()
│ with runtime model/tokenizer/data
└───────────────┬────────────────────┘
runtime components

Raw ConfigNode sections call .instantiate() inside the relevant recipe builder. .instantiate() does two things:

  1. Resolves _target_: Imports the Python path and obtains the callable (class or function).
  2. Calls it: Passes every other key in the section as a keyword argument.

Nested _target_ blocks (like a configured collate_fn inside dataloader) are recursively instantiated the same way. In contrast, RecipeConfig converts recognized sections to typed configs that own construction through .build(...): optimizer receives the model and device mesh, step scheduler receives the dataloader, DP size, and local batch size, and checkpoint receives runtime ranks and meshes.

The recipe key. Every config file includes a top-level recipe key that tells the CLI which recipe class to run. You can write it as a short name or as a fully-qualified Python path, and both resolve to the same class:

1# Short name (the CLI looks up the class automatically)
2recipe: TrainFinetuneRecipeForNextTokenPrediction
3
4# Fully-qualified path (used as-is)
5recipe: nemo_automodel.recipes.llm.train_ft.TrainFinetuneRecipeForNextTokenPrediction

The short name form is a convenience. The CLI scans all recipe modules under nemo_automodel.recipes and matches the bare class name. If you invoke the recipe script directly with torchrun instead of the automodel CLI, the recipe key is not required because the script itself is the recipe.

Not every section uses _target_. step_scheduler and checkpoint are converted into StepSchedulerConfig and CheckpointingConfig, whose .build(...) methods receive runtime-only values from the recipe. The distributed mapping is parsed into typed strategy, parallelism, pipeline, and MoE configs before DistributedSetup.build(...) creates the meshes.

Model

1model:
2 _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained
3 pretrained_model_name_or_path: meta-llama/Llama-3.2-1B
KeyRole
_target_Points to NeMoAutoModelForCausalLM.from_pretrained — a factory method that downloads (or loads from cache) a pretrained Hugging Face model and wraps it with NeMo distributed-training support.
pretrained_model_name_or_pathA keyword argument to from_pretrained. Any argument that from_pretrained accepts can be added here (e.g., cache_dir, torch_dtype).

This guide uses Meta Llama 3.2 1B as a running example. Replace pretrained_model_name_or_path with any supported Hugging Face model ID.

Llama is a family of decoder-only transformer models developed by Meta. The 1B variant is a compact model suitable for research and edge deployment, featuring RoPE positional embeddings, grouped-query attention (GQA), and SwiGLU activations.

Some Hugging Face models are gated. If the model page shows a “Request access” button:

  1. Log in with your Hugging Face account and accept the license.
  2. Ensure the token you use (from huggingface-cli login or HF_TOKEN) belongs to the approved account.

Pulling a gated model without an authorized token triggers a 403 error.

Dataset

1dataset:
2 _target_: nemo_automodel.components.datasets.llm.squad.make_squad_dataset
3 dataset_name: rajpurkar/squad # HF-Hub ID used to pull the dataset
4 split: train
5
6validation_dataset:
7 _target_: nemo_automodel.components.datasets.llm.squad.make_squad_dataset
8 dataset_name: rajpurkar/squad
9 split: validation
KeyRole
_target_Points to make_squad_dataset — a factory function that downloads the SQuAD dataset, tokenizes it, and returns a torch.utils.data.Dataset. To use a different dataset, change _target_ to a different factory function (see Integrate Your Own Text Dataset).
dataset_name, splitKeyword arguments passed to make_squad_dataset. Each dataset factory defines its own parameters — check the function signature to see what’s available.

This guide uses SQuAD v1.1 as a running example. Swap the dataset by changing _target_ and the dataset arguments. See Integrate Your Own Text Dataset and Dataset Overview: LLM, VLM, and Retrieval Datasets.

The Stanford Question Answering Dataset (SQuAD) is a reading comprehension dataset where each example consists of a Wikipedia passage, a question, and an answer span. SQuAD v1.1 guarantees all questions are answerable from the context, making it suitable for straightforward fine-tuning.

Example:

1{
2 "context": "Architecturally, the school has a Catholic character. ...",
3 "question": "To whom did the Virgin Mary allegedly appear in 1858 in Lourdes France?",
4 "answers": { "text": ["Saint Bernadette Soubirous"], "answer_start": [515] }
5}

Optional PEFT

1peft:
2 _target_: nemo_automodel.components._peft.lora.PeftConfig
3 target_modules: "*_proj" # glob pattern matching projection layer FQNs
4 dim: 8 # low-rank dimension of the adapters
5 alpha: 32 # scaling factor for learned weights
KeyRole
_target_Points to PeftConfig — a dataclass that describes which layers to adapt and how. Unlike the model and dataset sections, this instantiation produces a config object, not the adapter itself. The recipe passes the resulting PeftConfig into build_model, which applies LoRA adapters to the model.
target_modulesA glob pattern matched against fully-qualified layer names (e.g. "*_proj" matches Llama projection layers such as q_proj, k_proj, and down_proj).
dimThe low-rank dimension r — controls adapter capacity. Larger values learn more but use more memory.
alphaScaling factor applied to the adapter output (alpha / dim). Higher values give adapters more influence during training.

Including a peft: section enables LoRA fine-tuning. Remove it entirely to run SFT instead. See Switch Between SFT and PEFT.

QLoRA for Quantized Low-Rank Adaptation

If GPU memory is a constraint, QLoRA combines LoRA with 4-bit NormalFloat (NF4) quantization to reduce memory usage by up to 75% compared to full-parameter SFT in 16-bit precision, while maintaining comparable quality to standard LoRA.

QLoRA requires bitsandbytes, which is provided by AutoModel’s cuda_source install profile. Add it to the active environment created above before running a QLoRA recipe:

$uv pip install "nemo-automodel[cuda_source]"

To enable QLoRA, add a quantization: section alongside the peft: section in your config. The example keeps the "*_proj" projection-layer pattern and increases dim from 8 to 16 for additional adapter capacity.

1model:
2 _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained
3 pretrained_model_name_or_path: meta-llama/Llama-3.2-1B
4
5peft:
6 _target_: nemo_automodel.components._peft.lora.PeftConfig
7 target_modules: "*_proj" # match projection layers such as q_proj and down_proj
8 dim: 16 # LoRA rank (higher than default to offset quantization)
9 alpha: 32 # scaling factor
10 dropout: 0.1 # LoRA dropout rate
11
12quantization:
13 load_in_4bit: True # enable 4-bit quantization
14 load_in_8bit: False # use 4-bit, not 8-bit
15 bnb_4bit_compute_dtype: bfloat16 # compute dtype
16 bnb_4bit_use_double_quant: True # double quantization for extra savings
17 bnb_4bit_quant_type: nf4 # NormalFloat quantization type
18 bnb_4bit_quant_storage: bfloat16 # storage dtype for quantized weights

Training Schedule

1step_scheduler:
2 global_batch_size: 64 # samples per optimizer step across all DP ranks
3 local_batch_size: 8 # samples per forward/backward pass on each DP rank
4 num_epochs: 1 # train over the dataset once

step_scheduler has no _target_. RecipeConfig converts it to StepSchedulerConfig, and the recipe builds the runtime scheduler with the training dataloader, data-parallel group size, and local_batch_size. Gradient accumulation is derived rather than configured directly:

gradient accumulation steps = global_batch_size / (local_batch_size × DP size)

With the values above, a single DP rank accumulates 8 micro-batches; 8 DP ranks require no accumulation. global_batch_size must be divisible by local_batch_size × DP size.

Most example recipes use bf16 training by default for memory and throughput. If you are running long fine-tuning, especially full-parameter SFT, and need higher-precision optimizer state, configure it explicitly instead of assuming it from the mixed-precision compute policy. See the mixed-precision training guide for the recommended TE and torch AdamW patterns.

Full Config YAML

Save as finetune_config.yaml. This config runs PEFT (LoRA). To run SFT instead, remove the peft: section. Use this complete config rather than the current checked-in Llama SFT YAML: that file still nests the seed under rng:, which this recipe ignores, and its commented W&B block uses unsupported save_dir instead of dir.

1recipe: TrainFinetuneRecipeForNextTokenPrediction
2
3step_scheduler:
4 global_batch_size: 64
5 local_batch_size: 8
6 ckpt_every_steps: 10
7 val_every_steps: 10
8 num_epochs: 1
9
10dist_env:
11 backend: nccl
12 timeout_minutes: 1
13
14distributed:
15 strategy: fsdp2
16 dp_size: null
17 tp_size: 1
18 cp_size: 1
19
20model:
21 _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained
22 pretrained_model_name_or_path: meta-llama/Llama-3.2-1B
23
24peft:
25 _target_: nemo_automodel.components._peft.lora.PeftConfig
26 target_modules: "*_proj"
27 dim: 8
28 alpha: 32
29
30loss_fn:
31 _target_: nemo_automodel.components.loss.masked_ce.MaskedCrossEntropy
32
33dataset:
34 _target_: nemo_automodel.components.datasets.llm.squad.make_squad_dataset
35 dataset_name: rajpurkar/squad
36 split: train
37
38dataloader:
39 _target_: torchdata.stateful_dataloader.StatefulDataLoader
40 collate_fn:
41 _target_: nemo_automodel.components.datasets.utils.default_collater
42 shuffle: false
43
44validation_dataset:
45 _target_: nemo_automodel.components.datasets.llm.squad.make_squad_dataset
46 dataset_name: rajpurkar/squad
47 split: validation
48
49validation_dataloader:
50 _target_: torchdata.stateful_dataloader.StatefulDataLoader
51 collate_fn:
52 _target_: nemo_automodel.components.datasets.utils.default_collater
53
54optimizer:
55 _target_: torch.optim.Adam
56 lr: 1.0e-5
57 betas: [0.9, 0.999]
58 eps: 1.0e-8
59 weight_decay: 0
60
61checkpoint:
62 enabled: true
63 checkpoint_dir: checkpoints/
64 model_save_format: safetensors
65 save_consolidated: final

Fine-Tune the Model

You can run the recipe using the AutoModel CLI or directly with torchrun (advanced).

$automodel --nproc-per-node=8 finetune_config.yaml

The --nproc-per-node=8 flag specifies the number of GPUs per node. Adjust as needed. For a single GPU, omit the --nproc-per-node option.

Invoke the Recipe Script Directly (Advanced)

Alternatively, you can invoke the recipe script directly using torchrun as shown below.

$torchrun --nproc-per-node=8 nemo_automodel/recipes/llm/train_ft.py -c finetune_config.yaml

Sample Output

Running the recipe with the automodel app or by invoking the recipe script directly produces the following log:

$ automodel finetune_config.yaml
INFO:nemo_automodel.cli.app:Config: finetune_config.yaml
INFO:nemo_automodel.cli.app:Recipe: nemo_automodel.recipes.llm.train_ft.TrainFinetuneRecipeForNextTokenPrediction
INFO:nemo_automodel.cli.app:Launching job interactively (local)
cfg-path: finetune_config.yaml
INFO:root:step 4 | epoch 0 | loss 1.5514 | grad_norm 102.0000 | mem: 11.66 GiB | tps 6924.50
INFO:root:step 8 | epoch 0 | loss 0.7913 | grad_norm 46.2500 | mem: 14.58 GiB | tps 9328.79
Saving checkpoint to checkpoints/epoch_0_step_10
INFO:root:step 12 | epoch 0 | loss 0.4358 | grad_norm 23.8750 | mem: 15.48 GiB | tps 9068.99
INFO:root:step 16 | epoch 0 | loss 0.2057 | grad_norm 12.9375 | mem: 16.47 GiB | tps 9148.28
INFO:root:step 20 | epoch 0 | loss 0.2557 | grad_norm 13.4375 | mem: 12.35 GiB | tps 9196.97
Saving checkpoint to checkpoints/epoch_0_step_20
INFO:root:[val] step 20 | epoch 0 | loss 0.2469

Each log line reports the current loss, gradient norm, peak GPU memory, and tokens per second (TPS). Small fluctuations between steps (for example, 0.2057 to 0.2557) are normal. Look at the overall downward trend rather than individual values.

Checkpoint Contents

Checkpoints are saved as Hugging Face-compatible safetensors. For SFT, save_consolidated: final is the default: intermediate checkpoints contain sharded model weights plus a generated model/consolidate.sh helper, while the final checkpoint also contains model/consolidated/ for Transformers, vLLM, lm-eval-harness, and other Hugging Face ecosystem tools. Set save_consolidated: false to keep every checkpoint sharded and run the helper for whichever checkpoint you want to export. Use save_consolidated: every (or legacy true) only when you intentionally want inline HF export at every checkpoint save. PEFT checkpoints contain only the adapter weights (megabytes instead of gigabytes) and are saved directly under model/. They do not use model/consolidate.sh. At inference time you must load the original base model and apply the adapter on top. This distinction affects every downstream step (inference, publishing, and deployment).

SFT checkpoint:

$$ tree checkpoints/epoch_0_step_10/
$checkpoints/epoch_0_step_10/
$├── config.yaml
$├── dataloader.pt
$├── model
$ ├── consolidate.sh
$ ├── shard-00001-model-00001-of-00001.safetensors
$ └── shard-00002-model-00001-of-00001.safetensors
$├── optim
$ ├── __0_0.distcp
$ └── __1_0.distcp
$├── rng.pt
$└── step_scheduler.pt
$
$3 directories, 9 files

PEFT checkpoint:

$$ tree checkpoints/epoch_0_step_10/
$checkpoints/epoch_0_step_10/
$├── dataloader.pt
$├── config.yaml
$├── model
$ ├── adapter_config.json
$ ├── adapter_model.safetensors
$ └── automodel_peft_config.json
$├── optim
$ ├── __0_0.distcp
$ └── __1_0.distcp
$├── rng.pt
$└── step_scheduler.pt
$
$2 directories, 8 files

Run Inference

Inference uses the Hugging Face generate API. Because exported SFT checkpoints are self-contained while PEFT checkpoints store only adapter weights (see Checkpoint Contents), the loading procedure differs between the two modes.

SFT Inference

For an intermediate checkpoint, or any run with save_consolidated: false, first run the generated helper for the checkpoint you want to load:

$bash checkpoints/epoch_0_step_10/model/consolidate.sh

The exported SFT checkpoint at model/consolidated/ is a complete Hugging Face model and can be loaded directly:

1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer
3
4ckpt_path = "checkpoints/epoch_0_step_10/model/consolidated"
5tokenizer = AutoTokenizer.from_pretrained(ckpt_path)
6model = AutoModelForCausalLM.from_pretrained(ckpt_path)
7
8device = "cuda" if torch.cuda.is_available() else "cpu"
9model.to(device)
10
11prompt = (
12 "Context: Architecturally, the school has a Catholic character. "
13 "Atop the Main Building's gold dome is a golden statue of the Virgin Mary. "
14 "Immediately in front of the Main Building and facing it, is a copper statue of Christ "
15 "with arms upraised with the legend 'Venite Ad Me Omnes'.\n\n"
16 "Question: What is atop the Main Building?\n\n"
17 "Answer:"
18)
19inputs = tokenizer(prompt, return_tensors="pt").to(device)
20output = model.generate(**inputs, max_new_tokens=50)
21print(tokenizer.decode(output[0], skip_special_tokens=True))

PEFT Inference

PEFT adapters must be loaded on top of the base model. The base AutoModel install does not include the Hugging Face PEFT library, so add its tested version range to the same active environment before running the inference code:

$uv pip install "peft>=0.18.1"
1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer
3from peft import PeftModel
4
5base_model_name = "meta-llama/Llama-3.2-1B"
6tokenizer = AutoTokenizer.from_pretrained(base_model_name)
7model = AutoModelForCausalLM.from_pretrained(base_model_name)
8
9adapter_path = "checkpoints/epoch_0_step_10/model/"
10model = PeftModel.from_pretrained(model, adapter_path)
11
12device = "cuda" if torch.cuda.is_available() else "cpu"
13model.to(device)
14
15prompt = (
16 "Context: Architecturally, the school has a Catholic character. "
17 "Atop the Main Building's gold dome is a golden statue of the Virgin Mary. "
18 "Immediately in front of the Main Building and facing it, is a copper statue of Christ "
19 "with arms upraised with the legend 'Venite Ad Me Omnes'.\n\n"
20 "Question: What is atop the Main Building?\n\n"
21 "Answer:"
22)
23inputs = tokenizer(prompt, return_tensors="pt").to(device)
24output = model.generate(**inputs, max_new_tokens=50)
25print(tokenizer.decode(output[0], skip_special_tokens=True))

Evaluate the Fine-Tuned Model

Validation Loss During Training

The recipe automatically computes validation loss at the interval set by val_every_steps. Look for [val] lines in the training log:

INFO:root:[val] step 20 | epoch 0 | loss 0.2469

A decreasing validation loss across checkpoints indicates the model is learning. If validation loss plateaus or increases while training loss continues to drop, the model might be overfitting, so consider stopping earlier or reducing the learning rate.

Post-Training Evaluation Using lm-eval-harness

For task-specific benchmarks (e.g., MMLU, GSM8K, HellaSwag accuracy), use lm-eval-harness with the fine-tuned checkpoint. For an intermediate SFT checkpoint, or any SFT run with save_consolidated: false, run bash checkpoints/epoch_0_step_20/model/consolidate.sh before pointing evaluation at model/consolidated/. Create and activate a clean evaluation environment, then install both model backends before running either command:

$uv venv --python 3.12 --seed .venv-lm-eval
$source .venv-lm-eval/bin/activate
$uv pip install "lm_eval[hf,vllm]" --torch-backend=cu130
$
$# SFT checkpoint (using vLLM backend for faster evaluation)
$lm_eval --model vllm \
> --model_args pretrained=checkpoints/epoch_0_step_20/model/consolidated/ \
> --tasks hellaswag \
> --batch_size auto
$
$# PEFT adapter (using Hugging Face backend with built-in PEFT support)
$lm_eval --model hf \
> --model_args pretrained=meta-llama/Llama-3.2-1B,peft=checkpoints/epoch_0_step_20/model/ \
> --tasks hellaswag \
> --batch_size auto

The quoted lm_eval[hf,vllm] requirement installs both backends and their dependencies. The SFT example uses the vllm backend for faster evaluation. The PEFT example uses the hf backend with lm-eval’s built-in PEFT support to load the adapter on top of the base model.

Run lm-eval-harness on the base model before fine-tuning to establish a baseline, then compare against the fine-tuned checkpoint.

Publish to the Hugging Face Hub

Fine-tuned checkpoints and PEFT adapters are stored in Hugging Face-native format and can be uploaded directly to the Hub. For an intermediate SFT checkpoint, or any SFT run with save_consolidated: false, upload model/consolidated/ after running the generated consolidation helper.

  1. Install the Hugging Face Hub library (if not already installed):
$uv pip install "huggingface_hub"
  1. Log in to Hugging Face:
$huggingface-cli login
  1. Upload:

SFT checkpoint:

1from huggingface_hub import HfApi
2
3api = HfApi()
4api.upload_folder(
5 folder_path="checkpoints/epoch_0_step_10/model/consolidated",
6 repo_id="your-username/llama3.2_1b-finetuned-squad",
7 repo_type="model",
8)

PEFT adapter:

1from huggingface_hub import HfApi
2
3api = HfApi()
4api.upload_folder(
5 folder_path="checkpoints/epoch_0_step_10/model",
6 repo_id="your-username/llama3.2_1b-lora-squad",
7 repo_type="model",
8)

Once uploaded, load the checkpoint or adapter directly from the Hub:

SFT:

1from transformers import AutoModelForCausalLM
2
3model = AutoModelForCausalLM.from_pretrained("your-username/llama3.2_1b-finetuned-squad")

PEFT:

1from transformers import AutoModelForCausalLM
2from peft import PeftModel
3
4model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B")
5model = PeftModel.from_pretrained(model, "your-username/llama3.2_1b-lora-squad")

Deploy with vLLM

vLLM is an efficient inference engine for production deployment of LLMs.

Use a vLLM-compatible environment. For NVIDIA GPUs, create and activate a virtual environment before installing vLLM:

$uv venv --python 3.12 --seed .venv-vllm
$source .venv-vllm/bin/activate
$uv pip install "vllm" --torch-backend=cu130

You can skip these commands if your active environment already includes vLLM.

SFT Checkpoint with vLLM

For an intermediate checkpoint, or any run with save_consolidated: false, run the generated model/consolidate.sh helper before serving from model/consolidated/.

1from vllm import LLM, SamplingParams
2
3llm = LLM(model="checkpoints/epoch_0_step_10/model/consolidated/", model_impl="transformers")
4params = SamplingParams(max_tokens=20)
5outputs = llm.generate("Toronto is a city in Canada.", sampling_params=params)
6print(f"Generated text: {outputs[0].outputs[0].text}")
>>> Generated text: It is the capital of Ontario. Toronto is a global hub for cultural tourism. The City of Toronto

PEFT Adapter with vLLM

vLLM can load the AutoModel LoRA adapter directly on top of the original base model. Enable LoRA when constructing LLM, then select the adapter for the request with LoRARequest:

1from vllm import LLM, SamplingParams
2from vllm.lora.request import LoRARequest
3
4base_model = "meta-llama/Llama-3.2-1B"
5adapter_path = "checkpoints/epoch_0_step_10/model"
6
7llm = LLM(model=base_model, enable_lora=True, max_lora_rank=8)
8params = SamplingParams(temperature=0, max_tokens=20)
9outputs = llm.generate(
10 ["How are you doing?"],
11 params,
12 lora_request=LoRARequest("squad_adapter", 1, adapter_path),
13)
14print(outputs[0].outputs[0].text)

Full Configuration Reference

This section documents all available config fields for the fine-tuning recipe. For the quick-start config, see Configure Your Training Recipe.

Switch Between SFT and PEFT

The peft: section controls which mode runs:

ModeWhat to do in the YAML
PEFT (LoRA)Include the peft: section as shown below.
SFT (full-parameter)Remove/comment the peft: section entirely.

All other config sections remain the same for both modes.

Full Configuration

1# Recipe
2# Selects which recipe class runs the training loop.
3# Use a short name (auto-discovered) or a fully-qualified Python path:
4# recipe: nemo_automodel.recipes.llm.train_ft.TrainFinetuneRecipeForNextTokenPrediction
5recipe: TrainFinetuneRecipeForNextTokenPrediction
6
7# Training Schedule
8# RecipeConfig converts this mapping to StepSchedulerConfig. The recipe passes
9# runtime dataloader and DP values to its build() method.
10step_scheduler:
11 global_batch_size: 64 # samples per optimizer step across all DP ranks
12 local_batch_size: 8 # samples per forward/backward pass on each DP rank
13 # accumulation = global / (local × DP size)
14 ckpt_every_steps: 10 # save a checkpoint every N optimizer steps
15 val_every_steps: 10 # run the validation loop every N optimizer steps
16 num_epochs: 1 # how many full passes over the training dataset
17
18# Process Group
19# Initializes the PyTorch distributed process group.
20# No _target_ (backend and timeout are passed to initialize_distributed()).
21# You normally would not need to tune this.
22dist_env:
23 backend: nccl # communication backend: "nccl" (GPU, recommended) or "gloo" (CPU)
24 timeout_minutes: 1 # timeout for collective operations; increase for large models
25 # that take longer to initialize
26
27# Distributed Strategy
28# Determines how model weights, data, and compute are split across GPUs.
29# Parsed into typed strategy and parallelism configs before DistributedSetup.build().
30# See "Distributed Training: TP, PP, CP, and EP" in Advanced Topics for details.
31distributed:
32 strategy: fsdp2 # parallelism strategy: "fsdp2" (recommended), "megatron_fsdp",
33 # or "ddp". FSDP2 shards parameters and optimizer states across
34 # the data-parallel group.
35 dp_size: null # data-parallel group size. null = auto-detect from
36 # world_size ÷ (tp_size × cp_size × pp_size).
37 tp_size: 1 # tensor-parallel size: splits weight matrices across GPUs.
38 # Set to 2, 4, or 8 if the model doesn't fit on one GPU.
39 # Should divide evenly into the number of attention heads.
40 cp_size: 1 # context-parallel size: splits the input sequence across GPUs.
41 # Increase for very long contexts (e.g., 32k+ tokens).
42 sequence_parallel: false # when true, extends TP to also shard activations along
43 # the sequence dimension for additional memory savings
44
45# Random Number Generator
46# The recipe constructs its checkpointable, rank-aware StatefulRNG from this seed.
47seed: 1111 # global random seed for reproducibility
48
49# Model
50# _target_ → NeMoAutoModelForCausalLM.from_pretrained: downloads (or loads from
51# cache) a pretrained Hugging Face model and wraps it with NeMo distributed-training
52# support. Any from_pretrained kwarg is accepted (such as cache_dir and torch_dtype).
53model:
54 _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained
55 pretrained_model_name_or_path: meta-llama/Llama-3.2-1B
56
57# PEFT (remove / comment this entire section for full-parameter SFT)
58# _target_ → PeftConfig: a dataclass describing which layers get LoRA adapters.
59# The recipe passes this config into build_model(), which attaches adapters
60# to the matching layers.
61peft:
62 _target_: nemo_automodel.components._peft.lora.PeftConfig
63 target_modules: "*_proj" # glob pattern matched against fully-qualified layer names;
64 # matches Llama projection layers such as q_proj and down_proj
65 dim: 8 # low-rank dimension r — controls adapter capacity.
66 # Larger values are more expressive but use more memory.
67 alpha: 32 # LoRA scaling factor: adapter output is scaled by alpha/dim.
68 # Higher values give adapters more influence during training.
69 use_triton: True # use an optimized Triton kernel for LoRA forward/backward
70 # (requires the triton package)
71
72# Checkpointing
73# RecipeConfig converts this mapping to CheckpointingConfig; the recipe supplies
74# runtime ranks and meshes to its build() method.
75checkpoint:
76 enabled: true # set to false to skip saving checkpoints entirely
77 checkpoint_dir: checkpoints/ # output directory. Docker users: bind-mount this path
78 # (e.g. -v $(pwd)/checkpoints:/opt/Automodel/checkpoints)
79 # to persist checkpoints across container restarts.
80 model_save_format: safetensors # "safetensors" (recommended, faster and safer) or
81 # "torch_save" (legacy pickle-based format)
82 save_consolidated: final # default: intermediate checkpoints are sharded; the final checkpoint
83 # also exports consolidated HF weights. Other modes: false (sharded only)
84 # or every/true (export every checkpoint).
85
86# Training Dataset
87# _target_ → make_squad_dataset: a factory function that downloads the SQuAD
88# dataset, tokenizes it, and returns a torch Dataset. To use a different dataset,
89# change _target_ to another factory function (see the dataset guide).
90dataset:
91 _target_: nemo_automodel.components.datasets.llm.squad.make_squad_dataset
92 dataset_name: rajpurkar/squad # HuggingFace Hub dataset ID
93 split: train # which split to use (train, validation, test)
94
95# Validation Dataset
96validation_dataset:
97 _target_: nemo_automodel.components.datasets.llm.squad.make_squad_dataset
98 dataset_name: rajpurkar/squad
99 split: validation
100 limit_dataset_samples: 64 # cap validation set to 64 samples for faster eval loops;
101 # remove this line to use the full validation set
102
103# Training Dataloader
104# _target_ → StatefulDataLoader: a checkpointable DataLoader from torchdata that
105# saves and restores iteration state across training restarts, so resumed runs
106# don't re-process already-seen batches.
107dataloader:
108 _target_: torchdata.stateful_dataloader.StatefulDataLoader
109 collate_fn:
110 _target_: nemo_automodel.components.datasets.utils.default_collater
111 # pads and batches individual samples into tensors;
112 # can be swapped for custom collation
113 shuffle: true # whether to shuffle the dataset each epoch
114
115# The recipe sets the actual loader batch size from
116# step_scheduler.local_batch_size, overriding any dataloader.batch_size key.
117
118# Validation Dataloader
119validation_dataloader:
120 _target_: torchdata.stateful_dataloader.StatefulDataLoader
121 collate_fn:
122 _target_: nemo_automodel.components.datasets.utils.default_collater
123
124# Loss Function
125# _target_ → MaskedCrossEntropy: standard cross-entropy loss that automatically
126# ignores padding tokens so they don't affect the gradient.
127# Other available loss functions (swap _target_ to use):
128# - nemo_automodel.components.loss.chunked_ce.ChunkedCrossEntropy
129# Computes CE in chunks along the sequence dimension to reduce peak memory.
130# Useful for very long sequences. Accepts chunk_len (default 32).
131# - nemo_automodel.components.loss.linear_ce.FusedLinearCrossEntropy
132# Fuses the final linear projection (lm_head) with the CE computation,
133# avoiding the full logit tensor. Significant **memory savings** for large vocabs.
134# - nemo_automodel.components.loss.te_parallel_ce.TEParallelCrossEntropy
135# TE-based parallel CE with a Triton kernel. Designed for tensor-parallel
136# setups where logits are sharded across TP ranks.
137loss_fn:
138 _target_: nemo_automodel.components.loss.masked_ce.MaskedCrossEntropy
139
140# Optimizer
141# _target_ → torch.optim.Adam: RecipeConfig wraps this target and its fields in
142# an OptimizerConfig, which receives model parameters when the recipe calls build().
143optimizer:
144 _target_: torch.optim.Adam
145 lr: 1.0e-5 # learning rate — the most important hyperparameter to tune
146 betas: [0.9, 0.999] # Adam momentum coefficients (β₁ for mean, β₂ for variance)
147 eps: 1e-8 # small constant added to the denominator for numerical stability
148 weight_decay: 0 # L2 regularization strength (0 = no regularization)
149
150# Logging (optional)
151# Uncomment to enable Weights & Biases experiment tracking.
152# wandb:
153# project: <your_wandb_project> # W&B project name
154# entity: <your_wandb_entity> # W&B team or username
155# name: <your_wandb_exp_name> # display name for this run
156# dir: <your_wandb_dir> # local directory for W&B artifacts

Config Field Reference

SectionRequired?What to change
recipeCLI onlyRequired by automodel; select TrainFinetuneRecipeForNextTokenPrediction by short name or fully-qualified path.
modelYesSet pretrained_model_name_or_path to your Hugging Face model ID. Source: auto_model.py.
peftPEFT onlyRemove entirely for SFT. Adjust dim and alpha to tune adapter capacity. use_triton: True enables an optimized LoRA kernel (requires the triton package). For reduced memory usage, see QLoRA. Source: lora.py.
datasetYesChange _target_, dataset_name, and split for your data. Source: squad.py.
dataloaderYesConfigure the loader target, collation, workers, and shuffling. Set its effective batch size with step_scheduler.local_batch_size; the recipe injects that value into the loader. Uses StatefulDataLoader for checkpointable iteration. Collation: utils.py.
validation_dataset, validation_dataloaderFor validationConfigure both when running validation. The validation loader also receives step_scheduler.local_batch_size.
loss_fnYesMaskedCrossEntropy is the standard choice. Alternatives: ChunkedCrossEntropy (long sequences), FusedLinearCrossEntropy (large vocabs), TEParallelCrossEntropy (tensor-parallel).
seedOptionalControls the checkpointable, rank-aware RNG that the recipe constructs. Source: rng.py.
step_schedulerOptional (recommended)Set global_batch_size and local_batch_size; accumulation is derived as global / (local × DP size). ckpt_every_steps and val_every_steps are counted in optimizer steps.
distributedYesdp_size: null means auto-detect from world size. Adjust tp_size for tensor parallelism across GPUs.
checkpointRecommendedSet checkpoint_dir to a persistent path, especially in Docker. save_consolidated defaults to final: intermediate SFT checkpoints are sharded and the final checkpoint also exports consolidated HF weights.
optimizerYesAny supported torch.optim class can be selected via _target_; RecipeConfig normalizes it to an OptimizerConfig that owns construction. For long fine-tuning, especially full-parameter SFT, see the mixed-precision training guide before combining torch AdamW with bf16 resident parameters.
wandbOptionalUncomment and configure to enable Weights & Biases logging.

For the fine-tuning recipe itself, see train_ft.py.

Per-Parameter-Group Learning Rate

To give a subset of parameters a different learning rate or weight decay (for example, a smaller learning rate for MoE router parameters), add param_group_overrides to the optimizer block. Each entry matches parameters by name and applies a multiplier that the LR scheduler folds into that group every step:

1optimizer:
2 _target_: torch.optim.AdamW
3 lr: 2.0e-5
4 weight_decay: 0.01
5 param_group_overrides:
6 - pattern: "router" # Python regex matched against each parameter name (re.search)
7 lr_mult: 0.1 # this group trains at 0.1 x the scheduled learning rate
8 wd_mult: 1.0 # (optional) multiplier for weight decay; defaults to 1.0

Each parameter joins the group of the first override whose pattern matches its name; parameters matching no pattern keep the base learning rate. A pattern that matches nothing is skipped with a warning. This works for the standard torch.optim optimizers; Dion-family optimizers do their own parameter grouping and ignore param_group_overrides.

Distributed Training with TP, PP, CP, and EP

The distributed: section controls how the model and data are split across GPUs. NeMo AutoModel supports four parallelism dimensions, each of which slices the workload differently:

DimensionKeyWhat it shardsWhen to use
Data Parallelism (DP)dp_sizeReplicates the model on each group of GPUs; each replica trains on a different data batch.Default. Scales batch size linearly with GPU count.
Tensor Parallelism (TP)tp_sizeSplits individual weight matrices (attention, MLP) across GPUs within a node.Model is too large to fit on a single GPU, or you want to reduce per-GPU memory at the cost of extra communication.
Pipeline Parallelism (PP)pp_sizeAssigns different layers (stages) to different GPUs and pipelines micro-batches through them.Very deep models that do not fit even with TP, or multi-node training where TP’s all-reduce is too expensive across nodes.
Context Parallelism (CP)cp_sizeSplits the input sequence across GPUs so each GPU processes a portion of the context.Very long sequences that exceed single-GPU memory.
Expert Parallelism (EP)ep_sizeDistributes MoE experts across GPUs so each GPU holds a subset of experts.Mixture-of-Experts models only.

These dimensions compose with each other. The relationship between them and total GPU count is:

world_size = pp_size × dp_size × cp_size × tp_size

When dp_size is set to null (the default), it is inferred automatically:

dp_size = world_size ÷ (tp_size × cp_size × pp_size)

EP does not appear in this formula. Experts are distributed across the DP×CP rank groups, with the constraint that (dp_size × cp_size) must be divisible by ep_size.

Data Parallelism as the Default

Data parallelism is the default. With strategy: fsdp2, FSDP2 shards both model parameters and optimizer states across the DP group, so memory usage shrinks as you add GPUs:

1distributed:
2 strategy: fsdp2
3 dp_size: null # auto-detected from world_size ÷ (tp × cp × pp)
4 tp_size: 1
5 cp_size: 1

Tensor Parallelism

TP splits weight matrices across GPUs within a single node. Set tp_size to the number of GPUs you want to shard over (typically 2, 4, or 8, and should divide evenly into the number of attention heads):

1distributed:
2 strategy: fsdp2
3 dp_size: null
4 tp_size: 4
5 cp_size: 1
6 sequence_parallel: false # set to true for additional memory savings

sequence_parallel: true extends TP to also shard activation memory along the sequence dimension, further reducing per-GPU memory at the cost of additional communication.

Pipeline Parallelism

PP assigns groups of layers to different GPUs and streams micro-batches through the stages. It requires an additional nested pipeline: section:

1distributed:
2 strategy: fsdp2
3 dp_size: null
4 tp_size: 4
5 pp_size: 4
6 cp_size: 1
7 activation_checkpointing: true
8
9 pipeline:
10 pp_schedule: interleaved1f1b # pipeline schedule (1f1b or interleaved1f1b)
11 pp_microbatch_size: 1 # micro-batch size per pipeline step
12 layers_per_stage: 4 # how many layers each stage handles
13 scale_grads_in_schedule: false
KeyRole
pp_scheduleThe micro-batch schedule. 1f1b is simpler; interleaved1f1b overlaps compute and communication for better throughput.
pp_microbatch_sizeNumber of samples per micro-batch fed into the pipeline. Must satisfy: local_batch_size ÷ pp_microbatch_size ≥ pp_size.
layers_per_stageHow many transformer layers each pipeline stage contains. If omitted, the framework splits layers evenly across pp_size stages.

PP is available only when the model declares pipeline-parallel support and defines a _pp_plan that tells the framework how to split layers into stages. Do not assume every built-in model supports PP; check the model’s declared capabilities before setting pp_size > 1.

Context Parallelism

CP splits the sequence across GPUs, which is useful for very long contexts that exceed single-GPU memory. Set cp_size to the target split factor:

1distributed:
2 strategy: fsdp2
3 dp_size: null
4 tp_size: 1
5 cp_size: 2

When cp_size > 1, fused RoPE is automatically disabled. Some models also require the Transformer Engine (TE) attention backend for CP with packed sequences. The framework will raise an error with instructions if this applies.

Expert Parallelism for MoE Models

EP distributes MoE experts across GPUs. Set ep_size to the number of GPUs that share the full set of experts:

1distributed:
2 strategy: fsdp2
3 tp_size: 1
4 cp_size: 1
5 pp_size: 1
6 ep_size: 8
7 activation_checkpointing: true

EP only applies to Mixture-of-Experts models (such as Qwen3-MoE, Mixtral, and DeepSeek-V3). For dense models, leave ep_size at 1 or omit it.

Combine Multiple Dimensions

You can combine TP, PP, CP, and EP in a single config. For example, a large MoE model on a multi-node cluster might use:

1distributed:
2 strategy: fsdp2
3 dp_size: null
4 tp_size: 1
5 cp_size: 2
6 pp_size: 1
7 ep_size: 4
8 activation_checkpointing: true

When choosing a combination, keep these rules in mind:

  • world_size must be evenly divisible by pp_size × tp_size × cp_size so that the quotient is an integer dp_size.
  • (dp_size × cp_size) % ep_size == 0, where EP shares the DP×CP groups.
  • TP within a node, PP across nodes is the typical layout. TP requires fast NVLink bandwidth, whereas PP tolerates higher latency.
  • Start simple. Use DP-only first. Add TP if the model doesn’t fit on one GPU. Add PP for very large models. Add CP for long sequences. Add EP only for MoE architectures.

Next Steps