trinity.trainer

Subpackages

• trinity.trainer.verl
  • Submodules
  • trinity.trainer.verl.core_algos module
    • KLController
      • KLController.update()
    • AdaptiveKLController
      • AdaptiveKLController.__init__()
      • AdaptiveKLController.update()
    • FixedKLController
      • FixedKLController.__init__()
      • FixedKLController.update()
    • get_kl_controller()
    • compute_opmd_outcome_advantage()
    • compute_gae_advantage_return()
    • compute_grpo_outcome_advantage()
    • compute_rloo_outcome_advantage()
    • compute_reinforce_plus_plus_outcome_advantage()
    • compute_remax_outcome_advantage()
    • compute_rewards()
    • compute_policy_loss()
    • compute_policy_loss_dpo()
    • compute_policy_loss_pairwise_opmd()
    • compute_policy_loss_opmd()
    • compute_policy_loss_ppo()
    • compute_policy_loss_sft()
    • compute_entropy_loss()
    • compute_value_loss()
    • kl_penalty()
  • trinity.trainer.verl.dp_actor module
  • trinity.trainer.verl.fsdp_workers module
    • create_device_mesh()
    • get_sharding_strategy()
    • ActorRolloutRefWorker
      • ActorRolloutRefWorker.__init__()
      • ActorRolloutRefWorker.init_model()
      • ActorRolloutRefWorker.setup_weight_sync_group()
      • ActorRolloutRefWorker.sync_weight()
      • ActorRolloutRefWorker.set_mode()
      • ActorRolloutRefWorker.update_actor()
      • ActorRolloutRefWorker.generate_sequences()
      • ActorRolloutRefWorker.compute_log_prob()
      • ActorRolloutRefWorker.compute_ref_log_prob()
      • ActorRolloutRefWorker.save_checkpoint()
      • ActorRolloutRefWorker.load_checkpoint()
      • ActorRolloutRefWorker.clear_optimizer_state()
    • CriticWorker
      • CriticWorker.__init__()
      • CriticWorker.init_model()
      • CriticWorker.compute_values()
      • CriticWorker.update_critic()
      • CriticWorker.save_checkpoint()
      • CriticWorker.load_checkpoint()
      • CriticWorker.clear_optimizer_state()
    • RewardModelWorker
      • RewardModelWorker.__init__()
      • RewardModelWorker.init_model()
      • RewardModelWorker.compute_rm_score()
  • trinity.trainer.verl.ray_trainer module
    • Role
      • Role.Actor
      • Role.Rollout
      • Role.ActorRollout
      • Role.Critic
      • Role.RefPolicy
      • Role.RewardModel
      • Role.ActorRolloutRef
    • AdvantageEstimator
      • AdvantageEstimator.GAE
      • AdvantageEstimator.GRPO
      • AdvantageEstimator.REINFORCE_PLUS_PLUS
      • AdvantageEstimator.REMAX
      • AdvantageEstimator.RLOO
    • ResourcePoolManager
      • ResourcePoolManager.resource_pool_spec
      • ResourcePoolManager.mapping
      • ResourcePoolManager.resource_pool_dict
      • ResourcePoolManager.create_resource_pool()
      • ResourcePoolManager.get_resource_pool()
      • ResourcePoolManager.get_n_gpus()
      • ResourcePoolManager.__init__()
    • apply_kl_penalty()
    • compute_response_mask()
    • compute_advantage()
    • compute_advantage_opmd()
    • compute_advantage_ppo()
    • RayPPOTrainer
      • RayPPOTrainer.__init__()
      • RayPPOTrainer.init_workers()
      • RayPPOTrainer.fit()
  • Module contents

Submodules

trinity.trainer.trainer module

Trainer Class This file is modified from verl.trainer.main_ppo.py And is a reproduction code of Jiayi-Pan/TinyZero.

Note that we don’t combine the main with ray_trainer as ray_trainer is used by other main.

class trinity.trainer.trainer.TrainEngineWrapper

Bases: ABC

A wrapper class to wrap various training engines.

abstract prepare() → None

Do some preparation before training started.

abstract train_rft_step(experiences) → Tuple[bool, int]

Train on the RFT data.

abstract train_sft_step(experiences) → Tuple[bool, int]

Train on the SFT data.

abstract train_dpo_step(experiences) → Tuple[bool, int]

Train on the DPO data.

abstract save_checkpoint() → None

Save the checkpoint.

abstract sync_weight() → None

Sync the model weight.

abstract set_mode(algo_type: AlgorithmType) → None

Set training mode.

abstract shutdown() → None

Shutdown the engine.

trinity.trainer.trainer.get_trainer_wrapper(config: Config) → TrainEngineWrapper

Get a trainer wrapper.

trinity.trainer.verl_trainer module

veRL Trainer Class

Modified from verl/trainer/ppo/ray_trainer.py

class trinity.trainer.verl_trainer.VerlPPOTrainerWrapper(global_config: Config)

Bases: RayPPOTrainer, TrainEngineWrapper

A wrapper for verl.trainer.ppo.RayPPOTrainer.

__init__(global_config: Config)

reset_experiences_example_table()

prepare()

Do some preparation before training started.

train_dpo_step(experiences: Experiences) → Tuple[bool, int]

Train on the DPO data.

train_sft_step(experiences: Experiences) → Tuple[bool, int]

Train on the SFT data.

train_rft_step(experiences: Experiences) → Tuple[bool, int]

Train on the RFT data.

save_checkpoint() → None

Save the checkpoint.

sync_weight() → None

Sync the model weight.

set_mode(algorithm_type: AlgorithmType = AlgorithmType.PPO) → None

Set training mode.

sft_to_rft() → None

shutdown() → None

Shutdown the engine.

Module contents

class trinity.trainer.TrainEngineWrapper

Bases: ABC

A wrapper class to wrap various training engines.

abstract prepare() → None

Do some preparation before training started.

abstract train_rft_step(experiences) → Tuple[bool, int]

Train on the RFT data.

abstract train_sft_step(experiences) → Tuple[bool, int]

Train on the SFT data.

abstract train_dpo_step(experiences) → Tuple[bool, int]

Train on the DPO data.

abstract save_checkpoint() → None

Save the checkpoint.

abstract sync_weight() → None

Sync the model weight.

abstract set_mode(algo_type: AlgorithmType) → None

Set training mode.

abstract shutdown() → None

Shutdown the engine.

trinity.trainer.get_trainer_wrapper(config: Config) → TrainEngineWrapper

Get a trainer wrapper.