Skip to content

CONTRIBUTING.md

Latest commit

 

History

History
306 lines (223 loc) · 11.8 KB

CONTRIBUTING.md

File metadata and controls

306 lines (223 loc) · 11.8 KB
Raw

Contributing To NeMo-Automodel

Environment setup for Automodel

Common workflows used for setting up Automodel environment:

  1. Developing with Automodel container
  2. Developing with UV sync/pip install
  3. Developing with custom docker build

1. Developing with Automodel container

The latest Automodel container can be found: here

The container can be run with the following docker command:

docker run --gpus all --network=host -it --rm --shm-size=32g nvcr.io/nvidia/nemo-automodel:26.10.00 /bin/bash

Mounting local Automodel directory into the container

To sync local Automodel directory into the container, mount the local directory into /opt/Automodel and override the installed Automodel repository. Example docker command:

docker run --gpus all --network=host -it --rm -v <local-Automodel-path>:/opt/Automodel --shm-size=32g nvcr.io/nvidia/nemo-automodel:26.10.00 /bin/bash

Within the container, cd into /opt/Automodel/ and update the pyproject.toml and uv.lock file by running the following command:

bash docker/common/update_pyproject_pytorch.sh /opt/Automodel

Finally, run uv sync to sync the container with the updated repository:

uv sync --locked --extra all --all-groups

Warning

Ensure bash docker/common/update_pyproject_pytorch.sh /opt/Automodel is executed. Without this command, uv sync will attempt to reinstall torch. This leads to errors relating to CUDA version mismatch, TE import failures, etc. This work around is required as uv cannot recognize the torch installed in the PyTorch base container.

2. Developing with uv sync/pip install

Uv sync and pip install are both supported in Automodel. Uv sync is the recommened path.

From the local Automodel directory run the following command:

uv sync --locked --extra all

The following optional dependencies are available, please see [project.optional-dependencies] section in pyproject.toml:

  • cuda (all dependencies that require cuda)
  • extra (additional dependencies for model coverage)
  • fa (flash attention)
  • delta-databricks
  • moe
  • vlm
  • all (installs cuda, delta-databricks, extra and vlm)

Example, installing vlm dependencies:

uv sync --locked --extra vlm

3. Developing with custom docker build

For developers building a custom docker container with Automodel, please refer to Automodel's Dockerfile.

If Nvidia PyTorch container is used as a base image, please review Developing with Automodel container.

The image includes a pinned Mixture-of-Kittens build for GB200 (SM100) and GB300 (SM103). Pass --build-arg MOK_ARCH=SM100 or --build-arg MOK_ARCH=SM103 to build for a single target. MoK recipes require Blackwell GPUs; H100 GPUs are unsupported. Outside the image, install it with MOK_ARCH=ALL uv sync --extra mok on Linux with Python 3.12+, CUDA-enabled PyTorch 2.10+ and the matching CUDA 13 toolkit. Use MOK_ARCH=SM100 or MOK_ARCH=SM103 for a single target. Run uv cache clean mixture-of-kittens before switching architectures in an existing installation, then add --reinstall-package mixture-of-kittens to the sync command.

Command to build Automodel's Dockerfile:

# Set build arguments
export AUTOMODEL_INSTALL=all #(cuda, moe, vlm, ...)
export BASE_IMAGE=pytorch #(cuda, pytorch)

# Dependency install options [True, False]
export INSTALL_TE=True

docker build -f docker/Dockerfile \
    --build-arg AUTOMODEL_INSTALL=$AUTOMODEL_INSTALL \
    --build-arg BASE_IMAGE=$BASE_IMAGE \
    --build-arg INSTALL_MAMBA=$INSTALL_MAMBA \
    -t automodel --target=automodel_final .

MoE Dependency

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb
sudo apt-get update
sudo apt-get install cudnn-cuda-12
pip install transformer-engine[pytorch]==2.8.0
# Flash-attn version should be selected to satisfy TE requirements
# https://github.com/NVIDIA/TransformerEngine/blob/v2.4/transformer_engine/pytorch/attention/dot_product_attention/utils.py#L108
pip install flash-attn==2.7.4.post1

Development Dependencies

We use uv for managing dependencies.

New required dependencies can be added by uv add $DEPENDENCY. If pyproject.toml changes, update both lock files using the same flow as Generate Uv lock. This updates the default uv.lock, then updates the PyTorch-container lock file with the override dependencies from docker/common/uv-pytorch.toml. The PyTorch helper rewrites pyproject.toml for locking, so the commands below save and restore your intended pyproject.toml contents:

The workflow can update lock files automatically for same-repository PRs. For PRs opened from forks, the workflow runs in check-only mode because it cannot push fixes back to the fork. If the fork check fails, run this lock update flow locally and include the generated lock files in your branch.

set -e

uv lock --check || uv lock

