Paralelismo híbrido de datos particionados

El paralelismo de datos particionados es una técnica de entrenamiento distribuido que ahorra memoria y que divide el estado de un modelo (parámetros del modelo, gradientes y estados del optimizador) entre dispositivos. Esto ayuda a adaptarse a un modelo más grande o a aumentar el tamaño del lote utilizando la memoria liberada de la GPU. La biblioteca de SMP ofrece la capacidad de ejecutar paralelismo de datos particionados con paralelismo de datos totalmente particionados (FSDP) de PyTorch. El PyTorch FSDP particiona de forma predeterminada en todo el conjunto de GPU que se utilizan. En SMP v2, la biblioteca ofrece este paralelismo de datos particionados además del PyTorch FSDP extendiendo el particionamiento híbrido de PyTorch (HYBRID_SHARD), que es una de las estrategias de particionamiento que proporciona PyTorch FSDP: FULL_SHARD, SHARD_GRAD_OP, HYBRID_SHARD, _HYBRID_SHARD_ZERO2. Ampliar el particionamiento híbrido de esta manera ayuda a implementar particionamiento con reconocimiento de escala, como se describe en el blog Near-linear scaling of gigantic-model training en AWS para PyTorch FSDP.

La biblioteca de SMP facilita el uso de HYBRID_SHARD y _HYBRID_SHARD_ZERO2 en cualquier número configurable de GPU, lo que amplía el PyTorch FSDP nativo, que admite particionamiento en un solo nodo (HYBRID_SHARD) o en todas las GPU (FULL_SHARD). Las llamadas del PyTorch FSDP pueden permanecer como están y solo necesita agregar el argumento hybrid_shard_degree a la configuración de SMP, como se muestra en el siguiente código de ejemplo. No necesita cambiar el valor del argumento sharding_strategy en el encapsulador del PyTorch FSDP alrededor del modelo de PyTorch. Puede pasar ShardingStrategy.HYBRID_SHARD como valor. Como alternativa, la biblioteca de SMP anula la estrategia del script y la establece en ShardingStrategy.HYBRID_SHARD si especifica un valor igual o mayor que 2 en el parámetro hybrid_shard_degree.

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 hybrid_shard_degree, consulte Parámetros de configuración de las características esenciales de SMP v2.

Diccionario de configuración de SMP

{ "hybrid_shard_degree": 16 }

En script de entrenamiento

import torch.sagemaker as tsm
tsm.init()

# Set up a PyTorch model
model = ...

# Wrap the PyTorch model using the PyTorch FSDP module
model = FSDP(
    model,
    ...
)

# Optimizer needs to be created after FSDP wrapper
optimizer = ...