pyproject.toml

# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.

[build-system]
requires = ["setuptools>=80", "pybind11", "packaging>=24.2"]
build-backend = "setuptools.build_meta"

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

[tool.setuptools.packages.find]
include = [
    "megatron.core",
    "megatron.core.*",
    "megatron.training",
    "megatron.training.*",
]

[tool.setuptools.dynamic]
version = { attr = "megatron.core.package_info.__version__" }
readme = { file = "README.md", content-type = "text/markdown" }

[project]
name = "megatron-core"
dynamic = ["version", "readme"]
description = "Megatron Core - a library for efficient and scalable training of transformer based models"
requires-python = ">=3.12"
license = { text = "Apache 2.0" }
dependencies = ["torch>=2.6.0", "numpy", "packaging>=24.2"]
authors = [{ name = "NVIDIA", email = "nemo-toolkit@nvidia.com" }]
maintainers = [{ name = "NVIDIA", email = "nemo-toolkit@nvidia.com" }]
keywords = [
    "NLP",
    "NLU",
    "deep",
    "gpu",
    "language",
    "learning",
    "machine",
    "nvidia",
    "pytorch",
    "torch",
    "transformer",
]
classifiers = [
    "Development Status :: 5 - Production/Stable",
    "Environment :: Console",
    "Intended Audience :: Developers",
    "Intended Audience :: Information Technology",
    "Intended Audience :: Science/Research",
    "License :: OSI Approved :: BSD License",
    "Natural Language :: English",
    "Operating System :: OS Independent",
    "Programming Language :: Python :: 3",
    "Programming Language :: Python :: 3.10",
    "Programming Language :: Python :: 3.11",
    "Programming Language :: Python :: 3.12",
    "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",
]

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

[project.optional-dependencies]
training = [
    "flask-restful",
    "sentencepiece",
    "tiktoken",
    "wandb",
    "transformers",
    "accelerate",
    "omegaconf",
    "gigatoken",
]

### 'mlm' group is deprecated. please use 'training' instead ###
mlm = [
    "flask-restful",
    "sentencepiece",
    "tiktoken",
    "wandb",
    "transformers",
    "accelerate",
    "omegaconf",
    "gigatoken",
]

dev = [
    "nvidia-modelopt[torch]>=0.44; sys_platform != 'darwin'",
    "nvidia-resiliency-ext==0.7.0",
    "tqdm",
    "einops~=0.8",
    "tensorstore~=0.1,!=0.1.46,!=0.1.72",
    "multi-storage-client~=0.50",
    "opentelemetry-api>=1.33.1,<2",
    "nemo-lens",
    "megatron-energon[av_decode]~=7.0",
    "av",
    "Pillow",
    "flashinfer-python>=0.5.0,<0.7.0",
    "nvidia-cudnn-frontend[cutedsl]==1.29.0",
    "wget",
    "onnxscript",
    "fastapi~=0.50",                                    # Forcing a little bit more recent version of fastapi to be compatible with pydantic 2.0
    "datasets",
    "emerging_optimizers",
    "hypercorn",
    "quart",
    "openai[aiohttp]",
    "orjson",
    "fast-hadamard-transform",
    "zstandard",
]

inference = [
    "torch-memory-saver",
]

otel = [
    "nemo-lens[sdk]",
]

### 'lts' extra is deprecated                                                   ###
### LTS Python dependencies are now pinned directly in docker/Dockerfile.ci.lts ###
### so this pyproject.toml can host meaningful module-level extras without      ###
### colliding with the LTS pin set. Kept as an empty alias for one release      ###
### cycle so `pip install megatron-core[lts]` does not error out.               ###
lts = []

te = [
    "transformer-engine[pytorch,core_cu13]",
]

ssm = [
    "mamba-ssm",
    "causal-conv1d~=1.5",
    "flash-linear-attention==0.5.1",
]

[dependency-groups]
test = [
    "coverage",
    "nltk",
    "wrapt",
    "pytest==9.1.1",
    "pytest-mock",
    "mock",
    "pytest-cov",
    "pytest-random-order",
    "pytest-asyncio",
    "pygithub",
    "pydantic",
    "tensorboard",
    "pyyaml",
    "nemo-run",
]
# GitHub CI installs Testmon at runtime for baseline and selective test modes.
# Testmon remains inactive unless pytest receives a --testmon option.
testmon = ["pytest-testmon==2.2.0"]
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
    "myst_parser",         # For our markdown docs
    "nvidia-sphinx-theme", # Our NVIDIA theme
]
build = [
    "setuptools>=80",
    "packaging>=24.2",
    "hatchling",
    "pybind11",
    "Cython>=3.0.0",
    "torch",
    "nvidia-cudnn-frontend>=1.29.0",  # for TE
    "nvidia-mathdx",                  # for TE
]
linting = [
    "ruff~=0.9.0",
    "black==26.3.0",
    "isort==5.13.2",
    "flake8==7.1.0",
    "pylint==3.2.6",
]
ci = ["python-gitlab", "slack-sdk", "pandas"]
batch_invariant = ["deep_gemm"]
no_pypi_wheels = ["flash_mla", "emerging_optimizers", "deep_gemm"]

