Get Started

YAML Configuration

View as Markdown

NeMo AutoModel recipes use YAML configs. The YAML parser creates a ConfigNode that performs the following actions:

  • Translates common scalar strings into typed Python values (for example, "10" → 10)
  • Resolves _target_ and *_fn into Python callables and classes
  • Supports environment variable interpolation inside YAML strings
  • Makes config printing safer by preserving original placeholders to avoid leaking secrets

Load Model and Dataset Configs

Most recipes load YAML configs using nemo_automodel.components.config.loader.load_yaml_config(), which returns a ConfigNode.

A ConfigNode handles values as follows:

  • Nested dicts become nested ConfigNode objects
  • Lists are recursively wrapped
  • Scalars are translated with translate_value() when they are YAML strings

Typed Scalar Translation (translate_value)

Only strings are translated:

  • "123" → 123
  • "3.14" → 3.14
  • "true" / "false" → True / False
  • "None" / "none" → None

YAML-native types (such as step_size: 10 without quotes) are already typed by the YAML parser and remain unchanged.

Use _target_ for Instantiation

Any mapping containing a _target_ key can be instantiated using ConfigNode.instantiate().

1model:
2  _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained
3  pretrained_model_name_or_path: meta-llama/Llama-3.2-1B

ConfigNode can also resolve callables from these sources:

  • Dotted paths: pkg.module.symbol
  • Local file paths: /abs/path/to/file.py:symbol

Target Resolution and Code Execution

Target resolution is intentionally permissive and should be treated as code execution:

  • A dotted target can import any package discoverable through Python’s sys.path or reuse a package already present in sys.modules. It is not limited to a safe-prefix allowlist.
  • An existing file.py:symbol target loads that file directly without requiring an environment-variable opt-in. Loading either target form executes the top-level code of the imported module.
  • For file.py:symbol, private or dunder symbols (names beginning with _ or containing __) are rejected. For dotted targets, private attribute traversal is rejected by default and can be enabled with NEMO_ENABLE_USER_MODULES=1 or set_enable_user_modules(True).

Only load configuration files and _target_ values from trusted sources. A YAML file that names an importable package or Python file can execute arbitrary Python code during target resolution.

Distributed Section (Strategy-Based)

The distributed: section is not instantiated using _target_. Recipes parse it with a fixed schema. Use strategy: fsdp2, strategy: ddp, or strategy: megatron_fsdp. You can also configure parallelism sizes, such as dp_size, tp_size, and pp_size, and strategy-specific options. When pipeline parallelism is enabled (pp_size > 1), add a pipeline: subsection with options such as pp_schedule, pp_microbatch_size, and layers_per_stage. For examples, see the Pipeline Parallelism with AutoPipeline guide and the recipe configs.

Prewarm One-Time CUDA Initialization

The LLM training, LLM fine-tuning, and VLM fine-tuning recipes accept an optional prewarm: section.

1prewarm:
2  cublas_backward: true
3  fla_gdn_autotune: true
4  mamba_ssd_autotune: true
5  comm_groups: true

All options default to false. Enable only the warmup associated with an observed first-step failure:

  • cublas_backward initializes cuBLAS backward workspaces during setup.
  • fla_gdn_autotune autotunes flash-linear-attention gated-delta-net Triton kernels during setup.
  • mamba_ssd_autotune autotunes Mamba SSD Triton kernels during setup.
  • comm_groups initializes the process groups used by gradient-norm collectives.

Prewarming is most useful for long-context or otherwise memory-constrained runs. A run can have enough steady-state memory and still fail on its first training step. Activations and gradients are live while cuBLAS workspaces, Triton autotuning buffers, and NCCL communicators are initialized for the first time. Temporary autotuning allocations are released afterward. Some memory, including NCCL scratch space, is outside PyTorch’s caching allocator, so torch.cuda.memory_allocated() does not show the complete transient peak.

Adding GPUs helps only when it sufficiently reduces the relevant per-rank tensors. Each rank still initializes its own resources, and additional parallel dimensions can introduce more communication groups. Shorter contexts can leave enough memory for lazy initialization, while longer contexts might require the same initialization during setup. Long-context hybrid-attention MoE workloads are one common case. Enable each prewarm based on the observed first-step failure rather than a model name. Prewarming does not reduce steady-state memory or imply support for a particular context length.

