使用 Nova Forge SDK 启动 Nova 定制化实验

pip install amzn-nova-forge 

rom amzn_nova_forge import ( NovaModelCustomizer, SMTJRuntimeManager, TrainingMethod, EvaluationTask, CSVDatasetLoader, Model, )

# General Configuration
MODEL = Model.NOVA_LITE_2
INSTANCE_TYPE = 'ml.p5.48xlarge'
EXECUTION_ROLE = '<YOUR_EXECUTION_ROLE_ARN>'
TRAIN_INSTANCE_COUNT = 4
EVAL_INSTANCE_COUNT = 1
S3_BUCKET = '<YOUR_S3_BUCKET>'
S3_PREFIX = 'stack-overflow'
EVAL_DATA = './eval.csv' 
# Load data
# Note: 'query' maps to the question, 'response' to the classification label
loader = CSVDatasetLoader(
    query='Body', 
    # Question text column
    response='Y', 
    # Classification label column (HQ, LQ_EDIT, LQ_CLOSE)
    system='system' 
    # System prompt column
)
loader.load(EVAL_DATA)

# Transform to Nova format
loader.transform(method=TrainingMethod.EVALUATION, model=MODEL)
loader.show(n=3)

{
  "query": "<input data>",
  "response": "<output data>",
  "system": "<system prompt>"
}

# Validate data format
loader.validate(method=TrainingMethod.EVALUATION, model=MODEL)

# Save to S3
eval_s3_uri = loader.save_data(
    f"s3://{S3_BUCKET}/{S3_PREFIX}/data/eval.jsonl"
)

# Configure runtime infrastructure
runtime_manager = SMTJRuntimeManager(
    instance_type=INSTANCE_TYPE,
    instance_count=EVAL_INSTANCE_COUNT,
    execution_role=EXECUTION_ROLE
)
# Create baseline evaluator
baseline_customizer = NovaModelCustomizer(
    model=MODEL,
    method=TrainingMethod.EVALUATION,
    infra=runtime_manager,
    data_s3_path=eval_s3_uri,
    output_s3_path=f"s3://{S3_BUCKET}/{S3_PREFIX}/baseline-eval"
)
# Run evaluation
# GEN_QA task provides metrics like ROUGE, BLEU, F1, and Exact Match
baseline_result = baseline_customizer.evaluate(
    job_name="blogpost-baseline",
    eval_task=EvaluationTask.GEN_QA 
    # Use GEN_QA for classification
)

def classification_accuracy(samples):
    """Extract predicted class via substring match and compute accuracy."""
    correct, total, no_pred = 0, 0, 0
    for s in samples:
        gold = s["gold"].strip().upper()
        pred_raw = s["inference"][0] if isinstance(s["inference"], list) else s["inference"]
        pred_cat = extract_category(pred_raw)
        if pred_cat is None:
            no_pred += 1
            continue
        total += 1
        if pred_cat == gold:
            correct += 1
    acc = correct / total if total else 0
    print(f"Classification Accuracy: {correct}/{total} ({acc*100:.1f}%)")
    print(f" No valid prediction: {no_pred}/{total + no_pred}")
    return acc

print("???? Baseline Classification Accuracy (extracted class labels):")
baseline_accuracy = classification_accuracy(baseline_samples)

loader = CSVDatasetLoader(
    question='Body', 
    # Stack Overflow question text
    answer='Y', 
    # Classification label (HQ, LQ_EDIT, LQ_CLOSE)
    system='system'
)
loader.load('sft.csv')
loader.transform(method=TrainingMethod.SFT_LORA, model=Model.NOVA_LITE_2)
loader.show(n=3)

{
  "system": [ {"text": "<system prompt>"} ],
  "messages": [
    { "role": "user", "content": [ {"text": "<input data>"} ] },
    { "role": "assistant", "content": [ {"text": "<output data>"} ] }
  ]
}

loader.validate(method=TrainingMethod.SFT_LORA, model=Model.NOVA_LITE_2)

# Save to S3
train_path = loader.save_data(f"s3://{S3_BUCKET}/{S3_PREFIX}/data/train.jsonl")

runtime = SMTJRuntimeManager(
    instance_type=INSTANCE_TYPE,
    instance_count=TRAIN_INSTANCE_COUNT,
    execution_role=EXECUTION_ROLE
)

