使用 PennyLane 內嵌模擬器執行混合工作負載 - Amazon Braket

本文為英文版的機器翻譯版本,如內容有任何歧義或不一致之處,概以英文版為準。

使用 PennyLane 內嵌模擬器執行混合工作負載

讓我們來看看如何在 Amazon Braket Hybrid Jobs 上使用來自 PennyLane 的內嵌模擬器來執行混合工作負載。Pennylane 的 GPU 型內嵌模擬器 lightning.gpu使用 Nvidia cuQuantum 程式庫來加速電路模擬。嵌入式 GPU 模擬器已預先設定在所有 Braket 任務容器中,使用者可以立即使用。在此頁面中,我們會示範如何使用 lightning.gpu 來加速混合式工作負載。

使用 lightning.gpu處理 QAOA 工作負載

請考量此筆記本中的 Quantum 近似最佳化演算法 (QAOA) 範例。若要選取內嵌模擬器,請將device引數指定為格式的字串:"local:<provider>/<simulator_name>"。例如,您會"local:pennylane/lightning.gpu"為 設定 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)

在此頁面上,讓我們比較兩個內嵌的 PennyLane 狀態向量模擬器 lightning.qubit(以 CPU 為基礎) 和 lightning.gpu(以 GPU 為基礎)。您需要為模擬器提供一些自訂閘道分解,才能計算各種漸層。

現在您已準備好準備混合任務啟動指令碼。您將使用兩種執行個體類型執行 QAOA 演算法: m5.2xlargep3.2xlargem5.2xlarge 執行個體類型與標準開發人員筆記型電腦相當。p3.2xlarge 是具有單一 NVIDIA Volta GPU 和 16GB 記憶體的加速運算執行個體。

所有混合任務hyperparameters的 將相同。您只需依照下列方式變更兩行,即可嘗試不同的執行個體和模擬器。

# Specify device that the hybrid job will primarily be targeting device = "local:pennylane/lightning.qubit" # Run on a CPU based instance with about as much power as a laptop instance_config = InstanceConfig(instanceType='ml.m5.2xlarge')

或:

# Specify device that the hybrid job will primarily be targeting device = "local:pennylane/lightning.gpu" # Run on an inexpensive GPU based instance instance_config = InstanceConfig(instanceType='ml.p3.2xlarge')
注意

如果您將 指定instance_config為使用 GPU 型執行個體,但選擇 device做為內嵌的 CPU 型模擬器 (lightning.qubit),則不會使用 GPU。如果您想要鎖定 GPU,請務必使用內嵌的 GPU 型模擬器!

首先,您可以建立兩個混合任務,並在具有 18 個頂點的圖形上使用 QAOA 解決 Max-Cut。這會轉換為 18 qubit 電路,相對較小且可行,可在您的筆記型電腦或m5.2xlarge執行個體上快速執行。

num_nodes = 18 num_edges = 24 seed = 1967 graph = nx.gnm_random_graph(num_nodes, num_edges, seed=seed) # And similarly for the p3 job m5_job = AwsQuantumJob.create( device=device, source_module="qaoa_source", job_name="qaoa-m5-" + str(int(time.time())), image_uri=image_uri, # Relative to the source_module entry_point="qaoa_source.qaoa_algorithm_script", copy_checkpoints_from_job=None, instance_config=instance_config, # general parameters hyperparameters=hyperparameters, input_data={"input-graph": input_file_path}, wait_until_complete=True, )

m5.2xlarge 執行個體的平均反覆運算時間約為 25 秒,而p3.2xlarge執行個體約為 12 秒。對於此 18 qubit 工作流程,GPU 執行個體會給我們 2 倍的速度。如果您查看 Amazon Braket Hybrid Jobs 定價頁面,您可以看到m5.2xlarge執行個體的每分鐘成本為 0.00768 美元,而p3.2xlarge執行個體則為 0.06375 美元。若要執行總共 5 次的反覆運算,使用 CPU 執行個體將花費 0.016 美元,或使用 GPU 執行個體將花費 0.06375 美元,兩者都相當便宜!

現在,讓我們更難解決問題,並嘗試解決 24 頂點圖形上的 Max-Cut 問題,這會轉換為 24 個 qubit。在相同的兩個執行個體上再次執行混合式任務,並比較成本。

注意

您會看到在 CPU 執行個體上執行此混合式任務的時間可能大約為 5 小時!

num_nodes = 24 num_edges = 36 seed = 1967 graph = nx.gnm_random_graph(num_nodes, num_edges, seed=seed) # And similarly for the p3 job m5_big_job = AwsQuantumJob.create( device=device, source_module="qaoa_source", job_name="qaoa-m5-big-" + str(int(time.time())), image_uri=image_uri, # Relative to the source_module entry_point="qaoa_source.qaoa_algorithm_script", copy_checkpoints_from_job=None, instance_config=instance_config, # general parameters hyperparameters=hyperparameters, input_data={"input-graph": input_file_path}, wait_until_complete=True, )

m5.2xlarge 執行個體的平均反覆運算時間約為一小時,而p3.2xlarge執行個體則約為兩分鐘。對於這個較大的問題,GPU 執行個體的順序會更快!您只需要變更兩行程式碼,將執行個體類型和使用的本機模擬器替換出來,即可從此加速中獲益。若要執行總共 5 次反覆運算,如此處所述,使用 CPU 執行個體約需 2.27072 美元,或使用 GPU 執行個體約需 0.775625 美元。CPU 用量不僅更昂貴,而且需要更多時間才能執行。使用由 NVIDIA CuQuantum 支援的 PennyLane 內嵌模擬器 AWS,透過可用的 GPU 執行個體加速此工作流程,可讓您以更低的總成本和更短的時間,執行具有中繼 qubit 計數 (介於 20 到 30 之間) 的工作流程。這表示即使問題太大而無法在筆記型電腦或類似大小的執行個體上快速執行,您也可以使用量子運算進行實驗。

