Treine um PyTorch modelo

Este tópico explica o processo de treinamento de um PyTorch modelo usando HyperPod o.

Nesse cenário, vamos treinar um PyTorch modelo usando o hyp-pytorch-job modelo, o que simplifica a criação de trabalhos ao expor parâmetros comumente usados. Os artefatos do modelo serão armazenados em um bucket do S3 para uso posterior na inferência. No entanto, isso é opcional e você pode escolher seu local de armazenamento preferido.

Crie um trabalho de treinamento

Você pode treinar o modelo usando a CLI ou o SDK do Python.

Uso da CLI

Crie um trabalho de treinamento com o seguinte 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

Principais parâmetros necessários explicados:

• --job-name: Identificador exclusivo para seu trabalho de treinamento

• --image: imagem do Docker contendo seu ambiente de treinamento

Esse comando inicia um trabalho de treinamento chamadotest-pytorch-job. --output-s3-uriEspecifica onde os artefatos do modelo treinado serão armazenados, por exemplo,. s3://my-bucket/model-artifacts Observe esse local, pois você precisará dele para implantar o modelo personalizado.

Utilizar o Python SDK

Para controle programático, use o SDK. Crie um script Python para iniciar o mesmo trabalho de treinamento.

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()  
                

Monitore seu trabalho de treinamento

Monitore o progresso do seu trabalho com estes comandos:

Uso da 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

Observação: o tempo de treinamento varia de acordo com a complexidade do modelo e o tipo de instância. Monitore os registros para acompanhar o progresso.

Esses comandos ajudam você a verificar o status do trabalho e solucionar problemas. Depois que o trabalho for concluído com êxito, os artefatos do modelo serão salvos em. s3://my-bucket/model-artifacts

Utilizar o Python SDK

Adicione o código a seguir ao seu script 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())

Próximas etapas

Após o treinamento, os artefatos do modelo são armazenados no bucket do S3 que você especificou ()s3://my-bucket/model-artifacts. Você pode usar esses artefatos para implantar um modelo. Atualmente, você deve gerenciar manualmente a transição do treinamento para a inferência. Isso envolve:

• Localizando artefatos: verifique o bucket do S3 (s3://my-bucket/model-artifacts) para confirmar se os arquivos de modelo treinados estão presentes.

• Registrando o caminho: observe o caminho S3 exato (por exemplo,s3://my-bucket/model-artifacts/test-pytorch-job/model.tar.gz) para uso na configuração de inferência.

• Referência na implantação: forneça esse caminho do S3 ao configurar o endpoint personalizado para garantir que o modelo correto seja carregado.