trinity.algorithm.policy_loss_fn package

Submodules

• trinity.algorithm.policy_loss_fn.chord_policy_loss module
  • mu_schedule_function()
  • SFTISLossFn
    • SFTISLossFn.__init__()
    • SFTISLossFn.default_args()
    • SFTISLossFn.select_keys
  • phi_function()
  • SFTPhiLossFn
    • SFTPhiLossFn.__init__()
    • SFTPhiLossFn.default_args()
    • SFTPhiLossFn.select_keys
  • MIXCHORDPolicyLossFn
    • MIXCHORDPolicyLossFn.__init__()
    • MIXCHORDPolicyLossFn.default_args()
    • MIXCHORDPolicyLossFn.select_keys
• trinity.algorithm.policy_loss_fn.dpo_loss module
  • DPOLossFn
    • DPOLossFn.__init__()
    • DPOLossFn.default_args()
    • DPOLossFn.select_keys
• trinity.algorithm.policy_loss_fn.gspo_policy_loss module
  • GSPOLossFn
    • GSPOLossFn.__init__()
    • GSPOLossFn.default_args()
    • GSPOLossFn.select_keys
• trinity.algorithm.policy_loss_fn.mix_policy_loss module
  • MIXPolicyLossFn
    • MIXPolicyLossFn.__init__()
    • MIXPolicyLossFn.default_args()
    • MIXPolicyLossFn.select_keys
• trinity.algorithm.policy_loss_fn.opmd_policy_loss module
  • OPMDPolicyLossFn
    • OPMDPolicyLossFn.__init__()
    • OPMDPolicyLossFn.default_args()
    • OPMDPolicyLossFn.select_keys
• trinity.algorithm.policy_loss_fn.policy_loss_fn module
  • PolicyLossFnMeta
    • PolicyLossFnMeta.ignore_keys
  • PolicyLossFn
    • PolicyLossFn.__init__()
    • PolicyLossFn.default_args()
    • PolicyLossFn.select_keys
• trinity.algorithm.policy_loss_fn.ppo_policy_loss module
  • PPOPolicyLossFn
    • PPOPolicyLossFn.__init__()
    • PPOPolicyLossFn.select_keys
    • PPOPolicyLossFn.default_args()
• trinity.algorithm.policy_loss_fn.sft_loss module
  • SFTLossFn
    • SFTLossFn.__init__()
    • SFTLossFn.default_args()
    • SFTLossFn.select_keys

Module contents

class trinity.algorithm.policy_loss_fn.PolicyLossFn(backend: str = 'verl')

Bases: ABC

Abstract base class for policy loss functions.

This class provides the interface for implementing different policy gradient loss functions while handling parameter name mapping between different training frameworks.

__init__(backend: str = 'verl')

Initialize the policy loss function.

Parameters:

backend – The training framework/backend to use (e.g., “verl”)

abstract classmethod default_args() → Dict

Get default initialization arguments for this loss function.

Returns:

The default init arguments for the policy loss function.

Return type:

Dict

property select_keys

Returns parameter keys mapped to the specific training framework’s naming convention.

class trinity.algorithm.policy_loss_fn.PPOPolicyLossFn(backend: str = 'verl', clip_range: float | None = None, clip_range_low: float | None = None, clip_range_high: float | None = None)

Bases: PolicyLossFn

__init__(backend: str = 'verl', clip_range: float | None = None, clip_range_low: float | None = None, clip_range_high: float | None = None) → None

Initialize the policy loss function.

Parameters:

backend – The training framework/backend to use (e.g., “verl”)

property select_keys

Returns parameter keys mapped to the specific training framework’s naming convention.

classmethod default_args() → Dict

Get default initialization arguments for this loss function.

Returns:

The default init arguments for the policy loss function.

Return type:

Dict

class trinity.algorithm.policy_loss_fn.OPMDPolicyLossFn(backend: str = 'verl', tau: float = 1.0)

Bases: PolicyLossFn

__init__(backend: str = 'verl', tau: float = 1.0) → None

Initialize the policy loss function.

Parameters:

backend – The training framework/backend to use (e.g., “verl”)

classmethod default_args() → Dict

Get default initialization arguments for this loss function.

Returns:

The default init arguments for the policy loss function.

Return type:

Dict

property select_keys

Returns parameter keys mapped to the specific training framework’s naming convention.

class trinity.algorithm.policy_loss_fn.DPOLossFn(backend: str = 'verl', beta: float = 0.1, label_smoothing: float = 0.0)

Bases: PolicyLossFn

__init__(backend: str = 'verl', beta: float = 0.1, label_smoothing: float = 0.0) → None

Initialize the policy loss function.

Parameters:

backend – The training framework/backend to use (e.g., “verl”)

classmethod default_args() → Dict

Get default initialization arguments for this loss function.

Returns:

The default init arguments for the policy loss function.

Return type:

Dict

property select_keys

Returns parameter keys mapped to the specific training framework’s naming convention.

class trinity.algorithm.policy_loss_fn.SFTLossFn(backend: str = 'verl', use_token_level_loss: bool = True)

Bases: PolicyLossFn

__init__(backend: str = 'verl', use_token_level_loss: bool = True) → None

Initialize the policy loss function.

Parameters:

backend – The training framework/backend to use (e.g., “verl”)

classmethod default_args()

Get default initialization arguments for this loss function.

Returns:

The default init arguments for the policy loss function.

Return type:

Dict

property select_keys

Returns parameter keys mapped to the specific training framework’s naming convention.

