Benutzerdefinierte Belohnungsfunktionen in Ihrer Umgebung AWS

Benutzerdefinierte Belohnungsfunktionen in Ihrer AWS Umgebung unterstützen nur Single-Turn-RFT. Dadurch werden Modelle anhand von Aufgaben trainiert, bei denen eine einzelne Aufforderung eine einzige Antwort erhält, die unabhängig bewertet wird. Das Modell empfängt eine Aufforderung und generiert eine Antwort, die dann anhand Ihrer Belohnungsfunktion bewertet wird — es findet keine back-and-forth Konversation statt. Dies steht im Gegensatz zu Multi-Turn-RFT, bei der das Modell mehrere Interaktionsrunden mit einer Umgebung oder einem Benutzer durchläuft, bevor es eine endgültige Belohnung erhält.

Übersicht über die Architektur

Die Architektur besteht aus zwei Hauptkomponenten:

Schulung VPC:

Kunden-VPC:

Arbeitsablauf:

Konfiguration des Rezepts

Verwenden Sie dieses Rezept, wenn Ihre Belohnungsfunktion die Verarbeitung innerhalb von 15 Minuten abgeschlossen hat.

## Nova Lite RLVR Training (PEFT)
run:
  name: my-rft-run
  model_type: amazon.nova-2-lite-v1:0:256k
  model_name_or_path: nova-lite-2/prod
  data_s3_path: s3://example-bucket/train.jsonl
  output_s3_path: ""
  
  replicas: 2 # Number of compute instances for training. All supported values: {2, 4, 8, 16}
  generation_replicas: 2 # LLM inference replicas
  rollout_worker_replicas: 1
  
  # Lambda functions for RFT
  reward_lambda_arn: ""

## Training config - essential fields for all services
training_config:
  max_length: 10240
  global_batch_size: 256
  reasoning_effort: high
  
  data:
    shuffle: false
  
  rollout:
    rollout_strategy:
      type: off_policy_async
      age_tolerance: 2
    advantage_strategy:
      number_generation: 8
    generator:
      max_new_tokens: 8192
      set_random_seed: true
      temperature: 1
      top_k: 0
    rewards:
      api_endpoint:
        lambda_arn: ${run.reward_lambda_arn}
        lambda_concurrency_limit: 100 # Lambda should be able to handle (rollout_worker_replicas * 64) requests
  
  # Training configuration
  trainer:
    max_steps: 100
    save_steps: 5
    save_top_k: 5
    
    # RL parameters
    refit_freq: 4
    clip_ratio_high: 0.2
    ent_coeff: 0.001
    loss_scale: 1
  
    optim_config:                    # Optimizer settings
        lr: 7e-7                       # Learning rate
        weight_decay: 0.0              # L2 regularization strength (0.0–1.0)
        adam_beta1: 0.9
        adam_beta2: 0.95
  
    peft:                            # Parameter-efficient fine-tuning (LoRA)
        peft_scheme: "lora"            # Enable LoRA for PEFT
        lora_tuning:
            alpha: 32
            lora_plus_lr_ratio: 64.0     # LoRA+ learning rate scaling factor (0.0–100.0)

## Nova Lite RLVR Training
run:
  name: my-rft-run
  model_type: amazon.nova-2-lite-v1:0:256k
  model_name_or_path: nova-lite-2/prod
  data_s3_path: s3://example-bucket/train.jsonl
  output_s3_path: ""
  
  replicas: 2 # Number of compute instances for training. All supported values: {2, 4, 8, 16}
  generation_replicas: 2 # LLM inference replicas
  rollout_worker_replicas: 1
  
  # Lambda functions for RFT
  reward_lambda_arn: ""

