Amazon Nova 蒸馏 - Amazon Nova
关键组件使用案例先决条件设置数据增强启动训练任务CloudWatch Logs成功的训练验证增强的数据质量

Amazon Nova 蒸馏

本快速入门指南可帮助您开始基于 SageMaker AI 使用监督式微调（SFT）技术进行 Amazon Nova 模型蒸馏。

模型蒸馏是 SageMaker AI 中的一种方法，用于将知识从大型复杂模型迁移到小型高效模型。对于 Amazon Nova 模型，大型“教师式”模型（例如 Amazon Nova Pro）将其功能传递给小型“学生式”模型（例如 Amazon Nova Lite 或 Amazon Nova Micro）。这将创建一个自定义的模型，该模型在使用更少资源的同时保持高性能。

关键组件

蒸馏过程主要涉及两种类型的模型：

教师式模型，作为知识来源，包括：

  • Amazon Nova Pro（amazon.nova-pro-v1:0）

  • （amazon.nova-premier-v1:0）

学生式模型，接收和应用知识：

  • Amazon Nova Lite（amazon.nova-lite-v1:0:300k）

  • Amazon Nova Micro（amazon.nova-micro-v1:0:128k）

  • Amazon Nova Pro（amazon.nova-pro-v1:0:300k）：仅当作为教师式模型使用时可用

使用案例

模式蒸馏对于以下情况尤为有益：

  • 应用程序对延迟、成本和准确性有着严格的要求。

  • 针对特定任务需要使用自定义模型，但缺少质量足够高且带标注的训练数据。

  • 希望性能可以匹敌高级模型，同时维持较小模型的效率。

先决条件

  • 可以访问 Amazon Nova 模型和相应服务配额（最少 6 个 P5 和 1 个 R5 实例）的 AWS 账户。

  • 具有 SageMaker 训练作业权限的 IAM 角色。

  • 存储训练数据和输出的 Amazon S3 存储桶。

设置数据增强

数据增强阶段使用 SageMaker 训练作业，通过教师式模型生成高质量的训练数据。本节详细介绍了设置过程和要求。

IAM 角色

要创建 IAM 角色和附加策略，请参阅创建角色和附加策略（控制台）。如果您使用 AWS CLI，请按照 create-role 和 attach-role-policy 中的说明进行操作。有关更多信息，请参阅《SageMaker AI 开发人员指南》中的如何使用 SageMaker AI 执行角色。

下面是一些示例命令供您参考。

创建 SageMaker AI 执行角色

该角色是通过信任关系创建的，支持 SageMaker AI、Amazon Bedrock 担任此角色。这将使这些服务能够在执行批量推理作业时代表您执行操作。

aws iam create-role \
 --role-name NovaCustomizationRole \
 --assume-role-policy-document '{
 "Version": "2012-10-17",
 "Statement": [
 {
 "Effect": "Allow",
 "Principal": {
 "Service": ["sagemaker.amazonaws.com",
            "bedrock.amazonaws.com"]
 },
 "Action": "sts:AssumeRole"
 }
 ]
}'

附加必要的策略

# Attach AmazonSageMakerFullAccess
 aws iam attach-role-policy \
 --role-name NovaCustomizationRole \
 --policy-arn arn:aws:iam::aws:policy/AmazonSageMakerFullAccess

# Attach AmazonBedrockFullAccess
 aws iam attach-role-policy \
 --role-name NovaCustomizationRole \
 --policy-arn arn:aws:iam::aws:policy/AmazonBedrockFullAccess

 # Attach S3 access policy
 aws iam attach-role-policy \
 --role-name NovaCustomizationRole \
 --policy-arn arn:aws:iam::aws:policy/AmazonS3FullAccess
 
# Attach AmazonEC2FullAccess
 aws iam attach-role-policy \
 --role-name NovaCustomizationRole \
 --policy-arn arn:aws:iam::aws:policy/AmazonEC2FullAccess
 
# Attach AmazonEC2ContainerRegistryFullAccess
 aws iam attach-role-policy \
 --role-name NovaCustomizationRole \
 --policy-arn arn:aws:iam::aws:policy/AmazonEC2ContainerRegistryFullAccess
 
# Attach AmazonEC2ContainerRegistryFullAccess
 aws iam attach-role-policy \
 --role-name NovaCustomizationRole \
 --policy-arn arn:aws:iam::aws:policy/CloudWatchLogsFullAccess

将下面的内联策略附加到蒸馏容器所需的客户执行角色。

  • AWS KMS 权限：支持角色与 AWS 密钥管理服务进行交互，这是访问加密资源或管理加密密钥所必需的。

  • IAM:PassRole：当一个服务需要将此角色传递给另一个服务时，通常需要该权限，这是 AWS 服务集成中的常见模式。

aws iam put-role-policy \
 --role-name NovaCustomizationRole \
 --policy-name Distillation-Additional-Permissions\
 --policy-document '{
 "Version": "2012-10-17",
 "Statement": [
        {
            "Effect": "Allow",
            "Action": [
                "kms:*"
            ],
            "Resource": "*"
        },
        {
            "Effect": "Allow",
            "Action": [
                "iam:PassRole"
            ],
            "Resource": "*"
        }
    ]
}

Amazon VPC 配置

要使用 AWS 管理控制台为 SageMaker 训练作业创建 Amazon VPC 配置，请按照为 SageMaker AI 训练配置私有 VPC（控制台）中的说明进行操作。

创建新的 Amazon VPC

Name: Distillation-VPC
IPv4 CIDR: 10.0.0.0/16 (or your preferred range)
Availability Zones: 2
Public Subnets: 2
Private Subnets: 2
NAT Gateways: 1 (in one AZ)

创建安全组

Name: Distillation-SG
Description: Security group for data distillation jobs
Inbound Rules: Allow all traffic from self
Outbound Rules: Allow all traffic (0.0.0.0/0)

针对以下服务创建 VPC 端点

com.amazonaws.[region].s3
com.amazonaws.[region].sagemaker.api
com.amazonaws.[region].sagemaker.runtime
com.amazonaws.[region].bedrock.api
com.amazonaws.[region].bedrock.runtime
com.amazonaws.[region].sts
com.amazonaws.[region].logs
com.amazonaws.[region].ecr.api
com.amazonaws.[region].ecr.dkr

对于每个端点：

  • 选择您的 Distillation-VPC

  • 选择私有子网

  • 选择 Distillation-SG 安全组

AWS KMS 密钥

使用 Amazon Bedrock 批量推理时，需要 AWS KMS 密钥来确保数据安全和合规性。Amazon Bedrock 批量推理作业要求使用 AWS KMS 密钥对输入和输出 Amazon S3 存储桶进行加密，以确保静态数据保护。

通过 AWS CLI 使用该命令创建 KMS 密钥：

# Create KMS key
aws kms create-key \
  --description "KMS key for Nova distillation"

该命令将输出包括 ARN 在内的关键信息。输出示例：

{
    "KeyMetadata": {
        "KeyId": "1234abcd-12ab-34cd-56ef-1234567890ab",
        "Arn": "arn:aws:kms:us-east-1:111122223333:key/1234abcd-12ab-34cd-56ef-1234567890ab"
    }
}
注意

保存输出中的 KMS 密钥 ARN，因为在下一节中创建 Amazon S3 存储桶时需要使用它。

Amazon S3 存储桶

您需要两种类型的 Amazon S3 存储。客户托管 Amazon S3 存储桶用于存储您的输入数据和输出 manifest.json 文件。您可以创建和管理此存储桶，并且可以将单个存储桶同时用于输入和输出。此存储桶必须配置 KMS 加密，因为它将存储敏感输出数据，并且将由 Amazon Bedrock 批量推理作业使用。Amazon Bedrock 需要使用 KMS 加密的存储桶来处理批量推理任务。

