Paralelismo de datos partidos

El paralelismo de datos partidos 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 las GPU de un grupo de datos paralelos.

Al escalar verticalmente su trabajo de entrenamiento a un clúster de GPU grande, puede reducir el consumo de memoria por GPU del modelo partiendo el estado de entrenamiento del modelo en varias GPU. Esto ofrece dos ventajas: puede instalar modelos más grandes, que de otro modo se quedarían sin memoria con el paralelismo de datos estándar, o puede aumentar el tamaño del lote utilizando la memoria de la GPU liberada.

La técnica de paralelismo de datos estándar replica los estados de entrenamiento en las GPU del grupo paralelo de datos y realiza la agregación de gradientes en función de la operación AllReduce. El paralelismo de datos partidos modifica el procedimiento de entrenamiento distribuido de datos paralelos estándar para tener en cuenta la naturaleza partida de los estados del optimizador. Un grupo de rangos en los que se parten los estados del modelo y del optimizador se denomina grupo de partición. La técnica de paralelismo de datos partidos divide los parámetros que pueden entrenarse de un modelo y los gradientes y estados del optimizador correspondientes en las GPU del grupo de partición.

SageMaker logra un paralelismo de datos partidos mediante la implementación de MiCS, que se analiza en la entrada del blog AWS Escalado casi lineal del modelo gigantesco de entrenamiento en AWS. En esta implementación, puede establecer el grado de partición como un parámetro configurable, que debe ser inferior al grado de paralelismo de los datos. Durante cada pasada hacia adelante y hacia atrás, MiCS recombina temporalmente los parámetros del modelo en todas las GPU durante la operación AllGather. Tras pasar cada capa hacia adelante o hacia atrás, MiCS vuelve a partir los parámetros para ahorrar memoria en la GPU. Durante la pasada hacia atrás, MiCS reduce los gradientes y, al mismo tiempo, los parte entre las diferentes GPU a lo largo de la operación ReduceScatter. Por último, MiCS aplica los gradientes locales reducidos y partidos a sus correspondientes particiones de parámetros locales, utilizando las particiones locales de los estados del optimizador. Para reducir la sobrecarga de comunicación, la biblioteca de paralelismo de modelos de SageMaker busca previamente las siguientes capas en la pasada hacia adelante o hacia atrás y superpone la comunicación de red con el cálculo.

El estado de entrenamiento del modelo se replica en todos los grupos de partición. Esto significa que antes de aplicar los gradientes a los parámetros, la operación AllReduce debe realizarse en todos los grupos de partición, además de la operación ReduceScatter que se lleva a cabo dentro del grupo de partición.

En efecto, el paralelismo de datos partidos supone un equilibrio entre la sobrecarga de comunicación y la eficiencia de la memoria de la GPU. El uso del paralelismo de datos partidos aumenta el coste de la comunicación, pero el espacio de memoria por GPU (excluido el uso de memoria debido a las activaciones) se divide por el grado de paralelismo de los datos partidos, por lo que pueden caber modelos más grandes en el clúster de la GPU.

Seleccionar el grado de paralelismo de los datos partidos

Al seleccionar un valor para el grado de paralelismo de los datos partidos, el valor debe dividir uniformemente el grado de paralelismo de los datos. Por ejemplo, para un trabajo de paralelismo de datos de 8 vías, seleccione 2, 4 u 8 como grado de paralelismo de datos partidos. Al elegir el grado de paralelismo de los datos partidos, le recomendamos que comience con un número pequeño y vaya aumentado poco a poco hasta que el modelo quepa en la memoria y tenga el tamaño de lote deseado.

Seleccionar el tamaño del lote

Tras configurar el paralelismo de datos partidos, asegúrese de encontrar la configuración de entrenamiento más óptima que pueda ejecutarse correctamente en el clúster de la GPU. Para entrenar modelos de lenguaje grandes (LLM), empiece con el tamaño de lote 1 y vaya aumentando poco a poco hasta que llegue al punto en el que aparezca el error de falto de memoria (OOM). Si encuentra el error OOM incluso con el tamaño de lote más pequeño, aplique un mayor grado de paralelismo de datos partidos o una combinación de paralelismo de datos partidos y paralelismo de tensores.

Cómo aplicar el paralelismo de datos partidos a su trabajo de entrenamiento