## Training config - essential fields for all services
training_config:
  max_length: 10240
  global_batch_size: 256
  reasoning_effort: high
  
  data:
    shuffle: false
  
  rollout:
    rollout_strategy:
      type: off_policy_async
      age_tolerance: 2
    advantage_strategy:
      number_generation: 8
    generator:
      max_new_tokens: 8192
      set_random_seed: true
      temperature: 1
      top_k: 0
    rewards:
      api_endpoint:
        lambda_arn: ${run.reward_lambda_arn}
        lambda_concurrency_limit: 100 # Lambda should be able to handle (rollout_worker_replicas * 64) requests
  
  # Training configuration
  trainer:
    max_steps: 100
    save_steps: 5
    save_top_k: 5
    
    # RL parameters
    refit_freq: 4
    clip_ratio_high: 0.2
    ent_coeff: 0.001
    loss_scale: 1
  
    optim_config:                    # Optimizer settings
        lr: 7e-7                       # Learning rate
        weight_decay: 0.0              # L2 regularization strength (0.0–1.0)
        adam_beta1: 0.9
        adam_beta2: 0.95
          
    peft:                            # Parameter-efficient fine-tuning (LoRA)
        peft_scheme: "null"            # Disable LoRA for PEFT

Rezeptparameter

Auswahl des Argumentationsmodus

Wählen Sie je nach Komplexität Ihrer Aufgabe aus drei Stufen des Argumentationsaufwands:

Standardverhalten: Wenn ohne Wert angegeben reasoning_effort wird, wird standardmäßig auf. high

Richtlinien:

Implementierung der Belohnungsfunktion

Die Belohnungsfunktion (auch Scorer oder Grader genannt) ist die Kernkomponente, die Modellantworten bewertet und Feedbacksignale für das Training liefert. Sie muss als Lambda-Funktion implementiert werden, die Modellantworten akzeptiert und Belohnungspunkte zurückgibt.

Voraussetzungen

Stellen Sie sicher, dass Ihre Lambda-Funktionen und SQS-Warteschlangen dem erforderlichen Benennungsformat entsprechen und dass Ihre Ausführungsrolle über die erforderlichen Berechtigungen verfügt.

Lambda-ARN-Benennung:

Lambda ARN muss diesem Benennungsformat folgen:

arn:aws:lambda:*:*:function:*SageMaker*

SQS-Benennung (nur für Remote-Belohnungsfunktionen in Ihrer eigenen AWS Umgebung erforderlich):

IAM-Richtlinie für die Ausführungsrolle:

{
  "Action": "lambda:InvokeFunction",
  "Resource": [
    "arn:aws:lambda:*:*:function:*SageMaker*"
  ],
  "Effect": "Allow"
},
{
  "Action": [
    "sqs:DeleteMessage",
    "sqs:ReceiveMessage",
    "sqs:SendMessage"
  ],
  "Resource": [
    "arn:aws:sqs:*:*:*SageMaker*"
  ],
  "Effect": "Allow"
}

VPC-Endpunkt:

Damit der HyperPod Cluster Lambda-Funktionen aufrufen kann, müssen Sie:

Stellen Sie sicher, dass Sie in der mit dem EKS verbundenen VPC einen Lambda-Endpunkt sehen.

Format der Schnittstelle

Ihre Belohnungsfunktion muss Daten im folgenden Format akzeptieren und zurückgeben.

Beispiel für eine Eingabe zum Training:

[{
  "messages": [
    {
      "role": "user",
      "content": "Do you have a dedicated security team?"
    }
  ],
  "metadata": {
    "reference_answer": {
      "compliant": "No",
      "explanation": "As an AI developed by Company, I do not have a traditional security team..."
    },
    "my_key": "sample-001"
  }
}]

Beispielnutzlast für die Belohnung Lambda:

Das System fügt den Zug des Assistenten (generierte Antwort) an den letzten Zug des messages Feldes an und fügt einen eindeutigen Wert hinzu: id

[{
  "id": "123",
  "messages": [
    {
      "role": "user",
      "content": "Do you have a dedicated security team?"
    },
    {
      "role": "assistant",
      "content": "As an AI developed by Amazon, I do not have a dedicated security team..."
    }
  ],
  "metadata": {
    "reference_answer": {
      "compliant": "No",
      "explanation": "As an AI developed by Company, I do not have a traditional security team..."
    },
    "my_key": "sample-001"
  }
}]

Lambda-Vertrag belohnen:

def lambda_handler(event, context):
    return lambda_grader(event)