Quantum 機器學習和資料平行處理

如果您的工作負載類型是在資料集上訓練的量子機器學習 (QML),您可以使用資料平行處理進一步加速工作負載。在 QML 中,模型包含一或多個量子電路。模型可能也可能不包含傳統神經網路。使用資料集訓練模型時,模型中的參數會更新,以將損失函數降至最低。損失函數通常針對單一資料點定義,以及整個資料集的平均損失總損失。在 QML 中,損失通常會在平均至漸層運算的總損失之前,以序列方式計算。此程序很耗時,特別是有數百個資料點時。

由於某個資料點的損失不依賴其他資料點,因此可以平行評估損失!您可以同時評估與不同資料點相關聯的損失和梯度。這稱為資料平行處理。透過 SageMaker 的分散式資料平行程式庫,Amazon Braket Hybrid Jobs 可讓您更輕鬆地利用資料平行處理來加速訓練。

請考慮下列資料平行處理的 QML 工作負載,其使用來自已知 UCI 儲存庫的聲納資料集作為二進位分類的範例。聲納資料集各有 208 個資料點,其中包含從聲納訊號發出的 60 個功能。每個資料點都標記為「M」表示礦區,或「R」表示石頭。我們的 QML 模型包含輸入層、做為隱藏層的量子電路,以及輸出層。輸入和輸出層是在 PyTorch 中實作的傳統神經網路。量子電路使用 PennyLane 的 qml.qnn 模組與 PyTorch 神經網路整合。如需工作負載的詳細資訊,請參閱我們的範例筆記本。如同上述的 QAOA 範例,您可以使用 PennyLane 等嵌入式 GPU 模擬器來利用 GPU 的強大功能lightning.gpu,以改善效能。

若要建立混合任務,您可以透過其關鍵字引數呼叫AwsQuantumJob.create並指定演算法指令碼、裝置和其他組態。

instance_config = InstanceConfig(instanceType='ml.p3.2xlarge') 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套件在多個 GPUs 和多個執行個體之間分配工作負載時,執行大部分繁重作業。此套件已在 Braket PyTorch 和 TensorFlow 容器中預先設定。dist 模組會告知我們的演算法指令碼,訓練的 GPUs 總數 (world_size) 以及 GPU 核心local_rankrank和 。 rank 是 GPU 在所有執行個體的絕對索引,而 local_rank則是執行個體內的 GPU 索引。例如,如果有四個執行個體,每個都有八個 GPUs 配置給訓練,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)

接著,您可以DistributedSampler根據 定義 world_sizerank,然後將它傳遞到資料載入器。此取樣器會避免 GPUs 存取資料集的相同配量。

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 執行儲存和列印命令,日誌會充滿重複的資訊,而結果會互相覆寫。若要避免這種情況,您只能從具有 0 rank 的 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.p3.16xlarge執行個體類型。您可以透過混合任務中的InstanceConfig引數來設定執行個體類型。若要讓 SageMaker 分散式資料平行程式庫知道已啟用資料平行處理,您需要新增兩個額外的超參數,將 "sagemaker_distributed_dataparallel_enabled"設定為 "true" ,並將 "sagemaker_instance_type"設定為您正在使用的執行個體類型。smdistributed 套件會使用這兩個超參數。您的演算法指令碼不需要明確使用它們。在 Amazon Braket SDK 中,它提供了方便的關鍵字引數 distribution。在建立混合任務distribution="data_parallel"時,Amazon Braket SDK 會自動為您插入兩個超參數。如果您使用 Amazon Braket API,則需要包含這兩個超參數。

設定執行個體和資料平行處理後,您現在可以提交混合任務。ml.p3.16xlarge 執行個體中有 8 個 GPUs。當您設定 instanceCount=1 時,工作負載會分散到執行個體中的 8 個 GPUs。當您設定instanceCount大於一個時,工作負載會分散到所有執行個體中可用的 GPUs。使用多個執行個體時,每個執行個體都會根據您使用的時間產生費用。例如,當您使用四個執行個體時,計費時間為每個執行個體執行時間的四倍,因為有四個執行個體同時執行您的工作負載。

instance_config = InstanceConfig(instanceType='ml.p3.16xlarge', 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", ... )
注意

在上述混合任務建立中, train_dp.py是使用資料平行處理的修改演算法指令碼。請記住,只有在您根據上述章節修改演算法指令碼時,資料平行處理才能正常運作。如果啟用資料平行處理選項時沒有正確修改的演算法指令碼,混合任務可能會擲回錯誤,或者每個 GPU 可能會重複處理相同的資料配量,這很低效率。

讓我們比較 範例中的執行時間和成本,其中 使用 26 位數的量子電路訓練模型,以解決上述的二進位分類問題。此範例中使用的ml.p3.16xlarge執行個體每分鐘花費 0.4692 美元。如果沒有資料平行處理,模擬器大約需要 45 分鐘來訓練 1 個 Epoch 的模型 (即超過 208 個資料點),而且成本約為 20 美元。透過跨 1 個執行個體和 4 個執行個體的資料平行處理,分別只需要 6 分鐘和 1.5 分鐘,這兩者大約會轉換為 2.8 美元。透過在 4 個執行個體中使用資料平行處理,您不僅將執行時間提高 30 倍,也將成本大幅降低!