class trinity.algorithm.policy_loss_fn.MIXPolicyLossFn(backend: str = 'verl', mu: float = 0.1, clip_range: float | None = None, clip_range_low: float | None = None, clip_range_high: float | None = None, use_dynamic_bsz: bool | None = None, ppo_mini_batch_size: int = 1, ppo_micro_batch_size_per_gpu: int = 1, ngpus_trainer: int = 1, train_batch_size_usual: int = 1, train_batch_size_expert: int = 1, use_token_level_loss_in_sft: bool = True)

Bases: PolicyLossFn

Implements a mixed policy loss combining GRPO and SFT losses.

This loss function applies different loss components to data based on whether it comes from an expert or not, as indicated by expert_mask. It combines: - GRPO loss (self.grpo_loss_fn) for non-expert data - SFT loss (self.sft_loss_fn) for expert data - Weighting parameter mu

The per-sample weights are normalized using either experience_per_gpu or gradient_accumulation, depending on whether dynamic batch sizing is enabled, to ensure consistent weighting across different batches of the same type experiences.

__init__(backend: str = 'verl', mu: float = 0.1, clip_range: float | None = None, clip_range_low: float | None = None, clip_range_high: float | None = None, use_dynamic_bsz: bool | None = None, ppo_mini_batch_size: int = 1, ppo_micro_batch_size_per_gpu: int = 1, ngpus_trainer: int = 1, train_batch_size_usual: int = 1, train_batch_size_expert: int = 1, use_token_level_loss_in_sft: bool = True) → None

Initialize the policy loss function.

Parameters:

backend – The training framework/backend to use (e.g., “verl”)

classmethod default_args() → Dict

Get default initialization arguments for this loss function.

Returns:

The default init arguments for the policy loss function.

Return type:

Dict

property select_keys

Returns parameter keys mapped to the specific training framework’s naming convention.

class trinity.algorithm.policy_loss_fn.GSPOLossFn(backend: str = 'verl', clip_range: float | None = None, clip_range_low: float | None = None, clip_range_high: float | None = None)

Bases: PolicyLossFn

__init__(backend: str = 'verl', clip_range: float | None = None, clip_range_low: float | None = None, clip_range_high: float | None = None) → None

Initialize the policy loss function.

Parameters:

backend – The training framework/backend to use (e.g., “verl”)

classmethod default_args() → Dict

Get default initialization arguments for this loss function.

Returns:

The default init arguments for the policy loss function.

Return type:

Dict

property select_keys

Returns parameter keys mapped to the specific training framework’s naming convention.

class trinity.algorithm.policy_loss_fn.MIXCHORDPolicyLossFn(backend: str = 'verl', mu_warmup_steps: int = 0, mu_decay_steps: int = 0, mu_peak: float = 0.1, mu_valley: float = 0.1, enable_phi_function: bool = True, clip_range: float | None = None, clip_range_low: float | None = None, clip_range_high: float | None = None, use_dynamic_bsz: bool | None = None, ppo_mini_batch_size: int = 1, ppo_micro_batch_size_per_gpu: int = 1, ngpus_trainer: int = 1, train_batch_size_usual: int = 1, train_batch_size_expert: int = 1, use_token_level_loss_in_sft: bool = True)

Bases: PolicyLossFn

Implements a mixed policy loss combining GRPO and SFT losses.

This loss function applies different loss components to data based on whether it comes from an expert or not, as indicated by expert_mask. It combines:

• GRPO loss (self.grpo_loss_fn) for non-expert data

• SFT loss (self.sft_loss_fn) for expert data

  the weight of SFT loss is globally controled by mu_schedule function the tokenwise weights are calculated using different SFT loss formulas

The per-sample weights are normalized using either experience_per_gpu or gradient_accumulation, depending on whether dynamic batch sizing is enabled, to ensure consistent weighting across different batches of the same type experiences.

__init__(backend: str = 'verl', mu_warmup_steps: int = 0, mu_decay_steps: int = 0, mu_peak: float = 0.1, mu_valley: float = 0.1, enable_phi_function: bool = True, clip_range: float | None = None, clip_range_low: float | None = None, clip_range_high: float | None = None, use_dynamic_bsz: bool | None = None, ppo_mini_batch_size: int = 1, ppo_micro_batch_size_per_gpu: int = 1, ngpus_trainer: int = 1, train_batch_size_usual: int = 1, train_batch_size_expert: int = 1, use_token_level_loss_in_sft: bool = True) → None

Initialize the policy loss function.

Parameters:

backend – The training framework/backend to use (e.g., “verl”)

classmethod default_args() → Dict

mu_warmup_steps: int, mu_decay_steps: int, mu_peak: float, mu_valley: float

property select_keys

Returns parameter keys mapped to the specific training framework’s naming convention.

class trinity.algorithm.policy_loss_fn.SFTISLossFn(backend: str = 'verl', use_token_level_loss: bool = True)

Bases: PolicyLossFn

SFT loss with importance sampling

__init__(backend: str = 'verl', use_token_level_loss: bool = True) → None

Initialize the policy loss function.

Parameters:

backend – The training framework/backend to use (e.g., “verl”)

classmethod default_args()

Get default initialization arguments for this loss function.

Returns:

The default init arguments for the policy loss function.

Return type:

Dict

property select_keys

Returns parameter keys mapped to the specific training framework’s naming convention.

class trinity.algorithm.policy_loss_fn.SFTPhiLossFn(backend: str = 'verl', use_token_level_loss: bool = True, cutoff_prob: float = 1.0)

Bases: PolicyLossFn

SFT loss with transformed phi function

__init__(backend: str = 'verl', use_token_level_loss: bool = True, cutoff_prob: float = 1.0) → None

Initialize the policy loss function.

Parameters:

backend – The training framework/backend to use (e.g., “verl”)

classmethod default_args()

Get default initialization arguments for this loss function.

Returns:

The default init arguments for the policy loss function.

Return type:

Dict

property select_keys

Returns parameter keys mapped to the specific training framework’s naming convention.