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.

Entrena un PyTorch modelo

En este tema se explica el proceso de entrenamiento de un PyTorch modelo con HyperPod.

En este escenario, entrenemos un PyTorch modelo con la hyp-pytorch-job plantilla, que simplifica la creación de empleo al exponer los parámetros de uso común. Los artefactos del modelo se almacenarán en un bucket de S3 para usarlos posteriormente en la inferencia. Sin embargo, esto es opcional y puede elegir la ubicación de almacenamiento que prefiera.

Creación de un trabajo de entrenamiento

Puede entrenar el modelo mediante la CLI o el SDK de Python.

Uso de la CLI

Cree un trabajo de entrenamiento con el siguiente comando:

hyp create hyp-pytorch-job \
    --version 1.0 \
    --job-name test-pytorch-job \
    --image pytorch/pytorch:latest \
    --command '["python", "train.py"]' \
    --args '["--epochs", "10", "--batch-size", "32"]' \
    --environment '{"PYTORCH_CUDA_ALLOC_CONF": "max_split_size_mb:32"}' \
    --pull-policy "IfNotPresent" \
    --instance-type ml.p4d.24xlarge \
    --tasks-per-node 8 \
    --label-selector '{"accelerator": "nvidia", "network": "efa"}' \
    --deep-health-check-passed-nodes-only true \
    --scheduler-type "kueue" \
    --queue-name "training-queue" \
    --priority "high" \
    --max-retry 3 \
    --volumes '["data-vol", "model-vol", "checkpoint-vol"]' \
    --persistent-volume-claims '["shared-data-pvc", "model-registry-pvc"]' \
    --output-s3-uri s3://my-bucket/model-artifacts

Explicación de los principales parámetros necesarios:

Este comando inicia un trabajo de entrenamiento denominado test-pytorch-job. --output-s3-uri especifica dónde se almacenarán los artefactos del modelo entrenado, por ejemplo, s3://my-bucket/model-artifacts. Tome nota de esta ubicación, pues la necesitará para implementar el modelo personalizado.

Uso del SDK para Python

Use el SDK para el control programático. Cree un script de Python para lanzar el mismo trabajo de entrenamiento.

from sagemaker.hyperpod import HyperPodPytorchJob
from sagemaker.hyperpod.job 
import ReplicaSpec, Template, Spec, Container, Resources, RunPolicy, Metadata

# Define job specifications
nproc_per_node = "1"  # Number of processes per node
replica_specs = 
[
    ReplicaSpec
    (
        name = "pod",  # Replica name
        template = Template
        (
            spec = Spec
            (
                containers =
                [
                    Container
                    (
                        # Container name
                        name="container-name",  
                        
                        # Training image
                        image="448049793756.dkr.ecr.us-west-2.amazonaws.com/ptjob:mnist",  
                        
                        # Always pull image
                        image_pull_policy="Always",  
                        resources=Resources\
                        (
                            # No GPUs requested
                            requests={"nvidia.com/gpu": "0"},  
                            # No GPU limit
                            limits={"nvidia.com/gpu": "0"},   
                        ),
                        # Command to run
                        command=["python", "train.py"],  
                        # Script arguments
                        args=["--epochs", "10", "--batch-size", "32"],  
                    )
                ]
            )
        ),
    )
]
# Keep pods after completion
run_policy = RunPolicy(clean_pod_policy="None")  

# Create and start the PyTorch job
pytorch_job = HyperPodPytorchJob
(
    # Job name
    metadata = Metadata(name="demo"),  
    # Processes per node
    nproc_per_node = nproc_per_node,   
    # Replica specifications
    replica_specs = replica_specs,     
    # Run policy
    run_policy = run_policy,           
    # S3 location for artifacts
    output_s3_uri="s3://my-bucket/model-artifacts"  
)
# Launch the job
pytorch_job.create()  
                

Supervisión del trabajo de entrenamiento

Puede supervisar el progreso de su trabajo con estos comandos:

Uso de la CLI

# Check job status
hyp list hyp-pytorch-job

# Get detailed information
hyp describe hyp-pytorch-job --job-name test-pytorch-job

# View logs
hyp get-logs hyp-pytorch-job \
    --pod-name test-pytorch-job-pod-0 \
    --job-name test-pytorch-job

Nota: El tiempo de entrenamiento varía según la complejidad del modelo y el tipo de instancia. Supervise los registros para hacer un seguimiento del progreso.

Estos comandos le ayudan a verificar el estado del trabajo y a solucionar problemas. Cuando el trabajo se haya completado correctamente, los artefactos del modelo se guardarán en s3://my-bucket/model-artifacts.

Uso del SDK para Python

Incluya el siguiente código en el script de Python:

print("List all pods created for this job:")
print(pytorch_job.list_pods())

print("Check the logs from pod0:")
print(pytorch_job.get_logs_from_pod(pod_name="demo-pod-0"))

print("List all HyperPodPytorchJobs:")
print(HyperPodPytorchJob.list())

print("Describe job:")
print(HyperPodPytorchJob.get(name="demo").model_dump())

pytorch_job.refresh()
print(pytorch_job.status.model_dump())

Siguientes pasos

Tras el entrenamiento, los artefactos del modelo se almacenan en el bucket de S3 que haya especificado (s3://my-bucket/model-artifacts). Puede usar estos artefactos para implementar un modelo. Actualmente, debe administrar la transición del entrenamiento a la inferencia de forma manual. Esto implica: