Run with NeMo-Run

View as Markdown

In this guide, you will learn how to launch NeMo AutoModel training jobs using NeMo-Run. NeMo-Run itself supports several executor backends, but AutoModel’s current integration relies on a packaged config at /nemo_run/code/automodel_config.yaml; the Slurm container path is the end-to-end path documented here. The Local execution and Kubernetes executor sections describe current path-staging limitations. For cloud-based training, see Run on Any Cloud with SkyPilot. For direct sbatch usage, see Run on a Cluster (Slurm). For single-node workstation usage, see Run on Your Local Workstation.

NeMo-Run is an open-source tool from NVIDIA that manages job submission across different execution backends. You can define a compute configuration in Python and reuse it across compatible jobs.

Before You Begin

  1. Install the source checkout and its launcher dependencies. The all extra includes AutoModel’s cli extra, which depends on NeMo-Run:
$uv sync --locked --all-groups --extra all
  1. For a named non-local executor, create an executor definitions file at $NEMORUN_HOME/executors.py. NEMORUN_HOME defaults to ~/.nemo_run; set the environment variable to use a different location. Each custom executor name referenced by YAML must be defined here. executor: local is constructed directly and does not read this file. See Executor Setup for a complete example.

  2. Verify connectivity to the target in your executor (e.g. SSH for Slurm, kubeconfig for Kubernetes).

  3. Set required environment variables (if needed by your training config):

$export HF_TOKEN=hf_... # Required for gated models (e.g. Llama)
$export WANDB_API_KEY=... # Optional: Weights & Biases logging

Executor Setup

For a non-local custom executor, the executor: field in YAML names an entry in $NEMORUN_HOME/executors.py. This file must define a module-level EXECUTOR_MAP dictionary. The Slurm example below matches AutoModel’s current packaged-config path. The Kubeflow example documents its native fields and the additional integration constraints that currently prevent the unmodified AutoModel launcher from being an end-to-end Kubeflow path.

Slurm Executor

1import nemo_run as run
2
3def my_slurm_cluster():
4 executor = run.SlurmExecutor(
5 account="my_account",
6 partition="batch",
7 tunnel=run.SSHTunnel(
8 user="myuser",
9 host="login-node.example.com",
10 job_dir="/remote/path/nemo_run/jobs",
11 ),
12 nodes=1,
13 ntasks_per_node=8,
14 gpus_per_node=8,
15 mem="0",
16 exclusive=True,
17 packager=run.Packager(),
18 )
19 executor.container_image = "nvcr.io/nvidia/nemo-automodel:26.02"
20 executor.container_mounts = ["/data:/data", "/checkpoints:/checkpoints"]
21 executor.env_vars = {"HF_HOME": "/data/hf_cache"}
22 executor.time = "04:00:00"
23 return executor
24
25EXECUTOR_MAP = {
26 "my_slurm": my_slurm_cluster(),
27}

Kubernetes Executor

1import nemo_run as run
2
3def my_k8s_cluster():
4 return run.KubeflowExecutor(
5 namespace="training",
6 image="nvcr.io/nvidia/nemo-automodel:26.02",
7 num_nodes=1,
8 nprocs_per_node=8,
9 gpus_per_node=8,
10 workdir_pvc="nemo-run-workdir",
11 workdir_pvc_path="/nemo_run",
12 )
13
14EXECUTOR_MAP = {
15 "my_k8s": my_k8s_cluster(),
16}

KubeflowExecutor requires an environment with the nemo-run[kubeflow] extra, a working kubeconfig or in-cluster configuration, and an existing workdir_pvc. Without workdir_pvc, NeMo-Run’s Kubeflow package() method returns without copying any work directory. With the PVC, that method syncs the executor’s assigned task directory to the username-scoped executor.code_dir (by default /nemo_run/<username>/code); it does not consume AutoModel’s PatternPackager, and AutoModel does not copy its separately written timestamped config into that task directory. The generated command also still uses /nemo_run/code/automodel_config.yaml. The AutoModel launcher does not make that path equivalent to executor.code_dir, so this executor definition alone is not currently runnable through automodel; config staging and the Kubeflow work directory must first be aligned in the integration.

