使用 PennyLane 嵌入式仿真器运行混合工作负载 - Amazon Braket
lightning.gpu用于QAOA工作负载量子机器学习和数据并行性

本文属于机器翻译版本。若本译文内容与英语原文存在差异，则一律以英文原文为准。

使用 PennyLane 嵌入式仿真器运行混合工作负载

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

lightning.gpu用于QAOA工作负载

考虑本笔记本中的量子近似优化算法 (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.2xlarge。m5.2xlarge实例类型相当于标准的开发者笔记本电脑。p3.2xlarge是一个加速计算实例，它有一个GPU带有 16GB 内存的 NVIDIA Volta。

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')
注意

如果您使用GPU基于的实例指定为，但选择instance_config作为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 次迭代，使用实例的费用为 0.016 美元，使用CPU实例的费用为 0.06375 美元，两者GPU都非常便宜！

现在让我们让问题变得更难，尝试在 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 次迭代，使用该实例的费用约为 2.27072 美元，使用该CPU实例的费用约为 0.775625 美元。GPUCPU使用不仅更昂贵，而且运行时间也更长。使用由 PennyLane的嵌入式仿真器提供支持 AWS，使用上可用的GPU实例来加速此工作流程 NVIDIA 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通过使用GPU基于嵌入式的模拟器来利用 PennyLane的强大功能lightning.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是 a 在GPU所有实例中的绝对索引，而local_rank是实例GPU内的 a 的索引。例如，如果有四个实例，每个实例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运行保存和打印命令，则日志中将充斥重复的信息，结果将相互覆盖。为避免这种情况，您只能从有 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 BraketAPI，则需要包含这两个超参数。

配置好实例和数据并行度后，您现在可以提交混合作业了。一个ml.p3.16xlarge实例GPUs中有 8 个。设置后instanceCount=1，工作负载将分布在实例GPUs中的 8 个中。当您设置instanceCount大于 1 时，工作负载将分布在所有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 个纪元（即超过 208 个数据点）内训练模型，成本约为 20 美元。在 1 个实例和 4 个实例之间实现数据并行处理时，分别只需 6 分钟和 1.5 分钟，这意味着两者的费用约为 2.8 美元。通过在 4 个实例上使用数据并行性，您不仅可以将运行时间缩短 30 倍，还可以将成本降低一个数量级！