使用 SageMaker Python SDK 进行自定义

SageMaker Python SDK v3.0 引入了一种现代化的模块化 API，用于在 SageMaker 上训练、微调、部署和管理模型。该 SDK 支持多种训练方法，包括持续预训练（CPT）、监督式微调（SFT）、直接偏好优化（DPO）、强化微调（RFT）和多轮强化学习（MTRL）。您可以在 SageMaker Training Jobs 和 SageMaker HyperPod 上运行训练作业。

快速链接

按照以下步骤完成从安装到您的第一个训练作业的过程：

优势

要求

支持的 Python 版本

SageMaker Python SDK 支持 Python 3.10 及更高版本。

安装

要安装 SageMaker Python SDK，请运行以下命令：

pip install "sagemaker>=3.19.0"

支持的模型与技术

该 SDK 支持 Amazon Nova 系列中的以下模型与技术：

多轮强化学习输出

受限模型包（RMP）是一种 SageMaker AI 模型包，它将专有模型构件封装在由平台托管的托管存储中。RMP 让您可以通过 IAM 策略授权和控制这些模型的使用，而无需授予对底层构件的直接访问权限。无法直接下载、导出或查看模型数据。该数据仅可在授权的 AWS 服务中使用。RMP 位于标记为“StorageType: "Restricted"”的模型包组中。

当您在 SageMaker Training Jobs Serverless 中使用多轮强化学习 (MTRL) 训练模型时，输出结果将以 RMP ARN 的形式存储在模型包组中，而非 S3 路径中。这与其他训练方法（如 SFT、DPO 或 RFT）不同，后者的输出结果是模型检查点的 S3 路径。

要使用 MTRL，请使用 MultiTurnRLTrainer 类。在 SageMaker 训练作业（无服务器）上训练时，您可以选择指定 output_model_package_group 来控制输出 RMP 的注册位置。如果省略，SDK 会自动为您创建一个模型包组。有关更多信息和代码示例，请参阅受限模型包。

开始使用

1. 配置基础设施

SDK 支持三种计算平台。将适当的配置传递给训练器的 compute 参数。

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 训练作业（有服务器）

from sagemaker.core.training.configs import TrainingJobCompute

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

SageMaker 训练作业（无服务器）

完全托管，无需计算配置。省略 compute 参数，SDK 将默认使用无服务器模式：

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

2. 训练

使用 SFTTrainer 类开始监督式微调。提供您的模型、计算配置、训练数据集和输出路径。

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",
    compute=compute,
    training_dataset="s3://my-bucket/sft-data.jsonl",
    s3_output_path="s3://my-bucket/output/",
)

job = trainer.train(wait=False)

SDK 还提供 CPTTrainer（用于持续预训练）、DPOTrainer（用于直接偏好优化）、RLVRTrainer（用于强化微调）和 MultiTurnRLTrainer（用于多轮强化学习）。它们均遵循相同的模式：提供模型、计算配置、训练数据集和输出路径。

3. 监控

您可直接通过 SDK 追踪训练进度。使用 stream_logs() 实时流式传输 Amazon CloudWatch 日志，或使用 show_metrics() 在作业完成后绘制训练指标（例如损失和学习率）。

# 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. 评估

使用 BenchMarkEvaluator 类，对照内置基准测试任务评估训练好的模型。支持的基准测试包括 MMLU（大规模多任务语言理解）、BBH（高级推理任务）和 GPQA（研究生级是反谷歌问答）。有关其他评估选项，请参阅评估器。

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. 部署

训练后，将您的自定义模型部署到生产环境中。使用 SageMaker Python SDK，您可以部署到 SageMaker 实时推理端点和 Amazon Bedrock 按需推理。选择最适合您的延迟、吞吐量和成本要求的部署选项。

SageMaker 实时推理端

部署到 SageMaker 实时推理端点，以完全控制实例类型、扩缩策略和端点配置。使用 ModelBuilder 创建和部署 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 按需推理

按需推理提供按使用量付费的定价，无需预置容量。此选项适用于基于 LoRA 的自定义。当您的流量模式可变或不可预测时，请使用按需推理：

from sagemaker.serve import BedrockModelBuilder

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

deployment = builder.deploy()

关键功能

配方覆盖优先级

SageMaker Python SDK 将分层配置系统用于训练配方。启动训练作业时，参数按以下优先顺序（从高到低）进行解析：

例如，要在对所有其他参数使用 Hub 默认值时覆盖最大训练步数和学习率：

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",
    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)

在此示例中，max_epochs 和 optim.lr 通过覆盖显式设置。所有其他训练参数（批次大小、预热步骤、模型并行等）均回退到 nova-textgeneration-lite 模型的 Hub 默认配方。

企业级基础设施支持

SDK 支持多个计算平台，可自动管理基础设施配置、验证和作业编排：

在所有平台上，SDK 均会在提交作业之前验证实例类型、配方配置和数据集格式，从而在工作流早期防止错误发生。

全面评测

基于标准基准评测自定义模型。该 SDK 提供以下评估器：

生产环境部署

使用 SageMaker Python SDK，您可以通过多种部署选项部署自定义模型：

使用 ModelBuilder 或 BedrockModelBuilder 类来部署经过训练的模型。

数据混合

SageMaker Python SDK 提供了 DataMixingConfig 类，用于配置数据混合。

将 DataMixingConfig 与您的训练器配合使用，以指定客户数据的百分比及其在 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",
    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)

了解更多

准备好开始使用 SageMaker Python SDK 自定义 Nova 模型了吗？ 有关详细指南、API 参考及更多示例，请参阅 GitHub 上的 sagemaker-python-sdk。