Guide de compilation et dépendances

Le principe fondamental : compiler et développer à l'intérieur de conteneurs — le conteneur CI embarque le toolkit CUDA correct, la compilation PyTorch et les extensions natives pré-compilées (TransformerEngine, DeepEP, …) qui ne peuvent pas être reproduites sur un host nu.

Pourquoi les conteneurs

Megatron-LM dépend de CUDA, NCCL, PyTorch avec support GPU, TransformerEngine, et de composants optionnels comme ModelOpt et DeepEP. Installer ces éléments sur un host nu est fragile et difficile à reproduire. Le projet fournit des Dockerfiles qui figent chaque dépendance.

Utilisez le conteneur comme environnement de développement. Cela garantit :

• Des versions identiques de CUDA / NCCL / cuDNN sur tous les développeurs et en CI.
• uv.lock se résout de la même manière localement et en CI.
• Les opérations dépendantes du GPU (entraînement, tests) fonctionnent directement.

dev vs lts

Deux variantes d'image existent, contrôlées par l'argument de compilation IMAGE_TYPE et le label PR container::lts :

Utilisez dev pour tout sauf si vous avez une raison spécifique de tester lts. La CI exécute dev par défaut ; ajoutez container::lts à une PR seulement pour vérifier la compatibilité avec la pile stable (p. ex. une mise à niveau de dépendance qui ne doit pas casser les utilisateurs LTS). Le marqueur @pytest.mark.flaky_in_dev ignore les tests dans l'environnement dev ; @pytest.mark.flaky les ignore dans lts.

Étape 1 — Acquérir une image

Option A — NVIDIA-interne : récupérer une image construite par la CI

⚠️ Requiert l'accès à l'instance GitLab interne. Voir @tools/trigger_internal_ci.md pour la configuration (ajouter le remote git, obtenir un token).

La CI GitLab interne publie les images dans son registry de conteneurs. Dérivez l'hôte du registry depuis votre remote gitlab configuré — le même hôte que vous utilisez pour trigger_internal_ci.py :

# Dérivez l'hôte de votre remote 'gitlab' :
GITLAB_HOST=$(git remote get-url gitlab | sed 's/.*@\(.*\):.*/\1/')

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

Option B — Construire à partir de zéro (fonctionne pour tout le monde)

⚠️ Dockerfile.ci.dev a deux étapes : main et jet. L'étape jet requiert un secret de compilation interne et échouera sans lui. Passez toujours --target main pour s'arrêter à l'étape publique.

# image dev (par défaut)
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 .

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

La variante d'image utilisée est contrôlée par le label PR container::lts ; en l'absence de ce label, dev est utilisé.

Étape 2 — Lancer le conteneur

Option A — Runtime Docker local

docker run --rm --gpus all \
  -v $(pwd):/workspace \
  -w /workspace \
  megatron-lm:local \
  bash -c "<your command>"

Option B — Cluster Slurm (pour ceux sans runtime Docker local)

Les clusters NVIDIA utilisent généralement Pyxis + enroot. Demandez une session interactive :

srun \
  --nodes=1 --gpus-per-node=8 \
  --container-image megatron-lm:local \
  --container-mounts $(pwd):/workspace \
  --container-workdir /workspace \
  --pty bash

Pour les clusters qui requièrent une archive .sqsh d'abord :

enroot import -o megatron-lm.sqsh dockerd://megatron-lm:local
srun \
  --nodes=1 --gpus-per-node=8 \
  --container-image $(pwd)/megatron-lm.sqsh \
  --container-mounts $(pwd):/workspace \
  --container-workdir /workspace \
  --pty bash

Gestion des dépendances

Les dépendances sont déclarées dans pyproject.toml. Le venv se trouve à /opt/venv à l'intérieur du conteneur (déjà sur PATH).

Toutes les opérations uv doivent être exécutées à l'intérieur du conteneur. N'exécutez jamais uv sync / uv pip install sur le host.

Groupes de dépendances uv

Commandes d'installation (à l'intérieur du conteneur) :

# Environnement dev + test complet
uv sync --locked --group dev --group test

# Linting seulement
uv sync --locked --only-group linting

# Environnement LTS
uv sync --locked --group lts --group test

Plusieurs dépendances sont sourcées directement depuis git (TransformerEngine, nemo-run, FlashMLA, Emerging-Optimizers, nvidia-resiliency-ext). Le fichier verrouillé uv.lock épingle les révisions exactes ; mettez-le à jour avec uv lock quand vous modifiez pyproject.toml.

Ajouter une nouvelle dépendance

Suivez ce flux de travail en trois étapes :

1. Acquérir une image conteneur — voir Étape 1 ci-dessus.

2. Lancer le conteneur de façon interactive — voir Étape 2 ci-dessus.

3. Mettre à jour le fichier de verrou à l'intérieur du conteneur, puis le commiter :

   # À l'intérieur du conteneur :
   uv add <package>          # ajoute à pyproject.toml et résout
   uv lock                   # régénère uv.lock
   # Quittez le conteneur, puis sur le host :
   git add pyproject.toml uv.lock
   git commit -S -s -m "build: add <package> dependency"

Résoudre un conflit de fusion dans uv.lock

uv.lock est généré automatiquement ; ne résolvez jamais les conflits manuellement. À la place :

git checkout origin/main -- uv.lock   # prenez la version de main comme base
# puis à l'intérieur du conteneur :
uv lock                               # re-résolvez sur vos changements pyproject.toml

Pièges courants