Repository: NVIDIA/NeMo
Pull request #12672: build: Add trtllm
Author: login: ko3n1g
Labels: name: CI
Created: 2025-03-18T21:12:16Z | Merged: 2025-03-26T20:13:19Z

> [!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 a one line overview of what this PR aims to accomplish.

**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

```python
# 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](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [ ] 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](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md) contains specific people who can review PRs to various areas.

# Additional Information

- Related to # (issue)



Diff:
diff --git a/.github/workflows/cicd-main.yml b/.github/workflows/cicd-main.yml
index 05c2e19e56cd..ec91da143b3e 100644
--- a/.github/workflows/cicd-main.yml
+++ b/.github/workflows/cicd-main.yml
@@ -85,6 +85,8 @@ jobs:
         run: |
           BUILD_ARGS=$(cat << EOF
           BASE_IMAGE=$(cat requirements/manifest.json | jq -r '."ngc-pytorch"')
+          TRTLLM_REPO=$(cat requirements/manifest.json | jq -r '."vcs-dependencies"."trt-llm".repo')
+          TRTLLM_TAG=$(cat requirements/manifest.json | jq -r '."vcs-dependencies"."trt-llm".ref')
           MLM_REPO=$(cat requirements/manifest.json | jq -r '."vcs-dependencies"."megatron-lm".repo')
           MLM_TAG=$(cat requirements/manifest.json | jq -r '."vcs-dependencies"."megatron-lm".ref')
           TE_REPO=$(cat requirements/manifest.json | jq -r '."vcs-dependencies".transformer_engine.repo')
diff --git a/Dockerfile.ci b/Dockerfile.ci
index cb3340e685be..b4fa0e29da35 100644
--- a/Dockerfile.ci
+++ b/Dockerfile.ci
@@ -45,6 +45,15 @@ apt-get install -y bc libsox-fmt-all -y
 apt-get clean
 EOF
 
+ARG TRTLLM_REPO
+ARG TRTLLM_TAG
+RUN --mount=type=bind,from=nemo-bump,source=/opt/NeMo/reinstall.sh,target=/tmp/NeMo/reinstall.sh <<"EOF" bash -ex
+  
+  bash /tmp/NeMo/reinstall.sh --library trtllm --mode build
+  ls -al /tmp/Megatron-LM || true
+EOF
+
+
 ARG MLM_REPO
 ARG MLM_TAG
 RUN --mount=type=bind,from=nemo-bump,source=/opt/NeMo/reinstall.sh,target=/tmp/NeMo/reinstall.sh <<"EOF" bash -ex
diff --git a/reinstall.sh b/reinstall.sh
index 0326c929eebc..a54ac394cdc5 100755
--- a/reinstall.sh
+++ b/reinstall.sh
@@ -4,7 +4,7 @@ set -ex
 # List of all supported libraries (update this list when adding new libraries)
 # This also defines the order in which they will be installed by --libraries "all"
 ALL_LIBRARIES=(
-  # "te"
+  "trtllm"
   "mcore"
   "nemo"
   "vllm"
@@ -17,6 +17,48 @@ export INSTALL_DIR=${INSTALL_DIR:-"/opt"}
 export WHEELS_DIR=${WHEELS_DIR:-"$INSTALL_DIR/wheels"}
 export PIP=pip
 
+trtllm() {
+  local mode="$1"
+
+  curl -s https://packagecloud.io/install/repositories/github/git-lfs/script.deb.sh | bash &&
+    apt-get install git-lfs &&
+    git lfs install &&
+    apt-get clean
+
+  TRTLLM_REPO=${TRTLLM_REPO:-$(cat "$CURR/requirements/manifest.json" | jq -r '."vcs-dependencies"."trt_llm".repo')}
+  TRTLLM_TAG=${TRTLLM_TAG:-$(cat "$CURR/requirements/manifest.json" | jq -r '."vcs-dependencies"."trt_llm".ref')}
+  TRTLLM_DIR="$INSTALL_DIR/TensorRT-LLM"
+  if [ ! -d "$TRTLLM_DIR/.git" ]; then
+    rm -rf "$TRTLLM_DIR" &&
+      cd $(dirname "$TRTLLM_DIR") &&
+      git clone ${TRTLLM_REPO}
+  fi &&
+    pushd $TRTLLM_DIR &&
+    git checkout -f $TRTLLM_TAG &&
+    git lfs pull &&
+    popd
+
+  if [[ "$mode" == "build" ]]; then
+    if [[ -n "${NVIDIA_PYTORCH_VERSION}" ]]; then
+      cd $TRTLLM_DIR &&
+        . docker/common/install_tensorrt.sh \
+          --TRT_VER="10.8.0.43" \
+          --CUDA_VER="12.8" \
+          --CUDNN_VER="9.7.0.66-1" \
+          --NCCL_VER="2.25.1-1+cuda12.8" \
+          --CUBLAS_VER="12.8.3.14-1" \  &&
+        python3 ./scripts/build_wheel.py --job_count $(nproc) --trt_root /usr/local/tensorrt --python_bindings --benchmarks &&
+        pip wheel --wheel-dir $WHEELS_DIR/trtllm/ $TRTLLM_DIR
+    fi
+  else
+    if [ -d "$WHEELS_DIR" ] && [ -z "$(ls -A "$WHEELS_DIR")" ]; then
+      build
+    fi
+
+    pip install --no-cache-dir $WHEELS_DIR/trtllm/*.whl || true
+  fi
+}
+
 te() {
   local mode="$1"
 
@@ -174,7 +216,7 @@ nemo() {
   fi
 
   DEPS=(
-    "sox"                                                                                      # requires numpy to be there @URL: https://github.com/marl/pysox/issues/167
+    "sox<=1.5.0"                                                                               # v1.5.0 throws "Can not execute `setup.py` since setuptools is not available in the build environment."; requires numpy to be there @URL: https://github.com/marl/pysox/issues/167
     "llama-index==0.10.43"                                                                     # incompatible with nvidia-pytriton
     "ctc_segmentation==1.7.1 ; (platform_machine == 'x86_64' and platform_system != 'Darwin')" # requires numpy<2.0.0 to be installed before
     "nemo_run"                                                                                 # Not compatible in Python 3.12
diff --git a/requirements/manifest.json b/requirements/manifest.json
index 1feae1247dc8..2d9c29eb3913 100644
--- a/requirements/manifest.json
+++ b/requirements/manifest.json
@@ -12,6 +12,10 @@
         "megatron-lm": {
             "repo": "https://github.com/NVIDIA/Megatron-LM",
             "ref": "aa6207e2da6bca789375f591c9124aa408d2509a"
+        },
+        "trt-llm": {
+            "repo": "https://github.com/NVIDIA/TensorRT-LLM.git",
+            "ref": "v0.17.0"
         }
     },
     "pypi-dependencies": {}


Reviews:
user: login: chtruong814
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-19T02:59:36Z
user: login: chtruong814
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-19T03:00:58Z
user: login: ko3n1g
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T14:53:34Z
user: login: thomasdhc
body: LGTM Left a comment about tensorrt install script in trtllm
state: DISMISSED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T17:13:26Z
user: login: chtruong814
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-19T18:31:30Z
user: login: chtruong814
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-19T18:33:03Z
user: login: pablo-garay
state: DISMISSED
author_association: CONTRIBUTOR
submitted_at: 2025-03-23T23:19:37Z
user: login: thomasdhc
body: Approved, can address changes in follow up
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-26T20:13:17Z

Review comments:
diff_hunk: @@ -12,6 +12,10 @@
         "megatron-lm": {
             "repo": "https://github.com/NVIDIA/Megatron-LM",
             "ref": "aa6207e2da6bca789375f591c9124aa408d2509a"
+        },
+        "trt-llm": {
+            "repo": "https://github.com/NVIDIA/TensorRT-LLM.git",
+            "ref": "v0.13.0"
path: requirements/manifest.json
user: login: chtruong814
body: Aren't we on 0.17.0 now?
created_at: 2025-03-19T02:59:36Z
author_association: COLLABORATOR
diff_hunk: @@ -17,10 +17,47 @@ export INSTALL_DIR=${INSTALL_DIR:-"/opt"}
 export WHEELS_DIR=${WHEELS_DIR:-"$INSTALL_DIR/wheels"}
 export PIP=pip
 
+trtllm() {
+  local mode="$1"
+
+  curl -s https://packagecloud.io/install/repositories/github/git-lfs/script.deb.sh | bash &&
+    apt-get install git-lfs &&
+    git lfs install &&
+    apt-get clean
+
+  TRTLLM_REPO=${TRTLLM_REPO:-$(cat "$CURR/requirements/manifest.json" | jq -r '."vcs-dependencies"."trt_llm".repo')}
+  TRTLLM_TAG=${TRTLLM_TAG:-$(cat "$CURR/requirements/manifest.json" | jq -r '."vcs-dependencies"."trt_llm".ref')}
+  TRTLLM_DIR="$INSTALL_DIR/TensorRT-LLM"
+  if [ ! -d "$TRTLLM_DIR/.git" ]; then
+    rm -rf "$TRTLLM_DIR" &&
+      cd $(dirname "$TRTLLM_DIR") &&
+      git clone ${TRTLLM_REPO}
+  fi &&
+    pushd $TRTLLM_DIR &&
+    git checkout -f $TRTLLM_TAG &&
+    git lfs pull &&
+    popd
+
+  if [[ "$mode" == "build" ]]; then
+    if [[ -n "${NVIDIA_PYTORCH_VERSION}" ]]; then
+      cd $TRTLLM_DIR &&
+        . docker/common/install_tensorrt.sh &&
+        python3 ./scripts/build_wheel.py --job_count $(nproc) --trt_root /usr/local/tensorrt --python_bindings --benchmarks &&
+        pip wheel --wheel-dir $WHEELS_DIR/trtllm/ $TRTLLM_DIR
+    fi
+  else
+    if [ -d "$WHEELS_DIR" ] && [ -z "$(ls -A "$WHEELS_DIR")" ]; then
+      build
+    fi
+
+    pip install --no-cache-dir $WHEELS_DIR/trtllm/*.whl || true
+  fi
+}
+
 te() {
   local mode="$1"
 
-  TE_REPO=${TE_REPO:-$(cat "$CURR/requirements/manifest.json" | jq -r '."vcs-dependencies"."transformer_engine".repo')}
+  $(TE_)REPO=${TE_REPO:-$(cat "$CURR/requirements/manifest.json" | jq -r '."vcs-dependencies"."transformer_engine".repo')}
path: reinstall.sh
user: login: chtruong814
body: typo on `$(TE_)`?
created_at: 2025-03-19T03:00:58Z
author_association: COLLABORATOR
diff_hunk: @@ -12,6 +12,10 @@
         "megatron-lm": {
             "repo": "https://github.com/NVIDIA/Megatron-LM",
             "ref": "aa6207e2da6bca789375f591c9124aa408d2509a"
+        },
+        "trt-llm": {
+            "repo": "https://github.com/NVIDIA/TensorRT-LLM.git",
+            "ref": "v0.13.0"
path: requirements/manifest.json
user: login: ko3n1g
body: Yea, I had some issues initially with 0.17 yesterday and wanted to try a version where I knew i'd work. But now I also managed to get 0.17 up and running (presumably). 
created_at: 2025-03-19T14:53:34Z
author_association: CONTRIBUTOR
diff_hunk: @@ -45,6 +45,15 @@ apt-get install -y bc libsox-fmt-all -y
 apt-get clean
 EOF
 
+ARG TRTLLM_REPO
+ARG TRTLLM_TAG
+RUN --mount=type=bind,from=nemo-bump,source=/opt/NeMo/reinstall.sh,target=/tmp/NeMo/reinstall.sh <<"EOF" bash -ex
+  
+  bash /tmp/NeMo/reinstall.sh --library trtllm --mode build
+  ls -al /tmp/Megatron-LM || true
path: Dockerfile.ci
user: login: chtruong814
body: ah missed this the first time.  Do we want to change what the check is here after trtllm is bult?  Or do we not need this at all?
created_at: 2025-03-19T18:31:29Z
author_association: COLLABORATOR
diff_hunk: @@ -45,6 +45,15 @@ apt-get install -y bc libsox-fmt-all -y
 apt-get clean
 EOF
 
+ARG TRTLLM_REPO
+ARG TRTLLM_TAG
+RUN --mount=type=bind,from=nemo-bump,source=/opt/NeMo/reinstall.sh,target=/tmp/NeMo/reinstall.sh <<"EOF" bash -ex
+  
+  bash /tmp/NeMo/reinstall.sh --library trtllm --mode build
+  ls -al /tmp/Megatron-LM || true
path: Dockerfile.ci
user: login: chtruong814
body: btw, we'll have to see how well the cache works here.  But building trt-llm takes a while so hopefully it's mostly going to be from cache after the initial builds.  
created_at: 2025-03-19T18:33:02Z
author_association: COLLABORATOR

PR comments:
user: login: janekl
created_at: 2025-03-19T10:48:08Z
body: This is great
author_association: COLLABORATOR
user: login: github-actions[bot]
created_at: 2025-03-22T04:55:40Z
body: [🤖]: Hi @ko3n1g 👋,

We wanted to let you know that a [CICD pipeline](https://github.com/NVIDIA/NeMo/actions/runs/13976452247) for this PR just finished successfully

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

I'm just a bot so I'll leave it you what to do next.

//cc @pablo-garay @ko3n1g
author_association: CONTRIBUTOR
user: login: chtruong814
created_at: 2025-03-22T05:58:27Z
body: @ko3n1g tests originally passed but there were some merge conflicts unfortunately.  I resolved the conflicts and restarted the tests.  
author_association: COLLABORATOR
user: login: github-actions[bot]
created_at: 2025-03-22T19:41:03Z
body: [🤖]: Hi @ko3n1g 👋,

We wanted to let you know that a [CICD pipeline](https://github.com/NVIDIA/NeMo/actions/runs/14005807580) for this PR just finished successfully

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

I'm just a bot so I'll leave it you what to do next.

//cc @pablo-garay @ko3n1g
author_association: CONTRIBUTOR
user: login: github-actions[bot]
created_at: 2025-03-26T19:54:34Z
body: [🤖]: Hi @ko3n1g 👋,

We wanted to let you know that a [CICD pipeline](https://github.com/NVIDIA/NeMo/actions/runs/14086560852) for this PR just finished successfully

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

I'm just a bot so I'll leave it you what to do next.

//cc @pablo-garay @ko3n1g
author_association: CONTRIBUTOR
user: login: thomasdhc
created_at: 2025-03-26T20:12:30Z
body: @ko3n1g the pipeline succeeded but I see this line in the container build: https://github.com/NVIDIA/NeMo/actions/runs/14086560852/job/39454191754?pr=12672#step:6:690

can we double check installation on the container? Are there any tests that are in place right now for trtllm?
author_association: CONTRIBUTOR
user: login: ko3n1g
created_at: 2025-03-26T21:02:55Z
body: > @ko3n1g the pipeline succeeded but I see this line in the container build: https://github.com/NVIDIA/NeMo/actions/runs/14086560852/job/39454191754?pr=12672#step:6:690
> 
> can we double check installation on the container? Are there any tests that are in place right now for trtllm?

I think these tests will be onboarded next. Let's follow up there
author_association: CONTRIBUTOR