Acelere sus cargas de trabajo híbridas con simuladores integrados de PennyLane

Veamos cómo puede utilizar los simuladores integrados de PennyLane Amazon Braket Hybrid Jobs para ejecutar cargas de trabajo híbridas. El simulador integrado basado en la GPU de Pennylane utiliza la biblioteca Nvidia CuQuantum para lightning.gpu acelerar las simulaciones de circuitos. El simulador de GPU integrado viene preconfigurado en todos los contenedores de tareas de Braket, que los usuarios pueden utilizar de forma inmediata. En esta página, le mostramos cómo usarlo lightning.gpu para acelerar sus cargas de trabajo híbridas.

Utilización lightning.gpu para cargas de trabajo de algoritmos de optimización aproximada cuántica

Considere los ejemplos del algoritmo de optimización cuántica aproximada (QAOA) de este cuaderno. Para seleccionar un simulador integrado, debe especificar que el device argumento sea una cadena con la forma:. "local:<provider>/<simulator_name>" Por ejemplo, usted configuraría "local:pennylane/lightning.gpu" paralightning.gpu. La cadena de dispositivo que se proporciona al Hybrid Job al lanzarlo se transfiere al trabajo como variable de entorno"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)

En esta página, comparemos los dos simuladores vectoriales de PennyLane estado integrados lightning.qubit (que están basados en la CPU) y lightning.gpu (que están basados en la GPU). Deberás proporcionar a los simuladores algunas descomposiciones de compuertas personalizadas para poder calcular varios gradientes.

Ahora ya está listo para preparar el guion de lanzamiento de trabajos híbridos. Ejecutará el algoritmo QAOA mediante dos tipos de instancias: m5.2xlarge y. p3.2xlarge El tipo de m5.2xlarge instancia es comparable al de un portátil estándar para desarrolladores. p3.2xlargeSe trata de una instancia de computación acelerada que tiene una sola GPU NVIDIA Volta con 16 GB de memoria.

hyperparametersPara todos sus trabajos híbridos, será el mismo. Todo lo que necesita hacer para probar diferentes instancias y simuladores es cambiar dos líneas de la siguiente manera.

# 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')

o bien:

# 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')

nota

Si especificas el instance_config como mediante una instancia basada en la GPU, pero eliges device que sea el simulador integrado basado en la CPU (lightning.qubit), no se utilizará la GPU. ¡Asegúrate de usar el simulador integrado basado en la GPU si quieres apuntar a la GPU!

En primer lugar, puedes crear dos tareas híbridas y resolver Max-Cut con QAOA en un gráfico con 18 vértices. Esto se traduce en un circuito de 18 qubits, relativamente pequeño y fácil de ejecutar rápidamente en el portátil o en la instancia. 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,
)

El tiempo medio de iteración de la m5.2xlarge instancia es de unos 25 segundos, mientras que el de la p3.2xlarge instancia es de unos 12 segundos. Para este flujo de trabajo de 18 qubits, la instancia de GPU nos proporciona una aceleración del doble. Si echas un vistazo a la página de precios de Amazon Braket Hybrid Jobs, verás que el coste por minuto de una m5.2xlarge instancia es de 0,00768$, mientras que para la p3.2xlarge instancia es de 0,06375$. Ejecutarlo durante 5 iteraciones en total, como hiciste aquí, costaría 0,016$ con la instancia de CPU o 0,06375$ con la instancia de GPU, ¡ambas opciones bastante económicas!

Ahora vamos a complicar el problema e intentemos resolver un problema de corte máximo en un gráfico de 24 vértices, lo que se traducirá en 24 qubits. Vuelva a ejecutar los trabajos híbridos en las mismas dos instancias y compare el costo.

nota

¡Verá que el tiempo necesario para ejecutar este trabajo híbrido en la instancia de CPU puede ser de unas cinco horas!

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

