Paralelismo de contexto

El paralelismo de contexto es un tipo de paralelismo de modelos que divide las activaciones de modelos a lo largo de la dimensión secuencial. A diferencia de otras técnicas de paralelismo secuencial, que solo dividen LayerNorm y RMSNorm, el paralelismo de contexto divide las entradas de red y todas las activaciones intermedias a lo largo de la dimensión secuencial.

SMP v2 se integra con Transformer Engine para realizar paralelismo contextual y se puede usar junto con Paralelismo de tensores de FSDP y SMP de PyTorch. Puede habilitar los tres paralelismos simultáneamente para entrenamiento de modelos. El paralelismo de contexto es beneficioso para entrenar modelos con grandes tamaños de activación y longitudes de secuencia largas. Acelera la computación de puntuaciones de atención y salidas de atención, ya que permite que cada dispositivo compute solo una parte de las puntuaciones y las salidas a lo largo de la dimensión secuencial. Si bien el paralelismo de tensores también acelera la computación mediante la partición a lo largo de la dimensión oculta, la ventaja del paralelismo de contexto es más sustancial, ya que los requisitos de computación aumentan cuadráticamente con la dimensión secuencial.

Modelos de transformador Hugging Face compatibles con paralelismo de contexto de SMP

Actualmente, SMP v2 admite paralelismo de contexto para los siguientes modelos de transformador Hugging Face.

Configuración de paralelismo de contexto

Defina un valor entero para el parámetro context_parallel_degree que divide equitativamente el número de GPU del clúster. Por ejemplo, si tiene una instancia de 8 GPU, utilice 2, 4 u 8 para context_parallel_degree. Recomendamos empezar con un valor de context_parallel_degree pequeño y aumentarlo gradualmente hasta que el modelo se ajuste a la memoria de la GPU con la longitud de secuencia de entrada requerida.

En los siguientes fragmentos de código, se muestra cómo añadir el torch.sagemaker.init() del módulo de inicialización de SMP al script de entrenamiento y cómo configurar el diccionario de configuración de SMP en formato JSON para el lanzador de trabajos de entrenamiento mientras se sigue el proceso de dos pasos descrito en Uso de la biblioteca de paralelismo de modelos de SageMaker v2. No necesita hacer cambios en el modelo de PyTorch o en la configuración del PyTorch FSDP. Para obtener más información sobre el parámetro context_parallel_degree, consulte Parámetros de configuración de las características esenciales de SMP v2.

En el script de entrenamiento

Como parte del paso 1, inicialice el script con torch.sagemaker.init() para activar SMP v2 y encapsular el modelo con la API torch.sagemaker.transform.

A partir de SMP v2.6.0, puede usar el argumento cp_comm_type para determinar qué implementación de paralelismo de contexto debe usar. La biblioteca de SMP admite actualmente dos implementaciones: p2p y all_gather. La implementación p2p utiliza llamadas de envío y recepción entre pares para acumular claves-valores durante la implementación de atención y se ejecuta de forma asíncrona, lo que permite superposiciones con computación. La implementación de all_gather, por su parte, usa la operación colectiva AllGather y se ejecuta de forma sincrónica.

import torch.sagemaker as tsm
tsm.init()

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_config(..)
model = tsm.transform(model, cp_comm_type="p2p")

Configuración de SMP

Como parte del paso 2, añada el siguiente parámetro al diccionario de configuración de SMP para el estimador de PyTorch de SageMaker.

{   
    ..., # other SMP config parameters
    "context_parallel_degree": 2
}