Add Jetson platform support (by docker)

docker/jetson/Dockerfile

@@ -0,0 +1,52 @@
+ARG JETPACK_VERSION=36.2.0  # Jetpack 6.0 DP
+
+FROM "nvcr.io/nvidia/l4t-jetpack:r$JETPACK_VERSION"
+
+ARG JETPACK_VERSION=36.2.0
+ARG HTTP_PROXY=""
+ARG HTTPS_PROXY=""
+
+ENV JETPACK_VERSION=$JETPACK_VERSION
+
+RUN apt-get update && apt-get install -y --no-install-recommends rapidjson-dev libgoogle-glog-dev gdb python3-pip git libopenblas-dev libopenmpi-dev && \
+    pip3 config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple && \
+    python3 -m pip install --no-cache-dir --upgrade pip setuptools==69.5.1 &&\
+    python3 -m pip install --no-cache-dir cmake packaging wheel && \
+    mkdir -p /workspace
+
+COPY . /workspace/lmdeploy
+
+WORKDIR /workspace/lmdeploy
+
+RUN mkdir -p /workspace/lmdeploy/torch-jetson && \
+    chmod +x /workspace/lmdeploy/docker/jetson/install_pytorch.sh && \
+    /workspace/lmdeploy/docker/jetson/install_pytorch.sh $JETPACK_VERSION
+
+RUN cd /workspace/lmdeploy &&\
+    python3 -m pip install --no-cache-dir -r requirements/build.txt
+
+# proxy
+ENV http_proxy=$HTTP_PROXY
+ENV https_proxy=$HTTPS_PROXY
+
+RUN mkdir -p build && cd build &&\
+    cmake .. \
+        -DCMAKE_BUILD_TYPE=RelWithDebInfo \
+        -DCMAKE_EXPORT_COMPILE_COMMANDS=1 \
+        -DCMAKE_INSTALL_PREFIX=/workspace/lmdeploy/install \
+        -DBUILD_PY_FFI=ON \
+        -DBUILD_MULTI_GPU=OFF \
+        -DBUILD_CUTLASS_MOE=OFF \
+        -DBUILD_CUTLASS_MIXED_GEMM=OFF \
+        -DCMAKE_CUDA_FLAGS="-lineinfo" \
+        -DUSE_NVTX=ON &&\
+    make -j$(nproc) && make install &&\
+    cd .. &&\
+    python3 -m pip install -e . &&\
+    rm -rf build
+
+ENV LD_LIBRARY_PATH=/workspace/lmdeploy/install/lib:$LD_LIBRARY_PATH
+# # explicitly set ptxas path for triton
+# ENV TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas
+ENV http_proxy=""
+ENV https_proxy=""

docker/jetson/README.md

@@ -0,0 +1,22 @@
+# Jetson Support (beta)
+
+Build LMDeploy for NVIDIA Jetson platforms by docker.
+
+```sh
+export JETPACK_L4T="36.2.0" # Jetpack 6.0 DP
+
+cd lmdeploy
+docker build -t lmdeploy_jetson:r$JETPACK_L4T \
+    --build-arg JETPACK_VERSION=$JETPACK_L4T \
+    -f docker/jetson/Dockerfile .
+```
+
+Version Corresponding List:
+
+| $JETPACK_L4T | Jetpack Version | Python Version | Torch Version | CUDA VERSION |                                           Support Boards                                           |
+| :----------: | :-------------: | :------------: | :-----------: | :----------: | :------------------------------------------------------------------------------------------------: |
+|    35.2.1    |       5.1       |   Python 3.8   |     2.0.0     |     11.4     |            AGX Orin NX 32GB,<br>Orin NX 16GB,<br>Xavier NX series,<br>AGX Xavier Series            |
+|    35.3.1    |      5.1.1      |   Python 3.8   |     2.0.0     |     11.4     | AGX Orin Series,<br>Orin NX Series,<br>Orin Nano Series,<br>Xavier NX Series,<br>AGX Xavier Series |
+|    35.4.1    |      5.1.2      |   Python 3.8   |     2.1.0     |     11.4     | AGX Orin Series,<br>Orin NX Series,<br>Orin Nano Series,<br>Xavier NX Series,<br>AGX Xavier Series |
+|    36.2.0    |     6.0 DP      |  Python 3.10   |     2.2.0     |     12.2     |                      AGX Orin Series,<br>Orin NX Series,<br>Orin Nano Series                       |
+|    36.3.0    |       6.0       |  Python 3.10   |     2.4.0     |     12.2     |                      AGX Orin Series,<br>Orin NX Series,<br>Orin Nano Series                       |

docker/jetson/install_pytorch.sh