El tiempo medio de iteración de la m5.2xlarge instancia es de aproximadamente una hora, mientras que el de la p3.2xlarge instancia es de aproximadamente dos minutos. Para este problema mayor, ¡la instancia de la GPU es un orden de magnitud más rápida! Todo lo que tenías que hacer para beneficiarte de esta aceleración era cambiar dos líneas de código, sustituyendo el tipo de instancia por el simulador local utilizado. Ejecutarlo durante 5 iteraciones en total, como se hizo aquí, costaría unos 2.27072$ con la instancia de CPU o unos 0.775625$ con la instancia de GPU. El uso de la CPU no solo es más caro, sino que también lleva más tiempo ejecutarlo. Acelerar este flujo de trabajo con una instancia PennyLane de GPU disponible AWS, mediante el simulador integrado respaldado por NVIDIA CuQuantum, te permite ejecutar flujos de trabajo con recuentos de cúbits intermedios (entre 20 y 30) con un coste total menor y en menos tiempo. Esto significa que puedes experimentar con la computación cuántica incluso para problemas que son demasiado grandes como para ejecutarlos rápidamente en un portátil o en una instancia de tamaño similar.

Aprendizaje automático cuántico y paralelismo de datos

Si su tipo de carga de trabajo es el aprendizaje automático cuántico (QML) que se entrena con conjuntos de datos, puede acelerar aún más su carga de trabajo mediante el paralelismo de datos. En QML, el modelo contiene uno o más circuitos cuánticos. El modelo puede o no contener también redes neuronales clásicas. Al entrenar el modelo con el conjunto de datos, los parámetros del modelo se actualizan para minimizar la función de pérdida. Por lo general, se define una función de pérdida para un único punto de datos y la pérdida total para la pérdida media de todo el conjunto de datos. En QML, las pérdidas generalmente se calculan en serie antes de promediar la pérdida total para los cálculos de gradientes. Este procedimiento lleva mucho tiempo, especialmente cuando hay cientos de puntos de datos.

Como la pérdida de un punto de datos no depende de otros puntos de datos, ¡las pérdidas se pueden evaluar en paralelo! Las pérdidas y los gradientes asociados a diferentes puntos de datos se pueden evaluar al mismo tiempo. Esto se conoce como paralelismo de datos. Con SageMaker la biblioteca paralela de datos distribuida, Amazon Braket Hybrid Jobs le facilita aprovechar el paralelismo de datos para acelerar su entrenamiento.

Considere la siguiente carga de trabajo de QML para el paralelismo de datos, que utiliza el conjunto de datos del conjunto de datos Sonar del conocido repositorio de la UCI como ejemplo de clasificación binaria. El conjunto de datos del sonar tiene 208 puntos de datos, cada uno con 60 características que se recopilan a partir de las señales del sonar que rebotan en los materiales. Cada punto de datos se etiqueta con una «M» para las minas o con una «R» para las rocas. Nuestro modelo QML consta de una capa de entrada, un circuito cuántico como capa oculta y una capa de salida. Las capas de entrada y salida son redes neuronales clásicas implementadas en PyTorch. El circuito cuántico se integra con las redes PyTorch neuronales mediante PennyLane el módulo qml.qnn. Consulte nuestros ejemplos de cuadernos para obtener más información sobre la carga de trabajo. Al igual que en el ejemplo anterior de QAOA, puedes aprovechar la potencia de la GPU mediante el uso de simuladores integrados basados en la GPU, lightning.gpu para mejorar el rendimiento en comparación con los simuladores basados en la CPU integrada. PennyLane

Para crear un trabajo híbrido, puedes llamar AwsQuantumJob.create y especificar el script del algoritmo, el dispositivo y otras configuraciones mediante sus argumentos de palabras clave.

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

Para utilizar el paralelismo de datos, es necesario modificar algunas líneas de código en el script del algoritmo para que la biblioteca SageMaker distribuida paralelice correctamente el entrenamiento. En primer lugar, importa el smdistributed paquete que se encarga de la mayor parte del trabajo pesado de distribuir las cargas de trabajo entre varias GPU y varias instancias. Este paquete viene preconfigurado en Braket y en los contenedores. PyTorch TensorFlow El dist módulo indica a nuestro script de algoritmo cuál es el número total de GPU para el entrenamiento (world_size), así como el número rank y del núcleo local_rank de la GPU. rankes el índice absoluto de una GPU en todas las instancias, mientras que local_rank es el índice de una GPU dentro de una instancia. Por ejemplo, si hay cuatro instancias, cada una con ocho GPU asignadas para el entrenamiento, los rank rangos van de 0 a 31 y los local_rank rangos de 0 a 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)

