trinity.trainer

Subpackages

• trinity.trainer.verl
  • Submodules
  • trinity.trainer.verl.dp_actor module
    • DataParallelPPOActor
      • DataParallelPPOActor.__init__()
      • DataParallelPPOActor.set_algorithm()
      • DataParallelPPOActor.compute_log_prob()
      • DataParallelPPOActor.update_policy()
  • trinity.trainer.verl.fsdp_workers module
    • ActorRolloutRefWorker
      • ActorRolloutRefWorker.__init__()
      • ActorRolloutRefWorker.init_model()
      • ActorRolloutRefWorker.setup_weight_sync_group()
      • ActorRolloutRefWorker.sync_weight()
      • ActorRolloutRefWorker.set_algorithm()
      • ActorRolloutRefWorker.update_actor()
      • 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()
  • Module contents

Submodules

trinity.trainer.trainer module

Trainer Class

class trinity.trainer.trainer.Trainer(config: Config)

Bases: object

Consume the experience and train the model.

__init__(config: Config) → None

prepare() → None

Prepare the trainer.

train() → str

Train the model.

train_step() → bool

Train one step.

Returns:

Whether to continue training.

Return type:

bool

need_sync() → bool

Whether to sync the model weight.

sync_weight() → None

Sync the model weight.

flush_log(step: int) → None

Flush the log of the current step.

shutdown() → None

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 property train_step_num: int

Get the current training step number.

abstract train_step() → bool

Training.

abstract save_checkpoint() → None

Save the checkpoint.

abstract sync_weight() → None

Sync the model weight.

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)

Initialize distributed PPO trainer with Ray backend.

init_workers()

Initialize distributed training workers using Ray backend.

Creates:

1. Ray resource pools from configuration

2. Worker groups for each role (actor, critic, etc.)

reset_experiences_example_table()

property train_step_num: int

Get the current training step number.

prepare()

Do some preparation before training started.

train_step() → bool

Training.

save_checkpoint() → None

Save the checkpoint.

sync_weight() → None

Sync the model weight.

sft_to_rft() → None

shutdown() → None

Shutdown the engine.

Module contents

class trinity.trainer.Trainer(config: Config)

Bases: object

Consume the experience and train the model.

__init__(config: Config) → None

prepare() → None

Prepare the trainer.

train() → str

Train the model.

train_step() → bool

Train one step.

Returns:

Whether to continue training.

Return type:

bool

need_sync() → bool

Whether to sync the model weight.

sync_weight() → None

Sync the model weight.

flush_log(step: int) → None

Flush the log of the current step.

shutdown() → None

class trinity.trainer.TrainEngineWrapper

Bases: ABC

A wrapper class to wrap various training engines.

abstract prepare() → None

Do some preparation before training started.

abstract property train_step_num: int

Get the current training step number.

abstract train_step() → bool

Training.

abstract save_checkpoint() → None

Save the checkpoint.

abstract sync_weight() → None

Sync the model weight.

abstract shutdown() → None

Shutdown the engine.

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

Get a trainer wrapper.