Rayにおけるアクセラレータ(GPU)のサポート

Accelerator Support — Ray 2.35.0の翻訳です。

本書は著者が手動で翻訳したものであり内容の正確性を保証するものではありません。正確な内容に関しては原文を参照ください。

アクセラレータ(GPU)は多くの機械学習アプリケーションで重要なものです。Ray Coreはネイティブで事前定義済みのリソースタイプとして多くのアクセラレータをサポートしており、タスクやアクターでそれらのアクセラレータのリソース要件を指定することができます。

Ray Coreがネイティブでサポートしているアクセラレータは以下の通りです:

アクセラレータ	Rayのリソース名	サポートレベル
Nvidia GPU	GPU	完全にテストされておりRayチームがサポート
AMD GPU	GPU	実験段階でありコミュニティがサポート
Intel GPU	GPU	実験段階でありコミュニティがサポート
AWS Neuron Core	neuron_cores	実験段階でありコミュニティがサポート
Google TPU	TPU	実験段階でありコミュニティがサポート
Intel Gaudi	HPU	実験段階でありコミュニティがサポート
Huawei Ascend	NPU	実験段階でありコミュニティがサポート

アクセラレータを持つRayノードの起動

デフォルトでは、Rayは、ノードにおけるアクセラレータリソースの数量を、Rayによって自動検知された物理的なアクセラレータの数量に設定します。必要な場合には、これを上書きすることができます。

注意
Rayのリソースは論理的なので、マシンに搭載されているアクセラレータの真の数より大きい数量のアクセラレータのリソース(num_gpusなど)を指定することを妨げたりはしません。この場合、Rayはアクセラレータを必要とするタスクやアクターのスケジューリングの際に、あなたが指定した個数のアクセラレータがあるかのように動作します。タスクやアクターが実際には存在しないアクセラレータを使用しようとした際にのみトラブルが発生します。

Nvidia GPU

Rayで利用できるNvidia GPUの数を制限するために、Rayノードを起動する前に環境変数CUDA_VISIBLE_DEVICESを設定することができます。例えば、CUDA_VISIBLE_DEVICES=1,3 ray start --head --num-gpus=2によって、Rayはデバイス1と3のみを参照できます。

AMD GPU

Rayで利用できるAMD GPUの数を制限するために、Rayノードを起動する前に環境変数ROCR_VISIBLE_DEVICESを設定することができます。例えば、ROCR_VISIBLE_DEVICES=1,3 ray start --head --num-gpus=2によって、Rayはデバイス1と3のみを参照できます。

Intel GPU

Rayで利用できるIntel GPUの数を制限するために、Rayノードを起動する前に環境変数ONEAPI_DEVICE_SELECTORを設定することができます。例えば、ONEAPI_DEVICE_SELECTOR=1,3 ray start --head --num-gpus=2によって、Rayはデバイス1と3のみを参照できます。

AWS Neuron Core

Rayで利用できるAWS Neuro Coreの数を制限するために、Rayノードを起動する前に環境変数NEURON_RT_VISIBLE_CORESを設定することができます。例えば、NEURON_RT_VISIBLE_CORES=1,3 ray start --head --resources='{"neuron_cores": 2}'によって、Rayはデバイス1と3のみを参照できます。

オーケストレーションの基盤としてEKSとRay on Neuronの他のサンプルに関しては、AWSのドキュメントをご覧ください。

Google TPU

Rayで利用できるGoogle TPUの数を制限するために、Rayノードを起動する前に環境変数TPU_VISIBLE_CHIPSを設定することができます。例えば、TPU_VISIBLE_CHIPS=1,3 ray start --head --resources='{"TPU": 2}'によって、Rayはデバイス1と3のみを参照できます。

Intel Gaudi

Rayで利用できるIntel Gaudi HPUの数を制限するために、Rayノードを起動する前に環境変数HABANA_VISIBLE_MODULESを設定することができます。例えば、HABANA_VISIBLE_MODULES=1,3 ray start --head --resources='{"HPU": 2}'によって、Rayはデバイス1と3のみを参照できます。

