Restore Megatron extra-based setup

Author: Kovbo

pyproject.toml

@@ -41,6 +41,19 @@ backend = [
     "nvidia-cudnn-frontend<1.21 ; sys_platform == 'linux'",
     "vllm @ https://github.com/vivekkalyan/vllm/releases/download/v0.17.0-art1/vllm-0.17.0%2Bart1-cp38-abi3-manylinux_2_31_x86_64.whl ; sys_platform == 'linux'",
 ]
+megatron = [
+    "torch==2.10.0",
+    "quack-kernels==0.2.5",
+    "apex",
+    "transformer-engine==2.11.0",
+    "transformer-engine-cu12==2.11.0",
+    "transformer-engine-torch==2.11.0",
+    "megatron-core==0.16.0rc0",
+    "pybind11>=2.13.6",
+    "megatron-bridge",
+    "nvidia-ml-py==13.580.82",
+    "ml-dtypes>=0.5.0 ; python_full_version < '3.13'",
+]
 langgraph = [
     "langchain-core>=0.3.51",
     "langgraph>=0.6.2",
@@ -78,7 +91,6 @@ packages = ["src/art", "src/mp_actors"]
 
 [tool.hatch.build.targets.wheel.force-include]
 ".agents/skills" = "art/skills"
-"src/art_megatron_install.py" = "art_megatron_install.py"
 
 [tool.hatch.build]
 sources = ["src"]
@@ -226,3 +238,6 @@ dev = [
 
 [tool.uv.sources]
 panza = { git = "https://github.com/corbt/panza.git" }
+apex = { git = "https://github.com/NVIDIA/apex.git", branch = "25.09" }
+megatron-bridge = { git = "https://github.com/NVIDIA-NeMo/Megatron-Bridge.git", rev = "75f2c5ad4afb702b57b4781a00f5291a66bcf183" }
+transformer-engine-torch = { git = "https://github.com/NVIDIA/TransformerEngine.git", tag = "v2.11", subdirectory = "transformer_engine/pytorch" }

src/art/megatron/setup.sh

@@ -7,6 +7,9 @@ export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-9.0}"
 apt-get update
 apt-get install -y libcudnn9-headers-cuda-12 ninja-build
 
-# Python dependencies are installed through art_megatron_install so
-# downstream repos can reuse the same source of truth for versions and VCS pins.
-uv run python -m art_megatron_install
+# Python dependencies are declared in pyproject.toml extras.
+# Keep backend + megatron together so setup does not prune runtime deps (e.g. vllm).
+script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
+repo_root="$(cd -- "${script_dir}/../../.." && pwd)"
+cd "${repo_root}"
+uv sync --extra backend --extra megatron --frozen --active