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

使用 PennyLane 嵌入式模拟器运行混合工作负载

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

使用适用于 QAOA 工作负载的 lightning.gpu

考虑本 Notebook 中的量子近似优化算法（QAOA）示例。要选择嵌入式模拟器，请将 device 参数指定为以下形式的字符串："local:<provider>/<simulator_name>"。例如，您可以设置 lightning.gpu 的 "local:pennylane/lightning.gpu"。启动时提供给 Hybrid Jobs 的设备字符串将作为环境变量 "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 算法：ml.m5.2xlarge 和 ml.g4dn.xlarge。ml.m5.2xlarge 实例类型相当于标准的开发人员 Notebook。ml.g4dn.xlarge是一个加速计算实例，它有一个 NVIDIA T4 GPU 和 16GB 内存。

要运行 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 定价页面，你可以看到实例的每分钟费用为 0.00768 美元，而m5.2xlarge实例的每分钟费用为 0.01227 ml.g4dn.xlarge 美元。在这种情况下，在 GPU 实例上运行更快、更便宜。

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

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

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

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

要创建混合作业，您可以通过其关键字参数调用 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软件包，该软件包负责在多个 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)

接下来，您根据 world_size 和 rank 定义一个 DistributedSampler，然后将其传递到数据加载器中。此采样器可避免 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.g4dn.12xlarge实例类型。您可以通过 Hybrid Jobs 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.g4dn.12xlarge实例 GPUs 中有 4 个。设置后instanceCount=1，工作负载将分布在实例 GPUs 中的 8 个中。当您设置instanceCount大于 1 时，工作负载将分布在所有 GPUs 可用实例中。使用多个实例时，每个实例会根据您的使用时间产生费用。例如，当您使用四个实例时，计费时间是每个实例运行时的四倍，因为有四个实例同时运行您的工作负载。

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",
    ...
)

如果使用得当，使用多个实例可以使时间和成本减少几个数量级。有关更多详细信息，请参阅示例笔记本。