Huawei Ascend

Rayで利用できるHuawei Ascend NPUの数を制限するために、Rayノードを起動する前に環境変数ASCEND_RT_VISIBLE_DEVICESを設定することができます。例えば、ASCEND_RT_VISIBLE_DEVICES=1,3 ray start --head --resources='{"NPU": 2}'によって、Rayはデバイス1と3のみを参照できます。

タスクやアクターでのアクセラレータの使用

タスクやアクターでアクセラレータが必要な場合、対応するリソース要件を指定することができます(例: @ray.remote(num_gpus=1))。すると、Rayは十分なアクセラレータリソースを持つノードにタスクやアクターをスケジュールし、タスクやアクターのコードを実行する前に、対応する環境変数(CUDA_VISIBLE_DEVICESなど)を設定することで、タスクやアクターにアクセラレータを割り当てます。

Nvidia GPU

import os
import ray

ray.init(num_gpus=2)

@ray.remote(num_gpus=1)
class GPUActor:
    def ping(self):
        print("GPU IDs: {}".format(ray.get_runtime_context().get_accelerator_ids()["GPU"]))
        print("CUDA_VISIBLE_DEVICES: {}".format(os.environ["CUDA_VISIBLE_DEVICES"]))

@ray.remote(num_gpus=1)
def gpu_task():
    print("GPU IDs: {}".format(ray.get_runtime_context().get_accelerator_ids()["GPU"]))
    print("CUDA_VISIBLE_DEVICES: {}".format(os.environ["CUDA_VISIBLE_DEVICES"]))

gpu_actor = GPUActor.remote()
ray.get(gpu_actor.ping.remote())
# The actor uses the first GPU so the task uses the second one.
ray.get(gpu_task.remote())

(GPUActor pid=52420) GPU IDs: [0]
(GPUActor pid=52420) CUDA_VISIBLE_DEVICES: 0
(gpu_task pid=51830) GPU IDs: [1]
(gpu_task pid=51830) CUDA_VISIBLE_DEVICES: 1

AMD GPU

import os
import ray

ray.init(num_gpus=2)

@ray.remote(num_gpus=1)
class GPUActor:
    def ping(self):
        print("GPU IDs: {}".format(ray.get_runtime_context().get_accelerator_ids()["GPU"]))
        print("ROCR_VISIBLE_DEVICES: {}".format(os.environ["ROCR_VISIBLE_DEVICES"]))

@ray.remote(num_gpus=1)
def gpu_task():
    print("GPU IDs: {}".format(ray.get_runtime_context().get_accelerator_ids()["GPU"]))
    print("ROCR_VISIBLE_DEVICES: {}".format(os.environ["ROCR_VISIBLE_DEVICES"]))

gpu_actor = GPUActor.remote()
ray.get(gpu_actor.ping.remote())
# The actor uses the first GPU so the task uses the second one.
ray.get(gpu_task.remote())

(GPUActor pid=52420) GPU IDs: [0]
(GPUActor pid=52420) ROCR_VISIBLE_DEVICES: 0
(gpu_task pid=51830) GPU IDs: [1]
(gpu_task pid=51830) ROCR_VISIBLE_DEVICES: 1

Intel GPU

import os
import ray

ray.init(num_gpus=2)

@ray.remote(num_gpus=1)
class GPUActor:
    def ping(self):
        print("GPU IDs: {}".format(ray.get_runtime_context().get_accelerator_ids()["GPU"]))
        print("ONEAPI_DEVICE_SELECTOR: {}".format(os.environ["ONEAPI_DEVICE_SELECTOR"]))

@ray.remote(num_gpus=1)
def gpu_task():
    print("GPU IDs: {}".format(ray.get_runtime_context().get_accelerator_ids()["GPU"]))
    print("ONEAPI_DEVICE_SELECTOR: {}".format(os.environ["ONEAPI_DEVICE_SELECTOR"]))