Multiple Executors

You can define as many executors as you need for different backends, clusters, or resource configurations:

1EXECUTOR_MAP = {
2 "slurm_dev": my_slurm_dev(),
3 "slurm_prod": my_slurm_prod(),
4 "k8s": my_k8s_cluster(),
5}
  • Keys in EXECUTOR_MAP are names you reference in YAML (executor: slurm_dev).
  • Values can be executor instances or zero-argument callables that return one.
  • Override fields in YAML are applied on top of executor defaults. Use the executor’s actual field names: for example, Slurm and Local use ntasks_per_node, while Kubeflow uses nprocs_per_node (falling back to gpus_per_node). A devices key is not consumed by these executors.

Quickstart

To submit a compatible AutoModel YAML through the current packaged remote path, add a nemo_run: section that selects a configured Slurm executor. For example, given an existing config that you run locally:

$automodel examples/llm_finetune/qwen/qwen3_moe_30b_te_packed_sequence.yaml

Add a nemo_run: block to submit it to the named Slurm executor instead:

1# -- Add this section to a compatible config ----------------------------------
2nemo_run:
3 executor: my_slurm # Name from EXECUTOR_MAP in $NEMORUN_HOME/executors.py
4 container_image: /images/custom.sqsh # Override executor's default image
5 nodes: 1 # Override number of nodes
6 ntasks_per_node: 8 # Training processes per node
7 time: "04:00:00" # Override time limit
8 job_name: qwen3_moe_finetune # Experiment and job name
9
10# -- Everything below is your existing training config (unchanged) ------------
11recipe: TrainFinetuneRecipeForNextTokenPrediction
12
13step_scheduler:
14 global_batch_size: 32
15 # ... rest of your config ...

Then run the same command:

$automodel your_config.yaml

The CLI detects the nemo_run: key, strips it from the training config, loads the named Slurm executor from $NEMORUN_HOME/executors.py, and submits the job.

Configuration Reference

All nemo_run: Fields

FieldDefaultDescription
executor"local"Name from EXECUTOR_MAP for a custom non-local executor, or "local" to construct LocalExecutor (subject to the config-path limitation below)
job_name<recipe_class_name>Experiment and job name
detachtrueRequested NeMo-Run behavior. NeMo-Run supports detaching for Slurm and selected remote executors, but forces it to false for Local and Kubeflow executors
tail_logsfalseStream logs after submission
executors_file$NEMORUN_HOME/executors.pyPath to custom non-local executor definitions; unused for executor: local
job_dir./nemo_run_jobsLocal directory used to create the timestamped config snapshot and packaging input
(any other key)(from executor)Applied directly with setattr; dicts are merged and lists extended. Use exact native names such as Slurm/Local ntasks_per_node or Kubeflow nprocs_per_node. Unknown names can become inert attributes instead of raising an error.

Examples

Single-Node Fine-Tuning (1 x 8 GPUs)

1nemo_run:
2 executor: my_slurm
3 nodes: 1
4 ntasks_per_node: 8
5 job_name: single_node_finetune

Multi-Node Distributed Training (2 x 8 GPUs)

1nemo_run:
2 executor: my_slurm
3 nodes: 2
4 ntasks_per_node: 8
5 time: "08:00:00"
6 job_name: multinode_pretrain

AutoModel sets the executor’s launcher to NeMo-Run’s native torchrun launcher; AutoModel does not assemble torchrun rendezvous flags itself. For the generic multi-node path, NeMo-Run generates --nnodes, --nproc-per-node, --node-rank, --rdzv-backend, --rdzv-endpoint, and --rdzv-id. The endpoint uses NeMo-Run’s rank-zero host macro and rendezvous port.

