Personalizar com o SageMaker Python SDK

O SageMaker Python SDK v3 apresenta uma API moderna e modular para treinar, fazer o ajuste fino, implantar e gerenciar modelos no SageMaker. O SDK oferece suporte a vários métodos de treinamento, incluindo pré-treinamento contínuo (CPT), ajuste fino supervisionado (SFT), otimização de preferência direta (DPO), ajuste fino por reforço (RFT) e aprendizado por reforço de vários turnos (MTRL). Você pode executar tarefas de treinamento no SageMaker Training Jobs e no SageMaker HyperPod.

Links rápidos

Siga estas etapas da instalação até a primeira tarefa de treinamento:

Benefícios

Requisitos

Versões compatíveis do Python

O SageMaker Python SDK oferece suporte ao Python 3.10 e versões posteriores.

Instalação

Execute o seguinte comando para instalar o SageMaker Python SDK:

pip install "sagemaker>=3.19.0"

Modelos e técnicas compatíveis

O SDK é compatível com os seguintes modelos e técnicas dentro da família do Amazon Nova:

Resultado de aprendizado por reforço de várias turnos

Um Restricted Model Package (RMP) é um pacote de modelos do SageMaker AI que agrupa artefatos de modelos proprietários em armazenamento de depósito gerenciado pela plataforma. Os RMPs permitem que você autorize e controle o uso desses modelos por meio de políticas do IAM sem conceder acesso direto aos artefatos subjacentes. Os dados do modelo não podem ser baixados, exportados ou visualizados diretamente. Ele só pode ser usado em serviços da AWS autorizados. Os RMPs existem nos Model Package Groups marcados com StorageType: "Restricted".

Quando você treina um modelo usando o aprendizado por reforço de vários turnos (MTRL) no SageMaker Training Jobs Serverless, o resultado é entregue como um RMP ARN dentro de um Model Package Group, em vez de um caminho do S3. Isso difere de outros métodos de treinamento (como SFT, DPO ou RFT) em que a saída é um caminho do S3 até o ponto de verificação do modelo.

Para usar MTRL, use a classe MultiTurnRLTrainer. Ao treinar no SageMaker Training Jobs Serverless, é possível especificar, opcionalmente, um output_model_package_group para controlar onde a saída RMP é registrada. Se omitido, o SDK criará automaticamente um grupo de pacotes de modelos para você. Para obter mais informações e exemplos de código, consulte Restricted Model Packages.

Conceitos básicos

1. Configurar sua infraestrutura

O SDK oferece suporte a três plataformas de computação. Passe a configuração apropriada para o parâmetro compute do seu treinador.

SageMaker HyperPod

from sagemaker.core.training.configs import HyperPodCompute

compute = HyperPodCompute(
    cluster_name="my-hyperpod-cluster",
    instance_type="ml.p5.48xlarge",
    node_count=4,
)

SageMaker Training Jobs (com servidor)

from sagemaker.core.training.configs import TrainingJobCompute

compute = TrainingJobCompute(
    instance_type="ml.p5.48xlarge",
    instance_count=2,
)

SageMaker Training Jobs (sem servidor)

Totalmente gerenciado e sem necessidade de configuração computacional. Omita o parâmetro compute e o SDK usará a tecnologia sem servidor por padrão:

# No compute parameter needed as serverless is the default
trainer = SFTTrainer(
    model="nova-textgeneration-lite-v2",
    training_dataset="s3://my-bucket/sft-data.jsonl",
    s3_output_path="s3://my-bucket/output/",
)

2. Treinamento

Comece o ajuste fino supervisionado com a classe SFTTrainer. Forneça seu modelo, configuração computacional, conjunto de dados de treinamento e caminho de saída.

from sagemaker.train import SFTTrainer
from sagemaker.core.training.configs import HyperPodCompute

compute = HyperPodCompute(
    cluster_name="my-hyperpod-cluster",
    instance_type="ml.p5.48xlarge",
    node_count=4,
)

trainer = SFTTrainer(
    model="nova-textgeneration-lite-v2",
    compute=compute,
    training_dataset="s3://my-bucket/sft-data.jsonl",
    s3_output_path="s3://my-bucket/output/",
)

job = trainer.train(wait=False)

O SDK também fornece CPTTrainer para pré-treinamento contínuo, DPOTrainer para otimização direta de preferências, RLVRTrainer para ajuste fino por reforço e MultiTurnRLTrainer para aprendizado por reforço em vários turnos. Cada um segue o mesmo padrão: forneça um modelo, configuração computacional, conjunto de dados de treinamento e caminho de saída.

3. Monitor

Acompanhe seu progresso no treinamento diretamente no SDK. Use stream_logs() para transmitir os Amazon CloudWatch Logs em tempo real ou show_metrics() para traçar métricas de treinamento, como perda e taxa de aprendizado após a conclusão do trabalho.

# Stream CloudWatch logs in real-time (blocks until job completes)
trainer.stream_logs(poll=5)

# Or stream only the last N lines
trainer.stream_logs(tail_lines=50)

# Plot training metrics (training_loss, lr, reward_score)
df = trainer.show_metrics()

4. Avaliar

Avalie seu modelo treinado em relação às tarefas de benchmark integradas usando a classe BenchMarkEvaluator. Os benchmarks compatíveis incluem MMLU (Massive Multitask Language Understanding), BBH (Advanced Reasoning Tasks) e GPQA (Graduate-Level Google-Proof Q&A). Para outras opções de avaliação, consulte Avaliadores.

