翻訳は機械翻訳により提供されています。提供された翻訳内容と英語版の間で齟齬、不一致または矛盾がある場合、英語版が優先します。

PennyLane 埋め込みシミュレーターを使用したハイブリッドワークロードの実行

Amazon Braket Hybrid Jobs で PennyLane の埋め込みシミュレーターを使用してハイブリッドワークロードを実行する方法について見ていきましょう。Pennylane の GPU ベースの埋め込みシミュレーター lightning.gpu は、Nvidia cuQuantum ライブラリを使用して回路のシミュレーションを高速化します。埋め込み GPU シミュレーターは、ユーザーがすぐに使用できるすべての Braket ジョブコンテナに事前設定されています。このページでは、lightning.gpu を使用してハイブリッドワークロードを高速化する方法を示します。

QAOA ワークロードに lightning.gpu を使用する

こちらのノートブックにある量子近似最適化アルゴリズム (QAOA) の例について考えてみましょう。埋め込みシミュレーターを選択するには、引数 device を "local:<provider>/<simulator_name>" という形式の文字列として指定します。例えば、lightning.gpu の場合なら "local:pennylane/lightning.gpu" と設定します。ハイブリッドジョブの起動時にそれに渡すデバイス文字列は、環境変数 "AMZN_BRAKET_DEVICE_ARN" として渡されます。

device_string = os.environ["AMZN_BRAKET_DEVICE_ARN"]
prefix, device_name = device_string.split("/")
device = qml.device(simulator_name, wires=n_wires)

このページでは、2 つの埋め込み型 PennyLane 状態ベクトルシミュレーター lightning.qubit (CPU ベース) と lightning.gpu (GPU ベース) を比較します。各シミュレーターにカスタムゲート分解を提供して、さまざまな勾配を計算させます。

これで、ハイブリッドジョブ起動スクリプトを作成する準備ができました。ml.m5.2xlarge と ml.g4dn.xlarge の 2 つのインスタンスタイプを使用して QAOA アルゴリズムを実行します。ml.m5.2xlarge インスタンスタイプは、開発者用の標準的なラップトップに相当します。ml.g4dn.xlarge は、16GB のメモリを持つ単一の NVIDIA T4 GPU を持つ高速コンピューティングインスタンスです。

GPU を実行するには、まず互換性のあるイメージと正しいインスタンス (デフォルトでml.m5.2xlargeはインスタンス) を指定する必要があります。

from braket.aws import AwsSession
from braket.jobs.image_uris import Framework, retrieve_image

image_uri = retrieve_image(Framework.PL_PYTORCH, AwsSession().region)
instance_config = InstanceConfig(instanceType="ml.g4dn.xlarge")

次に、これらをハイブリッドジョブデコレータに入力し、システム引数とハイブリッドジョブ引数の両方で更新されたデバイスパラメータを入力する必要があります。

@hybrid_job(
        device="local:pennylane/lightning.gpu",
        input_data=input_file_path,
        image_uri=image_uri,
        instance_config=instance_config)
def run_qaoa_hybrid_job_gpu(p=1, steps=10):
    params = np.random.rand(2, p)

    braket_task_tracker = Tracker()

    graph = nx.read_adjlist(input_file_path, nodetype=int)
    wires = list(graph.nodes)
    cost_h, _mixer_h = qaoa.maxcut(graph)

    device_string = os.environ["AMZN_BRAKET_DEVICE_ARN"]
    prefix, device_name = device_string.split("/")
    dev= qml.device(simulator_name, wires=len(wires))
    ...

m5.2xlarge インスタンスの平均反復時間は約 73 秒で、ml.g4dn.xlargeインスタンスの平均反復時間は約 0.6 秒です。この 21 量子ビットのワークフローでは、GPU インスタンスによって 100 倍高速になります。Amazon Braket Hybrid Jobs の料金ページを見ると、m5.2xlargeインスタンスの 1 分あたりのコストは 0.00768 USD で、ml.g4dn.xlargeインスタンスのコストは 0.01227 USD であることがわかります。この例では、GPU インスタンスで実行する方が高速で安価です。

