[Feature]Add tis fall back for ppo_policy_loss #334

garyzhang99 · 2025-10-22T02:53:04Z

Description

The Truncate Large IS feature addresses training instability issues in PPO caused by computation errors between VLLM (used for old_logprob) and transformer (used for logprob) calculations.

Problem Statement

When probabilities are very small, even minor computation errors between VLLM and transformer can lead to extremely large importance sampling ratios. This is particularly problematic when:

Both old_logprob and logprob are small (representing low probabilities)
The advantage is negative
The ratio becomes very large (e.g., tens or hundreds) due to computation discrepancies

In the standard PPO implementation with one-sided clipping (multiply advantage first, then clip and take max), this can cause:

Large negative gradients
Training instability
Issues like repetition in generated text

Solution

The Truncate Large IS feature adds a pre-clipping step that truncates the importance sampling ratio to a configurable range before it's used in loss computation. This is similar to CISPO's approach to handling importance sampling ratios.

Configuration

Parameters

truncate_large_is (bool, default: False): Enable/disable the truncate large IS feature
truncate_is_range_low (float, default: 0.0): Lower bound for IS ratio truncation
truncate_is_range_high (float, default: 2.0): Upper bound for IS ratio truncation

Usage in YAML Config

algorithm:
  algorithm_type: grpo
  ...
  policy_loss_fn_config:
    clip_range: 0.2
    loss_agg_mode: token-mean
    # Enable truncate large IS
    truncate_large_is: true
    truncate_is_range_low: 0.0
    truncate_is_range_high: 2.0

Recommended Settings

For general use with potential VLLM/transformer computation errors:

truncate_large_is: true
truncate_is_range_low: 0.0
truncate_is_range_high: 2.0

Implementation Details

The truncation is applied as follows:

# Compute ratio
ratio = torch.exp(logprob - old_logprob)

# Apply truncation if enabled
if truncate_large_is:
    ratio = torch.clamp(ratio, truncate_is_range_low, truncate_is_range_high)

# Continue with standard PPO loss computation
pg_losses = -advantages * ratio
pg_losses2 = -advantages * torch.clamp(ratio, 1.0 - clip_range_low, 1.0 + clip_range_high)
pg_loss = torch.max(pg_losses, pg_losses2)

This ensures that extremely large (or small) ratios are bounded before they interact with the advantage values.

When to Use

Consider enabling this feature when:

Training exhibits instability (loss spikes, large grad norms)
Generated text shows repetition or degeneration
You observe extremely large IS ratios in training

Checklist

Please check the following items before code is ready to be reviewed.

Code has passed all tests
Docstrings have been added/updated in Google Style
Documentation has been updated
Code is ready for review

gemini-code-assist · 2025-10-22T02:53:22Z

Summary of Changes

Hello @garyzhang99, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a crucial feature to address training instability in Proximal Policy Optimization (PPO) algorithms, particularly when dealing with discrepancies in probability calculations between different model components like VLLM and transformer. By implementing a pre-clipping mechanism for the importance sampling ratio, the feature prevents extremely large or small ratios from causing large negative gradients and subsequent training issues such as repetition in generated text. This configurable truncation provides a robust way to stabilize PPO training, especially in scenarios where minor computational errors can have significant impacts.

Highlights

PPO Policy Loss Stabilization: Implemented a 'Truncate Large IS' feature for PPO policy loss to enhance training stability by addressing computation errors between VLLM and transformer calculations.
Configurable IS Ratio Truncation: Introduced new configurable parameters: 'truncate_large_is' (boolean), 'truncate_is_range_low' (float), and 'truncate_is_range_high' (float) to control the importance sampling ratio truncation bounds.
IS Truncation Metric: Added a new metric, 'is_truncate_frac', to monitor the frequency and impact of importance sampling ratio truncation during training.
New Test Case: Extended the PPO policy loss test suite with a new test case, 'test_ppo_policy_loss_with_truncate_is', to validate the functionality of the IS truncation feature.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request introduces a 'Truncate Large IS' feature to the PPO policy loss function, aimed at improving training stability. The implementation of the feature itself is sound, adding the necessary configuration and logic for clamping the importance sampling ratio. However, a critical issue is that the unit test for this new feature is incomplete, with the core assertions commented out. This means the correctness of the feature is not being verified. I have also provided a couple of suggestions to improve code robustness and clarity in the implementation. The incomplete test must be addressed before this PR can be merged.

tests/algorithm/policy_loss_test.py

trinity/algorithm/policy_loss_fn/ppo_policy_loss.py

garyzhang99 · 2025-10-22T05:53:20Z

/unittest-module-algorithm

github-actions · 2025-10-22T05:54:52Z

Summary

Tests 📝	Passed ✅	Failed ❌	Skipped ⏭️	Other ❓	Flaky 🍂	Duration ⏱️
14	14	0	0	0	0	4ms

Tests