tmp_pyproject="$(mktemp)"
cp pyproject.toml "$tmp_pyproject"
mv uv.lock uv_main.lock
bash docker/common/update_pyproject_pytorch.sh "$PWD"
uv lock --check || uv lock
mv uv.lock docker/common/uv-pytorch.lock
mv uv_main.lock uv.lock
cp "$tmp_pyproject" pyproject.toml
rm "$tmp_pyproject"

Only commit your intended pyproject.toml changes plus the generated lock files:

git add pyproject.toml uv.lock docker/common/uv-pytorch.lock
git commit -s -m "build: Adding dependencies"
git push

Linting and Formatting

We use ruff for linting and formatting.

Installation:

pip install ruff

Format:

ruff check --fix .
ruff format .

Ruff does not cover every rule we enforce. tools/lint_no_globals.py rejects globals() and its equivalents (module-scope vars(), sys.modules[__name__].__dict__, setattr(sys.modules[__name__], ...)). These are usually used to swap a module-level function so that another function picks up the replacement — a process-wide mutation that changes behavior for every concurrent or nested caller. Pass the collaborator in, or expose a method subclasses can override, instead. The one exception is the PEP 562 lazy-import cache: globals()[name] = attr written directly inside a module-level def __getattr__(name). The check covers the same paths as Bandit (app.py, nemo_automodel/, examples/, scripts/, tools/, tutorials/); tests/ is excluded, as it is for Ruff. Run it with:

python tools/lint_no_globals.py

Unit-test runtime

Unit-test time limits are enforced at three times the configured value to allow for CI load. The configured default runtime budget is 30 seconds, so new and modified tests have an effective 90-second budget, including setup, the test body, and teardown. A configured 15-second marker allows 45 seconds. If every phase passes but the test exceeds its effective budget, pytest runs it once more in the same worker with fresh function-scoped fixtures. Shared fixtures and compiler caches can remain warm. Both attempts must exceed the budget to fail the runtime gate; assertion failures, fixture errors, and hard timeouts are never retried. A warning records every confirmation.

Keep model shapes small and avoid downloads or unnecessary compilation. Tests that consistently need more than the effective budget should be optimized or moved to functional tests. An exact @pytest.mark.runtime_budget(...) may tighten the limit, but cannot raise it for a new or modified test. A changed test with a marker above 30 seconds stops collection with a usage error; update its marker and optimize the test, or move it to functional coverage. Existing larger allowances retain their hard watchdog on unrelated changes; they do not impose a runtime gate on those changes. The configured default hard watchdog is 70 seconds, enforced at 210 seconds. The same multiplier applies once to resolved pytest timeout markers and CLI, environment, or ini timeout settings, preserving their usual precedence, timeout method, and function-only scope. A zero timeout stays disabled. Global pytest timeout options still bypass the soft runtime policy. Diagnostics show both configured and effective soft budgets. These rules apply only under tests/unit_tests; workflow and subprocess timeouts are unchanged.

Changed tests are identified relative to the PR merge base (or origin/main locally). Set AUTOMODEL_RUNTIME_BUDGET_BASE=<commit> to explicitly select a base when reproducing a CI result. CPU and GPU unit-test jobs both limit host math libraries to one thread per process to avoid oversubscription by spawned ranks.

Pre-commit

We recommand to use prek to ensure code quality. It is a faster and more modern alternative to pre-commit.

Installation:

uv tool install prek

Usage:

# Install git hooks
prek install

# Run manually on all files
prek run --all-files

After installing the git hooks, git commit will automatically run incremental checks.

Adding Documentation

If your contribution involves documentation changes, please refer to the Fern docs maintenance guide guide for detailed instructions on building and serving the documentation.

Signing Your Work

  • We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license.

    • Any contribution which contains commits that are not Signed-Off will not be accepted.

  • To sign off on a commit you simply use the --signoff (or -s) option when committing your changes:

    git commit -s -m "Add cool feature."

    This will append the following to your commit message:

    Signed-off-by: Your Name <your@email.com>

  • Full text of the DCO:

    Developer Certificate of Origin
Version 1.1

Copyright (C) 2004, 2006 The Linux Foundation and its contributors.

Everyone is permitted to copy and distribute verbatim copies of this
license document, but changing it is not allowed.


Developer's Certificate of Origin 1.1

By making a contribution to this project, I certify that:

(a) The contribution was created in whole or in part by me and I
    have the right to submit it under the open source license
    indicated in the file; or

(b) The contribution is based upon previous work that, to the best
    of my knowledge, is covered under an appropriate open source
    license and I have the right under that license to submit that
    work with modifications, whether created in whole or in part
    by me, under the same open source license (unless I am
    permitted to submit under a different license), as indicated
    in the file; or

(c) The contribution was provided directly to me by some other
    person who certified (a), (b) or (c) and I have not modified
    it.

(d) I understand and agree that this project and the contribution
    are public and that a record of the contribution (including all
    personal information I submit with it, including my sign-off) is
    maintained indefinitely and may be redistributed consistent with
    this project or the open source license(s) involved.