AWS sagemaker documentation change
Summary
Expanded PPO documentation with detailed component descriptions, training data format requirements, reward categories, configuration parameters, and example recipes. Added sections about masking strategies, KL-divergence controls, and responsible AI (RAI) considerations.
Security assessment
While no specific security vulnerabilities are addressed, the documentation now explicitly includes a 'rai' reward category for responsible AI principles (fairness, transparency, ethics) and emphasizes KL-divergence controls to prevent model instability. These additions help users implement security-aware AI practices but don't fix existing vulnerabilities.
Diff
diff --git a/sagemaker/latest/dg/nova-ppo.md b/sagemaker/latest/dg/nova-ppo.md index 45585a5f4..4ddaec496 100644 --- a//sagemaker/latest/dg/nova-ppo.md +++ b//sagemaker/latest/dg/nova-ppo.md @@ -5 +5 @@ -# Proximal Policy Optimization (PPO) +# Proximal policy optimization (PPO) @@ -7 +7 @@ -Proximal Policy Optimization (PPO) is an advanced technique that employs multiple machine learning models working together to train and improve a language model. The PPO process involves five key components: +Proximal policy optimization (PPO) is the process of using several machine learning models to train and score a model. The following models are part of the PPO process: @@ -9 +9 @@ Proximal Policy Optimization (PPO) is an advanced technique that employs multipl - * The **Actor Train Model** (or policy model) is a supervised fine-tuned model that undergoes continuous updates during each training epoch. These updates are carefully controlled using a clipped-surrogate objective that limits how much the model can change at each step, ensuring training stability by keeping policy updates "proximal" to previous versions. + * **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. @@ -10,0 +11 @@ Proximal Policy Optimization (PPO) is an advanced technique that employs multipl + * **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. @@ -11,0 +13 @@ Proximal Policy Optimization (PPO) is an advanced technique that employs multipl + * **Reward model** : A model with frozen weights that's used to score the actor generation model. @@ -12,0 +15 @@ Proximal Policy Optimization (PPO) is an advanced technique that employs multipl + * **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. @@ -14 +17 @@ Proximal Policy Optimization (PPO) is an advanced technique that employs multipl - * The **Actor Generation Model** produces responses to prompts that are then evaluated by other models in the system. This model's weights are synchronized with the Actor Train Model at the beginning of each epoch. + * **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. @@ -19 +22 @@ Proximal Policy Optimization (PPO) is an advanced technique that employs multipl - * The **Reward Model** has fixed (frozen) weights and assigns scores to the outputs created by the Actor Generation Model, providing feedback on response quality. +The training data must be in JSONL format where each line contains a single JSON object that represents a training example. Here is an example: @@ -21,0 +25,6 @@ Proximal Policy Optimization (PPO) is an advanced technique that employs multipl + { + "turns": ["string", "string", ...], // Required + "turns_to_mask": [integer, integer, ...], // Required + "reward_category": "string", // Required + "meta_data": {} // Optional + } @@ -22,0 +32 @@ Proximal Policy Optimization (PPO) is an advanced technique that employs multipl + * `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"]]`. @@ -24 +34 @@ Proximal Policy Optimization (PPO) is an advanced technique that employs multipl - * The **Critic Model** has trainable weights and evaluates the Actor Generation Model's outputs, estimating the total reward the actor might receive for generating the remaining tokens in a sequence. + * `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). @@ -25,0 +36 @@ Proximal Policy Optimization (PPO) is an advanced technique that employs multipl + * `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`. @@ -26,0 +38 @@ Proximal Policy Optimization (PPO) is an advanced technique that employs multipl + * `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. @@ -29 +40,0 @@ Proximal Policy Optimization (PPO) is an advanced technique that employs multipl - * The **Anchor Model** is a frozen supervised fine-tuned model that helps calculate the Kullback-Leibler (KL) divergence between the Actor Train Model and the original base model. This component prevents the Actor Train Model from deviating too drastically from the base model's behavior, which could cause instability or performance issues. @@ -31,0 +43,345 @@ Proximal Policy Optimization (PPO) is an advanced technique that employs multipl +Here is an example record: + + + { + "turns": ["You are a helpful AI assistant.", + "User: What is ML? Bot:", + "Machine learning is...", "User: Examples? Bot:", + "Email spam filtering is..." + ], + "turns_to_mask": [0, 1, 3], + "reward_category": "default", + "meta_data": { + "messages": [{ + "role": "system", + "content": "You are a helpful AI assistant." + }, + { + "role": "user", + "content": "What is ML?" + }, + { + "role": "assistant", + "content": "Machine learning is..." + }, + { + "role": "user", + "content": "Examples?" + }, + { + "role": "assistant", + "content": "Email spam filtering is..." + } + ] + } + } + +The reward modeling framework implements multi-dimensional optimization across distinct categorical objectives to facilitate robust model convergence. The reward category should be selected based on the task that the model must be optimized for. + +We recommend the following guidelines for selecting the right framework for your tasks: + + * `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" + } + + + + +###### Masking turns + +In training datasets, the `turns_to_mask` parameter is crucial for controlling which conversation turns receive gradient updates during training. This array of indices determines which parts of the dialogue the model should learn to generate versus which parts should be treated as context only. Proper masking ensures the model learns appropriate response patterns while avoiding training on system prompts or user inputs that could degrade performance. + +We recommend the following guidance for masking: + + * **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. + + + + +The key to effective masking is monitoring training metrics and validation performance to identify whether your masking strategy preserves necessary context while focusing gradient updates on the desired model outputs. + +###### Enable KL-divergence loss + +For enabling KL-divergence loss, the anchor server needs to be enabled to compute the divergence of the current policy from the original distribution. The KL loss type needs to be specified, and coefficients need to be a value other than zero. Higher coefficient values help the model not deviate much from the original policy which results in lesser changes to general performance. Lower coefficient values allow larger deviations from previous policy, leading to better performance of target metrics but impacting the general performance. + + + ppo_anchor: + max_length: 8192 + trainer: + num_nodes: ${recipes.run.cm_replicas} + model: + global_batch_size: 32 + + ppo_actor_train: + model: + ######## Use KL in actor loss ######## + kl_loss_type: low_var_kl + kl_loss_coeff: 0.1 + + ######## Use KL in reward model ###### + kl_reward_penalty_coeff: 0.1 + +###### Learning rate + +The learning rate for the critic and policy models can be adjusted, with 3e-6 being the default balanced choice. Higher learning rates typically lead to training instabilities, which can be identified through KL divergence spikes and erratic policy behavior. Lower learning rates may cause convergence issues and slow learning, indicated by stagnant rewards and minimal policy updates. Regular monitoring of KL divergence, reward score, and value loss helps in determining whether to adjust the learning rate during training. + + + ppo_critic: