使用 lightning.gpu處理QAOA工作負載 Quantum 機器學習和資料平行處理

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

讓我們看看如何在 Amazon Braket Hybrid Jobs PennyLane 上使用的內嵌模擬器來執行混合工作負載。Pennylane GPU的型內嵌模擬器 lightning.gpu使用 Nvidia cuQuantum 程式庫來加速電路模擬。內嵌模擬器已預先在所有 Braket GPU 任務容器中設定，使用者可以立即使用。在此頁面中，我們會示範如何使用 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.2xlarge和 p3.2xlarge。m5.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 用量不僅更昂貴，而且需要更多時間來執行。使用支援的 PennyLane內嵌模擬器 AWS，透過上的GPU執行個體加速此工作流程NVIDIA CuQuantum，可讓您以較低的總成本和更短的時間，執行具有中繼 qubit 計數（介於 20 到 30 之間）的工作流程。這表示即使問題太大而無法在筆記型電腦或類似大小的執行個體上快速執行，您也可以試用量子運算。

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

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

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

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

若要建立混合任務，您可以透過其關鍵字引數呼叫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_rank的 rank和。 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 執行個體GPUs中有 8 個。當您設定 instanceCount=1 時，工作負載會分散到執行個體GPUs中的 8。當您設定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 倍，還可以將成本大幅降低！

您的瀏覽器已停用或無法使用 Javascript。

您必須啟用 Javascript，才能使用 AWS 文件。請參閱您的瀏覽器說明頁以取得說明。

文件慣用形式

使用混合任務和 PennyLane 執行QAOA演算法

攜帶您自己的容器 (BYOC)