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.0.0", "pybind11", "torch", "wheel"]
build-backend = "setuptools.build_meta"

[tool.setuptools]
package-dir = {"" = "src"}

[tool.setuptools.packages.find]
where = ["src"]
include = ["megatron.bridge", "megatron.bridge.*"]

[project]
name = "megatron-bridge"
description = "Megatron Bridge: Training Recipes for Megatron-based LLM and VLM models"
readme = "README.md"
dynamic = ["version"]
license = {file = "LICENSE"}
requires-python = ">=3.10"
authors = [{ name = "NVIDIA", email = "nemo-toolkit@nvidia.com" }]
maintainers = [{ name = "NVIDIA", email = "nemo-toolkit@nvidia.com" }]
keywords = [
    "NLP",
    "NeMo",
    "deep",
    "gpu",
    "language",
    "learning",
    "machine",
    "nvidia",
    "pytorch",
    "speech",
    "torch",
    "tts",
]
classifiers = [
    "Development Status :: 5 - Production/Stable",
    "Environment :: Console",
    "Intended Audience :: Developers",
    "Intended Audience :: Information Technology",
    "Intended Audience :: Science/Research",
    "License :: OSI Approved :: Apache Software License",
    "Natural Language :: English",
    "Operating System :: OS Independent",
    "Programming Language :: Python :: 3",
    "Programming Language :: Python :: 3.10",
    "Topic :: Scientific/Engineering :: Artificial Intelligence",
    "Topic :: Scientific/Engineering :: Image Recognition",
    "Topic :: Scientific/Engineering :: Mathematics",
    "Topic :: Scientific/Engineering",
    "Topic :: Software Development :: Libraries :: Python Modules",
    "Topic :: Software Development :: Libraries",
    "Topic :: Utilities",
]
dependencies = [
    "transformers<5.0.0",
    "datasets",
    "accelerate",
    "omegaconf>=2.3.0",
    "tensorboard>=2.19.0",
    "typing-extensions",
    "rich",
    "wandb>=0.19.10",
    "six>=1.17.0",
    "regex>=2024.11.6",
    "pyyaml>=6.0.2",
    "tqdm>=4.67.1",
    "hydra-core>1.3,<=1.3.2",
    "megatron-core[dev,mlm]>=0.15.0a0,<0.17.0",
    "qwen-vl-utils",
    "transformer-engine[pytorch]>=2.10.0a0,<2.12.0",
    "mamba-ssm",
    "nvidia-resiliency-ext",
    "causal-conv1d",
    "flash-linear-attention",
    "timm",
    "open-clip-torch>=3.2.0",
    "mlflow>=3.2.0",
]


[tool.setuptools.dynamic]
version = { attr = "megatron.bridge.package_info.__version__" }

[tool.uv]
# Currently, TE must be built with no build-isolation b/c it requires torch
no-build-isolation-package = [
    "transformer-engine",
    "transformer-engine-torch",
    "mamba-ssm",
    "causal-conv1d",
    "nv-grouped-gemm",
    "flash_mla",
    "flash-linear-attention",
]
prerelease = "allow"
override-dependencies = [
    "nvidia-modelopt[torch]>=0.37.0",
    "torch; sys_platform == 'never'",
    "torchvision; sys_platform == 'never'",
    "triton; sys_platform == 'never'",
    "transformer-engine[pytorch]>=2.9.0a0,<2.10.0",
]

# uv.sources allows us to override dependencies with VCS commits.
# Lets use this only for debugging purposes, but not for production (main).
[tool.uv.sources]
transformer-engine = { git = "https://github.com/NVIDIA/TransformerEngine.git", rev = "6a34b6574fa6c29d9d07fdcddf9812cbb1488878" }
megatron-core = { path = "3rdparty/Megatron-LM/" }
nvidia-resiliency-ext = { git = "https://github.com/NVIDIA/nvidia-resiliency-ext.git", rev = "54f85fe422d296cf04ea524130014bd3a2c3add1" }
nvidia-modelopt = { git = "https://github.com/NVIDIA/TensorRT-Model-Optimizer.git", rev = "0a4f0a8b933121f7af080261a0a5a7717f2c5d49" }
# mamba-ssm = { git = "https://github.com/yfw/mamba", branch = "general_stride_fix" }

[project.optional-dependencies]
recipes = [
    "nemo-run>=0.5.0a0,<0.6.0",
]
tensor-inspect = [
    "nvdlfw-inspect==0.2.1",
]

[dependency-groups]
docs = [
    "myst-parser>=4.0.1",
    "nvidia-sphinx-theme>=0.0.8",
    "sphinx>=8.1.3",
    "sphinx-autobuild>=2024.10.3",
    "sphinx-autodoc2>=0.5.0",
    "sphinx-copybutton>=0.5.2",
    "sphinxcontrib-mermaid",
]
test = [
    "coverage>=7.8.1",
    "flake8>=7.2.0",
    "pylint>=3.3.7",
    "pytest>=8.3.5",
    "pytest-timeout>=2.4.0",
    "pytest-mock>=3.14.0",
    "pytest-runner>=6.0.1",
    "pygithub",
    "click",
]
dev = [
    "pre-commit>=3.6.0",
    "ruff>=0.9.9",
    "mypy>=1.8.0",
]
build = ["setuptools", "torch", "pybind11", "Cython>=3.0.0", "numpy<2.0.0", "ninja", "nvidia-mathdx"]

[project.entry-points."nemo_run.cli"]
lm = "megatron.bridge"

[project.urls]
Download = "https://github.com/NVIDIA-NeMo/Megatron-Bridge/releases"
Homepage = "https://github.com/NVIDIA-NeMo/Megatron-Bridge"

[tool.pytest.ini_options]
# durations=0 will display all tests execution time, sorted in ascending order starting from from the slowest one.
# -vv will also display tests with durration = 0.00s
addopts = "--verbose --pyargs --durations=0 --strict-markers"  # always add these arguments to pytest
testpaths = ["tests"]
# directories to ignore when discovering tests
norecursedirs = [
    "src",
    "external",
    "examples",
    "docs",
    "scripts",
    "tools",
    "tutorials",
    "*.egg",
    ".*",
    "_darcs",
    "build",
    "CVS",
    "dist",
    "venv",
    "{arch}"
]
# markers to select tests, use `pytest --markers` to see all available markers, `pytest -m "<marker>"` to select tests
markers = [
    "unit: marks unit test, i.e. testing a single, well isolated functionality (deselect with '-m \"not unit\"')",
    "integration: marks test checking the elements when integrated into subsystems (deselect with '-m \"not integration\"')",
    "system: marks test working at the highest integration level (deselect with '-m \"not system\"')",
    "acceptance: marks test checking whether the developed product/model passes the user defined acceptance criteria (deselect with '-m \"not acceptance\"')",
    "docs: mark tests related to documentation (deselect with '-m \"not docs\"')",
    "skipduringci: marks tests that are skipped ci as they are addressed by Jenkins jobs but should be run to test user setups",
    "pleasefixme: marks tests that are broken and need fixing",
]

[tool.coverage.run]
concurrency = ["thread", "multiprocessing"]
omit = [
    "/tmp/*",
    "/opt/Megatron-Bridge/tests/*",
    "/opt/Megatron-Bridge/*.py",
    "*_ray.py",
]

[tool.coverage.paths]
source = [
    ".",
    "/opt/Megatron-Bridge/",
    "/home/runner/work/Megatron-Bridge/Megatron-Bridge",
]