Para empezar con el paralelismo de datos partidos, aplique las modificaciones necesarias al script de su entrenamiento y configure el estimador PyTorch de SageMaker con los parámetros específicos del paralelismo de datos partidos. Considere también tomar valores de referencia y cuadernos de ejemplo como punto de partida.

Adaptar su script de entrenamiento de PyTorch

Siga las instrucciones del Paso 1: modifique un script de entrenamiento de PyTorch para encapsular los objetos del modelo y del optimizador con los contenedores smdistributed.modelparallel.torch de los módulos torch.nn.parallel y torch.distributed.

(Opcional) Modificación adicional para registrar los parámetros externos del modelo

Si el modelo está creado con torch.nn.Module y usa parámetros que no están definidos en la clase de módulo, debe registrarlos manualmente en el módulo para que SMP recopile todos los parámetros. Para registrar los parámetros en un módulo, utilice smp.register_parameter(module, parameter).

class Module(torch.nn.Module):
    def __init__(self, *args):
        super().__init__(self, *args)
        self.layer1 = Layer1()
        self.layer2 = Layer2()
        smp.register_parameter(self, self.layer1.weight)

    def forward(self, input):
        x = self.layer1(input)
        # self.layer1.weight is required by self.layer2.forward
        y = self.layer2(x, self.layer1.weight)
        return y

Configurar el estimador PyTorch de SageMaker

Al configurar un estimador PyTorch de SageMaker en Paso 2: lance un trabajo de entrenamiento con el SageMaker Python SDK, agregue los parámetros para el paralelismo de datos partidos.

Para activar el paralelismo de datos partidos, agregue el parámetro sharded_data_parallel_degree al estimador PyTorch de SageMaker. Este parámetro especifica el número de GPU en las que se parte el estado de entrenamiento. El valor de sharded_data_parallel_degree debe ser un número entero entre uno y el grado de paralelismo de los datos y debe dividir el grado de paralelismo de los datos de manera uniforme. Tenga en cuenta que la biblioteca detecta automáticamente la cantidad de GPU y, por lo tanto, el grado de paralelismo de los datos. Los siguientes parámetros adicionales están disponibles para configurar el paralelismo de datos partidos.

El siguiente código muestra un ejemplo de cómo configurar el paralelismo de datos partidos.

import sagemaker
from sagemaker.pytorch import PyTorch

smp_options = {
    "enabled": True,
    "parameters": {
        # "pipeline_parallel_degree": 1,    # Optional, default is 1
        # "tensor_parallel_degree": 1,      # Optional, default is 1
        "ddp": True,
        # parameters for sharded data parallelism
        "sharded_data_parallel_degree": 2,              # Add this to activate sharded data parallelism
        "sdp_reduce_bucket_size": int(5e8),             # Optional
        "sdp_param_persistence_threshold": int(1e6),    # Optional
        "sdp_max_live_parameters": int(1e9),            # Optional
        "sdp_hierarchical_allgather": True,             # Optional
        "sdp_gradient_clipping": 1.0                    # Optional
    }
}

mpi_options = {
    "enabled" : True,                      # Required
    "processes_per_host" : 8               # Required
}

smp_estimator = PyTorch(
    entry_point="your_training_script.py", # Specify your train script
    role=sagemaker.get_execution_role(),
    instance_count=1,
    instance_type='ml.p3.16xlarge',
    framework_version='1.13.1',
    py_version='py3',
    distribution={
        "smdistributed": {"modelparallel": smp_options},
        "mpi": mpi_options
    },
    base_job_name="sharded-data-parallel-job"
)

smp_estimator.fit('s3://my_bucket/my_training_data/')

Configuraciones de referencia

El equipo de entrenamiento distribuido de SageMaker proporciona las siguientes configuraciones de referencia que puede utilizar como punto de partida. Puede extrapolar las siguientes configuraciones para experimentar y estimar el uso de memoria de la GPU para configurar su modelo.

Paralelismo de datos partidos con colectivos SMDDP

Por ejemplo, si aumenta la longitud de secuencia de un modelo de 20 000 millones de parámetros o aumenta el tamaño del modelo a 65 000 millones de parámetros, primero debe intentar reducir el tamaño del lote. Si el modelo sigue sin ajustarse al tamaño de lote más pequeño (el tamaño de lote de 1), intente aumentar el grado de paralelismo de modelos.

Paralelismo de datos partidos con paralelismo de tensores y colectivos NCCL

