Uso de los estimadores de marco de PyTorch en el SageMaker Python SDK - Amazon SageMaker AI

Las traducciones son generadas a través de traducción automática. En caso de conflicto entre la traducción y la version original de inglés, prevalecerá la version en inglés.

Uso de los estimadores de marco de PyTorch en el SageMaker Python SDK

Puede iniciar entrenamiento distribuido añadiendo el argumento distribution a los estimadores de marco de SageMaker AI, PyTorch o TensorFlow. Para obtener más información, elija uno de los marcos compatibles con la biblioteca de paralelismo de datos distribuidos (SMDDP) de SageMaker AI entre las siguientes selecciones.

PyTorch

Las siguientes opciones del iniciador están disponibles para iniciar entrenamiento distribuido de PyTorch.

  • pytorchddp: esta opción ejecuta mpirun y configura las variables de entorno necesarias para ejecutar el entrenamiento distribuido de PyTorch en SageMaker AI. Para usar esta opción, pase el siguiente diccionario al parámetro distribution.

    { "pytorchddp": { "enabled": True } }

  • torch_distributed: esta opción ejecuta torchrun y configura las variables de entorno necesarias para ejecutar el entrenamiento distribuido de PyTorch en SageMaker AI. Para usar esta opción, pase el siguiente diccionario al parámetro distribution.

    { "torch_distributed": { "enabled": True } }

  • smdistributed: esta opción también ejecuta mpirun pero con smddprun que configura las variables de entorno necesarias para ejecutar el entrenamiento distribuido de PyTorch en SageMaker AI.

    { "smdistributed": { "dataparallel": { "enabled": True } } }

Si opta por sustituir AllGather de NCCL por AllGather de SMDDP, puede utilizar las tres opciones. Seleccione una opción que se adapte a su caso de uso.

Si opta por sustituir AllReduce de NCCL por AllReduce de SMDDP, debe elegir una de las opciones basadas en mpirun: smdistributed o pytorchddp. También puede agregar opciones de MPI adicionales de la siguiente manera.

{ 
    "pytorchddp": {
        "enabled": True, 
        "custom_mpi_options": "-verbose -x NCCL_DEBUG=VERSION"
    }
}
{ 
    "smdistributed": { 
        "dataparallel": {
            "enabled": True, 
            "custom_mpi_options": "-verbose -x NCCL_DEBUG=VERSION"
        }
    }
}

El siguiente ejemplo de código muestra la estructura básica de un estimador de PyTorch con opciones de entrenamiento distribuido.

from sagemaker.pytorch import PyTorch

pt_estimator = PyTorch(
    base_job_name="training_job_name_prefix",
    source_dir="subdirectory-to-your-code",
    entry_point="adapted-training-script.py",
    role="SageMakerRole",
    py_version="py310",
    framework_version="2.0.1",

    # For running a multi-node distributed training job, specify a value greater than 1
    # Example: 2,3,4,..8
    instance_count=2,

    # Instance types supported by the SageMaker AI data parallel library: 
    # ml.p4d.24xlarge, ml.p4de.24xlarge
    instance_type="ml.p4d.24xlarge",

    # Activate distributed training with SMDDP
    distribution={ "pytorchddp": { "enabled": True } }  # mpirun, activates SMDDP AllReduce OR AllGather
    # distribution={ "torch_distributed": { "enabled": True } }  # torchrun, activates SMDDP AllGather
    # distribution={ "smdistributed": { "dataparallel": { "enabled": True } } }  # mpirun, activates SMDDP AllReduce OR AllGather
)

pt_estimator.fit("s3://bucket/path/to/training/data")
nota

PyTorch Lightning y sus bibliotecas de utilidades, como Lightning Bolts, no vienen preinstaladas en los DLC de PyTorch de SageMaker AI. Cree el siguiente archivo requirements.txt y guárdelo en el directorio de origen donde guarda el script de entrenamiento.

# requirements.txt
pytorch-lightning
lightning-bolts

Por ejemplo, el directorio con estructura de árbol debería tener el siguiente aspecto.

├── pytorch_training_launcher_jupyter_notebook.ipynb
└── sub-folder-for-your-code
    ├──  adapted-training-script.py
    └──  requirements.txt

Para obtener más información sobre cómo especificar el directorio de origen para colocar el archivo requirements.txt junto con su script de entrenamiento y el envío de un trabajo, consulte Uso de bibliotecas de terceros en la documentación del Amazon SageMaker AI Python SDK.

Consideraciones para activar las operaciones colectivas de SMDDP y utilizar las opciones del iniciador de entrenamiento distribuido adecuadas

  • Actualmente, AllReduce de SMDDP y AllGather de SMDDP no son compatibles entre sí.

  • AllReduce de SMDDP se activa de forma predeterminada cuando se utiliza smdistributed o pytorchddp, que son iniciadores basados en mpirun y se utiliza AllGather de NCCL.

  • AllGather de SMDDP se activa de forma predeterminada cuando se utiliza el iniciador de torch_distributed y AllReduce recurre a NCCL.

  • AllGather de SMDDP también se puede activar cuando se utilizan los iniciadores basados en mpirun con una variable de entorno adicional configurada de la siguiente manera.

    export SMDATAPARALLEL_OPTIMIZE_SDP=true
TensorFlow
importante

Se ha dejado de admitir la biblioteca de SMDDP para TensorFlow y ya no está disponible en los DLC de TensorFlow posteriores a v2.11.0. Para buscar DLC de TensorFlow anteriores con la biblioteca de SMDDP instalada, consulte TensorFlow (obsoleto).

from sagemaker.tensorflow import TensorFlow

tf_estimator = TensorFlow(
    base_job_name = "training_job_name_prefix",
    entry_point="adapted-training-script.py",
    role="SageMakerRole",
    framework_version="2.11.0",
    py_version="py38",

    # For running a multi-node distributed training job, specify a value greater than 1
    # Example: 2,3,4,..8
    instance_count=2,

    # Instance types supported by the SageMaker AI data parallel library: 
    # ml.p4d.24xlarge, ml.p3dn.24xlarge, and ml.p3.16xlarge
    instance_type="ml.p3.16xlarge",

    # Training using the SageMaker AI data parallel distributed training strategy
    distribution={ "smdistributed": { "dataparallel": { "enabled": True } } }
)

tf_estimator.fit("s3://bucket/path/to/training/data")