Paralelismo híbrido de datos fragmentados

El paralelismo de datos fragmentados 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 los dispositivos. Esto le ayuda a adaptarse a un modelo más grande o a aumentar el tamaño del lote utilizando la memoria de la GPU liberada. La biblioteca SMP ofrece la posibilidad de ejecutar paralelismos de datos fragmentados con PyTorch Fully Sharded Data Parallel (FSDP). 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, tal scale-aware-sharding como se describe en el blog, el escalado cercano al lineal de la formación en modelos gigantes para el FSDP. AWS PyTorch

La biblioteca SMP facilita su uso HYBRID_SHARD _HYBRID_SHARD_ZERO2 en cualquier número configurable de GPU, lo que amplía el PyTorch FSDP nativo que admite la fragmentación en un solo nodo () o en 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 Puedes pasar por el valorShardingStrategy.HYBRID_SHARD. Como alternativa, la biblioteca SMP anula la estrategia del script y la establece ShardingStrategy.HYBRID_SHARD si se especifica un valor igual o superior a 2 en el hybrid_shard_degree parámetro.

En los siguientes fragmentos de código, se muestra cómo añadir el módulo torch.sagemaker.init() de inicialización de SMP al guion de formación y cómo configurar el diccionario de configuración de SMP en formato JSON para el lanzador de tareas de formación, siguiendo el proceso de dos pasos descrito en el presente documento. 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 sobre el parámetro hybrid_shard_degree, consulte Parámetros de configuración de las funciones principales del SMP v2.

Diccionario de configuración SMP

{ "hybrid_shard_degree": 16 }

En el guion 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 = ...