Databricks Container ServiceでTensorRT-LLMを動かしてみた

本記事ではDatabricksのDatabricks Container Serviceを用いてNVIDIA社の推論ライブラリであるTensorRT-LLMを実行可能なNotebook環境を構築する方法を紹介します。

はじめに

こんにちは、NTTコミュニケーションズの露崎です。本記事ではDatabricksのDatabricks Container Serviceを用いてDatabricksのデフォルトRuntimeでサポートされていないNVIDIA社の推論ライブラリ、TensorRT-LLMを実行可能なNotebook環境を構築する方法を紹介します。

Databricks Container Service

Databricksはデータウェアハウスとデータレイクの両方の強みを兼ね備えたデータとAIのためのオープンな統合プラットフォーム「データ・インテリジェンス・プラットフォーム」を提供しています。DatabricksではDatabricks Runtimeと呼ばれるプリセットされたSparkの環境が提供されており、通常、ユーザは希望するRuntimeを選択するだけでNotebookの計算環境を構築できます。デフォルトのDatabricks RuntimeにはPythonやRといった実行用の言語環境、及び、機械学習や分析に必要な基本的なライブラリがプリインストールされています。

一方で、プリインストールに含まれていない任意のライブラリについてはNotebookの起動後にaptやpipコマンドを用いて別途インストール必要があります。依存するライブラリが多い場合やインストール中にコンパイルを含むパッケージがある場合、このインストール作業に時間がかかる場合があります。

Databricks Container Service はこうしたライブラリの追加に伴うリソース消費を避け、予め必要な設定を構成した環境を持ち込むための仕組みです。ユーザは事前にビルドしたContainer Imageを任意のレジストリサービス経由でDatabricksの実行環境として利用できます。

NVIDIA TensorRT-LLM

TensorRT-LLMはNVIDIA社の提供する大規模言語モデル(以下LLM)向けのライブラリです。

事前に学習されたLLMを最先端のLLM向け最適化を含んだ形でTensorRTエンジンに変換し、GPUを用いた高速な推論を実現します。TensorRT-LLMはTensorRTエンジンを簡単に生成する為のPython APIをユーザーに提供するだけでなく、TensorRTエンジンを実行するPythonおよびC++のランタイム作成の為のコンポーネントやTriton Inference Serverと統合する為のバックエンドも含まれており、GPUベースのLLM推論のデプロイを容易にします。

TensorRT-LLMを利用するためにはSDKとして利用するTensorRT、及び、TensorRT-LLMのパッケージをインストールする必要があります。

解決したいこと

TensorRT-LLMは強力なライブラリですがDatabricks Runtimeのデフォルトではサポートされていません。また、2024年4月現在Databricks RuntimeがサポートするCUDAのバージョンは11系であり、12系を利用する最新のTensorRT-LLMを利用できません。

本記事では、DatabricksのNotebook上でTensorRT-LLMを使った推論高速化を実現するためにDatabricks Container Serviceを用いてTensorRT-LLM v0.8.0の実行環境を持ち込む方法を紹介します。

TensorRT-LLM Container Imageの作成

Databricks社は、Databricks Container Serviceで実行可能なコンテナのサンプルをGitHubで公開していますので、これをベースにTensorRT-LLMをインストールしたContainer Imageを作成します。

TensorRT-LLMのリポジトリにもDocker Containerを構築するためのDockerfileが公開されていますがベースイメージであるnvcr.io/nvidia/pytorchのサイズが大きく、最終的なイメージサイズがDatabricksのContainerをベースにするよりも大きくなってしまうことから今回はDatabricksのContainerイメージをベースにする方法を紹介します。

Databricks Containers

実際の構築の前にベースとなるDatabricks社が公開しているコンテナのサンプルの構成を紹介します。Root directoryにある ubuntu diretoryで公開されているイメージがDatabricks Container Serviceでサポートされており、機能ごとにDockerfileが分かれています。

通常のRuntimeとしてはubuntu/standardが利用可能であり、Scala、Java、Pythonなど、Databricksで通常利用可能な機能がパッケージ化されています。

GPU向けRuntimeのDockerfileはubuntu/gpuにあり、対応するcudaのバージョン毎にdirectoryが分かれています。各cuda directoryはさらに対応する機能毎に base 、 venv 、 pytorch 、 tensorflow に分かれています。Notebookに必要な最低限の機能は venv で構成されており、これを継承した pytorch 、 tensorflow についてはユーザが予めインストールしておきたいフレームワークに応じて選択できます。

この後解説するTensorRT-LLM向けのイメージについてはこちらに動くものがありますが、実際の変更点については以降で解説していきます。

ベースイメージの変更

現在公開されているDatabricksのイメージはCUDA 11.8のためTensorRT-LLM用にベースイメージを変更します。また、TensorRT-LLMでは cudnn-runtime でなく devel を使いますのでベースイメージを nvidia/cuda:12.1.0-devel-ubuntu22.04 とします。