gpu_actor = GPUActor.remote()
ray.get(gpu_actor.ping.remote())
# The actor uses the first GPU so the task uses the second one.
ray.get(gpu_task.remote())

(GPUActor pid=52420) GPU IDs: [0]
(GPUActor pid=52420) ONEAPI_DEVICE_SELECTOR: 0
(gpu_task pid=51830) GPU IDs: [1]
(gpu_task pid=51830) ONEAPI_DEVICE_SELECTOR: 1

AWS Neuron Core

import os
import ray

ray.init(resources={"neuron_cores": 2})

@ray.remote(resources={"neuron_cores": 1})
class NeuronCoreActor:
    def ping(self):
        print("Neuron Core IDs: {}".format(ray.get_runtime_context().get_accelerator_ids()["neuron_cores"]))
        print("NEURON_RT_VISIBLE_CORES: {}".format(os.environ["NEURON_RT_VISIBLE_CORES"]))

@ray.remote(resources={"neuron_cores": 1})
def neuron_core_task():
    print("Neuron Core IDs: {}".format(ray.get_runtime_context().get_accelerator_ids()["neuron_cores"]))
    print("NEURON_RT_VISIBLE_CORES: {}".format(os.environ["NEURON_RT_VISIBLE_CORES"]))

neuron_core_actor = NeuronCoreActor.remote()
ray.get(neuron_core_actor.ping.remote())
# The actor uses the first Neuron Core so the task uses the second one.
ray.get(neuron_core_task.remote())

(NeuronCoreActor pid=52420) Neuron Core IDs: [0]
(NeuronCoreActor pid=52420) NEURON_RT_VISIBLE_CORES: 0
(neuron_core_task pid=51830) Neuron Core IDs: [1]
(neuron_core_task pid=51830) NEURON_RT_VISIBLE_CORES: 1

Google TPU

import os
import ray

ray.init(resources={"TPU": 2})

@ray.remote(resources={"TPU": 1})
class TPUActor:
    def ping(self):
        print("TPU IDs: {}".format(ray.get_runtime_context().get_accelerator_ids()["TPU"]))
        print("TPU_VISIBLE_CHIPS: {}".format(os.environ["TPU_VISIBLE_CHIPS"]))

@ray.remote(resources={"TPU": 1})
def tpu_task():
    print("TPU IDs: {}".format(ray.get_runtime_context().get_accelerator_ids()["TPU"]))
    print("TPU_VISIBLE_CHIPS: {}".format(os.environ["TPU_VISIBLE_CHIPS"]))

tpu_actor = TPUActor.remote()
ray.get(tpu_actor.ping.remote())
# The actor uses the first TPU so the task uses the second one.
ray.get(tpu_task.remote())

(TPUActor pid=52420) TPU IDs: [0]
(TPUActor pid=52420) TPU_VISIBLE_CHIPS: 0
(tpu_task pid=51830) TPU IDs: [1]
(tpu_task pid=51830) TPU_VISIBLE_CHIPS: 1

Intel Gaudi

import os
import ray

ray.init(resources={"HPU": 2})

@ray.remote(resources={"HPU": 1})
class HPUActor:
    def ping(self):
        print("HPU IDs: {}".format(ray.get_runtime_context().get_accelerator_ids()["HPU"]))
        print("HABANA_VISIBLE_MODULES: {}".format(os.environ["HABANA_VISIBLE_MODULES"]))

@ray.remote(resources={"HPU": 1})
def hpu_task():
    print("HPU IDs: {}".format(ray.get_runtime_context().get_accelerator_ids()["HPU"]))
    print("HABANA_VISIBLE_MODULES: {}".format(os.environ["HABANA_VISIBLE_MODULES"]))

hpu_actor = HPUActor.remote()
ray.get(hpu_actor.ping.remote())
# The actor uses the first HPU so the task uses the second one.
ray.get(hpu_task.remote())