A continuación, se define una en DistributedSampler función de world_size y rank y, a continuación, se pasa al cargador de datos. Este muestreador evita que las GPU accedan al mismo segmento de un conjunto de datos.

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

A continuación, usa la DistributedDataParallel clase para habilitar el paralelismo de datos.

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"])

Los cambios anteriores son los que necesita para usar el paralelismo de datos. En QML, a menudo quieres guardar los resultados e imprimir el progreso del entrenamiento. Si cada GPU ejecuta el comando de guardar e imprimir, el registro se inundará con la información repetida y los resultados se sobrescribirán entre sí. Para evitarlo, solo puedes guardar e imprimir desde la GPU que tenga rank 0.

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 admite tipos de ml.p3.16xlarge instancias para la biblioteca paralela de datos SageMaker distribuidos. El tipo de instancia se configura mediante el InstanceConfig argumento de Hybrid Jobs. Para que la biblioteca paralela de datos SageMaker distribuidos sepa que el paralelismo de datos está habilitado, debe agregar dos hiperparámetros adicionales: configurar "true" y "sagemaker_distributed_dataparallel_enabled" "sagemaker_instance_type" configurar el tipo de instancia que está utilizando. El paquete utiliza estos dos hiperparámetros. smdistributed No es necesario que el script de su algoritmo los utilice de forma explícita. En Amazon Braket SDK, proporciona un práctico argumento de palabra clave. distribution distribution="data_parallel"En el caso de la creación de empleo híbrido, el SDK de Amazon Braket inserta automáticamente los dos hiperparámetros. Si utilizas la API Amazon Braket, debes incluir estos dos hiperparámetros.

Con el paralelismo de instancias y datos configurado, ahora puede enviar su trabajo híbrido. Hay 8 GPU en una instancia. ml.p3.16xlarge Cuando lo configurasinstanceCount=1, la carga de trabajo se distribuye entre las 8 GPU de la instancia. Si configuras instanceCount más de una, la carga de trabajo se distribuye entre las GPU disponibles en todas las instancias. Al usar varias instancias, cada instancia conlleva un cargo en función del tiempo que la utilices. Por ejemplo, cuando usa cuatro instancias, el tiempo facturable es cuatro veces el tiempo de ejecución por instancia, ya que hay cuatro instancias que ejecutan sus cargas de trabajo al mismo tiempo.

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

nota

En la creación de empleos híbridos anterior, train_dp.py se encuentra el script de algoritmo modificado para utilizar el paralelismo de datos. Tenga en cuenta que el paralelismo de datos solo funciona correctamente cuando modifica el script del algoritmo de acuerdo con la sección anterior. Si la opción de paralelismo de datos está habilitada sin un script de algoritmo modificado correctamente, es posible que el trabajo híbrido genere errores o que cada GPU procese repetidamente el mismo segmento de datos, lo que resulta ineficiente.

Comparemos el tiempo de ejecución y el coste en un ejemplo en el que entrenamos un modelo con un circuito cuántico de 26 qubits para resolver el problema de clasificación binaria mencionado anteriormente. La ml.p3.16xlarge instancia utilizada en este ejemplo cuesta 0,4692$ por minuto. Sin el paralelismo de datos, el simulador tarda unos 45 minutos en entrenar el modelo para una época (es decir, más de 208 puntos de datos) y cuesta unos 20 dólares. Con el paralelismo de datos en 1 instancia y 4 instancias, solo se necesitan 6 minutos y 1,5 minutos respectivamente, lo que se traduce en unos 2,8$ para ambas. Al utilizar el paralelismo de datos en 4 instancias, no solo se mejora el tiempo de ejecución en 30 veces, ¡sino que también se reducen los costes en un orden de magnitud!