From 9ad81414c3f5ca606fe506340530aa5580309c77 Mon Sep 17 00:00:00 2001 From: Alexandros Koumparoulis Date: Fri, 3 Apr 2026 15:34:06 -0700 Subject: [PATCH] update the finetune guide to include more information to explain YAML fields more Signed-off-by: Alexandros Koumparoulis --- docs/guides/llm/finetune.md | 538 ++++++++++++++++++++++++++++-------- 1 file changed, 426 insertions(+), 112 deletions(-) diff --git a/docs/guides/llm/finetune.md b/docs/guides/llm/finetune.md index baa4417731..956b706c4a 100644 --- a/docs/guides/llm/finetune.md +++ b/docs/guides/llm/finetune.md @@ -15,7 +15,7 @@ NeMo AutoModel supports two fine-tuning modes: ┌──────────────┐ ┌──────────────┐ ┌──────────────┐ ┌──────────────┐ ┌──────────────┐ ┌──────────────┐ ┌──────────────┐ │ 1. Install │--->│ 2. Configure │--->│ 3. Train │--->│ 4. Inference │--->│ 5. Evaluate │--->│ 6. Publish │--->│ 7. Deploy │ │ │ │ │ │ │ │ │ │ │ │ (optional) │ │ (optional) │ -│ pip install │ │ YAML recipe │ │ automodel CLI│ │ HF generate │ │ Val loss + │ │ HF Hub │ │ vLLM serving │ +│ pip install │ │ YAML config │ │ automodel CLI│ │ HF generate │ │ Val loss + │ │ HF Hub │ │ vLLM serving │ │ or Docker │ │ Choose SFT │ │ or torchrun │ │ API │ │ lm-eval- │ │ upload │ │ │ │ │ │ or PEFT │ │ │ │ │ │ harness │ │ │ │ │ └──────────────┘ └──────────────┘ └──────────────┘ └──────────────┘ └──────────────┘ └──────────────┘ └──────────────┘ @@ -41,11 +41,11 @@ Alternatively, if you run into dependency or driver issues, use the pre-built Do ```bash docker pull nvcr.io/nvidia/nemo-automodel:26.02.00 -docker run --gpus all -it --rm --shm-size=8g nvcr.io/nvidia/nemo-automodel:26.02.00 +docker run --gpus all -it --rm --shm-size=8g -v $(pwd)/checkpoints:/tmp/checkpoints/ nvcr.io/nvidia/nemo-automodel:26.02.00 ``` :::{important} -**Docker users:** Checkpoints are lost when the container exits unless you bind-mount the checkpoint directory to the host. See [Install with NeMo Docker Container](../installation.md#install-with-nemo-docker-container) and [Saving Checkpoints When Using Docker](../checkpointing.md#saving-checkpoints-when-using-docker). +Docker containers are ephemeral — 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 [Saving Checkpoints When Using Docker](../checkpointing.md#saving-checkpoints-when-using-docker). ::: For the full set of installation methods, see the [installation guide](../installation.md). @@ -53,9 +53,88 @@ For the full set of installation methods, see the [installation guide](../instal ## Configure Your Training Recipe -Training is configured through a [YAML](https://en.wikipedia.org/wiki/YAML) file with three required sections — **model**, **dataset**, and **step_scheduler** — plus an optional **peft** section that fully describes a fine-tuning run. The sections below walk through each one. For the complete copy-pastable file, see [Full Recipe YAML](#full-recipe-yaml). Both SFT and PEFT are driven by a **recipe** — a self-contained Python module that wires together model loading, dataset preparation, training, checkpointing, and logging (see the `TrainFinetuneRecipeForNextTokenPrediction` [source](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/recipes/llm/train_ft.py)). +Training is configured through a [YAML](https://en.wikipedia.org/wiki/YAML) config file with three required sections — **model**, **dataset**, and **step_scheduler** — plus an optional **peft** section. The sections below walk through each one. For the complete copy-pastable file, see [Full Config YAML](#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`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/recipes/llm/train_ft.py). The config file tells the recipe *what* to build; the recipe decides *how* to build it. +:::{details} How the Config System Works +:class: note + +NeMo AutoModel configs use a convention borrowed from [Hydra](https://hydra.cc/): 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: + +```yaml +optimizer: + _target_: torch.optim.Adam + lr: 1.0e-5 + weight_decay: 0 +``` + +is equivalent to writing this Python code: + +```python +from torch.optim import Adam + +optimizer = Adam(lr=1.0e-5, weight_decay=0) +``` + +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. + +:::{tip} +To discover which parameters a section accepts, look up the Python signature of its `_target_`. For instance, `torch.optim.Adam` accepts `lr`, `betas`, `eps`, and `weight_decay` — those are the keys you can set in the YAML. +::: + +**From YAML to running code.** Here is the path a config takes through the framework: + +```text +finetune_config.yaml + │ + ▼ + ┌──────────────┐ load_yaml_config() parses the file into + │ ConfigNode │◄─── a tree of ConfigNode objects, one per + └──────┬───────┘ YAML section. + │ + ▼ + ┌──────────────┐ The recipe's setup() method reads + │ Recipe │◄─── each section from the ConfigNode tree + │ setup() │ and passes it to the matching builder. + └──────┬───────┘ + │ + ┌────┴─────────────────────────────────┐ + ▼ ▼ ▼ ▼ +build_model build_optimizer build_dataloader build_loss_fn ... + │ │ │ │ + ▼ ▼ ▼ ▼ +cfg.model cfg.optimizer cfg.dataset cfg.loss_fn + .instantiate() .instantiate() .instantiate() .instantiate() + │ │ │ │ + ▼ ▼ ▼ ▼ + Resolves Resolves Resolves Resolves + _target_, _target_, _target_, _target_, + calls it calls it calls it calls it + with kwargs with kwargs with kwargs with kwargs +``` + +Each builder function calls **`.instantiate()`** on its config section. `.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 `collate_fn` inside `dataloader`) are recursively instantiated the same way. + +**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** — both resolve to the same class: + +```yaml +# Short name (the CLI looks up the class automatically) +recipe: TrainFinetuneRecipeForNextTokenPrediction + +# Fully-qualified path (used as-is) +recipe: 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_`.** Some sections like `step_scheduler`, `distributed`, and `checkpoint` are plain key-value groups consumed directly by the recipe — they control training schedule, parallelism strategy, and checkpoint behavior without instantiating a Python object. +::: ### Model @@ -65,10 +144,15 @@ model: pretrained_model_name_or_path: meta-llama/Llama-3.2-1B ``` -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](https://github.com/NVIDIA-NeMo/Automodel/blob/main/docs/model-coverage/llm.md). +| Key | Role | +|-----|------| +| `_target_` | Points to [`NeMoAutoModelForCausalLM.from_pretrained`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/_transformers/auto_model.py) — 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_path` | A keyword argument to `from_pretrained`. Any argument that [`from_pretrained`](https://huggingface.co/docs/transformers/main_classes/model#transformers.PreTrainedModel.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](https://github.com/NVIDIA-NeMo/Automodel/blob/main/docs/model-coverage/llm.md). -:::{dropdown} About LLaMA 3.2 1B -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. +:::{dropdown} About Llama 3.2 1B +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. ::: :::{dropdown} Accessing gated models @@ -94,6 +178,11 @@ validation_dataset: split: validation ``` +| Key | Role | +|-----|------| +| `_target_` | Points to [`make_squad_dataset`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/components/datasets/llm/squad.py) — 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.md)). | +| `dataset_name`, `split` | Keyword 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](dataset.md) and [Dataset Overview](../dataset-overview.md). :::{dropdown} About SQuAD v1.1 @@ -119,8 +208,42 @@ peft: alpha: 32 # scaling factor for learned weights ``` +| Key | Role | +|-----|------| +| `_target_` | Points to [`PeftConfig`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/components/_peft/lora.py) — 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_modules` | A glob pattern matched against fully-qualified layer names (e.g. `"*.proj"` matches every layer whose name ends in `proj`). | +| `dim` | The low-rank dimension *r* — controls adapter capacity. Larger values learn more but use more memory. | +| `alpha` | Scaling 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 [Switching Between SFT and PEFT](#switching-between-sft-and-peft). +#### QLoRA (Quantized Low-Rank Adaptation) + +If GPU memory is a constraint, [QLoRA](https://arxiv.org/abs/2305.14314) 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. + +To enable QLoRA, add a `quantization:` section alongside the `peft:` section in your config. Note two differences from the standard PEFT config above: `target_modules` uses the broader `"*_proj"` pattern to apply LoRA to all projection layers (wider coverage compensates for precision loss from 4-bit weights), and `dim` is increased from 8 to 16 for additional adapter capacity. + +```yaml +model: + _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained + pretrained_model_name_or_path: meta-llama/Llama-3.2-1B + +peft: + _target_: nemo_automodel.components._peft.lora.PeftConfig + target_modules: "*_proj" # broader glob than "*.proj" to cover all projection layers + dim: 16 # LoRA rank (higher than default to offset quantization) + alpha: 32 # scaling factor + dropout: 0.1 # LoRA dropout rate + +quantization: + load_in_4bit: True # enable 4-bit quantization + load_in_8bit: False # use 4-bit, not 8-bit + bnb_4bit_compute_dtype: bfloat16 # compute dtype + bnb_4bit_use_double_quant: True # double quantization for extra savings + bnb_4bit_quant_type: nf4 # NormalFloat quantization type + bnb_4bit_quant_storage: bfloat16 # storage dtype for quantized weights +``` + ### Training Schedule ```yaml @@ -128,12 +251,14 @@ step_scheduler: num_epochs: 1 # Will train over the dataset once. ``` +Unlike the sections above, `step_scheduler` has **no `_target_`** — it is not instantiated into a Python object. Instead, the recipe reads its keys directly to control the training loop (how many epochs to run, when to checkpoint, when to validate). This is typical of sections that configure *behavior* rather than *components*. + All other settings (distributed strategy, optimizer, checkpointing, logging) use sensible defaults. See the [Full Configuration Reference](#full-configuration-reference) to customize them. -### Full Recipe YAML +### Full Config YAML :::{dropdown} finetune_config.yaml (click to expand) -Save as `finetune_config.yaml`. This config runs PEFT (LoRA). To run SFT instead, remove the `peft:` section. +Save as `finetune_config.yaml`. This config runs PEFT (LoRA). To run SFT instead, remove the `peft:` section. For production-ready examples, see the hosted configs: [Llama 3.2 1B SFT](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/llm_finetune/llama3_2/llama3_2_1b_squad.yaml) and [Llama 3.2 1B PEFT](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/llm_finetune/llama3_2/llama3_2_1b_squad_peft.yaml). ```yaml model: @@ -166,7 +291,7 @@ step_scheduler: You can run the recipe using the AutoModel CLI or directly with `torchrun` (advanced). ```bash -automodel --nproc-per-node=8 sft_guide.yaml +automodel --nproc-per-node=8 finetune_config.yaml ``` The `--nproc-per-node=8` flag specifies the number of GPUs per node. Adjust to your case (for a single GPU, omit the `--nproc-per-node` option). @@ -176,18 +301,18 @@ The `--nproc-per-node=8` flag specifies the number of GPUs per node. Adjust to y Alternatively, you can invoke the recipe [script](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/recipes/llm/train_ft.py) directly using [torchrun](https://docs.pytorch.org/docs/stable/elastic/run.html), as shown below. ``` bash -torchrun --nproc-per-node=8 nemo_automodel/recipes/llm/train_ft.py -c sft_guide.yaml +torchrun --nproc-per-node=8 nemo_automodel/recipes/llm/train_ft.py -c finetune_config.yaml ``` ### Sample Output Running the recipe using either the `automodel` app or by directly invoking the recipe script should produce the following log: ``` -$ automodel sft_guide.yaml -INFO:nemo_automodel.cli.app:Config: /mnt/4tb/auto/Automodel/sft_guide.yaml +$ 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: /mnt/4tb/auto/Automodel/sft_guide.yaml +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 @@ -202,7 +327,7 @@ Each log line reports the current loss, gradient norm, peak GPU memory, and toke ### Checkpoint Contents -Checkpoints are saved in native Hugging Face format, so no conversion is required — they work directly with Transformers, PEFT, vLLM, lm-eval-harness, and other tools in the Hugging Face ecosystem. SFT and PEFT produce different checkpoint layouts. **SFT checkpoints** contain the full model weights at `model/consolidated/` and can be loaded directly. **PEFT checkpoints** contain only the adapter weights (~MBs instead of GBs) — at inference time you must load the original base model and apply the adapter on top. This distinction affects every downstream step (inference, publishing, deployment). +Checkpoints are saved in native Hugging Face format, so no conversion is required — they work directly with Transformers, PEFT, vLLM, lm-eval-harness, and other tools in the Hugging Face ecosystem. SFT and PEFT produce different checkpoint layouts. **SFT checkpoints** contain the full model weights at `model/consolidated/` — a single, self-contained Hugging Face model directory created by gathering distributed shards into one location — and can be loaded directly. **PEFT checkpoints** contain only the adapter weights (~MBs instead of GBs) — at inference time you must load the original base model and apply the adapter on top. This distinction affects every downstream step (inference, publishing, deployment). :::{dropdown} Checkpoint directory structure **SFT checkpoint:** @@ -255,21 +380,17 @@ checkpoints/epoch_0_step_10/ Inference uses the Hugging Face `generate` API. Because SFT checkpoints are self-contained while PEFT checkpoints store only adapter weights (see [Checkpoint Contents](#checkpoint-contents)), the loading procedure differs between the two modes. -### PEFT Inference +### SFT Inference -PEFT adapters must be loaded on top of the base model: +The SFT checkpoint at `model/consolidated/` is a complete Hugging Face model and can be loaded directly: ```python import torch from transformers import AutoModelForCausalLM, AutoTokenizer -from peft import PeftModel -base_model_name = "meta-llama/Llama-3.2-1B" -tokenizer = AutoTokenizer.from_pretrained(base_model_name) -model = AutoModelForCausalLM.from_pretrained(base_model_name) - -adapter_path = "checkpoints/epoch_0_step_10/model/" -model = PeftModel.from_pretrained(model, adapter_path) +ckpt_path = "checkpoints/epoch_0_step_10/model/consolidated" +tokenizer = AutoTokenizer.from_pretrained(ckpt_path) +model = AutoModelForCausalLM.from_pretrained(ckpt_path) device = "cuda" if torch.cuda.is_available() else "cpu" model.to(device) @@ -287,17 +408,21 @@ output = model.generate(**inputs, max_new_tokens=50) print(tokenizer.decode(output[0], skip_special_tokens=True)) ``` -### SFT Inference +### PEFT Inference -The SFT checkpoint at `model/consolidated/` is a complete Hugging Face model and can be loaded directly: +PEFT adapters must be loaded on top of the base model: ```python import torch from transformers import AutoModelForCausalLM, AutoTokenizer +from peft import PeftModel -ckpt_path = "checkpoints/epoch_0_step_10/model/consolidated" -tokenizer = AutoTokenizer.from_pretrained(ckpt_path) -model = AutoModelForCausalLM.from_pretrained(ckpt_path) +base_model_name = "meta-llama/Llama-3.2-1B" +tokenizer = AutoTokenizer.from_pretrained(base_model_name) +model = AutoModelForCausalLM.from_pretrained(base_model_name) + +adapter_path = "checkpoints/epoch_0_step_10/model/" +model = PeftModel.from_pretrained(model, adapter_path) device = "cuda" if torch.cuda.is_available() else "cpu" model.to(device) @@ -479,7 +604,7 @@ The `peft:` section controls which mode runs: | Mode | What to do in the YAML | |------|----------------------| | **PEFT (LoRA)** | Include the `peft:` section as shown below. | -| **SFT (full-parameter)** | Remove the `peft:` section entirely. | +| **SFT (full-parameter)** | Remove/comment the `peft:` section entirely. | All other config sections remain the same for both modes. @@ -488,93 +613,162 @@ All other config sections remain the same for both modes. :::{dropdown} Full Config :open: ```yaml -# ── Model ── +# Recipe +# Selects which recipe class runs the training loop. +# Use a short name (auto-discovered) or a fully-qualified Python path: +# recipe: nemo_automodel.recipes.llm.train_ft.TrainFinetuneRecipeForNextTokenPrediction +recipe: TrainFinetuneRecipeForNextTokenPrediction + +# Training Schedule +# Controls epoch count, batch sizes, and how often to checkpoint / validate. +# No _target_ — these are plain values read directly by the recipe. +step_scheduler: + grad_acc_steps: 4 # number of micro-batches accumulated before each optimizer + # step. Effective batch = grad_acc_steps × batch_size. + ckpt_every_steps: 10 # save a checkpoint every N gradient steps + val_every_steps: 10 # run the validation loop every N gradient steps + num_epochs: 1 # how many full passes over the training dataset + +# Process Group +# Initializes the PyTorch distributed process group. +# No _target_ — consumed directly by the recipe. +# You normally would not need to tune this. +dist_env: + backend: nccl # communication backend: "nccl" (GPU, recommended) or "gloo" (CPU) + timeout_minutes: 1 # timeout for collective operations; increase for large models + # that take longer to initialize + +# Distributed Strategy +# Determines how model weights, data, and compute are split across GPUs. +# No _target_ — consumed directly by the recipe. +# See "Distributed Training: TP, PP, CP, and EP" in Advanced Topics for details. +distributed: + strategy: fsdp2 # parallelism strategy: "fsdp2" (recommended), "megatron_fsdp", + # or "ddp". FSDP2 shards parameters and optimizer states across + # the data-parallel group. + dp_size: null # data-parallel group size. null = auto-detect from + # world_size ÷ (tp_size × cp_size × pp_size). + tp_size: 1 # tensor-parallel size: splits weight matrices across GPUs. + # Set to 2, 4, or 8 if the model doesn't fit on one GPU. + # Should divide evenly into the number of attention heads. + cp_size: 1 # context-parallel size: splits the input sequence across GPUs. + # Increase for very long contexts (e.g. 32k+ tokens). + sequence_parallel: false # when true, extends TP to also shard activations along + # the sequence dimension for additional memory savings + +# Random Number Generator +# _target_ → StatefulRNG: a checkpointable RNG that ensures identical sequences +# across training restarts. Seed and ranked are kwargs to StatefulRNG(). +rng: + _target_: nemo_automodel.components.training.rng.StatefulRNG + seed: 1111 # global random seed for reproducibility + ranked: true # when true, each GPU rank gets a unique RNG stream derived + # from the seed, so data shuffling differs per GPU + +# Model +# _target_ → NeMoAutoModelForCausalLM.from_pretrained: downloads (or loads from +# cache) a pretrained HuggingFace model and wraps it with NeMo distributed-training +# support. Any from_pretrained kwarg is accepted (cache_dir, torch_dtype, etc.). model: _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained pretrained_model_name_or_path: meta-llama/Llama-3.2-1B -# ── PEFT (remove this section entirely for SFT) ── +# PEFT (remove / comment this entire section for full-parameter SFT) +# _target_ → PeftConfig: a dataclass describing which layers get LoRA adapters. +# The recipe passes this config into build_model(), which attaches adapters +# to the matching layers. peft: _target_: nemo_automodel.components._peft.lora.PeftConfig - target_modules: "*.proj" # glob pattern matching linear layer FQNs - dim: 8 # low-rank dimension of the adapters - alpha: 32 # scaling factor for learned weights - use_triton: True # use optimized Triton-based LoRA kernel (requires triton) - -# ── Dataset ── + target_modules: "*.proj" # glob pattern matched against fully-qualified layer names; + # "*.proj" matches every layer ending in "proj" + dim: 8 # low-rank dimension r — controls adapter capacity. + # Larger values are more expressive but use more memory. + alpha: 32 # LoRA scaling factor: adapter output is scaled by alpha/dim. + # Higher values give adapters more influence during training. + use_triton: True # use an optimized Triton kernel for LoRA forward/backward + # (requires the triton package) + +# Checkpointing +# No _target_ — plain key-value group consumed by the recipe. +checkpoint: + enabled: true # set to false to skip saving checkpoints entirely + checkpoint_dir: checkpoints/ # output directory. Docker users: bind-mount this path + # (e.g. -v $(pwd)/checkpoints:/workspace/checkpoints) + # to persist checkpoints across container restarts. + model_save_format: safetensors # "safetensors" (recommended, faster and safer) or + # "torch_save" (legacy pickle-based format) + save_consolidated: True # when true, writes a single HuggingFace-compatible checkpoint + # to model/consolidated/ that can be loaded directly by + # Transformers, vLLM, etc. Requires safetensors format. + +# Training Dataset +# _target_ → make_squad_dataset: a factory function that downloads the SQuAD +# dataset, tokenizes it, and returns a torch Dataset. To use a different dataset, +# change _target_ to another factory function (see the dataset guide). dataset: _target_: nemo_automodel.components.datasets.llm.squad.make_squad_dataset - dataset_name: rajpurkar/squad - split: train + dataset_name: rajpurkar/squad # HuggingFace Hub dataset ID + split: train # which split to use (train, validation, test) +# Validation Dataset validation_dataset: _target_: nemo_automodel.components.datasets.llm.squad.make_squad_dataset dataset_name: rajpurkar/squad split: validation - limit_dataset_samples: 64 - -# ── Training schedule ── -step_scheduler: - grad_acc_steps: 4 # micro-batches accumulated before each optimizer step - ckpt_every_steps: 10 # save checkpoint every N gradient steps - val_every_steps: 10 # run validation every N gradient steps - num_epochs: 1 + limit_dataset_samples: 64 # cap validation set to 64 samples for faster eval loops; + # remove this line to use the full validation set -# ── Distributed ── -dist_env: - backend: nccl - timeout_minutes: 1 - -distributed: - strategy: fsdp2 - dp_size: null # auto-detected from world size, tp_size, and cp_size - tp_size: 1 - cp_size: 1 - sequence_parallel: false - -# ── RNG ── -rng: - _target_: nemo_automodel.components.training.rng.StatefulRNG - seed: 1111 - ranked: true - -# ── Loss ── -loss_fn: - _target_: nemo_automodel.components.loss.masked_ce.MaskedCrossEntropy - -# ── Dataloaders ── +# Training Dataloader +# _target_ → StatefulDataLoader: a checkpointable DataLoader from torchdata that +# saves and restores iteration state across training restarts, so resumed runs +# don't re-process already-seen batches. dataloader: _target_: torchdata.stateful_dataloader.StatefulDataLoader collate_fn: nemo_automodel.components.datasets.utils.default_collater - batch_size: 8 - shuffle: false + # function that pads and batches individual samples + # into tensors; can be swapped for custom collation + batch_size: 8 # samples per micro-batch per GPU + shuffle: true # whether to shuffle the dataset each epoch +# Validation Dataloader validation_dataloader: _target_: torchdata.stateful_dataloader.StatefulDataLoader collate_fn: nemo_automodel.components.datasets.utils.default_collater batch_size: 8 -# ── Checkpointing ── -checkpoint: - enabled: true - checkpoint_dir: checkpoints/ - model_save_format: safetensors - save_consolidated: True # single HF-compatible bundle (requires safetensors format) +# Loss Function +# _target_ → MaskedCrossEntropy: standard cross-entropy loss that automatically +# ignores padding tokens so they don't affect the gradient. +# Other available loss functions (swap _target_ to use): +# - nemo_automodel.components.loss.chunked_ce.ChunkedCrossEntropy +# Computes CE in chunks along the sequence dimension to reduce peak memory. +# Useful for very long sequences. Accepts chunk_len (default 32). +# - nemo_automodel.components.loss.linear_ce.FusedLinearCrossEntropy +# Fuses the final linear projection (lm_head) with the CE computation, +# avoiding the full logit tensor. Significant **memory savings** for large vocabs. +# - nemo_automodel.components.loss.te_parallel_ce.TEParallelCrossEntropy +# TE-based parallel CE with a Triton kernel. Designed for tensor-parallel +# setups where logits are sharded across TP ranks. +loss_fn: + _target_: nemo_automodel.components.loss.masked_ce.MaskedCrossEntropy -# ── Optimizer ── +# Optimizer +# _target_ → torch.optim.Adam: any torch.optim class can be used here (e.g. +# AdamW, SGD). All remaining keys become kwargs to the constructor. optimizer: _target_: torch.optim.Adam - betas: [0.9, 0.999] - eps: 1e-8 - lr: 1.0e-5 - weight_decay: 0 + lr: 1.0e-5 # learning rate — the most important hyperparameter to tune + betas: [0.9, 0.999] # Adam momentum coefficients (β₁ for mean, β₂ for variance) + eps: 1e-8 # small constant added to the denominator for numerical stability + weight_decay: 0 # L2 regularization strength (0 = no regularization) -# ── Logging (optional) ── +# Logging (optional) +# Uncomment to enable Weights & Biases experiment tracking. # wandb: -# project: -# entity: -# name: -# save_dir: +# project: # W&B project name +# entity: # W&B team or username +# name: # display name for this run +# save_dir: # local directory for W&B artifacts ``` ::: @@ -582,40 +776,160 @@ optimizer: | Section | Required? | What to change | |---------|-----------|----------------| -| `model` | Yes | Set `pretrained_model_name_or_path` to your Hugging Face model ID. | -| `peft` | PEFT only | Remove 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](#qlora-quantized-low-rank-adaptation) below. | -| `dataset` | Yes | Change `_target_`, `dataset_name`, and `split` for your data. | +| `model` | Yes | Set `pretrained_model_name_or_path` to your Hugging Face model ID. Source: [`auto_model.py`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/_transformers/auto_model.py). | +| `peft` | PEFT only | Remove 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](#qlora-quantized-low-rank-adaptation). Source: [`lora.py`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/components/_peft/lora.py). | +| `dataset` | Yes | Change `_target_`, `dataset_name`, and `split` for your data. Source: [`squad.py`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/components/datasets/llm/squad.py). | +| `dataloader` | Optional | Adjust `batch_size` and `shuffle`. Uses [`StatefulDataLoader`](https://pytorch.org/data/stable/stateful_dataloader.html) for checkpointable iteration. Collation: [`utils.py`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/components/datasets/utils.py). | +| `loss_fn` | Optional | Default is [`MaskedCrossEntropy`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/components/loss/masked_ce.py). Alternatives: [`ChunkedCrossEntropy`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/components/loss/chunked_ce.py) (long sequences), [`FusedLinearCrossEntropy`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/components/loss/linear_ce.py) (large vocabs), [`TEParallelCrossEntropy`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/components/loss/te_parallel_ce.py) (tensor-parallel). | +| `rng` | Optional | Controls reproducibility. Source: [`rng.py`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/components/training/rng.py). | | `step_scheduler` | Yes | `grad_acc_steps` sets how many micro-batches accumulate per gradient step. `ckpt_every_steps` and `val_every_steps` are counted in gradient steps. | | `distributed` | Yes | `dp_size: null` means auto-detect from world size. Adjust `tp_size` for tensor parallelism across GPUs. | | `checkpoint` | Recommended | Set `checkpoint_dir` to a persistent path, especially in Docker. | | `optimizer` | Optional | Defaults are reasonable. Any `torch.optim` class can be substituted via `_target_`. | | `wandb` | Optional | Uncomment and configure to enable Weights & Biases logging. | -## Advanced Topics +For the fine-tuning recipe itself, see [`train_ft.py`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/recipes/llm/train_ft.py). For more example configs, browse [`examples/llm_finetune/`](https://github.com/NVIDIA-NeMo/Automodel/tree/main/examples/llm_finetune). -### QLoRA (Quantized Low-Rank Adaptation) +## Distributed Training: TP, PP, CP, and EP -If GPU memory is a constraint, [QLoRA](https://arxiv.org/abs/2305.14314) combines LoRA with 4-bit NormalFloat (NF4) quantization to reduce memory usage by up to 75% compared to full-precision fine-tuning, while maintaining comparable quality to standard LoRA. +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: -To enable QLoRA, add a `quantization:` section alongside the `peft:` section in your config. Note two differences from the standard PEFT config above: `target_modules` uses the broader `"*_proj"` pattern to apply LoRA to all projection layers (wider coverage compensates for precision loss from 4-bit weights), and `dim` is increased from 8 to 16 for additional adapter capacity. +| Dimension | Key | What it shards | When to use | +|-----------|-----|---------------|-------------| +| **Data Parallel (DP)** | `dp_size` | Replicates the model on each group of GPUs; each replica trains on a different data batch. | Default. Scales batch size linearly with GPU count. | +| **Tensor Parallel (TP)** | `tp_size` | Splits 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 Parallel (PP)** | `pp_size` | Assigns different *layers* (stages) to different GPUs and pipelines micro-batches through them. | Very deep models that don't fit even with TP, or multi-node training where TP's all-reduce is too expensive across nodes. | +| **Context Parallel (CP)** | `cp_size` | Splits the input *sequence* across GPUs so each GPU processes a portion of the context. | Very long sequences that exceed single-GPU memory. | +| **Expert Parallel (EP)** | `ep_size` | Distributes 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: + +```text +world_size = pp_size × dp_size × cp_size × tp_size +``` + +When `dp_size` is set to `null` (the default), it is inferred automatically: + +```text +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 Parallel (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: ```yaml -model: - _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained - pretrained_model_name_or_path: meta-llama/Llama-3.2-1B +distributed: + strategy: fsdp2 + dp_size: null # auto-detected from world_size ÷ (tp × cp × pp) + tp_size: 1 + cp_size: 1 +``` -peft: - _target_: nemo_automodel.components._peft.lora.PeftConfig - target_modules: "*_proj" # broader glob than "*.proj" to cover all projection layers - dim: 16 # LoRA rank (higher than default to offset quantization) - alpha: 32 # scaling factor - dropout: 0.1 # LoRA dropout rate +#### Tensor Parallelism -quantization: - load_in_4bit: True # enable 4-bit quantization - load_in_8bit: False # use 4-bit, not 8-bit - bnb_4bit_compute_dtype: bfloat16 # compute dtype - bnb_4bit_use_double_quant: True # double quantization for extra savings - bnb_4bit_quant_type: nf4 # NormalFloat quantization type - bnb_4bit_quant_storage: bfloat16 # storage dtype for quantized weights +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 — should divide evenly into the number of attention heads): + +```yaml +distributed: + strategy: fsdp2 + dp_size: null + tp_size: 4 + cp_size: 1 + 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: + +```yaml +distributed: + strategy: fsdp2 + dp_size: null + tp_size: 4 + pp_size: 4 + cp_size: 1 + activation_checkpointing: true + + pipeline: + pp_schedule: interleaved1f1b # pipeline schedule (1f1b or interleaved1f1b) + pp_microbatch_size: 1 # micro-batch size per pipeline step + layers_per_stage: 4 # how many layers each stage handles + scale_grads_in_schedule: false +``` + +| Key | Role | +|-----|------| +| `pp_schedule` | The micro-batch schedule. `1f1b` is simpler; `interleaved1f1b` overlaps compute and communication for better throughput. | +| `pp_microbatch_size` | Number of samples per micro-batch fed into the pipeline. Must satisfy: `local_batch_size ÷ pp_microbatch_size ≥ pp_size`. | +| `layers_per_stage` | How many transformer layers each pipeline stage contains. If omitted, the framework splits layers evenly across `pp_size` stages. | + +:::{note} +PP requires the model to define a `_pp_plan` that tells the framework how to split layers into stages. All built-in models include this plan; custom models must add one. +::: + +#### Context Parallelism + +CP splits the sequence across GPUs — useful for very long contexts that exceed single-GPU memory. Set `cp_size` to the desired split factor: + +```yaml +distributed: + strategy: fsdp2 + dp_size: null + tp_size: 1 + cp_size: 2 +``` + +:::{important} +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 (MoE models) + +EP distributes MoE experts across GPUs. Set `ep_size` to the number of GPUs that share the full set of experts: + +```yaml +distributed: + strategy: fsdp2 + tp_size: 1 + cp_size: 1 + pp_size: 1 + ep_size: 8 + activation_checkpointing: true ``` + +EP only applies to Mixture-of-Experts models (e.g. Qwen3-MoE, Mixtral, DeepSeek-V3). For dense models, leave `ep_size` at `1` or omit it. + +#### Combining 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: + +```yaml +distributed: + strategy: fsdp2 + dp_size: null + tp_size: 1 + cp_size: 2 + pp_size: 1 + ep_size: 4 + activation_checkpointing: true +``` + +When choosing a combination, keep these rules in mind: + +- **`world_size` must divide evenly** into `pp_size × tp_size × cp_size` (the remainder becomes `dp_size`). +- **`(dp_size × cp_size) % ep_size == 0`** — EP shares the DP×CP groups. +- **TP within a node, PP across nodes** is the typical layout — TP requires fast NVLink bandwidth, while 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 + +- [Integrate Your Own Text Dataset](dataset.md) — swap the SQuAD example for your own data. +- [Recipes and End-to-End Examples](../overview.md) — browse the full set of recipes available in NeMo AutoModel. See also the [`examples/llm_finetune/`](https://github.com/NVIDIA-NeMo/Automodel/tree/main/examples/llm_finetune) directory for ready-to-run configs. +- [Dataset Overview](../dataset-overview.md) — see all supported dataset types across LLM, VLM, and retrieval tasks. +- [Knowledge Distillation](knowledge-distillation.md) — distill a fine-tuned model into a smaller one.