Move LTS dependencies from pyproject.toml to Dockerfile.ci.lts

.github/workflows/cicd-main.yml

@@ -598,10 +598,12 @@ jobs:
             NGC_VERSION=$(cat docker/.ngc_version.lts)
             echo "version=$NGC_VERSION" | tee -a $GITHUB_OUTPUT
             echo "image_type=lts" | tee -a $GITHUB_OUTPUT
+            echo "dockerfile=./docker/Dockerfile.ci.lts" | tee -a $GITHUB_OUTPUT
           else
             NGC_VERSION=$(cat docker/.ngc_version.dev)
             echo "version=$NGC_VERSION" | tee -a $GITHUB_OUTPUT
             echo "image_type=dev" | tee -a $GITHUB_OUTPUT
+            echo "dockerfile=./docker/Dockerfile.ci.dev" | tee -a $GITHUB_OUTPUT
           fi
 
       - name: Set up Docker Buildx
@@ -610,7 +612,7 @@ jobs:
       - name: Build and push
         uses: docker/build-push-action@v7.1.0
         with:
-          file: ./docker/Dockerfile.ci.dev
+          file: ${{ steps.base-image.outputs.dockerfile }}
           push: true
           context: .
           target: main

.gitlab/stages/01.build.yml

@@ -52,13 +52,11 @@ test:pre_build_image:
   parallel:
     matrix:
       - IMAGE: CI_MCORE_LTS_IMAGE
-        FILE: Dockerfile.ci.dev
-        IMAGE_TYPE: lts
+        FILE: Dockerfile.ci.lts
         BASE_IMAGE: nvcr.io/nvidia/pytorch:25.09-py3
         PLATFORM: amd64
       - IMAGE: CI_MCORE_LTS_IMAGE
-        FILE: Dockerfile.ci.dev
-        IMAGE_TYPE: lts
+        FILE: Dockerfile.ci.lts
         BASE_IMAGE: nvcr.io/nvidia/pytorch:25.09-py3
         PLATFORM: arm64
       - IMAGE: CI_MCORE_DEV_IMAGE

docker/Dockerfile.ci.lts

