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 SMP ofrece la posibilidad de ejecutar el paralelismo de datos fragmentados con Fully Sharded Data Parallel (FSDP). PyTorch PyTorch De forma predeterminada, el FSDP se divide en todo el conjunto de GPU que se utilizan. En SMP v2, la biblioteca ofrece este paralelismo de datos fragmentados además del PyTorch FSDP al extender la fragmentación PyTorch híbrida (HYBRID_SHARD), que es una de las estrategias de fragmentación que ofrece FSDP:,,,. PyTorch FULL_SHARD SHARD_GRAD_OP HYBRID_SHARD _HYBRID_SHARD_ZERO2 Ampliar la fragmentación híbrida de esta manera ayuda a implementar la fragmentación basada en la escalabilidad, tal como se describe en el blog Scaling of Gigantic-model training sobre el FSDP. Near-linear AWS PyTorch

La biblioteca SMP facilita su uso HYBRID_SHARD en cualquier cantidad configurable de GPU, lo que amplía el PyTorch FSDP nativo que admite la fragmentación en un solo nodo () o en _HYBRID_SHARD_ZERO2 todas las GPU (). HYBRID_SHARD FULL_SHARD PyTorch Las llamadas del FSDP pueden permanecer como están y solo necesitas añadir el hybrid_shard_degree argumento a la configuración del SMP, tal y como se muestra en el siguiente ejemplo de código. No necesitas cambiar el valor del sharding_strategy argumento en el contenedor PyTorch FSDP que rodea a tu modelo. 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 Utilice la biblioteca de paralelismo de SageMaker modelos v2. No es necesario realizar ningún cambio en el PyTorch modelo ni PyTorch en la configuración del 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 = ...