chore: Bump pytorch to 26.06 and vllm to 0.24.0 (#718)

Signed-off-by: Charlie Truong <chtruong@nvidia.com>

Author: chtruong814

docker/Dockerfile.pytorch

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-FROM nvcr.io/nvidian/nemo:export-deploy-pyt-26.04-vllm-v0.20.1
+FROM nvcr.io/nvidian/nemo:export-deploy-pyt-26.06-vllm-v0.24.0
 WORKDIR /workspace
 COPY . .
 ARG INFERENCE_FRAMEWORK

pyproject.toml

@@ -80,7 +80,7 @@ dependencies = [
 
 [project.optional-dependencies]
 inframework = []
-vllm = ["vllm~=0.20.1", "pandas", "timm"]
+vllm = ["vllm~=0.24.0", "pandas", "timm"]
 trt-onnx = [
     "tensorrt==10.16.1.11",
     "onnx==1.21.0",
@@ -147,23 +147,25 @@ override-dependencies = [
     "urllib3>1.27.0",
     "tiktoken>=0.9.0",                                                                   # because nemo-toolkit and megatron-bridge disagree on tiktoken, we need to pin it here,
     "fsspec[http]>=2023.1.0,<=2024.9.0",
-    "transformer-engine[pytorch]>=2.14.0a0,<2.17.0",
+    "transformer-engine[pytorch]>=2.16.0a0,<2.19.0",
     "open-clip-torch>=3.2.0",
     "megatron-energon[av-decode]>=6.0,<7.dev0",
     "flashinfer-python==0.6.8.post1",
     "datasets>=3.3.0",
     "flash-linear-attention>=0.3.0,<0.4.dev0",
     "patchelf; sys_platform=='never'",
     "nvidia-resiliency-ext>=0.6.0",
-    "transformer-engine[pytorch,core_cu13]>=2.14.0a0,<2.17.0; sys_platform != 'darwin'",
-    "transformer-engine-cu13>=2.14.0a0,<2.17.0; sys_platform != 'darwin'",
+    "transformer-engine[pytorch,core_cu13]>=2.16.0a0,<2.19.0; sys_platform != 'darwin'",
+    "transformer-engine-cu13>=2.16.0a0,<2.19.0; sys_platform != 'darwin'",
     "transformer-engine-cu12; sys_platform == 'never'",
     "mamba-ssm>=2.3.0,<2.4.0",
-    "transformers>=5.0.0",
+    "transformers==5.8.1",
     "protobuf~=6.33.5",
     "opencv-python-headless; sys_platform == 'never'",
     "cryptography>=43.0.0,<47",
     "onnx<1.20.0",
+    "flashinfer-python==0.6.8.post1",
+    "flashinfer-cubin==0.6.8.post1",
 ]
 prerelease = "allow"
 

tests/functional_tests/utils/run_onnx_trt_reranking_export.py

@@ -52,7 +52,7 @@ def get_args():
     )
     parser.add_argument(
         "--trt_version_compatible",
-        default=True,
+        default=False,
         action="store_true",
         help="Whether to generate version compatible TensorRT models.",
     )