使用来自的嵌入式模拟器加速混合工作负载 PennyLane

让我们来看看如何使用 Amazon Braket Hybrid Jobs PennyLane 上的嵌入式模拟器来运行混合工作负载。Pennylane 基于 GPU 的嵌入式仿真器使用 Nvidia cuQuantum 库来加速电路仿真。lightning.gpu嵌入式 GPU 模拟器已在所有 Braket 作业容器中进行了预配置，用户可以开箱即用。在本页中，我们将向您展示如何使用lightning.gpu来加快混合工作负载的速度。

lightning.gpu用于量子近似优化算法工作负载

考虑本笔记本中的量子近似优化算法 (QAOA) 示例。要选择嵌入式模拟器，请将device参数指定为以下形式的字符串："local:<provider>/<simulator_name>"。例如，您可以设置"local:pennylane/lightning.gpu"lightning.gpu。启动时提供给 Hybrid Job 的设备字符串将作为环境变量传递给该作业"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.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 的实例，但选择作为基于 CPU 的嵌入式模拟器 (lightning.qubit)，则不会使用 GPU。device如果你想瞄准 GPU，一定要使用基于 GPU 的嵌入式模拟器！

首先，您可以创建两个混合作业，并在具有 18 个顶点的图表上使用 QAOA 求解 Max-Cut。这意味着一个 18 量子比特的电路——相对较小，在笔记本电脑或实例上快速运行是可行的。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 量子比特的工作流程，GPU 实例为我们提供了 2 倍的加速。如果你看一下Amazon Braket Hybrid Jobs的定价页面，你可以看到一个m5.2xlarge实例的每分钟费用为0.00768美元，而实例的每分钟费用为0.06375美元。p3.2xlarge像你在这里所做的那样，总共运行 5 次迭代，使用 CPU 实例的费用为 0.016 美元，使用 GPU 实例的费用为 0.06375 美元，两者都非常便宜！

现在让我们让问题变得更难，尝试在 24 个顶点图上求解一个 Max-Cut 问题，该问题将转换为 24 个量子比特。在相同的两个实例上再次运行混合作业并比较成本。

注意

您将看到，在 CPU 实例上运行此混合作业的时间可能约为五个小时！

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 支持的嵌入式仿真器 AWS，使用上可用 PennyLane的 GPU 实例加速这一工作流程 CuQuantum，使您能够以更低的总成本和更短的时间运行具有中间量子比特计数（介于 20 到 30 之间）的工作流程。这意味着，即使问题太大，无法在笔记本电脑或类似大小的实例上快速运行，您也可以尝试量子计算。

量子机器学习和数据并行性

如果您的工作负载类型是基于数据集训练的量子机器学习 (QML)，则可以使用数据并行性进一步加快工作负载。在 QML 中，模型包含一个或多个量子电路。该模型可能还包含也可能不包含经典神经网络。使用数据集训练模型时，会更新模型中的参数以最小化损失函数。损失函数通常是针对单个数据点定义的，以及整个数据集的平均损失的总损失。在 QML 中，通常先串行计算损耗，然后再求平均为梯度计算的总损耗。此过程非常耗时，尤其是在有数百个数据点的情况下。

由于一个数据点的损失不依赖于其他数据点，因此可以并行评估损失！可以同时评估与不同数据点相关的损失和梯度。这就是所谓的数据并行性。借助 SageMaker分布式数据并行库，Amazon Braket Hybrid Jobs 使您可以更轻松地利用数据并行性来加速训练。

考虑以下 QML 数据并行工作负载，该工作负载使用知名 UCI 存储库中的 Sonar 数据集数据集作为二进制分类的示例。声纳数据集有 208 个数据点，每个数据点有 60 个特征，这些特征是从材料上反弹的声纳信号中收集的。每个数据点要么被标记为 “M”（代表地雷），要么标记为 “R”（表示岩石）。我们的 QML 模型由输入层、作为隐藏层的量子电路和输出层组成。输入层和输出层是中实现的经典神经网络 PyTorch。量子电路使用的 qml.q PennyLane nn 模块与 PyTorch 神经网络集成。有关工作负载的更多详细信息，请参阅我们的示例笔记本。就像上面的 QAOA 示例一样，你可以利用基于 GPU 的嵌入式模拟器（比如）lightning.gpu来提高基于 CPU PennyLane 的嵌入式模拟器的性能。

要创建混合作业，您可以通过其关键字参数调用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软件包，该软件包负责将工作负载分配到多个 GPU 和多个实例。此软件包已在 Braket PyTorch 和 TensorFlow容器中预先配置。该dist模块告诉我们的算法脚本用于训练 (world_size) 的 GPU 总数，以及 GPU 内核local_rank的rank和。 rank是 GPU 在所有实例中的绝对索引，而local_rank是 GPU 在实例中的索引。例如，如果有四个实例，每个实例都为训练分配了八个 GPU，则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然后将其传递到数据加载器中。此采样器可避免 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 混合任务支持 SageMaker 分布式数据并行库的ml.p3.16xlarge实例类型。您可以通过 Hybrid Job InstanceConfig s 中的参数配置实例类型。要使 SageMaker 分布式数据并行库知道数据并行性已启用，您需要再添加两个超参数，即"sagemaker_distributed_dataparallel_enabled"设置为正在使用的实例类型，"true"并"sagemaker_instance_type"设置为正在使用的实例类型。这两个超参数由smdistributed软件包使用。您的算法脚本无需明确使用它们。在 Amazon Braket SDK 中，它提供了一个方便的关键字参数。distributiondistribution="data_parallel"在混合任务创建中，Amazon Braket SDK 会自动为您插入两个超参数。如果您使用 Amazon Braket API，则需要包含这两个超参数。

配置好实例和数据并行度后，您现在可以提交混合作业了。一个ml.p3.16xlarge实例中有 8 个 GPU。设置后instanceCount=1，工作负载将分布在实例中的 8 个 GPU 上。如果设置instanceCount大于 1，则工作负载将分布在所有实例中可用的 GPU 上。使用多个实例时，每个实例会根据您的使用时间收取费用。例如，当您使用四个实例时，计费时间是每个实例运行时间的四倍，因为有四个实例同时运行您的工作负载。

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 个纪元（即超过 208 个数据点）内训练模型，成本约为 20 美元。在 1 个实例和 4 个实例之间实现数据并行处理时，分别只需 6 分钟和 1.5 分钟，这意味着两者都花费大约 2.8 美元。通过在 4 个实例上使用数据并行性，您不仅可以将运行时间缩短 30 倍，还可以将成本降低一个数量级！