Dockerfile: build from the official PyTorch image by artbataev · Pull Request #15638 · NVIDIA-NeMo/NeMo

artbataev · 2026-04-22T20:09:33Z

Important

The Update branch button must only be pressed in very rare occassions.
An outdated branch is never blocking the merge of a PR.
Please reach out to the automation team before pressing that button.

What does this PR do ?

Add Dockerfile.stable to build a container with NeMo from official PyTorch container https://hub.docker.com/r/pytorch/pytorch/tags

Collection: [Note which collection this PR will affect]

Changelog

Add specific line by line info of high level changes in this PR.

Usage

You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

GitHub Actions CI

The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.

The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first click "Approve and run".

Before your PR is "Ready for review"

Pre checks:

Make sure you read and followed Contributor guidelines
Did you write any new necessary tests?
Did you add or update any necessary documentation?
Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
- Reviewer: Does the PR have correct import guards for all optional libraries?

PR Type:

New Feature
Bugfix
Documentation

If you haven't finished some of the above items you can still open "Draft" PR.

Who can review?

Anyone in the community is free to review the PR once the checks have passed.
Contributor guidelines contains specific people who can review PRs to various areas.

Additional Information

Related to # (issue)

copy-pr-bot · 2026-04-22T20:09:37Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

pzelasko · 2026-04-22T20:15:49Z

+COPY tutorials /workspace/nemo/tutorials
+# COPY README.rst LICENSE /workspace/nemo/
+
+# Fix for Megatron: set a permanent location inside the container (NOT /tmp, which might get wiped)


Where is Megatron coming from? We removed it

Automodel installs megatron-fsdp
https://github.com/NVIDIA-NeMo/Automodel/blob/e34a12c68021f540c85b6e5103ccc1ed2cd04119/pyproject.toml#L78

I thought megatron-fsdp is a standalone lib with no dependency on Megatron Core / LM

pzelasko · 2026-04-22T20:16:28Z

+# Check install
+# NB: adjusting LD_LIBRARY_PATH (only here, should not be persistent!) is a temporary hack
+# to avoid failure if CUDA is unavailable (`docker build` does not expose GPUs)
+# The error is raised in NeMo Core, and the main reason is reinstalled Transformer-Engine;


What is installing Transformer Engine? We can install it for Automodel's fast-path in the container, but I'm surprised it's present somehow

pzelasko · 2026-04-22T20:16:36Z

+# The error is raised in NeMo Core, and the main reason is reinstalled Transformer-Engine;
+RUN CHECK_MSG=$(export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:${CUDA_HOME}/compat/lib.real && \
+  python -c "import nemo.collections.asr as nemo_asr" && \
+  python -c "import nemo.collections.nlp as nemo_nlp" && \


This collection no longer exists

pzelasko · 2026-04-22T20:17:56Z

+  apt-get install -y \
+  libtool \
+  libltdl-dev \
+  automake \


What do we need various build tools listed here for? BTW I thought we removed K2 dependency (mostly), and it's not in this container - do we want to keep supporting it?

pzelasko · 2026-04-22T20:18:49Z

+
+# base image: official PyTorch image
+ARG BASE_IMAGE=pytorch/pytorch:2.10.0-cuda13.0-cudnn9-devel
+# NB: automodel requires torch<=2.10.0, cannot use 2.11.0 for now


@adil-a is this true? if yes, why is max torch version capped in Automodel?

https://github.com/NVIDIA-NeMo/Automodel/blob/e34a12c68021f540c85b6e5103ccc1ed2cd04119/pyproject.toml#L84 (current main)
in dependencies: "torch>=2.6.0,<=2.10.0"

Change for max version is from here: https://github.com/NVIDIA-NeMo/Automodel/pull/1004/changes (was 2.9.0 before this PR)

Updated automodel -> now using pytorch 2.11.0

pzelasko · 2026-04-22T20:19:08Z

+  else echo "Import check success"; fi
+
+
+# copy scripts/examples/tests into container for end user


Why not copy the entire repo?

I think that the idea is to copy only necessary separate parts of the repo. NeMo is installed in non-editable mode. But if copied, it can be used or not used (through automatic addition to PYTHONPATH) depending on the directory used.
I would prefer to keep it as is (copy only scripts/examples/test/tutorials). If someone wants to change NeMo code - it is possible through mounting the directory and manually adjusting PYTHONPATH.

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

pzelasko

LGTM

github-actions · 2026-04-25T03:09:54Z

[🤖]: Hi @artbataev 👋,

We wanted to let you know that a CICD pipeline for this PR just finished successfully.

So it might be time to merge this PR or get some approvals.

artbataev requested a review from pzelasko April 22, 2026 20:09

artbataev force-pushed the vbataev/updated_dockerfile branch from 5e37348 to 76217fb Compare April 22, 2026 20:17

pzelasko reviewed Apr 22, 2026

View reviewed changes

artbataev added 3 commits April 22, 2026 13:53

Add Dockerfile from official PyTorch container

2c66126

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

Fix build

3ec0a7b

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

Fix collections check

64b7733

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

artbataev force-pushed the vbataev/updated_dockerfile branch from 0e9ca32 to 64b7733 Compare April 22, 2026 20:53

artbataev added 2 commits April 23, 2026 23:28

Merge branch 'main' into vbataev/updated_dockerfile

4ae8c6e

Clean up

0463c7d

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

artbataev marked this pull request as ready for review April 23, 2026 19:50

artbataev changed the title ~~Update Dockerfile: build from the official PyTorch image~~ Dockerfile: build from the official PyTorch image Apr 23, 2026

pzelasko approved these changes Apr 24, 2026

View reviewed changes

artbataev added the Run CICD label Apr 24, 2026

artbataev temporarily deployed to test April 24, 2026 21:27 — with GitHub Actions Inactive

artbataev enabled auto-merge (squash) April 24, 2026 21:27

pzelasko mentioned this pull request Apr 24, 2026

Upgrade base image to pytorch:26.03-py3 and remove Apex/TE/Megatron #15611

Closed

8 tasks

github-actions Bot removed the Run CICD label Apr 25, 2026

artbataev merged commit 2223816 into main Apr 25, 2026
267 checks passed

artbataev deleted the vbataev/updated_dockerfile branch April 25, 2026 03:10

		else echo "Import check success"; fi


		# copy scripts/examples/tests into container for end user

Conversation

artbataev commented Apr 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Changelog

Usage

GitHub Actions CI

Before your PR is "Ready for review"

Who can review?

Additional Information

Uh oh!

copy-pr-bot Bot commented Apr 22, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

artbataev Apr 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pzelasko left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Apr 25, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

artbataev commented Apr 22, 2026 •

edited

Loading

artbataev Apr 22, 2026 •

edited

Loading