El uso combinado del paralelismo de datos partidos y el paralelismo de tensores resulta útil cuando se quiere encajar un modelo de lenguaje grande (LLM) en un clúster a gran escala y, al mismo tiempo, utilizar datos de texto con una longitud de secuencia más larga, lo que lleva a utilizar un tamaño de lote más pequeño y, por lo tanto, gestionar el uso de la memoria de la GPU para entrenar los LLM con secuencias de texto más largas. Para obtener más información, consulte Paralelismo de datos partidos con paralelismo de tensores.

Para ver casos prácticos, pruebas comparativas y más ejemplos de configuración, consulte la entrada del blog Nuevas mejoras de rendimiento en la biblioteca paralela de modelos de Amazon SageMaker.

Paralelismo de datos partidos con colectivos SMDDP

La biblioteca de paralelismo de datos de SageMaker ofrece primitivas de comunicación colectiva (colectivos SMDDP) optimizadas para la infraestructura AWS. Logra la optimización al adoptar un patrón de comunicación de todo tipo mediante el uso del Elastic Fabric Adapter (EFA), lo que da como resultado colectivos de alto rendimiento y menos sensibles a la latencia, descargando el procesamiento relacionado con la comunicación a la CPU y liberando los ciclos de la GPU para el cálculo. En clústeres grandes, los colectivos SMDDP pueden ofrecer mejoras de hasta un 40 % en el rendimiento de los entrenamientos distribuidos en comparación con los NCCL. Para ver casos prácticos y resultados comparativos, consulte el blog Nuevas mejoras de rendimiento en la biblioteca de paralelismo de modelos de Amazon SageMaker.

En el paralelismo de datos partidos, que es una técnica de uso común en el entrenamiento distribuido a gran escala, el colectivo AllGather se utiliza para reconstituir los parámetros de la capa partida para los cálculos de la pasada hacia adelante y hacia atrás, en paralelo con el cálculo de la GPU. En el caso de los modelos de gran tamaño, es fundamental realizar la operación AllGather de forma eficiente para evitar problemas con la GPU y reducir la velocidad de entrenamiento. Cuando se activa el paralelismo de datos partidos, los colectivos SMDDP se agrupan en estos colectivos AllGather fundamentales para el rendimiento, lo que mejora el rendimiento del entrenamiento.

Entrenar con colectivos SMDDP

Cuando su trabajo de entrenamiento tenga activado el paralelismo de datos partidos y cumpla con Supported configurations, los colectivos SMDDP se activarán automáticamente. De manera interna, los colectivos SMDDP optimizan el colectivo AllGatherpara rendir en la infraestructura AWS y recurren a los NCCL para el resto de colectivos. Además, en configuraciones no compatibles, todos los colectivos, incluidos AllGather, utilizan automáticamente el backend de los NCCL.

Desde la versión 1.13.0 de la biblioteca de paralelismo de modelos de SageMaker, el parámetro "ddp_dist_backend" se añade a las opciones modelparallel. El valor predeterminado de este parámetro de configuración es "auto", que utiliza los colectivos SMDDP siempre que sea posible y, de lo contrario, recurre a los NCCL. Para forzar a la biblioteca a utilizar siempre los NCCL, especifique "nccl" en el parámetro de configuración "ddp_dist_backend".

El siguiente ejemplo de código muestra cómo configurar un estimador PyTorch utilizando el paralelismo de datos partidos con el parámetro "ddp_dist_backend", que está establecido de forma predeterminada en "auto" y, por lo tanto, su adición es opcional.

import sagemaker
from sagemaker.pytorch import PyTorch

smp_options = {
    "enabled":True,
    "parameters": {                        
        "partitions": 1,
        "ddp": True,
        "sharded_data_parallel_degree": 64
        "bf16": True,
        "ddp_dist_backend": "auto"  # Specify "nccl" to force to use NCCL.
    }
}

mpi_options = {
    "enabled" : True,                      # Required
    "processes_per_host" : 8               # Required
}

smd_mp_estimator = PyTorch(
    entry_point="your_training_script.py", # Specify your train script
    source_dir="location_to_your_script",
    role=sagemaker.get_execution_role(),
    instance_count=8,
    instance_type='ml.p4d.24xlarge',
    framework_version='1.13.1',
    py_version='py3',
    distribution={
        "smdistributed": {"modelparallel": smp_options},
        "mpi": mpi_options
    },
    base_job_name="sharded-data-parallel-demo",
)

smd_mp_estimator.fit('s3://my_bucket/my_training_data/')