(HPUActor pid=52420) HPU IDs: [0]
(HPUActor pid=52420) HABANA_VISIBLE_MODULES: 0
(hpu_task pid=51830) HPU IDs: [1]
(hpu_task pid=51830) HABANA_VISIBLE_MODULES: 1

Huawei Ascend

import os
import ray

ray.init(resources={"NPU": 2})

@ray.remote(resources={"NPU": 1})
class NPUActor:
    def ping(self):
        print("NPU IDs: {}".format(ray.get_runtime_context().get_accelerator_ids()["NPU"]))
        print("ASCEND_RT_VISIBLE_DEVICES: {}".format(os.environ["ASCEND_RT_VISIBLE_DEVICES"]))

@ray.remote(resources={"NPU": 1})
def npu_task():
    print("NPU IDs: {}".format(ray.get_runtime_context().get_accelerator_ids()["NPU"]))
    print("ASCEND_RT_VISIBLE_DEVICES: {}".format(os.environ["ASCEND_RT_VISIBLE_DEVICES"]))

npu_actor = NPUActor.remote()
ray.get(npu_actor.ping.remote())
# The actor uses the first NPU so the task uses the second one.
ray.get(npu_task.remote())

(NPUActor pid=52420) NPU IDs: [0]
(NPUActor pid=52420) ASCEND_RT_VISIBLE_DEVICES: 0
(npu_task pid=51830) NPU IDs: [1]
(npu_task pid=51830) ASCEND_RT_VISIBLE_DEVICES: 1

タスクやアクター内では、ray.get_runtime_context().get_accelerator_ids()によって、タスクやアクターで利用可能なアクセラレータIDのリストが返却されます。通常は、Rayは多くのMLフレームワークがアクセラレータの割り当ての目的のために考慮する、対応する環境変数を自動で設定するので、get_accelerator_ids()を呼び出す必要はありません。

注意: 上で定義したリモート関数やアクターでは、実際にどのアクセラレータも使用しません。Rayは少なくとも1つのアクセラレータを持つノードにそれをスケジュールし、実行される期間一つのアクセラレータを予約しますが、実際にアクセラレータを使うかどうかは関数に依存します。これは、通常TensorFlowのような外部ライブラリを通じて行われます。いかに、実際にアクセラレータを使う例を示します。このサンプルを動作させるには、GPUバージョンのTensorFlowをインストールする必要があります。

@ray.remote(num_gpus=1)
def gpu_task():
    import tensorflow as tf

    # Create a TensorFlow session. TensorFlow restricts itself to use the
    # GPUs specified by the CUDA_VISIBLE_DEVICES environment variable.
    tf.Session()

注意: 割り当てられたアクセラレータを無視したり、マシンのすべてのアクセラレータを使うようにすること自体は可能です。Rayはこれを防がず、同時に同じアクセラレータを使う大量のタスクやアクターが存在する状況に陥ることがあります。しかし、Rayはユーザーに上書きされないものと仮定して、多くのディープラーニングフレームワークによって使われるアクセラレータを制限するための環境変数(CUDA_VISIBLE_DEVICESなど)を自動で設定します。

部分的アクセラレータ

Rayでは、部分的リソース要件をサポートしているので、同じアクセラレータを複数のタスクやアクターで共有することができます。

Nvidia GPU

ray.init(num_cpus=4, num_gpus=1)

@ray.remote(num_gpus=0.25)
def f():
    import time

    time.sleep(1)

# The four tasks created here can execute concurrently
# and share the same GPU.
ray.get([f.remote() for _ in range(4)])

AMD GPU

ray.init(num_cpus=4, num_gpus=1)

@ray.remote(num_gpus=0.25)
def f():
    import time

    time.sleep(1)

# The four tasks created here can execute concurrently
# and share the same GPU.
ray.get([f.remote() for _ in range(4)])

Intel GPU

