😎

VertexAIPipelinesのコンポーネントでGPU版XGBoostを使う

に公開
GPU
xgboost
Vertex AI Pipelines
tech

結論

コンポーネントの処理実装をContainerizeする際に使用するDockerfileのベースイメージを
こちらを参考に設定すれば良い。

Dockerfileの例
FROM rapidsai/rapidsai-core:22.06-cuda11.5-base-ubuntu20.04-py3.8

WORKDIR /src

RUN pip install poetry==1.2.0
COPY pyproject.toml poetry.lock ./
RUN poetry config virtualenvs.create false \
  && poetry install --no-dev

COPY . .

（おまけ）OOMエラーに遭遇したら

こちらに書いてあるようにxgboost.QuantileDMatrixを使えばよい。
検証セットorテストセットにQuantileDMatrixを使用する場合は訓練セットのQuantileDMatrixorDMatrixへの参照を渡さないとマズい点に注意（訓練セットと同様の量子化を行うために必要）。

Do not use QuantileDMatrix as validation/test dataset without supplying a reference (the training dataset) QuantileDMatrix using ref as some information may be lost in quantisation.
（ref: https://xgboost.readthedocs.io/en/latest/python/python_api.html#xgboost.QuantileDMatrix）

参考資料

Discussion