Megatron-LM Container and Dependency Setup

Build & Dependency Guide

The core principle: build and develop inside containers — the CI container ships the correct CUDA toolkit, PyTorch build, and pre-compiled native extensions (TransformerEngine, DeepEP, …) that cannot be reproduced on a bare host.

Answer-First Constants

For text-only dependency or container questions, give these repo-specific facts up front before the longer workflow:

• Run dependency work inside the Megatron-LM CI container, not on the host.
• The container venv is /opt/venv, already on PATH.
• Default dev uses docker/.ngc_version.dev and the dev uv group; lts uses docker/.ngc_version.lts and the lts uv group. The container::lts PR label selects the LTS path; otherwise CI uses dev.
• lts is opt-in only when the user explicitly asks for it. It is the older long-term-support base, not a routine second lane — never attach container::lts, build the LTS image, or run the lts uv group on your own initiative, not even for a container or dependency change.
• Install commands inside the container: uv sync --locked --group dev --group test, uv sync --locked --only-group linting, or uv sync --locked --group lts --group test.
• Dependency edits use uv add <package> followed by uv lock, both inside the container.
• docker/Dockerfile.ci.dev has main and jet stages. The jet stage needs an internal secret; local/public builds should pass --target main.

Why Containers

Megatron-LM depends on CUDA, NCCL, PyTorch with GPU support, TransformerEngine, and optional components like ModelOpt and DeepEP. Installing these on a bare host is fragile and hard to reproduce. The project ships Dockerfiles that pin every dependency.

Use the container as your development environment. This guarantees:

• Identical CUDA / NCCL / cuDNN versions across all developers and CI.
• uv.lock resolves the same way locally and in CI.
• GPU-dependent operations (training, testing) work out of the box.

dev vs lts

Two image variants exist, each with its own Dockerfile, selected by the container::lts PR label. The defining difference is the base container: dev tracks the latest NGC PyTorch release, while lts ("long-term support") pins the previous, still-supported NGC PyTorch/CUDA release. container::lts exists to verify a change still works on that older base — the dependency differences below follow from it, they are not the point.

LTS deps used to live in [project.optional-dependencies].lts in pyproject.toml. They were moved into docker/lts/requirements.txt so pyproject.toml can host meaningful module-level extras without colliding with the LTS pin set. To bump an LTS dependency, edit the version in docker/lts/requirements.txt and rebuild docker/Dockerfile.ci.lts.

Use dev for everything. lts is off-limits unless the user explicitly asks for it. CI runs dev by default, and that is the only variant you touch on your own initiative. Treat container::lts as a high barrier, not a fallback: do not attach the label, build docker/Dockerfile.ci.lts, or run the lts uv group unless the user has explicitly requested LTS validation — not even for a container or dependency change. When they do ask, container::lts verifies the change still works on the older long-term-support PyTorch/CUDA base that LTS users run. The @pytest.mark.flaky_in_dev marker skips tests in the dev environment; @pytest.mark.flaky skips them in lts.

Step 1 — Acquire an Image

Option A — NVIDIA-internal: pull a CI-built image

⚠️ Requires access to the internal GitLab instance. See @tools/trigger_internal_ci.md for setup (adding the git remote, obtaining a token).

The internal GitLab CI publishes images to its container registry. Derive the registry host from your configured gitlab remote — the same host you use for trigger_internal_ci.py:

# Derive host from your 'gitlab' remote:
GITLAB_HOST=$(git remote get-url gitlab | sed 's/.*@\(.*\):.*/\1/')

docker pull ${GITLAB_HOST}/adlr/megatron-lm/mcore_ci_dev:main

Option B — Build from scratch (works for everyone)

⚠️ Dockerfile.ci.dev has two stages: main and jet. The jet stage requires an internal build secret and will fail without it. Always pass --target main to stop at the public stage.

# dev image (default)
docker build \
  --target main \
  --build-arg FROM_IMAGE_NAME=$(cat docker/.ngc_version.dev) \
  --build-arg IMAGE_TYPE=dev \
  -f docker/Dockerfile.ci.dev \
  -t megatron-lm:local .

# lts image (uses a dedicated Dockerfile; no IMAGE_TYPE arg)
docker build \
  --target main \
  --build-arg FROM_IMAGE_NAME=$(cat docker/.ngc_version.lts) \
  -f docker/Dockerfile.ci.lts \
  -t megatron-lm:local-lts .

Which image variant is used is controlled by the PR label container::lts; absent that label, dev is used.