Deprecate ModelOpt custom docker and directly use TRT-LLM / PyTorch / TRT docker (NVIDIA#346)

Signed-off-by: Keval Morabia <[email protected]>

Author: kevalmorabia97

.github/CODEOWNERS

@@ -30,7 +30,6 @@ modelopt/torch/trace @NVIDIA/modelopt-torch-nas-prune-codeowners
 modelopt/torch/utils @NVIDIA/modelopt-torch-utils-codeowners
 
 # Examples
-/docker @NVIDIA/modelopt-docker-codeowners
 /README.md @NVIDIA/modelopt-examples-codeowners
 /examples @NVIDIA/modelopt-examples-codeowners
 /examples/chained_optimizations @NVIDIA/modelopt-torch-nas-prune-codeowners

.github/workflows/example_tests.yml

@@ -68,15 +68,17 @@ jobs:
     container: &example_container
       image: nvcr.io/nvidia/tensorrt-llm/release:1.1.0rc2.post2
       env:
-        LD_LIBRARY_PATH: "/usr/lib/x86_64-linux-gnu:/usr/local/tensorrt/targets/x86_64-linux-gnu/lib:${LD_LIBRARY_PATH}"
-        # PATH: "/usr/local/tensorrt/targets/x86_64-linux-gnu/bin:${PATH}"
         PIP_CONSTRAINT: "" # Disable pip constraint for upgrading packages
     steps: &example_steps
       - uses: actions/checkout@v4
       - uses: nv-gha-runners/setup-proxy-cache@main
+      - name: Setup environment variables
+        run: |
+          echo "LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/include:/usr/lib/x86_64-linux-gnu:/usr/local/tensorrt/targets/x86_64-linux-gnu/lib" >> $GITHUB_ENV
+          echo "PATH=${PATH}:/usr/local/tensorrt/targets/x86_64-linux-gnu/bin" >> $GITHUB_ENV
       - name: Run example tests
         run: |
-          pip install ".[all,dev-test]"
+          pip install ".[hf,dev-test]"
           find examples/${{ matrix.EXAMPLE }} -name "requirements.txt" | while read req_file; do pip install -r "$req_file" || exit 1; done
           pytest -s tests/examples/${{ matrix.EXAMPLE }}
   example-tests-non-pr:

.github/workflows/gpu_tests.yml

@@ -66,11 +66,14 @@ jobs:
       image: nvcr.io/nvidia/pytorch:25.06-py3
       env:
         GIT_DEPTH: 1000 # For correct version for tests/gpu/torch/quantization/plugins/test_megatron.py
-        LD_LIBRARY_PATH: "/usr/lib/x86_64-linux-gnu:${LD_LIBRARY_PATH}" # Add libcudnn*.so and libnv*.so to path.
         PIP_CONSTRAINT: "" # Disable pip constraint for upgrading packages
     steps: &gpu_steps
       - uses: actions/checkout@v4
       - uses: nv-gha-runners/setup-proxy-cache@main
+      - name: Setup environment variables
+        run: |
+          echo "LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/include:/usr/lib/x86_64-linux-gnu:/usr/local/tensorrt/targets/x86_64-linux-gnu/lib" >> $GITHUB_ENV
+          echo "PATH=${PATH}:/usr/local/tensorrt/targets/x86_64-linux-gnu/bin" >> $GITHUB_ENV
       - name: Run gpu tests
         run: pip install tox-current-env && tox -e py312-cuda12-gpu --current-env
   gpu-tests-non-pr:

.gitlab/tests.yml

@@ -1,11 +1,12 @@
-# NOTE: Make sure this file is consistent with .github/workflows/{unit,gpu}_tests.yml
+# NOTE: Make sure this file is consistent with .github/workflows/{unit,gpu,example}_tests.yml
 .tests-default:
+  variables:
+    PIP_CONSTRAINT: "" # Disable pip constraint for upgrading packages
   stage: tests
   rules:
     - if: $CI_PIPELINE_SOURCE == "schedule"
-      when: always
-    - if: $CI_PIPELINE_SOURCE != "schedule"
-      when: manual
+    - if: $CI_COMMIT_TAG =~ /^\d+\.\d+\.\d+$/
+    - when: manual
 
 ##### Unit Tests #####
 unit:
@@ -24,50 +25,74 @@ unit:
     - tox -e py3$PYTHON-torch$TORCH-tf_$TRANSFORMERS-unit
 
 ##### GPU Tests #####
-gpu:
+.multi-gpu-tests-default:
   extends: .tests-default
   timeout: 60m
   image: nvcr.io/nvidia/pytorch:25.06-py3
   variables:
     GIT_DEPTH: 1000 # For correct version for tests/gpu/torch/quantization/plugins/test_megatron.py
-    LD_LIBRARY_PATH: "/usr/lib/x86_64-linux-gnu:${LD_LIBRARY_PATH}" # Add libcudnn*.so and libnv*.so to path.
-    PIP_CONSTRAINT: "" # Disable pip constraint for upgrading packages
   tags: [docker, linux, 2-gpu]
+  before_script:
+    # Add libcudnn*.so and libnv*.so to path
+    - export LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/usr/include:/usr/lib/x86_64-linux-gnu:/usr/local/tensorrt/targets/x86_64-linux-gnu/lib"
+    # Add trtexec to path
+    - export PATH="${PATH}:/usr/local/tensorrt/targets/x86_64-linux-gnu/bin"
+    # Install git-lfs for Daring-Anteater dataset
+    - apt-get update && apt-get install -y git-lfs
+    - git lfs install --system
+
+multi-gpu:
+  extends: .multi-gpu-tests-default
   script:
     # Use pre-installed packages without a new venv with tox-current-env
     - pip install tox-current-env
     - tox -e py312-cuda12-gpu --current-env
 
 ##### Example Tests #####
-example:
-  extends: .tests-default
-  stage: tests
-  timeout: 45m
-  image: gitlab-master.nvidia.com:5005/omniml/modelopt/modelopt_examples:latest
-  variables:
-    TEST_TYPE: pytest
-  tags: [docker, linux, 2-gpu, sm<89]
+example-torch:
+  extends: .multi-gpu-tests-default
+  timeout: 30m
   parallel:
     matrix:
-      - EXAMPLE: [diffusers, llm_distill, llm_qat, llm_sparsity, onnx_ptq, speculative_decoding]
-  allow_failure: true # Allow to continue next stages even if job is canceled (e.g. during release)
-  before_script:
-    - pip install ".[all,dev-test]"
+      - EXAMPLE: [llm_distill, llm_sparsity, speculative_decoding]
   script:
-    # Uninstall apex since T5 Int8 (PixArt) + Apex is not supported as per https://github.com/huggingface/transformers/issues/21391
-    - if [ "$EXAMPLE" = "diffusers" ]; then pip uninstall -y apex; fi
+    - pip install ".[hf,dev-test]"
     - find examples/$EXAMPLE -name "requirements.txt" | while read req_file; do pip install -r "$req_file" || exit 1; done
-    - if [ "$TEST_TYPE" = "pytest" ]; then pytest -s tests/examples/$EXAMPLE; else bash tests/examples/test_$EXAMPLE.sh; fi
+    - pytest -s tests/examples/$EXAMPLE
 
-example-ada:
-  extends: example
+# TODO: Fix llm_qat test hang in GitLab CI
+example-failing:
+  extends: example-torch
+  allow_failure: true
+  parallel:
+    matrix:
+      - EXAMPLE: [llm_qat]
+
+example-trtllm:
+  extends: example-torch
   timeout: 60m
+  image: nvcr.io/nvidia/tensorrt-llm/release:1.1.0rc2.post2
+  tags: [docker, linux, 2-gpu, sm>=89]
+  parallel:
+    matrix:
+      - EXAMPLE: [llm_autodeploy, llm_eval, llm_ptq, vlm_ptq]
+
+example-onnx:
+  extends: example-torch
+  image: nvcr.io/nvidia/tensorrt:25.08-py3
   tags: [docker, linux, 2-gpu, sm>=89]
   parallel:
     matrix:
-      - EXAMPLE: [llm_eval, llm_ptq, vlm_ptq, llm_autodeploy]
+      - EXAMPLE: [diffusers, onnx_ptq]
+        TEST_TYPE: pytest
       - EXAMPLE: [onnx_ptq]
         TEST_TYPE: bash
+  script:
+    # Uninstall apex since T5 Int8 (PixArt) + Apex is not supported as per https://github.com/huggingface/transformers/issues/21391
+    - if [ "$EXAMPLE" = "diffusers" ]; then pip uninstall -y apex; fi
+    - pip install ".[all,dev-test]"
+    - find examples/$EXAMPLE -name "requirements.txt" | while read req_file; do pip install -r "$req_file" || exit 1; done
+    - if [ "$TEST_TYPE" = "pytest" ]; then pytest -s tests/examples/$EXAMPLE; else bash tests/examples/test_$EXAMPLE.sh; fi
 
 ##### Megatron / NeMo Integration Tests #####
 megatron-nemo-integration:

CHANGELOG.rst

@@ -6,6 +6,7 @@ Model Optimizer Changelog (Linux)
 
 **Deprecations**
 
+- Deprecated ModelOpt's custom docker images. Please use the PyTorch, TensorRT-LLM or TensorRT docker image directly or refer to the `installation guide <https://nvidia.github.io/TensorRT-Model-Optimizer/getting_started/2_installation.html>`_ for more details.
 - Deprecated ``quantize_mode`` argument in ``examples/onnx_ptq/evaluate.py`` to support strongly typing. Use ``engine_precision`` instead.
 - Deprecated TRT-LLM's TRT backend in ``examples/llm_ptq`` and ``examples/vlm_ptq``. Tasks ``build`` and ``benchmark`` support are removed and replaced with ``quant``. For performance evaluation, please use ``trtllm-bench`` directly.
 - ``--export_fmt`` flag in ``examples/llm_ptq`` is removed. By default we export to the unified Hugging Face checkpoint format.

CONTRIBUTING.md

@@ -11,7 +11,7 @@ pip install -e ".[dev]"
 ```
 
 If you are working on features that require dependencies like TensorRT-LLM or Megatron-Core, consider using a docker container to simplify the setup process.
-See [docker README](./README.md#installation--docker) for more details.
+Visit our [installation docs](https://nvidia.github.io/TensorRT-Model-Optimizer/getting_started/2_installation.html) for more information.
 
 ## 🧹 Code linting and formatting
 

README.md

@@ -61,10 +61,10 @@ Model Optimizer is also integrated with [NVIDIA NeMo](https://github.com/NVIDIA-
 To install stable release packages for Model Optimizer with `pip` from [PyPI](https://pypi.org/project/nvidia-modelopt/):
 
 ```bash
-pip install nvidia-modelopt[all]
+pip install -U nvidia-modelopt[all]
 ```
 
-To install from source in editable mode with all development dependencies or to test the latest changes, run:
+To install from source in editable mode with all development dependencies or to use the latest features, run:
 
 ```bash
 # Clone the Model Optimizer repository
@@ -74,7 +74,11 @@ cd TensorRT-Model-Optimizer
 pip install -e .[dev]
 ```
 
-Visit our [installation guide](https://nvidia.github.io/TensorRT-Model-Optimizer/getting_started/2_installation.html) for more fine-grained control on installed dependencies or view our pre-made [dockerfiles](docker/README.md) for more information.
+You can also directly use the [TensorRT-LLM docker images](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/tensorrt-llm/containers/release/tags)
+(e.g., `nvcr.io/nvidia/tensorrt-llm/release:<version>`), which have Model Optimizer pre-installed.
+Make sure to upgrade Model Optimizer to the latest version using ``pip`` as described above.
+Visit our [installation guide](https://nvidia.github.io/TensorRT-Model-Optimizer/getting_started/2_installation.html) for
+more fine-grained control on installed dependencies or for alternative docker images and environment variables to setup.
 
 ## Techniques