服务托管的 Amazon S3 存储桶用于存储模型权重。服务托管的 Amazon S3 存储桶是在您的第一次训练作业期间自动创建的。它具有受限的访问控制，只能通过清单文件访问特定路径。

要在特定 AWS 区域中创建存储桶，请使用 create-bucket CLI 命令。

创建使用 AWS KMS 加密的 Amazon S3 存储桶的示例命令。将 {kms_key_arn} 替换为您的 AWS KMS 密钥 ARN。您需要先创建 AWS KMS 密钥（如果尚未执行此操作）。

aws s3api create-bucket \
    "Rules": [
        {
            "ApplyServerSideEncryptionByDefault": {
                "SSEAlgorithm": "aws:kms",
                "KMSMasterKeyID": "{kms_key_arn}"
            },
            "BucketKeyEnabled": true
        }
    ]
}'

开始 SageMaker 训练作业

在开始训练作业之前，请准备好数据。

数据格式要求：您的输入数据集必须采用 JSONL 格式，每行都包含一个 Converse 格式的示例，有关更多信息，请参阅准备数据以蒸馏理解模型。

数据集限制

  • 最小提示数：100

  • 最大文件大小：2 GB

  • 行最大长度：180 KB

  • 文件格式：仅限 JSONL

要上传输入数据，请运行以下命令：

aws s3 cp /path/to/input-data/ s3://customer-input-data-bucket/ —recursive

数据增强配方

您可以从配方存储库中获取蒸馏配方。蒸馏配方位于以下目录中：recipes-collection/recipes/fine-tuning/nova。数据增强过程由 YAML 配置文件控制。下面是每个参数的详细说明。所有字段均为必填字段。

参数 说明
name

训练作业的描述性名称。这有助于您在 AWS 管理控制台中识别您的作业。
distillation_data

支持数据蒸馏作业，请勿修改此字段。
maxNumberOfPrompts

数据集中的最大提示数量。
maxResponseLength

每个提示的最大响应长度（词元）。
maxInputFileSizeInGB

输入文件的最大大小（以 GB 为单位)。
maxLineLengthInKB

输入文件中单行的最大大小（以 KB 为单位）。
maxStudentModelFineTuningContextLengthInTokens

学生式模型的最大上下文窗口大小（词元）。该值不得超过学生式模型的容量。您可以根据学生式模型的容量将该值设置为 32k 或 64k。
teacherModelId

设置教师式模型 ID 时，请从以下两个选项中进行选择：

  • 对于 Amazon Nova Premier：IAD 区域为“us.amazon.nova-premier-v1:0”。注意：这仅在 IAD 区域可用。

  • 对于 Amazon Nova Pro：IAD 区域为“us.amazon.nova-pro-v1:0”，ARN 区域为“eu.amazon.nova-pro-v1:0”。
temperature

控制响应随机性（为平衡起见，建议使用 0.7）。
top_p

词元采样的累积概率阈值（建议使用 0.9）。
customer_bucket

用于输入/输出数据的 Amazon S3 存储桶。
kms_key

加密 S3 中输出的 AWS KMS 密钥，Bedrock 批量推理需要使用此密钥来存储推理作业返回的输出。

限制

对于 Nova Premier 教师式模型：由于 Amazon Bedrock 批量推理在 ARN（eu-north-1）区域不可用，因此仅在 IAD 区域（us-east-1）中支持。

最佳实践

数据准备

  • 包括 100 个带标注的高质量示例来指导教师式模型

  • 提交前请移除质量差的标注

  • 遵循文本理解提示最佳实践

  • 在开始蒸馏之前，请使用教师式模型测试提示

模型选择

  • 使用 Nova Pro 作为一般使用案例的教师式模型

  • 考虑使用 Nova Premier 获取专业领域知识

  • 根据延迟和成本要求选择学生式模型

性能优化

  • 从推荐温度（0.7）和 top_p（0.9）开始

  • 在微调之前验证增强的数据质量

  • 按照选择超参数中的指导调整超参数