Configuraciones admitidas

La operación AllGather con los colectivos SMDDP se activa en los trabajos de entrenamiento cuando se cumplen todos los requisitos de configuración siguientes.

Ajuste del rendimiento y de la memoria

Los colectivos SMDDP utilizan memoria de GPU adicional. Existen dos variables de entorno para configurar el uso de la memoria de la GPU en función de los diferentes casos de uso de entrenamiento de modelos.

Guía de ajuste de las variables de tamaño del búfer

Los valores predeterminados de las variables de entorno deberían funcionar bien en la mayoría de los casos de uso. Recomendamos ajustar estas variables solo si el entrenamiento se produce con el error de falta de memoria (OOM).

En la siguiente lista, se describen algunos consejos de ajuste para reducir el consumo de memoria de la GPU de los colectivos SMDDP y, al mismo tiempo, conservar la ganancia de rendimiento que generan.

Es posible que algunos colectivos recurran a los NCCL; por lo tanto, es posible que no consiga el rendimiento de los colectivos SMDDP optimizados. Si hay memoria de GPU adicional disponible para su uso, puede considerar la posibilidad de aumentar los valores de SMDDP_AG_SCRATCH_BUFFER_SIZE_BYTES y SMDDP_AG_SORT_BUFFER_SIZE_BYTES para aprovechar la ganancia de rendimiento.

El siguiente código muestra cómo configurar las variables de entorno agregándolas a mpi_options en el parámetro de distribución del estimador PyTorch.

import sagemaker
from sagemaker.pytorch import PyTorch

smp_options = {
    .... # All modelparallel configuration options go here
}

mpi_options = {
    "enabled" : True,                      # Required
    "processes_per_host" : 8               # Required
}

# Use the following two lines to tune values of the environment variables for buffer
mpioptions += " -x SMDDP_AG_SCRATCH_BUFFER_SIZE_BYTES=8192" 
mpioptions += " -x SMDDP_AG_SORT_BUFFER_SIZE_BYTES=8192"

smd_mp_estimator = PyTorch(
    entry_point="your_training_script.py", # Specify your train script
    source_dir="location_to_your_script",
    role=sagemaker.get_execution_role(),
    instance_count=8,
    instance_type='ml.p4d.24xlarge',
    framework_version='1.13.1',
    py_version='py3',
    distribution={
        "smdistributed": {"modelparallel": smp_options},
        "mpi": mpi_options
    },
    base_job_name="sharded-data-parallel-demo-with-tuning",
)

smd_mp_estimator.fit('s3://my_bucket/my_training_data/')

Entrenamiento de precisión mixto con paralelismo de datos partidos

Para ahorrar aún más memoria en la GPU con números de coma flotante de precisión media y paralelismo de datos partidos, puede activar el formato de coma flotante de 16 bits (FP16) o el formato de punto flotante Brain (BF16) añadiendo un parámetro adicional a la configuración de entrenamiento distribuido.

Para el entrenamiento de FP16 con paralelismo de datos partidos

Para ejecutar el entrenamiento de FP16 con paralelismo de datos partidos, agregue "fp16": True" al diccionario de configuración smp_options. En el script de su entrenamiento, puede elegir entre las opciones de escalado de pérdidas estáticas y dinámicas a través del módulo smp.DistributedOptimizer. Para obtener más información, consulte Entrenamiento con el FP16 con paralelismo de modelos.

smp_options = {
    "enabled": True,
    "parameters": {
        "ddp": True,
        "sharded_data_parallel_degree": 2,
        "fp16": True
    }
}

Para el entrenamiento del BF16 con paralelismo de datos partidos

La característica de paralelismo de datos partidos de SageMaker permite el entrenamiento en el tipo de datos BF16. El tipo de datos BF16 usa 8 bits para representar el exponente de un número de coma flotante, mientras que el tipo de datos FP16 usa 5 bits. Preservar los 8 bits del exponente permite mantener la misma representación del exponente de un número de coma flotante de precisión única (FP32) de 32 bits. Esto hace que la conversión entre el FP32 y el BF16 sea más sencilla y mucho menos propensa a provocar problemas de desbordamiento o subdesbordamiento que suelen surgir en el entrenamiento con el FP16, especialmente cuando se entrenan modelos más grandes. Si bien ambos tipos de datos utilizan 16 bits en total, este aumento del rango de representación del exponente en el formato BF16 se produce a expensas de una menor precisión. Para el entrenamiento de modelos grandes, esta precisión reducida suele considerarse una compensación aceptable entre el alcance y la estabilidad del entrenamiento.

