pyproject.toml

# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

[build-system]
requires = ["setuptools >= 80.10.2"]
build-backend = "setuptools.build_meta"

[tool.setuptools]
include-package-data = true

[tool.setuptools.packages.find]
include = ["nemo_automodel", "nemo_automodel.*"]

[tool.setuptools.package-data]
nemo_automodel = [
    "components/datasets/llm/megatron/Makefile",
]

[tool.setuptools.dynamic]
version = { attr = "nemo_automodel.package_info.__version__" } # any module attribute compatible with ast.literal_eval
readme = { file = "README.md", content-type = "text/markdown" }

[project]
description = "DTensor-native pretraining and fine-tuning for LLMs/VLMs with day-0 Hugging Face support, GPU-acceleration, and memory efficiency."
name = "nemo-automodel"
dynamic = ["version"]
readme = "README.md"
authors = [{ name = "NVIDIA", email = "nemo-toolkit@nvidia.com" }]
maintainers = [{ name = "NVIDIA", email = "nemo-toolkit@nvidia.com" }]
requires-python = ">=3.10"
classifiers = [
    # How mature is this project? Common values are
    #  1 - Planning
    #  2 - Pre-Alpha
    #  3 - Alpha
    #  4 - Beta
    #  5 - Production/Stable
    #  6 - Mature
    #  7 - Inactive
    'Development Status :: 4 - Beta',
    # Indicate who your project is intended for
    'Intended Audience :: Developers',
    'Intended Audience :: Science/Research',
    'Intended Audience :: Information Technology',
    # Indicate what your project relates to
    'Topic :: Scientific/Engineering',
    'Topic :: Scientific/Engineering :: Mathematics',
    'Topic :: Scientific/Engineering :: Image Recognition',
    'Topic :: Scientific/Engineering :: Artificial Intelligence',
    'Topic :: Software Development :: Libraries',
    'Topic :: Software Development :: Libraries :: Python Modules',
    'Topic :: Utilities',
    # Pick your license as you wish (should match "license" above)
    'License :: OSI Approved :: Apache Software License',
    # Supported python versions
    'Programming Language :: Python :: 3',
    'Programming Language :: Python :: 3.10',
    'Programming Language :: Python :: 3.11',
    'Programming Language :: Python :: 3.12',
    # Additional Setting
    'Environment :: Console',
    'Natural Language :: English',
    'Operating System :: OS Independent',
]
dependencies = [
    "datasets>=4.0.0",
    "megatron-fsdp>=0.2.3",
    "mistral-common[image,audio,hf-hub,sentencepiece]",
    "opencv-python-headless==4.10.0.84",
    "pybind11",
    "pyyaml",
    "torch>=2.6.0,<=2.10.0",
    "torchdata",
    "transformers>=5.3.0,<5.4.0",
    "wandb",
    "torchao",
    "mlflow",
    "flashoptim>=0.1.3",
]

[project.optional-dependencies]
diffusion = [
    "diffusers>=0.36.0",
    "ftfy",
    "imageio",
    "imageio-ffmpeg",
    "kernels",
    "opencv-python-headless",
    "torchvision",
]
# nvidia-cudnn-cu12 pin: This is required for GPTOSS TE support + faster cudnn attention (only on Linux where CUDA is available)
# "nvidia-cudnn-cu12>=9.18.0.0; sys_platform == 'linux'",
# 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
# "flash-attn==2.8.1",
cuda = [
    "causal-conv1d",
    "mamba-ssm",
    "nv-grouped-gemm",
    # This is required to prevent ONNX errors arising from TE
    "onnxscript>=0.5.6",
    "transformer-engine[pytorch]<=2.11.0",
]
cuda_source = [
    "bitsandbytes",
]
extra = ["perceptron", "sentencepiece"]
fa = ["flash-attn<=2.8.3"]
delta-databricks = [
    "deltalake>=1.0.0",  # For Delta Lake metadata
    "databricks-sql-connector>=3.0.0",  # For Unity Catalog streaming access
]
moe = [
    "nemo_automodel[cuda]",
    "deep_ep",
]
vlm = [
    "albumentations",
    "backoff",
    "mistral_common[opencv]>=1.9.0",
    "numpy",
    "numba",
    "open-clip-torch",
    "pillow",
    "qwen-omni-utils",
    "qwen-vl-utils[decord]; (platform_machine == 'x86_64' and platform_system != 'Darwin')",
    "timm<=1.0.22",
    "torchcodec; (platform_machine == 'x86_64' and platform_system != 'Darwin')",
]
cli = [
    "pyyaml",
]
all = [
    "nemo_automodel[cuda]",
    "nemo_automodel[delta-databricks]",
    "nemo_automodel[diffusion]",
    "nemo_automodel[extra]",
    "nemo_automodel[vlm]",
]