[tool.uv]
managed = true
default-groups = ["linting", "build", "test"]
# CUTLASS 4.8 moved CuTe APIs used by older Quack releases.
constraint-dependencies = ["quack-kernels>=0.6.5"]
no-build-isolation-package = [
    "causal-conv1d",
    "flash_mla",
    "mamba-ssm",
    "transformer-engine",
    "transformer-engine-torch",
    "deep_gemm",
    "fast-hadamard-transform",
]
link-mode = "copy"
# The Mamba revision retained from main pins an older TileLang/TVM FFI family.
# Override those transitive constraints for cuDNN Frontend 1.28 compatibility.
# We don't define override-dependencies globally but rather locally where we need it.
# For instance, when installing into a PyTorch base image, we don't want to install torch,
# torchvision, and triton.
override-dependencies = [
    "torch; sys_platform == 'never'",
    "torchvision; sys_platform == 'never'",
    "triton; sys_platform == 'never'",
    "tilelang==0.1.12",
    "apache-tvm-ffi==0.1.11",
    # TorchX 0.7.0 caps urllib3 below 1.27; require the fix for GHSA-38jv-5279-wg99.
    "urllib3>=2.6.3",
]

[[tool.uv.dependency-metadata]]
name = "flash-mla"
version = "1.0.0+9edee0c"
requires-dist = []

[[tool.uv.dependency-metadata]]
name = "transformer-engine"
version = "2.20.2+6ea2a74a"
requires-dist = [
    # Cap below 2.14: pydantic 2.14 breaks langchain_core's module-level
    # RunnablePassthrough() instantiation, which is imported transitively in
    # downstream consumers (e.g. Megatron-Bridge via nvidia-resiliency-ext).
    "pydantic<2.14",
    "importlib-metadata>=1.0",
    "packaging",
    "torch>=2.1",
    "einops",
    "onnxscript",
    "onnx",
    "nvdlfw-inspect",
    "nvidia-cudnn-frontend>=1.29.0",
]

[[tool.uv.dependency-metadata]]
name = "fast-hadamard-transform"
version = "1.0.4.post1"
requires-dist = ["torch", "packaging", "ninja"]

[tool.uv.sources]

flash_mla = [
    { git = "https://github.com/deepseek-ai/FlashMLA", rev = "nv_dev" },
]
deep_gemm = { git = "https://github.com/deepseek-ai/DeepGEMM.git", rev = "714dd1a4a980f7937a74343d19a8eba4fe321480" }
transformer-engine = { git = "https://github.com/NVIDIA/TransformerEngine.git", rev = "6ea2a74a9e98c99e6d7b164a33775cc457520027" }
nemo-run = { git = "https://github.com/NVIDIA-NeMo/Run.git", rev = "e3935393a290aed1822af52139b4b8ee270fed1f" }
emerging_optimizers = { git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git", rev = "v0.3.0" }
fast-hadamard-transform = { git = "https://github.com/Dao-AILab/fast-hadamard-transform.git", rev = "f134af63deb2df17e1171a9ec1ea4a7d8604d5ca" }
mamba-ssm = { git = "https://github.com/state-spaces/mamba.git", rev = "0048fbf2e7b2f214dcbe703ea3dec2b9647595e1" }
nemo-lens = { git = "https://github.com/NVIDIA-NeMo/Lens.git", rev = "b85578fc2b736a1804705e537001b5f45e9c715d" }
torch-memory-saver = { git = "https://github.com/fzyzcjy/torch_memory_saver.git", rev = "9bc9a442e6d108c7b7903def199896a005143aaf" }

[tool.isort]
profile = "black"                                                          # black-compatible
line_length = 100                                                          # should match black parameters
py_version = 312                                                           # python 3.12 as a target version
known_first_party = ["megatron"]                                           # FIRSTPARTY section
known_third_party = ["transformer_engine"]                                 # THIRDPARTY section
sections = ["FUTURE", "STDLIB", "THIRDPARTY", "FIRSTPARTY", "LOCALFOLDER"]
default_section = "THIRDPARTY"
extend_skip = ["setup.py"]

[tool.black]
line_length = 100
skip_string_normalization = true
# recognized by future versions, disallows to reformat code with incompatible versions
# Matches NeMO version so people working on both codebases don't need two different version of black installed
required_version = "26"
skip_magic_trailing_comma = true
include = '\.pyi?$'
exclude = '''
/(
    \.git
  | \.venv
  | build
)/
'''

[tool.pytest.ini_options]
addopts = "--durations=15 -s -rA"
testpaths = ["tests"]
python_files = "test_*.py"
markers = [
    "internal: mark a test as a test to private/internal functions.",
    "flaky: mark flaky tests for LTS environment",
    "flaky_in_dev: mark flaky tests for DEV environment",
    "launch_on_gb200: mark a unit test to be launched on GB200 hardware (4 GPUs/node)",
]

[tool.coverage.run]
concurrency = ["thread", "multiprocessing"]
omit = [
    "/tmp/*",
    "/opt/megatron-lm/tests/*",
    "/opt/megatron-lm/tools/*",
    "/usr/local/lib/python3.12/dist-packages/*",
    "/opt/megatron-lm/_remote_module_non_scriptable",
    # These fused paths are mostly Triton DSL kernels and PyTorch
    # autograd-engine callbacks. The CUDA code is exercised by unit tests,
    # but coverage.py cannot reliably trace those execution contexts.
    "/opt/megatron-lm/megatron/core/fusions/*",
]
parallel = true
sigterm = false

[tool.coverage.paths]
source = [
    ".",
    "/opt/megatron-lm/",
    "/opt/megatron-lm-legacy/",
    "/home/runner/work/Megatron-LM/Megatron-LM/",
]

[tool.ruff.lint]
# Enable all `pydocstyle` rules, limiting to those that adhere to the
# Google convention via `convention = "google"`, below.
# select = ["D", "F"]
select = ["S506"]

# - On top of the Google convention, disable `D417`, which requires
#   documentation for every function parameter.
# - F841: local variable assigned but never used (exluced to favor readability)
# TODO: Remove D10 once we are about to release to get all the docstrings written
ignore = ["D417", "D10", "F841"]

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

# Section to exclude errors for different file types
[tool.ruff.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"]