Paralelismo de tensores

El paralelismo de tensores es un tipo de paralelismo de modelos en el que las ponderaciones, gradientes y estados del optimizador específicos del modelo se dividen entre los dispositivos. A diferencia del paralelismo por canalización, que mantiene intactos los pesos individuales pero divide el conjunto de pesos, gradientes u optimizador entre dispositivos, el paralelismo tensorial fragmenta los pesos individuales. Por lo general, esto implica el cálculo distribuido de operaciones, módulos o capas específicos del modelo.

El paralelismo de tensores es necesario en los casos en los que un único parámetro consume la mayor parte de la memoria de la GPU (por ejemplo, tablas de incrustación grandes con un gran tamaño de vocabulario o una capa softmax grande con un gran número de clases). En este caso, tratar este gran tensor u operación como una unidad atómica es ineficiente e impide equilibrar la carga de memoria.

SMP v2 se integra con Transformer Engine para implementar el paralelismo tensorial y se ejecuta sobre las API del FSDP. PyTorch Puede habilitar el paralelismo tensorial PyTorch FSDP y SMP simultáneamente y determinar el mejor paralelismo del modelo para obtener el mejor rendimiento.

En la práctica, el paralelismo tensorial es especialmente útil en los siguientes escenarios.

• Cuando se entrena con contextos largos, se obtiene una memoria de activación elevada solo con el FSDP.

• Cuando entrenas con clústeres muy grandes en los que el tamaño del lote global supera los límites deseados.

Modelos de Hugging Face Transformer compatibles con el paralelismo tensorial SMP

Actualmente, SMP v2 ofrece soporte de paralelismo tensorial para los siguientes modelos de transformadores Hugging Face.

• GPT-NeoX

• Llama 2

Para ver la configuración de referencia para aplicar el paralelismo tensorial en estos modelos, consulte. Sugerencias de configuración

Configure el paralelismo tensorial

Paratensor_parallel_degree, selecciona un valor para el grado de paralelismo tensorial. El valor debe dividir en partes iguales el número de GPU del clúster. Por ejemplo, para fragmentar el modelo mientras usas una instancia con 8 GPU, elige 2, 4 u 8. Te recomendamos empezar con un número pequeño y aumentarlo gradualmente hasta que el modelo quepa en la memoria de la GPU.

En los siguientes fragmentos de código se muestra cómo añadir el módulo de inicialización de SMP torch.sagemaker.init() a su guion de formación y cómo configurar el diccionario de configuración de SMP en formato JSON para el lanzador de tareas de entrenamiento, siguiendo el proceso de dos pasos descrito en la presente descripción. Comience con la biblioteca de paralelismo de SageMaker modelos v2 No es necesario realizar ningún cambio en el modelo ni en la configuración del FSDP. PyTorch PyTorch Para obtener más información acerca de los parámetros tensor_parallel_degree y random_seed, consulte Parámetros de configuración de las funciones principales del SMP v2.

Configuración SMP

{
    "tensor_parallel_degree": 8,
    "random_seed": 0 
}

En tu guion de entrenamiento

Inicialice con torch.sagemaker.init() para activar SMP v2 y empaquetar su modelo con la torch.sagemaker.transform API.

import torch.sagemaker as tsm
tsm.init()

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_config(..)
model = tsm.transform(model)

Guardar y cargar los puntos de control de Hugging Face Transformer

Una vez que la biblioteca SMP transforma un modelo, cambia el diccionario de estados (state_dict) del modelo. Esto significa que el modelo deja de ser compatible con las funcionalidades de puntos de control originales del Hugging Face Transformer. Para ello, la biblioteca SMP proporciona API para guardar los puntos de control de un modelo transformado en la representación de Hugging Face Transformer, y la API para cargar un punto de control torch.sagemaker.transform del modelo Hugging Face Transformer para realizar ajustes precisos.

Para obtener más información sobre cómo guardar puntos de control mientras se utiliza la función de paralelismo tensorial de SMP v2, consulte. Guarde y cargue los puntos de control mientras usa SMP

Para obtener más información sobre cómo ajustar un modelo aplicando la función de paralelismo tensorial de SMP v2, consulte. Microajuste