def lambda_grader(samples: list[dict]) -> list[dict]:
    """
    Args:
        samples: List of dictionaries in OpenAI format
        
        Example input (List of such sample):
        {
            "id": "123",
            "messages": [
                {
                    "role": "user",
                    "content": "Do you have a dedicated security team?"
                },
                {
                    "role": "assistant",
                    "content": "As an AI developed by Company, I do not have a dedicated security team..."
                }
            ],
            "metadata": {
                "reference_answer": {
                    "compliant": "No",
                    "explanation": "As an AI developed by Company, I do not have a traditional security team..."
                },
                "my_key": "sample-001"
            }
        }
    
    Returns:
        List of dictionaries with reward scores:
        {
            "id": str,                              # Same id as input sample
            "aggregate_reward_score": float,        # Overall score for the sample
            "metrics_list": [                       # OPTIONAL: Component scores
                {
                    "name": str,                    # Name of the component score
                    "value": float,                 # Value of the component score
                    "type": str                     # "Reward" or "Metric"
                }
            ]
        }
    """

Eingabefelder:

Ausgabefelder:

Technische Einschränkungen

Bewährte Methoden:

Verwenden von benutzerdefinierten Belohnungsfunktionen

Implementieren Sie benutzerdefinierte Belohnungsfunktionen, wenn Sie aufgabenspezifische Bewertungskriterien haben:

Beispiel-Lambda-Funktion

In diesem Beispiel wird das Eingabeformat validiert und die Modellausgabe mit den Referenzantworten verglichen. Ersetzen Sie die Bewertungslogik durch Ihre tatsächlichen Bewertungskriterien.

from typing import List
import json
from dataclasses import asdict, dataclass


@dataclass
class RewardOutput:
    """Reward service output."""
    id: str
    aggregate_reward_score: float


def lambda_handler(event, context):
    """ Main lambda handler """
    return lambda_grader(event)


def lambda_grader(samples: list[dict]) -> list[dict]:
    """ Core grader function """
    scores: List[RewardOutput] = []
    for sample in samples:

        # Extract components
        idx = sample["id"]
        ground_truth = sample.get("metadata", {}).get("reference_answer")

        if "messages" not in sample:
            print(f"Messages is None/empty for id: {idx}")
            ro = RewardOutput(id=idx, aggregate_reward_score=0.0)
            scores.append(ro)

        if ground_truth is None:
            print(f"No answer found in ground truth for id: {idx}")
            ro = RewardOutput(id=idx, aggregate_reward_score=0.0)
            scores.append(ro)

        # Get model's response (last turn is assistant turn)
        last_message = sample["messages"][-1]
        assert last_message["role"] == "assistant", "Last message must be from assistant"
        model_text = last_message["content"]

        ground_truth_text = _extract_ground_truth_text(ground_truth)
        
        if model_text.lower() == ground_truth_text.lower():
            score = 1.0
        else:
            score = 0.0

        ro = RewardOutput(id=idx, aggregate_reward_score=score)
        scores.append(ro)

    # Convert to dict format for JSON serialization
    return [asdict(score) for score in scores]


def _extract_ground_truth_text(ground_truth) -> str:
    """
    Turn the `ground_truth` field into a plain string.
    """
    if isinstance(ground_truth, str):
        return ground_truth

    if isinstance(ground_truth, dict):
        # Common patterns: { "explanation": "...", "answer": "..." }
        if "explanation" in ground_truth and isinstance(ground_truth["explanation"], str):
            return ground_truth["explanation"]
        if "answer" in ground_truth and isinstance(ground_truth["answer"], str):
            return ground_truth["answer"]
        # Fallback: stringify the whole dict
        return json.dumps(ground_truth, ensure_ascii=False)

    # Fallback: stringify anything else
    return str(ground_truth)

Verwendung von LLM als Richter für Belohnungsfunktionen

Large Language Models (LLMs) werden zunehmend als Richter in Workflows zur Verstärkungsfeinabstimmung (RFT) eingesetzt. Sie liefern automatisierte Belohnungssignale, die als Leitfaden für die Modelloptimierung dienen. Bei diesem Ansatz bewertet ein LLM die Ergebnisse des Modells anhand bestimmter Kriterien — sei es bei der Bewertung von Korrektheit, Qualität, Stiltreue oder semantischer Gleichwertigkeit — und vergibt Belohnungen, die den Reinforcement-Learning-Prozess vorantreiben.