使用 PySDK 开始作业

以下示例笔记本演示如何运行 SageMaker 训练作业进行蒸馏。有关更多信息，请参阅使用 SageMaker AI 估算器运行训练作业。

import os
import sagemaker,boto3
from sagemaker.pytorch import PyTorch
from sagemaker.inputs import TrainingInput

sagemaker_session = sagemaker.Session()
role = sagemaker.get_execution_role()

# SETUP
job_name = <Your_job_name> # Must be unique for every run

input_s3_uri = <S3 URI to your input dataset> # Must end in .jsonl file
output_s3_uri = <S3 URI to your output bucket> + job_name

image_uri = "708977205387.dkr.ecr.us-east-1.amazonaws.com/nova-distillation-repo:SM-TJ-DISTILL-LATEST" # Do not change
instance_type = "ml.r5.4xlarge" # Recommedation is to use cpu instances 
instance_count = 1 # Must be 1, do not change 
role_arn = <IAM role to execute the job with>
recipe_path = <Local path to your recipe>

# Execution

estimator = PyTorch(
    output_path=output_s3_uri,
    base_job_name=job_name,
    role=role_arn,
    instance_count=instance_count,
    instance_type=instance_type,
    training_recipe=recipe_path,
    max_run=432000,
    sagemaker_session=sagemaker_session,
    image_uri=image_uri,
    subnets= ['subnet-xxxxxxxxxxxxxxxxx','subnet-xxxxxxxxxxxxxxxxx'], # Add subnet groups created in previous steps
    security_group_ids= ['sg-xxxxxxxxxxxxxxxxx'], # Add security group created in previous steps
    disable_profiler=True,
    debugger_hook_config=False
)

trainingInput = TrainingInput(
    s3_data=input_s3_uri,
    distribution='FullyReplicated',
    s3_data_type='Converse'
)

# The keys must be "train".
estimator.fit(inputs={"train": trainingInput})

CloudWatch Logs

日志位于 Amazon CloudWatch 中您 AWS 账户的 /aws/sagemaker/TrainingJobs 日志组下。每台主机都有一个用于训练作业的日志文件。

成功的训练

对于成功的训练作业，您将会在日志末尾看到“训练已完成”的日志消息。

输出存储桶包含以下文件：

  • distillation_data/manifest.json：包含增强数据的位置。您可以使用此数据集启动 Amazon Nova 微调作业。此数据集仅支持 SFT 训练。

    {
  "distillation_data": "s3://customer_escrow_bucket/job_id/distillation_data/"
}

  • distillation_data/sample_training_data.jsonl：此 JSONL 文件包含 50 个增强数据样本可供预览，有助于您确定数据质量。

  • distillation_data/training_config.json：此文件包含推荐用于 Amazon Nova 微调作业的超参数。以下是示例文件：

    {
    "epochCount": 5, 
    "learningRate": 1e-05, 
    "batchSize": 1, 
    "learningRateWarmupSteps": 1
}

验证增强的数据质量

在继续进行微调之前，验证增强数据的质量至关重要：

  1. 查看输出存储桶中的 sample_training_data.jsonl 文件。该文件包含来自增强数据集的 50 个随机样本。

  2. 手动检查这些样本的相关性、连贯性以及与您的使用案例的一致性。

  3. 如果质量不符合您的预期，则可能需要调整输入数据或蒸馏参数，然后重新运行数据增强过程。

数据增强完成后，第二阶段涉及微调学生式模型。有关更多信息，请参阅 监督式微调（SFT）。

在 SFT 训练配方中，您可以传递从上一个作业返回的数据集路径。

data_s3_path: "s3://[escrow-bucket]/[job-name]/distillation_data/training_data.jsonl"

还要覆盖上一步中生成的推荐训练配置。

限制

  • 仅对于此增强数据支持 SFT Nova 微调技术。

  • 不支持多模态蒸馏。

  • 不支持自定义教师式模型。