量子機械学習とデータ並列処理

ワークロードタイプがデータセットでトレーニングする量子機械学習 (QML) である場合は、データ並列処理を使用した方がワークロードを高速化できます。QML では、モデルには 1 つ以上の量子回路が含まれています。モデルには古典ニューラルネットが含まれる場合と含まれない場合があります。データセットを使用してモデルをトレーニングすると、損失関数が最小限に抑えられるようにモデルのパラメータが更新されます。損失関数は通常、単一のデータポイントに対して定義されるのに対し、バッチ平均損失はデータセット全体に対する平均損失として定義されます。QML では、各損失が通常、逐次的に計算されてから、バッチ平均損失が勾配計算用に算出されます。このプロシージャは、特に数百のデータポイントがある場合に時間がかかります。

1 つのデータポイントからの損失は他のデータポイントに依存しないため、複数の損失を並列処理で評価できます。複数のデータポイントに関連する損失と勾配を同時に評価できるのです。これはデータ並列処理と呼ばれます。データ並列処理を使用してトレーニングを高速化しやすくするには、Amazon Braket Hybrid Jobs と SageMaker の分散データ並列ライブラリを使用します。

二項分類の例としてよく知られている UCI リポジトリの ソナーデータセットを使用するデータ並列処理については、以下の QML ワークロードを検討してください。ソナーデータセットには 208 個のデータポイントがあります。各データポイントには、対象物から反射したソナー信号から収集された 60 個の特徴量があります。各データポイントには、機雷の場合は「M」、岩の場合は「R」というラベルが付けられます。当社の QML モデルは、入力レイヤー、非表示レイヤーとしての量子回路、および出力レイヤーで構成されています。入力レイヤーと出力レイヤーは、PyTorch 内に実装されている古典ニューラルネットです。量子回路は、PennyLane の qml.qnn モジュールを使用して PyTorch ニューラルネットと統合されています。ワークロードの詳細については、「example」のノートブックを参照してください。上記の QAOA の例と同様に、PennyLane の lightning.gpu などの埋め込み GPU ベースのシミュレーターを使用して GPU の能力を活用することで、埋め込み CPU ベースのシミュレーターよりもパフォーマンスを向上させることができます。

ハイブリッドジョブを作成するには、アルゴリズムスクリプト、デバイス、およびその他の設定をキーワード引数に指定して AwsQuantumJob.create を呼び出します。

instance_config = InstanceConfig(instanceType='ml.g4dn.xlarge')

hyperparameters={"nwires": "10",
                 "ndata": "32",
                 ...
}

job = AwsQuantumJob.create(
    device="local:pennylane/lightning.gpu",
    source_module="qml_source",
    entry_point="qml_source.train_single",
    hyperparameters=hyperparameters,
    instance_config=instance_config,
    ...
)

データ並列処理を使用するには、トレーニングを正しく並列化するように、SageMaker 分散ライブラリのアルゴリズムスクリプトで数行のコードを変更する必要があります。まず、smdistributed パッケージをインポートします。このパッケージが、ワークロードを複数の GPU と複数のインスタンスに分散するための力仕事のほとんどを行います。このパッケージは Braket PyTorch コンテナと TensorFlow コンテナに事前設定されています。dist モジュールは、トレーニング (world_size) の GPU の合計数と GPU コアの rank と local_rank をアルゴリズムスクリプトに通知します。rank はすべてのインスタンスにおける GPU の絶対インデックスで、local_rank は特定のインスタンスにおける GPU のインデックスです。例えば、それぞれに 8 つの GPU が割り当てられた 4 つのインスタンスがある場合、rank の範囲は 0～31 で、local_rank の範囲は 0～7 です。

import smdistributed.dataparallel.torch.distributed as dist

dp_info = {
    "world_size": dist.get_world_size(),
    "rank": dist.get_rank(),
    "local_rank": dist.get_local_rank(),
}
batch_size //= dp_info["world_size"] // 8
batch_size = max(batch_size, 1)