customizer = NovaModelCustomizer(
    model=MODEL,
    method=TrainingMethod.SFT_LORA,
    infra=runtime,
    data_s3_path=train_path,
    output_s3_path=f"s3://{S3_BUCKET}/{S3_PREFIX}/sft-output"
)
training_config = {
    "lr": 5e-6, 
    # Learning rate
    "warmup_steps": 17, 
    # Gradual LR ramp-up
    "max_steps": 100, 
    # Total training steps
    "global_batch_size": 64, 
    # Samples per gradient update
    "max_length": 8192, 
    # Maximum sequence length in tokens
}
result = customizer.train(
    job_name="blogpost-sft",
    overrides=training_config
)

result = customizer.train(
    job_name="blogpost-sft",
    overrides=training_config,
    dry_run=True
)

result = customizer.train(
    job_name="blogpost-sft",
    overrides=training_config
)

# Save to a file
result.dump(file_path=".", file_name="training_result.json")
# Load from a file
result = TrainingResult.load("training_result.json")

monitor = CloudWatchLogMonitor.from_job_result(result)
monitor.show_logs(limit=50)

customizer.get_logs(limit=20)

result.get_job_status() 
# Returns (JobStatus.IN_PROGRESS, ...) or (JobStatus.COMPLETED, ...)

# Configure runtime infrastructure
runtime_manager = SMTJRuntimeManager(
    instance_type=INSTANCE_TYPE,
    instance_count=EVAL_INSTANCE_COUNT,
    execution_role=EXECUTION_ROLE
)
# Create baseline evaluator
baseline_customizer = NovaModelCustomizer(
    model=MODEL,
    method=TrainingMethod.EVALUATION,
    infra=runtime_manager,
    data_s3_path=eval_s3_uri,
    output_s3_path=f"s3://{S3_BUCKET}/{S3_PREFIX}/sft-eval"
)
# Run evaluation
baseline_result = baseline_customizer.evaluate(
    job_name="blogpost-eval",
    eval_task=EvaluationTask.GEN_QA, 
    job_result=result, 
    # Automatically derives checkpoint path from training result
)

print("Post-SFT Classification Accuracy (extracted class labels):")
sft_accuracy = classification_accuracy(sft_samples)

"""Binary reward function for classification: +1 correct, -1 wrong.
Simple and clear signal:
- Correct prediction: +1.0
- Wrong prediction: -1.0
""" 
def calculate_reward(prediction: str, ground_truth: str) -> float:
    """ Calculates binary reward """
    extracted = extract_category(prediction) 
    # Extracts category from prediction and normalize it
    truth_norm = normalize_text(ground_truth) 
    # Normalize the groundtruth
    # Correct prediction
    if extracted and extracted == truth_norm:
        return 1.0
    # Wrong prediction
    return -1.0

def lambda_handler(event, context):
    """ Lambda handler with binary rewards. """
    scores: List[RewardOutput] = []
    for sample in event:
        idx = sample.get("id", "no_id")
        ground_truth = sample.get("reference_answer", "")
        prediction = last_message.get("content", "")
        # Calculate binary reward
        reward = calculate_reward(prediction, ground_truth)
        scores.append(RewardOutput(id=idx, aggregate_reward_score=reward))
    return [asdict(score) for score in scores]

RFT_DATA = './rft.csv'
rft_loader = CSVDatasetLoader(
    query='Body',
    response='Y',
    system='system'
)
rft_loader.load(RFT_DATA)
# Transform for RFT
rft_loader.transform(method=TrainingMethod.RFT_LORA, model=MODEL)
rft_loader.validate(method=TrainingMethod.RFT_LORA, model=MODEL)
# Save to S3
rft_s3_uri = rft_loader.save_data(
    f"s3://{S3_BUCKET}/{S3_PREFIX}/data/rft.jsonl"
)

{
  "system": [ { "text": "<system prompt>" } ],
  "messages": [
    { "role": "user", "content": [ { "text": "<input data>" } ] }
  ],
  "reference_answer": "<reference answer>"
  #any other metadata field you use in data loader mapping
}

REWARD_LAMBDA_ARN = "arn:aws:lambda:us-east-1:ACCOUNT:function:classification-reward"
# Configure RFT infrastructure
RFT_INSTANCE_COUNT = 2
rft_runtime = SMTJRuntimeManager(
    instance_type=INSTANCE_TYPE,
    instance_count=RFT_INSTANCE_COUNT,
    execution_role=EXECUTION_ROLE
)
# Create RFT customizer
rft_customizer = NovaModelCustomizer(
    model=MODEL,
    method=TrainingMethod.RFT_LORA,
    infra=rft_runtime,
    data_s3_path=rft_s3_uri,
    output_s3_path=f"s3://{S3_BUCKET}/{S3_PREFIX}/rft-output",
    model_path=sft_checkpoint 
    # Start from SFT checkpoint
)