Dies ist besonders nützlich für Aufgaben, bei denen traditionelle Belohnungsfunktionen programmatisch schwer zu definieren sind, z. B. um festzustellen, ob verschiedene Repräsentationen (wie „1/3“, „0,333“ und „ein Drittel“) semantisch gleichwertig sind, oder um nuancierte Eigenschaften wie Kohärenz und Relevanz zu bewerten. Durch die Nutzung von LLM-basierten Juroren als Belohnungsfunktionen können Sie RFT auf komplexe Bereiche skalieren, ohne dass umfangreiche menschliche Anmerkungen erforderlich sind. Dies ermöglicht eine schnelle Iteration und kontinuierliche Verbesserung Ihrer Modelle für verschiedene Anwendungsfälle, die über herkömmliche Ausrichtungsprobleme hinausgehen.

Bevor Sie eines LLM-as-a-Judge in der Produktion einsetzen, überprüfen Sie, ob die Bewertungen des Richtermodells dem menschlichen Urteilsvermögen entsprechen. Dabei werden die Übereinstimmungsquoten zwischen dem LLM-Richter und den menschlichen Gutachtern anhand repräsentativer Stichproben Ihrer Aufgabe gemessen. Im Idealfall wird sichergestellt, dass die Übereinstimmung des LLM-Teams mit Menschen den Übereinstimmungsraten zwischen den Menschen entspricht oder diese übertrifft. Dieser Validierungsschritt hilft dabei, mögliche Verzerrungen zu erkennen, stellt sicher, dass das Belohnungssignal Ihr Modell in die beabsichtigte Richtung lenkt, und schafft Vertrauen, dass der automatisierte Bewertungsprozess Modelle hervorbringt, die Ihren Qualitätskriterien für die Produktion entsprechen.

Using LLM-as-a-Judge ist eine einfache Erweiterung der Verwendung von Lambda-Funktionen für Reinforcement Learning with Verifiable Rewards (RLVR). In der Lambda-Funktion rufen Sie eines der in Amazon Bedrock gehosteten Modelle auf. Um sicherzustellen, dass die Schulung und Bewertung mit dem Judge-Modell gut funktionieren, stellen Sie sicher, dass Ihre Durchsatzquote für das verwendete Amazon Bedrock-Modell ausreichend ist.

Konfigurieren Sie Ihre Lambda-Funktion so, dass das Timeout lang ist, bis zu einem Maximum von 15 Minuten. Die Standardeinstellung für Lambda ist 3 Sekunden, und die Änderung des Timeouts in der Lambda-Konfiguration ist wichtig, um längere Antwortzeiten von Amazon Bedrock-Modellen im Vergleich zu logikbasierten Belohnungsfunktionen zu berücksichtigen. Das Lambda wird auch während des Trainings parallel aufgerufen. Erhöhen Sie also die Parallelität, um den verfügbaren Durchsatz voll auszuschöpfen. Beachten Sie, dass das Parallelitätslimit sowohl in der Lambda-Konfiguration als auch im Trainingsjob-Rezept festgelegt werden muss.

Beispiel für ein Trainingsrezept:

display_name: "Nova Lite V2 LoRA RLVR SMTJ training on GPU"
version: "1.0"
instance_types: ["ml.p5.48xlarge", "ml.p5en.48xlarge"]

run:
  name: <experiment_name>
  model_type: amazon.nova-2-lite-v1:0:256k
  model_name_or_path: "nova-lite-2/prod"
  data_s3_path: s3://<path>/<training_data>.jsonl
  replicas: 4
  reward_lambda_arn: arn:aws:lambda:<region>:<account>:function:<lambda-name>