Test Name	Status	Duration
tests/algorithm/advantage_fn_test.py::TestGroupedAdvantageFn::test_batch_level_std_grpo	✅	1ms
tests/algorithm/advantage_fn_test.py::TestGroupedAdvantageFn::test_batch_level_step_wise_grpo_advantage	✅	1ms
tests/algorithm/advantage_fn_test.py::TestGroupedAdvantageFn::test_duplicate_grpo	✅	1ms
tests/algorithm/advantage_fn_test.py::TestGroupedAdvantageFn::test_grpo_advantage	✅	1ms
tests/algorithm/advantage_fn_test.py::TestGroupedAdvantageFn::test_grpo_correct_bias	✅	1ms
tests/algorithm/advantage_fn_test.py::TestGroupedAdvantageFn::test_grpo_reward_std	✅	1ms
tests/algorithm/advantage_fn_test.py::TestGroupedAdvantageFn::test_step_wise_grpo_advantage	✅	1ms
tests/algorithm/policy_loss_test.py::VerlPolicyLossTest::test_dpo_policy_loss	✅	1ms
tests/algorithm/policy_loss_test.py::VerlPolicyLossTest::test_gspo_policy_loss	✅	1ms
tests/algorithm/policy_loss_test.py::VerlPolicyLossTest::test_mix_policy_loss	✅	1ms
tests/algorithm/policy_loss_test.py::VerlPolicyLossTest::test_opmd_policy_loss	✅	1ms
tests/algorithm/policy_loss_test.py::VerlPolicyLossTest::test_ppo_policy_loss	✅	1ms
tests/algorithm/policy_loss_test.py::VerlPolicyLossTest::test_ppo_policy_loss_with_truncate_is	✅	1ms
tests/algorithm/policy_loss_test.py::VerlPolicyLossTest::test_sft_policy_loss	✅	1ms

Github Test Reporter by CTRF 💚

trinity/algorithm/policy_loss_fn/ppo_policy_loss.py

garyzhang99 · 2025-10-27T02:54:52Z

/unittest-module-algorithm

github-actions · 2025-10-27T02:56:07Z

Summary

Tests 📝	Passed ✅	Failed ❌	Skipped ⏭️	Other ❓	Flaky 🍂	Duration ⏱️
14	14	0	0	0	0	4ms

Tests

Test Name	Status	Duration
tests/algorithm/advantage_fn_test.py::TestGroupedAdvantageFn::test_batch_level_std_grpo	✅	1ms
tests/algorithm/advantage_fn_test.py::TestGroupedAdvantageFn::test_batch_level_step_wise_grpo_advantage	✅	1ms
tests/algorithm/advantage_fn_test.py::TestGroupedAdvantageFn::test_duplicate_grpo	✅	1ms
tests/algorithm/advantage_fn_test.py::TestGroupedAdvantageFn::test_grpo_advantage	✅	1ms
tests/algorithm/advantage_fn_test.py::TestGroupedAdvantageFn::test_grpo_correct_bias	✅	1ms
tests/algorithm/advantage_fn_test.py::TestGroupedAdvantageFn::test_grpo_reward_std	✅	1ms
tests/algorithm/advantage_fn_test.py::TestGroupedAdvantageFn::test_step_wise_grpo_advantage	✅	1ms
tests/algorithm/policy_loss_test.py::VerlPolicyLossTest::test_dpo_policy_loss	✅	1ms
tests/algorithm/policy_loss_test.py::VerlPolicyLossTest::test_gspo_policy_loss	✅	1ms
tests/algorithm/policy_loss_test.py::VerlPolicyLossTest::test_mix_policy_loss	✅	1ms
tests/algorithm/policy_loss_test.py::VerlPolicyLossTest::test_opmd_policy_loss	✅	1ms
tests/algorithm/policy_loss_test.py::VerlPolicyLossTest::test_ppo_policy_loss	✅	1ms
tests/algorithm/policy_loss_test.py::VerlPolicyLossTest::test_ppo_policy_loss_with_truncate_is	✅	1ms
tests/algorithm/policy_loss_test.py::VerlPolicyLossTest::test_sft_policy_loss	✅	1ms

Github Test Reporter by CTRF 💚

add tis fall back for ppo_policy_losss

d71878e

fix pre-commit

334ff8f

gemini-code-assist bot reviewed Oct 22, 2025

View reviewed changes

tests/algorithm/policy_loss_test.py Outdated Show resolved Hide resolved

trinity/algorithm/policy_loss_fn/ppo_policy_loss.py Show resolved Hide resolved

trinity/algorithm/policy_loss_fn/ppo_policy_loss.py Show resolved Hide resolved

add correct test values

2db4d99

yanxi-chen reviewed Oct 22, 2025

View reviewed changes

问昊 added 2 commits October 27, 2025 10:42

resolve comments

db021f5

fix style

a0e805d

garyzhang99 mentioned this pull request Oct 27, 2025

[Feature] Truncate based on the sign of advantage after clipping #340

Open

remove unneccsary lines

781e06b

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[Feature]Add tis fall back for ppo_policy_loss #334

[Feature]Add tis fall back for ppo_policy_loss #334

Uh oh!

garyzhang99 commented Oct 22, 2025 •

edited

Loading

Uh oh!

gemini-code-assist bot commented Oct 22, 2025

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

garyzhang99 commented Oct 22, 2025

Uh oh!

github-actions bot commented Oct 22, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

garyzhang99 commented Oct 27, 2025

Uh oh!

github-actions bot commented Oct 27, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

[Feature]Add tis fall back for ppo_policy_loss #334

Are you sure you want to change the base?

[Feature]Add tis fall back for ppo_policy_loss #334

Uh oh!

Conversation

garyzhang99 commented Oct 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Problem Statement

Solution

Configuration

Parameters

Usage in YAML Config

Recommended Settings

Implementation Details

When to Use

Checklist

Uh oh!

gemini-code-assist bot commented Oct 22, 2025

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

garyzhang99 commented Oct 22, 2025

Uh oh!

github-actions bot commented Oct 22, 2025

Summary

Tests

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

garyzhang99 commented Oct 27, 2025

Uh oh!

github-actions bot commented Oct 27, 2025

Summary

Tests

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

garyzhang99 commented Oct 22, 2025 •

edited

Loading