Para ejecutar el entrenamiento de BF16 con paralelismo de datos partidos, agregue "bf16": True al diccionario de configuración smp_options.

smp_options = {
    "enabled": True,
    "parameters": {
        "ddp": True,
        "sharded_data_parallel_degree": 2,
        "bf16": True
    }
}

Paralelismo de datos partidos con paralelismo de tensores

Si utiliza el paralelismo de datos partidos y también necesita reducir el tamaño del lote global, considere la posibilidad de utilizar el paralelismo de tensores con el paralelismo de datos partidos. Al entrenar un modelo grande con paralelismo de datos partidos en un clúster de procesamiento muy grande (normalmente 128 nodos o más), incluso un tamaño de lote pequeño por GPU da como resultado un tamaño de lote global muy grande. Puede provocar problemas de convergencia o problemas de bajo rendimiento computacional. En ocasiones, no es posible reducir el tamaño del lote por GPU solo con el paralelismo de datos partidos si un solo lote ya es grande y no puede reducirse más. En estos casos, el uso del paralelismo de datos partidos en combinación con el paralelismo de tensores ayuda a reducir el tamaño del lote global.

La elección de los grados paralelos y tensoriales paralelos óptimos para los datos partidos depende de la escala del modelo, el tipo de instancia y el tamaño del lote global que sea razonable para que el modelo converja. Le recomendamos que comience desde un grado de paralelismo de tensor bajo para ajustar el tamaño del lote global al clúster de cómputo a fin de resolver los errores de falta de memoria de CUDA y lograr el mejor rendimiento. Consulte los dos casos de ejemplo siguientes para saber cómo la combinación de paralelismo de tensores y paralelismo de datos partidos le ayuda a ajustar el tamaño del lote global al agrupar las GPU para el paralelismo del modelo, lo que resulta en un menor número de réplicas de modelos y un tamaño de lote global más pequeño.

Ejemplo 1

Supongamos que queremos entrenar un modelo en un clúster de 1536 GPU (192 nodos con 8 GPU en cada uno), estableciendo el grado de paralelismo de los datos partidos en 32 (sharded_data_parallel_degree=32) y el tamaño del lote por GPU en 1, donde cada lote tiene una longitud de secuencia de 4096 fichas. En este caso, hay 1536 réplicas de modelos, el tamaño del lote global pasa a ser 1536 y cada lote global contiene unos 6 millones de tokens.

(1536 GPUs) * (1 batch per GPU) = (1536 global batches)
(1536 batches) * (4096 tokens per batch) = (6,291,456 tokens)

Si se le añade un paralelismo de tensores, se puede reducir el tamaño del lote global. Un ejemplo de configuración puede ser establecer el grado de paralelismo del tensor en 8 y el tamaño del lote por GPU en 4. Esto forma 192 grupos de tensores paralelos o 192 réplicas de modelos, donde cada réplica de modelo se distribuye en 8 GPU. El tamaño del lote de 4 es la cantidad de datos de entrenamiento por iteración y por grupo de tensores paralelos; es decir, cada réplica del modelo consume 4 lotes por iteración. En este caso, el tamaño del lote global pasa a ser 768 y cada lote global contiene unos 3 millones de tokens. Por lo tanto, el tamaño del lote global se reduce a la mitad en comparación con el caso anterior solo con el paralelismo de datos partidos.

(1536 GPUs) / (8 tensor parallel degree) = (192 tensor parallelism groups)
(192 tensor parallelism groups) * (4 batches per tensor parallelism group) = (768 global batches)
(768 batches) * (4096 tokens per batch) = (3,145,728 tokens)

Ejemplo 2

Cuando se activan tanto el paralelismo de datos partidos como el paralelismo de tensores, la biblioteca primero aplica el paralelismo de tensores y parte el modelo en esta dimensión. Para cada rango paralelo tensorial, el paralelismo de datos se aplica según sharded_data_parallel_degree.