## SMTJ RFT Training specific configs
training_config:
  max_length: 1200                             # Context window (tokens) for inputs+prompt
  global_batch_size: 64                         # Total samples per optimizer step across all replicas (16/32/64/128/256)
  reasoning_effort: high                        # Enables reasoning mode High / Low / or null for non-reasoning
  test_freq: 10
  
  rollout:                                      # How responses are generated for GRPO/advantage calc
    advantage_strategy:
      number_generation: 4                      # N samples per prompt to estimate advantages (variance vs cost)
    generator:
      max_new_tokens: 1024                     # Cap on tokens generated per sample
      set_random_seed: true                     # Seed generation for reproducibility across runs
      temperature: 1                            # Softmax temperature
      top_k: 1                                  # Sample only from top-K logits
    rewards:
      preset_reward_function: null              # Usage of reward functions built into Verl [exact_match, code_executions, math_answers]
      api_endpoint:
        lambda_arn: arn:aws:lambda:<region>:<account>:function:<lambda-name>
        lambda_concurrency_limit: 12             # Max concurrent Lambda invocations (throughput vs. throttling)
  
  trainer:
    max_steps: 100                                 # Steps to train for. One Step = global_batch_size
    save_steps: 20
    test_freq:10
    
    # RL parameters
    ent_coeff: 0.0                              # A bonus added to the policy loss that rewards higher-output entropy
    kl_loss_coef: 0.0                         # Weight on the KL penalty between the actor (trainable policy) and a frozen reference model
    
    optim_config:                    # Optimizer settings
        lr: 1e-6                       # Learning rate
        weight_decay: 0.0              # L2 regularization strength (0.0–1.0)
        adam_beta1: 0.9
        adam_beta2: 0.95

Beispiel Lambda:

Diese Lambda-Funktion implementiert ein LLM-as-a-Judge Belohnungssystem für die Feinabstimmung von Verstärkungen. Sie verarbeitet Stapel von modellgenerierten Antworten, indem sie Antworten aus gut formatierten Ausgaben extrahiert (nach \boxed{} Notation sucht) und verwendet dann Claude Haiku als Richtermodell, um die semantische Ähnlichkeit zwischen der extrahierten Antwort und der Ground-Truth-Referenzantwort auf einer Skala von 0,0-1,0 zu bewerten. Der Richter vergleicht die Antworten, um festzustellen, ob sie semantisch gleichwertig sind (auch wenn sie unterschiedlich dargestellt werden, wie „1/3“ gegenüber „0,333“), und behandelt Fälle, in denen Antworten unterschiedlich formatiert sein können. Die Funktion beinhaltet Wiederholungslogik für die Drosselung, überprüft die Nachrichtenstruktur und gibt eine Liste von Prämienwerten zurück, die als Trainingssignale im Reinforcement-Learning-Prozess verwendet werden können. Wenn Antworten nicht extrahiert werden können oder die Validierung fehlschlägt, wird ein Wert von 0,0 zugewiesen.

import json
import random

from dataclasses import asdict, dataclass

import re

from typing import Dict, Optional, Any, List
import boto3

from botocore.exceptions import ClientError

from copy import deepcopy
import time

import base64


def extract_solution_nova(solution_str: str) -> Optional[str]:
    """
    Extract solution from Nova-formatted response.
    
    Args:
        solution_str: The solution text from Nova model
        method: "strict" or "flexible" extraction method
    
    Returns:
        Extracted numerical answer or None
    """
    boxed_matches = re.findall(r'\\boxed\{([^}]+)\}', solution_str)
    if boxed_matches:
        final_answer = boxed_matches[-1].replace(",", "").replace("$", "")
        return final_answer
    
    return 0.0

bedrock_runtime = boto3.client('bedrock-runtime', region_name='us-east-1')
JUDGE_MODEL_ID = "global.anthropic.claude-haiku-4-5-20251001-v1:0"
SYSTEM_PROMPT = "You must output ONLY a number between 0.0 and 1.0. No explanations, no text, just the number."

JUDGE_PROMPT_TEMPLATE = """Compare the following two responses and rate how similar they are on a scale of 0.0 to 1.0, where:
- 1.0 means the responses are semantically equivalent (same meaning, even if worded differently)
- 0.5 means the responses are partially similar
- 0.0 means the responses are completely different or contradictory

Response A: {response_a}

Response B: {response_b}

Output ONLY a number between 0.0 and 1.0. No explanations."""