[project.urls]
Homepage = "https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/"
Repository = "https://github.com/NVIDIA/NeMo-Automodel"
Download = "https://github.com/NVIDIA/NeMo-Automodel/releases"

[project.scripts]
automodel = "nemo_automodel.cli.app:main"
am = "nemo_automodel.cli.app:main"

[dependency-groups]
build = ["setuptools", "torch<=2.10.0", "packaging", "psutil"]
docs = [
    "sphinx",
    "sphinx-autobuild",    # For live doc serving while editing docs
    "sphinx-autodoc2",     # For documenting Python API
    "sphinx-copybutton",   # Adds a copy button for code blocks
    "sphinx-design",       # For grid layout and card components
    "myst_parser",         # For our markdown docs
    "nvidia-sphinx-theme", # Our NVIDIA theme
]
linting = [
    "pre-commit>=4.2.0",
    "ruff~=0.9.0",
    "import-linter~=2.4",
    "ty>=0.0.20",
]
test = ["coverage", "pytest", "peft>=0.18.1", "onnxruntime-gpu; (platform_machine == 'x86_64' and platform_system != 'Darwin')"]
# CP for linear attention (fla.ops.cp) requires fla > 0.4.1; pin to main until a release includes it.
dev = [
    "cut-cross-entropy",
    "flash-linear-attention @ git+https://github.com/fla-org/flash-linear-attention.git@7cbe461b10bde60180c32c5f9381a04d41887514",
    "liger-kernel; (platform_machine == 'x86_64' and platform_system != 'Darwin')",
    "dion @ git+https://github.com/microsoft/dion.git; platform_machine == 'x86_64' and platform_system != 'Darwin'",
]

[tool.uv]
default-groups = ["build", "docs", "test"]
no-build-isolation-package = [
    "bitsandbytes",
    "causal-conv1d",
    "deep_ep",
    "flash-attn",
    "mamba-ssm",
    "nv-grouped-gemm",
    "transformer-engine",
    "transformer-engine-torch",
]
constraint-dependencies = [
    "aiohttp>=3.13.3", # Address CVE GHSA-6mq8-rvhq-8wgg
    "cryptography>=46.0.5", # Address CVE GHSA-r6ph-v2qm-q3c2
    "jaraco-context>=6.1.0", # Address CVE GHSA-58pv-8j8x-9vj2
    "pillow>=12.1.1", # Address CVE GHSA-cfh3-3jmp-rvhc
    "protobuf>=6.33.5", # Address CVE GHSA-7gcm-g887-7qv7
    "pyasn1>=0.6.2", # Address CVE GHSA-63vm-454h-vhhq
    "starlette>=0.49.1", # Address CVE GHSA-7f5h-v6xp-fcq8
    "setuptools>=80.10.2",
    "urllib3>=2.6.3", # Address CVE GHSA-gm62-xv2j-4w53
    "wheel>=0.46.2", # Address CVE GHSA-8rrh-rw8j-w5fx
    "fonttools>=4.60.2", # Address CVE-2025-66034
    "mlflow>=3.10.0rc0", # Address CVE
]
override-dependencies = [
    "nvidia-cudnn-cu12==9.19.0.56; sys_platform == 'linux'"
]
prerelease = "allow"