次に、world_size と rank に従って DistributedSampler を定義し、それをデータローダーに渡します。次の例のサンプラー (sampler) は、データセットの同じスライスにアクセスする GPU を避けます。

train_sampler = torch.utils.data.distributed.DistributedSampler(
    train_dataset,
    num_replicas=dp_info["world_size"],
    rank=dp_info["rank"]
)
train_loader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=batch_size,
    shuffle=False,
    num_workers=0,
    pin_memory=True,
    sampler=train_sampler,
)

次に、DistributedDataParallel クラスを使用してデータ並列処理を有効にします。

from smdistributed.dataparallel.torch.parallel.distributed import DistributedDataParallel as DDP

model = DressedQNN(qc_dev).to(device)
model = DDP(model)
torch.cuda.set_device(dp_info["local_rank"])
model.cuda(dp_info["local_rank"])

上記は、データ並列処理を使用するために必要な変更です。QML では、結果を保存し、トレーニングの進行状況を出力する必要が頻繁に生じます。各 GPU が保存および出力コマンドを実行すると、ログは繰り返しの情報でいっぱいになり、結果は互いに上書きされます。これを回避するには、rank が 0 の GPU においてのみ、保存および出力を行います。

if dp_info["rank"]==0:
    print('elapsed time: ', elapsed)
    torch.save(model.state_dict(), f"{output_dir}/test_local.pt")
    save_job_result({"last loss": loss_before})

Amazon Braket Hybrid Jobs は、SageMaker 分散データ並列ライブラリの ml.g4dn.12xlarge インスタンスタイプをサポートしています。Hybrid Jobs で InstanceConfig 引数を使用してインスタンスタイプを設定します。SageMaker 分散データ並列ライブラリがデータ並列処理が有効になっていることを知るには、2 つのハイパーパラメータを追加する必要があります ("sagemaker_distributed_dataparallel_enabled" を "true" に設定し、"sagemaker_instance_type" を使用予定のインスタンスタイプに設定)。これら 2 つのハイパーパラメータは smdistributed パッケージで使用されます。これらのハイパーパラメータは、アルゴリズムスクリプトで明示的に使用する必要はありません。Amazon Braket SDK のアルゴリズムスクリプトには、便利なキーワード引数 distribution が備わっているためです。ハイブリッドジョブの作成で distribution="data_parallel" を使用することで、これら 2 つのハイパーパラメータが Amazon Braket SDK により自動的に挿入されるのです。Amazon Braket API を使用する場合は、これら 2 つのハイパーパラメータをユーザーが指定する必要があります。

インスタンスおよびデータ並列処理を設定することで、ハイブリッドジョブを送信できるようになりました。ml.g4dn.12xlarge インスタンスには 4 つの GPUsがあります。instanceCount=1 と設定すると、ワークロードはインスタンス内の 8 つの GPU に分散されます。instanceCount を 1 より大きい値に設定すると、ワークロードはすべてのインスタンスに存在する GPU 間で分散されます。複数のインスタンスを使用する場合は、各インスタンスの使用時間に基づいて料金が発生します。例えば、4 つのインスタンスを使用する場合は、ワークロードを同時に実行しているインスタンスが 4 つあるため、請求対象時間はインスタンスあたりの実行時間の 4 倍になります。

instance_config = InstanceConfig(instanceType='ml.g4dn.12xlarge',
                                 instanceCount=1,
)

hyperparameters={"nwires": "10",
                 "ndata": "32",
                 ...,
}

job = AwsQuantumJob.create(
    device="local:pennylane/lightning.gpu",
    source_module="qml_source",
    entry_point="qml_source.train_dp",
    hyperparameters=hyperparameters,
    instance_config=instance_config,
    distribution="data_parallel",
    ...
)

正しく使用すると、複数のインスタンスを使用すると、時間とコストの両方が大幅に削減される可能性があります。 詳細については、ノートブックの例を参照してください。