ray.init(num_cpus=4, num_gpus=1)

@ray.remote(num_gpus=0.25)
def f():
    import time

    time.sleep(1)

# The four tasks created here can execute concurrently
# and share the same GPU.
ray.get([f.remote() for _ in range(4)])

AWS Neuron Core

AWS Neuron Coreでは、部分的リソースをサポートしていません。

Google TPU

Google TPUでは、部分的リソースをサポートしていません。

Intel Gaudi

Intel Gaudiでは、部分的リソースをサポートしていません。

Huawei Ascend

ray.init(num_cpus=4, resources={"NPU": 1})

@ray.remote(resources={"NPU": 0.25})
def f():
    import time

    time.sleep(1)

# The four tasks created here can execute concurrently
# and share the same NPU.
ray.get([f.remote() for _ in range(4)])

注意: 個々のタスクがアクセラレータの取り分以上を使わないようにするかどうかはユーザーの責任です。PytorchやTensorFlowではメモリー使用量を制限するように設定することができます。

Rayが部分的リソース要件を持つタスクやアクターにノードのアクセラレータを割り当てる際、断片化を回避するために次のアクセラレータに移動する前に、一つのアクセラレータをパッキングします。

ray.init(num_gpus=3)

@ray.remote(num_gpus=0.5)
class FractionalGPUActor:
    def ping(self):
        print("GPU id: {}".format(ray.get_runtime_context().get_accelerator_ids()["GPU"]))

fractional_gpu_actors = [FractionalGPUActor.remote() for _ in range(3)]
# Ray tries to pack GPUs if possible.
[ray.get(fractional_gpu_actors[i].ping.remote()) for i in range(3)]

(FractionalGPUActor pid=57417) GPU id: [0]
(FractionalGPUActor pid=57416) GPU id: [0]
(FractionalGPUActor pid=57418) GPU id: [1]

GPUリソースを解放しないワーカー

現時点では、GPUを使用するタスク(TensorFlowを通じてなど)をワーカーが実行する際、タスクはGPUのメモリーを配分しますが、タスクが実行を終えた際に解放しない場合があります。これは、次のタイミングで、あるタスクが同じGPUを使おうとした際に問題につながる場合があります。この問題に対処するために、タスクの処理が終了した際にGPUリソースが解放されるように、RayはデフォルトでGPUタスク間でのワーカープロセスの再利用を無効化します。これは、GPUタスクのスケジュールにオーバーヘッドを引き起こするので、ray.remoteデコレーターで、max_calls=0を設定することで、ワーカーの再利用の再度有効化することができます。

# By default, ray does not reuse workers for GPU tasks to prevent
# GPU resource leakage.
@ray.remote(num_gpus=1)
def leak_gpus():
    import tensorflow as tf

    # This task allocates memory on the GPU and then never release it.
    tf.Session()

アクセラレータのタイプ

Rayでは、リソース固有のアクセラレータのタイプをサポートします。accelerator_typeオプションを用いることで、特定のタイプのアクセラレータを持つノードでのタスクやアクターの実行を強制することができます。内部では、このアクセラレータタイプオプションは、"accelerator_type:<type>": 0.001のカスタムリソース要件として実装されています。これによって、利用できる特定のアクセラレータタイプを持つノードへのタスクやアクターの配置を強制することができます。また、これによってマルチノードタイプのオートスケーラーは、当該タイプのリソースに対する需要を認識することができ、場合によっては当該アクセラレータを提供する新たなノードを起動することができます。

from ray.util.accelerators import NVIDIA_TESLA_V100

@ray.remote(num_gpus=1, accelerator_type=NVIDIA_TESLA_V100)
def train(data):
    return "This function was run on a node with a Tesla V100 GPU"

ray.get(train.remote(1))

利用できるアクセラレータタイプに関しては、ray.util.acceleratorsをご覧ください。

はじめてのDatabricks

はじめてのDatabricks

Databricks無料トライアル

Databricks無料トライアル