# 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.package-data] "megatron.core.allocator" = ["csrc/*.cpp"] [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.47; 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>=13.0.0", "flashinfer-python>=0.5.0,<0.8.0", "nvidia-cudnn-frontend[cutedsl]==1.29.1", "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.29 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+44f31bbe" 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 = "ea2b7805b97956268bed1adf94591f428ded6cc1" } transformer-engine = { git = "https://github.com/NVIDIA/TransformerEngine.git", rev = "44f31bbeaf99bd30a67b928000d4b980026debd2" } 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"]