def lambda_graded(id: str, response_a: str, response_b: str, max_retries: int = 50) -> float:
    """Call Bedrock to compare responses and return similarity score."""
    prompt = JUDGE_PROMPT_TEMPLATE.format(response_a=response_a, response_b=response_b)
    print(f"Calling judge: {JUDGE_MODEL_ID}")
    for attempt in range(max_retries):
        try:
            print(f"Attempt: {attempt}")
            response = bedrock_runtime.converse(
                modelId=JUDGE_MODEL_ID,
                messages=[{"role": "user", "content": [{"text": prompt}]}],
                system=[{"text": SYSTEM_PROMPT}],
                inferenceConfig={"temperature": 0.0, "maxTokens": 10}
                )
            print(f"Bedrock call successful: {response}")
            output = response['output']['message']['content'][0]['text'].strip()
            score = float(output)

            print(f"Score parsed: {score}")

            return max(0.0, min(1.0, score))

        except Exception as e:
            if "ThrottlingException" in str(e) and attempt < max_retries - 1:
                time.sleep(2 ** attempt)
                print(f"Throttling {id}")
            else:
                print(f"Bedrock call failed: {e}")
                return 0.0
    print("Max retries reached. Unable to complete the request.")
    return 0.0


def compute_score(id: str, solution_str: str, ground_truth: str, method: str = "strict",
                      format_score: float = 0.0, score: float = 1.0,
                      data_source: str ='dataset_name', extra_info: Optional[dict] = None) -> float:
    """
    The scoring function for PandaLM with Nova format.
    
    Args:
        solution_str: The solution text from Nova model
        ground_truth: JSON string containing the ground truth answer
        method: The method to extract the solution, choices are 'strict' and 'flexible'
        format_score: The score for format compliance
        score: The score for correct answer
        data_source: Should match the data_source in the given dataset
        extra_info: Optional dict with additional fields. Required in function signature.
    
    Returns:
        Score between 0 and 1
    """
    import json
    
    answer = extract_solution_nova(solution_str=solution_str, method=method)
    if answer is None:
        return 0.0
    print(f"Answer: {str(answer)}, Reference: {str(ground_truth)}")
    
    # Clean both answers for comparison
    clean_answer = str(answer)
    clean_ground_truth = str(ground_truth)
    
    score = lambda_graded(id, response_a=clean_answer, response_b=clean_ground_truth)
    print(f"Raw score: {score}")
    return score


@dataclass
class RewardOutput:
    """Reward service."""

    id: str
    aggregate_reward_score: float

def lambda_handler(event, context):

    scores: List[RewardOutput] = []

    samples = event
    print(len(samples))
    for sample in samples:
        # Extract the ground truth key. In the current dataset it's answer
        print("Sample: ", json.dumps(sample, indent=2))
        ground_truth = sample["reference_answer"]
        
        idx = "no id"
        # print(sample)
        if not "id" in sample:
            print(f"ID is None/empty for sample: {sample}")
        else:
            idx = sample["id"]

        ro = RewardOutput(id=idx, aggregate_reward_score=0.0)

        if not "messages" in sample:
            print(f"Messages is None/empty for id: {idx}")
            scores.append(RewardOutput(id="0", aggregate_reward_score=0.0))
            continue
        
        # Extract answer from ground truth dict
        if ground_truth is None:
            print(f"No answer found in ground truth for id: {idx}")
            scores.append(RewardOutput(id="0", aggregate_reward_score=0.0))
            continue
        
        # Get completion from last message (assistant message)
        last_message = sample["messages"][-1]
        completion_text = last_message["content"]
        
        if last_message["role"] not in ["assistant", "nova_assistant"]:
            print(f"Last message is not from assistant for id: {idx}")
            scores.append(RewardOutput(id="0", aggregate_reward_score=0.0))
            continue

        if not "content" in last_message:
            print(f"Completion text is empty for id: {idx}")
            scores.append(RewardOutput(id="0", aggregate_reward_score=0.0))
            continue

        random_score = compute_score(id=id, solution_str=completion_text, ground_truth=ground_truth)
        ro = RewardOutput(id=idx, aggregate_reward_score=random_score)

        print(f"Response for id: {idx} is {ro}")
        scores.append(ro)

    return [asdict(score) for score in scores]