from sagemaker.train.evaluate import BenchMarkEvaluator, get_benchmarks

# Get the trained model s3 path from the completed training job
s3_path = job.model_artifacts.s3_model_artifacts

Benchmark = get_benchmarks()

evaluator = BenchMarkEvaluator(
    benchmark=Benchmark.MMLU,
    model=s3_path,
    s3_output_path="s3://my-bucket/eval-output/",
)

execution = evaluator.evaluate()

5. Implantar

Após o treinamento, implante seu modelo personalizado na produção. Com o SageMaker Python SDK, você pode implantar nos endpoints de inferência em tempo real do SageMaker e no Amazon Bedrock On-Demand. Escolha a opção de implantação que melhor se adapta aos seus requisitos de latência, throughput e custo.

Inferência em tempo real do SageMaker

Implante em um endpoint de inferência em tempo real do SageMaker para ter controle total sobre os tipos de instância, as políticas de escalabilidade e a configuração do endpoint. Use ModelBuilder para criar e implantar um endpoint do SageMaker:

from sagemaker.serve import ModelBuilder

# Get the trained model checkpoint path
s3_path = job.model_artifacts.s3_model_artifacts

# Deploy to SageMaker Real-time Inference endpoint
builder = ModelBuilder(
    model=s3_path,
    instance_type="ml.p5.48xlarge",
    env_vars={
        "CONTEXT_LENGTH": "8000",
        "MAX_CONCURRENCY": "2",
    },
)

builder.build().deploy()  # Build the model and deploy to an endpoint

Bedrock On-Demand

A inferência sob demanda fornece preços de pagamento por uso sem capacidade provisionada. Essa opção é aplicável às personalizações baseadas em LoRA. Use o On-Demand quando você tiver padrões de tráfego variáveis ou imprevisíveis:

from sagemaker.serve import BedrockModelBuilder

# Deploy with Bedrock On-Demand
builder = BedrockModelBuilder(
    model=s3_path,
    throughput_type="on-demand",
)

deployment = builder.deploy()

Principais recursos

Precedência de substituição de fórmula

O SageMaker Python SDK usa um sistema de configuração em camadas para fórmulas de treinamento. Quando você inicia uma tarefa de treinamento, os parâmetros são resolvidos na seguinte ordem de precedência (da maior para a menor):

Por exemplo, para substituir as etapas máximas de treinamento e a taxa de aprendizado ao usar os padrões do Hub para todos os outros parâmetros:

from sagemaker.train import SFTTrainer
from sagemaker.core.training.configs import HyperPodCompute

compute = HyperPodCompute(
    cluster_name="my-hyperpod-cluster",
    instance_type="ml.p5.48xlarge",
    node_count=2,
)

trainer = SFTTrainer(
    model="nova-textgeneration-lite-v2",
    training_dataset="s3://my-bucket/sft-data.jsonl",
    s3_output_path="s3://my-bucket/output/",
    base_job_name="my-sft-training-job",
    overrides={
        "training_config.trainer.max_epochs": 1,
        "training_config.model.optim.lr": 1e-5,
    },
)

job = trainer.train(wait=False)

Neste exemplo, max_epochs e optim.lr são definidos explicitamente por meio de substituições. Todos os outros parâmetros de treinamento (tamanho do lote, etapas de aquecimento, modelo paralelo e assim por diante) se enquadram na receita padrão do Hub para o modelo nova-textgeneration-lite-v2.

Suporte à infraestrutura de empresas

O SDK oferece suporte a várias plataformas de computação, gerenciando automaticamente a configuração, a validação e a orquestração de tarefas da infraestrutura:

Em todas as plataformas, o SDK valida tipos de instância, configurações de receitas e formatos de conjunto de dados antes de enviar trabalhos, evitando erros no início do fluxo de trabalho.

Avaliação abrangente

Avalie seus modelos personalizados em relação às avaliações comparativas padrão. O SDK fornece os seguintes avaliadores:

Implantação na produção

Com o SageMaker Python SDK, você pode implantar modelos personalizados usando várias opções de implantação:

Use as classes ModelBuilder ou BedrockModelBuilder para implantar modelos treinados.

Combinação de dados

O SageMaker Python SDK fornece a classe DataMixingConfig para configuração da combinação de dados.

Use DataMixingConfig com seu treinador para especificar a porcentagem de dados do cliente e a distribuição nas categorias de dados da Nova:

from sagemaker.train import SFTTrainer
from sagemaker.train.data_mixing_config import DataMixingConfig
from sagemaker.core.training.configs import HyperPodCompute

data_mixing = DataMixingConfig(
    customer_data_percent=70.0,
    nova_data_percentages={
        "code": 40.0,
        "reasoning": 30.0,
        "instruction-following": 30.0,
    },
)

trainer = SFTTrainer(
    model="nova-textgeneration-lite-v2",
    compute=HyperPodCompute(
        cluster_name="my-cluster",
        instance_type="ml.p5.48xlarge",
        node_count=4,
    ),
    training_dataset="s3://my-bucket/sft-data.jsonl",
    s3_output_path="s3://my-bucket/output/",
    data_mixing_config=data_mixing,
)

job = trainer.train(wait=False)

Saiba mais

Tudo pronto para começar a personalizar os modelos do Nova com o SageMaker Python SDK? Para obter guias detalhados, referências de API e exemplos adicionais, consulte sagemaker-python-sdk no GitHub.