Por ejemplo, supongamos que queremos configurar 32 GPU con un grado paralelo tensorial de 4 (formando grupos de 4 GPU), un grado paralelo de datos partidos de 4 y terminando con un grado de replicación de 2. La asignación crea ocho grupos de GPU en función del grado paralelo del tensor de la siguiente manera: (0,1,2,3), (4,5,6,7), (8,9,10,11), (12,13,14,15), (16,17,18,19), (20,21,22,23), (24,25,26,27), (28,29,30,31). Es decir, cuatro GPU forman un grupo de tensores paralelos. En este caso, el grupo paralelo de datos reducido para las GPU de rango 0 de los grupos paralelos de tensores sería (0,4,8,12,16,20,24,28). El grupo paralelo de datos reducido se parte en función del grado de paralelo de datos partidos de 4, lo que da como resultado dos grupos de replicación para el paralelismo de datos. Las GPU (0,4,8,12) forman un grupo de partición, que en conjunto contiene una copia completa de todos los parámetros del rango paralelo del tensor 0, y las GPU (16,20,24,28) forman otro grupo de este tipo. Otros rangos paralelos tensoriales también tienen grupos de partición y replicación similares.

Figura 1: grupos de paralelismo de tensores para (nodos, grado de paralelismo de datos particionados, grado de paralelismo de tensores) = (4, 4, 4), donde cada rectángulo representa una GPU con índices de 0 a 31. Las GPU forman grupos de paralelismo de tensores de TPG₀ a TPG₇. Los grupos de replicación son ({TPG₀, TPG₄}, {TPG₁, TPG₅}, {TPG₂, TPG₆} y {TPG₃, TPG₇}); cada par de grupos de replicación comparte el mismo color pero se rellena de forma diferente.

Figura 2: grupos de paralelismo de datos particionados para (nodos, grado de paralelismo de datos particionados, grado de paralelismo de tensores) = (4, 4, 4), donde cada rectángulo representa una GPU con índices de 0 a 31. Las GPU forman grupos de paralelismo de datos particionados de SDPG₀ a SDPG₇. Los grupos de replicación son ({SDPG₀, SDPG₄}, {SDPG₁, SDPG₅}, {SDPG₂, SDPG₆} y {SDPG₃, SDPG₇}); cada par de grupos de replicación comparte el mismo color pero se rellena de forma diferente.

Cómo activar el paralelismo de datos partidos con el paralelismo de tensores

Para utilizar el paralelismo de datos partidos con el paralelismo de tensores, debe establecer tanto sharded_data_parallel_degree como tensor_parallel_degree en la configuración para distribution al crear un objeto de la clase de estimador PyTorch de SageMaker.

También debe activar prescaled_batch. Esto significa que, en lugar de que cada GPU lea su propio lote de datos, cada grupo paralelo de tensores leerá colectivamente un lote combinado del tamaño de lote elegido. Efectivamente, en lugar de dividir el conjunto de datos en partes iguales al número de GPU (o tamaño paralelo de datos, smp.dp_size()), lo divide en partes iguales al número de GPU dividido por tensor_parallel_degree (también denominado tamaño paralelo de datos reducido, smp.rdp_size()). Para obtener más información sobre los lotes preescalados, consulte Lote preescalado en la documentación del SageMaker Python SDK. Consulte también el ejemplo de script de entrenamiento train_gpt_simple.pypara GPT-2 en el repositorio GitHub de ejemplos de SageMaker.

El siguiente fragmento de código muestra un ejemplo de creación de un objeto estimador de PyTorch basado en el escenario mencionado anteriormente en Ejemplo 2.

mpi_options = "-verbose --mca orte_base_help_aggregate 0 "
smp_parameters = {
    "ddp": True,
    "fp16": True,
    "prescaled_batch": True,
    "sharded_data_parallel_degree": 4,
    "tensor_parallel_degree": 4
}

pytorch_estimator = PyTorch(
    entry_point="your_training_script.py",
    role=role,
    instance_type="ml.p4d.24xlarge",
    volume_size=200,
    instance_count=4,
    sagemaker_session=sagemaker_session,
    py_version="py3",
    framework_version="1.13.1",
    distribution={
        "smdistributed": {
            "modelparallel": {
                "enabled": True, 
                "parameters": smp_parameters,
            }
        },
        "mpi": {
            "enabled": True,
            "processes_per_host": 8,
            "custom_mpi_options": mpi_options,
        },
    },
    source_dir="source_directory_of_your_code",
    output_path=s3_output_location
)

Consejos y consideraciones a tener en cuenta para utilizar el paralelismo de datos partidos

Tenga en cuenta lo siguiente a la hora de utilizar el paralelismo de datos partidos de la biblioteca de modelos de SageMaker.