Entrenamiento de precisión mixta

La biblioteca de paralelismo de modelos de SageMaker (SMP) v2 admite entrenamiento mixto de precisión lista para usar mediante la integración con marcos de código abierto como PyTorch FSDP y Transformer Engine. Para obtener más información, consulte los temas siguientes.

Entrenamiento de precisión mixto con FP8 en instancias P5 utilizando Transformer Engine

A partir de la biblioteca de paralelismo de modelos de SageMaker (SMP) v2.2.0, la biblioteca de SMP se integra con Transformer Engine y admite entrenamiento de precisión mixta de FP8 listo para usar, manteniendo la compatibilidad con PyTorch FSDP MixedPrecision. Esto significa que puede usar PyTorch FSDP para entrenamiento de precisión mixta y Transformer Engine para el entrenamiento de FP8. En el caso de las capas de modelos que no admite la característica de entrenamiento FP8 de Transformer Engine, esas capas recurren a precisión mixta de PyTorch FSDP.

FP8 (precisión de coma flotante de 8 bits) es un tipo de datos que se ha convertido en otro paradigma para acelerar el entrenamiento con aprendizaje profundo de los modelos de LLM. Con el lanzamiento de las GPU H100 de NVIDIA que admiten tipos de datos FP8, puede beneficiarse de las ventajas de las mejoras de rendimiento de las instancias P5 equipadas con GPU H100 y, al mismo tiempo, acelerar el entrenamiento distribuido con entrenamiento de precisión mixta de FP8.

El tipo de datos FP8 se ramifica aún más a los formatos E4M3 y E5M2. E4M3 ofrece mayor precisión, tiene un rango dinámico limitado y es ideal para la pasada hacia delante en entrenamiento de modelos. E5M2 tiene un rango dinámico más amplio pero menor precisión, y es más adecuado para la pasada hacia atrás, donde la precisión es menos crítica y resulta beneficioso un rango dinámico más amplio. Por tanto, recomendamos utilizar la receta de estrategia híbrida FP8 para aprovechar estas características de forma eficaz.

Para tipos de datos de precisión media (FP16 y BF16), las técnicas de escalado de pérdidas globales, como el escalado de pérdidas estático o el escalado de pérdidas dinámico, gestionan los problemas de convergencia que surgen de la pérdida de información debido al redondeo de los gradientes con precisión media. Sin embargo, el rango dinámico de FP8 es aún más estrecho y las técnicas de escalado de pérdidas globales no son suficientes. En este punto necesitamos una técnica de escalado por tensor más detallada. El escalado diferido es una estrategia que selecciona un factor de escalado en función de los valores absolutos máximos observados en varios tensores de iteraciones anteriores. Esta estrategia tiene su contrapartida: aprovecha todas las ventajas de rendimiento de la computación FP8, pero requiere memoria para mantener el historial de valores máximos de los tensores. Para obtener más información sobre la estrategia de escalado diferido en general, consulte el documento FP8 Formats for Deep Learning.

En la práctica, el uso del FP8 resulta útil en todos los escenarios de entrenamiento de instancias P5. Recomendamos encarecidamente activar FP8 siempre que sea posible para mejorar el rendimiento del entrenamiento.

SMP v2 es compatible con Transformer Engine tal como se entrega. Por tanto, cuando ejecute entrenamiento de FP8 con SMP v2 en instancias P5 de SageMaker (ml.p5.48xlarge), lo único que debe hacer es importar torch.sagemaker en el script de entrenamiento y seguir utilizando el paquete Python nativo de Transformer Engine. Para obtener más información sobre el uso de Transformer Engine para el entrenamiento de FP8 en general, consulte Using FP8 with Transformer Engine en la documentación de NVIDIA Transformer Engine. El siguiente fragmento de código muestra el aspecto que deberían tener las líneas de código para importar la biblioteca de SMP y configurar FP8 en un script de entrenamiento.

import torch.sagemaker as tsm
import transformer_engine.pytorch as te
from transformer_engine.common.recipe import DelayedScaling, Format

# Initialize the SMP torch.sagemaker API.
tsm.init()

# Define a transformer model and wrap it with the torch.sagemaker.transform API.
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_config(ModelConfig)
model = tsm.transform(model)

# Enable E4M3 during forward pass, E5M2 during backward pass.
fp8_format = Format.HYBRID

# Create an FP8 recipe.
fp8_recipe = DelayedScaling(fp8_format=fp8_format, amax_history_len=32, amax_compute_algo="max")

# Enable FP8 autocasting.
with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe, fp8_group=tsm.state.world_process_group):
    out = model(inp)

loss = out.sum()
loss.backward()

Para obtener un ejemplo práctico de entrenamiento de FP8 con SMP v2 en instancias P5, consulte el cuaderno de ejemplo en Accelerate SageMaker PyTorch FSDP Training of Llama-v2 (or GPT-NeoX) with FP8 on P5 instances.

Entrenamiento de precisión mixto con tipos de datos de media precisión utilizando PyTorch FSDP

SMP v2 es compatible con MixedPrecision de FSDP PyTorch para trabajos de entrenamiento en instancias P4 y P5. PyTorch FSDP proporciona varias configuraciones para una precisión mixta tanto para mejorar el rendimiento como para reducir la memoria.

La forma estándar de configurar un modelo para precisión mixta consiste en crear el modelo en float32 y, a continuación, permitir que FSDP envíe los parámetros a float16 o bfloat16 sobre la marcha pasando una política de MixedPrecision, como se muestra en el siguiente fragmento de código. Para obtener más información sobre las opciones para cambiar el dtype de parámetros, reducción o búferes para conseguir precisión mixta en PyTorch, consulte PyTorch FSDP MixedPrecision API en la documentación de PyTorch.

# Native PyTorch API
from torch.distributed.fsdp import MixedPrecision

dtype = torch.bfloat16
mixed_precision_policy = MixedPrecision(
    param_dtype=dtype, reduce_dtype=dtype, buffer_dtype=dtype
)

model = FSDP(
    model,
    ...,
    mixed_precision=mixed_precision_policy
)

Tenga en cuenta que algunos modelos (como el modelo Transformers Llama de Hugging Face) esperan búferes como float32. Para usar float32, sustituya torch.bfloat16 por torch.float32 en la línea que define el objeto dtype.