Pytorch バージョンの変更

ベースイメージの変更でCUDAのバージョンを上げたのでPytorchについてもバージョンを上げます。今回、紹介するTensorRT-LLMのバージョン (v0.8.0) では Pytorchのバージョンを2.2.0a以下の制約があるため、利用できる範囲の最新バージョンである2.1.2を利用します。またベースイメージで設定したCUDA 12.1向けのPytorchを利用するため、Dockerfile内で torch==2.1.2 torchvision==1.6.2 --index-url https://download.pytorch.org/whl/cu121 のように --index-url にPytorchのリポジトリを設定します。 ¹

TensorRT-LLMのインストール

ここまで設定するとCUDA 12.1、PyTorch 2.1.2に対応したDatabricks Container(便宜的に databricksruntime/gpu-pytorch:cuda12.1 とします)がビルドできます。次にこのContainer Imageに以下のDockerfileでTensorRT-LLMのパッケージを追加します。

FROM databricksruntime/gpu-pytorch:cuda12.1 

RUN apt-get update && \
    apt-get -y install openmpi-bin libopenmpi-dev && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*

# install the tenserrt llm version
RUN /databricks/python3/bin/pip install \
    tensorrt_llm==0.8.0 --extra-index-url https://pypi.nvidia.com \
    && /databricks/python3/bin/pip cache purge

これでDatabricks上で実行できるContainerのDockerfileができました。ビルドするとおおよそ以下のサイズのイメージが作成されます。

REPOSITORY                                  TAG                                 IMAGE ID       CREATED         SIZE
databricksruntime/gpu-trt-llm               cuda12.1                            348ca75258f6   3 hours ago     18.6GB

後はこれを持ち込みたい環境からアクセスできるレジストリ (Docker Hub, Google Container Registry, Azure Container Registry, Amazon Elastic Container Registryなど)にpushしましょう。

動作確認

Databricks環境設定

Databricks Container Serviceのドキュメントに従い、クラスタを起動します。今回は以下の設定で動作確認しています。

Databricks Runtimeバージョン: 15.0 (Apache Spark 3.5.0, Scala 2.12)
ノードタイプ: Standard_NC24ads_A100_v4 (Azure)
高度なオプション
- Spark構成
```
  spark.master local[*, 4]
  spark.databricks.cluster.profile singleNode
  spark.databricks.unityCatalog.volumes.enabled true
```
- Docker
  - イメージURL: PrivateのAzure Container Registry
  - 認証：ユーザ名とパスワード
  - ユーザ名、パスワード：(省略)

TensorRT-LLMのインポート

以下のセルを実行しTensorRT-LLM 0.8.0が正しく組み込めていることが確認できました。

Llama2 HF-7b-instruct モデルの変換

公式のdocs を参考にhuggingfaceで提供されているMeta社のLlama2を変換します。checkpointのconvert及び、エンジン化の変換手順は公式のdocs に記載があるのでここでは省略します。

手順によって生成された config.json と rank0.engine (GPUの枚数や環境次第でengineファイルの名前は変更されます)を次のTensorRT-LLMの呼び出しで利用します。

TensorRT-LLMの呼び出し

以下は、TensorRT-LLMで生成したエンジンを呼び出して推論するためのサンプルです。

環境変数である BASE_DIR には変換したエンジンが保存されているディレクトリを設定します。Databricksの環境の場合、通常のワークスペースファイルには200MBの制限(Azureの場合)がありますのでLLMを扱う場合にはUnity Catalog上のボリュームを利用するのが良いでしょう。

from pathlib import Path
import os

import torch
from transformers import LlamaTokenizer

import tensorrt_llm
from tensorrt_llm.runtime import ModelRunner

EOS_TOKEN = 2
PAD_TOKEN = 2

def generate(input_text):
    tensorrt_llm.logger.set_level("info")
    tokenizer_dir = "meta-llama/Llama-2-7b-chat-hf"
    engine_dir = Path(os.environ['BASE_DIR'])
    runtime_rank = tensorrt_llm.mpi_rank()
    torch.cuda.set_device(0)
    tokenizer = LlamaTokenizer.from_pretrained(tokenizer_dir, legacy=False)

    runner = ModelRunner.from_dir(engine_dir=engine_dir, rank=runtime_rank)

    input_tokens = [tokenizer.encode(input_text, add_special_tokens=False)]
    input_ids = torch.nested.to_padded_tensor(
        torch.nested.nested_tensor(input_tokens, dtype=torch.int32),
        EOS_TOKEN).cuda()

    output_gen_ids = runner.generate(
        input_ids,
        pad_id=PAD_TOKEN,
        end_id=EOS_TOKEN,
        max_new_tokens=100,
        streaming=False)
    torch.cuda.synchronize()
    output_text = tokenizer.decode(output_gen_ids[0][0])
    return output_text

text = generate("hello, could you tell a story for the galaxy?")
print(text)