@@ -0,0 +1,138 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# syntax=docker/dockerfile:1.3-labs
+#
+# LTS CI image for Megatron-LM.
+#
+# The LTS image is bumped only once a year and intentionally lags behind the
+# floating dev tag. Its Python dependency set is therefore pinned directly in
+# this Dockerfile (rather than in pyproject.toml) so the pyproject can host
+# meaningful module-level extras (inference, RL, MoE, ...) without colliding
+# with the LTS pin set. See AUT-479 for background.
+#
+# To bump the LTS dependency set:
+#   1. Update the LTS_REQUIREMENTS list below.
+#   2. Rebuild and run the LTS CI lane.
+#   3. (Optional) Bump docker/.ngc_version.lts if also moving the base image.
+
+ARG FROM_IMAGE_NAME
+FROM ${FROM_IMAGE_NAME} as main
+ENV PIP_CONSTRAINT=""
+ENV DEBIAN_FRONTEND=noninteractive
+ARG UV_VERSION=0.7.2
+ARG YQ_VERSION=4.44.1
+ENV PATH="/root/.local/bin:$PATH"
+ARG UV_PROJECT_ENVIRONMENT=/opt/venv
+ENV UV_PROJECT_ENVIRONMENT=${UV_PROJECT_ENVIRONMENT}
+ENV VIRTUAL_ENV=$UV_PROJECT_ENVIRONMENT
+ENV PATH="$UV_PROJECT_ENVIRONMENT/bin:$PATH"
+ENV UV_LINK_MODE=copy
+
+RUN bash -ex <<"EOF"
+    apt-get update
+    apt-get install -y --no-install-recommends gettext python3-venv psmisc uuid-runtime
+    apt-get clean
+    python -m venv /opt/jet
+    ARCH=$(uname -m)
+    case "${ARCH}" in \
+      "x86_64") YQ_ARCH=amd64 ;; \
+      "aarch64") YQ_ARCH=arm64 ;; \
+      "armv7l") YQ_ARCH=arm ;; \
+      *) echo "Unsupported architecture: ${ARCH}" && exit 1 ;; \
+    esac
+    wget https://github.com/mikefarah/yq/releases/download/v${YQ_VERSION}/yq_linux_${YQ_ARCH} -O /usr/local/bin/yq
+    chmod a+x /usr/local/bin/yq
+    curl -LsSf https://astral.sh/uv/${UV_VERSION}/install.sh | sh
+EOF
+
+COPY README.md pyproject.toml uv.lock /workspace/
+COPY megatron/core/__init__.py /workspace/megatron/core/
+COPY megatron/core/package_info.py /workspace/megatron/core/
+ENV NVTE_BUILD_NUM_PHILOX_ROUNDS=3
+RUN --mount=type=cache,target=/root/.cache/uv \
+    bash -ex <<"EOF"
+    export NVTE_CUDA_ARCHS="80;90;100"
+    uv venv ${UV_PROJECT_ENVIRONMENT} --system-site-packages
+    uv sync --only-group build
+    uv sync --extra mlm --extra ssm --extra te --link-mode copy --locked \
+        --no-install-package torch \
+        --no-install-package torchvision \
+        --no-install-package triton \
+        --no-install-package transformer-engine-cu12 \
+        --no-install-package nvidia-cublas-cu12 \
+        --no-install-package nvidia-cuda-cupti-cu12 \
+        --no-install-package nvidia-cuda-nvrtc-cu12 \
+        --no-install-package nvidia-cuda-runtime-cu12 \
+        --no-install-package nvidia-cudnn-cu12 \
+        --no-install-package nvidia-cufft-cu12 \
+        --no-install-package nvidia-cufile-cu12 \
+        --no-install-package nvidia-curand-cu12 \
+        --no-install-package nvidia-cusolver-cu12 \
+        --no-install-package nvidia-cusparse-cu12 \
+        --no-install-package nvidia-cusparselt-cu12 \
+        --no-install-package nvidia-nccl-cu12
+EOF
+
+# LTS-specific Python dependencies.
+#
+# These used to live in `[project.optional-dependencies].lts` in pyproject.toml,
+# but were moved here (AUT-479) so pyproject.toml can host meaningful per-module
+# extras. Pins are reviewed at LTS bump time only.
+RUN --mount=type=cache,target=/root/.cache/uv \
+    bash -ex <<"EOF"
+    uv pip install \
+        "tqdm" \
+        "einops~=0.8" \
+        "tensorstore~=0.1,!=0.1.46,!=0.1.72" \
+        "multi-storage-client~=0.27" \
+        "opentelemetry-api~=1.33.1" \
+        "megatron-energon[av_decode]~=6.0" \
+        "av" \
+        "flashinfer-python>=0.5.0,<0.7.0" \
+        "wget" \
+        "onnxscript" \
+        "fastapi~=0.50" \
+        "datasets" \
+        "emerging_optimizers @ git+https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git@v0.2.0"
+EOF
+
+# Install DeepEP
+COPY docker/patches/deepep.patch /workspace/deepep.patch
+RUN bash -ex <<"EOF"
+    cd /workspace
+    uv pip install nvidia-nvshmem-cu13==3.4.5
+    pushd /opt/venv/lib/python3.12/site-packages/nvidia/nvshmem/lib/
+        ln -s libnvshmem_host.so.3 libnvshmem_host.so
+    popd
+
+    git clone --branch hybrid-ep https://github.com/deepseek-ai/DeepEP.git
+    pushd DeepEP
+        git checkout 34152ae28f80bcc3ee38d7a12cb2ad87cfd4ea72
+        patch -p1 < /workspace/deepep.patch
+    popd
+    TORCH_CUDA_ARCH_LIST="9.0 10.0 12.0" uv pip install --no-build-isolation -v DeepEP/.
+    rm -rf DeepEP
+EOF
+
+COPY assets/ /opt/data/
+ENV UV_PYTHON=$UV_PROJECT_ENVIRONMENT/bin/python
+
+##### For NVIDIANS only #####
+FROM main as jet
+ARG JET_API_VERSION
+ENV PATH="$PATH:/opt/jet/bin"
+RUN --mount=type=secret,id=JET_INDEX_URLS bash -ex <<"EOF"
+    JET_INDEX_URLS=$(cat /run/secrets/JET_INDEX_URLS)
+    python -m venv /opt/jet
+    /opt/jet/bin/pip install --no-cache-dir $JET_INDEX_URLS \
+        "jet-api==$JET_API_VERSION" "setuptools<82.0.0"
+EOF
+
+RUN --mount=type=secret,id=JET_INDEX_URLS \
+    --mount=type=secret,id=LOGGER_INDEX_URL bash -ex <<"EOF"
+    JET_INDEX_URLS=$(cat /run/secrets/JET_INDEX_URLS)
+    LOGGER_INDEX_URL=$(cat /run/secrets/LOGGER_INDEX_URL)
+    uv pip install --no-cache-dir --upgrade $LOGGER_INDEX_URL "one-logger"
+    uv pip install --no-cache-dir --upgrade "setuptools>=80"
+    uv pip install --no-cache-dir --upgrade $JET_INDEX_URLS "jet-client~=4.0"
+EOF
+###