Custom Container Image and Mounts

1nemo_run:
2 executor: my_slurm
3 container_image: /images/automodel_nightly.sqsh
4 container_mounts:
5 - /scratch/datasets:/datasets
6 - /scratch/checkpoints:/checkpoints
7 env_vars:
8 HF_HOME: /datasets/hf_cache
9 NCCL_DEBUG: INFO

Local Execution (No Cluster)

The launcher can construct LocalExecutor without an $NEMORUN_HOME/executors.py entry, and its process-count field is ntasks_per_node. However, the current AutoModel script always opens /nemo_run/code/automodel_config.yaml, while LocalExecutor does not stage or mount the PatternPackager output at /nemo_run/code. Therefore the following shape is not currently an end-to-end runnable AutoModel launch:

1nemo_run:
2 executor: local
3 ntasks_per_node: 2
4 job_name: local_test

For a working local two-process launch, omit nemo_run: and use AutoModel’s interactive launcher:

$uv run automodel your_config.yaml --nproc-per-node 2

Monitor and Manage Jobs

NeMo-Run stores experiment metadata under $NEMORUN_HOME/experiments/. Set tail_logs: true in the YAML to stream job output after submission.

For Slurm-based executors, standard Slurm commands also work:

$squeue -u $USER # List your queued and running jobs
$JOB_ID=12345 # Replace with the Slurm job ID
$sacct -j "$JOB_ID" # View job accounting information
$scancel "$JOB_ID" # Cancel a running or pending job

For Kubeflow jobs launched after independently resolving the packaging/path constraints above, use kubectl to monitor pods and jobs.

How It Works

  1. The automodel CLI detects the nemo_run: key and imports NemoRunLauncher.
  2. The nemo_run: section is popped from the config, and the remaining training config is passed to the launcher.
  3. For a named non-local executor, the launcher loads a pre-configured executor from $NEMORUN_HOME/executors.py; for executor: local, it directly creates LocalExecutor. Override fields are applied on top of executor defaults.
  4. The launcher writes the training config to <job_dir>/<timestamp>/automodel_config.yaml. It then assigns a NeMo-Run PatternPackager to the executor, with the timestamped directory as its relative root. Whether that package is extracted or mounted at the path used by the script is backend-specific; it is not done for Local, and Kubeflow uses its separate PVC-backed code_dir flow.
  5. AutoModel creates a script that launches the recipe module with -c /nemo_run/code/automodel_config.yaml and selects NeMo-Run’s native torchrun launcher. NeMo-Run, not AutoModel, derives the torchrun arguments from executor.nnodes() and executor.nproc_per_node(); the generic path uses --rdzv-endpoint.
  6. The script is submitted via nemo_run.Experiment with the requested detach value. NeMo-Run forces detach=False for executor types that do not support detachment, including Local and Kubeflow.

Customize Configuration

Override any training parameter from the command line, same as with local runs:

$automodel config_with_nemo_run.yaml \
> --model.pretrained_model_name_or_path meta-llama/Llama-3.2-3B

When to Use NeMo-Run vs. SkyPilot vs. Slurm

NeMo-RunSkyPilotSlurm (sbatch)
InfrastructureCurrent documented AutoModel path: Slurm; NeMo-Run has other backends with the limitations abovePublic cloud (AWS, GCP, Azure)On-prem HPC
Container supportBackend-specific; the documented Slurm path uses Pyxis/EnrootN/A (cloud VMs)Manual (in sbatch script)
Setup requiredLocked AutoModel environment + a named non-local executor definition; Kubeflow needs its extra and PVC path workCloud credentials + sky checkCluster access + sbatch script
Job submissionuv run automodel config.yaml for the compatible Slurm pathuv run automodel config.yamlsbatch slurm.sub
Good forReusable Slurm executor configuration through the current AutoModel integrationCloud burst, cost optimization, spot instancesDirect Slurm scripts, full control over sbatch