Triton reuses an autotuned configuration according to each kernel’s explicit cache key. The Mamba SSD warmup therefore matches the real head, state, group, chunk, and dtype geometry but uses batch size 1 and two chunks. Batch size and sequence length change the launch grid and activation footprint, but they are not cache keys in the pinned Mamba SSD kernels. Keeping those dimensions small is intentional: using the full training batch and sequence would recreate the peak-memory condition that setup prewarming is meant to avoid.

Repair Near-Zero Input Embeddings

The LLM training and fine-tuning recipe can repair a small number of damaged input-embedding rows after loading the base checkpoint and before creating the optimizer.

1embedding_row_repair:
2  min_norm: 1.0e-4
3  max_rows: 256

Rows with a non-finite L2 norm or a norm at or below min_norm are replaced with the corresponding output-embedding direction, scaled to the RMS norm of healthy input rows. Setup logs the affected token IDs. max_rows is a safety limit: setup aborts instead of rewriting a broadly damaged or mismatched checkpoint.

This option is intended for diagnosed checkpoint defects where rare token IDs produce extreme input-embedding gradients. It is disabled when the section is absent and is not currently supported with pipeline parallelism.

For diagnosis, safety behavior, verification steps, and performance impact, see Repair Damaged Input-Embedding Rows.

Frozen Multimodal FSDP Policy

FSDP2 uses the correctness-first root policy for fully frozen vision/audio towers and multimodal projectors.

1distributed:
2  strategy: fsdp2
3  multimodal:
4    frozen_sharding: root

The supported policies are:

  • root (default): the always-run outer FSDP root owns frozen multimodal parameters. This keeps collective ordering aligned when some ranks execute a modality branch and others skip it.
  • per_layer: layers inside the frozen multimodal module get their normal FSDP units. This can reduce peak unshard memory, but it is an expert option: every rank in the FSDP group must execute or skip those units the same number of times and in the same order on every microbatch.
  • replicate: frozen multimodal parameters are excluded from FSDP and copied on every rank. This removes their collectives at the cost of higher per-rank parameter memory.

This setting does not change modules that contain any trainable parameters. For nested MoE or VLM models, root requires distributed.moe.wrap_outer_model: true when a fully frozen multimodal module is outside the inner text-model FSDP root. Use per_layer or replicate if the outer model must remain unwrapped.

Interpolate Environment Variables in YAML

NeMo AutoModel supports env var interpolation inside YAML string values.

Supported Forms

  • Braced:
    • ${VAR}
    • ${VAR,default}
    • ${var.dot.var} (dots are treated as part of the env var name)
  • Dollar:
    • $VAR
    • $var.dot.var
  • Back-compat:
    • ${oc.env:VAR}
    • ${oc.env:VAR,default}

Interpolation Behavior

  • Interpolation occurs when values are wrapped into a ConfigNode.
  • If a referenced env var is missing and no default is provided, config loading raises a KeyError.
  • Defaults are supported only for braced forms using the first comma: ${VAR,default_value}.

Example (Databricks Delta)

The following configuration shows an example of environment variable interpolation for a Databricks Delta dataset.

1dataset:
2  _target_: nemo_automodel.components.datasets.llm.column_mapped_text_instruction_iterable_dataset.ColumnMappedTextInstructionIterableDataset
3  path_or_dataset_id: delta://catalog.schema.training_data
4  delta_storage_options:
5    DATABRICKS_HOST: ${DATABRICKS_HOST}
6    DATABRICKS_TOKEN: ${DATABRICKS_TOKEN}
7    DATABRICKS_HTTP_PATH: ${DATABRICKS_HTTP_PATH}

Prevent Secret Leakage in Logs

When an env var placeholder is resolved, the config keeps the original placeholder in an internal ._orig_value field for safe printing.

  • str(cfg) or repr(cfg) prints placeholders (for example, ${DATABRICKS_TOKEN}), not resolved secrets.
  • cfg.to_yaml_dict(use_orig_values=True, redact_sensitive=True) is the recommended way to produce a loggable YAML dict.

Printing a leaf value (for example, print(cfg.dataset.delta_storage_options.DATABRICKS_TOKEN)) outputs the resolved secret. Instead, print the full config or use a redacted YAML dict.

Configure Slurm

Slurm jobs are submitted directly with sbatch. No YAML section is required. Copy the reference script, edit CONFIG and the cluster settings, and then submit the job:

$cp slurm.sub my_cluster.sub
$vim my_cluster.sub
$sbatch my_cluster.sub

All cluster-specific configuration, including SBATCH directives, the container image, mounts, secrets, and environment variables, resides in your sbatch script. See Run on a Cluster for complete examples.