docker/common/install.sh

@@ -55,6 +55,19 @@ if [[ "$ENVIRONMENT" != "dev" && "$ENVIRONMENT" != "lts" ]]; then
     exit 1
 fi
 
+# AUT-479: LTS Python dependencies were moved out of pyproject.toml into
+# docker/Dockerfile.ci.lts. This script targets the floating dev stack and no
+# longer builds a working LTS environment by itself.
+if [[ "$ENVIRONMENT" == "lts" ]]; then
+    echo "Error: --environment lts is no longer supported by install.sh."
+    echo "       LTS dependencies are pinned in docker/Dockerfile.ci.lts."
+    echo "       Build the LTS image directly with:"
+    echo "         docker build --target main \\"
+    echo "           --build-arg FROM_IMAGE_NAME=\$(cat docker/.ngc_version.lts) \\"
+    echo "           -f docker/Dockerfile.ci.lts -t megatron-lm:local-lts ."
+    exit 1
+fi
+
 main() {
     if [[ -n "${PAT:-}" ]]; then
         echo -e "machine github.com\n  login token\n  password $PAT" >~/.netrc

pyproject.toml

@@ -112,21 +112,12 @@ dev = [
     "orjson",
 ]
 
-lts = [
-    "tqdm",
-    "einops~=0.8",
-    "tensorstore~=0.1,!=0.1.46,!=0.1.72",
-    "multi-storage-client~=0.27",
-    "opentelemetry-api~=1.33.1",
-    "megatron-energon[av_decode]~=6.0",
-    "av",
-    "flashinfer-python>=0.5.0,<0.7.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",
-]
+### '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]",
@@ -191,7 +182,6 @@ no-build-isolation-package = [
     "transformer-engine-torch",
 ]
 link-mode = "copy"
-conflicts = [[{ extra = "lts" }, { extra = "dev" }]]
 # 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.

skills/build-and-dependency/SKILL.md

@@ -29,13 +29,19 @@ dependency.
 
 ## dev vs lts
 
-Two image variants exist, controlled by the `IMAGE_TYPE` build arg and the
+Two image variants exist, each with its own Dockerfile, selected by the
 `container::lts` PR label:
 
-| Variant | Base image pin | uv group | When used |
-|---------|---------------|----------|-----------|
-| **`dev`** | `docker/.ngc_version.dev` | `dev` | Default — CI, local development, most PRs |
-| **`lts`** | `docker/.ngc_version.lts` | `lts` | Stability testing; excludes ModelOpt and other bleeding-edge extras |
+| Variant | Base image pin | Dockerfile | Where deps live | When used |
+|---------|---------------|------------|-----------------|-----------|
+| **`dev`** | `docker/.ngc_version.dev` | `docker/Dockerfile.ci.dev` | `pyproject.toml` `dev` extra (uv-resolved) | Default — CI, local development, most PRs |
+| **`lts`** | `docker/.ngc_version.lts` | `docker/Dockerfile.ci.lts` | Inline `uv pip install` list at the top of `Dockerfile.ci.lts` | Stability testing; excludes ModelOpt and other bleeding-edge extras |
+
+> LTS deps used to live in `[project.optional-dependencies].lts` in
+> `pyproject.toml`. They were moved into `Dockerfile.ci.lts` so
+> `pyproject.toml` can host meaningful module-level extras without colliding
+> with the LTS pin set. To bump an LTS dependency, edit the `uv pip install`
+> block in `docker/Dockerfile.ci.lts`.
 
 **Use `dev` for everything unless you have a specific reason to test `lts`.**
 CI runs `dev` by default; attach `container::lts` to a PR only when verifying
@@ -78,12 +84,11 @@ docker build \
   -f docker/Dockerfile.ci.dev \
   -t megatron-lm:local .
 
-# lts image
+# lts image (uses a dedicated Dockerfile; no IMAGE_TYPE arg)
 docker build \
   --target main \
   --build-arg FROM_IMAGE_NAME=$(cat docker/.ngc_version.lts) \
-  --build-arg IMAGE_TYPE=lts \
-  -f docker/Dockerfile.ci.dev \
+  -f docker/Dockerfile.ci.lts \
   -t megatron-lm:local-lts .
 ```
 
@@ -146,11 +151,14 @@ inside the container (already on `PATH`).
 |-------|---------|
 | `training` | Runtime training extras |
 | `dev` | Full dev environment (TransformerEngine, ModelOpt, …) |
-| `lts` | LTS-safe subset (no ModelOpt) |
 | `test` | pytest, coverage, nemo-run |
 | `linting` | ruff, black, isort, pylint |
 | `build` | Cython, pybind11, nvidia-mathdx |
 
+> The previous `lts` extra has been emptied. LTS deps are pinned in
+> `docker/Dockerfile.ci.lts` rather than `pyproject.toml`. Do not add new
+> packages under `[project.optional-dependencies].lts`.
+
 Install commands (inside the container):
 
 ```bash
@@ -159,11 +167,12 @@ uv sync --locked --group dev --group test
 
 # Linting only
 uv sync --locked --only-group linting
-
-# LTS environment
-uv sync --locked --group lts --group test
 ```
 
+The LTS environment is reproduced by building `docker/Dockerfile.ci.lts`
+end-to-end; there is no `uv sync`-only equivalent because the LTS deps no
+longer live in `pyproject.toml`.
+
 Several dependencies are sourced directly from git (TransformerEngine, nemo-run,
 FlashMLA, Emerging-Optimizers, nvidia-resiliency-ext). The locked `uv.lock` file
 pins exact revisions; update it with `uv lock` when changing `pyproject.toml`.

skills/bump-base-image/SKILL.md

@@ -11,7 +11,7 @@ End-to-end workflow for moving Megatron-LM's CI to a newer `nvcr.io/nvidia/pytor
 ## Inputs to gather from the user
 
 1. **Target tag**, e.g. `26.04-py3`. NVIDIA NGC PyTorch containers are released as `nvcr.io/nvidia/pytorch:YY.MM-py3`.
-2. **Scope** — usually `dev` only. The `lts` pin (`docker/.ngc_version.lts`, plus the `IMAGE_TYPE: lts` rows in GitLab) is bumped on a different cadence; only touch it if the user explicitly asks.
+2. **Scope** — usually `dev` only. The `lts` pin (`docker/.ngc_version.lts`, plus the `FILE: Dockerfile.ci.lts` rows in GitLab) is bumped on a different cadence; only touch it if the user explicitly asks.
 3. **Workflow run ID** (optional but typical) — after the first CI run, the user will provide a GitHub Actions run ID for golden-value refresh.
 
 ## Workflow
@@ -54,7 +54,7 @@ GitLab CI does **not** read `docker/.ngc_version.dev`. It hardcodes `BASE_IMAGE`
   PLATFORM: arm64
 ```
 
-Leave the `IMAGE_TYPE: lts` rows alone. Quick sanity check before commit:
+Leave the `FILE: Dockerfile.ci.lts` rows alone. Quick sanity check before commit:
 
 ```bash
 rg -n '^\s*BASE_IMAGE: nvcr\.io/nvidia/pytorch:' .gitlab/stages/01.build.yml
@@ -132,12 +132,12 @@ All three lines should show `nvcr.io/nvidia/pytorch:<YY.MM>-py3`. If they don't,
 | `.gitlab/stages/01.build.yml`                                                                 | Update both `IMAGE_TYPE: dev` `BASE_IMAGE:` rows (amd64 + arm64)                       |
 | `tests/functional_tests/test_cases/**/golden_values_dev_dgx_{h100,gb200}.json`                | Refresh via the `update-golden-values` skill                                           |
 | `tests/test_utils/recipes/<arch>/<suite>.yaml`                                                | Flip drifting / hanging cases to `mr-broken` / `mr-github-broken` with an issue link    |
-| `docker/.ngc_version.lts`, `.gitlab/stages/01.build.yml` `IMAGE_TYPE: lts` rows               | **Skip unless explicitly bumping LTS.** LTS has its own release cadence.                |
+| `docker/.ngc_version.lts`, `.gitlab/stages/01.build.yml` `FILE: Dockerfile.ci.lts` rows       | **Skip unless explicitly bumping LTS.** LTS has its own release cadence and its own Dockerfile (`docker/Dockerfile.ci.lts`) where LTS Python deps are pinned inline. |
 
 ## Gotchas
 
 - **GitHub vs GitLab pins are independent.** `docker/.ngc_version.dev` only drives GitHub CI's local container build via `Dockerfile.ci.dev`. GitLab CI has its own hardcoded `BASE_IMAGE:` matrix in `.gitlab/stages/01.build.yml`. PR #4688 existed solely because #4611 forgot the second one — don't repeat this.
-- **Don't bump LTS along with dev.** The `IMAGE_TYPE: lts` rows and `docker/.ngc_version.lts` are stability-pinned for the `container::lts` label path. Bump them in a dedicated PR with its own LTS validation.
+- **Don't bump LTS along with dev.** The `FILE: Dockerfile.ci.lts` rows, `docker/Dockerfile.ci.lts`, and `docker/.ngc_version.lts` are stability-pinned for the `container::lts` label path. Bump them in a dedicated PR with its own LTS validation. LTS Python deps are pinned inline in `docker/Dockerfile.ci.lts` (not in `pyproject.toml`) — edit that block when an LTS dependency needs to move.
 - **Don't fix golden-value drift by hand.** Use `tests/test_utils/python_scripts/download_golden_values.py` via the `update-golden-values` skill. Hand-editing the JSONs invites diff noise and relative-difference regressions on subsequent bumps.
 - **`mr-broken` is a real scope, not a comment marker.** It keeps the recipe wired into the matrix (so it stays discoverable and runnable on demand) without gating merges. Don't delete the test case from the recipe.
 - **`/ok to test` is per-commit.** A new force-push or fixup commit needs a fresh `/ok to test <sha>` comment to re-trigger NVIDIA-runner CI on a fork PR.