[tool.uv.sources]
deep_ep = { git = "https://github.com/deepseek-ai/DeepEP.git", rev = "7febc6e25660af0f54d95dd781ecdcd62265ecca" }
cut-cross-entropy = { git = "https://github.com/apple/ml-cross-entropy.git", rev = "87a86aba72cfd2f0d8abecaf81c13c4528ea07d8" }
liger-kernel = { git = "https://github.com/linkedin/Liger-Kernel.git", rev = "1f1a4b8d6a6c3c1ddb3573a46480c41cef25bde6" }
onnxruntime-gpu = { index = "onnxruntime-cuda-13" }
torch = [
  { index = "pytorch-cpu", marker = "sys_platform != 'linux' and sys_platform != 'darwin'" },
  { index = "pytorch-cu129", marker = "sys_platform == 'linux'" },
  { index = "pypi", marker = "sys_platform == 'darwin'" },
]
torchvision = [
  { index = "pytorch-cpu", marker = "sys_platform != 'linux' and sys_platform != 'darwin'" },
  { index = "pytorch-cu129", marker = "sys_platform == 'linux'" },
  { index = "pypi", marker = "sys_platform == 'darwin'" },
]

[[tool.uv.index]]
name = "pypi"
url = "https://pypi.org/simple"
explicit = true

