Proximal policy optimization (PPO)

Actor train or policy model: A supervised fine-tuning (SFT) model that gets fine-tuned and updated every epoch. The updates are made by sampling prompts, generating completions, and updating weights using a clipped-surrogate objective. This limits the per-token log-profitability change so that each policy step is proximal to the previous one, preserving training stability.

Actor generation model: A model that generates prompt completions or responses to be judged by the reward model and critic model. The weights of this model are updated from the actor train or policy model each epoch.

Reward model: A model with frozen weights that's used to score the actor generation model.

Critic model: A model with unfrozen weights that's used to score the actor generation model. This score is often viewed as an estimate of the total reward the actor receives when generating the remaining tokens.

Anchor model: An SFT model with frozen weights that is used to calculate the KL divergence between the actor train model and the base model. The anchor model ensures that the updates to the actor model are not too drastic compared to the base model. Drastic changes can lead to instability or performance degradation.

turns is an array of conversation string arrays that represent the dialog sequence. This line contains system prompts, user messages, and bot responses. User messages typically end with "Bot: " to indicate where the model output begins. For example, [["System prompt"], ["User: Question Bot:"], ["Bot response"]].

turns_to_mask is an array of 0-based indices that identify which turns should not receive gradient updates. The masked turns are typically system prompts and user turns. For example, [0, 1, 3] masks the system prompt and user messages (the first and third messages).

reward_category is a string that identifies what aspects of model performance to evaluate. It's used to select the appropriate reward model category during training. The reward category is available for the following reward categories: default, math, coding, if, rag, and rai.

meta_data is an optional object that contains additional contextual or ground-truth information. This can include identifiers, source information, or conversation context. The structure is flexible based on your dataset needs.

default: A general purpose optimizer for standard conversational tasks and basic interactions. Used for general conversations and discussions, basic writing tasks, simple question answering, and non-specialized knowledge queries.

Here is an example:

{
    "turns": ["Write a summary of climate change"],
    "turns_to_mask": [0],
    "reward_category": "default"
}

math: A specialized optimizer for mathematical computations and numerical reasoning tasks. Used for mathematical problem-solving, arithmetic calculations, algebraic equations, geometric problems, and statistical analysis.

Here is an example:

{
    "turns": ["Calculate the derivative of x²"],
    "turns_to_mask": [0],
    "reward_category": "math"
}

coding: A dedicated category for programming and software development-related queries. Used for code implementation, debugging assistance, algorithm design, technical documentation, and system architecture questions.

Here is an example:

{
    "turns": ["Write a function to check if a string is palindrome"],
    "turns_to_mask": [0],
    "reward_category": "coding"
}

if: A category for tasks that require precise procedural execution and step-by-step guidance. Used for multi-step procedures, sequential instructions, complex task decomposition, and process documentation.

Here is an example:

{
    "turns": ["Provide steps to deploy a web application"],
    "turns_to_mask": [0],
    "reward_category": "if"
}

rag: A reward category for tasks that require answering queries based specifically on retrieved contextual information. Used when responses should be derived directly from provided reference materials, synthesizing factual content without going beyond the scope of retrieved information, ensuring answers are grounded in the supplied context rather than general knowledge.

Here is an example:

{
            "turns": ["The Synthesis Report integrates findings from all six IPCC assessment cycles, revealing that global surface temperature has increased 1.1°C from 1850-1900 to 2011-2020, with human activities unequivocally identified as the cause of this warming. Alarmingly, current policies put the world on track for 3.2°C warming by 2100. The document identifies 5 key climate system "tipping points" approaching and emphasizes that greenhouse gas emissions must decline 43% by 2030 (compared to 2019 levels) to limit warming to 1.5°C. Climate-related risks will escalate with every increment of warming, with loss and damage disproportionately affecting vulnerable populations. Despite some progress, climate adaptation remains uneven with significant gaps, and financial flows continue to fall below levels needed for mitigation goals.",
            "What were the key findings of the latest IPCC climate report?"],
            "turns_to_mask": [0, 0],
            "reward_category": "rag"
            }

rai: A reward category for tasks that require applying responsible AI principles such as fairness, transparency, and ethics. Used for evaluating potential biases in AI systems, ensuring privacy considerations, addressing ethical dilemmas, and promoting inclusive design principles.

Here is an example:

{
            "turns": ["Identify potential bias concerns when developing a loan approval algorithm and suggest mitigation strategies"],
            "turns_to_mask": [0],
            "reward_category": "rai"
            }

Always mask index 0 - System prompts should never receive gradient updates.

Always mask user turns - Prevent the model from learning to generate user inputs.

Pattern consistency - Use identical masking patterns for similar conversation structures, such as (0, 1, 3, 5) for multi-turn dialogues.

Selective training - Mask early bot responses to focus training on improved final responses.

Chain-of-thought preservation - Only mask system and user turns when training on reasoning sequences.

Quality filtering - Mask low-quality assistant responses to prevent performance degradation.

Context optimization - Ensure masked turns don't remove essential context needed for subsequent responses.

Run configuration

Actor train configuration (policy config)

Reward model configuration

Critic model configuration

Anchor model configuration

Actor generation model configuration

Intermediate checkpoints are not saved for evaluation and you can't resume from an intermediate checkpoint. Only the last checkpoint is saved.

Multimodal datasets aren't supported.

Training jobs aren't automatically stopped. You have to stop the job using the HyperPod CLI.

MLflow logging isn't supported.

Critic training metrics are not supported on TensorBoard.

To adjust the hyperparameters, follow the guidance in Selecting hyperparameters.