Build and Dependency
Quick Start
Clone and install:
git clone https://github.com/NVIDIA-NeMo/Automodel.git && cd Automodel
uv sync --locked --all-groups --extra all
Or use the NeMo-AutoModel container from NVIDIA NGC (pick a published tag from
the NGC catalog —
e.g. 26.04):
docker pull nvcr.io/nvidia/nemo-automodel:26.04
docker run --gpus all -it nvcr.io/nvidia/nemo-automodel:26.04
Installation Options
Option 1: NeMo-AutoModel Container (NGC)
The container ships with all dependencies pre-installed at /opt/Automodel
(WORKDIR) with the venv at /opt/venv. Run as-is:
docker run --gpus all --network=host -it --rm --shm-size=32g \
nvcr.io/nvidia/nemo-automodel:26.04 /bin/bash
Mounting your local checkout into the container
To develop against your host checkout, bind-mount it over /opt/Automodel to
override the installed source:
docker run --gpus all --network=host -it --rm --shm-size=32g \
-v <local-Automodel-path>:/opt/Automodel \
nvcr.io/nvidia/nemo-automodel:26.04 /bin/bash
Inside the container, patch pyproject.toml / uv.lock for the PyTorch base
image, then re-sync:
cd /opt/Automodel
bash docker/common/update_pyproject_pytorch.sh /opt/Automodel
uv sync --locked --all-groups --extra all
Warning: the
update_pyproject_pytorch.shstep is required. Without it,uv syncwill try to reinstalltorch, which leads to CUDA version mismatches and TE import failures — uv cannot recognize the torch baked into the PyTorch base container.
Option 2: uv (Recommended for Local Development)
--all-groups pulls the build, docs, and test dev groups (defined in
pyproject.toml); drop it for a runtime-only install.
uv sync --locked --all-groups # base + dev groups
uv sync --locked --all-groups --extra cuda # CUDA support
uv sync --locked --all-groups --extra fa # flash-attention
uv sync --locked --all-groups --extra moe # mixture-of-experts
uv sync --locked --all-groups --extra vlm # vision-language models (core)
uv sync --locked --all-groups --extra vlm-media # + video/Qwen/Mistral decode (opencv, decord, qwen-utils; FFmpeg-bearing)
uv sync --locked --all-groups --extra diffusion # diffusion models
uv sync --locked --all-groups --extra diffusion-media # + diffusion preprocessing/export (imageio-ffmpeg, opencv)
uv sync --locked --all-groups --extra media # vlm-media + diffusion-media (union)
uv sync --locked --all-groups --extra delta-databricks # Delta Lake / Databricks
uv sync --locked --all-groups --extra all # all standard extras (EXCLUDES media — FFmpeg kept opt-in)
The media extras (vlm-media, diffusion-media, media) bundle FFmpeg and are
deliberately excluded from all and from the container image — add them
explicitly for video/image decode.
Option 3: uv pip
Full install (matches uv sync --extra all):
uv venv
source .venv/bin/activate
uv pip install -e ".[all]"
To add NeMo Run submission support to the base package:
uv venv
source .venv/bin/activate
uv pip install "nemo-automodel[cli]"
The cli extra is additive: it adds nemo-run but does not remove the base
package's core training dependencies, including PyTorch.
Package Management
Always use uv. Do not introduce pip install commands in scripts or docs.
| Task | Command |
|---|---|
| Install from lockfile | uv sync --locked |
| Add a new dependency | uv add <package> |
| Add an optional dependency | uv add --optional --extra <group> <package> |
| Regenerate the lockfile | uv lock |
Environment Variables
export HF_TOKEN="hf_..." # Hugging Face token for gated models
export WANDB_API_KEY="..." # Weights & Biases logging
export HF_HOME="/path/to/hf_cache" # Hugging Face cache directory
CLI Usage
The entry point is automodel (defined at nemo_automodel.cli.app:main).
Pattern: uv run automodel <config.yaml> [--nproc-per-node N] [--key.subkey value ...]
# The YAML's recipe field selects LLM, VLM, diffusion, or retrieval behavior.
uv run automodel examples/llm_finetune/llama3_2/llama3_2_1b_squad.yaml --nproc-per-node 8
Override any config value from the CLI:
uv run automodel examples/llm_finetune/llama3_2/llama3_2_1b_squad.yaml \
--model.pretrained_model_name_or_path meta-llama/Llama-3.2-1B
Common Pitfalls
| Problem | Cause | Fix |
|---|---|---|
Stale .venv after switching branches |
Cached environment out of sync | Delete .venv and re-run uv sync --locked |
| Import errors for optional features (TE, flash-attn, MoE) | Missing extras | Install the matching uv extra (--extra fa, --extra moe, etc.) |
Import errors for media (cv2, decord, qwen_vl_utils, imageio_ffmpeg) |
Media extras are opt-in (not in all) |
Install --extra vlm-media (VLM/Qwen/Mistral) or --extra diffusion-media (diffusion) |
| TransformerEngine version mismatch | The TE installed by uv sync takes precedence over the version baked into the container |
Set the desired TE version in pyproject.toml / uv.lock and re-run uv sync — the venv's TE wins, not the container's |