Partición del estado del optimizador

La partición del estado del optimizador es una técnica útil para ahorrar memoria que parte el estado del optimizador (el conjunto de ponderaciones que describe el estado del optimizador) entre grupos de dispositivos paralelos de datos. Puede utilizar la partición del estado del optimizador siempre que utilice un optimizador con estado (como Adam) o un optimizador FP16 (que almacena copias de los parámetros tanto en FP16 como en FP32).

nota

La fragmentación del estado del optimizador está disponible PyTorch en la biblioteca de paralelismo de SageMaker modelos, versión 1.6.0 y versiones posteriores.

Cómo utilizar la partición de estados del optimizador

Puede activar la partición del estado del optimizador estableciendo "shard_optimizer_state": True en la configuración de modelparallel.

Cuando esta característica está activada, la biblioteca divide el conjunto de parámetros del modelo en función del grado de paralelismo de los datos. Los gradientes correspondientes a la división número i se reducen solo en el rango paralelo de datos número i. Al final de la primera llamada a una característica del decorador smp.step, el optimizador encapsulado con smp.DistributedOptimizer redefine sus parámetros para limitarse solo a los parámetros correspondientes a la división del rango paralelo de datos actual. Los parámetros redefinidos se denominan parámetros virtuales y comparten el almacenamiento subyacente con los parámetros originales. Durante la primera llamada a optimizer.step, los estados del optimizador se crean en función de estos parámetros redefinidos, que están partidos debido a la partición original. Tras la actualización del optimizador, la AllGather operación (como parte de la optimizer.step llamada) se ejecuta en los rangos paralelos de datos para lograr estados de parámetros consistentes.

sugerencia

La partición del estado del optimizador puede resultar útil cuando el grado de paralelismo de los datos es superior a 1 y el modelo tiene más de mil millones de parámetros.

El grado de paralelismo de los datos se calcula mediante (processes_per_host * instance_count / pipeline_parallel_degree) y la característica smp.dp_size() gestiona el tamaño en segundo plano.

Configure un estimador SageMaker PyTorch

mpi_options = {
    "enabled" : True,
    "processes_per_host" : 8,               # 8 processes
    "custom_mpi_options" : "--mca btl_vader_single_copy_mechanism none "
}

smp_options = {
    "enabled":True,
    "parameters": {
        "microbatches": 4,
        "pipeline_parallel_degree": 2,    # alias for "partitions"
        "placement_strategy": "cluster",
        "tensor_parallel_degree": 2,      # tp over 2 devices
        "ddp": True,
        "shard_optimizer_state": True
    }
}

Adapta tu guion de entrenamiento PyTorch

Consulta Cómo adaptar tu guion de PyTorch entrenamiento en la sección Paralelismo tensorial combinado con paralelismo de canalización. No es necesario realizar ninguna modificación adicional en el script.