[[tool.uv.index]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
explicit = true

[[tool.uv.index]]
name = "pytorch-cu129"
url = "https://download.pytorch.org/whl/cu129"
explicit = true

[[tool.uv.index]]
name = "pytorch-cu130"
url = "https://download.pytorch.org/whl/cu130"
explicit = true

[[tool.uv.index]]
name = "onnxruntime-cuda-13"
url = "https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-13/pypi/simple/"
explicit = true

# Augment build dependencies for packages that need torch at build time
[tool.uv.extra-build-dependencies]
deep_ep = [{ requirement = "torch", match-runtime = true }]
transformer-engine = [{ requirement = "torch", match-runtime = true }]
transformer-engine-torch = [{ requirement = "torch", match-runtime = true }]
nv_grouped_gemm = ["setuptools"]

# Needed when building from source
[[tool.uv.dependency-metadata]]
name = "deep_ep"
# This version has to match the version in the commit/rev/tag used
version = "v1.2.1+7febc6e"
requires-dist = ["torch", "packaging", "ninja"]

[[tool.uv.dependency-metadata]]
name = "nv_grouped_gemm"
# This version has to match the version in the commit/rev/tag used
version = "1.1.4.post8"
requires-dist = ["torch", "packaging", "ninja", "setuptools"]

[tool.pytest.ini_options]
pythonpath = ["."]

[tool.coverage.paths]
source = [".", "/opt/Automodel", "/home/runner/work/Automodel/Automodel"]

[tool.coverage.run]
source = ["nemo_automodel"]
parallel = true
concurrency = ["thread", "multiprocessing"]
omit = [
    "nemo_automodel/components/checkpoint/_backports/*.py",
    "tests/*",
    "nemo_automodel/components/_peft/lora_kernel.py",
    "nemo_automodel/components/loss/triton/te_cross_entropy.py",
    "nemo_automodel/components/moe/megatron/*.py",
    # Vendored from https://github.com/uccl-project/uccl/tree/main/ep/deep_ep_wrapper
    "nemo_automodel/components/moe/uccl_ep/_buffer.py",
    "nemo_automodel/components/moe/uccl_ep/_utils.py",
    "config.py",
    "config-3.py",
    # transformer v5 backports
    "nemo_automodel/_transformers/tokenization/tokenization_mistral_common.py",
    # requires databricks
    "nemo_automodel/components/datasets/llm/delta_lake_dataset.py",
    # Weird file suddenly appearing in coverage
    "_remote_module_non_scriptable",
    # deleted file; may still appear in coverage data from cache
    "nemo_automodel/_transformers/retrieval.py",
]

[tool.ruff]
line-length = 120
target-version = "py310"
exclude = [
   "tests/",
]

[tool.ruff.format]
quote-style = "double"
skip-magic-trailing-comma = false
line-ending = "auto"

[tool.ruff.lint]
# Enable all `pydocstyle` rules, limiting to those that adhere to the
# Google convention via `convention = "google"`, below.
select = [
    "F541", # f-string without any placeholders
    "F841", # local variable assigned but never used
    "F401", # imported but unused
    "E741", # ambiguous variable name
    "F821", # undefined name
    "E266", # too many leading '#' for block comment
    "I",    # isort
    "D101", # docstring
    "D103",
    "S301", # no pickle
    "S403", # no pickle
]

ignore = [
    "E501", # Line too long - handled by formatter
    "D101",
    "D103",
]

[tool.ruff.lint.pydocstyle]
convention = "google"

# Section to exclude errors for different file types
[tool.ruff.lint.per-file-ignores]
# Ignore all directories named `tests`.
"tests/**" = ["D"]
# Ignore all files that end in `_test.py`.
"*_test.py" = ["D"]
# Ignore F401 (import but unused) in __init__.py
"__init__.py" = ["F401"]


[tool.ruff.lint.isort]
known-third-party = ["datasets", "transformers", "wandb"]

[tool.importlinter]
root_package = "nemo_automodel"
exclude_type_checking_imports = true

[[tool.importlinter.contracts]]
name = "Components must not import each other"
type = "independence"
modules = [
    "nemo_automodel.components.checkpoint",
    "nemo_automodel.components.config",
    "nemo_automodel.components.datasets",
    "nemo_automodel.components.distributed",
    "nemo_automodel.components.flow_matching",
    "nemo_automodel.components.launcher",
    "nemo_automodel.components.loggers",
    "nemo_automodel.components.loss",
    "nemo_automodel.components.optim",
    "nemo_automodel.components.training",
    "nemo_automodel.components.utils",
]
ignore_imports = [
    "nemo_automodel.components.distributed.optimized_tp_plans -> nemo_automodel.components.models.llama.model",
]

# ty: type checking (https://github.com/astral-sh/ty)
# Default: permissive. Strict checks are opt-in per-file via overrides below.
# Phase 1 modules: launcher, attention, config, optim, loggers
# To advance to the next phase, move module paths from the "Phase 2" comment
# into the active overrides list and fix any new errors.
[tool.ty.environment]
python-version = "3.10"

[tool.ty.src]
include = ["nemo_automodel"]
exclude = [
    "tests/**",
    "nemo_automodel/components/checkpoint/_backports/**",
]

[tool.ty.analysis]
allowed-unresolved-imports = ["**"]

[tool.ty.rules]
unused-ignore-comment = "warn"
possibly-missing-attribute = "ignore"

# Phase 1: small, well-typed, leaf-dependency modules
[[tool.ty.overrides]]
include = [
    "nemo_automodel/components/attention/**",
    "nemo_automodel/components/config/**",
    "nemo_automodel/components/launcher/**",
    "nemo_automodel/components/loggers/**",
    "nemo_automodel/components/optim/**",
]

[tool.ty.overrides.rules]
# Stricter checking for Phase 1 modules
possibly-unresolved-reference = "error"
unresolved-import = "error"

# Phase 2 (uncomment when ready):
# [[tool.ty.overrides]]
# include = [
#     "nemo_automodel/components/loss/**",
#     "nemo_automodel/components/quantization/**",
# ]
#
# [tool.ty.overrides.rules]
# possibly-unresolved-reference = "error"
# unresolved-import = "error"

# Phase 3 (uncomment when ready):
# [[tool.ty.overrides]]
# include = [
#     "nemo_automodel/components/training/**",
#     "nemo_automodel/components/distributed/**",
#     "nemo_automodel/components/utils/**",
# ]
#
# [tool.ty.overrides.rules]
# possibly-unresolved-reference = "error"
# unresolved-import = "error"