@@ -0,0 +1,25 @@
+#!/bin/bash
+
+JETPACK_VERSION=$1
+
+echo $JETPACK_VERSION
+
+if [ "$JETPACK_VERSION" = "35.2.1" ] ; then     # Jetpack 5.1
+    wget "https://developer.download.nvidia.cn/compute/redist/jp/v51/pytorch/torch-2.0.0a0+8aa34602.nv23.03-cp38-cp38-linux_aarch64.whl" -P /workspace/lmdeploy/torch-jetson
+    python3 -m pip install "/workspace/lmdeploy/torch-jetson/torch-2.0.0a0+8aa34602.nv23.03-cp38-cp38-linux_aarch64.whl"
+elif [ "$JETPACK_VERSION" = "35.3.1" ] ; then   # Jetpack 5.1.1
+    wget "https://developer.download.nvidia.cn/compute/redist/jp/v511/pytorch/torch-2.0.0a0+fe05266f.nv23.04-cp38-cp38-linux_aarch64.whl" -P /workspace/lmdeploy/torch-jetson
+    python3 -m pip install "/workspace/lmdeploy/torch-jetson/torch-2.0.0a0+fe05266f.nv23.04-cp38-cp38-linux_aarch64.whl"
+elif [ "$JETPACK_VERSION" = "35.4.1" ] ; then   # Jetpack 5.1.2
+    wget "https://developer.download.nvidia.cn/compute/redist/jp/v512/pytorch/torch-2.1.0a0+41361538.nv23.06-cp38-cp38-linux_aarch64.whl" -P /workspace/lmdeploy/torch-jetson
+    python3 -m pip install "/workspace/lmdeploy/torch-jetson/torch-2.1.0a0+41361538.nv23.06-cp38-cp38-linux_aarch64.whl"
+elif [ "$JETPACK_VERSION" = "36.2.0" ] ; then   # Jetpack 6.0 DP
+    wget "https://developer.download.nvidia.cn/compute/redist/jp/v60dp/pytorch/torch-2.2.0a0+81ea7a4.nv24.01-cp310-cp310-linux_aarch64.whl" -P /workspace/lmdeploy/torch-jetson
+    python3 -m pip install "/workspace/lmdeploy/torch-jetson/torch-2.2.0a0+81ea7a4.nv24.01-cp310-cp310-linux_aarch64.whl"
+elif [ "$JETPACK_VERSION" = "36.3.0" ] ; then   # Jetpack 6.0
+    wget https://developer.download.nvidia.cn/compute/redist/jp/v60/pytorch/torch-2.4.0a0+07cecf4168.nv24.05.14710581-cp310-cp310-linux_aarch64.whl -P /workspace/lmdeploy/torch-jetson
+    python3 -m pip install "/workspace/lmdeploy/torch-jetson/torch-2.4.0a0+07cecf4168.nv24.05.14710581-cp310-cp310-linux_aarch64.whl"
+else
+    echo "We currently do not support Jetpack v$JETPACK_VERSION. Please try 35.2.1, 35.3.1, 35.4.1, 36.2.0, or 36.3.0"
+    exit 1
+fi

requirements/jetson.txt

@@ -0,0 +1,18 @@
+# runtime.txt without torch, torchvision and triton
+accelerate>=0.29.3
+einops
+fastapi
+fire
+mmengine-lite
+numpy<2.0.0
+peft<=0.11.1
+pillow
+protobuf
+pydantic>2.0.0
+pynvml
+safetensors
+sentencepiece
+shortuuid
+tiktoken
+transformers
+uvicorn

setup.py

@@ -128,8 +128,23 @@ def gen_packages_items():
     return packages
 
 
+def check_jetson_platform():
+    """Assuming Linux/arch64 is the NVIDIA Jetson platform.
+
+    If support for other Linux/Arrch64 platforms is introduced in the future,
+    please make appropriate modifications.
+    """
+    import platform
+    current_system = platform.system().lower()
+    current_arch = platform.machine().lower()
+    suspected_jetson = current_system == 'linux' and current_arch == 'aarch64'
+    if suspected_jetson: return True
+    return False
+
+
 if __name__ == '__main__':
     lmdeploy_package_data = ['lmdeploy/bin/llama_gemm']
+    jetson_support = check_jetson_platform()
     setup(
         name='lmdeploy',
         version=get_version(),
@@ -145,7 +160,8 @@ def gen_packages_items():
         include_package_data=True,
         setup_requires=parse_requirements('requirements/build.txt'),
         tests_require=parse_requirements('requirements/test.txt'),
-        install_requires=parse_requirements('requirements/runtime.txt'),
+        install_requires=parse_requirements('requirements/runtime.txt') if
+        not jetson_support else parse_requirements('requirements/jetson.txt'),
         extras_require={
             'all': parse_requirements('requirements.txt'),
             'lite': parse_requirements('requirements/lite.txt'),